Support segmented reductions and null mask reductions

cpp/include/cudf/detail/reduction.cuh

@@ -18,6 +18,7 @@
 
 #include "reduction_operators.cuh"
 
+#include <cudf/column/column_factories.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -26,6 +27,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cub/device/device_reduce.cuh>
+#include <cub/device/device_segmented_reduce.cuh>
 
 #include <thrust/for_each.h>
 #include <thrust/iterator/iterator_traits.h>
@@ -218,6 +220,87 @@ std::unique_ptr<scalar> reduce(InputIterator d_in,
   return std::unique_ptr<scalar>(result);
 }
 
+/**
+ * @brief Compute the specified simple reduction over each of the segments in the
+ * input range of elements.
+ *
+ * @param[in] d_in          the begin iterator to input
+ * @param[in] d_offset      the begin iterator to offset
+ * @param[in] num_items     the number of items
+ * @param[in] num_segments  the number of segments
+ * @param[in] op            the reduction operator
+ * @param[in] stream        CUDA stream used for device memory operations and kernel launches.
+ * @returns   Output scalar in device memory
+ *
+ * @tparam Op               the reduction operator with device binary operator
+ * @tparam InputIterator    the input column iterator
+ * @tparam OutputType       the output type of reduction
+ */
+template <typename Op,
+          typename InputIterator,
+          typename OffsetIterator,
+          typename OutputType = typename thrust::iterator_value<InputIterator>::type,
+          typename std::enable_if_t<is_fixed_width<OutputType>() &&
+                                    not cudf::is_fixed_point<OutputType>()>* = nullptr>
+std::unique_ptr<column> segmented_reduce(InputIterator d_in,
+                                         OffsetIterator d_offset_begin,
+                                         cudf::size_type num_segments,
+                                         op::simple_op<Op> sop,
+                                         rmm::cuda_stream_view stream,
+                                         rmm::mr::device_memory_resource* mr)
+{
+  auto binary_op = sop.get_binary_op();
+  auto identity  = sop.template get_identity<OutputType>();
+  // auto dev_result = rmm::device_scalar<OutputType>{identity, stream, mr};
+  auto dev_result = make_fixed_width_column(
+    data_type{type_to_id<OutputType>()}, num_segments, mask_state::UNALLOCATED, stream, mr);
+  auto dev_result_mview = dev_result->mutable_view();
+
+  // Allocate temporary storage
+  rmm::device_buffer d_temp_storage;
+  size_t temp_storage_bytes = 0;
+  cub::DeviceSegmentedReduce::Reduce(d_temp_storage.data(),
+                                     temp_storage_bytes,
+                                     d_in,
+                                     dev_result_mview.data<OutputType>(),
+                                     num_segments,
+                                     d_offset_begin,
+                                     d_offset_begin + 1,
+                                     binary_op,
+                                     identity,
+                                     stream.value());
+  d_temp_storage = rmm::device_buffer{temp_storage_bytes, stream};
+
+  // Run reduction
+  cub::DeviceSegmentedReduce::Reduce(d_temp_storage.data(),
+                                     temp_storage_bytes,
+                                     d_in,
+                                     dev_result_mview.data<OutputType>(),
+                                     num_segments,
+                                     d_offset_begin,
+                                     d_offset_begin + 1,
+                                     binary_op,
+                                     identity,
+                                     stream.value());
+
+  return dev_result;
+}
+
+template <typename Op,
+          typename InputIterator,
+          typename OffsetIterator,
+          typename OutputType = typename thrust::iterator_value<InputIterator>::type,
+          typename std::enable_if_t<std::is_same_v<OutputType, string_view>>* = nullptr>
+std::unique_ptr<column> segmented_reduce(InputIterator d_in,
+                                         OffsetIterator d_offset_begin,
+                                         cudf::size_type num_segments,
+                                         op::simple_op<Op> sop,
+                                         rmm::cuda_stream_view stream,
+                                         rmm::mr::device_memory_resource* mr)
+{
+  CUDF_FAIL("Not supported.");
+}
+
 }  // namespace detail
 }  // namespace reduction
 }  // namespace cudf

cpp/include/cudf/detail/reduction_functions.hpp

@@ -21,8 +21,13 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <optional>
+#include <variant>
 namespace cudf {
 namespace reduction {
+
+using scalar_or_column_ptr = std::variant<std::unique_ptr<scalar>, std::unique_ptr<column>>;
+
 /**
  * @brief Computes sum of elements in input column
  *
@@ -37,8 +42,9 @@ namespace reduction {
  * @param mr Device memory resource used to allocate the returned scalar's device memory
  * @return Sum as scalar of type `output_dtype`.
  */
-std::unique_ptr<scalar> sum(
+scalar_or_column_ptr sum(
   column_view const& col,
+  std::optional<std::reference_wrapper<column_view const>> offsets,
   data_type const output_dtype,
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
@@ -135,8 +141,9 @@ std::unique_ptr<scalar> all(
  * @param mr Device memory resource used to allocate the returned scalar's device memory
  * @return Product as scalar of type `output_dtype`.
  */
-std::unique_ptr<scalar> product(
+scalar_or_column_ptr product(
   column_view const& col,
+  std::optional<std::reference_wrapper<column_view const>> offsets,
   data_type const output_dtype,
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
@@ -155,8 +162,9 @@ std::unique_ptr<scalar> product(
  * @param mr Device memory resource used to allocate the returned scalar's device memory
  * @return Sum of squares as scalar of type `output_dtype`.
  */
-std::unique_ptr<scalar> sum_of_squares(
+scalar_or_column_ptr sum_of_squares(
   column_view const& col,
+  std::optional<std::reference_wrapper<column_view const>> offsets,
   data_type const output_dtype,
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());

cpp/include/cudf/reduction.hpp

@@ -69,6 +69,13 @@ std::unique_ptr<scalar> reduce(
   data_type output_dtype,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+std::unique_ptr<column> segment_reduce(
+  column_view const& col,
+  column_view const& offsets,
+  std::unique_ptr<aggregation> const& agg,
+  data_type output_dtype,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /**
  * @brief  Computes the scan of a column.
  *

cpp/src/reductions/product.cu

@@ -23,15 +23,18 @@
 namespace cudf {
 namespace reduction {
 
-std::unique_ptr<cudf::scalar> product(column_view const& col,
-                                      cudf::data_type const output_dtype,
-                                      rmm::cuda_stream_view stream,
-                                      rmm::mr::device_memory_resource* mr)
+std::variant<std::unique_ptr<scalar>, std::unique_ptr<column>> product(
+  column_view const& col,
+  std::optional<std::reference_wrapper<column_view const>> offsets,
+  cudf::data_type const output_dtype,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr)
 {
   return cudf::type_dispatcher(
     cudf::is_dictionary(col.type()) ? dictionary_column_view(col).keys().type() : col.type(),
     simple::element_type_dispatcher<cudf::reduction::op::product>{},
     col,
+    offsets,
     output_dtype,
     stream,
     mr);

cpp/src/reductions/reductions.cpp

@@ -26,36 +26,42 @@
 #include <cudf/scalar/scalar_factories.hpp>
 
 #include <cudf/structs/structs_column_view.hpp>
+#include <functional>
 #include <rmm/cuda_stream_view.hpp>
 
 namespace cudf {
 namespace detail {
 struct reduce_dispatch_functor {
   column_view const col;
+  std::optional<std::reference_wrapper<column_view const>> offsets;
   data_type output_dtype;
   rmm::mr::device_memory_resource* mr;
   rmm::cuda_stream_view stream;
 
   reduce_dispatch_functor(column_view const& col,
+                          std::optional<std::reference_wrapper<column_view const>> offsets,
                           data_type output_dtype,
                           rmm::cuda_stream_view stream,
                           rmm::mr::device_memory_resource* mr)
-    : col(col), output_dtype(output_dtype), mr(mr), stream(stream)
+    : col(col), offsets(offsets), output_dtype(output_dtype), mr(mr), stream(stream)
   {
   }
 
   template <aggregation::Kind k>
-  std::unique_ptr<scalar> operator()(std::unique_ptr<aggregation> const& agg)
+  std::variant<std::unique_ptr<scalar>, std::unique_ptr<column>> operator()(
+    std::unique_ptr<aggregation> const& agg)
   {
     switch (k) {
-      case aggregation::SUM: return reduction::sum(col, output_dtype, stream, mr); break;
-      case aggregation::PRODUCT: return reduction::product(col, output_dtype, stream, mr); break;
+      case aggregation::SUM: return reduction::sum(col, offsets, output_dtype, stream, mr); break;
+      case aggregation::PRODUCT:
+        return reduction::product(col, offsets, output_dtype, stream, mr);
+        break;
       case aggregation::MIN: return reduction::min(col, output_dtype, stream, mr); break;
       case aggregation::MAX: return reduction::max(col, output_dtype, stream, mr); break;
       case aggregation::ANY: return reduction::any(col, output_dtype, stream, mr); break;
       case aggregation::ALL: return reduction::all(col, output_dtype, stream, mr); break;
       case aggregation::SUM_OF_SQUARES:
-        return reduction::sum_of_squares(col, output_dtype, stream, mr);
+        return reduction::sum_of_squares(col, offsets, output_dtype, stream, mr);
         break;
       case aggregation::MEAN: return reduction::mean(col, output_dtype, stream, mr); break;
       case aggregation::VARIANCE: {
@@ -107,8 +113,9 @@ struct reduce_dispatch_functor {
   }
 };
 
-std::unique_ptr<scalar> reduce(
+std::variant<std::unique_ptr<scalar>, std::unique_ptr<column>> reduce(
   column_view const& col,
+  std::optional<std::reference_wrapper<column_view const>> offsets,
   std::unique_ptr<aggregation> const& agg,
   data_type output_dtype,
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
@@ -124,7 +131,7 @@ std::unique_ptr<scalar> reduce(
   }
 
   return aggregation_dispatcher(
-    agg->kind, reduce_dispatch_functor{col, output_dtype, stream, mr}, agg);
+    agg->kind, reduce_dispatch_functor{col, offsets, output_dtype, stream, mr}, agg);
 }
 }  // namespace detail
 
@@ -134,7 +141,18 @@ std::unique_ptr<scalar> reduce(column_view const& col,
                                rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::reduce(col, agg, output_dtype, rmm::cuda_stream_default, mr);
+  return std::get<0>(
+    detail::reduce(col, std::nullopt, agg, output_dtype, rmm::cuda_stream_default, mr));
+}
+
+std::unique_ptr<column> segment_reduce(column_view const& col,
+                                       column_view const& offsets,
+                                       std::unique_ptr<aggregation> const& agg,
+                                       data_type output_dtype,
+                                       rmm::mr::device_memory_resource* mr)
+{
+  CUDF_FUNC_RANGE();
+  return std::get<1>(detail::reduce(col, offsets, agg, output_dtype, rmm::cuda_stream_default, mr));
 }
 
 }  // namespace cudf