diff --git a/cpp/src/core/column.cu b/cpp/src/core/column.cu index 3a1ee08..ab4027b 100644 --- a/cpp/src/core/column.cu +++ b/cpp/src/core/column.cu @@ -311,7 +311,7 @@ std::unique_ptr 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 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 column) diff --git a/cpp/src/sort.cu b/cpp/src/sort.cu index c73570b..fb21cc4 100644 --- a/cpp/src/sort.cu +++ b/cpp/src/sort.cu @@ -68,8 +68,10 @@ std::unique_ptr create_column(cudf::size_type num_rows, { if (num_rows == 0) { return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); } 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 @@ -270,15 +272,18 @@ std::vector find_splits_for_distribution( ctx, sorted_table, global_split_values->view(), keys_idx, column_order, null_precedence); } -static std::unique_ptr apply_limit(std::unique_ptr tbl, int64_t limit) +static std::unique_ptr apply_limit(TaskContext& ctx, + std::unique_ptr tbl, + int64_t limit) { if (limit != INT64_MIN && std::abs(limit) < tbl->num_rows()) { cudf::size_type cudf_limit = static_cast(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(slice); } @@ -324,7 +329,7 @@ static std::unique_ptr apply_limit(std::unique_ptr 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()); diff --git a/cpp/src/strings.cu b/cpp/src/strings.cu index be7d6ca..31c1b72 100644 --- a/cpp/src/strings.cu +++ b/cpp/src/strings.cu @@ -41,7 +41,7 @@ namespace legate::dataframe::task { } std::unique_ptr 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());