Improve subquery planning

kuzudb · Dec 20, 2024 · 3530c01 · 3530c01
1 parent 86f1f43
commit 3530c01
Show file tree

Hide file tree

Showing 7 changed files with 196 additions and 80 deletions.
diff --git a/src/include/common/data_chunk/sel_vector.h b/src/include/common/data_chunk/sel_vector.h
@@ -24,7 +24,6 @@ class SelectionVector {
     enum class State {
         DYNAMIC,
         STATIC,
-        STATIC_FILTERED,
     };
 
 public:
@@ -50,9 +49,12 @@ class SelectionVector {
     }
     void setRange(sel_t startPos, sel_t size) {
         KU_ASSERT(startPos + size <= capacity);
-        selectedPositions = const_cast<sel_t*>(INCREMENTAL_SELECTED_POS.data()) + startPos;
+        selectedPositions = selectedPositionsBuffer.get();
+        for (auto i = 0u; i < size; ++i) {
+            selectedPositions[i] = startPos + i;
+        }
         selectedSize = size;
-        state = State::STATIC_FILTERED;
+        state = State::DYNAMIC;
     }
 
     // Set to filtered is not very accurate. It sets selectedPositions to a mutable array.

diff --git a/src/include/common/string_utils.h b/src/include/common/string_utils.h
@@ -60,7 +60,7 @@ class KUZU_API StringUtils {
     }
     static std::string_view rtrim(std::string_view input) {
         auto end = input.size();
-        while (end > 0 && isspace(input[end - 1])) {
+        while (end > 0 && isSpace(input[end - 1])) {
             end--;
         }
         return input.substr(0, end);

diff --git a/src/include/planner/join_order/cost_model.h b/src/include/planner/join_order/cost_model.h
@@ -10,8 +10,12 @@ class CostModel {
     static uint64_t computeExtendCost(const LogicalPlan& childPlan);
     static uint64_t computeRecursiveExtendCost(uint8_t upperBound, double extensionRate,
         const LogicalPlan& childPlan);
+    static uint64_t computeHashJoinCost(const std::vector<binder::expression_pair>& joinConditions,
+        const LogicalPlan& probe, const LogicalPlan& build);
     static uint64_t computeHashJoinCost(const binder::expression_vector& joinNodeIDs,
         const LogicalPlan& probe, const LogicalPlan& build);
+    static uint64_t computeMarkJoinCost(const std::vector<binder::expression_pair>& joinConditions,
+        const LogicalPlan& probe, const LogicalPlan& build);
     static uint64_t computeMarkJoinCost(const binder::expression_vector& joinNodeIDs,
         const LogicalPlan& probe, const LogicalPlan& build);
     static uint64_t computeIntersectCost(const LogicalPlan& probePlan,

diff --git a/src/include/planner/planner.h b/src/include/planner/planner.h
@@ -270,12 +270,18 @@ class Planner {
     void appendHashJoin(const binder::expression_vector& joinNodeIDs, common::JoinType joinType,
         std::shared_ptr<binder::Expression> mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
         LogicalPlan& resultPlan);
-    void appendAccHashJoin(const binder::expression_vector& joinNodeIDs, common::JoinType joinType,
-        std::shared_ptr<binder::Expression> mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
-        LogicalPlan& resultPlan);
+    void appendHashJoin(const std::vector<binder::expression_pair>& joinConditions,
+        common::JoinType joinType, std::shared_ptr<binder::Expression> mark, LogicalPlan& probePlan,
+        LogicalPlan& buildPlan, LogicalPlan& resultPlan);
+    void appendAccHashJoin(const std::vector<binder::expression_pair>& joinConditions,
+        common::JoinType joinType, std::shared_ptr<binder::Expression> mark, LogicalPlan& probePlan,
+        LogicalPlan& buildPlan, LogicalPlan& resultPlan);
     void appendMarkJoin(const binder::expression_vector& joinNodeIDs,
         const std::shared_ptr<binder::Expression>& mark, LogicalPlan& probePlan,
-        LogicalPlan& buildPlan);
+        LogicalPlan& buildPlan, LogicalPlan& resultPlan);
+    void appendMarkJoin(const std::vector<binder::expression_pair>& joinConditions,
+        const std::shared_ptr<binder::Expression>& mark, LogicalPlan& probePlan,
+        LogicalPlan& buildPlan, LogicalPlan& resultPlan);
     void appendIntersect(const std::shared_ptr<binder::Expression>& intersectNodeID,
         binder::expression_vector& boundNodeIDs, LogicalPlan& probePlan,
         std::vector<std::unique_ptr<LogicalPlan>>& buildPlans);

diff --git a/src/planner/join_order/cost_model.cpp b/src/planner/join_order/cost_model.cpp
@@ -18,6 +18,23 @@ uint64_t CostModel::computeRecursiveExtendCost(uint8_t upperBound, double extens
            upperBound;
 }
 
+binder::expression_vector getJoinNodeIDs(
+    const std::vector<binder::expression_pair>& joinConditions) {
+    binder::expression_vector joinNodeIDs;
+    for (auto& [left, _] : joinConditions) {
+        if (left->expressionType == ExpressionType::PROPERTY &&
+            left->getDataType().getLogicalTypeID() == LogicalTypeID::INTERNAL_ID) {
+            joinNodeIDs.push_back(left);
+        }
+    }
+    return joinNodeIDs;
+}
+
+uint64_t CostModel::computeHashJoinCost(const std::vector<binder::expression_pair>& joinConditions,
+    const LogicalPlan& probe, const LogicalPlan& build) {
+    return computeHashJoinCost(getJoinNodeIDs(joinConditions), probe, build);
+}
+
 uint64_t CostModel::computeHashJoinCost(const binder::expression_vector& joinNodeIDs,
     const LogicalPlan& probe, const LogicalPlan& build) {
     uint64_t cost = 0ul;
@@ -29,6 +46,11 @@ uint64_t CostModel::computeHashJoinCost(const binder::expression_vector& joinNod
     return cost;
 }
 
+uint64_t CostModel::computeMarkJoinCost(const std::vector<binder::expression_pair>& joinConditions,
+    const LogicalPlan& probe, const LogicalPlan& build) {
+    return computeMarkJoinCost(getJoinNodeIDs(joinConditions), probe, build);
+}
+
 uint64_t CostModel::computeMarkJoinCost(const binder::expression_vector& joinNodeIDs,
     const LogicalPlan& probe, const LogicalPlan& build) {
     return computeHashJoinCost(joinNodeIDs, probe, build);

diff --git a/src/planner/plan/append_join.cpp b/src/planner/plan/append_join.cpp
@@ -21,6 +21,12 @@ void Planner::appendHashJoin(const expression_vector& joinNodeIDs, JoinType join
     for (auto& joinNodeID : joinNodeIDs) {
         joinConditions.emplace_back(joinNodeID, joinNodeID);
     }
+    appendHashJoin(joinConditions, joinType, mark, probePlan, buildPlan, resultPlan);
+}
+
+void Planner::appendHashJoin(const std::vector<expression_pair>& joinConditions, JoinType joinType,
+    std::shared_ptr<Expression> mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
+    LogicalPlan& resultPlan) {
     auto hashJoin = make_shared<LogicalHashJoin>(joinConditions, joinType, mark,
         probePlan.getLastOperator(), buildPlan.getLastOperator());
     // Apply flattening to probe side
@@ -38,26 +44,33 @@ void Planner::appendHashJoin(const expression_vector& joinNodeIDs, JoinType join
     // Update cost
     hashJoin->setCardinality(cardinalityEstimator.estimateHashJoin(joinConditions,
         probePlan.getLastOperatorRef(), buildPlan.getLastOperatorRef()));
-    resultPlan.setCost(CostModel::computeHashJoinCost(joinNodeIDs, probePlan, buildPlan));
+    resultPlan.setCost(CostModel::computeHashJoinCost(joinConditions, probePlan, buildPlan));
     resultPlan.setLastOperator(std::move(hashJoin));
 }
 
-void Planner::appendAccHashJoin(const expression_vector& joinNodeIDs, JoinType joinType,
-    std::shared_ptr<Expression> mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
-    LogicalPlan& resultPlan) {
+void Planner::appendAccHashJoin(const std::vector<binder::expression_pair>& joinConditions,
+    JoinType joinType, std::shared_ptr<Expression> mark, LogicalPlan& probePlan,
+    LogicalPlan& buildPlan, LogicalPlan& resultPlan) {
     KU_ASSERT(probePlan.hasUpdate());
     tryAppendAccumulate(probePlan);
-    appendHashJoin(joinNodeIDs, joinType, mark, probePlan, buildPlan, resultPlan);
+    appendHashJoin(joinConditions, joinType, mark, probePlan, buildPlan, resultPlan);
     auto& sipInfo = probePlan.getLastOperator()->cast<LogicalHashJoin>().getSIPInfoUnsafe();
     sipInfo.direction = SIPDirection::PROBE_TO_BUILD;
 }
 
 void Planner::appendMarkJoin(const expression_vector& joinNodeIDs,
-    const std::shared_ptr<Expression>& mark, LogicalPlan& probePlan, LogicalPlan& buildPlan) {
+    const std::shared_ptr<Expression>& mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
+    LogicalPlan& resultPlan) {
     std::vector<join_condition_t> joinConditions;
     for (auto& joinNodeID : joinNodeIDs) {
         joinConditions.emplace_back(joinNodeID, joinNodeID);
     }
+    appendMarkJoin(joinConditions, mark, probePlan, buildPlan, resultPlan);
+}
+
+void Planner::appendMarkJoin(const std::vector<expression_pair>& joinConditions,
+    const std::shared_ptr<Expression>& mark, LogicalPlan& probePlan, LogicalPlan& buildPlan,
+    LogicalPlan& resultPlan) {
     auto hashJoin = make_shared<LogicalHashJoin>(joinConditions, JoinType::MARK, mark,
         probePlan.getLastOperator(), buildPlan.getLastOperator());
     // Apply flattening to probe side
@@ -69,8 +82,8 @@ void Planner::appendMarkJoin(const expression_vector& joinNodeIDs,
     hashJoin->computeFactorizedSchema();
     // update cost. Mark join does not change cardinality.
     hashJoin->setCardinality(probePlan.getCardinality());
-    probePlan.setCost(CostModel::computeMarkJoinCost(joinNodeIDs, probePlan, buildPlan));
-    probePlan.setLastOperator(std::move(hashJoin));
+    resultPlan.setCost(CostModel::computeMarkJoinCost(joinConditions, probePlan, buildPlan));
+    resultPlan.setLastOperator(std::move(hashJoin));
 }
 
 void Planner::appendIntersect(const std::shared_ptr<Expression>& intersectNodeID,