Skip to content

fix(optimizer): Drop a redundant COALESCE in decorrelation #1782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mbasmanova wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(optimizer): Drop a redundant COALESCE in decorrelation #1782

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5eccf07
test: Pin v2 plan shapes in the enforceSingleRow tests
mbasmanova Aug 26, 2026
7e5a91d
fix(optimizer): Drop a redundant COALESCE in decorrelation
mbasmanova Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions axiom/optimizer/FunctionRegistry.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -526,12 +526,9 @@ class FunctionRegistry {
const std::vector<std::string>& names,
AggregateEmptyResultResolver resolver);

/// Returns the result of an aggregate function over empty input.
/// If registerCount() was called, returns 0 (as BIGINT) for 'count' function.
/// Otherwise, uses the resolver registered via
/// Returns the result of an aggregate function over empty input. Uses
/// 'count' registered via registerCount() or a resolver registered via
/// registerAggregateEmptyResultResolver().
/// Returns the result of an aggregate function over empty input. Returns 0
/// as BIGINT for the 'count' registered via registerCount(); otherwise uses
/// a resolver registered via registerAggregateEmptyResultResolver().
/// @param name The aggregate function name.
/// @param argTypes The argument types of the aggregate function.
/// @return Non-null Variant with the result for empty input, or null Variant
Expand Down
2 changes: 1 addition & 1 deletion axiom/optimizer/tests/SqlTestBase.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ std::shared_ptr<runner::LocalRunner> SqlTestBase::makeRunner(
/*user=*/"test",
::axiom::sql::presto::ParserOptions{},
connector::ConnectorProperties{}));
auto statement = parser.parse(sql, true);
auto statement = parser.parse(sql);

VELOX_CHECK(
statement->isSelect(), "Only SELECT statements are supported: {}", sql);
Expand Down
128 changes: 65 additions & 63 deletions axiom/optimizer/tests/SubqueryTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1123,31 +1123,34 @@ TEST_P(SubqueryTest, enforceSingleRow) {
auto logicalPlan = parseSelect(query);

{
auto matcher =
matchHiveScan("region")
.nestedLoopJoin(matchHiveScan("nation").enforceSingleRow())
.filter()
.project()
.build();
auto matcher = matchHiveScan("region")
.nestedLoopJoin(
matchHiveScan("nation").enforceSingleRow(),
core::JoinType::kInner,
"r_regionkey > n_regionkey")
.build();

auto plan = toSingleNodePlan(logicalPlan);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);

VELOX_ASSERT_THROW(runVelox(plan), "Expected single row of input.");
}

{
auto matcher =
matchHiveScan("region")
.nestedLoopJoin(
matchHiveScan("nation").gather().enforceSingleRow().broadcast())
.filter()
.project()
.gather()
.build();
auto matcher = matchHiveScan("region")
.nestedLoopJoin(
matchHiveScan("nation")
.gather()
.localPartition()
.enforceSingleRow()
.broadcast(),
core::JoinType::kInner,
"r_regionkey > n_regionkey")
.gather()
.build();

auto distributedPlan = planVelox(logicalPlan);
AXIOM_ASSERT_DISTRIBUTED_PLAN_V1(distributedPlan.plan, matcher);
AXIOM_ASSERT_DISTRIBUTED_PLAN_V2(distributedPlan.plan, matcher);
}
}

Expand All @@ -1158,32 +1161,33 @@ TEST_P(SubqueryTest, enforceSingleRowInProjection) {
auto logicalPlan = parseSelect(query);

{
auto matcher = core::PlanMatcherBuilder()
.hiveScan("region", test::eq("r_name", "AFRICA"))
.enforceSingleRow()
.nestedLoopJoin(matchHiveScan("nation"))
auto matcher = matchHiveScan("nation")
.nestedLoopJoin(
matchHiveScan("region", test::eq("r_name", "AFRICA"))
.enforceSingleRow())
.build();

auto plan = toSingleNodePlan(logicalPlan);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

{
// TODO: Plan has an extra fragment and an extra shuffle. A
// broadcast-then-EnforceSingleRow shape would collapse the gather
// + broadcast pair into a single broadcast with EnforceSingleRow
// running on each consumer.
auto matcher =
core::PlanMatcherBuilder()
.hiveScan("region", test::eq("r_name", "AFRICA"))
.gather()
.enforceSingleRow()
.nestedLoopJoin(
core::PlanMatcherBuilder().tableScan("nation").broadcast())
.build();
auto matcher = matchHiveScan("nation")
.nestedLoopJoin(
matchHiveScan("region", test::eq("r_name", "AFRICA"))
.gather()
.localPartition()
.enforceSingleRow()
.broadcast())
.gather()
.build();

auto distributedPlan = planVelox(logicalPlan);
AXIOM_ASSERT_DISTRIBUTED_PLAN_V1(distributedPlan.plan, matcher);
AXIOM_ASSERT_DISTRIBUTED_PLAN_V2(distributedPlan.plan, matcher);
}
}

Expand Down Expand Up @@ -1235,8 +1239,9 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalar) {
.assignUniqueId("unique_id")
.nestedLoopJoin(
matchHiveScan("nation").project(
{"true as marker", "n_regionkey"}),
velox::core::JoinType::kLeft)
{"n_regionkey", "true as marker"}),
velox::core::JoinType::kLeft,
"r_regionkey > n_regionkey")
.streamingAggregation(
{"unique_id"},
{
Expand All @@ -1250,7 +1255,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalar) {
.build();

auto plan = toSingleNodePlan(logicalPlan);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

{
Expand Down Expand Up @@ -1338,7 +1343,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalarWithNestedAggregation) {
.hashJoin(
matchScan("u")
.singleAggregation({"c"}, {"count(*) as inner_cnt"})
.project({"true as marker", "inner_cnt", "c"}),
.project({"c", "inner_cnt", "true as marker"}),
velox::core::JoinType::kLeft)
.streamingAggregation(
{"unique_id"},
Expand All @@ -1352,7 +1357,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalarWithNestedAggregation) {
.build();

auto plan = toSingleNodePlan(parseSelect(query, kTestConnectorId));
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

TEST_P(SubqueryTest, nonEquiCorrelatedProject) {
Expand All @@ -1369,21 +1374,20 @@ TEST_P(SubqueryTest, nonEquiCorrelatedProject) {
.assignUniqueId("unique_id")
.nestedLoopJoin(
matchHiveScan("nation").project(
{"true as marker", "n_regionkey"}),
velox::core::JoinType::kLeft)
{"n_regionkey", "true as marker"}),
velox::core::JoinType::kLeft,
"r_regionkey > n_regionkey")
.streamingAggregation(
{"unique_id"},
{
"count(*) filter (where marker) as cnt",
"arbitrary(r_regionkey)",
"arbitrary(r_name) as r_name",
})
.project({"length(r_name)", "cnt"})
.project()
.build();

auto plan = toSingleNodePlan(logicalPlan);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

{
Expand Down Expand Up @@ -1561,7 +1565,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalarThenCorrelatedExists) {
.assignUniqueId("unique_id")
.nestedLoopJoin(
matchHiveScan("nation").project(
{"true as marker", "n_regionkey"}),
{"n_regionkey", "true as marker"}),
velox::core::JoinType::kLeft)
.streamingAggregation(
{"unique_id"},
Expand All @@ -1577,7 +1581,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalarThenCorrelatedExists) {
.build();

auto plan = toSingleNodePlan(query);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

// Non-equi correlated scalar (count(*)) followed by an uncorrelated
Expand All @@ -1590,27 +1594,25 @@ TEST_P(SubqueryTest, nonEquiCorrelatedThenUncorrelatedScalar) {
"FROM region";
SCOPED_TRACE(query);

auto matcher = matchHiveScan("region")
.assignUniqueId("unique_id")
.nestedLoopJoin(
matchHiveScan("nation").project(
{"true as marker", "n_regionkey"}),
velox::core::JoinType::kLeft)
.streamingAggregation(
{"unique_id"},
{
"count(*) filter (where marker) as cnt",
"arbitrary(r_regionkey) as r_regionkey",
})
.project()
.nestedLoopJoin(matchHiveScan("supplier")
.singleAggregation(
{}, {"max(s_suppkey) as max_key"}))
.project({"cnt as x", "max_key as y"})
.build();
auto matcher =
matchHiveScan("region")
.assignUniqueId("unique_id")
.nestedLoopJoin(
matchHiveScan("nation").project(
{"n_regionkey", "true as marker"}),
velox::core::JoinType::kLeft,
"r_regionkey < n_regionkey")
.streamingAggregation(
{"unique_id"}, {"count(*) filter (where marker) as cnt"})
.project()
.nestedLoopJoin(
matchHiveScan("supplier")
.singleAggregation({}, {"max(s_suppkey) as max_key"}))
.project({"cnt as x", "max_key as y"})
.build();

auto plan = toSingleNodePlan(query);
AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

// Correlated EXISTS combined with an uncorrelated IN in the same SELECT.
Expand Down Expand Up @@ -2041,12 +2043,12 @@ TEST_P(SubqueryTest, nonEquiLeftJoinWithScalarSubquery) {
// nested-loop join.
auto matcher =
matchScan("t")
.nestedLoopJoin(matchScan("u"), velox::core::JoinType::kLeft)
.nestedLoopJoin(matchScan("v").enforceSingleRow())
.nestedLoopJoin(matchScan("u"), velox::core::JoinType::kLeft, "b < c")
.project()
.build();

AXIOM_ASSERT_PLAN_V1(plan, matcher);
AXIOM_ASSERT_PLAN_V2(plan, matcher);
}

// LEFT JOIN with a post-join WHERE equality referencing both sides, where one
Expand Down
5 changes: 5 additions & 0 deletions axiom/optimizer/tests/sql/subquery.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,11 @@ FROM (SELECT 20 AS x, 30 AS y) v
-- constant, for an outer row the subquery has no row for.
SELECT a, (SELECT 1 FROM v WHERE v.a = t.a) AS one FROM t
----
-- A correlated count(*) reads 0, not NULL, for an outer row the subquery
-- has no row for, so a HAVING on that count still sees 0.
-- error_v1: (0 vs. 1)
SELECT a, (SELECT count(*) FROM u WHERE u.a > t.a HAVING count(*) = 0) AS c FROM t
----
-- Multiple correlated scalar count(*) subqueries with non-equi predicates
-- in the same SELECT list, each correlating on a different outer column.
SELECT
Expand Down
25 changes: 18 additions & 7 deletions axiom/optimizer/v2/DecorrelatePass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1683,7 +1683,8 @@ class Decorrelator : public NodeRewriter<> {
input, aggregate, std::move(filterPreConjuncts));
NodeCP decorrelatedInner = rewrite(innerApply.apply);

auto wraps = buildAggregateWraps(aggregate, numGroupingKeys);
auto wraps = buildAggregateWraps(
aggregate, numGroupingKeys, /*everyOuterRowHasGroup=*/true);
NodeCP liftedAggregate = buildLiftedAggregate(
input,
aggregate,
Expand Down Expand Up @@ -1798,8 +1799,8 @@ class Decorrelator : public NodeRewriter<> {
AggregateCP aggregate,
EquiCorrelation correlation,
const ExprVector& filterPostConjuncts) {
std::vector<AggregateWrap> wraps =
buildAggregateWraps(aggregate, /*numGroupingKeys=*/0);
std::vector<AggregateWrap> wraps = buildAggregateWraps(
aggregate, /*numGroupingKeys=*/0, /*everyOuterRowHasGroup=*/false);

// The correlation keys become the new Aggregate's grouping keys and the
// join-back's right keys: reuse the body column for a plain column, mint
Expand Down Expand Up @@ -1874,6 +1875,13 @@ class Decorrelator : public NodeRewriter<> {
// FunctionRegistry: aggregates whose empty value is non-NULL (count,
// count_if, etc.) need COALESCE; others pass through.
//
// 'everyOuterRowHasGroup' is true when the lifted Aggregate groups by the
// outer row id above a kLeft join. An outer row with no matches still forms
// a group there, and a masked aggregate over that empty group already
// returns its empty-input value, so no COALESCE is needed. Only the
// join-back shape, where such an outer row has no group at all and the join
// pads it with NULL, needs one.
//
// Slot-identity invariant: a Column* must carry the same value
// across all output positions. For COALESCE-needing aggregates,
// the lifted Aggregate's raw output
Expand All @@ -1893,7 +1901,8 @@ class Decorrelator : public NodeRewriter<> {

std::vector<AggregateWrap> buildAggregateWraps(
AggregateCP aggregate,
size_t numGroupingKeys) {
size_t numGroupingKeys,
bool everyOuterRowHasGroup) {
std::vector<AggregateWrap> wraps;
wraps.reserve(aggregate->aggregates().size());
const auto* registry = FunctionRegistry::instance();
Expand All @@ -1910,7 +1919,7 @@ class Decorrelator : public NodeRewriter<> {
velox::Variant emptyValue = registry->aggregateResultForEmptyInput(
aggregateCall->name(), argumentTypes);

if (emptyValue.isNull()) {
if (everyOuterRowHasGroup || emptyValue.isNull()) {
wraps.push_back({originalOutput, originalOutput});
} else {
ColumnCP rawOutput =
Expand Down Expand Up @@ -2480,7 +2489,8 @@ class Decorrelator : public NodeRewriter<> {
AggregateRecovery::validateAggregateArgs(
aggregate->aggregates(), node->correlationColumns());

auto wraps = buildAggregateWraps(aggregate, numGroupingKeys);
auto wraps = buildAggregateWraps(
aggregate, numGroupingKeys, /*everyOuterRowHasGroup=*/true);

AggregateCallVector stage1Aggregates = recovery.rewriteCountStar(
aggregate->aggregates(), innerApply.includeMarker);
Expand Down Expand Up @@ -2668,7 +2678,8 @@ class Decorrelator : public NodeRewriter<> {
AggregateRecovery::validateAggregateArgs(
aggregate->aggregates(), node->correlationColumns());

auto wraps = buildAggregateWraps(aggregate, numGroupingKeys);
auto wraps = buildAggregateWraps(
aggregate, numGroupingKeys, /*everyOuterRowHasGroup=*/true);

AggregateRecovery recovery(builder(), exprFactory_);
auto innerApply = buildAggregateInnerApply(
Expand Down
Loading