GH-48679: [C++] Fix pivot_wider with non-monotonic group ids

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();

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();

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<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 +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<std::string_view> 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> grouper_;
+  Datum group_ids_to_key_indices_;
   PivotWiderOptions::UnexpectedKeyBehavior unexpected_key_behavior_;
   std::shared_ptr<Buffer> last_group_ids_;
 };