diff --git a/axiom/optimizer/FunctionRegistry.h b/axiom/optimizer/FunctionRegistry.h index 8a078ca44..7497b8654 100644 --- a/axiom/optimizer/FunctionRegistry.h +++ b/axiom/optimizer/FunctionRegistry.h @@ -526,12 +526,9 @@ class FunctionRegistry { const std::vector& 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 diff --git a/axiom/optimizer/tests/SqlTestBase.cpp b/axiom/optimizer/tests/SqlTestBase.cpp index 57cbef370..2c8d97b3d 100644 --- a/axiom/optimizer/tests/SqlTestBase.cpp +++ b/axiom/optimizer/tests/SqlTestBase.cpp @@ -241,7 +241,7 @@ std::shared_ptr 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); diff --git a/axiom/optimizer/tests/SubqueryTest.cpp b/axiom/optimizer/tests/SubqueryTest.cpp index 929c8758e..03c8b0633 100644 --- a/axiom/optimizer/tests/SubqueryTest.cpp +++ b/axiom/optimizer/tests/SubqueryTest.cpp @@ -768,20 +768,17 @@ TEST_P(SubqueryTest, correlatedScalar) { // The correlated scalar subquery is transformed into a LEFT JOIN with // aggregation grouped by the correlation key, then filtered. - auto matcher = - matchHiveScan("region") - .hashJoin( - matchHiveScan("nation") - .singleAggregation({"n_regionkey"}, {"min(n_nationkey)"}) - .projectIf(!useV2_), - velox::core::JoinType::kLeft) - .filter("r_regionkey = min") - .project() - .build(); + auto matcher = matchHiveScan("region") + .hashJoinInner(matchHiveScan("nation") + .singleAggregation( + {"n_regionkey"}, + {"min(n_nationkey) as min_key"}) + .filter("n_regionkey = min_key")) + .build(); SCOPED_TRACE(query); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } { @@ -793,19 +790,20 @@ TEST_P(SubqueryTest, correlatedScalar) { // aggregation grouped by the correlation key. The count result is wrapped // with COALESCE to return 0 for unmatched rows (instead of NULL from // LEFT JOIN). - auto matcher = matchHiveScan("region") - .hashJoin( - matchHiveScan("nation") - .singleAggregation({"n_regionkey"}, {"count(*)"}) - .projectIf(!useV2_), - velox::core::JoinType::kLeft) - .filter("r_regionkey = coalesce(count, 0)") - .project() - .build(); + auto matcher = + matchHiveScan("region") + .hashJoin( + matchHiveScan("nation") + .singleAggregation({"n_regionkey"}, {"count(*) as cnt"}) + .projectIf(!useV2_), + velox::core::JoinType::kLeft) + .filter("r_regionkey = coalesce(cnt, 0)") + .project() + .build(); SCOPED_TRACE(query); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN(plan, matcher); } { @@ -813,20 +811,22 @@ TEST_P(SubqueryTest, correlatedScalar) { "SELECT * FROM region " "WHERE r_regionkey = (SELECT approx_distinct(n_name) FROM nation WHERE n_regionkey = r_regionkey)"; - auto matcher = matchHiveScan("region") - .hashJoin( - matchHiveScan("nation") - .singleAggregation( - {"n_regionkey"}, {"approx_distinct(n_name)"}) - .projectIf(!useV2_), - velox::core::JoinType::kLeft) - .filter("r_regionkey = coalesce(approx_distinct, 0)") - .project() - .build(); + auto matcher = + matchHiveScan("region") + .hashJoin( + matchHiveScan("nation") + .singleAggregation( + {"n_regionkey"}, + {"approx_distinct(n_name) as distinct_names"}) + .projectIf(!useV2_), + velox::core::JoinType::kLeft) + .filter("r_regionkey = coalesce(distinct_names, 0)") + .project() + .build(); SCOPED_TRACE(query); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN(plan, matcher); } } @@ -1127,7 +1127,7 @@ TEST_P(SubqueryTest, enforceSingleRow) { .nestedLoopJoin( matchHiveScan("nation").enforceSingleRow(), core::JoinType::kInner, - "gt(r_regionkey, n_regionkey)") + "r_regionkey > n_regionkey") .build(); auto plan = toSingleNodePlan(logicalPlan); @@ -1145,7 +1145,7 @@ TEST_P(SubqueryTest, enforceSingleRow) { .enforceSingleRow() .broadcast(), core::JoinType::kInner, - "gt(r_regionkey, n_regionkey)") + "r_regionkey > n_regionkey") .gather() .build(); @@ -1239,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"}, { @@ -1254,7 +1255,7 @@ TEST_P(SubqueryTest, nonEquiCorrelatedScalar) { .build(); auto plan = toSingleNodePlan(logicalPlan); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } { @@ -1342,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"}, @@ -1356,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) { @@ -1373,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); } { @@ -1565,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"}, @@ -1581,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 @@ -1594,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. @@ -1684,35 +1682,35 @@ TEST_P(SubqueryTest, correlatedScalarWithoutAggregation) { auto query = "SELECT * FROM t WHERE a > (SELECT c + d FROM u WHERE d < b)"; SCOPED_TRACE(query); - auto matcher = matchScan("t") - .assignUniqueId("unique_id") - .nestedLoopJoin( - matchScan("u").project({"c + d as cd", "d"}), - velox::core::JoinType::kLeft) - .enforceDistinct({"unique_id"}) - .filter("a > cd") - .project() - .build(); + auto matcher = + matchScan("t") + .assignUniqueId("unique_id") + .nestedLoopJoin( + matchScan("u"), velox::core::JoinType::kLeft, "b > d") + .enforceDistinct({"unique_id"}) + .filter("a > c + d") + .project() + .build(); auto plan = toSingleNodePlan(parseSelect(query, kTestConnectorId)); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } { auto query = "SELECT a + (SELECT c + d FROM u WHERE d < b) FROM t"; SCOPED_TRACE(query); - auto matcher = matchScan("t") - .assignUniqueId("unique_id") - .nestedLoopJoin( - matchScan("u").project({"c + d as cd", "d"}), - velox::core::JoinType::kLeft) - .enforceDistinct({"unique_id"}) - .project({"a + cd"}) - .build(); + auto matcher = + matchScan("t") + .assignUniqueId("unique_id") + .nestedLoopJoin( + matchScan("u"), velox::core::JoinType::kLeft, "b > d") + .enforceDistinct({"unique_id"}) + .project({"a + (c + d)"}) + .build(); auto plan = toSingleNodePlan(parseSelect(query, kTestConnectorId)); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } } @@ -1728,17 +1726,15 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { " AND n.n_nationkey = (SELECT min(s_nationkey) FROM supplier)"; SCOPED_TRACE(query); - auto matcher = - matchHiveScan("nation") - .hashJoin( - matchHiveScan("supplier") - .singleAggregation({}, {"min(s_nationkey)"}), - velox::core::JoinType::kInner) - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) - .build(); + auto matcher = matchHiveScan("region") + .hashJoin(matchHiveScan("nation").hashJoinInner( + matchHiveScan("supplier") + .singleAggregation({}, {"min(s_nationkey)"}))) + .project() + .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } // IN subquery in ON clause. @@ -1747,16 +1743,14 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { baseJoin + " AND n.n_nationkey IN (SELECT s_nationkey FROM supplier)"; SCOPED_TRACE(query); - auto matcher = - matchHiveScan("supplier") - .hashJoin( - matchHiveScan("nation"), - velox::core::JoinType::kRightSemiFilter) - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) - .build(); + auto matcher = matchHiveScan("supplier") + .hashJoinRightSemiFilter(matchHiveScan("nation")) + .hashJoinInner(matchHiveScan("region")) + .project() + .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } // NOT IN subquery in ON clause. @@ -1766,17 +1760,17 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { SCOPED_TRACE(query); auto matcher = - matchHiveScan("supplier") - .hashJoin( - matchHiveScan("nation"), - velox::core::JoinType::kRightSemiProject, - {.nullAware = true}) - .filter() - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) + matchHiveScan("region") + .hashJoinInner(matchHiveScan("supplier") + .hashJoinRightSemiProject( + matchHiveScan("nation"), {.nullAware = true}) + .filter() + .project()) + .project() .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } // EXISTS subquery in ON clause. @@ -1786,16 +1780,14 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { " WHERE s.s_nationkey = n.n_nationkey)"; SCOPED_TRACE(query); - auto matcher = - matchHiveScan("supplier") - .hashJoin( - matchHiveScan("nation"), - velox::core::JoinType::kRightSemiFilter) - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) - .build(); + auto matcher = matchHiveScan("supplier") + .hashJoinRightSemiFilter(matchHiveScan("nation")) + .hashJoinInner(matchHiveScan("region")) + .project() + .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } // NOT EXISTS subquery in ON clause. @@ -1805,18 +1797,18 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { " WHERE s.s_nationkey = n.n_nationkey)"; SCOPED_TRACE(query); - auto matcher = - matchHiveScan("supplier") - .hashJoin( - matchHiveScan("nation"), - velox::core::JoinType::kRightSemiProject, - {.nullAware = false}) - .filter() - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) - .build(); + auto matcher = matchHiveScan("region") + .hashJoinInner(matchHiveScan("supplier") + .hashJoinRightSemiProject( + matchHiveScan("nation"), + {.nullAware = false}) + .filter() + .project()) + .project() + .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } // Correlated scalar @@ -1848,19 +1840,17 @@ TEST_P(SubqueryTest, innerJoinOnSubquery) { "(SELECT min(s_nationkey) FROM supplier)"; SCOPED_TRACE(query); - auto matcher = - matchHiveScan("nation") - .hashJoin(matchHiveScan("region"), velox::core::JoinType::kInner) - .nestedLoopJoin( - matchHiveScan("supplier") - .singleAggregation({}, {"min(s_nationkey)"}), - velox::core::JoinType::kInner) - .filter() - .project() - .build(); + auto matcher = matchHiveScan("nation") + .hashJoinInner(matchHiveScan("region")) + .project() + .nestedLoopJoin( + matchHiveScan("supplier") + .singleAggregation({}, {"min(s_nationkey)"}), + velox::core::JoinType::kInner) + .build(); auto plan = toSingleNodePlan(query); - AXIOM_ASSERT_PLAN_V1(plan, matcher); + AXIOM_ASSERT_PLAN_V2(plan, matcher); } } @@ -2045,12 +2035,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 diff --git a/axiom/optimizer/tests/sql/subquery.sql b/axiom/optimizer/tests/sql/subquery.sql index e7af80af7..dde56a11d 100644 --- a/axiom/optimizer/tests/sql/subquery.sql +++ b/axiom/optimizer/tests/sql/subquery.sql @@ -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 diff --git a/axiom/optimizer/v2/DecorrelatePass.cpp b/axiom/optimizer/v2/DecorrelatePass.cpp index 11bb19f15..fec66ec67 100644 --- a/axiom/optimizer/v2/DecorrelatePass.cpp +++ b/axiom/optimizer/v2/DecorrelatePass.cpp @@ -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, @@ -1798,8 +1799,8 @@ class Decorrelator : public NodeRewriter<> { AggregateCP aggregate, EquiCorrelation correlation, const ExprVector& filterPostConjuncts) { - std::vector wraps = - buildAggregateWraps(aggregate, /*numGroupingKeys=*/0); + std::vector 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 @@ -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 @@ -1893,7 +1901,8 @@ class Decorrelator : public NodeRewriter<> { std::vector buildAggregateWraps( AggregateCP aggregate, - size_t numGroupingKeys) { + size_t numGroupingKeys, + bool everyOuterRowHasGroup) { std::vector wraps; wraps.reserve(aggregate->aggregates().size()); const auto* registry = FunctionRegistry::instance(); @@ -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 = @@ -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); @@ -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(