Add group-key predicate pushdown optimizer

src/include/optimizer/group_key_predicate_push_down_optimizer.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_set>
+
+#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<planner::LogicalOperator> visitOperator(
+        std::shared_ptr<planner::LogicalOperator> op);
+    std::shared_ptr<planner::LogicalOperator> tryRewriteFilter(
+        std::shared_ptr<planner::LogicalOperator> op);
+
+    static std::shared_ptr<planner::LogicalOperator> findAggregateThroughProjections(
+        const std::shared_ptr<planner::LogicalOperator>& op);
+    static bool isGroupKeyOnlyPredicate(const binder::Expression& predicate,
+        const std::unordered_set<std::string>& groupKeyNames);
+    static bool isComposedFromGroupKeys(const binder::Expression& expression,
+        const std::unordered_set<std::string>& groupKeyNames);
+    static bool containsRandomFunction(const binder::Expression& expression);
+
+    std::shared_ptr<planner::LogicalOperator> pushPredicate(
+        std::shared_ptr<planner::LogicalOperator> op,
+        const std::shared_ptr<binder::Expression>& predicate);
+    static std::shared_ptr<planner::LogicalOperator> appendFilter(
+        std::shared_ptr<planner::LogicalOperator> child,
+        const std::shared_ptr<binder::Expression>& predicate,
+        common::cardinality_t cardinality = 0);
+    static void recomputeFlatSchemas(const std::shared_ptr<planner::LogicalOperator>& op);
+};
+
+} // namespace optimizer
+} // namespace lbug

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

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<LogicalOperator> GroupKeyPredicatePushDownOptimizer::visitOperator(
+    std::shared_ptr<LogicalOperator> 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<LogicalOperator> GroupKeyPredicatePushDownOptimizer::tryRewriteFilter(
+    std::shared_ptr<LogicalOperator> op) {
+    auto& filter = op->cast<LogicalFilter>();
+    auto aggregateOp = findAggregateThroughProjections(filter.getChild(0));
+    if (aggregateOp == nullptr) {
+        return op;
+    }
+    auto& aggregate = aggregateOp->cast<LogicalAggregate>();
+    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<std::string> 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<LogicalOperator>
+GroupKeyPredicatePushDownOptimizer::findAggregateThroughProjections(
+    const std::shared_ptr<LogicalOperator>& 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<std::string>& groupKeyNames) {
+    return !containsRandomFunction(predicate) && isComposedFromGroupKeys(predicate, groupKeyNames);
+}
+
+bool GroupKeyPredicatePushDownOptimizer::isComposedFromGroupKeys(const Expression& expression,
+    const std::unordered_set<std::string>& 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<LogicalOperator> GroupKeyPredicatePushDownOptimizer::pushPredicate(
+    std::shared_ptr<LogicalOperator> op, const std::shared_ptr<Expression>& 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<LogicalAccumulate>();
+        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<LogicalHashJoin>();
+        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<LogicalOperator> GroupKeyPredicatePushDownOptimizer::appendFilter(
+    std::shared_ptr<LogicalOperator> child, const std::shared_ptr<Expression>& predicate,
+    cardinality_t cardinality) {
+    auto filter = std::make_shared<LogicalFilter>(predicate, std::move(child), cardinality);
+    filter->computeFlatSchema();
+    return filter;
+}
+
+void GroupKeyPredicatePushDownOptimizer::recomputeFlatSchemas(
+    const std::shared_ptr<LogicalOperator>& op) {
+    for (auto i = 0u; i < op->getNumChildren(); ++i) {
+        recomputeFlatSchemas(op->getChild(i));
+    }
+    op->computeFlatSchema();
+}
+
+} // namespace optimizer
+} // namespace lbug

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);
 

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' "