diff --git a/axiom/optimizer/v2/EmitPass.cpp b/axiom/optimizer/v2/EmitPass.cpp index 1e5c02be8..e149777b4 100644 --- a/axiom/optimizer/v2/EmitPass.cpp +++ b/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 diff --git a/axiom/optimizer/v2/PhysicalProperties.cpp b/axiom/optimizer/v2/PhysicalProperties.cpp index 6e055bdf8..2ed922bab 100644 --- a/axiom/optimizer/v2/PhysicalProperties.cpp +++ b/axiom/optimizer/v2/PhysicalProperties.cpp @@ -16,6 +16,8 @@ #include "axiom/optimizer/v2/PhysicalProperties.h" +#include + #include #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) { diff --git a/axiom/optimizer/v2/PhysicalProperties.h b/axiom/optimizer/v2/PhysicalProperties.h index dbc66a5b0..49ac0bb02 100644 --- a/axiom/optimizer/v2/PhysicalProperties.h +++ b/axiom/optimizer/v2/PhysicalProperties.h @@ -222,6 +222,16 @@ struct LocalProperty { /// A relation's per-driver local properties, outermost first. using LocalPropertyVector = QGVector; +/// 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 diff --git a/axiom/optimizer/v2/PlanPhysicalPass.cpp b/axiom/optimizer/v2/PlanPhysicalPass.cpp index cf19c861e..c8cfa2742 100644 --- a/axiom/optimizer/v2/PlanPhysicalPass.cpp +++ b/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(); diff --git a/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp b/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp index afb573b0a..331696fe3 100644 --- a/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp +++ b/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( - {.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::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); } diff --git a/axiom/optimizer/v2/PrecomputeProjectionsPass.h b/axiom/optimizer/v2/PrecomputeProjectionsPass.h index b2f8beeae..47518e554 100644 --- a/axiom/optimizer/v2/PrecomputeProjectionsPass.h +++ b/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.