diff --git a/cpp/src/groupby/sort/group_merge_m2.cu b/cpp/src/groupby/sort/group_merge_m2.cu index e3d086d75a9..578f3bb1fcd 100644 --- a/cpp/src/groupby/sort/group_merge_m2.cu +++ b/cpp/src/groupby/sort/group_merge_m2.cu @@ -49,11 +49,23 @@ struct merge_fn { if (partial_n == 0) { continue; } auto const partial_avg = d_means[idx]; auto const partial_m2 = d_M2s[idx]; - auto const new_n = n + partial_n; - auto const delta = partial_avg - avg; - m2 += partial_m2 + delta * delta * n * partial_n / new_n; - avg = (avg * n + partial_avg * partial_n) / new_n; - n = new_n; + + // Merging an empty accumulator with a non-empty partial is an identity operation. Running + // the generic formula for this case can evaluate inf * 0 and turn extreme finite partials + // into NaN. + if (n == 0) { + n = partial_n; + avg = partial_avg; + m2 = partial_m2; + continue; + } + + auto const new_n = n + partial_n; + auto const delta = partial_avg - avg; + auto const delta_n = delta / new_n; + m2 += partial_m2 + delta * delta_n * n * partial_n; + avg += delta_n * partial_n; + n = new_n; } return {n, avg, m2}; diff --git a/cpp/tests/groupby/merge_m2_tests.cpp b/cpp/tests/groupby/merge_m2_tests.cpp index d9c9b35dd00..1bcc405a0fc 100644 --- a/cpp/tests/groupby/merge_m2_tests.cpp +++ b/cpp/tests/groupby/merge_m2_tests.cpp @@ -14,6 +14,8 @@ #include #include +#include + using namespace cudf::test::iterators; namespace { @@ -87,11 +89,119 @@ auto merge_M2(vcol_views const& keys_cols, vcol_views const& values_cols) auto result = gb_obj.aggregate(requests); return std::pair(std::move(result.first->release()[0]), std::move(result.second[0].results[0])); } + +template +void test_extreme_finite_first_partial() +{ + auto const key = keys_col{1}; + auto counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto means = means_col{std::numeric_limits::max()}; + auto m2s = M2s_col{0.0}; + auto const vals = structs_col{counts, means, m2s}; + + auto const [out_key, out_vals] = merge_M2({key}, {vals}); + + auto const expected_keys = keys_col{1}; + auto expected_counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto expected_means = means_col{std::numeric_limits::max()}; + auto expected_m2s = M2s_col{0.0}; + auto const expected_values = structs_col{expected_counts, expected_means, expected_m2s}; + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_keys, *out_key, verbosity); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_values, *out_vals, verbosity); +} + +template +void test_extreme_finite_merged_partials() +{ + auto const keys = keys_col{1, 1}; + auto counts = cudf::test::fixed_width_column_wrapper{CountType{1}, CountType{1}}; + auto means = means_col{std::numeric_limits::max(), 0.0}; + auto m2s = M2s_col{0.0, 0.0}; + auto const vals = structs_col{counts, means, m2s}; + + auto const [out_keys, out_vals] = merge_M2({keys}, {vals}); + + auto const expected_keys = keys_col{1}; + auto expected_counts = cudf::test::fixed_width_column_wrapper{CountType{2}}; + auto expected_means = means_col{std::numeric_limits::max() / 2}; + auto expected_m2s = M2s_col{std::numeric_limits::infinity()}; + auto const expected_values = structs_col{expected_counts, expected_means, expected_m2s}; + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_keys, *out_keys, verbosity); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_values, *out_vals, verbosity); +} + +// A non-finite mean in a partial must propagate, not be coerced. This pins the +// identity-branch behavior: a single non-empty partial with NaN/Inf statistics +// should be preserved as-is, since coercing to NaN would discard upstream +// signal (e.g. an upstream overflow that already produced +Inf). +template +void test_nan_mean_first_partial() +{ + auto const key = keys_col{1}; + auto counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto means = means_col{std::numeric_limits::quiet_NaN()}; + auto m2s = M2s_col{std::numeric_limits::quiet_NaN()}; + auto const vals = structs_col{counts, means, m2s}; + + auto const [out_key, out_vals] = merge_M2({key}, {vals}); + + auto const expected_keys = keys_col{1}; + auto expected_counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto expected_means = means_col{std::numeric_limits::quiet_NaN()}; + auto expected_m2s = M2s_col{std::numeric_limits::quiet_NaN()}; + auto const expected_values = structs_col{expected_counts, expected_means, expected_m2s}; + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_keys, *out_key, verbosity); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_values, *out_vals, verbosity); +} + +template +void test_inf_mean_first_partial() +{ + auto const key = keys_col{1}; + auto counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto means = means_col{std::numeric_limits::infinity()}; + auto m2s = M2s_col{0.0}; + auto const vals = structs_col{counts, means, m2s}; + + auto const [out_key, out_vals] = merge_M2({key}, {vals}); + + auto const expected_keys = keys_col{1}; + auto expected_counts = cudf::test::fixed_width_column_wrapper{CountType{1}}; + auto expected_means = means_col{std::numeric_limits::infinity()}; + auto expected_m2s = M2s_col{0.0}; + auto const expected_values = structs_col{expected_counts, expected_means, expected_m2s}; + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_keys, *out_key, verbosity); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_values, *out_vals, verbosity); +} + +// Once the accumulator holds a NaN, any subsequent merge step must keep +// propagating NaN (NaN ⊕ anything = NaN), regardless of partial order. +template +void test_nan_mean_merged_with_finite() +{ + auto const keys = keys_col{1, 1}; + auto counts = cudf::test::fixed_width_column_wrapper{CountType{10}, CountType{10}}; + auto means = means_col{std::numeric_limits::quiet_NaN(), 5.0}; + auto m2s = M2s_col{std::numeric_limits::quiet_NaN(), 20.0}; + auto const vals = structs_col{counts, means, m2s}; + + auto const [out_keys, out_vals] = merge_M2({keys}, {vals}); + + auto const expected_keys = keys_col{1}; + auto expected_counts = cudf::test::fixed_width_column_wrapper{CountType{20}}; + auto expected_means = means_col{std::numeric_limits::quiet_NaN()}; + auto expected_m2s = M2s_col{std::numeric_limits::quiet_NaN()}; + auto const expected_values = structs_col{expected_counts, expected_means, expected_m2s}; + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_keys, *out_keys, verbosity); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_values, *out_vals, verbosity); +} } // namespace template struct GroupbyMergeM2TypedTest : public cudf::test::BaseFixture {}; +struct GroupbyMergeM2ExtremeTest : public cudf::test::BaseFixture {}; + using TestTypes = cudf::test::Concat, cudf::test::FloatingPointTypes>; TYPED_TEST_SUITE(GroupbyMergeM2TypedTest, TestTypes); @@ -145,6 +255,56 @@ TYPED_TEST(GroupbyMergeM2TypedTest, EmptyInput) CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(vals, *out_vals, verbosity); } +TEST_F(GroupbyMergeM2ExtremeTest, ExtremeFiniteFirstPartialInt64Count) +{ + test_extreme_finite_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, ExtremeFiniteFirstPartialDoubleCount) +{ + test_extreme_finite_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, ExtremeFiniteMergedPartialsInt64Count) +{ + test_extreme_finite_merged_partials(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, ExtremeFiniteMergedPartialsDoubleCount) +{ + test_extreme_finite_merged_partials(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, NanMeanFirstPartialInt64Count) +{ + test_nan_mean_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, NanMeanFirstPartialDoubleCount) +{ + test_nan_mean_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, InfMeanFirstPartialInt64Count) +{ + test_inf_mean_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, InfMeanFirstPartialDoubleCount) +{ + test_inf_mean_first_partial(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, NanMeanMergedWithFiniteInt64Count) +{ + test_nan_mean_merged_with_finite(); +} + +TEST_F(GroupbyMergeM2ExtremeTest, NanMeanMergedWithFiniteDoubleCount) +{ + test_nan_mean_merged_with_finite(); +} + TYPED_TEST(GroupbyMergeM2TypedTest, SimpleInput) { using T = TypeParam;