diff --git a/src/query/frontend/logical/planner.cpp b/src/query/frontend/logical/planner.cpp index 9d08892a3..c924c09e2 100644 --- a/src/query/frontend/logical/planner.cpp +++ b/src/query/frontend/logical/planner.cpp @@ -1,5 +1,6 @@ #include "query/frontend/logical/planner.hpp" +#include #include #include "query/frontend/ast/ast.hpp" @@ -176,6 +177,59 @@ auto GenMatch(Match &match, LogicalOperator *input_op, return last_op; } +// Ast tree visitor which collects all the symbols referenced by identifiers. +class SymbolCollector : public TreeVisitorBase { + public: + SymbolCollector(const SymbolTable &symbol_table) + : symbol_table_(symbol_table) {} + + using TreeVisitorBase::Visit; + using TreeVisitorBase::PostVisit; + + void Visit(Identifier &ident) override { + symbols_.insert(symbol_table_.at(ident)); + } + + const auto &symbols() const { return symbols_; } + + private: + // Calculates the Symbol hash based on its position. + struct SymbolHash { + size_t operator()(const Symbol &symbol) const { + return std::hash{}(symbol.position_); + } + }; + + const SymbolTable &symbol_table_; + std::unordered_set symbols_; +}; + +auto GenWith(With &with, LogicalOperator *input_op, + const query::SymbolTable &symbol_table) { + if (with.distinct_) { + // TODO: Plan disctint with, when operator available. + throw NotYetImplemented(); + } + // WITH clause is Accumulate/Aggregate (advance_command) + Produce. + SymbolCollector symbol_collector(symbol_table); + // Collect used symbols so that accumulate doesn't copy the whole frame. + for (auto &named_expr : with.named_expressions_) { + named_expr->expression_->Accept(symbol_collector); + } + auto symbols = symbol_collector.symbols(); + // TODO: Check whether we need aggregate instead of accumulate. + LogicalOperator *last_op = + new Accumulate(std::shared_ptr(input_op), + std::vector(symbols.begin(), symbols.end()), true); + last_op = new Produce(std::shared_ptr(last_op), + with.named_expressions_); + if (with.where_) { + last_op = new Filter(std::shared_ptr(last_op), + with.where_->expression_); + } + return last_op; +} + } // namespace std::unique_ptr MakeLogicalPlan( @@ -223,6 +277,8 @@ std::unique_ptr MakeLogicalPlan( input_op = new plan::RemoveLabels(std::shared_ptr(input_op), input_symbol, rem->labels_); + } else if (auto *with = dynamic_cast(clause_ptr)) { + input_op = GenWith(*with, input_op, symbol_table); } else { throw NotYetImplemented(); } diff --git a/src/query/frontend/semantic/symbol_generator.cpp b/src/query/frontend/semantic/symbol_generator.cpp index 2ffe84ed8..62e7a9b35 100644 --- a/src/query/frontend/semantic/symbol_generator.cpp +++ b/src/query/frontend/semantic/symbol_generator.cpp @@ -39,6 +39,36 @@ void SymbolGenerator::PostVisit(Return &ret) { } } +void SymbolGenerator::SetWithSymbols(With &with) { + // WITH clause removes declarations of all the previous variables and declares + // only those established through named expressions. New declarations must not + // be visible inside named expressions themselves. + scope_.symbols.clear(); + for (auto &named_expr : with.named_expressions_) { + symbol_table_[*named_expr] = CreateSymbol(named_expr->name_); + } +} + +void SymbolGenerator::Visit(With &with) { + scope_.with = &with; +} + +void SymbolGenerator::Visit(Where &where) { + if (scope_.with) { + // New symbols must be visible in WHERE clause, so this must be done here + // and not in PostVisit(With&). + SetWithSymbols(*scope_.with); + } +} + +void SymbolGenerator::PostVisit(With &with) { + if (!with.where_) { + // This wasn't done when visiting Where, so do it here. + SetWithSymbols(with); + } + scope_.with = nullptr; +} + // Expressions void SymbolGenerator::Visit(Identifier &ident) { diff --git a/src/query/frontend/semantic/symbol_generator.hpp b/src/query/frontend/semantic/symbol_generator.hpp index 6f0ca345c..ab6947aaf 100644 --- a/src/query/frontend/semantic/symbol_generator.hpp +++ b/src/query/frontend/semantic/symbol_generator.hpp @@ -27,6 +27,9 @@ class SymbolGenerator : public TreeVisitorBase { void Visit(Create &create) override; void PostVisit(Create &create) override; void PostVisit(Return &ret) override; + void Visit(With &with) override; + void PostVisit(With &with) override; + void Visit(Where &where) override; // Expressions void Visit(Identifier &ident) override; @@ -53,6 +56,8 @@ class SymbolGenerator : public TreeVisitorBase { bool in_node_atom{false}; bool in_edge_atom{false}; bool in_property_map{false}; + // Pointer to With clause if we are inside it, otherwise nullptr. + With *with{nullptr}; std::map symbols; }; @@ -68,6 +73,9 @@ class SymbolGenerator : public TreeVisitorBase { auto GetOrCreateSymbol(const std::string &name, Symbol::Type type = Symbol::Type::Any); + // Clear old symbol bindings and establish new from WITH clause. + void SetWithSymbols(With &with); + SymbolTable &symbol_table_; Scope scope_; }; diff --git a/tests/unit/query_common.hpp b/tests/unit/query_common.hpp index e7be8b311..bfea7de4a 100644 --- a/tests/unit/query_common.hpp +++ b/tests/unit/query_common.hpp @@ -78,10 +78,15 @@ auto GetQuery(AstTreeStorage &storage, Clause *clause) { storage.query()->clauses_.emplace_back(clause); return storage.query(); } -template -auto GetQuery(AstTreeStorage &storage, Clause *clause, T *... clauses) { - storage.query()->clauses_.emplace_back(clause); - return GetQuery(storage, clauses...); +auto GetQuery(AstTreeStorage &storage, Match *match, Where *where) { + match->where_ = where; + storage.query()->clauses_.emplace_back(match); + return storage.query(); +} +auto GetQuery(AstTreeStorage &storage, With *with, Where *where) { + with->where_ = where; + storage.query()->clauses_.emplace_back(with); + return storage.query(); } template auto GetQuery(AstTreeStorage &storage, Match *match, Where *where, @@ -90,6 +95,18 @@ auto GetQuery(AstTreeStorage &storage, Match *match, Where *where, storage.query()->clauses_.emplace_back(match); return GetQuery(storage, clauses...); } +template +auto GetQuery(AstTreeStorage &storage, With *with, Where *where, + T *... clauses) { + with->where_ = where; + storage.query()->clauses_.emplace_back(with); + return GetQuery(storage, clauses...); +} +template +auto GetQuery(AstTreeStorage &storage, Clause *clause, T *... clauses) { + storage.query()->clauses_.emplace_back(clause); + return GetQuery(storage, clauses...); +} /// /// Create the return clause with given named expressions. @@ -123,6 +140,38 @@ auto GetReturn(AstTreeStorage &storage, T *... exprs) { return GetReturn(ret, exprs...); } +/// +/// Create the with clause with given named expressions. +/// +auto GetWith(With *with, NamedExpression *named_expr) { + with->named_expressions_.emplace_back(named_expr); + return with; +} +auto GetWith(With *with, Expression *expr, NamedExpression *named_expr) { + // This overload supports `RETURN(expr, AS(name))` construct, since + // NamedExpression does not inherit Expression. + named_expr->expression_ = expr; + with->named_expressions_.emplace_back(named_expr); + return with; +} +template +auto GetWith(With *with, Expression *expr, NamedExpression *named_expr, + T *... rest) { + named_expr->expression_ = expr; + with->named_expressions_.emplace_back(named_expr); + return GetWith(with, rest...); +} +template +auto GetWith(With *with, NamedExpression *named_expr, T *... rest) { + with->named_expressions_.emplace_back(named_expr); + return GetWith(with, rest...); +} +template +auto GetWith(AstTreeStorage &storage, T *... exprs) { + auto with = storage.Create(); + return GetWith(with, exprs...); +} + /// /// Create the delete clause with given named expressions. /// @@ -207,10 +256,11 @@ auto GetRemove(AstTreeStorage &storage, const std::string &name, query::test_common::GetPropertyLookup(storage, __VA_ARGS__) #define NEXPR(name, expr) storage.Create((name), (expr)) // AS is alternative to NEXPR which does not initialize NamedExpression with -// Expression. It should be used with RETURN. For example: +// Expression. It should be used with RETURN or WITH. For example: // RETURN(IDENT("n"), AS("n")) vs. RETURN(NEXPR("n", IDENT("n"))). #define AS(name) storage.Create((name)) #define RETURN(...) query::test_common::GetReturn(storage, __VA_ARGS__) +#define WITH(...) query::test_common::GetWith(storage, __VA_ARGS__) #define DELETE(...) query::test_common::GetDelete(storage, {__VA_ARGS__}) #define DETACH_DELETE(...) \ query::test_common::GetDelete(storage, {__VA_ARGS__}, true) diff --git a/tests/unit/query_planner.cpp b/tests/unit/query_planner.cpp index fdb118e2d..fff7205af 100644 --- a/tests/unit/query_planner.cpp +++ b/tests/unit/query_planner.cpp @@ -7,8 +7,8 @@ #include "query/frontend/ast/ast.hpp" #include "query/frontend/logical/operator.hpp" #include "query/frontend/logical/planner.hpp" -#include "query/frontend/semantic/symbol_table.hpp" #include "query/frontend/semantic/symbol_generator.hpp" +#include "query/frontend/semantic/symbol_table.hpp" #include "query_common.hpp" @@ -47,6 +47,7 @@ class PlanChecker : public LogicalOperatorVisitor { void Visit(ExpandUniquenessFilter &op) override { AssertType(op); } + void Visit(Accumulate &op) override { AssertType(op); } std::list types_; @@ -258,4 +259,24 @@ TEST(TestLogicalPlanner, MatchEdgeCycle) { CheckPlan(*query); } +TEST(TestLogicalPlanner, MatchWithReturn) { + // Test MATCH (old) WITH old AS new RETURN new AS new + AstTreeStorage storage; + auto query = QUERY(MATCH(PATTERN(NODE("old"))), WITH(IDENT("old"), AS("new")), + RETURN(IDENT("new"), AS("new"))); + CheckPlan(*query); +} + +TEST(TestLogicalPlanner, MatchWithWhereReturn) { + // Test MATCH (old) WITH old AS new WHERE new.prop < 42 RETURN new AS new + Dbms dbms; + auto dba = dbms.active(); + auto prop = dba->property("prop"); + AstTreeStorage storage; + auto query = QUERY(MATCH(PATTERN(NODE("old"))), WITH(IDENT("old"), AS("new")), + WHERE(LESS(PROPERTY_LOOKUP("new", prop), LITERAL(42))), + RETURN(IDENT("new"), AS("new"))); + CheckPlan(*query); +} + } // namespace diff --git a/tests/unit/query_semantic.cpp b/tests/unit/query_semantic.cpp index 2328fbaf2..1ff97c42c 100644 --- a/tests/unit/query_semantic.cpp +++ b/tests/unit/query_semantic.cpp @@ -2,6 +2,7 @@ #include "gtest/gtest.h" +#include "dbms/dbms.hpp" #include "query/frontend/ast/ast.hpp" #include "query/frontend/interpret/interpret.hpp" #include "query/frontend/semantic/symbol_table.hpp" @@ -266,4 +267,77 @@ TEST(TestSymbolGenerator, CreateDeleteUnbound) { EXPECT_THROW(query->Accept(symbol_generator), UnboundVariableError); } +TEST(TestSymbolGenerator, MatchWithReturn) { + // Test MATCH (old) WITH old AS n RETURN n AS n + AstTreeStorage storage; + auto node = NODE("old"); + auto old_ident = IDENT("old"); + auto with_as_n = AS("n"); + auto n_ident = IDENT("n"); + auto ret_as_n = AS("n"); + auto query = + QUERY(MATCH(PATTERN(node)), WITH(old_ident, with_as_n), RETURN(n_ident, ret_as_n)); + SymbolTable symbol_table; + SymbolGenerator symbol_generator(symbol_table); + query->Accept(symbol_generator); + EXPECT_EQ(symbol_table.max_position(), 3); + auto node_symbol = symbol_table.at(*node->identifier_); + auto old = symbol_table.at(*old_ident); + EXPECT_EQ(node_symbol, old); + auto with_n = symbol_table.at(*with_as_n); + EXPECT_NE(old, with_n); + auto n = symbol_table.at(*n_ident); + EXPECT_EQ(n, with_n); + auto ret_n = symbol_table.at(*ret_as_n); + EXPECT_NE(n, ret_n); +} + +TEST(TestSymbolGenerator, MatchWithReturnUnbound) { + // Test MATCH (old) WITH old AS n RETURN old AS old + AstTreeStorage storage; + auto query = QUERY(MATCH(PATTERN(NODE("old"))), WITH(IDENT("old"), AS("n")), + RETURN(IDENT("old"), AS("old"))); + SymbolTable symbol_table; + SymbolGenerator symbol_generator(symbol_table); + EXPECT_THROW(query->Accept(symbol_generator), UnboundVariableError); +} + +TEST(TestSymbolGenerator, MatchWithWhere) { + // Test MATCH (old) WITH old AS n WHERE n.prop < 42 + Dbms dbms; + auto dba = dbms.active(); + auto prop = dba->property("prop"); + AstTreeStorage storage; + auto node = NODE("old"); + auto old_ident = IDENT("old"); + auto with_as_n = AS("n"); + auto n_prop = PROPERTY_LOOKUP("n", prop); + auto query = QUERY(MATCH(PATTERN(node)), WITH(old_ident, with_as_n), + WHERE(LESS(n_prop, LITERAL(42)))); + SymbolTable symbol_table; + SymbolGenerator symbol_generator(symbol_table); + query->Accept(symbol_generator); + EXPECT_EQ(symbol_table.max_position(), 2); + auto node_symbol = symbol_table.at(*node->identifier_); + auto old = symbol_table.at(*old_ident); + EXPECT_EQ(node_symbol, old); + auto with_n = symbol_table.at(*with_as_n); + EXPECT_NE(old, with_n); + auto n = symbol_table.at(*n_prop->expression_); + EXPECT_EQ(n, with_n); +} + +TEST(TestSymbolGenerator, MatchWithWhereUnbound) { + // Test MATCH (old) WITH old AS n WHERE old.prop < 42 + Dbms dbms; + auto dba = dbms.active(); + auto prop = dba->property("prop"); + AstTreeStorage storage; + auto query = QUERY(MATCH(PATTERN(NODE("old"))), WITH(IDENT("old"), AS("n")), + WHERE(LESS(PROPERTY_LOOKUP("old", prop), LITERAL(42)))); + SymbolTable symbol_table; + SymbolGenerator symbol_generator(symbol_table); + EXPECT_THROW(query->Accept(symbol_generator), UnboundVariableError); +} + }