diff --git a/cpp/src/groupby/streaming_groupby/common.cuh b/cpp/src/groupby/streaming_groupby/common.cuh index ceb521cb102b..3df6f98dc683 100644 --- a/cpp/src/groupby/streaming_groupby/common.cuh +++ b/cpp/src/groupby/streaming_groupby/common.cuh @@ -295,6 +295,7 @@ struct streaming_groupby::impl { std::unique_ptr> _key_set; [[nodiscard]] size_type num_keys() const { return static_cast(_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), diff --git a/cpp/src/groupby/streaming_groupby/impl.cu b/cpp/src/groupby/streaming_groupby/impl.cu index 32af7580b64e..edcee5ca5dcb 100644 --- a/cpp/src/groupby/streaming_groupby/impl.cu +++ b/cpp/src/groupby/streaming_groupby/impl.cu @@ -10,6 +10,7 @@ #include "groupby/hash/output_utils.hpp" #include "partial_aggregates.cuh" +#include #include #include #include @@ -19,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -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 streaming_groupby::impl::make_aggregate_results( size_type capacity, size_type first_column, @@ -263,9 +270,9 @@ std::unique_ptr
streaming_groupby::impl::make_aggregate_results( std::span{_is_agg_intermediate}.subspan(first_column, num_columns); std::vector> schema_columns; schema_columns.reserve(num_columns); - for (auto const type : - std::span{_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
keys, cuda::stream_ref stream, cudf::memory_resources mr) const { + std::unordered_map> exact_sums; + auto gathered_cols = agg_gathered->release(); + for (size_type i = 0; i < static_cast(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
(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
keys, ar.aggregations.push_back(std::unique_ptr{ dynamic_cast(_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( + 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 = diff --git a/cpp/src/groupby/streaming_groupby/warp_aggregate.cu b/cpp/src/groupby/streaming_groupby/warp_aggregate.cu index 23d0edf26ef8..c4e92ec0c74f 100644 --- a/cpp/src/groupby/streaming_groupby/warp_aggregate.cu +++ b/cpp/src/groupby/streaming_groupby/warp_aggregate.cu @@ -28,15 +28,24 @@ namespace { namespace hash = detail::hash; namespace cg = cooperative_groups; +template +constexpr bool is_wide_integer_sum = K == aggregation::SUM && cudf::is_integral_not_bool(); + +template +using warp_result_t = cuda::std::conditional_t, + __int128_t, + hash::rep_type_t>>; + template constexpr bool is_warp_aggregation_supported = [] { - if constexpr (cuda::std::same_as && K == aggregation::SUM) { + if constexpr ((cuda::std::same_as && K == aggregation::SUM) || + is_wide_integer_sum) { return true; // This kind has a dedicated existing 128-bit atomic addition. } if constexpr (hash::is_reduction_supported(K) && (K == aggregation::SUM || K == aggregation::SUM_OF_SQUARES || K == aggregation::PRODUCT || K == aggregation::MIN || K == aggregation::MAX)) { - using Result = hash::rep_type_t>; + using Result = warp_result_t; return sizeof(hash::rep_type_t) <= sizeof(int64_t) && sizeof(Result) <= sizeof(int64_t) && cudf::has_atomic_support(); } @@ -107,7 +116,7 @@ __device__ void aggregate_value(hash::rep_type_t input, { if constexpr (is_warp_aggregation_supported) { if (column.output == nullptr) { return; } - using Result = hash::rep_type_t>; + using Result = warp_result_t; using Op = cudf::detail::corresponding_operator_t; auto value = static_cast(input); if constexpr (K == aggregation::SUM_OF_SQUARES) { value *= value; } diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 1b0b4f4729ba..52e0e2a8cc6b 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -30,13 +30,16 @@ #include #include +#include #include #include #include +#include #include +#include #include #include -#include +#include #include static std::vector const KEY_COL{0}; @@ -448,8 +451,20 @@ TEST_F(StreamingGroupbyTest, ColumnwiseGrowthPreservesResultsOnAllocationFailure constexpr cudf::size_type num_rows = 65536; constexpr std::size_t num_columns = 3; // Four groups initially have 64 partials each. Growth to 1,024 groups collapses them to one. - constexpr std::size_t old_column_bytes = 4 * 64 * sizeof(int64_t); - constexpr std::size_t new_column_bytes = 1024 * sizeof(int64_t); + // Integer SUM state is 128-bit; MIN and MAX keep the int64 input width. + constexpr std::size_t old_rows = 4 * 64; + constexpr std::size_t new_rows = 1024; + constexpr std::array element_bytes{ + sizeof(__int128_t), sizeof(int64_t), sizeof(int64_t)}; + auto const is_column_size = [&](std::size_t bytes, std::size_t rows) { + return std::any_of(element_bytes.begin(), element_bytes.end(), [&](auto element) { + return bytes == rows * element; + }); + }; + auto const max_old_column_bytes = + old_rows * *std::max_element(element_bytes.begin(), element_bytes.end()); + auto const total_new_bytes = + new_rows * std::accumulate(element_bytes.begin(), element_bytes.end(), std::size_t{0}); std::vector key_data(num_rows); std::vector value_data(num_rows); std::vector validity(num_rows); @@ -480,9 +495,11 @@ TEST_F(StreamingGroupbyTest, ColumnwiseGrowthPreservesResultsOnAllocationFailure for (std::size_t const fail_after : {0, 1, 2}) { SCOPED_TRACE(fail_after); - std::unordered_set old_buffers; + std::unordered_map old_buffers; + std::size_t old_bytes = 0; std::size_t retired_columns = 0; std::size_t grown_columns = 0; + std::size_t grown_bytes = 0; std::size_t peak_state_bytes = 0; bool growing = false; rmm::mr::cuda_memory_resource upstream; @@ -494,18 +511,24 @@ TEST_F(StreamingGroupbyTest, ColumnwiseGrowthPreservesResultsOnAllocationFailure throw rmm::bad_alloc{"Injected streaming groupby growth failure"}; } auto* ptr = upstream.allocate(stream, bytes, alignment); - if (!growing && bytes == old_column_bytes) { old_buffers.insert(ptr); } - if (growing && bytes == new_column_bytes) { + if (!growing && is_column_size(bytes, old_rows)) { + old_buffers.emplace(ptr, bytes); + old_bytes += bytes; + } + if (growing && is_column_size(bytes, new_rows)) { ++grown_columns; + grown_bytes += bytes; EXPECT_EQ(retired_columns + 1, grown_columns); - peak_state_bytes = - std::max(peak_state_bytes, - old_buffers.size() * old_column_bytes + grown_columns * new_column_bytes); + peak_state_bytes = std::max(peak_state_bytes, old_bytes + grown_bytes); } return ptr; }, [&](cuda::stream_ref stream, void* ptr, std::size_t bytes, std::size_t alignment, void*) { - if (old_buffers.erase(ptr) != 0 && growing) { ++retired_columns; } + if (auto const it = old_buffers.find(ptr); it != old_buffers.end()) { + old_bytes -= it->second; + old_buffers.erase(it); + if (growing) { ++retired_columns; } + } upstream.deallocate(stream, ptr, bytes, alignment); }}; cudf::groupby::streaming_groupby agg{ @@ -524,7 +547,7 @@ TEST_F(StreamingGroupbyTest, ColumnwiseGrowthPreservesResultsOnAllocationFailure auto const expected_columns = fail_after == 0 ? num_columns : fail_after; EXPECT_EQ(retired_columns, expected_columns); EXPECT_EQ(grown_columns, expected_columns); - EXPECT_LE(peak_state_bytes, num_columns * new_column_bytes + old_column_bytes); + EXPECT_LE(peak_state_bytes, total_new_bytes + max_old_column_bytes); EXPECT_EQ(agg.distinct_keys(), fail_after == 0 ? 1024 : 4); if (fail_after != 0) { EXPECT_THROW(agg.aggregate(repeated_batch), cudf::logic_error); } @@ -1404,6 +1427,225 @@ TEST_F(StreamingGroupbyTest, ManySmallBatches) verify_against_groupby(keys, results, batches, KEY_COL, reqs); } +TEST_F(StreamingGroupbyTest, MeanInt64OverflowAcrossBatches) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr max = std::numeric_limits::max(); + auto constexpr min = std::numeric_limits::min(); + + auto reqs = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + for (int b = 0; b < 3; ++b) { + cudf::test::fixed_width_column_wrapper keys{0, 1, 2}; + cudf::test::fixed_width_column_wrapper vals{max, min, V{b}}; + streaming_agg.aggregate(cudf::table_view{{keys, vals}}); + } + auto [keys, results] = streaming_agg.finalize(); + + cudf::test::fixed_width_column_wrapper ek{0, 1, 2}; + cudf::test::fixed_width_column_wrapper ev{ + static_cast(max), static_cast(min), 1.0}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, MeanInt64OverflowWithinBatch) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr max = std::numeric_limits::max(); + auto constexpr min = std::numeric_limits::min(); + auto constexpr num_rows = 1000; + + std::vector key_data(num_rows); + std::vector val_data(num_rows); + for (int i = 0; i < num_rows; ++i) { + key_data[i] = i % 2; + val_data[i] = i % 2 == 0 ? max : min; + } + cudf::test::fixed_width_column_wrapper keys(key_data.begin(), key_data.end()); + cudf::test::fixed_width_column_wrapper vals(val_data.begin(), val_data.end()); + + auto reqs = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(cudf::table_view{{keys, vals}}); + auto [result_keys, results] = streaming_agg.finalize(); + + cudf::test::fixed_width_column_wrapper ek{0, 1}; + cudf::test::fixed_width_column_wrapper ev{static_cast(max), + static_cast(min)}; + check(result_keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, MeanUint64Overflow) +{ + using K = int32_t; + using V = uint64_t; + + auto constexpr max = std::numeric_limits::max(); + + auto reqs = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + for (int b = 0; b < 3; ++b) { + cudf::test::fixed_width_column_wrapper keys{0}; + cudf::test::fixed_width_column_wrapper vals{max}; + streaming_agg.aggregate(cudf::table_view{{keys, vals}}); + } + auto [keys, results] = streaming_agg.finalize(); + + cudf::test::fixed_width_column_wrapper ek{0}; + cudf::test::fixed_width_column_wrapper ev{static_cast(max)}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, SumAndMeanInt64OverflowShareState) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr max = std::numeric_limits::max(); + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(1, cudf::make_mean_aggregation())); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + for (int b = 0; b < 3; ++b) { + cudf::test::fixed_width_column_wrapper keys{0, 1}; + cudf::test::fixed_width_column_wrapper vals{max, V{b}}; + streaming_agg.aggregate(cudf::table_view{{keys, vals}}); + } + auto [keys, results] = streaming_agg.finalize(); + + auto const wrapped_sum = static_cast(3 * static_cast(max)); + cudf::test::fixed_width_column_wrapper ek{0, 1}; + cudf::test::fixed_width_column_wrapper e_sum{wrapped_sum, V{3}}; + cudf::test::fixed_width_column_wrapper e_mean{static_cast(max), 1.0}; + check(keys, results, cudf::table_view{{ek}}, {e_sum, e_mean}); +} + +TEST_F(StreamingGroupbyTest, MeanInt64OverflowWithNulls) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr max = std::numeric_limits::max(); + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(1, cudf::make_mean_aggregation())); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + for (int b = 0; b < 2; ++b) { + cudf::test::fixed_width_column_wrapper keys{0, 0, 1}; + cudf::test::fixed_width_column_wrapper vals{{max, V{7}, V{7}}, {true, false, false}}; + streaming_agg.aggregate(cudf::table_view{{keys, vals}}); + } + auto [keys, results] = streaming_agg.finalize(); + + auto const wrapped_sum = static_cast(2 * static_cast(max)); + cudf::test::fixed_width_column_wrapper ek{0, 1}; + cudf::test::fixed_width_column_wrapper e_sum{{wrapped_sum, V{0}}, {true, false}}; + cudf::test::fixed_width_column_wrapper e_mean{{static_cast(max), 0.0}, + {true, false}}; + check(keys, results, cudf::table_view{{ek}}, {e_sum, e_mean}); +} + +TEST_F(StreamingGroupbyTest, MeanInt64OverflowMergeTwoObjects) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr max = std::numeric_limits::max(); + + cudf::test::fixed_width_column_wrapper keys1{0, 1}; + cudf::test::fixed_width_column_wrapper vals1{max, V{2}}; + cudf::test::fixed_width_column_wrapper keys2{0, 0, 2}; + cudf::test::fixed_width_column_wrapper vals2{max, max, V{4}}; + + auto reqs1 = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby worker1(KEY_COL, reqs1, DEFAULT_MAX_DISTINCT_KEYS); + worker1.aggregate(cudf::table_view{{keys1, vals1}}); + + auto reqs2 = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby worker2(KEY_COL, reqs2, DEFAULT_MAX_DISTINCT_KEYS); + worker2.aggregate(cudf::table_view{{keys2, vals2}}); + + worker1.merge(worker2); + auto [keys, results] = worker1.finalize(); + + cudf::test::fixed_width_column_wrapper ek{0, 1, 2}; + cudf::test::fixed_width_column_wrapper ev{static_cast(max), 2.0, 4.0}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, WideIntegerSumGrowthMatchesGroupby) +{ + using K = int32_t; + using V = int64_t; + + auto constexpr num_batches = 4; + auto constexpr rows_per_batch = 2000; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(1, cudf::make_mean_aggregation())); + reqs.push_back(make_req(1, cudf::make_variance_aggregation())); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 4 * rows_per_batch); + + std::vector batches; + std::vector> owners; + for (int b = 0; b < num_batches; ++b) { + std::vector key_data(rows_per_batch); + std::vector val_data(rows_per_batch); + for (int i = 0; i < rows_per_batch; ++i) { + key_data[i] = (b + 1) * i; // Each batch adds new groups and revisits earlier ones. + val_data[i] = (i % 7) - 3 + b * 1000; + } + owners.push_back( + cudf::test::fixed_width_column_wrapper(key_data.begin(), key_data.end()).release()); + owners.push_back( + cudf::test::fixed_width_column_wrapper(val_data.begin(), val_data.end()).release()); + batches.emplace_back( + std::vector{owners[owners.size() - 2]->view(), owners.back()->view()}); + streaming_agg.aggregate(batches.back()); + } + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, batches, KEY_COL, reqs); +} + +template +struct StreamingGroupbyIntegralTest : public cudf::test::BaseFixture {}; + +TYPED_TEST_SUITE(StreamingGroupbyIntegralTest, cudf::test::IntegralTypesNotBool); + +TYPED_TEST(StreamingGroupbyIntegralTest, SumMeanVarianceMatchGroupby) +{ + using K = int32_t; + using V = TypeParam; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 3}; + cudf::test::fixed_width_column_wrapper vals1{{10, 20, 30, 5}, + {true, true, true, false}}; + cudf::test::fixed_width_column_wrapper keys2{1, 2, 3, 2}; + cudf::test::fixed_width_column_wrapper vals2{50, 40, 7, 60}; + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(1, cudf::make_mean_aggregation())); + reqs.push_back(make_req(1, cudf::make_variance_aggregation())); + reqs.push_back(make_req(1, cudf::make_min_aggregation())); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + // Test that exceeding distinct-key capacity throws. TEST_F(StreamingGroupbyTest, ExceedsDistinctKeyCapacityThrows) {