From 9e177b65f01906855d5451cf68a280b021dadf37 Mon Sep 17 00:00:00 2001 From: Wei He Date: Fri, 28 Aug 2026 19:55:31 -0700 Subject: [PATCH 1/3] refactor(optimizer): Prepare DISTINCT planning (#1804) Summary: Extract aggregate input preparation and physical stage selection so DISTINCT lowering can reuse them before aggregate stage selection without recursively replanning inputs. Move the existing pre-grouped-key calculation to shared physical properties without changing its behavior. Differential Revision: D117896114 --- axiom/optimizer/v2/EmitPass.cpp | 49 ------- axiom/optimizer/v2/PhysicalProperties.cpp | 45 ++++++ axiom/optimizer/v2/PhysicalProperties.h | 10 ++ axiom/optimizer/v2/PlanPhysicalPass.cpp | 22 ++- .../v2/PrecomputeProjectionsPass.cpp | 130 +++++++++--------- .../optimizer/v2/PrecomputeProjectionsPass.h | 10 ++ 6 files changed, 148 insertions(+), 118 deletions(-) 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. From 7e23e5e2fd15e6c32fcdcbb0e3d5ba83d2409f9c Mon Sep 17 00:00:00 2001 From: Wei He Date: Fri, 28 Aug 2026 19:55:31 -0700 Subject: [PATCH 2/3] feat(optimizer): Plan DISTINCT with MarkDistinct (#1805) Summary: Distributed v2 queries with `MarkDistinct` now choose marker distribution before selecting split vs. single steps for the aggregation. Previously, the late expansion pass inserted `MarkDistinct` after physical planning, leaving the DISTINCT aggregation to always take single-stage. Moves DISTINCT lowering into physical aggregate planning. Plans with one worker and one driver, or with pre-grouped input, continue using native DISTINCT. Removes the separate late expansion pass. Differential Revision: D117940699 --- .../tests/BucketedExecutionPlanTest.cpp | 6 +- .../tests/DistinctAggregationTest.cpp | 238 ++++++++++-------- axiom/optimizer/v2/CMakeLists.txt | 1 - axiom/optimizer/v2/ExpandAggregatePass.cpp | 234 ----------------- axiom/optimizer/v2/ExpandAggregatePass.h | 39 --- axiom/optimizer/v2/Optimize.h | 4 +- axiom/optimizer/v2/PhysicalPlanAndEmit.cpp | 6 +- axiom/optimizer/v2/PhysicalPlanAndEmit.h | 3 +- axiom/optimizer/v2/PlanPhysicalPass.cpp | 175 ++++++++++++- .../v2/PrecomputeProjectionsPass.cpp | 32 +++ 10 files changed, 343 insertions(+), 395 deletions(-) delete mode 100644 axiom/optimizer/v2/ExpandAggregatePass.cpp delete mode 100644 axiom/optimizer/v2/ExpandAggregatePass.h diff --git a/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp b/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp index 459917e59..cd875de31 100644 --- a/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp +++ b/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp @@ -353,8 +353,10 @@ TEST_P(BucketedExecutionTest, aggregation) { AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("a_orders") - .localPartition({"customer_id"}) - .singleAggregation() + .localPartition({"customer_id", "amount"}) + .markDistinct({"customer_id", "amount"}, {"m0"}) + .localAggregation( + {"customer_id"}, {"count(amount) filter (where m0)"}) .fragment({.width = 4, .bucketedScans = 1}) .gather() .build()); diff --git a/axiom/optimizer/tests/DistinctAggregationTest.cpp b/axiom/optimizer/tests/DistinctAggregationTest.cpp index 5a8acc474..0958c7c85 100644 --- a/axiom/optimizer/tests/DistinctAggregationTest.cpp +++ b/axiom/optimizer/tests/DistinctAggregationTest.cpp @@ -350,10 +350,6 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctDifferentArgSets) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -367,7 +363,7 @@ TEST_P(DistinctAggregationTest, markDistinctDifferentArgSets) { .aggregate({"a"}, {"count(DISTINCT b)", "sum(DISTINCT d % 5)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"a", "b as p0", "d % 5 as p1"}) @@ -378,7 +374,7 @@ TEST_P(DistinctAggregationTest, markDistinctDifferentArgSets) { {"count(p0) filter (where m0)", "sum(p1) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"a", "b", "d % 5 as p0"}) @@ -386,10 +382,6 @@ TEST_P(DistinctAggregationTest, markDistinctDifferentArgSets) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctGlobalWithMultipleSets) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -403,7 +395,7 @@ TEST_P(DistinctAggregationTest, markDistinctGlobalWithMultipleSets) { .aggregate({}, {"count(DISTINCT b)", "sum(DISTINCT d % 5)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"b as p0", "d % 5 as p1"}) @@ -415,7 +407,7 @@ TEST_P(DistinctAggregationTest, markDistinctGlobalWithMultipleSets) { .localGather() .finalAggregation() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"b", "d % 5 as p0"}) @@ -423,10 +415,6 @@ TEST_P(DistinctAggregationTest, markDistinctGlobalWithMultipleSets) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctMixedDistinctAndNonDistinct) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -441,7 +429,7 @@ TEST_P(DistinctAggregationTest, markDistinctMixedDistinctAndNonDistinct) { {"a"}, {"count(DISTINCT b)", "sum(DISTINCT d % 5)", "avg(b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"a", "b as p0", "d % 5 as p1"}) @@ -454,7 +442,7 @@ TEST_P(DistinctAggregationTest, markDistinctMixedDistinctAndNonDistinct) { "avg(p0)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"a", "b", "d % 5 as p0"}) @@ -463,10 +451,81 @@ TEST_P(DistinctAggregationTest, markDistinctMixedDistinctAndNonDistinct) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. +TEST_P(DistinctAggregationTest, markDistinctSingleWorkerMultipleDrivers) { + testConnector_->addTable( + "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); + SCOPE_EXIT { + testConnector_->dropTableIfExists("t"); + }; + + auto logicalPlan = lp::PlanBuilder(makeContext()) + .tableScan("t") + .aggregate({"a"}, {"count(DISTINCT b)", "avg(c)"}) + .build(); + auto plan = planVelox( + logicalPlan, {.numWorkers = 1, .numDrivers = 4}, optimizerOptions_); + AXIOM_ASSERT_DISTRIBUTED_PLAN( + plan.plan, + matchScan("t") + .localPartition({"a", "b"}) + .markDistinct({"a", "b"}, {"m0"}) + .localAggregation({"a"}, {"count(b) filter (where m0)", "avg(c)"}) + .build()); +} + +TEST_P(DistinctAggregationTest, markDistinctAboveWindow) { + testConnector_->addTable("t", ROW("b", BIGINT())); + SCOPE_EXIT { + testConnector_->dropTableIfExists("t"); + }; + + auto logicalPlan = parseSelect( + "SELECT count(DISTINCT rn), sum(b) " + "FROM (" + " SELECT b, row_number() OVER (ORDER BY b + 1) AS rn " + " FROM t" + ")", + kTestConnectorId); + auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); + AXIOM_ASSERT_DISTRIBUTED_PLAN( + plan.plan, + matchScan("t") + .projectIf(!useV2_, {"b", "b + 1 as p0"}) + .gather() + .projectIf(useV2_, {"b", "b + 1 as p0"}) + .localGather() + .window({"row_number() OVER (ORDER BY p0) as rn"}) + .projectIf(!useV2_, {"b", "rn"}) + .localPartition({"rn"}) + .markDistinct({"rn"}, {"m0"}) + .localAggregation({}, {"count(rn) filter (where m0)", "sum(b)"}) + .build()); +} + +TEST_P(DistinctAggregationTest, preGroupedInputUsesNativeDistinct) { + testConnector_->addTable("t", ROW({"a", "b", "c"}, BIGINT())); + SCOPE_EXIT { + testConnector_->dropTableIfExists("t"); + }; + + auto logicalPlan = lp::PlanBuilder(makeContext()) + .tableScan("t") + .orderBy({"a"}) + .limit(100) + .aggregate({"a", "b"}, {"count(DISTINCT c)", "sum(c)"}) + .build(); + auto plan = planVelox( + logicalPlan, {.numWorkers = 1, .numDrivers = 4}, optimizerOptions_); + AXIOM_ASSERT_DISTRIBUTED_PLAN( + plan.plan, + matchScan("t") + .topN(100) + .localMerge() + .finalLimit(0, 100) + .singleAggregation({"a", "b"}, {"count(DISTINCT c)", "sum(c)"}) + .build()); +} + TEST_P(DistinctAggregationTest, markDistinctMultiArgAggregates) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -480,7 +539,7 @@ TEST_P(DistinctAggregationTest, markDistinctMultiArgAggregates) { .aggregate({"a"}, {"covar_pop(DISTINCT b, c)", "count(DISTINCT d)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b", "c"}, {"m0"}) @@ -491,7 +550,7 @@ TEST_P(DistinctAggregationTest, markDistinctMultiArgAggregates) { "count(d) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -499,10 +558,6 @@ TEST_P(DistinctAggregationTest, markDistinctMultiArgAggregates) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -521,7 +576,7 @@ TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { {"count(DISTINCT c)", "covar_pop(DISTINCT b, c)", "sum(c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"b", "c"}, {"m0"}) @@ -532,7 +587,7 @@ TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { "sum(c)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -550,7 +605,7 @@ TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { .aggregate({"b"}, {"covar_pop(DISTINCT b, c)", "count(DISTINCT b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"b", "c"}, {"m0"}) @@ -559,7 +614,7 @@ TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { {"covar_pop(b, c) filter (where m0)", "count(DISTINCT b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -568,10 +623,6 @@ TEST_P(DistinctAggregationTest, markDistinctSharedMarkers) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctOrderBy) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -590,7 +641,7 @@ TEST_P(DistinctAggregationTest, markDistinctOrderBy) { "array_agg(b ORDER BY b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"a", "b as p0", "d % 5 as p1"}) @@ -603,7 +654,7 @@ TEST_P(DistinctAggregationTest, markDistinctOrderBy) { "array_agg(p0 ORDER BY p0 ASC NULLS LAST)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"a", "b as p0", "d % 5 as p1"}) @@ -626,7 +677,7 @@ TEST_P(DistinctAggregationTest, markDistinctOrderBy) { "array_agg(DISTINCT d % 5 ORDER BY d % 5)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"b as p0", "d % 5 as p1"}) @@ -637,7 +688,7 @@ TEST_P(DistinctAggregationTest, markDistinctOrderBy) { {"array_agg(p0 ORDER BY p0 ASC NULLS LAST) filter (where m0)", "array_agg(p1 ORDER BY p1 ASC NULLS LAST) filter (where m1)"}) .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"b as p0", "d % 5 as p1"}) @@ -649,10 +700,6 @@ TEST_P(DistinctAggregationTest, markDistinctOrderBy) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctLiterals) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -669,7 +716,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { .aggregate({"a"}, {"count(DISTINCT b)", "max_by(DISTINCT d, 1)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0"}) @@ -680,7 +727,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { "max_by(d, 1) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -697,7 +744,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { .aggregate({"a"}, {"count(DISTINCT b)", "max_by(DISTINCT a, 1)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0"}) @@ -705,7 +752,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { {"a"}, {"count(b) filter (where m0)", "max_by(DISTINCT a, 1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -722,7 +769,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { .aggregate({"a"}, {"count(DISTINCT b)", "count(DISTINCT 1)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0"}) @@ -730,7 +777,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { {"a"}, {"count(b) filter (where m0)", "count(DISTINCT 1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -749,7 +796,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { {"count(DISTINCT b)", "count(DISTINCT 1) FILTER (WHERE d > 0)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"a", "b", "d > 0 as p0"}) @@ -760,7 +807,7 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { "count(DISTINCT 1) filter (where p0)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"a", "b", "d > 0 as p0"}) @@ -771,10 +818,6 @@ TEST_P(DistinctAggregationTest, markDistinctLiterals) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. // TODO: Track emitted local partitioning so compatible consecutive // MarkDistinct nodes reuse one local exchange. TEST_P(DistinctAggregationTest, multipleMarkDistinctWithNoShuffleInBetween) { @@ -811,10 +854,6 @@ TEST_P(DistinctAggregationTest, multipleMarkDistinctWithNoShuffleInBetween) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { testConnector_->addTable( "t", @@ -834,7 +873,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { "count(DISTINCT b) FILTER (WHERE e)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0", "m1", "m2"}) @@ -843,7 +882,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { {"count(b) filter (where m1)", "count(b) filter (where m2)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -863,7 +902,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { "count(DISTINCT c) FILTER (WHERE e)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0", "m1"}) @@ -873,7 +912,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { {"count(b) filter (where m1)", "count(c) filter (where m3)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -884,10 +923,6 @@ TEST_P(DistinctAggregationTest, markDistinctFilterDifferentArgSets) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterGlobalAggregation) { testConnector_->addTable("t", ROW({"a", "b"}, {BIGINT(), BOOLEAN()})); SCOPE_EXIT { @@ -899,7 +934,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterGlobalAggregation) { .aggregate({}, {"count(DISTINCT a) FILTER (WHERE b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a"}, {"m0", "m1"}) @@ -908,17 +943,13 @@ TEST_P(DistinctAggregationTest, markDistinctFilterGlobalAggregation) { .localGather() .finalAggregation() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation({}, {"count(DISTINCT a) filter (where b)"}) .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), BOOLEAN()})); @@ -936,7 +967,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { "sum(DISTINCT b) FILTER (WHERE c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0", "m1"}) @@ -945,7 +976,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { {"count(b) filter (where m1)", "sum(b) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -965,7 +996,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { {"count(DISTINCT b)", "count(DISTINCT b) FILTER (WHERE c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0", "m1"}) @@ -974,7 +1005,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { {"count(b) filter (where m0)", "count(b) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -984,10 +1015,6 @@ TEST_P(DistinctAggregationTest, markDistinctFilterSharedMarkers) { } } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterOrderBy) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), BOOLEAN()})); @@ -1002,7 +1029,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterOrderBy) { {"a"}, {"array_agg(DISTINCT b ORDER BY b) FILTER (WHERE c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0", "m1"}) @@ -1011,7 +1038,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterOrderBy) { {"array_agg(b ORDER BY b ASC NULLS LAST) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -1019,10 +1046,6 @@ TEST_P(DistinctAggregationTest, markDistinctFilterOrderBy) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterMixedDistinctAndNonDistinct) { testConnector_->addTable( "t", @@ -1040,7 +1063,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterMixedDistinctAndNonDistinct) { {"sum(b) FILTER (WHERE e)", "count(DISTINCT c) FILTER (WHERE d)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "c"}, {"m0", "m1"}) @@ -1048,7 +1071,7 @@ TEST_P(DistinctAggregationTest, markDistinctFilterMixedDistinctAndNonDistinct) { {"a"}, {"sum(b) filter (where e)", "count(c) filter (where m1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -1086,10 +1109,6 @@ TEST_P(DistinctAggregationTest, markDistinctFilterRedundantKeys) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P(DistinctAggregationTest, markDistinctFilterExpressionCondition) { testConnector_->addTable( "t", @@ -1109,20 +1128,24 @@ TEST_P(DistinctAggregationTest, markDistinctFilterExpressionCondition) { "count(DISTINCT b) FILTER (WHERE c > 0.0)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") - .project({"a", "d", "b", "c > 0.0"}) + .project( + useV2_ ? std::vector{"a", "b", "d", "c > 0.0"} + : std::vector{"a", "d", "b", "c > 0.0"}) .distributedMarkDistinct({"a", "b"}, {"m0", "m1", "m2"}) .distributedAggregation( {"a"}, {"count(b) filter (where m1)", "count(b) filter (where m2)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") - .project({"a", "d", "b", "c > 0.0 as p0"}) + .project( + useV2_ ? std::vector{"a", "b", "d", "c > 0.0 as p0"} + : std::vector{"a", "d", "b", "c > 0.0 as p0"}) .singleAggregation( {"a"}, {"count(DISTINCT b) filter (where d)", @@ -1130,10 +1153,6 @@ TEST_P(DistinctAggregationTest, markDistinctFilterExpressionCondition) { .build()); } -// V1 is better: it plans MarkDistinct distribution before selecting the outer -// Aggregate stages, enabling distributed deduplication and partial aggregation. -// TODO: Make V2 lower DISTINCT-to-MarkDistinct inside physical aggregation -// planning. TEST_P( DistinctAggregationTest, markDistinctAllLiteralDistinctMixColumnDistinct) { @@ -1152,14 +1171,14 @@ TEST_P( .aggregate({}, {"count(DISTINCT b)", "count(DISTINCT 1)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"b"}, {"m0"}) .distributedSingleAggregation( {}, {"count(b) filter (where m0)", "count(DISTINCT 1)"}) .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation({}, {"count(DISTINCT b)", "count(DISTINCT 1)"}) @@ -1175,7 +1194,7 @@ TEST_P( .aggregate({"a"}, {"count(DISTINCT b)", "count(DISTINCT 1)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedMarkDistinct({"a", "b"}, {"m0"}) @@ -1183,7 +1202,7 @@ TEST_P( {"a"}, {"count(b) filter (where m0)", "count(DISTINCT 1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -1235,11 +1254,10 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctToGroupBy) { } } -// V1 is better: it produces a valid grouping-set plan. V2 generates duplicate -// output column `a1` and fails during plan validation. +// V2 currently rejects this plan because GroupId emits grouping key `a1`, +// which collides with the aggregate result carrying that name. // TODO: Allocate collision-proof optimizer-owned output names for GroupId -// grouping-key outputs in V2 TranslatePass. Afterward, apply the shared V2 -// DISTINCT-to-MarkDistinct physical-planning fix. +// grouping-key outputs in V2 TranslatePass. TEST_P(DistinctAggregationTest, groupingSetsDistinctToMarkDistinct) { testConnector_->addTable("t", ROW({"a", "b", "c"}, BIGINT())); SCOPE_EXIT { diff --git a/axiom/optimizer/v2/CMakeLists.txt b/axiom/optimizer/v2/CMakeLists.txt index 3446899b2..5d8e8e5d0 100644 --- a/axiom/optimizer/v2/CMakeLists.txt +++ b/axiom/optimizer/v2/CMakeLists.txt @@ -27,7 +27,6 @@ add_library( EmitPass.cpp EstimateLeafStatsPass.cpp EstimateProvider.cpp - ExpandAggregatePass.cpp ExprEmitter.cpp ExprFactory.cpp ExprSimplifier.cpp diff --git a/axiom/optimizer/v2/ExpandAggregatePass.cpp b/axiom/optimizer/v2/ExpandAggregatePass.cpp deleted file mode 100644 index c2b3aaeec..000000000 --- a/axiom/optimizer/v2/ExpandAggregatePass.cpp +++ /dev/null @@ -1,234 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and its affiliates. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "axiom/optimizer/v2/ExpandAggregatePass.h" - -#include - -#include - -#include "axiom/optimizer/PlanObject.h" -#include "axiom/optimizer/v2/AppendAll.h" -#include "axiom/optimizer/v2/NodeRewriter.h" - -namespace facebook::axiom::optimizer::v2 { - -namespace { - -// Appends each non-literal column reference in 'args' to 'keys' if not -// already present. Used to build a MarkDistinct key set as -// `groupingKeys ∪ aggregate.args()`. Literals contribute nothing. -ExprVector unionColumnArgs(const ExprVector& keys, const ExprVector& args) { - ExprVector merged = keys; - PlanObjectSet seen = PlanObjectSet::fromObjects(keys); - for (ExprCP arg : args) { - if (arg->is(PlanType::kLiteralExpr)) { - continue; - } - VELOX_CHECK( - arg->is(PlanType::kColumnExpr), - "Expected column or literal aggregate arg: {}", - arg->toString()); - if (seen.contains(arg)) { - continue; - } - seen.add(arg); - merged.push_back(arg); - } - return merged; -} - -// True when every aggregate shares a single distinct signature: all DISTINCT, -// no FILTER, no ORDER BY, and the same column arguments. Velox then dedups them -// in one native distinct aggregation pass (aggregates keep `distinct=true`). -// Any other mix needs MarkDistinct. -bool canUseNativeDistinct(const AggregateCallVector& aggregates) { - if (aggregates.empty()) { - return false; - } - std::optional commonArgs; - for (const auto* aggregate : aggregates) { - if (!aggregate->isDistinct() || aggregate->condition() != nullptr || - !aggregate->orderKeys().empty()) { - return false; - } - PlanObjectSet columnArgs = - PlanObjectSet::fromObjects(unionColumnArgs({}, aggregate->args())); - if (!commonArgs.has_value()) { - commonArgs = std::move(columnArgs); - } else if (columnArgs != *commonArgs) { - return false; - } - } - return true; -} - -// One MarkDistinct group: all distinct aggregates whose key set -// `(groupingKeys ∪ args)` equals 'keys'. Within a group, each unique FILTER -// condition gets its own per-mask marker; aggregates with no FILTER share -// `markers[0]` (the no-mask marker). -struct MarkDistinctGroup { - ExprVector keys; - ColumnVector markers; - ColumnVector masks; - // Maps each FILTER condition to the marker that records first occurrence - // among rows where that condition is true. `nullptr` keys the no-mask - // marker (`markers[0]`). - folly::F14FastMap filterToMarker; -}; - -struct DistinctExpansion { - NodeCP input; - AggregateCallVector aggregates; -}; - -class AggregateExpander : public NodeRewriter<> { - public: - explicit AggregateExpander(Builder& builder) : NodeRewriter(builder) {} - - protected: - NodeCP rewriteAggregate(const Aggregate* node, NoContext& context) override { - NodeCP newInput = rewrite(node->input(), context); - - // Grouping-set lowering already ran in translate, so the group-id column - // (if any) is one of the grouping keys and MarkDistinct dedup is per - // grouping set. - AggregateCallVector aggregates = node->aggregates(); - DistinctExpansion distinct = - expandDistinct(newInput, node->groupingKeys(), aggregates); - newInput = distinct.input; - aggregates = std::move(distinct.aggregates); - - if (newInput == node->input() && aggregates == node->aggregates()) { - return node; - } - return builder().make( - {.input = newInput, - .groupingKeys = node->groupingKeys(), - .aggregates = std::move(aggregates), - .outputColumns = node->outputColumns(), - .step = node->step(), - .groupId = node->groupId(), - .globalGroupingSets = node->globalGroupingSets()}); - } - - private: - // Lowers DISTINCT aggregates. When they share a single distinct signature - // (see `canUseNativeDistinct`) the aggregates are left as-is for Velox's - // native distinct aggregation. Otherwise each unique `(groupingKeys ∪ args)` - // set gets a `MarkDistinct` and its aggregates are rewritten as non-distinct - // with the marker as their FILTER. Distinct aggregates whose args ⊆ - // `groupingKeys` are redundant (GROUP BY already dedups) and keep a native - // distinct flag without a marker. - DistinctExpansion expandDistinct( - NodeCP input, - const ExprVector& groupingKeys, - const AggregateCallVector& aggregates) { - if (canUseNativeDistinct(aggregates)) { - return {input, aggregates}; - } - - PlanObjectSet groupingKeySet = PlanObjectSet::fromObjects(groupingKeys); - folly::F14VectorMap groups; - folly::F14FastMap aggregateToMarker; - - bool anyDistinct = false; - for (const auto* aggregate : aggregates) { - if (!aggregate->isDistinct()) { - continue; - } - anyDistinct = true; - - ExprVector keys = unionColumnArgs(groupingKeys, aggregate->args()); - PlanObjectSet keySet = PlanObjectSet::fromObjects(keys); - if (keySet == groupingKeySet) { - continue; - } - - auto [groupIt, isNewGroup] = groups.try_emplace(keySet); - auto& group = groupIt->second; - if (isNewGroup) { - group.keys = std::move(keys); - group.markers.push_back(Column::createBoolean("mark")); - group.filterToMarker[nullptr] = group.markers.back(); - } - - ExprCP filter = aggregate->condition(); - auto [filterIt, isNewFilter] = - group.filterToMarker.try_emplace(filter, nullptr); - if (isNewFilter) { - group.markers.push_back(Column::createBoolean("mark")); - filterIt->second = group.markers.back(); - - ColumnCP maskColumn = filter->as(); - VELOX_CHECK_NOT_NULL( - maskColumn, - "MarkDistinct mask must be a Column reference; got: {}", - filter->toString()); - group.masks.push_back(maskColumn); - } - aggregateToMarker[aggregate] = filterIt->second; - } - - if (!anyDistinct) { - return {input, aggregates}; - } - - NodeCP currentInput = input; - // F14VectorMap iterates in LIFO; reverse to keep insertion order so the - // first encountered key set sits closest to the original input. - for (auto it = groups.rbegin(); it != groups.rend(); ++it) { - auto& group = it->second; - ColumnVector outputColumns; - outputColumns.reserve( - currentInput->outputColumns().size() + group.markers.size()); - appendAll(outputColumns, currentInput->outputColumns()); - appendAll(outputColumns, group.markers); - - currentInput = builder().make({ - currentInput, - group.markers, - group.keys, - group.masks, - std::move(outputColumns), - }); - } - - AggregateCallVector newAggregates; - newAggregates.reserve(aggregates.size()); - for (const auto* aggregate : aggregates) { - if (auto it = aggregateToMarker.find(aggregate); - it != aggregateToMarker.end()) { - newAggregates.push_back( - aggregate->replaceDistinctAndFilterByMarker(it->second)); - } else { - // Either non-distinct, or distinct whose args are all in - // `groupingKeys` (per-group dedup is implicit; Velox handles - // `distinct=true` natively for the trivial case). - newAggregates.push_back(aggregate); - } - } - return {currentInput, std::move(newAggregates)}; - } -}; - -} // namespace - -NodeCP ExpandAggregatePass::run(NodeCP root, Builder& builder) { - AggregateExpander expander{builder}; - return expander.rewrite(root); -} - -} // namespace facebook::axiom::optimizer::v2 diff --git a/axiom/optimizer/v2/ExpandAggregatePass.h b/axiom/optimizer/v2/ExpandAggregatePass.h deleted file mode 100644 index 973c7b99c..000000000 --- a/axiom/optimizer/v2/ExpandAggregatePass.h +++ /dev/null @@ -1,39 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and its affiliates. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "axiom/optimizer/v2/Builder.h" -#include "axiom/optimizer/v2/Node.h" - -namespace facebook::axiom::optimizer::v2 { - -/// Lowers DISTINCT aggregates into a shape Velox can execute. -class ExpandAggregatePass { - public: - /// Returns `root` with each DISTINCT aggregate lowered to - /// `Aggregate(MarkDistinct(input, ...))`, where `MarkDistinct` emits a - /// boolean marker per unique `(groupingKeys, args)` tuple and each - /// `agg(DISTINCT x)` is rewritten to `agg(x) FILTER (marker)`. Aggregates - /// that share one distinct signature keep native Velox distinct instead. The - /// group-id column of a grouping-set aggregate is already one of the grouping - /// keys (grouping sets are lowered to `GroupId` in translate), so distinct - /// dedup is per grouping set. Other node shapes pass through unchanged. Runs - /// after `PlanPhysicalPass` and `PrecomputeProjectionsPass`, before - /// `EmitPass`. - static NodeCP run(NodeCP root, Builder& builder); -}; - -} // namespace facebook::axiom::optimizer::v2 diff --git a/axiom/optimizer/v2/Optimize.h b/axiom/optimizer/v2/Optimize.h index c92ed1ebb..5d6990a68 100644 --- a/axiom/optimizer/v2/Optimize.h +++ b/axiom/optimizer/v2/Optimize.h @@ -83,11 +83,11 @@ class Optimizer { /// - PushdownAndPrune — push predicates down, prune unused columns; /// - EstimateLeafStats — populate base-table cardinalities from the /// connector; - /// - PlanPhysical — cost-based join order and distribution; + /// - PlanPhysical — cost-based join order and distribution, including + /// lowering distinct aggregates when needed; /// - PrecomputeProjections — lift compound expressions into `Project`s /// where /// Velox needs a column or literal; - /// - ExpandAggregate — lower distinct aggregates to `MarkDistinct`; /// - Emit — lower to Velox `PlanNode`s. /// /// `options.numWorkers` / `options.numDrivers` (each >= 1) are the target diff --git a/axiom/optimizer/v2/PhysicalPlanAndEmit.cpp b/axiom/optimizer/v2/PhysicalPlanAndEmit.cpp index a8e7d61c0..6c5ed0e0e 100644 --- a/axiom/optimizer/v2/PhysicalPlanAndEmit.cpp +++ b/axiom/optimizer/v2/PhysicalPlanAndEmit.cpp @@ -16,7 +16,6 @@ #include "axiom/optimizer/v2/PhysicalPlanAndEmit.h" -#include "axiom/optimizer/v2/ExpandAggregatePass.h" #include "axiom/optimizer/v2/PlanPhysicalPass.h" #include "axiom/optimizer/v2/PrecomputeProjectionsPass.h" @@ -33,11 +32,8 @@ EmitPass::Result physicalPlanAndEmit( NodeCP physicalPlanned = PlanPhysicalPass::run( root, builder, session.options(), options.numWorkers, options.numDrivers); NodeCP precomputed = PrecomputeProjectionsPass::run(physicalPlanned, builder); - // Distinct aggregates lower to MarkDistinct here, after physical planning - // (grouping sets were already lowered to GroupId in translate). - NodeCP expanded = ExpandAggregatePass::run(precomputed, builder); return EmitPass::run( - expanded, outputColumns, outputNames, session, evaluator, options); + precomputed, outputColumns, outputNames, session, evaluator, options); } } // namespace facebook::axiom::optimizer::v2 diff --git a/axiom/optimizer/v2/PhysicalPlanAndEmit.h b/axiom/optimizer/v2/PhysicalPlanAndEmit.h index 91736db76..d3265f313 100644 --- a/axiom/optimizer/v2/PhysicalPlanAndEmit.h +++ b/axiom/optimizer/v2/PhysicalPlanAndEmit.h @@ -22,7 +22,8 @@ namespace facebook::axiom::optimizer::v2 { /// Runs the physical-planning and emit passes over 'root' and returns the emit -/// result: PlanPhysical -> PrecomputeProjections -> ExpandAggregate -> Emit. +/// result: PlanPhysical -> PrecomputeProjections -> Emit. PlanPhysical lowers +/// DISTINCT aggregates as part of selecting their physical execution. /// Shared by full optimization and the translate-time constant fold so the two /// cannot drift. 'outputColumns' / 'outputNames' pin the emitted output layout. EmitPass::Result physicalPlanAndEmit( diff --git a/axiom/optimizer/v2/PlanPhysicalPass.cpp b/axiom/optimizer/v2/PlanPhysicalPass.cpp index c8cfa2742..4ea726de6 100644 --- a/axiom/optimizer/v2/PlanPhysicalPass.cpp +++ b/axiom/optimizer/v2/PlanPhysicalPass.cpp @@ -32,6 +32,7 @@ #include "axiom/optimizer/v2/JoinCluster.h" #include "axiom/optimizer/v2/JoinTreeEmitter.h" #include "axiom/optimizer/v2/NodeRewriter.h" +#include "axiom/optimizer/v2/PrecomputeProjectionsPass.h" namespace facebook::axiom::optimizer::v2 { @@ -188,6 +189,123 @@ bool satisfies( VELOX_UNREACHABLE(); } +// Returns true when any result call requires DISTINCT semantics. +bool hasDistinct(const Aggregate& aggregate) { + return std::ranges::any_of(aggregate.aggregates(), [](const auto* call) { + return call->isDistinct(); + }); +} + +// Appends each non-literal argument Column not already present in `keys`. +ExprVector unionColumnArgs(const ExprVector& keys, const ExprVector& args) { + ExprVector merged = keys; + PlanObjectSet seen = PlanObjectSet::fromObjects(keys); + for (ExprCP arg : args) { + if (arg->is(PlanType::kLiteralExpr)) { + continue; + } + VELOX_CHECK( + arg->is(PlanType::kColumnExpr), + "Expected column or literal aggregate arg: {}", + arg->toString()); + if (seen.contains(arg)) { + continue; + } + seen.add(arg); + merged.push_back(arg); + } + return merged; +} + +// Describes one MarkDistinct node independently of its physical input. +struct MarkDistinctSpec { + ExprVector keys; + ColumnVector markers; + ColumnVector masks; +}; + +// Carries the marker groups and rewritten calls for one result Aggregate. +struct MarkDistinctLowering { + std::vector groups; + AggregateCallVector rewrittenCalls; +}; + +// Accumulates marker sharing while DISTINCT calls are analyzed. +struct MarkDistinctGroupState { + ExprVector keys; + ColumnVector markers; + ColumnVector masks; + folly::F14FastMap filterToMarker; +}; + +// Groups DISTINCT calls by marker keys and rewrites them to consume markers. +MarkDistinctLowering analyzeMarkDistinct(const Aggregate& aggregate) { + const PlanObjectSet groupingKeySet = + PlanObjectSet::fromObjects(aggregate.groupingKeys()); + folly::F14VectorMap groups; + folly::F14FastMap aggregateToMarker; + + for (const auto* call : aggregate.aggregates()) { + if (!call->isDistinct()) { + continue; + } + + ExprVector keys = unionColumnArgs(aggregate.groupingKeys(), call->args()); + const PlanObjectSet keySet = PlanObjectSet::fromObjects(keys); + if (keySet == groupingKeySet) { + continue; + } + + auto [groupIt, isNewGroup] = groups.try_emplace(keySet); + auto& group = groupIt->second; + if (isNewGroup) { + group.keys = std::move(keys); + group.markers.push_back(Column::createBoolean("__mark")); + group.filterToMarker[nullptr] = group.markers.back(); + } + + ExprCP filter = call->condition(); + auto [filterIt, isNewFilter] = + group.filterToMarker.try_emplace(filter, nullptr); + if (isNewFilter) { + group.markers.push_back(Column::createBoolean("__mark")); + filterIt->second = group.markers.back(); + + ColumnCP maskColumn = filter->as(); + VELOX_CHECK_NOT_NULL( + maskColumn, + "MarkDistinct mask must be a Column reference; got: {}", + filter->toString()); + group.masks.push_back(maskColumn); + } + aggregateToMarker[call] = filterIt->second; + } + + MarkDistinctLowering lowering; + lowering.groups.reserve(groups.size()); + // F14VectorMap iterates in LIFO order. Reverse it so the first encountered + // key set is planned closest to the original input. + for (auto it = groups.rbegin(); it != groups.rend(); ++it) { + auto& group = it->second; + lowering.groups.push_back( + MarkDistinctSpec{ + .keys = std::move(group.keys), + .markers = std::move(group.markers), + .masks = std::move(group.masks)}); + } + + lowering.rewrittenCalls.reserve(aggregate.aggregates().size()); + for (const auto* call : aggregate.aggregates()) { + if (auto it = aggregateToMarker.find(call); it != aggregateToMarker.end()) { + lowering.rewrittenCalls.push_back( + call->replaceDistinctAndFilterByMarker(it->second)); + } else { + lowering.rewrittenCalls.push_back(call); + } + } + return lowering; +} + // True when regrouping the scans under 'node' can make it bucketed. Mirrors // GroupedScanRewriter's two stopping rules -- only a scan of a bucketed table // contributes, and nothing past an exchange does -- but reads the tree instead @@ -765,7 +883,62 @@ class PhysicalPlanRewriter : public NodeRewriter<> { } NodeCP rewriteAggregate(const Aggregate* node, NoContext& context) override { - return planAggregateStages(node, rewrite(node->input(), context)); + NodeCP input = rewrite(node->input(), context); + if (!hasDistinct(*node) || (numWorkers_ == 1 && numDrivers_ == 1)) { + return planAggregateStages(node, input); + } + + AggregateCP prepared = PrecomputeProjectionsPass::prepareAggregateInputs( + node, input, builder()); + if (!computePreGroupedKeys( + prepared->input()->physicalProperties().local, + prepared->groupingKeys()) + .empty()) { + return planAggregateStages(prepared, prepared->input()); + } + + return planDistinctToMarkDistinct(prepared); + } + + // Plans the MarkDistinct groups and the rewritten result Aggregate without + // recursively rewriting any generated node. + NodeCP planDistinctToMarkDistinct(const Aggregate* aggregate) { + MarkDistinctLowering lowering = analyzeMarkDistinct(*aggregate); + // No marker group is needed when every DISTINCT call's non-literal + // arguments are already grouping keys, including literal-only DISTINCT. + if (lowering.groups.empty()) { + return planAggregateStages(aggregate, aggregate->input()); + } + + NodeCP input = aggregate->input(); + for (const auto& group : lowering.groups) { + input = planMarkDistinct(input, group); + } + + AggregateCP outer = builder().make(Aggregate::Key{ + .input = input, + .groupingKeys = aggregate->groupingKeys(), + .aggregates = std::move(lowering.rewrittenCalls), + .outputColumns = aggregate->outputColumns(), + .step = aggregate->step(), + .groupId = aggregate->groupId(), + .globalGroupingSets = aggregate->globalGroupingSets()}); + return planAggregateStages(outer, input); + } + + // Co-locates a marker group's input across workers and constructs the + // corresponding MarkDistinct node. + NodeCP planMarkDistinct(NodeCP input, const MarkDistinctSpec& spec) { + auto [coLocated, keys] = ensureCoLocated(input, spec.keys); + ColumnVector outputColumns{coLocated->outputColumns()}; + outputColumns.insert( + outputColumns.end(), spec.markers.begin(), spec.markers.end()); + return builder().make(MarkDistinct::Key{ + .input = coLocated, + .markers = spec.markers, + .distinctKeys = std::move(keys), + .masks = spec.masks, + .outputColumns = std::move(outputColumns)}); } // Selects single-stage or partial/final execution for an Aggregate whose diff --git a/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp b/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp index 331696fe3..a187f3f2d 100644 --- a/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp +++ b/axiom/optimizer/v2/PrecomputeProjectionsPass.cpp @@ -338,6 +338,9 @@ class Rewriter : public NodeRewriter<> { protected: NodeCP rewriteAggregate(const Aggregate* aggregate, NoContext& context) override; + NodeCP rewriteMarkDistinct( + const MarkDistinct* markDistinct, + NoContext& context) override; NodeCP rewriteWindow(const Window* window, NoContext& context) override; NodeCP rewriteRowNumber(const RowNumber* rowNumber, NoContext& context) override; @@ -363,6 +366,35 @@ NodeCP Rewriter::rewriteAggregate( aggregate, rewrite(aggregate->input(), context), builder()); } +NodeCP Rewriter::rewriteMarkDistinct( + const MarkDistinct* markDistinct, + NoContext& context) { + NodeCP rewrittenInput = rewrite(markDistinct->input(), context); + const PlanObjectSet rewrittenInputColumns = + PlanObjectSet::fromObjects(rewrittenInput->outputColumns()); + PrecomputeProjections precompute{ + rewrittenInput, builder(), /*projectAllInputs=*/true}; + + ExprVector distinctKeys; + distinctKeys.reserve(markDistinct->distinctKeys().size()); + for (ExprCP key : markDistinct->distinctKeys()) { + distinctKeys.push_back(precompute.toColumn(key)); + } + + NodeCP finalInput = std::move(precompute).node(); + const PlanObjectSet finalInputColumns = + PlanObjectSet::fromObjects(finalInput->outputColumns()); + return builder().make(MarkDistinct::Key{ + .input = finalInput, + .markers = markDistinct->markers(), + .distinctKeys = std::move(distinctKeys), + .masks = markDistinct->masks(), + .outputColumns = replacePrefix( + markDistinct->outputColumns(), + markDistinct->input()->outputColumns().size(), + finalInput->outputColumns())}); +} + NodeCP Rewriter::rewriteWindow(const Window* window, NoContext& context) { NodeCP newInput = rewrite(window->input(), context); PrecomputeProjections precompute{newInput, builder()}; From 35e52f852795df7ace1ef5882fb09750ba4e7699 Mon Sep 17 00:00:00 2001 From: Wei He Date: Fri, 28 Aug 2026 19:55:31 -0700 Subject: [PATCH 3/3] feat(optimizer): Plan DISTINCT aggregation with nested GroupBy Summary: V2 now plans aggregates that share one unfiltered DISTINCT argument set as an inner grouping Aggregate followed by a regular result Aggregate. This reduces rows before the final grouping exchange instead of carrying marker columns through the result aggregation. The inner and outer Aggregates select their physical stages independently. Other DISTINCT shapes continue to use `MarkDistinct`, while single-driver and pre-grouped inputs retain native DISTINCT. This PR also fixes a correctness bug in V1 where a keyless inner Aggregate used to change the empty-input result. Differential Revision: D117954163 --- axiom/optimizer/AggregationPlanner.cpp | 13 +++ .../tests/BucketedExecutionPlanTest.cpp | 6 +- .../tests/DistinctAggregationTest.cpp | 99 +++++++++---------- .../tests/sql/distinctAggregation.sql | 3 + axiom/optimizer/v2/PlanPhysicalPass.cpp | 87 ++++++++++++++++ 5 files changed, 153 insertions(+), 55 deletions(-) diff --git a/axiom/optimizer/AggregationPlanner.cpp b/axiom/optimizer/AggregationPlanner.cpp index 27a7f6945..4072f282b 100644 --- a/axiom/optimizer/AggregationPlanner.cpp +++ b/axiom/optimizer/AggregationPlanner.cpp @@ -619,6 +619,19 @@ std::pair AggregationPlanner::makeDistinctAggregation( "groupId must be one of the groupingKeys when present"); if (auto distinctArgs = getCommonDistinctArgs(aggregates)) { + if (distinctArgs->empty() && groupingKeys.empty()) { + // A keyless inner aggregation would create a row on empty input and + // change the result aggregation's empty-input semantics. + return makeSingleAggregationPlan( + std::move(plan), + groupingKeys, + aggregates, + aggPlan->intermediateColumns(), + aggPlan->columns(), + std::move(globalGroupingSets), + groupId, + state); + } return makeDistinctToGroupByPlan( std::move(plan), groupingKeys, diff --git a/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp b/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp index cd875de31..65b474dec 100644 --- a/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp +++ b/axiom/optimizer/tests/BucketedExecutionPlanTest.cpp @@ -353,10 +353,8 @@ TEST_P(BucketedExecutionTest, aggregation) { AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("a_orders") - .localPartition({"customer_id", "amount"}) - .markDistinct({"customer_id", "amount"}, {"m0"}) - .localAggregation( - {"customer_id"}, {"count(amount) filter (where m0)"}) + .localAggregation({"customer_id", "amount"}, {}) + .localAggregation({"customer_id"}, {"count(amount)"}) .fragment({.width = 4, .bucketedScans = 1}) .gather() .build()); diff --git a/axiom/optimizer/tests/DistinctAggregationTest.cpp b/axiom/optimizer/tests/DistinctAggregationTest.cpp index 0958c7c85..f97b70228 100644 --- a/axiom/optimizer/tests/DistinctAggregationTest.cpp +++ b/axiom/optimizer/tests/DistinctAggregationTest.cpp @@ -53,10 +53,6 @@ class DistinctAggregationTest : public test::QueryTestBase, // 1. Inner: GROUP BY (original_keys + distinct_args) - for deduplication // 2. Outer: Regular aggregation without DISTINCT flag // This avoids the overhead of tracking distinct values in each aggregate. -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, singleDistinctToGroupBy) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); @@ -73,13 +69,13 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupBy) { .aggregate({}, {"count(DISTINCT b)", "covar_pop(DISTINCT b, b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"b"}, {}) .distributedAggregation({}, {"count(b)", "covar_pop(b, b)"}) .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -94,14 +90,14 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupBy) { .aggregate({"a"}, {"count(DISTINCT b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a", "b"}, {}) .distributedAggregation({"a"}, {"count(b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t").singleAggregation({"a"}, {"count(DISTINCT b)"}).build()); } @@ -114,14 +110,14 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupBy) { .aggregate({"a"}, {"count(DISTINCT b)", "covar_pop(DISTINCT b, b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a", "b"}, {}) .distributedAggregation({"a"}, {"count(b)", "covar_pop(b, b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -130,10 +126,6 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupBy) { } } -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); @@ -150,7 +142,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { {"a + 1"}, {"count(DISTINCT b + c)", "sum(DISTINCT b + c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .project({"a + 1 as p0", "b + c as p1"}) @@ -158,7 +150,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { .distributedAggregation({"p0"}, {"count(p1)", "sum(p1)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .project({"a + 1 as p0", "b + c as p1"}) @@ -178,7 +170,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { {"covar_pop(DISTINCT b, c)", "covar_samp(DISTINCT c, b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a", "b", "c"}, {}) @@ -186,7 +178,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { {"a"}, {"covar_pop(b, c)", "covar_samp(c, b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -202,14 +194,14 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { .aggregate({"b"}, {"covar_pop(DISTINCT b, c)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"b", "c"}, {}) .distributedAggregation({"b"}, {"covar_pop(b, c)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation({"b"}, {"covar_pop(DISTINCT b, c)"}) @@ -217,10 +209,6 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithExpressionInputs) { } } -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); @@ -239,7 +227,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { "min_by(DISTINCT a, b ORDER BY b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"c", "a", "b"}, {}) @@ -247,7 +235,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { {"c"}, {"max_by(a, b ORDER BY a)", "min_by(a, b ORDER BY b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -267,7 +255,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { "min_by(DISTINCT b, 2 ORDER BY b)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a", "b"}, {}) @@ -275,7 +263,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { {"a"}, {"max_by(b, 1 ORDER BY b)", "min_by(b, 2 ORDER BY b)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -286,10 +274,6 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithOrderBy) { } } -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { testConnector_->addTable( "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); @@ -306,14 +290,14 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { {"a"}, {"max_by(DISTINCT b, 1)", "min_by(DISTINCT b, 2)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a", "b"}, {}) .distributedAggregation({"a"}, {"max_by(b, 1)", "min_by(b, 2)"}) .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -332,7 +316,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { .aggregate({"a"}, {"count(DISTINCT 1)", "count(DISTINCT 2)"}) .build(); auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .distributedAggregation({"a"}, {}) @@ -341,7 +325,7 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { .finalAggregation() .shuffle() .build()); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( toSingleNodePlan(logicalPlan), matchScan("t") .singleAggregation( @@ -350,6 +334,29 @@ TEST_P(DistinctAggregationTest, singleDistinctToGroupByWithLiterals) { } } +TEST_P( + DistinctAggregationTest, + globalAllLiteralDistinctKeepsNativeAggregation) { + testConnector_->addTable( + "t", ROW({"a", "b", "c"}, {BIGINT(), DOUBLE(), DOUBLE()})); + SCOPE_EXIT { + testConnector_->dropTableIfExists("t"); + }; + + auto logicalPlan = + lp::PlanBuilder(makeContext()) + .tableScan("t") + .aggregate({}, {"count(DISTINCT 1)", "sum(DISTINCT 1)"}) + .build(); + auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); + AXIOM_ASSERT_DISTRIBUTED_PLAN( + plan.plan, + matchScan("t") + .distributedSingleAggregation( + {}, {"count(DISTINCT 1)", "sum(DISTINCT 1)"}) + .build()); +} + TEST_P(DistinctAggregationTest, markDistinctDifferentArgSets) { testConnector_->addTable( "t", ROW({"a", "b", "c", "d"}, {BIGINT(), DOUBLE(), DOUBLE(), BIGINT()})); @@ -1211,10 +1218,6 @@ TEST_P( } } -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, groupingSetsDistinctToGroupBy) { testConnector_->addTable("t", ROW({"a", "b"}, {BIGINT(), BIGINT()})); SCOPE_EXIT { @@ -1230,7 +1233,7 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctToGroupBy) { // aggregation is forced split (partial + final). { auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .groupId({{"a"}, {}}, {"b"}, "gid") @@ -1244,7 +1247,7 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctToGroupBy) { // Single-node plan: a single aggregation computes DISTINCT natively. { auto plan = toSingleNodePlan(logicalPlan); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( plan, matchScan("t") .groupId({{"a"}, {}}, {"b"}, "gid") @@ -1316,10 +1319,6 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctToMarkDistinct) { // DISTINCT aggregate combined with ORDER BY and grouping sets without a global // set. ORDER BY forces the result-producing aggregation to single-step. -// V1 is better: it distributes common DISTINCT through an inner deduplication -// Aggregate and a non-DISTINCT outer Aggregate. -// TODO: Make V2 apply the V1 DISTINCT-to-GroupBy transformation before -// physical aggregation planning. TEST_P(DistinctAggregationTest, groupingSetsDistinctWithOrderBy) { testConnector_->addTable("t", ROW({"a", "b", "c"}, BIGINT())); SCOPE_EXIT { @@ -1336,14 +1335,12 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctWithOrderBy) { // Distributed plan. { auto plan = planVelox(logicalPlan, runnerOptions_, optimizerOptions_); - AXIOM_ASSERT_DISTRIBUTED_PLAN_V1( + AXIOM_ASSERT_DISTRIBUTED_PLAN( plan.plan, matchScan("t") .groupId({{"a"}, {"b"}}, {"c"}, "gid") .distributedAggregation({"a", "b", "gid", "c"}, {}) - .shuffle() - .localPartition() - .singleAggregation( + .distributedSingleAggregation( {"a", "b", "gid"}, {"array_agg(c ORDER BY c) as a0"}) .project({"a", "b", "a0", "gid"}) .gather() @@ -1354,7 +1351,7 @@ TEST_P(DistinctAggregationTest, groupingSetsDistinctWithOrderBy) { // natively. { auto plan = toSingleNodePlan(logicalPlan); - AXIOM_ASSERT_PLAN_V1( + AXIOM_ASSERT_PLAN( plan, matchScan("t") .groupId({{"a"}, {"b"}}, {"c"}, "gid") diff --git a/axiom/optimizer/tests/sql/distinctAggregation.sql b/axiom/optimizer/tests/sql/distinctAggregation.sql index ab3521ed4..0baa920ef 100644 --- a/axiom/optimizer/tests/sql/distinctAggregation.sql +++ b/axiom/optimizer/tests/sql/distinctAggregation.sql @@ -117,6 +117,9 @@ SELECT a, count(DISTINCT 1), count(DISTINCT a), count(DISTINCT b) FROM t GROUP B -- DISTINCT: global aggregation mixing column DISTINCT and all-literal DISTINCT. SELECT count(DISTINCT c), count(DISTINCT 1) FROM t ---- +-- DISTINCT: global all-literal aggregation. +SELECT count(DISTINCT 1), sum(DISTINCT 1) FROM t +---- -- DISTINCT: the first MarkDistinct key is a subset of the second, hence no shuffle between multiple MarkDistincts. SELECT a, count(DISTINCT b), covar_pop(DISTINCT b, c) FROM t GROUP BY a ---- diff --git a/axiom/optimizer/v2/PlanPhysicalPass.cpp b/axiom/optimizer/v2/PlanPhysicalPass.cpp index 4ea726de6..ec623adf1 100644 --- a/axiom/optimizer/v2/PlanPhysicalPass.cpp +++ b/axiom/optimizer/v2/PlanPhysicalPass.cpp @@ -217,6 +217,34 @@ ExprVector unionColumnArgs(const ExprVector& keys, const ExprVector& args) { return merged; } +// Returns the shared set of non-literal DISTINCT arguments when every call is +// unfiltered DISTINCT over that same set, or nullopt otherwise. ORDER BY does +// not affect eligibility because the outer Aggregate preserves each call's +// original ordering. +std::optional commonDistinctArgs(const Aggregate& aggregate) { + if (aggregate.aggregates().empty()) { + return std::nullopt; + } + + std::optional commonArgSet; + ExprVector commonArgs; + for (const auto* call : aggregate.aggregates()) { + if (!call->isDistinct() || call->condition() != nullptr) { + return std::nullopt; + } + + ExprVector args = unionColumnArgs({}, call->args()); + PlanObjectSet argSet = PlanObjectSet::fromObjects(args); + if (!commonArgSet.has_value()) { + commonArgs = std::move(args); + commonArgSet = std::move(argSet); + } else if (argSet != *commonArgSet) { + return std::nullopt; + } + } + return commonArgs; +} + // Describes one MarkDistinct node independently of its physical input. struct MarkDistinctSpec { ExprVector keys; @@ -897,9 +925,68 @@ class PhysicalPlanRewriter : public NodeRewriter<> { return planAggregateStages(prepared, prepared->input()); } + if (auto commonArgs = commonDistinctArgs(*prepared)) { + if (commonArgs->empty() && prepared->groupingKeys().empty()) { + // A keyless inner Aggregate would create a row on empty input and + // change the result Aggregate's empty-input semantics. + return planAggregateStages(prepared, prepared->input()); + } + return planDistinctToGroupBy(prepared, *commonArgs); + } return planDistinctToMarkDistinct(prepared); } + // Plans common-signature DISTINCT calls as an inner deduplication Aggregate + // followed by the non-DISTINCT result Aggregate. + NodeCP planDistinctToGroupBy( + const Aggregate* aggregate, + const ExprVector& commonArgs) { + ExprVector innerKeys = + unionColumnArgs(aggregate->groupingKeys(), commonArgs); + ColumnVector innerColumns; + innerColumns.reserve(innerKeys.size()); + for (ExprCP key : innerKeys) { + ColumnCP column = key->as(); + VELOX_CHECK_NOT_NULL( + column, + "DISTINCT-to-GroupBy key must be a Column reference; got: {}", + key->toString()); + innerColumns.push_back(column); + } + + AggregateCP inner = builder().make(Aggregate::Key{ + .input = aggregate->input(), + .groupingKeys = std::move(innerKeys), + .aggregates = {}, + .outputColumns = std::move(innerColumns), + .step = AggregateStep::kSingle, + .groupId = nullptr, + .globalGroupingSets = {}}); + NodeCP physicalInner = planAggregateStages(inner, inner->input()); + + ExprVector outerKeys; + outerKeys.reserve(aggregate->groupingKeys().size()); + for (size_t i = 0; i < aggregate->groupingKeys().size(); ++i) { + outerKeys.push_back(physicalInner->outputColumns()[i]); + } + + AggregateCallVector outerCalls; + outerCalls.reserve(aggregate->aggregates().size()); + for (const auto* call : aggregate->aggregates()) { + outerCalls.push_back(call->dropDistinct()); + } + + AggregateCP outer = builder().make(Aggregate::Key{ + .input = physicalInner, + .groupingKeys = std::move(outerKeys), + .aggregates = std::move(outerCalls), + .outputColumns = aggregate->outputColumns(), + .step = aggregate->step(), + .groupId = aggregate->groupId(), + .globalGroupingSets = aggregate->globalGroupingSets()}); + return planAggregateStages(outer, physicalInner); + } + // Plans the MarkDistinct groups and the rewritten result Aggregate without // recursively rewriting any generated node. NodeCP planDistinctToMarkDistinct(const Aggregate* aggregate) {