Implement HOST_UDF aggregation for reduction and segmented reduction

cpp/include/cudf/aggregation/host_udf.hpp

@@ -100,6 +100,26 @@ struct host_udf_base {
   host_udf_base()          = default;
   virtual ~host_udf_base() = default;
 
+  /**
+   * @brief Define the possible data needed for reduction.
+   */
+  enum class reduction_data_attribute : int32_t {
+    INPUT_VALUES,  ///< The input values column.
+    OUTPUT_DTYPE,  ///< Data type for the output result.
+    INIT_VALUE     ///< The initial value for computing reduction.
+  };
+
+  /**
+   * @brief Define the possible data needed for segmented reduction.
+   */
+  enum class segmented_reduction_data_attribute : int32_t {
+    INPUT_VALUES,  ///< The input values column.
+    OUTPUT_DTYPE,  ///< Data type for the output result.
+    INIT_VALUE,    ///< The initial value for computing reduction.
+    NULL_POLICY,   ///< To control null handling (INCLUDE or EXCLUDE nulls).
+    OFFSETS        ///< The offsets defining segments for segmented reduction.
+  };
+
   /**
    * @brief Define the possible data needed for groupby aggregations.
    *
@@ -131,7 +151,10 @@ struct host_udf_base {
     /**
      * @brief Hold all possible data types for the input of the aggregation in the derived class.
      */
-    using value_type = std::variant<groupby_data_attribute, std::unique_ptr<aggregation>>;
+    using value_type = std::variant<reduction_data_attribute,
+                                    segmented_reduction_data_attribute,
+                                    groupby_data_attribute,
+                                    std::unique_ptr<aggregation>>;
     value_type value;  ///< The actual data attribute, wrapped by this struct
                        ///< as a wrapper is needed to define `hash` and `equal_to` functors.
 
@@ -142,7 +165,10 @@ struct host_udf_base {
      * @brief Construct a new data attribute from an aggregation attribute.
      * @param value_ An aggregation attribute
      */
-    template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, groupby_data_attribute>)>
+    template <typename T,
+              CUDF_ENABLE_IF(std::is_same_v<T, reduction_data_attribute> ||
+                             std::is_same_v<T, segmented_reduction_data_attribute> ||
+                             std::is_same_v<T, groupby_data_attribute>)>
     data_attribute(T value_) : value{value_}
     {
     }
@@ -220,7 +246,12 @@ struct host_udf_base {
    * @brief Hold all possible types of the data that is passed to the derived class for executing
    * the aggregation.
    */
-  using input_data_t = std::variant<column_view, size_type, device_span<size_type const>>;
+  using input_data_t = std::variant<column_view,
+                                    data_type,
+                                    std::optional<std::reference_wrapper<scalar const>>,
+                                    null_policy,
+                                    size_type,
+                                    device_span<size_type const>>;
 
   /**
    * @brief Input to the aggregation, mapping from each data attribute to its actual data.
@@ -229,20 +260,18 @@ struct host_udf_base {
     unordered_map<data_attribute, input_data_t, data_attribute::hash, data_attribute::equal_to>;
 
   /**
-   * @brief Output type of the aggregation.
-   *
-   * Currently only a single type is supported as the output of the aggregation, but it will hold
-   * more type in the future when reduction is supported.
+   * @brief Output type of the aggregation. It can be either a scalar (for reduction) or a column
+   * (for segmented reduction or groupby aggregations).
    */
-  using output_t = std::variant<std::unique_ptr<column>>;
+  using output_t = std::variant<std::unique_ptr<scalar>, std::unique_ptr<column>>;
 
   /**
    * @brief Get the output when the input values column is empty.
    *
    * This is called in libcudf when the input values column is empty. In such situations libcudf
    * tries to generate the output directly without unnecessarily evaluating the intermediate data.
    *
-   * @param output_dtype The expected output data type
+   * @param output_dtype The output data type for reduction and segmented reduction results
    * @param stream The CUDA stream to use for any kernel launches
    * @param mr Device memory resource to use for any allocations
    * @return The output result of the aggregation when input values is empty

cpp/include/cudf/detail/aggregation/aggregation.hpp

@@ -967,7 +967,9 @@ class udf_aggregation final : public rolling_aggregation {
 /**
  * @brief Derived class for specifying host-based UDF aggregation.
  */
-class host_udf_aggregation final : public groupby_aggregation {
+class host_udf_aggregation final : public groupby_aggregation,
+                                   public reduce_aggregation,
+                                   public segmented_reduce_aggregation {
  public:
   std::unique_ptr<host_udf_base> udf_ptr;
 

cpp/src/groupby/sort/host_udf_aggregation.cpp

@@ -99,5 +99,9 @@ template CUDF_EXPORT std::unique_ptr<aggregation> make_host_udf_aggregation<aggr
   std::unique_ptr<host_udf_base>);
 template CUDF_EXPORT std::unique_ptr<groupby_aggregation>
   make_host_udf_aggregation<groupby_aggregation>(std::unique_ptr<host_udf_base>);
+template CUDF_EXPORT std::unique_ptr<reduce_aggregation>
+  make_host_udf_aggregation<reduce_aggregation>(std::unique_ptr<host_udf_base>);
+template CUDF_EXPORT std::unique_ptr<segmented_reduce_aggregation>
+  make_host_udf_aggregation<segmented_reduce_aggregation>(std::unique_ptr<host_udf_base>);
 
 }  // namespace cudf

cpp/src/reductions/reductions.cpp

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cudf/aggregation/host_udf.hpp>
 #include <cudf/column/column.hpp>
 #include <cudf/detail/aggregation/aggregation.hpp>
 #include <cudf/detail/copy.hpp>
@@ -144,6 +145,39 @@ struct reduce_dispatch_functor {
         auto td_agg = static_cast<cudf::detail::merge_tdigest_aggregation const&>(agg);
         return tdigest::detail::reduce_merge_tdigest(col, td_agg.max_centroids, stream, mr);
       }
+      case aggregation::HOST_UDF: {
+        auto const& udf_ptr = dynamic_cast<cudf::detail::host_udf_aggregation const&>(agg).udf_ptr;
+        auto const data_attrs = [&]() -> host_udf_base::data_attribute_set_t {
+          if (auto tmp = udf_ptr->get_required_data(); !tmp.empty()) { return tmp; }
+          // Empty attribute set means everything.
+          return {host_udf_base::reduction_data_attribute::INPUT_VALUES,
+                  host_udf_base::reduction_data_attribute::OUTPUT_DTYPE,
+                  host_udf_base::reduction_data_attribute::INIT_VALUE};
+        }();
+
+        // Do not cache udf_input, as the actual input data may change from run to run.
+        host_udf_base::input_map_t udf_input;
+        for (auto const& attr : data_attrs) {
+          CUDF_EXPECTS(std::holds_alternative<host_udf_base::reduction_data_attribute>(attr.value),
+                       "Invalid input data attribute for HOST_UDF reduction.");
+          switch (std::get<host_udf_base::reduction_data_attribute>(attr.value)) {
+            case host_udf_base::reduction_data_attribute::INPUT_VALUES:
+              udf_input.emplace(attr, col);
+              break;
+            case host_udf_base::reduction_data_attribute::OUTPUT_DTYPE:
+              udf_input.emplace(attr, output_dtype);
+              break;
+            case host_udf_base::reduction_data_attribute::INIT_VALUE:
+              udf_input.emplace(attr, init);
+              break;
+            default: CUDF_UNREACHABLE("Invalid input data attribute for HOST_UDF reduction.");
+          }
+        }
+        auto output = (*udf_ptr)(udf_input, stream, mr);
+        CUDF_EXPECTS(std::holds_alternative<std::unique_ptr<scalar>>(output),
+                     "Invalid output type from HOST_UDF reduction.");
+        return std::get<std::unique_ptr<scalar>>(std::move(output));
+      }  // case aggregation::HOST_UDF
       default: CUDF_FAIL("Unsupported reduction operator");
     }
   }
@@ -161,9 +195,11 @@ std::unique_ptr<scalar> reduce(column_view const& col,
                cudf::data_type_error);
   if (init.has_value() && !(agg.kind == aggregation::SUM || agg.kind == aggregation::PRODUCT ||
                             agg.kind == aggregation::MIN || agg.kind == aggregation::MAX ||
-                            agg.kind == aggregation::ANY || agg.kind == aggregation::ALL)) {
+                            agg.kind == aggregation::ANY || agg.kind == aggregation::ALL ||
+                            agg.kind == aggregation::HOST_UDF)) {
     CUDF_FAIL(
-      "Initial value is only supported for SUM, PRODUCT, MIN, MAX, ANY, and ALL aggregation types");
+      "Initial value is only supported for SUM, PRODUCT, MIN, MAX, ANY, ALL, and HOST_UDF "
+      "aggregation types");
   }
 
   // Returns default scalar if input column is empty or all null

cpp/src/reductions/segmented/reductions.cpp

@@ -13,6 +13,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
+#include <cudf/aggregation/host_udf.hpp>
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/aggregation/aggregation.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
@@ -98,6 +100,49 @@ struct segmented_reduce_dispatch_functor {
       }
       case segmented_reduce_aggregation::NUNIQUE:
         return segmented_nunique(col, offsets, null_handling, stream, mr);
+      case aggregation::HOST_UDF: {
+        auto const& udf_ptr = dynamic_cast<cudf::detail::host_udf_aggregation const&>(agg).udf_ptr;
+        auto const data_attrs = [&]() -> host_udf_base::data_attribute_set_t {
+          if (auto tmp = udf_ptr->get_required_data(); !tmp.empty()) { return tmp; }
+          // Empty attribute set means everything.
+          return {host_udf_base::segmented_reduction_data_attribute::INPUT_VALUES,
+                  host_udf_base::segmented_reduction_data_attribute::OUTPUT_DTYPE,
+                  host_udf_base::segmented_reduction_data_attribute::INIT_VALUE,
+                  host_udf_base::segmented_reduction_data_attribute::NULL_POLICY,
+                  host_udf_base::segmented_reduction_data_attribute::OFFSETS};
+        }();
+
+        // Do not cache udf_input, as the actual input data may change from run to run.
+        host_udf_base::input_map_t udf_input;
+        for (auto const& attr : data_attrs) {
+          CUDF_EXPECTS(
+            std::holds_alternative<host_udf_base::segmented_reduction_data_attribute>(attr.value),
+            "Invalid input data attribute for HOST_UDF segmented reduction.");
+          switch (std::get<host_udf_base::segmented_reduction_data_attribute>(attr.value)) {
+            case host_udf_base::segmented_reduction_data_attribute::INPUT_VALUES:
+              udf_input.emplace(attr, col);
+              break;
+            case host_udf_base::segmented_reduction_data_attribute::OUTPUT_DTYPE:
+              udf_input.emplace(attr, output_dtype);
+              break;
+            case host_udf_base::segmented_reduction_data_attribute::INIT_VALUE:
+              udf_input.emplace(attr, init);
+              break;
+            case host_udf_base::segmented_reduction_data_attribute::NULL_POLICY:
+              udf_input.emplace(attr, null_handling);
+              break;
+            case host_udf_base::segmented_reduction_data_attribute::OFFSETS:
+              udf_input.emplace(attr, offsets);
+              break;
+            default:
+              CUDF_UNREACHABLE("Invalid input data attribute for HOST_UDF segmented reduction.");
+          }
+        }
+        auto output = (*udf_ptr)(udf_input, stream, mr);
+        CUDF_EXPECTS(std::holds_alternative<std::unique_ptr<column>>(output),
+                     "Invalid output type from HOST_UDF segmented reduction.");
+        return std::get<std::unique_ptr<column>>(std::move(output));
+      }  // case aggregation::HOST_UDF
       default: CUDF_FAIL("Unsupported aggregation type.");
     }
   }
@@ -117,9 +162,11 @@ std::unique_ptr<column> segmented_reduce(column_view const& segmented_values,
                cudf::data_type_error);
   if (init.has_value() && !(agg.kind == aggregation::SUM || agg.kind == aggregation::PRODUCT ||
                             agg.kind == aggregation::MIN || agg.kind == aggregation::MAX ||
-                            agg.kind == aggregation::ANY || agg.kind == aggregation::ALL)) {
+                            agg.kind == aggregation::ANY || agg.kind == aggregation::ALL ||
+                            agg.kind == aggregation::HOST_UDF)) {
     CUDF_FAIL(
-      "Initial value is only supported for SUM, PRODUCT, MIN, MAX, ANY, and ALL aggregation types");
+      "Initial value is only supported for SUM, PRODUCT, MIN, MAX, ANY, ALL, and HOST_UDF "
+      "aggregation types");
   }
 
   if (segmented_values.is_empty() && offsets.empty()) {