Merge remote-tracking branch 'patryk/aggregate_sequencer' into combined2

Rafał Hibner · Rafał Hibner · commit d7b29d2be2b0 · 2025-08-12T11:19:25.000Z
diff --git a/cpp/src/arrow/acero/aggregate_internal.cc b/cpp/src/arrow/acero/aggregate_internal.cc
@@ -214,12 +214,12 @@ Result<std::shared_ptr<Schema>> MakeOutputSchema(
     ARROW_ASSIGN_OR_RAISE(auto args,
                           ScalarAggregateNode::MakeAggregateNodeArgs(
                               input_schema, keys, segment_keys, aggregates, exec_ctx,
-                              /*concurrency=*/1, {}));
+                              /*concurrency=*/1));
     return std::move(args.output_schema);
   } else {
     ARROW_ASSIGN_OR_RAISE(
         auto args, GroupByNode::MakeAggregateNodeArgs(input_schema, keys, segment_keys,
-                                                      aggregates, exec_ctx, {}));
+                                                      aggregates, exec_ctx));
     return std::move(args.output_schema);
   }
 }
diff --git a/cpp/src/arrow/acero/aggregate_internal.h b/cpp/src/arrow/acero/aggregate_internal.h
@@ -98,7 +98,7 @@ struct AggregateNodeArgs {
   std::vector<const KernelType*> kernels;
   std::vector<std::vector<TypeHolder>> kernel_intypes;
   std::vector<std::vector<std::unique_ptr<KernelState>>> states;
-  Ordering ordering;
+  bool requires_ordering;
 };
 
 std::vector<TypeHolder> ExtendWithGroupIdType(const std::vector<TypeHolder>& in_types);
@@ -194,7 +194,7 @@ class ScalarAggregateNode : public ExecNode,
   static Result<AggregateNodeArgs<ScalarAggregateKernel>> MakeAggregateNodeArgs(
       const std::shared_ptr<Schema>& input_schema, const std::vector<FieldRef>& keys,
       const std::vector<FieldRef>& segment_keys, const std::vector<Aggregate>& aggs,
-      ExecContext* exec_ctx, size_t concurrency, std::vector<ExecNode*> inputs);
+      ExecContext* exec_ctx, size_t concurrency);
 
   static Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
                                 const ExecNodeOptions& options);
@@ -291,7 +291,7 @@ class GroupByNode : public ExecNode,
   static Result<AggregateNodeArgs<HashAggregateKernel>> MakeAggregateNodeArgs(
       const std::shared_ptr<Schema>& input_schema, const std::vector<FieldRef>& keys,
       const std::vector<FieldRef>& segment_keys, const std::vector<Aggregate>& aggs,
-      ExecContext* ctx, std::vector<ExecNode*> inputs);
+      ExecContext* ctx);
 
   static Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
                                 const ExecNodeOptions& options);
diff --git a/cpp/src/arrow/acero/groupby_aggregate_node.cc b/cpp/src/arrow/acero/groupby_aggregate_node.cc
@@ -74,7 +74,7 @@ Status GroupByNode::Init() {
 Result<AggregateNodeArgs<HashAggregateKernel>> GroupByNode::MakeAggregateNodeArgs(
     const std::shared_ptr<Schema>& input_schema, const std::vector<FieldRef>& keys,
     const std::vector<FieldRef>& segment_keys, const std::vector<Aggregate>& aggs,
-    ExecContext* ctx, std::vector<ExecNode*> inputs) {
+    ExecContext* ctx) {
   // Find input field indices for key fields
   std::vector<int> key_field_ids(keys.size());
   for (size_t i = 0; i < keys.size(); ++i) {
@@ -167,37 +167,6 @@ Result<AggregateNodeArgs<HashAggregateKernel>> GroupByNode::MakeAggregateNodeArg
   for (size_t i = 0; i < aggs.size(); ++i) {
     output_fields[base + i] = agg_result_fields[i]->WithName(aggs[i].name);
   }
-  Ordering out_ordering = Ordering::Unordered();
-  if (requires_ordering && inputs[0]->ordering().is_implicit()) {
-    out_ordering = Ordering::Implicit();
-  } else if (requires_ordering) {
-    std::vector<compute::SortKey> out_sort_keys;
-    std::unordered_set<int> segmented_key_field_id_set(segment_key_field_ids.begin(),
-                                                       segment_key_field_ids.end());
-    // Propagate output sorting only by segmented keys excluding sorting by regualr keys
-    // since this will break the segmentation.
-    for (auto key : inputs[0]->ordering().sort_keys()) {
-      ARROW_ASSIGN_OR_RAISE(auto match, key.target.FindOne(*input_schema));
-      if (segmented_key_field_id_set.find(match[0]) != segmented_key_field_id_set.end()) {
-        out_sort_keys.emplace_back(key);
-      } else {
-        break;
-      }
-    }
-    if (out_sort_keys.size() > 0)
-      out_ordering = Ordering(out_sort_keys);
-    else
-      out_ordering = Ordering::Implicit();
-  }
-
-  if (!out_ordering.is_unordered()) {
-    if (inputs[0]->ordering().is_unordered()) {
-      return Status::Invalid(
-          "Aggregate node's input has no meaningful ordering and so limit/offset will be "
-          "non-deterministic.  Please establish order in some way (e.g. by inserting an "
-          "order_by node)");
-    }
-  }
 
   return AggregateNodeArgs<HashAggregateKernel>{schema(std::move(output_fields)),
                                                 std::move(key_field_ids),
@@ -208,7 +177,7 @@ Result<AggregateNodeArgs<HashAggregateKernel>> GroupByNode::MakeAggregateNodeArg
                                                 std::move(agg_kernels),
                                                 std::move(agg_src_types),
                                                 /*states=*/{},
-                                                std::move(out_ordering)};
+                                                requires_ordering};
 }
 
 Result<ExecNode*> GroupByNode::Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
@@ -224,14 +193,45 @@ Result<ExecNode*> GroupByNode::Make(ExecPlan* plan, std::vector<ExecNode*> input
 
   const auto& input_schema = input->output_schema();
   auto exec_ctx = plan->query_context()->exec_context();
-  ARROW_ASSIGN_OR_RAISE(auto args, MakeAggregateNodeArgs(input_schema, keys, segment_keys,
-                                                         aggs, exec_ctx, inputs));
+  ARROW_ASSIGN_OR_RAISE(
+      auto args, MakeAggregateNodeArgs(input_schema, keys, segment_keys, aggs, exec_ctx));
 
+  Ordering out_ordering = Ordering::Unordered();
+  if (args.requires_ordering && input->ordering().is_implicit()) {
+    out_ordering = Ordering::Implicit();
+  } else if (args.requires_ordering) {
+    std::vector<compute::SortKey> out_sort_keys;
+    std::unordered_set<int> segmented_key_field_id_set(args.segment_key_field_ids.begin(),
+                                                       args.segment_key_field_ids.end());
+    // Propagate output sorting only by segmented keys excluding sorting by regular keys
+    // since this will break the segmentation.
+    for (auto key : input->ordering().sort_keys()) {
+      ARROW_ASSIGN_OR_RAISE(auto match, key.target.FindOne(*input_schema));
+      if (segmented_key_field_id_set.find(match[0]) != segmented_key_field_id_set.end()) {
+        out_sort_keys.emplace_back(key);
+      } else {
+        break;
+      }
+    }
+    if (out_sort_keys.size() > 0) {
+      out_ordering = Ordering(out_sort_keys);
+    } else {
+      out_ordering = Ordering::Implicit();
+    }
+  }
+  if (!out_ordering.is_unordered()) {
+    if (inputs[0]->ordering().is_unordered()) {
+      return Status::Invalid(
+          "Aggregate node's input has no meaningful ordering and so limit/offset will be "
+          "non-deterministic.  Please establish order in some way (e.g. by inserting an "
+          "order_by node)");
+    }
+  }
   return input->plan()->EmplaceNode<GroupByNode>(
       input, std::move(args.output_schema), std::move(args.grouping_key_field_ids),
       std::move(args.segment_key_field_ids), std::move(args.segmenter),
       std::move(args.kernel_intypes), std::move(args.target_fieldsets),
-      std::move(args.aggregates), std::move(args.kernels), std::move(args.ordering));
+      std::move(args.aggregates), std::move(args.kernels), std::move(out_ordering));
 }
 
 Status GroupByNode::ResetKernelStates() {
diff --git a/cpp/src/arrow/acero/hash_aggregate_test.cc b/cpp/src/arrow/acero/hash_aggregate_test.cc
@@ -39,6 +39,7 @@
 #include "arrow/compute/cast.h"
 #include "arrow/compute/exec.h"
 #include "arrow/compute/exec_internal.h"
+#include "arrow/compute/ordering.h"
 #include "arrow/compute/registry.h"
 #include "arrow/compute/row/grouper.h"
 #include "arrow/table.h"
@@ -315,7 +316,8 @@ Result<Datum> RunGroupBy(const BatchesWithSchema& input,
       Declaration::Sequence(
           {
               {"source",
-               SourceNodeOptions{input.schema, input.gen(use_threads, /*slow=*/false)}},
+               SourceNodeOptions{input.schema, input.gen(use_threads, /*slow=*/false),
+                                 Ordering::Implicit()}},
               {"aggregate", AggregateNodeOptions{aggregates, std::move(keys),
                                                  std::move(segment_keys)}},
               {"sink", SinkNodeOptions{&sink_gen}},
diff --git a/cpp/src/arrow/acero/plan_test.cc b/cpp/src/arrow/acero/plan_test.cc
@@ -1651,12 +1651,13 @@ TEST(ExecPlanExecution, SegmentedAggregationWithMultiThreading) {
   data.schema = schema({field("i32", int32())});
   Declaration plan = Declaration::Sequence(
       {{"source",
-        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false)}},
+        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false),
+                          Ordering::Unordered()}},
        {"aggregate", AggregateNodeOptions{/*aggregates=*/{
                                               {"count", nullptr, "i32", "count(i32)"},
                                           },
                                           /*keys=*/{}, /*segment_keys=*/{"i32"}}}});
-  EXPECT_RAISES_WITH_MESSAGE_THAT(NotImplemented, HasSubstr("multi-threaded"),
+  EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, HasSubstr("meaningful ordering"),
                                   DeclarationToExecBatches(std::move(plan)));
 }
 
@@ -1674,7 +1675,8 @@ TEST(ExecPlanExecution, SegmentedAggregationWithOneSegment) {
 
   Declaration plan = Declaration::Sequence(
       {{"source",
-        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false)}},
+        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false),
+                          Ordering::Implicit()}},
        {"aggregate", AggregateNodeOptions{/*aggregates=*/{
                                               {"hash_sum", nullptr, "c", "sum(c)"},
                                               {"hash_mean", nullptr, "c", "mean(c)"},
@@ -1703,7 +1705,8 @@ TEST(ExecPlanExecution, SegmentedAggregationWithTwoSegments) {
 
   Declaration plan = Declaration::Sequence(
       {{"source",
-        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false)}},
+        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false),
+                          Ordering::Implicit()}},
        {"aggregate", AggregateNodeOptions{/*aggregates=*/{
                                               {"hash_sum", nullptr, "c", "sum(c)"},
                                               {"hash_mean", nullptr, "c", "mean(c)"},
@@ -1733,7 +1736,8 @@ TEST(ExecPlanExecution, SegmentedAggregationWithBatchCrossingSegment) {
 
   Declaration plan = Declaration::Sequence(
       {{"source",
-        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false)}},
+        SourceNodeOptions{data.schema, data.gen(/*parallel=*/false, /*slow=*/false),
+                          Ordering::Implicit()}},
        {"aggregate", AggregateNodeOptions{/*aggregates=*/{
                                               {"hash_sum", nullptr, "c", "sum(c)"},
                                               {"hash_mean", nullptr, "c", "mean(c)"},
diff --git a/cpp/src/arrow/acero/scalar_aggregate_node.cc b/cpp/src/arrow/acero/scalar_aggregate_node.cc
@@ -62,8 +62,7 @@ ScalarAggregateNode::MakeAggregateNodeArgs(const std::shared_ptr<Schema>& input_
                                            const std::vector<FieldRef>& keys,
                                            const std::vector<FieldRef>& segment_keys,
                                            const std::vector<Aggregate>& aggs,
-                                           ExecContext* exec_ctx, size_t concurrency,
-                                           std::vector<ExecNode*> inputs) {
+                                           ExecContext* exec_ctx, size_t concurrency) {
   // Copy (need to modify options pointer below)
   std::vector<Aggregate> aggregates(aggs);
   std::vector<int> segment_field_ids(segment_keys.size());
@@ -158,27 +157,61 @@ ScalarAggregateNode::MakeAggregateNodeArgs(const std::shared_ptr<Schema>& input_
     fields[base + i] = field(aggregates[i].name, out_type.GetSharedPtr());
   }
 
+  return AggregateNodeArgs<ScalarAggregateKernel>{schema(std::move(fields)),
+                                                  /*grouping_key_field_ids=*/{},
+                                                  std::move(segment_field_ids),
+                                                  std::move(segmenter),
+                                                  std::move(target_fieldsets),
+                                                  std::move(aggregates),
+                                                  std::move(kernels),
+                                                  std::move(kernel_intypes),
+                                                  std::move(states),
+                                                  requires_ordering};
+}
+
+Result<ExecNode*> ScalarAggregateNode::Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
+                                            const ExecNodeOptions& options) {
+  RETURN_NOT_OK(ValidateExecNodeInputs(plan, inputs, 1, "ScalarAggregateNode"));
+
+  const auto& aggregate_options = checked_cast<const AggregateNodeOptions&>(options);
+  auto aggregates = aggregate_options.aggregates;
+  const auto& keys = aggregate_options.keys;
+  const auto& segment_keys = aggregate_options.segment_keys;
+  const auto input = inputs[0];
+  const auto concurrency = plan->query_context()->max_concurrency();
+
+  if (keys.size() > 0) {
+    return Status::Invalid("Scalar aggregation with some key");
+  }
+
+  const auto& input_schema = inputs[0]->output_schema();
+  auto exec_ctx = plan->query_context()->exec_context();
+
+  ARROW_ASSIGN_OR_RAISE(
+      auto args, MakeAggregateNodeArgs(input_schema, keys, segment_keys, aggregates,
+                                       exec_ctx, concurrency));
   Ordering out_ordering = Ordering::Unordered();
-  if (requires_ordering && inputs[0]->ordering().is_implicit()) {
+  if (args.requires_ordering && input->ordering().is_implicit()) {
     out_ordering = Ordering::Implicit();
-  } else if (requires_ordering) {
+  } else if (args.requires_ordering) {
     std::vector<compute::SortKey> out_sort_keys;
-    std::unordered_set<int> segmented_key_field_id_set(segment_field_ids.begin(),
-                                                       segment_field_ids.end());
-    // Propagate output sorting only by segmented keys excluding sorting by regualr keys
+    std::unordered_set<int> segmented_key_field_id_set(args.segment_key_field_ids.begin(),
+                                                       args.segment_key_field_ids.end());
+    // Propagate output sorting only by segmented keys excluding sorting by regular keys
     // since this will break the segmentation.
-    for (auto key : inputs[0]->ordering().sort_keys()) {
+    for (auto key : input->ordering().sort_keys()) {
       ARROW_ASSIGN_OR_RAISE(auto match, key.target.FindOne(*input_schema));
       if (segmented_key_field_id_set.find(match[0]) != segmented_key_field_id_set.end()) {
         out_sort_keys.emplace_back(key);
       } else {
         break;
       }
     }
-    if (out_sort_keys.size() > 0)
+    if (out_sort_keys.size() > 0) {
       out_ordering = Ordering(out_sort_keys);
-    else
+    } else {
       out_ordering = Ordering::Implicit();
+    }
   }
 
   if (!out_ordering.is_unordered()) {
@@ -189,42 +222,11 @@ ScalarAggregateNode::MakeAggregateNodeArgs(const std::shared_ptr<Schema>& input_
           "order_by node)");
     }
   }
-
-  return AggregateNodeArgs<ScalarAggregateKernel>{
-      schema(std::move(fields)),
-      /*grouping_key_field_ids=*/{}, std::move(segment_field_ids),
-      std::move(segmenter),          std::move(target_fieldsets),
-      std::move(aggregates),         std::move(kernels),
-      std::move(kernel_intypes),     std::move(states),
-      std::move(out_ordering)};
-}
-
-Result<ExecNode*> ScalarAggregateNode::Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
-                                            const ExecNodeOptions& options) {
-  RETURN_NOT_OK(ValidateExecNodeInputs(plan, inputs, 1, "ScalarAggregateNode"));
-
-  const auto& aggregate_options = checked_cast<const AggregateNodeOptions&>(options);
-  auto aggregates = aggregate_options.aggregates;
-  const auto& keys = aggregate_options.keys;
-  const auto& segment_keys = aggregate_options.segment_keys;
-  const auto concurrency = plan->query_context()->max_concurrency();
-
-  if (keys.size() > 0) {
-    return Status::Invalid("Scalar aggregation with some key");
-  }
-
-  const auto& input_schema = inputs[0]->output_schema();
-  auto exec_ctx = plan->query_context()->exec_context();
-
-  ARROW_ASSIGN_OR_RAISE(
-      auto args, MakeAggregateNodeArgs(input_schema, keys, segment_keys, aggregates,
-                                       exec_ctx, concurrency, inputs));
-
   return plan->EmplaceNode<ScalarAggregateNode>(
       plan, std::move(inputs), std::move(args.output_schema), std::move(args.segmenter),
       std::move(args.segment_key_field_ids), std::move(args.target_fieldsets),
       std::move(args.aggregates), std::move(args.kernels), std::move(args.kernel_intypes),
-      std::move(args.states), std::move(args.ordering));
+      std::move(args.states), std::move(out_ordering));
 }
 
 Status ScalarAggregateNode::DoConsume(const ExecSpan& batch, size_t thread_index) {
