Compute mean aggregations in int128 type

cpp/src/groupby/streaming_groupby/common.cuh

@@ -295,6 +295,7 @@ struct streaming_groupby::impl {
   std::unique_ptr<cuda::device_buffer<size_type>> _key_set;
 
   [[nodiscard]] size_type num_keys() const { return static_cast<size_type>(_key_indices.size()); }
+  [[nodiscard]] bool is_wide_sum(size_type i) const;
   void ensure_not_invalidated() const
   {
     CUDF_EXPECTS(!_invalidated.load(std::memory_order_relaxed),

cpp/src/groupby/streaming_groupby/impl.cu

@@ -10,6 +10,7 @@
 #include "groupby/hash/output_utils.hpp"
 #include "partial_aggregates.cuh"
 
+#include <cudf/aggregation.hpp>
 #include <cudf/column/column_factories.hpp>
 #include <cudf/column/column_stream.hpp>
 #include <cudf/column/column_view.hpp>
@@ -19,6 +20,7 @@
 #include <cudf/detail/aggregation/result_cache.hpp>
 #include <cudf/detail/copy.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/detail/unary.hpp>
 #include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/fixed_point/fixed_point.hpp>
@@ -250,6 +252,11 @@ void streaming_groupby::impl::initialize(table_view const& data, cuda::stream_re
   _initialized = true;
 }
 
+bool streaming_groupby::impl::is_wide_sum(size_type i) const
+{
+  return _agg_kinds[i] == aggregation::SUM && cudf::is_integral_not_bool(_aggregate_input_types[i]);
+}
+
 std::unique_ptr<table> streaming_groupby::impl::make_aggregate_results(
   size_type capacity,
   size_type first_column,
@@ -263,9 +270,9 @@ std::unique_ptr<table> streaming_groupby::impl::make_aggregate_results(
     std::span<int8_t const>{_is_agg_intermediate}.subspan(first_column, num_columns);
   std::vector<std::unique_ptr<column>> schema_columns;
   schema_columns.reserve(num_columns);
-  for (auto const type :
-       std::span<data_type const>{_aggregate_input_types}.subspan(first_column, num_columns)) {
-    schema_columns.push_back(cudf::make_empty_column(type));
+  for (auto i = first_column; i < first_column + num_columns; ++i) {
+    schema_columns.push_back(cudf::make_empty_column(
+      is_wide_sum(i) ? data_type{type_id::DECIMAL128, 0} : _aggregate_input_types[i]));
   }
   auto const schema = table{std::move(schema_columns)};
   auto results      = detail::hash::create_results_table(
@@ -438,6 +445,26 @@ streaming_groupby::impl::finalize_gathered(std::unique_ptr<table> keys,
                                            cuda::stream_ref stream,
                                            cudf::memory_resources mr) const
 {
+  std::unordered_map<size_type, std::unique_ptr<column>> exact_sums;
+  auto gathered_cols = agg_gathered->release();
+  for (size_type i = 0; i < static_cast<size_type>(gathered_cols.size()); ++i) {
+    if (!is_wide_sum(i)) { continue; }
+    auto const input_column = _value_col_indices[i];
+    auto const has_mean =
+      std::any_of(_requests_clone.begin(), _requests_clone.end(), [&](auto const& req) {
+        return req.column_index == input_column && req.aggregation->kind == aggregation::MEAN;
+      });
+    if (has_mean && !exact_sums.contains(input_column)) {
+      exact_sums.emplace(
+        input_column,
+        cudf::detail::cast(
+          gathered_cols[i]->view(), data_type{type_id::FLOAT64}, stream, mr.get_temporary_mr()));
+    }
+    gathered_cols[i] = cudf::detail::cast(
+      gathered_cols[i]->view(), data_type{type_id::INT64}, stream, mr.get_output_mr());
+  }
+  agg_gathered = std::make_unique<table>(std::move(gathered_cols));
+
   // Each input column needs one cache identity shared by all of its simple and compound
   // aggregations. Prefer a result retaining the input type (e.g. SUM for decimal/duration
   // MEAN), even when COUNT was requested first. For promoted numeric results, the compound
@@ -470,6 +497,26 @@ streaming_groupby::impl::finalize_gathered(std::unique_ptr<table> keys,
     ar.aggregations.push_back(std::unique_ptr<groupby_aggregation>{
       dynamic_cast<groupby_aggregation*>(_requests_clone[i].aggregation->clone().release())});
     auto const& agg = *ar.aggregations.front();
+    if (auto const exact_sum = exact_sums.find(_requests_clone[i].column_index);
+        agg.kind == aggregation::MEAN && exact_sum != exact_sums.end() &&
+        !cache.has_result(ar.values, agg)) {
+      // Finalize MEAN from the exact sum in a scratch cache, then share it under the request key.
+      auto const sum_agg   = make_sum_aggregation();
+      auto const count_agg = make_count_aggregation();
+      auto const sum_view  = exact_sum->second->view();
+      cudf::detail::result_cache exact_cache(1);
+      exact_cache.add_result(
+        sum_view,
+        *count_agg,
+        std::make_unique<column>(
+          cache.get_result(ar.values, *count_agg), stream, mr.get_temporary_mr()));
+      exact_cache.add_result(sum_view, *sum_agg, std::move(exact_sum->second));
+      exact_sums.erase(exact_sum);
+      auto const finalizer =
+        detail::hash::hash_compound_agg_finalizer(sum_view, &exact_cache, true, stream, mr);
+      cudf::detail::aggregation_dispatcher(agg.kind, finalizer, agg);
+      cache.add_result(ar.values, agg, exact_cache.release_result(sum_view, agg));
+    }
     if (_has_compound_aggs && !cache.has_result(ar.values, agg)) {
       // Gathered intermediates may lack null masks, so MEAN derives nulls from the counts.
       auto const finalizer =

cpp/src/groupby/streaming_groupby/warp_aggregate.cu

@@ -28,15 +28,24 @@ namespace {
 namespace hash = detail::hash;
 namespace cg   = cooperative_groups;
 
+template <typename T, aggregation::Kind K>
+constexpr bool is_wide_integer_sum = K == aggregation::SUM && cudf::is_integral_not_bool<T>();
+
+template <typename T, aggregation::Kind K>
+using warp_result_t = cuda::std::conditional_t<is_wide_integer_sum<T, K>,
+                                               __int128_t,
+                                               hash::rep_type_t<cudf::detail::target_type_t<T, K>>>;
+
 template <typename T, aggregation::Kind K>
 constexpr bool is_warp_aggregation_supported = [] {
-  if constexpr (cuda::std::same_as<T, numeric::decimal128> && K == aggregation::SUM) {
+  if constexpr ((cuda::std::same_as<T, numeric::decimal128> && K == aggregation::SUM) ||
+                is_wide_integer_sum<T, K>) {
     return true;  // This kind has a dedicated existing 128-bit atomic addition.
   }
   if constexpr (hash::is_reduction_supported<T>(K) &&
                 (K == aggregation::SUM || K == aggregation::SUM_OF_SQUARES ||
                  K == aggregation::PRODUCT || K == aggregation::MIN || K == aggregation::MAX)) {
-    using Result = hash::rep_type_t<cudf::detail::target_type_t<T, K>>;
+    using Result = warp_result_t<T, K>;
     return sizeof(hash::rep_type_t<T>) <= sizeof(int64_t) && sizeof(Result) <= sizeof(int64_t) &&
            cudf::has_atomic_support<Result>();
   }
@@ -107,7 +116,7 @@ __device__ void aggregate_value(hash::rep_type_t<T> input,
 {
   if constexpr (is_warp_aggregation_supported<T, K>) {
     if (column.output == nullptr) { return; }
-    using Result = hash::rep_type_t<cudf::detail::target_type_t<T, K>>;
+    using Result = warp_result_t<T, K>;
     using Op     = cudf::detail::corresponding_operator_t<K>;
     auto value   = static_cast<Result>(input);
     if constexpr (K == aggregation::SUM_OF_SQUARES) { value *= value; }