This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 34647d278a GH-48679: [C++] Fix pivot_wider with non-monotonic group
ids (#50423)
34647d278a is described below
commit 34647d278ab997007ca70aa880a6f1a2128905de
Author: Antoine Pitrou <[email protected]>
AuthorDate: Wed Jul 22 11:07:57 2026 +0200
GH-48679: [C++] Fix pivot_wider with non-monotonic group ids (#50423)
### Rationale for this change
The pivot_wider implementation was assuming that the `Grouper` produces
monotonic group ids for the pivot key names.
However, that turns out to not always be true when the `FastGrouperImpl` is
involved (perhaps when hash collisions are involved?).
### What changes are included in this PR?
Compute and apply a mapping from pivot key group ids back to pivot key
indices.
### Are these changes tested?
Yes, by additional tests.
### Are there any user-facing changes?
Only a bugfix.
* GitHub Issue: #48679
Authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/acero/hash_aggregate_test.cc | 77 +++++++++++++++++++++++++
cpp/src/arrow/compute/kernels/aggregate_test.cc | 34 +++++++++++
cpp/src/arrow/compute/kernels/pivot_internal.cc | 37 ++++++++++--
3 files changed, 144 insertions(+), 4 deletions(-)
diff --git a/cpp/src/arrow/acero/hash_aggregate_test.cc
b/cpp/src/arrow/acero/hash_aggregate_test.cc
index 442dcd5883..2f498f3a7b 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<std::string> 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<TypeHolder> types = {int32(), utf8(), float32()};
+ std::vector<ArgShape> 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>(
+ PivotWiderOptions(/*key_names=*/{"m", "n", "o"}));
+ Aggregate aggregate{"hash_pivot_wider", options,
+ std::vector<FieldRef>{"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 bfc3442a73..4ef2b2f88a 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 c9ccc7d607..8350edba12 100644
--- a/cpp/src/arrow/compute/kernels/pivot_internal.cc
+++ b/cpp/src/arrow/compute/kernels/pivot_internal.cc
@@ -23,10 +23,11 @@
#include "arrow/array/array_primitive.h"
#include "arrow/array/builder_binary.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 +46,7 @@ struct ConcretePivotWiderKeyMapper : public
PivotWiderKeyMapper {
static_cast<size_t>(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 +63,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<std::string_view> seen;
@@ -75,6 +77,23 @@ 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 of group ids to key indices.
+ // TODO: revisit this if the fast grouper is amended to guarantee monotonic
+ // group ids.
+ auto key_indices_to_group_ids_data =
key_indices_to_group_ids.array()->Copy();
+ // 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(auto group_ids_to_key_indices,
+ InversePermutation(key_indices_to_group_ids_data,
+
InversePermutationOptions::Defaults(), ctx));
+ auto group_ids_to_key_indices_data =
group_ids_to_key_indices.array()->Copy();
+ group_ids_to_key_indices_data->type = uint32();
+ DCHECK_EQ(group_ids_to_key_indices.length(), grouper_->num_groups());
+ DCHECK_EQ(group_ids_to_key_indices.null_count(), 0);
+ group_ids_to_key_indices_ =
Datum(std::move(group_ids_to_key_indices_data));
return Status::OK();
}
@@ -134,12 +153,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> grouper_;
+ Datum group_ids_to_key_indices_;
PivotWiderOptions::UnexpectedKeyBehavior unexpected_key_behavior_;
std::shared_ptr<Buffer> last_group_ids_;
};
