From afa0d42b121f05adbb87d851d0605ce986a65c92 Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Wed, 8 Jul 2026 16:35:51 +0200 Subject: [PATCH] GH-48679: [C++] Fix pivot_wider bug with non-monotonic group ids --- cpp/src/arrow/acero/hash_aggregate_test.cc | 77 +++++++++++++++++++ .../arrow/compute/kernels/aggregate_test.cc | 34 ++++++++ .../arrow/compute/kernels/pivot_internal.cc | 33 +++++++- 3 files changed, 141 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/acero/hash_aggregate_test.cc b/cpp/src/arrow/acero/hash_aggregate_test.cc index 442dcd588345..2f498f3a7bb2 100644 --- a/cpp/src/arrow/acero/hash_aggregate_test.cc +++ b/cpp/src/arrow/acero/hash_aggregate_test.cc @@ -4819,6 +4819,83 @@ TEST_P(GroupBy, PivotScalarKey) { } } +TEST_P(GroupBy, PivotNonMonotonicGroupId) { + // Hard-coded test for GH-48679: FastGrouperImpl can yield non-mononotonic group ids + // and the pivot_wider implementation has to account for that. + + // NOTE The precise keys to trigger this situation rely on implementation details + // of FastGrouperImpl. Any internal change might lead to this test not exercising + // the desired situation anymore. + auto key_type = utf8(); + auto value_type = float32(); + std::vector table_json = { + R"([ + [1, "k", 10.5], + [2, "l", 11.5] + ])", + R"([ + [2, "m", 12.5] + ])", + R"([ + [3, "k", 13.5], + [1, "n", 14.5], + [1, "o", 15.5] + ])"}; + std::string expected_json = R"([ + [1, {"k": 10.5, "n": 14.5, "o": 15.5} ], + [2, {"l": 11.5, "m": 12.5} ], + [3, {"k": 13.5} ] + ])"; + for (auto unexpected_key_behavior : + {PivotWiderOptions::kIgnore, PivotWiderOptions::kRaise}) { + PivotWiderOptions options(/*key_names=*/{"k", "l", "m", "n", "o"}, + unexpected_key_behavior); + TestPivot(key_type, value_type, options, table_json, expected_json); + } +} + +TEST_P(GroupBy, PivotNonMonotonicGroupIdWithScalarKey) { + // Like PivotNonMonotonicGroupId, but with a scalar key. + BatchesWithSchema input; + std::vector types = {int32(), utf8(), float32()}; + std::vector shapes = {ArgShape::ARRAY, ArgShape::SCALAR, ArgShape::ARRAY}; + input.batches = { + ExecBatchFromJSON(types, shapes, R"([ + [1, "m", 10.5], + [2, "m", 11.5] + ])"), + ExecBatchFromJSON(types, shapes, R"([ + [2, "o", 12.5] + ])"), + ExecBatchFromJSON(types, shapes, R"([ + [3, "n", 13.5], + [1, "n", 14.5] + ])"), + }; + input.schema = schema({field("group_key", int32()), field("pivot_key", utf8()), + field("pivot_value", float32())}); + Datum expected = ArrayFromJSON( + struct_({field("group_key", int32()), + field("pivoted", struct_({field("m", float32()), field("n", float32()), + field("o", float32())}))}), + R"([ + [1, {"m": 10.5, "n": 14.5} ], + [2, {"m": 11.5, "o": 12.5} ], + [3, {"n": 13.5} ] + ])"); + auto options = std::make_shared( + PivotWiderOptions(/*key_names=*/{"m", "n", "o"})); + Aggregate aggregate{"hash_pivot_wider", options, + std::vector{"pivot_key", "pivot_value"}, "pivoted"}; + for (bool use_threads : {false, true}) { + SCOPED_TRACE(use_threads ? "parallel/merged" : "serial"); + ASSERT_OK_AND_ASSIGN(Datum actual, + RunGroupBy(input, {"group_key"}, {aggregate}, use_threads)); + ValidateOutput(actual); + AssertDatumsApproxEqual(expected, actual, /*verbose=*/true); + } +} + TEST_P(GroupBy, PivotUnusedKeyName) { auto key_type = utf8(); auto value_type = float32(); diff --git a/cpp/src/arrow/compute/kernels/aggregate_test.cc b/cpp/src/arrow/compute/kernels/aggregate_test.cc index bfc3442a73d5..4ef2b2f88a51 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_test.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc @@ -4762,6 +4762,40 @@ TEST_F(TestPivotKernel, ScalarValue) { PivotWiderOptions(/*key_names=*/{"height", "width"})); } +TEST_F(TestPivotKernel, NonMonotonicGroupId) { + // Hard-coded test for GH-48679: FastGrouperImpl can yield non-mononotonic group ids + // and the pivot_wider implementation has to account for that. + + // NOTE The precise keys to trigger this situation rely on implementation details + // of FastGrouperImpl. Any internal change might lead to this test not exercising + // the desired situation anymore. + // (see similar test in hash_aggregate_test.cc) + + auto key_type = utf8(); + auto value_type = int16(); + auto keys = ArrayFromJSON(key_type, R"(["m", "n", "o"])"); + auto values = ArrayFromJSON(value_type, "[10, 11, 12]"); + auto expected = ScalarFromJSON( + struct_({field("m", value_type), field("n", value_type), field("o", value_type)}), + "[10, 11, 12]"); + AssertPivot(keys, values, *expected, PivotWiderOptions(/*key_names=*/{"m", "n", "o"})); +} + +TEST_F(TestPivotKernel, NonMonotonicGroupIdWithScalarKey) { + // Like NonMonotonicGroupId, but with a scalar key. + // Even with a single key in the data, the presence of several keys in key_names + // can still trigger the issue. + auto key_type = utf8(); + auto value_type = int16(); + + auto keys = ScalarFromJSON(key_type, R"("o")"); + auto values = ArrayFromJSON(value_type, "[null, 11, null]"); + auto expected = ScalarFromJSON( + struct_({field("m", value_type), field("n", value_type), field("o", value_type)}), + "[null, null, 11]"); + AssertPivot(keys, values, *expected, PivotWiderOptions(/*key_names=*/{"m", "n", "o"})); +} + TEST_F(TestPivotKernel, EmptyInput) { auto key_type = utf8(); auto value_type = float32(); diff --git a/cpp/src/arrow/compute/kernels/pivot_internal.cc b/cpp/src/arrow/compute/kernels/pivot_internal.cc index c9ccc7d60766..8dfbecf23690 100644 --- a/cpp/src/arrow/compute/kernels/pivot_internal.cc +++ b/cpp/src/arrow/compute/kernels/pivot_internal.cc @@ -23,10 +23,13 @@ #include "arrow/array/array_primitive.h" #include "arrow/array/builder_binary.h" +#include "arrow/array/util.h" +#include "arrow/compute/api_vector.h" #include "arrow/compute/cast.h" #include "arrow/compute/exec.h" #include "arrow/compute/kernels/codegen_internal.h" #include "arrow/compute/row/grouper.h" +#include "arrow/result.h" #include "arrow/scalar.h" #include "arrow/type_traits.h" #include "arrow/util/bit_run_reader.h" @@ -45,6 +48,7 @@ struct ConcretePivotWiderKeyMapper : public PivotWiderKeyMapper { static_cast(kMaxPivotKey), " columns: got ", options->key_names.size()); } + ctx_ = ctx; unexpected_key_behavior_ = options->unexpected_key_behavior; ARROW_ASSIGN_OR_RAISE(grouper_, Grouper::Make({&key_type}, ctx)); // Build a binary array of the pivot key values, and cast it to the desired key type @@ -61,9 +65,9 @@ struct ConcretePivotWiderKeyMapper : public PivotWiderKeyMapper { ARROW_ASSIGN_OR_RAISE(auto binary_key_array, builder.Finish()); ARROW_ASSIGN_OR_RAISE(auto key_array, Cast(*binary_key_array, &key_type, CastOptions::Safe(), ctx)); - // Populate the grouper with the keys from the array + // Populate the grouper with the keys from the array, and get the key group ids ExecSpan batch({ExecValue(*key_array->data())}, key_array->length()); - RETURN_NOT_OK(grouper_->Populate(batch)); + ARROW_ASSIGN_OR_RAISE(auto key_indices_to_group_ids, grouper_->Consume(batch)); if (grouper_->num_groups() != options->key_names.size()) { // There's a duplicate key, find it to emit a nicer error message std::unordered_set seen; @@ -75,6 +79,19 @@ struct ConcretePivotWiderKeyMapper : public PivotWiderKeyMapper { } Unreachable("Grouper doesn't agree with std::unordered_set"); } + // GH-48679: the fast grouper implementation may produce non-monotonic + // group ids, for example [0,1,2,4,3] rather than [0,1,2,3,4]. + // Therefore, we need to produce a mapping a mapping of group ids to key indices. + auto key_indices_to_group_ids_data = key_indices_to_group_ids.array(); + // InversePermutation doesn't allow unsigned integers, patch to signed. + DCHECK_EQ(key_indices_to_group_ids_data->type->id(), Type::UINT32); + key_indices_to_group_ids_data->type = int32(); + ARROW_ASSIGN_OR_RAISE(group_ids_to_key_indices_, + InversePermutation(key_indices_to_group_ids_data, + InversePermutationOptions::Defaults(), ctx)); + group_ids_to_key_indices_.array()->type = uint32(); + DCHECK_EQ(group_ids_to_key_indices_.length(), grouper_->num_groups()); + DCHECK_EQ(group_ids_to_key_indices_.null_count(), 0); return Status::OK(); } @@ -134,12 +151,22 @@ struct ConcretePivotWiderKeyMapper : public PivotWiderKeyMapper { } return Status::KeyError("Unexpected pivot key: ", key_scalar->ToString()); } - return group_id_array; + // Map back the group ids to indices in the original keys array + // NOTE Instead of materializing the Take result here, we could instead expose + // the group_ids_to_key_indices_ mapping to the caller and let them + // apply the mapping as needed. This would spare a memory allocation. + ARROW_ASSIGN_OR_RAISE(result, Take(group_ids_to_key_indices_, result, + TakeOptions::NoBoundsCheck(), ctx_)); + DCHECK(result.is_array()); + DCHECK_EQ(result.type()->id(), Type::UINT32); + return result.array(); } Status NullKeyName() { return Status::KeyError("pivot key name cannot be null"); } + ExecContext* ctx_; std::unique_ptr grouper_; + Datum group_ids_to_key_indices_; PivotWiderOptions::UnexpectedKeyBehavior unexpected_key_behavior_; std::shared_ptr last_group_ids_; };