Fix groupby reductions that perform operations on source type instead…

… of target type (#10250)

In groupby reductions, some operations such as SUM on integers have the output type different from the source type. For example, target type is int64 while the source type is int32. This is to prevent overflow of the computation result.

This fixes a bug in groupby reductions that perform computation using the data in source type instead of target type.

Closes #10246.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Mike Wilson (https://github.com/hyperbolic2346)
  - Conor Hoekstra (https://github.com/codereport)

URL: #10250

cpp/src/groupby/sort/group_single_pass_reduction_util.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -74,12 +74,16 @@ struct element_arg_minmax_fn {
 /**
  * @brief Value accessor for column which supports dictionary column too.
  *
+ * This is similar to `value_accessor` in `column_device_view.cuh` but with support of dictionary
+ * type.
+ *
  * @tparam T Type of the underlying column. For dictionary column, type of the key column.
  */
 template <typename T>
 struct value_accessor {
   column_device_view const col;
   bool const is_dict;
+
   value_accessor(column_device_view const& col) : col(col), is_dict(cudf::is_dictionary(col.type()))
   {
   }
@@ -93,27 +97,36 @@ struct value_accessor {
       return col.element<T>(i);
     }
   }
+
   __device__ auto operator()(size_type i) const { return value(i); }
 };
 
 /**
  * @brief Null replaced value accessor for column which supports dictionary column too.
  * For null value, returns null `init` value
  *
- * @tparam T Type of the underlying column. For dictionary column, type of the key column.
+ * @tparam SourceType Type of the underlying column. For dictionary column, type of the key column.
+ * @tparam TargetType Type that is used for computation.
  */
-template <typename T>
-struct null_replaced_value_accessor : value_accessor<T> {
-  using super_t = value_accessor<T>;
+template <typename SourceType, typename TargetType>
+struct null_replaced_value_accessor : value_accessor<SourceType> {
+  using super_t = value_accessor<SourceType>;
+
+  TargetType const init;
   bool const has_nulls;
-  T const init;
-  null_replaced_value_accessor(column_device_view const& col, T const& init, bool const has_nulls)
+
+  null_replaced_value_accessor(column_device_view const& col,
+                               TargetType const& init,
+                               bool const has_nulls)
     : super_t(col), init(init), has_nulls(has_nulls)
   {
   }
-  __device__ T operator()(size_type i) const
+
+  __device__ TargetType operator()(size_type i) const
   {
-    return has_nulls && super_t::col.is_null_nocheck(i) ? init : super_t::value(i);
+    return has_nulls && super_t::col.is_null_nocheck(i)
+             ? init
+             : static_cast<TargetType>(super_t::value(i));
   }
 };
 
@@ -168,7 +181,7 @@ struct group_reduction_functor<K, T, std::enable_if_t<is_group_reduction_support
                                         rmm::mr::device_memory_resource* mr)
 
   {
-    using DeviceType  = device_storage_type_t<T>;
+    using SourceDType = device_storage_type_t<T>;
     using ResultType  = cudf::detail::target_type_t<T, K>;
     using ResultDType = device_storage_type_t<ResultType>;
 
@@ -203,9 +216,11 @@ struct group_reduction_functor<K, T, std::enable_if_t<is_group_reduction_support
       do_reduction(count_iter, result_begin, binop);
     } else {
       using OpType    = cudf::detail::corresponding_operator_t<K>;
-      auto init       = OpType::template identity<DeviceType>();
+      auto init       = OpType::template identity<ResultDType>();
       auto inp_values = cudf::detail::make_counting_transform_iterator(
-        0, null_replaced_value_accessor{*d_values_ptr, init, values.has_nulls()});
+        0,
+        null_replaced_value_accessor<SourceDType, ResultDType>{
+          *d_values_ptr, init, values.has_nulls()});
       do_reduction(inp_values, result_begin, OpType{});
     }
 

cpp/tests/groupby/sum_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -156,6 +156,27 @@ TYPED_TEST(groupby_sum_test, dictionary)
                   force_use_sort_impl::YES);
 }
 
+struct overflow_test : public cudf::test::BaseFixture {
+};
+TEST_F(overflow_test, overflow_integer)
+{
+  using int32_col = fixed_width_column_wrapper<int32_t>;
+  using int64_col = fixed_width_column_wrapper<int64_t>;
+
+  auto const keys        = int32_col{0, 0};
+  auto const vals        = int32_col{-2147483648, -2147483648};
+  auto const expect_keys = int32_col{0};
+  auto const expect_vals = int64_col{-4294967296L};
+
+  auto test_sum = [&](auto const use_sort) {
+    auto agg = make_sum_aggregation<groupby_aggregation>();
+    test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg), use_sort);
+  };
+
+  test_sum(force_use_sort_impl::NO);
+  test_sum(force_use_sort_impl::YES);
+}
+
 template <typename T>
 struct FixedPointTestAllReps : public cudf::test::BaseFixture {
 };