refactor(optimizer): Prepare DISTINCT planning

axiom/optimizer/v2/EmitPass.cpp

@@ -872,55 +872,6 @@ velox::core::AggregationNode::Aggregate toVeloxFinalAggregate(
   return out;
 }
 
-// Subset of `groupingKeys` that `local` guarantees are pre-grouped (equal-key
-// rows contiguous), so an aggregation over them can stream rather than build a
-// full hash table. The larger of: the leading `Sorted` run that are grouping
-// keys, and a `Grouped(S)` whose columns are all grouping keys (the whole `S`,
-// since rows are contiguous on the set jointly, not on a bare subset). Empty
-// when nothing is pre-grouped.
-ExprVector computePreGroupedKeys(
-    const LocalPropertyVector& local,
-    const ExprVector& groupingKeys) {
-  if (groupingKeys.empty()) {
-    return {};
-  }
-
-  const auto isGroupingKey = [&](ColumnCP column) {
-    for (ExprCP key : groupingKeys) {
-      if (column->sameOrEqual(*key)) {
-        return true;
-      }
-    }
-    return false;
-  };
-
-  ExprVector best;
-  for (const LocalProperty& property : local) {
-    ExprVector candidate;
-    if (property.kind == LocalPropertyKind::kSorted) {
-      // Sorted on c1..cn implies grouped on any leading prefix; take the
-      // leading run of sort keys that are grouping keys.
-      for (ColumnCP column : property.columns) {
-        if (!isGroupingKey(column)) {
-          break;
-        }
-        candidate.push_back(column);
-      }
-    } else if (std::ranges::all_of(property.columns, isGroupingKey)) {
-      // Grouped on the whole set jointly; usable only if every member is a
-      // grouping key.
-      for (ColumnCP column : property.columns) {
-        candidate.push_back(column);
-      }
-    }
-    if (candidate.size() > best.size()) {
-      best = std::move(candidate);
-    }
-  }
-
-  return best;
-}
-
 // Builds the (output name, Velox Aggregate) lists for an AggregationNode. The
 // per-step call construction differs (single/partial output the result vs the
 // intermediate type; final reads its input accumulator column), so the caller

axiom/optimizer/v2/PhysicalProperties.cpp

@@ -16,6 +16,8 @@
 
 #include "axiom/optimizer/v2/PhysicalProperties.h"
 
+#include <algorithm>
+
 #include <folly/container/F14Map.h>
 
 #include "axiom/connectors/ConnectorMetadata.h"
@@ -58,6 +60,49 @@ AXIOM_DEFINE_ENUM_NAME(PropertyScope, propertyScopeNames);
 AXIOM_DEFINE_ENUM_NAME(PartitionKind, partitionKindNames);
 AXIOM_DEFINE_ENUM_NAME(LocalPropertyKind, localPropertyKindNames);
 
+ExprVector computePreGroupedKeys(
+    const LocalPropertyVector& local,
+    const ExprVector& groupingKeys) {
+  if (groupingKeys.empty()) {
+    return {};
+  }
+
+  const auto isGroupingKey = [&](ColumnCP column) {
+    for (ExprCP key : groupingKeys) {
+      if (column->sameOrEqual(*key)) {
+        return true;
+      }
+    }
+    return false;
+  };
+
+  ExprVector best;
+  for (const LocalProperty& property : local) {
+    ExprVector candidate;
+    if (property.kind == LocalPropertyKind::kSorted) {
+      // Sorted on c1..cn implies grouped on any leading prefix; take the
+      // leading run of sort keys that are grouping keys.
+      for (ColumnCP column : property.columns) {
+        if (!isGroupingKey(column)) {
+          break;
+        }
+        candidate.push_back(column);
+      }
+    } else if (std::ranges::all_of(property.columns, isGroupingKey)) {
+      // Grouped on the whole set jointly; usable only if every member is a
+      // grouping key.
+      for (ColumnCP column : property.columns) {
+        candidate.push_back(column);
+      }
+    }
+    if (candidate.size() > best.size()) {
+      best = std::move(candidate);
+    }
+  }
+
+  return best;
+}
+
 Partitioning Partitioning::globalHash(
     const ExprVector& keys,
     bool replicateNullsAndAny) {

axiom/optimizer/v2/PhysicalProperties.h

@@ -222,6 +222,16 @@ struct LocalProperty {
 /// A relation's per-driver local properties, outermost first.
 using LocalPropertyVector = QGVector<LocalProperty>;
 
+/// Subset of `groupingKeys` that `local` guarantees are pre-grouped (equal-key
+/// rows contiguous), so an aggregation over them can stream rather than build a
+/// full hash table. The larger of: the leading `Sorted` run that are grouping
+/// keys, and a `Grouped(S)` whose columns are all grouping keys (the whole `S`,
+/// since rows are contiguous on the set jointly, not on a bare subset). Empty
+/// when nothing is pre-grouped.
+ExprVector computePreGroupedKeys(
+    const LocalPropertyVector& local,
+    const ExprVector& groupingKeys);
+
 /// A set of columns that is unique across the relation — i.e., functionally
 /// determines the row — at `scope`. Stored minimal: a key-set whose columns are
 /// a superset of another stored key-set is redundant and not kept, but

axiom/optimizer/v2/PlanPhysicalPass.cpp

@@ -765,7 +765,12 @@ class PhysicalPlanRewriter : public NodeRewriter<> {
   }
 
   NodeCP rewriteAggregate(const Aggregate* node, NoContext& context) override {
-    NodeCP input = rewrite(node->input(), context);
+    return planAggregateStages(node, rewrite(node->input(), context));
+  }
+
+  // Selects single-stage or partial/final execution for an Aggregate whose
+  // input is already physically planned.
+  NodeCP planAggregateStages(const Aggregate* node, NodeCP input) {
     if (isSplittableAggregate(node)) {
       // Remote two-stage: the input must shuffle across workers to co-locate
       // its groups, so the partial reduces rows before that remote exchange.
@@ -775,7 +780,7 @@ class PhysicalPlanRewriter : public NodeRewriter<> {
                 input, node->groupingKeys(), Alignment::kCoLocated)) {
           input = grouped;
         } else {
-          return rewriteAggregateSplit(node, input, /*remoteExchange=*/true);
+          return planAggregateSplit(node, input, /*remoteExchange=*/true);
         }
       }
       // Local two-stage: the input is already co-located (e.g. a bucketed
@@ -785,9 +790,14 @@ class PhysicalPlanRewriter : public NodeRewriter<> {
       // The local exchange itself is not materialized here — emit inserts it at
       // numDrivers > 1 (local exchanges are implicit).
       if (numDrivers_ > 1) {
-        return rewriteAggregateSplit(node, input, /*remoteExchange=*/false);
+        return planAggregateSplit(node, input, /*remoteExchange=*/false);
       }
     }
+    return planSingleAggregate(node, input);
+  }
+
+  // Plans a single-stage Aggregate whose input is already physically planned.
+  NodeCP planSingleAggregate(const Aggregate* node, NodeCP input) {
     // A global () grouping set emits a default row over empty input; a
     // single-stage aggregate must gather (empty keys) so that row is produced
     // once, not once per worker.
@@ -896,10 +906,8 @@ class PhysicalPlanRewriter : public NodeRewriter<> {
   // aggregate (e.g. array_agg) gains nothing and pays an extra hash pass; not
   // splitting it needs a reducing/non-reducing classification that does not yet
   // exist, so that pessimization is deferred.
-  NodeCP rewriteAggregateSplit(
-      const Aggregate* node,
-      NodeCP input,
-      bool remoteExchange) {
+  NodeCP
+  planAggregateSplit(const Aggregate* node, NodeCP input, bool remoteExchange) {
     const size_t numKeys = node->groupingKeys().size();
     const auto& finalColumns = node->outputColumns();
 

axiom/optimizer/v2/PrecomputeProjectionsPass.cpp

@@ -359,68 +359,8 @@ class Rewriter : public NodeRewriter<> {
 NodeCP Rewriter::rewriteAggregate(
     const Aggregate* aggregate,
     NoContext& context) {
-  NodeCP newInput = rewrite(aggregate->input(), context);
-  // An Aggregate reads only its grouping keys and aggregate inputs, so the
-  // lifting project outputs just those — dropping any input column kept solely
-  // to feed a lifted aggregate expression.
-  PrecomputeProjections precompute{
-      newInput, builder(), /*projectAllInputs=*/false};
-
-  ExprVector newGroupingKeys;
-  newGroupingKeys.reserve(aggregate->groupingKeys().size());
-  for (size_t i = 0; i < aggregate->groupingKeys().size(); ++i) {
-    // Reuse the existing output column as the projection alias so the
-    // Aggregate's outputColumns identity is preserved.
-    newGroupingKeys.push_back(precompute.toColumn(
-        aggregate->groupingKeys()[i], aggregate->outputColumns()[i]));
-  }
-
-  // A kFinal aggregate's args reference the Partial's raw inputs, which are
-  // absent at the Final's input (it consumes intermediate accumulators), so
-  // leave them untouched rather than precompute them here.
-  AggregateCallVector newAggregates;
-  if (aggregate->step() == AggregateStep::kFinal) {
-    newAggregates = aggregate->aggregates();
-  } else {
-    newAggregates.reserve(aggregate->aggregates().size());
-    for (const auto* call : aggregate->aggregates()) {
-      ExprVector newArgs;
-      newArgs.reserve(call->args().size());
-      for (ExprCP arg : call->args()) {
-        newArgs.push_back(precompute.toColumn(
-            arg, /*alias=*/nullptr, /*allowConstant=*/true));
-      }
-      ExprCP newCondition = call->condition() != nullptr
-          ? precompute.toColumn(
-                call->condition(), /*alias=*/nullptr, /*allowConstant=*/true)
-          : nullptr;
-      ExprVector newOrderKeys;
-      newOrderKeys.reserve(call->orderKeys().size());
-      for (ExprCP key : call->orderKeys()) {
-        newOrderKeys.push_back(precompute.toColumn(key));
-      }
-      newAggregates.push_back(
-          builder().makeAggregate(
-              call->name(),
-              call->value(),
-              std::move(newArgs),
-              call->functions(),
-              call->isDistinct(),
-              newCondition,
-              call->intermediateType(),
-              std::move(newOrderKeys),
-              call->orderTypes()));
-    }
-  }
-
-  return builder().make<Aggregate>(
-      {.input = std::move(precompute).node(),
-       .groupingKeys = std::move(newGroupingKeys),
-       .aggregates = std::move(newAggregates),
-       .outputColumns = aggregate->outputColumns(),
-       .step = aggregate->step(),
-       .groupId = aggregate->groupId(),
-       .globalGroupingSets = aggregate->globalGroupingSets()});
+  return PrecomputeProjectionsPass::prepareAggregateInputs(
+      aggregate, rewrite(aggregate->input(), context), builder());
 }
 
 NodeCP Rewriter::rewriteWindow(const Window* window, NoContext& context) {
@@ -780,6 +720,72 @@ NodeCP Rewriter::rewriteUnionAll(const UnionAll* unionAll, NoContext& context) {
 
 } // namespace
 
+AggregateCP PrecomputeProjectionsPass::prepareAggregateInputs(
+    const Aggregate* aggregate,
+    NodeCP rewrittenInput,
+    Builder& builder) {
+  // An Aggregate reads only its grouping keys and aggregate inputs, so the
+  // lifting project outputs just those — dropping any input column kept solely
+  // to feed a lifted aggregate expression.
+  PrecomputeProjections precompute{
+      rewrittenInput, builder, /*projectAllInputs=*/false};
+
+  ExprVector newGroupingKeys;
+  newGroupingKeys.reserve(aggregate->groupingKeys().size());
+  for (size_t i = 0; i < aggregate->groupingKeys().size(); ++i) {
+    // Reuse the existing output column as the projection alias so the
+    // Aggregate's outputColumns identity is preserved.
+    newGroupingKeys.push_back(precompute.toColumn(
+        aggregate->groupingKeys()[i], aggregate->outputColumns()[i]));
+  }
+
+  // A kFinal aggregate's args reference the Partial's raw inputs, which are
+  // absent at the Final's input (it consumes intermediate accumulators), so
+  // leave them untouched rather than precompute them here.
+  AggregateCallVector newAggregates;
+  if (aggregate->step() == AggregateStep::kFinal) {
+    newAggregates = aggregate->aggregates();
+  } else {
+    newAggregates.reserve(aggregate->aggregates().size());
+    for (const auto* call : aggregate->aggregates()) {
+      ExprVector newArgs;
+      newArgs.reserve(call->args().size());
+      for (ExprCP arg : call->args()) {
+        newArgs.push_back(precompute.toColumn(
+            arg, /*alias=*/nullptr, /*allowConstant=*/true));
+      }
+      ExprCP newCondition = call->condition() != nullptr
+          ? precompute.toColumn(
+                call->condition(), /*alias=*/nullptr, /*allowConstant=*/true)
+          : nullptr;
+      ExprVector newOrderKeys;
+      newOrderKeys.reserve(call->orderKeys().size());
+      for (ExprCP key : call->orderKeys()) {
+        newOrderKeys.push_back(precompute.toColumn(key));
+      }
+      newAggregates.push_back(builder.makeAggregate(
+          call->name(),
+          call->value(),
+          std::move(newArgs),
+          call->functions(),
+          call->isDistinct(),
+          newCondition,
+          call->intermediateType(),
+          std::move(newOrderKeys),
+          call->orderTypes()));
+    }
+  }
+
+  return builder.make<Aggregate>(Aggregate::Key{
+      .input = std::move(precompute).node(),
+      .groupingKeys = std::move(newGroupingKeys),
+      .aggregates = std::move(newAggregates),
+      .outputColumns = aggregate->outputColumns(),
+      .step = aggregate->step(),
+      .groupId = aggregate->groupId(),
+      .globalGroupingSets = aggregate->globalGroupingSets()});
+}
+
 NodeCP PrecomputeProjectionsPass::run(NodeCP node, Builder& builder) {
   return Rewriter{builder}.rewrite(node);
 }

axiom/optimizer/v2/PrecomputeProjectionsPass.h

@@ -24,6 +24,16 @@ namespace facebook::axiom::optimizer::v2 {
 /// Moves expressions a consumer references into a `Project` over its input.
 class PrecomputeProjectionsPass {
  public:
+  /// Lifts grouping expressions, aggregate arguments, filters, and ordering
+  /// keys of 'aggregate' into a Project where required. Aggregate arguments
+  /// comes from 'rewrittenInput' provided by caller in replacement for
+  /// `aggregate->input()`. Returns a new equivalent Aggregate node with lifted
+  /// fields.
+  static AggregateCP prepareAggregateInputs(
+      const Aggregate* aggregate,
+      NodeCP rewrittenInput,
+      Builder& builder);
+
   /// Returns the tree rooted at 'node' rewritten so the expressions listed
   /// below are computed by a `Project` inserted between the consumer and its
   /// input, with the consumer rebuilt to reference the projected column.