Refactor replace and sort

cpp/include/legate_dataframe/replace.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,9 +19,19 @@
 #include <legate.h>
 
 #include <legate_dataframe/core/column.hpp>
+#include <legate_dataframe/core/library.hpp>
 
 namespace legate::dataframe {
 
+namespace task {
+class ReplaceNullScalarTask : public Task<ReplaceNullScalarTask, OpCode::ReplaceNullsWithScalar> {
+ public:
+  static void cpu_variant(legate::TaskContext context);
+  static void gpu_variant(legate::TaskContext context);
+};
+
+}  // namespace task
+
 /**
  * @brief Replace nulls in a column with a scalar
  *

cpp/include/legate_dataframe/sort.hpp

@@ -19,9 +19,37 @@
 #include <string>
 #include <vector>
 
+#include <legate_dataframe/core/library.hpp>
 #include <legate_dataframe/core/table.hpp>
 
 namespace legate::dataframe {
+namespace task {
+
+/**
+ * @brief Return points at which to split a dataset.
+ *
+ * @param nvalues The total number of values to split.
+ * @param nsplits the number of splits (and split values as last is included)
+ * @param include_start Whether to include the starting 0.
+ * @returns column selecting containing nsplits indices.
+ */
+std::vector<std::size_t> get_split_ind(TaskContext& ctx,
+                                       std::size_t nvalues,
+                                       int nsplits,
+                                       bool include_start);
+class SortTask : public Task<SortTask, OpCode::Sort> {
+ public:
+  static constexpr auto GPU_VARIANT_OPTIONS = legate::VariantOptions{}
+                                                .with_has_allocations(true)
+                                                .with_concurrent(true)
+                                                .with_elide_device_ctx_sync(true);
+  static constexpr auto CPU_VARIANT_OPTIONS =
+    legate::VariantOptions{}.with_has_allocations(true).with_concurrent(true);
+
+  static void cpu_variant(legate::TaskContext context);
+  static void gpu_variant(legate::TaskContext context);
+};
+}  // namespace task
 
 /**
  * @brief Sort a logical table.

cpp/src/replace.cpp

@@ -18,11 +18,7 @@
 
 #include <legate.h>
 
-#include <cudf/replace.hpp>
-
 #include <arrow/compute/api.h>
-#include <legate_dataframe/core/column.hpp>
-#include <legate_dataframe/core/library.hpp>
 #include <legate_dataframe/core/table.hpp>
 #include <legate_dataframe/core/task_argument.hpp>
 #include <legate_dataframe/core/task_context.hpp>
@@ -31,48 +27,25 @@
 namespace legate::dataframe {
 namespace task {
 
-class ReplaceNullScalarTask : public Task<ReplaceNullScalarTask, OpCode::ReplaceNullsWithScalar> {
- public:
-  static void cpu_variant(legate::TaskContext context)
-  {
-    TaskContext ctx{context};
-
-    const auto input = argument::get_next_input<PhysicalColumn>(ctx);
-    auto scalar_col  = argument::get_next_input<PhysicalColumn>(ctx);
-
-    auto arrow_input = input.arrow_array_view();
-
-    auto scalar = ARROW_RESULT(scalar_col.arrow_array_view()->GetScalar(0));
-    auto output = argument::get_next_output<PhysicalColumn>(ctx);
-
-    auto datum_result =
-      ARROW_RESULT(arrow::compute::CallFunction("coalesce", {arrow_input, scalar}));
-    if (get_prefer_eager_allocations()) {
-      output.copy_into(std::move(datum_result.make_array()));
-    } else {
-      output.move_into(std::move(datum_result.make_array()));
-    }
-  }
-
-  static void gpu_variant(legate::TaskContext context)
-  {
-    TaskContext ctx{context};
+/*static*/ void ReplaceNullScalarTask::cpu_variant(legate::TaskContext context)
+{
+  TaskContext ctx{context};
 
-    const auto input = argument::get_next_input<PhysicalColumn>(ctx);
-    auto scalar_col  = argument::get_next_input<PhysicalColumn>(ctx);
-    auto output      = argument::get_next_output<PhysicalColumn>(ctx);
+  const auto input = argument::get_next_input<PhysicalColumn>(ctx);
+  auto scalar_col  = argument::get_next_input<PhysicalColumn>(ctx);
 
-    auto cudf_scalar = scalar_col.cudf_scalar();
+  auto arrow_input = input.arrow_array_view();
 
-    auto ret = cudf::replace_nulls(input.column_view(), *cudf_scalar, ctx.stream(), ctx.mr());
+  auto scalar = ARROW_RESULT(scalar_col.arrow_array_view()->GetScalar(0));
+  auto output = argument::get_next_output<PhysicalColumn>(ctx);
 
-    if (get_prefer_eager_allocations()) {
-      output.copy_into(std::move(ret));
-    } else {
-      output.move_into(std::move(ret));
-    }
+  auto datum_result = ARROW_RESULT(arrow::compute::CallFunction("coalesce", {arrow_input, scalar}));
+  if (get_prefer_eager_allocations()) {
+    output.copy_into(std::move(datum_result.make_array()));
+  } else {
+    output.move_into(std::move(datum_result.make_array()));
   }
-};
+}
 
 }  // namespace task
 
@@ -85,8 +58,8 @@ LogicalColumn replace_nulls(const LogicalColumn& col, const LogicalColumn& scala
   std::optional<size_t> size{};
   if (get_prefer_eager_allocations()) { size = col.num_rows(); }
   auto ret =
-    LogicalColumn::empty_like(col.cudf_type(), col.nullable() && scalar.nullable(), false, size);
-  if (col.cudf_type() != scalar.cudf_type()) {
+    LogicalColumn::empty_like(col.arrow_type(), col.nullable() && scalar.nullable(), false, size);
+  if (!col.arrow_type()->Equals(scalar.arrow_type())) {
     throw std::invalid_argument("Scalar type does not match column type.");
   }
   if (!scalar.is_scalar()) {

cpp/src/replace.cu

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <stdexcept>
+
+#include <legate.h>
+
+#include <cudf/replace.hpp>
+#include <legate_dataframe/core/table.hpp>
+#include <legate_dataframe/core/task_argument.hpp>
+#include <legate_dataframe/core/task_context.hpp>
+#include <legate_dataframe/replace.hpp>
+
+namespace legate::dataframe::task {
+/*static*/ void ReplaceNullScalarTask::gpu_variant(legate::TaskContext context)
+{
+  TaskContext ctx{context};
+
+  const auto input = argument::get_next_input<PhysicalColumn>(ctx);
+  auto scalar_col  = argument::get_next_input<PhysicalColumn>(ctx);
+  auto output      = argument::get_next_output<PhysicalColumn>(ctx);
+
+  auto cudf_scalar = scalar_col.cudf_scalar();
+
+  auto ret = cudf::replace_nulls(input.column_view(), *cudf_scalar, ctx.stream(), ctx.mr());
+
+  if (get_prefer_eager_allocations()) {
+    output.copy_into(std::move(ret));
+  } else {
+    output.move_into(std::move(ret));
+  }
+}
+
+}  // namespace legate::dataframe::task