From e5131f39fcd15c5283105191f6bd74f81b2ddc09 Mon Sep 17 00:00:00 2001 From: ShengHuang Date: Mon, 31 Aug 2026 17:36:30 +0800 Subject: [PATCH] Add group-key predicate pushdown optimizer --- .../group_key_predicate_push_down_optimizer.h | 49 +++++ src/optimizer/CMakeLists.txt | 1 + ...roup_key_predicate_push_down_optimizer.cpp | 204 ++++++++++++++++++ src/optimizer/optimizer.cpp | 7 + test/optimizer/optimizer_test.cpp | 87 ++++++++ .../group_key_predicate_push_down.test | 32 +++ 6 files changed, 380 insertions(+) create mode 100644 src/include/optimizer/group_key_predicate_push_down_optimizer.h create mode 100644 src/optimizer/group_key_predicate_push_down_optimizer.cpp create mode 100644 test/test_files/optimizer/group_key_predicate_push_down.test diff --git a/src/include/optimizer/group_key_predicate_push_down_optimizer.h b/src/include/optimizer/group_key_predicate_push_down_optimizer.h new file mode 100644 index 000000000..1e1c9f74d --- /dev/null +++ b/src/include/optimizer/group_key_predicate_push_down_optimizer.h @@ -0,0 +1,49 @@ +#pragma once + +#include +#include +#include + +#include "binder/expression/expression.h" +#include "planner/operator/logical_plan.h" + +namespace lbug { +namespace optimizer { + +// Pushes predicates from a WITH ... WHERE clause below an aggregate when every value referenced +// by the predicate is a grouping key of that aggregate. In particular, this turns plans shaped as +// +// FILTER(a = b) -> PROJECTION -> AGGREGATE -> ... -> CROSS_PRODUCT +// +// into plans where the filter is directly above the CROSS_PRODUCT. The regular filter push-down +// pass can then rewrite the cross product plus equality predicate into a hash join. +class GroupKeyPredicatePushDownOptimizer { +public: + void rewrite(planner::LogicalPlan* plan); + +private: + std::shared_ptr visitOperator( + std::shared_ptr op); + std::shared_ptr tryRewriteFilter( + std::shared_ptr op); + + static std::shared_ptr findAggregateThroughProjections( + const std::shared_ptr& op); + static bool isGroupKeyOnlyPredicate(const binder::Expression& predicate, + const std::unordered_set& groupKeyNames); + static bool isComposedFromGroupKeys(const binder::Expression& expression, + const std::unordered_set& groupKeyNames); + static bool containsRandomFunction(const binder::Expression& expression); + + std::shared_ptr pushPredicate( + std::shared_ptr op, + const std::shared_ptr& predicate); + static std::shared_ptr appendFilter( + std::shared_ptr child, + const std::shared_ptr& predicate, + common::cardinality_t cardinality = 0); + static void recomputeFlatSchemas(const std::shared_ptr& op); +}; + +} // namespace optimizer +} // namespace lbug diff --git a/src/optimizer/CMakeLists.txt b/src/optimizer/CMakeLists.txt index 545ffec30..b24b1614e 100644 --- a/src/optimizer/CMakeLists.txt +++ b/src/optimizer/CMakeLists.txt @@ -9,6 +9,7 @@ add_library(lbug_optimizer factorization_rewriter.cpp filter_push_down_optimizer.cpp foreign_join_push_down_optimizer.cpp + group_key_predicate_push_down_optimizer.cpp logical_operator_collector.cpp logical_operator_visitor.cpp optimizer.cpp diff --git a/src/optimizer/group_key_predicate_push_down_optimizer.cpp b/src/optimizer/group_key_predicate_push_down_optimizer.cpp new file mode 100644 index 000000000..6c8767e8a --- /dev/null +++ b/src/optimizer/group_key_predicate_push_down_optimizer.cpp @@ -0,0 +1,204 @@ +#include "optimizer/group_key_predicate_push_down_optimizer.h" + +#include "binder/expression_visitor.h" +#include "common/enums/accumulate_type.h" +#include "common/enums/join_type.h" +#include "planner/operator/logical_accumulate.h" +#include "planner/operator/logical_aggregate.h" +#include "planner/operator/logical_filter.h" +#include "planner/operator/logical_hash_join.h" + +using namespace lbug::binder; +using namespace lbug::common; +using namespace lbug::planner; + +namespace lbug { +namespace optimizer { + +void GroupKeyPredicatePushDownOptimizer::rewrite(LogicalPlan* plan) { + plan->setLastOperator(visitOperator(plan->getLastOperator())); +} + +std::shared_ptr GroupKeyPredicatePushDownOptimizer::visitOperator( + std::shared_ptr op) { + // Rewrite bottom-up so a filter sees the final shape of its aggregate input. + for (auto i = 0u; i < op->getNumChildren(); ++i) { + op->setChild(i, visitOperator(op->getChild(i))); + } + op->computeFlatSchema(); + if (op->getOperatorType() == LogicalOperatorType::FILTER) { + return tryRewriteFilter(std::move(op)); + } + return op; +} + +std::shared_ptr GroupKeyPredicatePushDownOptimizer::tryRewriteFilter( + std::shared_ptr op) { + auto& filter = op->cast(); + auto aggregateOp = findAggregateThroughProjections(filter.getChild(0)); + if (aggregateOp == nullptr) { + return op; + } + auto& aggregate = aggregateOp->cast(); + if (!aggregate.hasKeys()) { + // A keyless aggregate always emits one row, even for empty input. Moving a predicate below + // it could therefore change an empty result into a single aggregate row. + return op; + } + + std::unordered_set groupKeyNames; + for (auto& key : aggregate.getAllKeys()) { + groupKeyNames.insert(key->getUniqueName()); + } + + expression_vector predicatesToPush; + expression_vector predicatesToKeep; + for (auto& predicate : filter.getPredicate()->splitOnAND()) { + if (isGroupKeyOnlyPredicate(*predicate, groupKeyNames)) { + predicatesToPush.push_back(predicate); + } else { + predicatesToKeep.push_back(predicate); + } + } + if (predicatesToPush.empty()) { + return op; + } + + auto aggregateChild = aggregate.getChild(0); + for (auto& predicate : predicatesToPush) { + aggregateChild = pushPredicate(std::move(aggregateChild), predicate); + } + aggregate.setChild(0, std::move(aggregateChild)); + + // The moved filters preserve every operator's output schema, but recomputing keeps schema + // ownership and group positions consistent after rewiring. + recomputeFlatSchemas(filter.getChild(0)); + + auto result = filter.getChild(0); + for (auto& predicate : predicatesToKeep) { + result = appendFilter(std::move(result), predicate, filter.getCardinality()); + } + return result; +} + +std::shared_ptr +GroupKeyPredicatePushDownOptimizer::findAggregateThroughProjections( + const std::shared_ptr& op) { + auto current = op; + while (current->getOperatorType() == LogicalOperatorType::PROJECTION) { + if (current->getNumChildren() != 1) { + return nullptr; + } + current = current->getChild(0); + } + return current->getOperatorType() == LogicalOperatorType::AGGREGATE ? current : nullptr; +} + +bool GroupKeyPredicatePushDownOptimizer::isGroupKeyOnlyPredicate(const Expression& predicate, + const std::unordered_set& groupKeyNames) { + return !containsRandomFunction(predicate) && isComposedFromGroupKeys(predicate, groupKeyNames); +} + +bool GroupKeyPredicatePushDownOptimizer::isComposedFromGroupKeys(const Expression& expression, + const std::unordered_set& groupKeyNames) { + // An expression projected by an earlier query part may itself be an aggregate expression, but + // it is an ordinary grouping key in the current aggregate. Check key identity before checking + // the expression kind for that reason. + if (groupKeyNames.contains(expression.getUniqueName())) { + return true; + } + switch (expression.expressionType) { + case ExpressionType::LITERAL: + case ExpressionType::PARAMETER: + return true; + case ExpressionType::AGGREGATE_FUNCTION: + case ExpressionType::SUBQUERY: + return false; + default: + break; + } + auto children = ExpressionChildrenCollector::collectChildren(expression); + if (children.empty()) { + return false; + } + for (auto& child : children) { + if (!isComposedFromGroupKeys(*child, groupKeyNames)) { + return false; + } + } + return true; +} + +bool GroupKeyPredicatePushDownOptimizer::containsRandomFunction(const Expression& expression) { + if (ExpressionVisitor::isRandom(expression)) { + return true; + } + for (auto& child : ExpressionChildrenCollector::collectChildren(expression)) { + if (containsRandomFunction(*child)) { + return true; + } + } + return false; +} + +std::shared_ptr GroupKeyPredicatePushDownOptimizer::pushPredicate( + std::shared_ptr op, const std::shared_ptr& predicate) { + switch (op->getOperatorType()) { + case LogicalOperatorType::PROJECTION: + case LogicalOperatorType::FILTER: + case LogicalOperatorType::NODE_LABEL_FILTER: { + if (op->getNumChildren() == 1 && op->getChild(0)->getSchema()->evaluable(*predicate)) { + op->setChild(0, pushPredicate(op->getChild(0), predicate)); + op->computeFlatSchema(); + return op; + } + } break; + case LogicalOperatorType::ACCUMULATE: { + auto& accumulate = op->cast(); + if (accumulate.getAccumulateType() == AccumulateType::REGULAR && !accumulate.hasMark() && + op->getChild(0)->getSchema()->evaluable(*predicate)) { + op->setChild(0, pushPredicate(op->getChild(0), predicate)); + op->computeFlatSchema(); + return op; + } + } break; + case LogicalOperatorType::HASH_JOIN: { + auto& join = op->cast(); + auto inProbe = op->getChild(0)->getSchema()->evaluable(*predicate); + auto inBuild = op->getChild(1)->getSchema()->evaluable(*predicate); + if (inProbe && !inBuild && + (join.getJoinType() == JoinType::INNER || join.getJoinType() == JoinType::LEFT)) { + op->setChild(0, pushPredicate(op->getChild(0), predicate)); + op->computeFlatSchema(); + return op; + } + if (!inProbe && inBuild && join.getJoinType() == JoinType::INNER) { + op->setChild(1, pushPredicate(op->getChild(1), predicate)); + op->computeFlatSchema(); + return op; + } + } break; + default: + break; + } + return appendFilter(std::move(op), predicate); +} + +std::shared_ptr GroupKeyPredicatePushDownOptimizer::appendFilter( + std::shared_ptr child, const std::shared_ptr& predicate, + cardinality_t cardinality) { + auto filter = std::make_shared(predicate, std::move(child), cardinality); + filter->computeFlatSchema(); + return filter; +} + +void GroupKeyPredicatePushDownOptimizer::recomputeFlatSchemas( + const std::shared_ptr& op) { + for (auto i = 0u; i < op->getNumChildren(); ++i) { + recomputeFlatSchemas(op->getChild(i)); + } + op->computeFlatSchema(); +} + +} // namespace optimizer +} // namespace lbug diff --git a/src/optimizer/optimizer.cpp b/src/optimizer/optimizer.cpp index b00a491d6..bc7617d36 100644 --- a/src/optimizer/optimizer.cpp +++ b/src/optimizer/optimizer.cpp @@ -10,6 +10,7 @@ #include "optimizer/factorization_rewriter.h" #include "optimizer/filter_push_down_optimizer.h" #include "optimizer/foreign_join_push_down_optimizer.h" +#include "optimizer/group_key_predicate_push_down_optimizer.h" #include "optimizer/limit_push_down_optimizer.h" #include "optimizer/order_by_push_down_optimizer.h" #include "optimizer/projection_push_down_optimizer.h" @@ -59,6 +60,12 @@ void Optimizer::optimize(planner::LogicalPlan* plan, main::ClientContext* contex auto boolFoldingOptimizer = BoolFoldingOptimizer(); boolFoldingOptimizer.rewrite(plan); + // A WITH ... WHERE predicate that depends only on aggregate grouping keys can be + // evaluated before aggregation. Moving it first allows regular filter push-down to turn + // otherwise quadratic cross products into joins. + auto groupKeyPredicatePushDownOptimizer = GroupKeyPredicatePushDownOptimizer(); + groupKeyPredicatePushDownOptimizer.rewrite(plan); + auto filterPushDownOptimizer = FilterPushDownOptimizer(context, &cardinalityEstimator); filterPushDownOptimizer.rewrite(plan); diff --git a/test/optimizer/optimizer_test.cpp b/test/optimizer/optimizer_test.cpp index c0a019349..f1fbeed43 100644 --- a/test/optimizer/optimizer_test.cpp +++ b/test/optimizer/optimizer_test.cpp @@ -182,6 +182,93 @@ TEST_F(OptimizerTest, FilterPushDownTest) { ASSERT_STREQ(getEncodedPlan(q1).c_str(), "E(b)Filter()Filter()S(a)"); } +TEST_F(OptimizerTest, GroupKeyPredicatePushDownEliminatesCrossProduct) { + auto query = "MATCH (a:person) " + "WITH a.ID AS entity_id, COUNT(a.ID) AS metric_0 " + "MATCH (b:person) " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.ID) AS metric_1 " + "WHERE metric_1_entity_id = entity_id " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + auto explicitJoinQuery = + "MATCH (a:person) " + "WITH a.ID AS entity_id, COUNT(a.ID) AS metric_0 " + "MATCH (b:person) WHERE b.ID = entity_id " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.ID) AS metric_1 " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + + auto plan = getRoot(query); + ASSERT_FALSE(hasOperatorType(plan->getLastOperator().get(), + planner::LogicalOperatorType::CROSS_PRODUCT)); + ASSERT_TRUE( + hasOperatorType(plan->getLastOperator().get(), planner::LogicalOperatorType::HASH_JOIN)); + + auto result = conn->query(query); + auto expected = conn->query(explicitJoinQuery); + ASSERT_TRUE(result->isSuccess()) << result->getErrorMessage(); + ASSERT_TRUE(expected->isSuccess()) << expected->getErrorMessage(); + ASSERT_EQ(TestHelper::convertResultToString(*result), + TestHelper::convertResultToString(*expected)); +} + +TEST_F(OptimizerTest, GroupKeyPredicatePushDownThroughOptionalMatch) { + auto query = "MATCH (a:person) " + "OPTIONAL MATCH (a)-[:knows]->(aFriend:person) " + "WITH a.ID AS entity_id, COUNT(aFriend.ID) AS metric_0 " + "MATCH (b:person) " + "OPTIONAL MATCH (b)-[:knows]->(bFriend:person) " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.age) AS metric_1 " + "WHERE metric_1_entity_id = entity_id " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + auto explicitJoinQuery = + "MATCH (a:person) " + "OPTIONAL MATCH (a)-[:knows]->(aFriend:person) " + "WITH a.ID AS entity_id, COUNT(aFriend.ID) AS metric_0 " + "MATCH (b:person) WHERE b.ID = entity_id " + "OPTIONAL MATCH (b)-[:knows]->(bFriend:person) " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.age) AS metric_1 " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + + auto plan = getRoot(query); + ASSERT_FALSE(hasOperatorType(plan->getLastOperator().get(), + planner::LogicalOperatorType::CROSS_PRODUCT)); + + auto result = conn->query(query); + auto expected = conn->query(explicitJoinQuery); + ASSERT_TRUE(result->isSuccess()) << result->getErrorMessage(); + ASSERT_TRUE(expected->isSuccess()) << expected->getErrorMessage(); + ASSERT_EQ(TestHelper::convertResultToString(*result), + TestHelper::convertResultToString(*expected)); +} + +TEST_F(OptimizerTest, GroupKeyPredicatePushDownKeepsAggregatePredicate) { + auto query = "MATCH (a:person) " + "WITH a.ID AS entity_id, COUNT(a.ID) AS metric_0 " + "MATCH (b:person) " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.ID) AS metric_1 " + "WHERE metric_1_entity_id = entity_id AND metric_1 > 3 " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + auto explicitJoinQuery = + "MATCH (a:person) " + "WITH a.ID AS entity_id, COUNT(a.ID) AS metric_0 " + "MATCH (b:person) WHERE b.ID = entity_id " + "WITH entity_id, metric_0, b.ID AS metric_1_entity_id, AVG(b.ID) AS metric_1 " + "WHERE metric_1 > 3 " + "RETURN entity_id, metric_0, metric_1 ORDER BY entity_id"; + + auto plan = getRoot(query); + ASSERT_FALSE(hasOperatorType(plan->getLastOperator().get(), + planner::LogicalOperatorType::CROSS_PRODUCT)); + ASSERT_TRUE( + hasOperatorType(plan->getLastOperator().get(), planner::LogicalOperatorType::FILTER)); + + auto result = conn->query(query); + auto expected = conn->query(explicitJoinQuery); + ASSERT_TRUE(result->isSuccess()) << result->getErrorMessage(); + ASSERT_TRUE(expected->isSuccess()) << expected->getErrorMessage(); + ASSERT_EQ(TestHelper::convertResultToString(*result), + TestHelper::convertResultToString(*expected)); +} + TEST_F(OptimizerTest, IndexScanTest) { auto q1 = "MATCH (a:person) " "WHERE a.ID = 0 AND a.fName='Alice' " diff --git a/test/test_files/optimizer/group_key_predicate_push_down.test b/test/test_files/optimizer/group_key_predicate_push_down.test new file mode 100644 index 000000000..6a6b50255 --- /dev/null +++ b/test/test_files/optimizer/group_key_predicate_push_down.test @@ -0,0 +1,32 @@ +-DATASET CSV empty +-BUFFER_POOL_SIZE 67108864 + +-- + +# Without group-key predicate push-down, the second aggregate builds 25 million groups and +# exhausts this 64 MiB buffer pool before the post-aggregate equality filter can run. +-CASE GroupKeyPredicatePushDownAvoidsAggregateOOM +-SKIP_PAGE_SIZE_TESTS +-STATEMENT CREATE NODE TABLE Person(id INT64, PRIMARY KEY(id)); +---- ok +-STATEMENT UNWIND range(1, 5000) AS i CREATE (:Person {id: i}); +---- ok +-STATEMENT MATCH (p0:Person) + WITH p0.id AS entity_id, COUNT(p0.id) AS metric_0_value + MATCH (p1:Person) + WITH entity_id, metric_0_value, p1.id AS metric_1_entity_id, + AVG(p1.id) AS metric_1_value + WHERE metric_1_entity_id = entity_id + WITH entity_id, metric_0_value, metric_1_value + MATCH (p2:Person) + WITH entity_id, metric_0_value, metric_1_value, + p2.id AS metric_2_entity_id, SUM(p2.id) AS metric_2_value + WHERE metric_2_entity_id = entity_id + WITH entity_id, metric_0_value, metric_1_value, metric_2_value + MATCH (p3:Person) + WITH entity_id, metric_0_value, metric_1_value, metric_2_value, + p3.id AS metric_3_entity_id, COUNT(p3.id) AS metric_3_value + WHERE metric_3_entity_id = entity_id + RETURN COUNT(*); +---- 1 +5000