BUG: Pass stream to get_element() and others in cudf (#147)

Maybe `get_element()` didn't support it when this was first added (not
sure).

These failures are now very reliable on multi-gpu runs only, though. I
audited them and the tests are passing now locally with multi-gpu.

Signed-off-by: Sebastian Berg <[email protected]>

Author: seberg

cpp/src/core/column.cu

@@ -311,7 +311,7 @@ std::unique_ptr<cudf::scalar> LogicalColumn::get_cudf_scalar(
   if (col->size() != 1) {
     throw std::invalid_argument("only length 1/scalar columns can be converted to scalar.");
   }
-  return std::move(cudf::get_element(col->view(), 0));
+  return std::move(cudf::get_element(col->view(), 0, stream, mr));
 }
 
 namespace task {
@@ -367,7 +367,7 @@ std::unique_ptr<cudf::scalar> PhysicalColumn::cudf_scalar() const
   if (num_rows() != 1) {
     throw std::invalid_argument("can only convert length one columns to scalar.");
   }
-  return cudf::get_element(column_view(), 0);
+  return cudf::get_element(column_view(), 0, ctx_->stream(), ctx_->mr());
 }
 
 void PhysicalColumn::copy_into(std::unique_ptr<cudf::column> column)

cpp/src/sort.cu

@@ -68,8 +68,10 @@ std::unique_ptr<cudf::column> create_column(cudf::size_type num_rows,
 {
   if (num_rows == 0) { return cudf::make_empty_column(cudf::data_type{cudf::type_to_id<T>()}); }
   return cudf::sequence(num_rows,
-                        *cudf::make_fixed_width_scalar(fill_value),
-                        *cudf::make_fixed_width_scalar(int32_t{0}));
+                        *cudf::make_fixed_width_scalar(fill_value, ctx.stream(), ctx.mr()),
+                        *cudf::make_fixed_width_scalar(int32_t{0}, ctx.stream(), ctx.mr()),
+                        ctx.stream(),
+                        ctx.mr());
 }
 
 template <typename T>
@@ -270,15 +272,18 @@ std::vector<cudf::size_type> find_splits_for_distribution(
     ctx, sorted_table, global_split_values->view(), keys_idx, column_order, null_precedence);
 }
 
-static std::unique_ptr<cudf::table> apply_limit(std::unique_ptr<cudf::table> tbl, int64_t limit)
+static std::unique_ptr<cudf::table> apply_limit(TaskContext& ctx,
+                                                std::unique_ptr<cudf::table> tbl,
+                                                int64_t limit)
 {
   if (limit != INT64_MIN && std::abs(limit) < tbl->num_rows()) {
     cudf::size_type cudf_limit = static_cast<cudf::size_type>(limit);
     cudf::table_view slice;
     if (limit < 0) {
-      slice = cudf::slice(tbl->view(), {tbl->num_rows() + cudf_limit, tbl->num_rows()})[0];
+      slice =
+        cudf::slice(tbl->view(), {tbl->num_rows() + cudf_limit, tbl->num_rows()}, ctx.stream())[0];
     } else {
-      slice = cudf::slice(tbl->view(), {0, cudf_limit})[0];
+      slice = cudf::slice(tbl->view(), {0, cudf_limit}, ctx.stream())[0];
     }
     tbl = std::make_unique<cudf::table>(slice);
   }
@@ -324,7 +329,7 @@ static std::unique_ptr<cudf::table> apply_limit(std::unique_ptr<cudf::table> tbl
   auto sorted_table =
     sort_func(cudf_tbl, key, column_order, null_precedence, ctx.stream(), ctx.mr());
 
-  sorted_table = apply_limit(std::move(sorted_table), limit);
+  sorted_table = apply_limit(ctx, std::move(sorted_table), limit);
 
   if (ctx.nranks == 1) {
     output.move_into(sorted_table->release());

cpp/src/strings.cu

@@ -41,7 +41,7 @@ namespace legate::dataframe::task {
   }
 
   std::unique_ptr<cudf::column> ret;
-  auto cudf_pattern = cudf::string_scalar(pattern);
+  auto cudf_pattern = cudf::string_scalar(pattern, true, ctx.stream(), ctx.mr());
 
   if (match_func == "starts_with") {
     ret = cudf::strings::starts_with(input.column_view(), cudf_pattern, ctx.stream(), ctx.mr());