Implement HOST_UDF aggregation for reduction and groupby

cpp/include/cudf/aggregation.hpp

@@ -16,11 +16,20 @@
 
 #pragma once
 
+#include <cudf/detail/utilities/visitor_overload.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/export.hpp>
+#include <cudf/utilities/span.hpp>
+#include <cudf/utilities/traits.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
 
 #include <functional>
 #include <memory>
+#include <optional>
+#include <unordered_map>
+#include <unordered_set>
+#include <variant>
 #include <vector>
 
 /**
@@ -70,6 +79,11 @@ enum class rank_percentage : int32_t {
   ONE_NORMALIZED    ///< (rank - 1) / (count - 1)
 };
 
+/**
+ *  @brief Enum to describe scan operation type.
+ */
+enum class scan_type : bool { INCLUSIVE, EXCLUSIVE };
+
 /**
  * @brief Abstract base class for specifying the desired aggregation in an
  * `aggregation_request`.
@@ -110,8 +124,9 @@ class aggregation {
     COLLECT_SET,     ///< collect values into a list without duplicate entries
     LEAD,            ///< window function, accesses row at specified offset following current row
     LAG,             ///< window function, accesses row at specified offset preceding current row
-    PTX,             ///< PTX  UDF based reduction
-    CUDA,            ///< CUDA UDF based reduction
+    PTX,             ///< PTX  based UDF aggregation
+    CUDA,            ///< CUDA based UDF aggregation
+    HOST_UDF,        ///< host based UDF aggregation
     MERGE_LISTS,     ///< merge multiple lists values into one list
     MERGE_SETS,      ///< merge multiple lists values into one list then drop duplicate entries
     MERGE_M2,        ///< merge partial values of M2 aggregation,
@@ -120,7 +135,7 @@ class aggregation {
     TDIGEST,         ///< create a tdigest from a set of input values
     MERGE_TDIGEST,   ///< create a tdigest by merging multiple tdigests together
     HISTOGRAM,       ///< compute frequency of each element
-    MERGE_HISTOGRAM  ///< merge partial values of HISTOGRAM aggregation,
+    MERGE_HISTOGRAM  ///< merge partial values of HISTOGRAM aggregation
   };
 
   aggregation() = delete;
@@ -599,6 +614,271 @@ std::unique_ptr<Base> make_udf_aggregation(udf_type type,
                                            std::string const& user_defined_aggregator,
                                            data_type output_type);
 
+/**
+ * @brief The interface for host-based UDF implementation.
+ *
+ * An implementation of host-based UDF needs to be derived from this base class, defining
+ * its own version of the required functions.
+ */
+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     ///< Initial value
+  };
+
+  /**
+   * @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,    ///< Initial value
+    NULL_POLICY,   ///< To control null handling
+    OFFSETS        ///< The offsets defining segments
+  };
+
+  /**
+   * @brief Define the possible data needed for groupby aggregations.
+   *
+   * Note that only sort-based groupby aggregations are supported.
+   */
+  enum class groupby_data_attribute : int32_t {
+    INPUT_VALUES,    ///< The input values column
+    GROUPED_VALUES,  ///< The input values grouped according to the input `keys` for which the
+                     ///< values within each group maintain their original order
+    SORTED_GROUPED_VALUES,  ///< The input values grouped according to the input `keys` and
+                            ///< sorted within each group
+    GROUP_OFFSETS,          ///< The offsets separating groups
+    GROUP_LABELS            ///< Group labels (which is also the same as group indices)
+  };
+
+  /**
+   * @brief Describe possible data that may be needed in the derived class for its operations.
+   *
+   * Such data can be either intermediate data such as sorted values or group labels etc, or the
+   * results of other aggregations.
+   *
+   * Each derived host-based UDF class may need a different set of data. It is inefficient to
+   * evaluate and pass down all these possible data at once from libcudf. A solution for that is,
+   * the derived class can define a subset of data that it needs and libcudf will evaluate
+   * and pass down only data requested from that set.
+   */
+  struct data_attribute {
+    /**
+     * @brief Hold all possible data types for the input of the derived class.
+     */
+    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.
+
+    data_attribute()                 = default;  ///< Default constructor
+    data_attribute(data_attribute&&) = default;  ///< Move constructor
+
+    /**
+     * @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, 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_}
+    {
+    }
+
+    /**
+     * @brief Construct a new data attribute from another aggregation request.
+     * @param value_ An aggregation request
+     */
+    template <typename T,
+              CUDF_ENABLE_IF(std::is_same_v<T, aggregation> ||
+                             std::is_same_v<T, groupby_aggregation>)>
+    data_attribute(std::unique_ptr<T> value_) : value{std::move(value_)}
+    {
+      if constexpr (std::is_same_v<T, aggregation>) {
+        CUDF_EXPECTS(
+          dynamic_cast<groupby_aggregation*>(std::get<std::unique_ptr<T>>(value).get()) != nullptr,
+          "Requesting results from other aggregations is only supported in groupby "
+          "aggregations.");
+      }
+      CUDF_EXPECTS(std::get<std::unique_ptr<aggregation>>(value) != nullptr,
+                   "Invalid aggregation request.");
+    }
+
+    /**
+     * @brief Copy constructor.
+     * @param other The other data attribute to copy from
+     */
+    data_attribute(data_attribute const& other)
+      : value{std::visit(
+          cudf::detail::visitor_overload{
+            [](auto const& val) { return value_type{val}; },
+            [](std::unique_ptr<aggregation> const& val) { return value_type{val->clone()}; }},
+          other.value)}
+    {
+    }
+
+    /**
+     * @brief Hash functor for `data_attribute`.
+     */
+    struct hash {
+      /**
+       * @brief Compute the hash value of a data attribute.
+       * @param attr The data attribute to hash
+       * @return The hash value of the data attribute
+       */
+      std::size_t operator()(data_attribute const& attr) const
+      {
+        auto const& value = attr.value;
+        auto const hash_value =
+          std::visit(cudf::detail::visitor_overload{
+                       [](auto const& val) { return std::hash<int>{}(static_cast<int>(val)); },
+                       [](std::unique_ptr<aggregation> const& val) { return val->do_hash(); }},
+                     value);
+        return std::hash<std::size_t>{}(value.index()) ^ hash_value;
+      }
+    };  // struct hash
+
+    /**
+     * @brief Equality comparison functor for `data_attribute`.
+     */
+    struct equal_to {
+      /**
+       * @brief Check if two data attributes are equal.
+       * @param lhs The left-hand side data attribute
+       * @param rhs The right-hand side data attribute
+       * @return True if the two data attributes are equal
+       */
+      bool operator()(data_attribute const& lhs, data_attribute const& rhs) const
+      {
+        auto const& lhs_val = lhs.value;
+        auto const& rhs_val = rhs.value;
+        if (lhs_val.index() != rhs_val.index()) { return false; }
+        return std::visit(cudf::detail::visitor_overload{
+                            [](auto const& lhs_val, auto const& rhs_val) {
+                              if constexpr (std::is_same_v<decltype(lhs_val), decltype(rhs_val)>) {
+                                return lhs_val == rhs_val;
+                              }
+                              return false;
+                            },
+                            [](std::unique_ptr<aggregation> const& lhs_val,
+                               std::unique_ptr<aggregation> const& rhs_val) {
+                              return lhs_val->is_equal(*rhs_val);
+                            }},
+                          lhs_val,
+                          rhs_val);
+      }
+    };  // struct equal_to
+  };    // struct data_attribute
+
+  /**
+   * @brief Set of attributes for the input data that is needed for computing the aggregation.
+   */
+  using input_data_attributes =
+    std::unordered_set<data_attribute, data_attribute::hash, data_attribute::equal_to>;
+
+  /**
+   * @brief Return a set of attributes for the data that is needed for computing the aggregation.
+   *
+   * If this function is not overridden, all the data attributes (except results from other
+   * aggregations in groupby) are assumed to be needed.
+   *
+   * @return A set of `data_attribute`.
+   */
+  [[nodiscard]] virtual input_data_attributes get_required_data() const { return {}; }
+
+  /**
+   * @brief Hold all possible types of the data that is passed to the derived class for executing
+   * the aggregation.
+   */
+  using input_data_type = std::variant<column_view,
+                                       data_type,
+                                       std::optional<std::reference_wrapper<scalar const>>,
+                                       null_policy,
+                                       device_span<size_type const>>;
+
+  /**
+   * @brief Input to the aggregation, mapping from each data attribute to its actual data.
+   */
+  using host_udf_input = std::
+    unordered_map<data_attribute, input_data_type, data_attribute::hash, data_attribute::equal_to>;
+
+  /**
+   * @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_type = 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 for reduction (if specified)
+   * @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
+   */
+  [[nodiscard]] virtual output_type get_empty_output(std::optional<data_type> output_dtype,
+                                                     rmm::cuda_stream_view stream,
+                                                     rmm::device_async_resource_ref mr) const = 0;
+
+  /**
+   * @brief Perform the main computation for the host-based UDF.
+   *
+   * @param input The input data needed for performing all computation
+   * @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
+   */
+  [[nodiscard]] virtual output_type operator()(host_udf_input const& input,
+                                               rmm::cuda_stream_view stream,
+                                               rmm::device_async_resource_ref mr) const = 0;
+
+  /**
+   * @brief Computes hash value of the derived class's instance.
+   * @return The hash value of the instance
+   */
+  [[nodiscard]] virtual std::size_t do_hash() const = 0;
+
+  /**
+   * @brief Compares two instances of the derived class for equality.
+   * @param other The other derived class's instance to compare with
+   * @return True if the two instances are equal
+   */
+  [[nodiscard]] virtual bool is_equal(host_udf_base const& other) const = 0;
+
+  /**
+   * @brief Clones the instance.
+   *
+   * A class derived from `host_udf_base` should not store too much data such that its instances
+   * remain lightweight for efficient cloning.
+   *
+   * @return A new instance cloned from this
+   */
+  [[nodiscard]] virtual std::unique_ptr<host_udf_base> clone() const = 0;
+};
+
+/**
+ * @brief Factory to create a HOST_UDF aggregation
+ *
+ * @param host_udf An instance of a class derived from `host_udf_base` to perform aggregation
+ * @return A HOST_UDF aggregation object
+ */
+template <typename Base = aggregation>
+std::unique_ptr<Base> make_host_udf_aggregation(std::unique_ptr<host_udf_base> host_udf);
+
 /**
  * @brief Factory to create a MERGE_LISTS aggregation.
  *