Skip sum updates once SUM_WITH_OVERFLOW flag is set

cpp/include/cudf/detail/aggregation/device_aggregators.cuh

@@ -144,15 +144,15 @@ struct update_target_element<Source, aggregation::SUM_WITH_OVERFLOW> {
     auto sum_column      = target.child(0);
     auto overflow_column = target.child(1);
 
+    auto overflow_ref = cuda::atomic_ref<bool, cuda::thread_scope_device>{
+      *(overflow_column.data<bool>() + target_index)};
+
+    if (overflow_ref.load(cuda::memory_order_relaxed)) { return; }
+
     auto const source_value = source.element<DeviceType>(source_index);
     auto const old_sum =
       cudf::detail::atomic_add(&sum_column.element<DeviceType>(target_index), source_value);
 
-    // Early exit if overflow is already set
-    auto bool_ref = cuda::atomic_ref<bool, cuda::thread_scope_device>{
-      *(overflow_column.data<bool>() + target_index)};
-    if (bool_ref.load(cuda::memory_order_relaxed)) { return; }
-
     if (cuda::add_overflow<DeviceType>(old_sum, source_value).overflow) {
       cudf::detail::atomic_max(&overflow_column.element<bool>(target_index), true);
     }

cpp/include/cudf/reduction.hpp

@@ -34,7 +34,8 @@ enum class scan_type : bool { INCLUSIVE, EXCLUSIVE };
  *
  * The `SUM_WITH_OVERFLOW` aggregation is a special case that detects integer
  * overflow during summation of signed integer or decimal values and returns a struct
- * containing both the sum result and an overflow flag.
+ * containing both the sum result and an overflow flag. On overflow the sum value
+ * is unspecified; the boolean flag is the source of truth.
  *
  * Only `min` and `max` ops are supported for reduction of non-arithmetic
  * types (e.g. timestamp or string).

cpp/src/reductions/sum_with_overflow.cu

@@ -150,10 +150,9 @@ std::unique_ptr<cudf::scalar> sum_with_overflow_impl(
       overflow_sum_op<DeviceType>{});
   }
 
-  // On overflow, zero the sum value; the boolean flag is the source of truth.
-  auto const overflowed = result.wraps != 0;
+  // On overflow the sum value is unspecified; the boolean flag is the source of truth.
   return make_sum_overflow_struct_scalar<Source>(
-    overflowed ? DeviceType{0} : result.sum, overflowed, true, col.type(), stream, mr);
+    result.sum, result.wraps != 0, true, col.type(), stream, mr);
 }
 
 struct sum_with_overflow_dispatcher {

cpp/tests/groupby/sum_with_overflow_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,6 +14,8 @@
 #include <cudf/column/column_factories.hpp>
 #include <cudf/fixed_point/fixed_point.hpp>
 #include <cudf/groupby.hpp>
+#include <cudf/sorting.hpp>
+#include <cudf/structs/structs_column_view.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
@@ -208,183 +210,103 @@ TYPED_TEST(groupby_sum_with_overflow_test, overflow_detection)
   using K = int32_t;
   using V = TypeParam;
 
-  if constexpr (cudf::is_fixed_point<V>()) {
-    using namespace numeric;
-    using RepType = cudf::device_storage_type_t<V>;
-
-    // Test decimal overflow detection for all decimal types
-    auto constexpr scale = scale_type{0};  // Use scale 0 for simplicity
+  auto check_overflow_flags = [](cudf::column_view const& keys,
+                                 cudf::column_view const& vals,
+                                 cudf::column_view const& expect_keys,
+                                 cudf::column_view const& expect_overflow) {
+    std::vector<cudf::groupby::aggregation_request> requests;
+    requests.emplace_back();
+    requests[0].values = vals;
+    requests[0].aggregations.push_back(
+      cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>());
 
-    // Use type-specific values that will cause overflow for decimal types
-    auto constexpr type_max = cuda::std::numeric_limits<RepType>::max();
-    auto constexpr type_min = cuda::std::numeric_limits<RepType>::min();
+    auto result = cudf::groupby::groupby(cudf::table_view{{keys}}).aggregate(requests);
+    auto const overflow_child =
+      cudf::structs_column_view{result.second[0].results[0]->view()}.get_sliced_child(1);
 
-    cudf::test::fixed_width_column_wrapper<K> keys{1, 2, 3, 4, 1, 2, 2, 1, 3, 3, 2, 4, 4};
+    auto sorted = cudf::sort_by_key(
+      cudf::table_view{{result.first->get_column(0).view(), overflow_child}}, result.first->view());
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted->view().column(0), expect_keys);
+    CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted->view().column(1), expect_overflow);
+  };
 
-    // Create values that will cause overflow for this specific decimal type
-    RepType large_positive  = type_max - 5;  // Close to max
-    RepType small_increment = 10;            // Will cause overflow when added to large_positive
-    RepType large_negative  = type_min + 5;  // Close to min (decimal types are always signed)
-    RepType small_decrement = -10;           // Will cause underflow when added to large_negative
+  cudf::test::fixed_width_column_wrapper<K> keys{1, 2, 3, 4, 1, 2, 2, 1, 3, 3, 2, 4, 4};
+  cudf::test::fixed_width_column_wrapper<K> expect_keys{1, 2, 3, 4};
+  cudf::test::fixed_width_column_wrapper<bool> expect_overflow{true, false, true, true};
 
-    // Use values that fit within the type range for non-overflowing groups
-    RepType small_val1 = 10;
-    RepType small_val2 = 20;
-    RepType small_val3 = 30;
-    RepType small_val4 = 40;
+  if constexpr (cudf::is_fixed_point<V>()) {
+    using RepType        = cudf::device_storage_type_t<V>;
+    auto constexpr scale = scale_type{0};
+
+    auto constexpr type_max       = cuda::std::numeric_limits<RepType>::max();
+    auto constexpr type_min       = cuda::std::numeric_limits<RepType>::min();
+    RepType const large_positive  = type_max - 5;
+    RepType const small_increment = 10;
+    RepType const large_negative  = type_min + 5;
+    RepType const small_decrement = -10;
+    RepType const small_val1      = 10;
+    RepType const small_val2      = 20;
+    RepType const small_val3      = 30;
+    RepType const small_val4      = 40;
 
     cudf::test::fixed_point_column_wrapper<RepType> vals{
-      {large_positive,   // Group 1: Close to max
-       small_val1,       // Group 2: Small value
-       small_val2,       // Group 3: Small value
-       large_negative,   // Group 4: Close to min
-       small_increment,  // Group 1: Will cause positive overflow
-       small_val2,       // Group 2: Small value
-       small_val3,       // Group 2: Small value
-       large_positive,   // Group 1: Close to max (second occurrence)
-       large_positive,   // Group 3: Close to max
-       1,                // Group 3: Small value
-       small_val4,       // Group 2: Small value
-       small_decrement,  // Group 4: Will cause negative overflow
-       large_negative},  // Group 4: Close to min (second occurrence)
+      {large_positive,   // Group 1
+       small_val1,       // Group 2
+       small_val2,       // Group 3
+       large_negative,   // Group 4
+       small_increment,  // Group 1: positive overflow
+       small_val2,       // Group 2
+       small_val3,       // Group 2
+       large_positive,   // Group 1
+       large_positive,   // Group 3
+       1,                // Group 3
+       small_val4,       // Group 2
+       small_decrement,  // Group 4: negative overflow
+       large_negative},  // Group 4
       scale};
 
-    cudf::test::fixed_width_column_wrapper<K> expect_keys{1, 2, 3, 4};
-
-    // Expected sums (with overflow handled by wrapping)
-    auto overflow_sum_1 = static_cast<RepType>(static_cast<RepType>(large_positive) +
-                                               static_cast<RepType>(small_increment) +
-                                               static_cast<RepType>(large_positive));
-    auto normal_sum_2   = static_cast<RepType>(small_val1 + small_val2 + small_val3 + small_val4);
-    auto overflow_sum_3 =
-      static_cast<RepType>(static_cast<RepType>(small_val2) + static_cast<RepType>(large_positive) +
-                           static_cast<RepType>(1));
-    auto overflow_sum_4 = static_cast<RepType>(static_cast<RepType>(large_negative) +
-                                               static_cast<RepType>(small_decrement) +
-                                               static_cast<RepType>(large_negative));
-
-    cudf::test::fixed_point_column_wrapper<RepType> expect_sum_vals{
-      {overflow_sum_1, normal_sum_2, overflow_sum_3, overflow_sum_4}, scale};
-    cudf::test::fixed_width_column_wrapper<bool> expect_overflow_vals{true, false, true, true};
-
-    std::vector<std::unique_ptr<cudf::column>> children;
-    children.push_back(expect_sum_vals.release());
-    children.push_back(expect_overflow_vals.release());
-    auto expect_vals = cudf::create_structs_hierarchy(4, std::move(children), 0, {});
-
-    auto agg = cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>();
-    test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg));
-
-    // Verify that sort-based groupby throws for decimals
-    auto agg2 = cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>();
-    EXPECT_THROW(
-      test_single_agg(
-        keys, vals, expect_keys, *expect_vals, std::move(agg2), force_use_sort_impl::YES),
-      cudf::logic_error);
+    check_overflow_flags(keys, vals, expect_keys, expect_overflow);
+
+    // Adding nth_element forces sort-based groupby, which must throw for SUM_WITH_OVERFLOW.
+    std::vector<cudf::groupby::aggregation_request> sort_requests;
+    sort_requests.emplace_back();
+    sort_requests[0].values = vals;
+    sort_requests[0].aggregations.push_back(
+      cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>());
+    sort_requests[0].aggregations.push_back(
+      cudf::make_nth_element_aggregation<cudf::groupby_aggregation>(0));
+    EXPECT_THROW(cudf::groupby::groupby(cudf::table_view{{keys}}).aggregate(sort_requests),
+                 cudf::logic_error);
   } else {
     using DeviceType = cudf::device_storage_type_t<V>;
 
-    // Use type-specific values that will cause overflow for each integer type
-    auto constexpr type_max = cuda::std::numeric_limits<DeviceType>::max();
-    auto constexpr type_min = cuda::std::numeric_limits<DeviceType>::min();
-
-    cudf::test::fixed_width_column_wrapper<K> keys{1, 2, 3, 4, 1, 2, 2, 1, 3, 3, 2, 4, 4};
-
-    // Create values that will cause overflow for this specific type
-    // Use smaller increments for smaller types to avoid immediate wrapping
-    DeviceType large_positive  = type_max - 5;  // Close to max
-    DeviceType small_increment = 10;            // Will cause overflow when added to large_positive
-    DeviceType large_negative, small_decrement;
-    if constexpr (cuda::std::is_signed_v<DeviceType>) {
-      large_negative  = type_min + 5;  // Close to min (only for signed types)
-      small_decrement = -10;  // Will cause underflow when added to large_negative (signed only)
-    }
-
-    // Use values that fit within the type range for non-overflowing groups
-    DeviceType small_val1 = 10;
-    DeviceType small_val2 = 20;
-    DeviceType small_val3 = 30;
-    DeviceType small_val4 = 40;
-
-    if constexpr (cuda::std::is_signed_v<DeviceType>) {
-      // For signed types: test both positive and negative overflow
-      cudf::test::fixed_width_column_wrapper<V> vals{
-        static_cast<V>(large_positive),   // Group 1: Close to max
-        static_cast<V>(small_val1),       // Group 2: Small value
-        static_cast<V>(small_val2),       // Group 3: Small value
-        static_cast<V>(large_negative),   // Group 4: Close to min
-        static_cast<V>(small_increment),  // Group 1: Will cause positive overflow
-        static_cast<V>(small_val2),       // Group 2: Small value
-        static_cast<V>(small_val3),       // Group 2: Small value
-        static_cast<V>(large_positive),   // Group 1: Close to max (second occurrence)
-        static_cast<V>(large_positive),   // Group 3: Close to max
-        static_cast<V>(1),                // Group 3: Small value
-        static_cast<V>(small_val4),       // Group 2: Small value
-        static_cast<V>(small_decrement),  // Group 4: Will cause negative overflow
-        static_cast<V>(large_negative)};  // Group 4: Close to min (second occurrence)
-
-      cudf::test::fixed_width_column_wrapper<K> expect_keys{1, 2, 3, 4};
-
-      // Expected results: Groups 1, 3, and 4 overflow; Group 2 does not
-      auto sum_col = cudf::test::fixed_width_column_wrapper<V>{
-        static_cast<V>(static_cast<DeviceType>(large_positive) + small_increment +
-                       static_cast<DeviceType>(large_positive)),  // Group 1: overflowed result
-        static_cast<V>(small_val1 + small_val2 + small_val3 + small_val4),  // Group 2: no overflow
-        static_cast<V>(static_cast<DeviceType>(small_val2) +
-                       static_cast<DeviceType>(large_positive) +
-                       static_cast<DeviceType>(1)),  // Group 3: overflowed result
-        static_cast<V>(static_cast<DeviceType>(large_negative) + small_decrement +
-                       static_cast<DeviceType>(large_negative))  // Group 4: overflowed result
-      };
-      auto overflow_col = cudf::test::fixed_width_column_wrapper<bool>{true, false, true, true};
-      std::vector<std::unique_ptr<cudf::column>> children;
-      children.push_back(sum_col.release());
-      children.push_back(overflow_col.release());
-      auto expect_vals = cudf::create_structs_hierarchy(4, std::move(children), 0, {});
-
-      auto agg = cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>();
-      test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg));
-    } else {
-      // For unsigned types: only test positive overflow
-      cudf::test::fixed_width_column_wrapper<V> vals{
-        static_cast<V>(large_positive),   // Group 1: Close to max
-        static_cast<V>(small_val1),       // Group 2: Small value
-        static_cast<V>(small_val2),       // Group 3: Small value
-        static_cast<V>(small_val1),       // Group 4: Small value
-        static_cast<V>(small_increment),  // Group 1: Will cause positive overflow
-        static_cast<V>(small_val2),       // Group 2: Small value
-        static_cast<V>(small_val3),       // Group 2: Small value
-        static_cast<V>(large_positive),   // Group 1: Close to max (second occurrence)
-        static_cast<V>(large_positive),   // Group 3: Close to max
-        static_cast<V>(1),                // Group 3: Small value
-        static_cast<V>(small_val4),       // Group 2: Small value
-        static_cast<V>(small_val2),       // Group 4: Small value
-        static_cast<V>(small_val3)};      // Group 4: Small value
-
-      cudf::test::fixed_width_column_wrapper<K> expect_keys{1, 2, 3, 4};
-
-      // Expected results: Groups 1 and 3 overflow; Groups 2 and 4 do not
-      auto sum_col = cudf::test::fixed_width_column_wrapper<V>{
-        static_cast<V>(static_cast<DeviceType>(large_positive) + small_increment +
-                       static_cast<DeviceType>(large_positive)),  // Group 1: overflowed result
-        static_cast<V>(small_val1 + small_val2 + small_val3 + small_val4),  // Group 2: no overflow
-        static_cast<V>(static_cast<DeviceType>(small_val2) +
-                       static_cast<DeviceType>(large_positive) +
-                       static_cast<DeviceType>(1)),           // Group 3: overflowed result
-        static_cast<V>(small_val1 + small_val2 + small_val3)  // Group 4: no overflow
-      };
-      auto overflow_col = cudf::test::fixed_width_column_wrapper<bool>{true, false, true, false};
-      std::vector<std::unique_ptr<cudf::column>> children;
-      children.push_back(sum_col.release());
-      children.push_back(overflow_col.release());
-      auto expect_vals = cudf::create_structs_hierarchy(4, std::move(children), 0, {});
-
-      auto agg = cudf::make_sum_with_overflow_aggregation<cudf::groupby_aggregation>();
-      test_single_agg(keys, vals, expect_keys, *expect_vals, std::move(agg));
-    }
-
-    // Note: SUM_WITH_OVERFLOW only works with hash groupby, not sort groupby
-  }  // end else block for non-decimal types
+    auto constexpr type_max          = cuda::std::numeric_limits<DeviceType>::max();
+    auto constexpr type_min          = cuda::std::numeric_limits<DeviceType>::min();
+    DeviceType const large_positive  = type_max - 5;
+    DeviceType const small_increment = 10;
+    DeviceType const large_negative  = type_min + 5;
+    DeviceType const small_decrement = -10;
+    DeviceType const small_val1      = 10;
+    DeviceType const small_val2      = 20;
+    DeviceType const small_val3      = 30;
+    DeviceType const small_val4      = 40;
+
+    cudf::test::fixed_width_column_wrapper<V> vals{static_cast<V>(large_positive),
+                                                   static_cast<V>(small_val1),
+                                                   static_cast<V>(small_val2),
+                                                   static_cast<V>(large_negative),
+                                                   static_cast<V>(small_increment),
+                                                   static_cast<V>(small_val2),
+                                                   static_cast<V>(small_val3),
+                                                   static_cast<V>(large_positive),
+                                                   static_cast<V>(large_positive),
+                                                   static_cast<V>(1),
+                                                   static_cast<V>(small_val4),
+                                                   static_cast<V>(small_decrement),
+                                                   static_cast<V>(large_negative)};
+
+    check_overflow_flags(keys, vals, expect_keys, expect_overflow);
+  }
 }
 
 // Test that SUM_WITH_OVERFLOW throws an error for bool type (which is not supported)

java/src/main/java/ai/rapids/cudf/Aggregation.java

@@ -545,7 +545,7 @@ private SumWithOverflowAggregation() {
      * children {sum: same type as input, overflow: BOOL8}. The input may be any
      * signed integer type (INT8/16/32/64) or fixed-point decimal
      * (DECIMAL32/64/128), for both column reductions and hash-based groupby.
-     * On overflow the sum value is zeroed; the boolean flag is the source of
+     * On overflow the sum value is unspecified; the boolean flag is the source of
      * truth. Sort-based groupby, scan, segmented reduce, and rolling are not
      * supported by cudf.
      */

java/src/main/java/ai/rapids/cudf/GroupByAggregation.java

@@ -75,8 +75,9 @@ public static GroupByAggregation sum() {
    * types are signed integers (INT8/16/32/64) and fixed-point decimal
    * (DECIMAL32/64/128). The sum-child has the same type AND scale as the input
    * column -- e.g. a DECIMAL64 input at scale -4 produces a DECIMAL64 sum-child
-   * at scale -4; cudf does not widen or rescale. Only hash-based groupby is
-   * supported; sort-based groupby will throw.
+   * at scale -4; cudf does not widen or rescale. On overflow the sum value is
+   * unspecified; the boolean flag is the source of truth. Only hash-based
+   * groupby is supported; sort-based groupby will throw.
    */
   public static GroupByAggregation sumWithOverflow() {
     return new GroupByAggregation(Aggregation.sumWithOverflow());

java/src/main/java/ai/rapids/cudf/ReductionAggregation.java

@@ -56,7 +56,7 @@ public static ReductionAggregation sum() {
    * Sum reduction that also reports overflow. The result is a struct scalar
    * with children {sum: same type as input, overflow: BOOL8}. The input may
    * be any signed integer type (INT8/16/32/64) or fixed-point decimal
-   * (DECIMAL32/64/128). On overflow the sum value is zeroed; the boolean flag
+   * (DECIMAL32/64/128). On overflow the sum value is unspecified; the boolean flag
    * is the source of truth.
    */
   public static ReductionAggregation sumWithOverflow() {