Merge pull request #1906 from kuzudb/return-after-create

Support return after create
kuzudb · Aug 8, 2023 · 924d22e · 924d22e
2 parents 4a12985 + 51ae0ab
commit 924d22e
Show file tree

Hide file tree

Showing 15 changed files with 167 additions and 47 deletions.
diff --git a/src/include/planner/logical_plan/persistent/logical_create.h b/src/include/planner/logical_plan/persistent/logical_create.h
@@ -9,12 +9,16 @@ namespace planner {
 struct LogicalCreateNodeInfo {
     std::shared_ptr<binder::NodeExpression> node;
     std::shared_ptr<binder::Expression> primaryKey;
+    binder::expression_vector propertiesToReturn;
 
     LogicalCreateNodeInfo(std::shared_ptr<binder::NodeExpression> node,
-        std::shared_ptr<binder::Expression> primaryKey)
-        : node{std::move(node)}, primaryKey{std::move(primaryKey)} {}
+        std::shared_ptr<binder::Expression> primaryKey,
+        binder::expression_vector propertiesToReturn)
+        : node{std::move(node)}, primaryKey{std::move(primaryKey)}, propertiesToReturn{std::move(
+                                                                        propertiesToReturn)} {}
     LogicalCreateNodeInfo(const LogicalCreateNodeInfo& other)
-        : node{other.node}, primaryKey{other.primaryKey} {}
+        : node{other.node}, primaryKey{other.primaryKey}, propertiesToReturn{
+                                                              other.propertiesToReturn} {}
 
     inline std::unique_ptr<LogicalCreateNodeInfo> copy() const {
         return std::make_unique<LogicalCreateNodeInfo>(*this);
@@ -27,12 +31,14 @@ struct LogicalCreateNodeInfo {
 struct LogicalCreateRelInfo {
     std::shared_ptr<binder::RelExpression> rel;
     std::vector<binder::expression_pair> setItems;
+    binder::expression_vector propertiesToReturn;
 
-    LogicalCreateRelInfo(
-        std::shared_ptr<binder::RelExpression> rel, std::vector<binder::expression_pair> setItems)
-        : rel{std::move(rel)}, setItems{std::move(setItems)} {}
+    LogicalCreateRelInfo(std::shared_ptr<binder::RelExpression> rel,
+        std::vector<binder::expression_pair> setItems, binder::expression_vector propertiesToReturn)
+        : rel{std::move(rel)}, setItems{std::move(setItems)}, propertiesToReturn{
+                                                                  std::move(propertiesToReturn)} {}
     LogicalCreateRelInfo(const LogicalCreateRelInfo& other)
-        : rel{other.rel}, setItems{other.setItems} {}
+        : rel{other.rel}, setItems{other.setItems}, propertiesToReturn{other.propertiesToReturn} {}
 
     inline std::unique_ptr<LogicalCreateRelInfo> copy() const {
         return std::make_unique<LogicalCreateRelInfo>(*this);
@@ -76,8 +82,8 @@ class LogicalCreateRel : public LogicalOperator {
         : LogicalOperator{LogicalOperatorType::CREATE_REL, std::move(child)}, infos{std::move(
                                                                                   infos)} {}
 
-    inline void computeFactorizedSchema() final { copyChildSchema(0); }
-    inline void computeFlatSchema() final { copyChildSchema(0); }
+    void computeFactorizedSchema() final;
+    void computeFlatSchema() final;
 
     std::string getExpressionsForPrinting() const final;
 

diff --git a/src/include/planner/query_planner.h b/src/include/planner/query_planner.h
@@ -165,7 +165,7 @@ class QueryPlanner {
 
     // Append scan operators
     void appendScanNodeID(std::shared_ptr<NodeExpression>& node, LogicalPlan& plan);
-    void appendScanNodePropIfNecessary(const expression_vector& propertyExpressions,
+    void appendScanNodeProperties(const expression_vector& propertyExpressions,
         std::shared_ptr<NodeExpression> node, LogicalPlan& plan);
 
     // Append extend operators
@@ -219,8 +219,7 @@ class QueryPlanner {
 
     static std::vector<std::unique_ptr<LogicalPlan>> getInitialEmptyPlans();
 
-    expression_vector getPropertiesForNode(NodeExpression& node);
-    expression_vector getPropertiesForRel(RelExpression& rel);
+    expression_vector getProperties(const binder::Expression& nodeOrRel);
 
     std::unique_ptr<JoinOrderEnumeratorContext> enterContext(
         binder::expression_vector nodeIDsToScanFromInnerAndOuter);

diff --git a/src/include/processor/operator/persistent/insert_executor.h b/src/include/processor/operator/persistent/insert_executor.h
@@ -42,9 +42,11 @@ class RelInsertExecutor {
 public:
     RelInsertExecutor(storage::RelsStatistics& relsStatistics, storage::RelTable* table,
         const DataPos& srcNodePos, const DataPos& dstNodePos,
+        std::vector<DataPos> lhsVectorPositions,
         std::vector<std::unique_ptr<evaluator::ExpressionEvaluator>> evaluators)
         : relsStatistics{relsStatistics}, table{table}, srcNodePos{srcNodePos},
-          dstNodePos{dstNodePos}, evaluators{std::move(evaluators)} {}
+          dstNodePos{dstNodePos}, lhsVectorPositions{std::move(lhsVectorPositions)},
+          evaluators{std::move(evaluators)} {}
     RelInsertExecutor(const RelInsertExecutor& other);
 
     void init(ResultSet* resultSet, ExecutionContext* context);
@@ -63,11 +65,13 @@ class RelInsertExecutor {
     storage::RelTable* table;
     DataPos srcNodePos;
     DataPos dstNodePos;
+    std::vector<DataPos> lhsVectorPositions;
     std::vector<std::unique_ptr<evaluator::ExpressionEvaluator>> evaluators;
 
     common::ValueVector* srcNodeIDVector = nullptr;
     common::ValueVector* dstNodeIDVector = nullptr;
-    std::vector<common::ValueVector*> propertyVectors;
+    std::vector<common::ValueVector*> lhsVectors;
+    std::vector<common::ValueVector*> rhsVectors;
 };
 
 } // namespace processor

diff --git a/src/include/processor/plan_mapper.h b/src/include/processor/plan_mapper.h
@@ -113,7 +113,8 @@ class PlanMapper {
         storage::RelsStore* relsStore, planner::LogicalCreateNodeInfo* info,
         const planner::Schema& inSchema, const planner::Schema& outSchema);
     std::unique_ptr<RelInsertExecutor> getRelInsertExecutor(storage::RelsStore* relsStore,
-        planner::LogicalCreateRelInfo* info, const planner::Schema& inSchema);
+        planner::LogicalCreateRelInfo* info, const planner::Schema& inSchema,
+        const planner::Schema& outSchema);
     std::unique_ptr<NodeSetExecutor> getNodeSetExecutor(storage::NodesStore* store,
         planner::LogicalSetPropertyInfo* info, const planner::Schema& inSchema);
     std::unique_ptr<RelSetExecutor> getRelSetExecutor(storage::RelsStore* store,

diff --git a/src/planner/operator/persistent/logical_create.cpp b/src/planner/operator/persistent/logical_create.cpp
@@ -30,14 +30,20 @@ void LogicalCreateNode::computeFactorizedSchema() {
     for (auto& info : infos) {
         auto groupPos = schema->createGroup();
         schema->setGroupAsSingleState(groupPos);
-        schema->insertToGroupAndScope(info->node->getInternalIDProperty(), groupPos);
+        for (auto& property : info->propertiesToReturn) {
+            schema->insertToGroupAndScope(property, groupPos);
+        }
+        schema->insertToGroupAndScopeMayRepeat(info->node->getInternalIDProperty(), groupPos);
     }
 }
 
 void LogicalCreateNode::computeFlatSchema() {
     copyChildSchema(0);
     for (auto& info : infos) {
-        schema->insertToGroupAndScope(info->node->getInternalIDProperty(), 0);
+        for (auto& property : info->propertiesToReturn) {
+            schema->insertToGroupAndScope(property, 0);
+        }
+        schema->insertToGroupAndScopeMayRepeat(info->node->getInternalIDProperty(), 0);
     }
 }
 
@@ -57,6 +63,26 @@ f_group_pos_set LogicalCreateNode::getGroupsPosToFlatten() {
         childSchema->getGroupsPosInScope(), childSchema);
 }
 
+void LogicalCreateRel::computeFactorizedSchema() {
+    copyChildSchema(0);
+    for (auto& info : infos) {
+        auto groupPos = schema->createGroup();
+        schema->setGroupAsSingleState(groupPos);
+        for (auto& property : info->propertiesToReturn) {
+            schema->insertToGroupAndScope(property, groupPos);
+        }
+    }
+}
+
+void LogicalCreateRel::computeFlatSchema() {
+    copyChildSchema(0);
+    for (auto& info : infos) {
+        for (auto& property : info->propertiesToReturn) {
+            schema->insertToGroupAndScope(property, 0);
+        }
+    }
+}
+
 std::string LogicalCreateRel::getExpressionsForPrinting() const {
     std::string result;
     for (auto& info : infos) {

diff --git a/src/planner/plan/append_create.cpp b/src/planner/plan/append_create.cpp
@@ -11,14 +11,18 @@ namespace planner {
 std::unique_ptr<LogicalCreateNodeInfo> QueryPlanner::createLogicalCreateNodeInfo(
     BoundCreateInfo* boundCreateInfo) {
     auto node = std::static_pointer_cast<NodeExpression>(boundCreateInfo->nodeOrRel);
+    auto propertiesToReturn = getProperties(*node);
     auto extraInfo = (ExtraCreateNodeInfo*)boundCreateInfo->extraInfo.get();
-    return std::make_unique<LogicalCreateNodeInfo>(node, extraInfo->primaryKey);
+    return std::make_unique<LogicalCreateNodeInfo>(
+        node, extraInfo->primaryKey, std::move(propertiesToReturn));
 }
 
 std::unique_ptr<LogicalCreateRelInfo> QueryPlanner::createLogicalCreateRelInfo(
     BoundCreateInfo* boundCreateInfo) {
     auto rel = std::static_pointer_cast<RelExpression>(boundCreateInfo->nodeOrRel);
-    return std::make_unique<LogicalCreateRelInfo>(rel, boundCreateInfo->setItems);
+    auto propertiesToReturn = getProperties(*rel);
+    return std::make_unique<LogicalCreateRelInfo>(
+        rel, boundCreateInfo->setItems, std::move(propertiesToReturn));
 }
 
 std::vector<std::unique_ptr<BoundSetPropertyInfo>> QueryPlanner::createLogicalSetPropertyInfo(

diff --git a/src/planner/plan/append_extend.cpp b/src/planner/plan/append_extend.cpp
@@ -176,7 +176,7 @@ void QueryPlanner::createPathNodePropertyScanPlan(
     for (auto& property : recursiveNode->getPropertyExpressions()) {
         properties.push_back(property->copy());
     }
-    appendScanNodePropIfNecessary(properties, recursiveNode, plan);
+    appendScanNodeProperties(properties, recursiveNode, plan);
     auto expressionsToProject = properties;
     expressionsToProject.push_back(recursiveNode->getInternalIDProperty());
     expressionsToProject.push_back(recursiveNode->getLabelExpression());

diff --git a/src/planner/plan/append_scan_node.cpp b/src/planner/plan/append_scan_node.cpp
@@ -14,7 +14,7 @@ void QueryPlanner::appendScanNodeID(std::shared_ptr<NodeExpression>& node, Logic
     plan.setLastOperator(std::move(scan));
 }
 
-void QueryPlanner::appendScanNodePropIfNecessary(const expression_vector& propertyExpressions,
+void QueryPlanner::appendScanNodeProperties(const expression_vector& propertyExpressions,
     std::shared_ptr<NodeExpression> node, LogicalPlan& plan) {
     expression_vector propertyExpressionToScan;
     for (auto& propertyExpression : propertyExpressions) {

diff --git a/src/planner/plan/plan_join_order.cpp b/src/planner/plan/plan_join_order.cpp
@@ -146,8 +146,8 @@ void QueryPlanner::planNodeScan(uint32_t nodePos) {
     // query, we only scan internal ID of "a".
     if (!context->nodeToScanFromInnerAndOuter(node.get())) {
         appendScanNodeID(node, *plan);
-        auto properties = getPropertiesForNode(*node);
-        appendScanNodePropIfNecessary(properties, node, *plan);
+        auto properties = getProperties(*node);
+        appendScanNodeProperties(properties, node, *plan);
         auto predicates = getNewlyMatchedExpressions(
             context->getEmptySubqueryGraph(), newSubgraph, context->getWhereExpressions());
         appendFilters(predicates, *plan);
@@ -188,7 +188,7 @@ void QueryPlanner::appendExtendAndFilter(std::shared_ptr<NodeExpression> boundNo
     ExtendDirection direction, const expression_vector& predicates, LogicalPlan& plan) {
     switch (rel->getRelType()) {
     case common::QueryRelType::NON_RECURSIVE: {
-        auto properties = getPropertiesForRel(*rel);
+        auto properties = getProperties(*rel);
         appendNonRecursiveExtend(boundNode, nbrNode, rel, direction, properties, plan);
     } break;
     case common::QueryRelType::VARIABLE_LENGTH:

diff --git a/src/planner/query_planner.cpp b/src/planner/query_planner.cpp
@@ -94,22 +94,15 @@ std::vector<std::unique_ptr<LogicalPlan>> QueryPlanner::getInitialEmptyPlans() {
     return plans;
 }
 
-expression_vector QueryPlanner::getPropertiesForNode(NodeExpression& node) {
-    expression_vector result;
-    for (auto& expression : propertiesToScan) {
-        auto property = (PropertyExpression*)expression.get();
-        if (property->getVariableName() == node.getUniqueName()) {
-            result.push_back(expression);
-        }
+expression_vector QueryPlanner::getProperties(const binder::Expression& nodeOrRel) {
+    auto typeID = nodeOrRel.getDataType().getLogicalTypeID();
+    if (typeID != common::LogicalTypeID::NODE && typeID != common::LogicalTypeID::REL) {
+        throw common::NotImplementedException("QueryPlanner::getProperties");
     }
-    return result;
-}
-
-expression_vector QueryPlanner::getPropertiesForRel(RelExpression& rel) {
     expression_vector result;
     for (auto& expression : propertiesToScan) {
         auto property = (PropertyExpression*)expression.get();
-        if (property->getVariableName() == rel.getUniqueName()) {
+        if (property->getVariableName() == nodeOrRel.getUniqueName()) {
             result.push_back(expression);
         }
     }

diff --git a/src/processor/map/map_create.cpp b/src/processor/map/map_create.cpp
@@ -48,30 +48,38 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapCreateNode(LogicalOperator* log
 }
 
 std::unique_ptr<RelInsertExecutor> PlanMapper::getRelInsertExecutor(storage::RelsStore* relsStore,
-    planner::LogicalCreateRelInfo* info, const planner::Schema& inSchema) {
+    planner::LogicalCreateRelInfo* info, const planner::Schema& inSchema,
+    const planner::Schema& outSchema) {
     auto rel = info->rel;
     auto relTableID = rel->getSingleTableID();
     auto table = relsStore->getRelTable(relTableID);
     auto srcNode = rel->getSrcNode();
     auto dstNode = rel->getDstNode();
     auto srcNodePos = DataPos(inSchema.getExpressionPos(*srcNode->getInternalIDProperty()));
     auto dstNodePos = DataPos(inSchema.getExpressionPos(*dstNode->getInternalIDProperty()));
+    std::vector<DataPos> lhsVectorPositions;
     std::vector<std::unique_ptr<ExpressionEvaluator>> evaluators;
     for (auto& [lhs, rhs] : info->setItems) {
+        if (outSchema.isExpressionInScope(*lhs)) {
+            lhsVectorPositions.emplace_back(outSchema.getExpressionPos(*lhs));
+        } else {
+            lhsVectorPositions.emplace_back(INVALID_DATA_CHUNK_POS, INVALID_DATA_CHUNK_POS);
+        }
         evaluators.push_back(expressionMapper.mapExpression(rhs, inSchema));
     }
-    return std::make_unique<RelInsertExecutor>(
-        relsStore->getRelsStatistics(), table, srcNodePos, dstNodePos, std::move(evaluators));
+    return std::make_unique<RelInsertExecutor>(relsStore->getRelsStatistics(), table, srcNodePos,
+        dstNodePos, std::move(lhsVectorPositions), std::move(evaluators));
 }
 
 std::unique_ptr<PhysicalOperator> PlanMapper::mapCreateRel(LogicalOperator* logicalOperator) {
     auto logicalCreateRel = (LogicalCreateRel*)logicalOperator;
     auto inSchema = logicalCreateRel->getChild(0)->getSchema();
+    auto outSchema = logicalCreateRel->getSchema();
     auto prevOperator = mapOperator(logicalOperator->getChild(0).get());
     std::vector<std::unique_ptr<RelInsertExecutor>> executors;
     for (auto& info : logicalCreateRel->getInfosRef()) {
-        executors.push_back(
-            getRelInsertExecutor(&storageManager.getRelsStore(), info.get(), *inSchema));
+        executors.push_back(getRelInsertExecutor(
+            &storageManager.getRelsStore(), info.get(), *inSchema, *outSchema));
     }
     return std::make_unique<InsertRel>(std::move(executors), std::move(prevOperator),
         getOperatorID(), logicalCreateRel->getExpressionsForPrinting());

diff --git a/src/processor/map/map_merge.cpp b/src/processor/map/map_merge.cpp
@@ -26,7 +26,8 @@ std::unique_ptr<PhysicalOperator> PlanMapper::mapMerge(planner::LogicalOperator*
     }
     std::vector<std::unique_ptr<RelInsertExecutor>> relInsertExecutors;
     for (auto& info : logicalMerge->getCreateRelInfosRef()) {
-        relInsertExecutors.push_back(getRelInsertExecutor(relsStore, info.get(), *inSchema));
+        relInsertExecutors.push_back(
+            getRelInsertExecutor(relsStore, info.get(), *inSchema, *outSchema));
     }
     std::vector<std::unique_ptr<NodeSetExecutor>> onCreateNodeSetExecutors;
     for (auto& info : logicalMerge->getOnCreateSetNodeInfosRef()) {

diff --git a/src/processor/operator/persistent/insert_executor.cpp b/src/processor/operator/persistent/insert_executor.cpp
@@ -1,5 +1,7 @@
 #include "processor/operator/persistent/insert_executor.h"
 
+using namespace kuzu::common;
+
 namespace kuzu {
 namespace processor {
 
@@ -45,7 +47,7 @@ std::vector<std::unique_ptr<NodeInsertExecutor>> NodeInsertExecutor::copy(
 
 RelInsertExecutor::RelInsertExecutor(const RelInsertExecutor& other)
     : relsStatistics{other.relsStatistics}, table{other.table}, srcNodePos{other.srcNodePos},
-      dstNodePos{other.dstNodePos} {
+      dstNodePos{other.dstNodePos}, lhsVectorPositions{other.lhsVectorPositions} {
     for (auto& evaluator : other.evaluators) {
         evaluators.push_back(evaluator->clone());
     }
@@ -54,21 +56,45 @@ RelInsertExecutor::RelInsertExecutor(const RelInsertExecutor& other)
 void RelInsertExecutor::init(ResultSet* resultSet, ExecutionContext* context) {
     srcNodeIDVector = resultSet->getValueVector(srcNodePos).get();
     dstNodeIDVector = resultSet->getValueVector(dstNodePos).get();
+    for (auto& pos : lhsVectorPositions) {
+        if (pos.dataChunkPos != INVALID_DATA_CHUNK_POS) {
+            lhsVectors.push_back(resultSet->getValueVector(pos).get());
+        } else {
+            lhsVectors.push_back(nullptr);
+        }
+    }
     for (auto& evaluator : evaluators) {
         evaluator->init(*resultSet, context->memoryManager);
-        propertyVectors.push_back(evaluator->resultVector.get());
+        rhsVectors.push_back(evaluator->resultVector.get());
     }
 }
 
 void RelInsertExecutor::insert(transaction::Transaction* tx) {
     auto offset = relsStatistics.getNextRelOffset(tx, table->getRelTableID());
-    propertyVectors[0]->setValue(0, offset);
-    propertyVectors[0]->setNull(0, false);
+    rhsVectors[0]->setValue(0, offset);
+    rhsVectors[0]->setNull(0, false);
     for (auto i = 1; i < evaluators.size(); ++i) {
         evaluators[i]->evaluate();
     }
-    table->insertRel(srcNodeIDVector, dstNodeIDVector, propertyVectors);
+    table->insertRel(srcNodeIDVector, dstNodeIDVector, rhsVectors);
     relsStatistics.updateNumRelsByValue(table->getRelTableID(), 1);
+    for (auto i = 0u; i < lhsVectors.size(); ++i) {
+        auto lhsVector = lhsVectors[i];
+        auto rhsVector = rhsVectors[i];
+        if (lhsVector == nullptr) {
+            continue;
+        }
+        assert(lhsVector->state->selVector->selectedSize == 1 &&
+               rhsVector->state->selVector->selectedSize == 1);
+        auto lhsPos = lhsVector->state->selVector->selectedPositions[0];
+        auto rhsPos = rhsVector->state->selVector->selectedPositions[0];
+        if (rhsVector->isNull(rhsPos)) {
+            lhsVector->setNull(lhsPos, true);
+        } else {
+            lhsVector->setNull(lhsPos, false);
+            lhsVector->copyFromVectorData(lhsPos, rhsVector, rhsPos);
+        }
+    }
 }
 
 std::vector<std::unique_ptr<RelInsertExecutor>> RelInsertExecutor::copy(

diff --git a/test/test_files/tinysnb/update_node/create_read.test b/test/test_files/tinysnb/update_node/create_read.test
@@ -0,0 +1,19 @@
+-GROUP TinySnbUpdateTest
+-DATASET CSV tinysnb
+
+--
+
+-CASE CreateNodeRead1
+-STATEMENT CREATE (a:person {ID:80, isWorker:true,age:22,eyeSight:1.1}) RETURN a.ID, a.age, a.fName, a.eyeSight;
+---- 1
+80|22||1.100000
+-STATEMENT CREATE (a:organisation {ID:0, name:'test'}) RETURN a, a.history;
+---- 1
+{_ID: 1:3, _LABEL: organisation, ID: 0, name: test}|
+
+-CASE CreateNodeRead2
+-STATEMENT MATCH (a:person) WHERE a.ID < 3 CREATE (b:person {ID: a.ID + 11, fName: 'new', age:a.age * 2})
+           RETURN a.ID, a.fName, a.age, b.ID, b.fName, b.age
+---- 2
+0|Alice|35|11|new|70
+2|Bob|30|13|new|60