Convert bag and distinct from antlr to AST

Reviewers: teon.banek, florijan

Reviewed By: teon.banek

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D308

src/query/frontend/ast/cypher_main_visitor.cpp

@@ -15,8 +15,7 @@
 #include "utils/assert.hpp"
 #include "utils/exceptions.hpp"
 
-namespace query {
-namespace frontend {
+namespace query::frontend {
 
 const std::string CypherMainVisitor::kAnonPrefix = "anon";
 
@@ -146,28 +145,36 @@ antlrcpp::Any CypherMainVisitor::visitCreate(CypherParser::CreateContext *ctx) {
 antlrcpp::Any CypherMainVisitor::visitCypherReturn(
     CypherParser::CypherReturnContext *ctx) {
   auto *return_clause = storage_.Create<Return>();
+  return_clause->body_ = ctx->returnBody()->accept(this);
   if (ctx->DISTINCT()) {
-    // TODO: implement other clauses.
-    throw utils::NotYetImplemented();
+    return_clause->body_.distinct = true;
   }
-  return_clause->body_.named_expressions =
-      ctx->returnBody()->accept(this).as<std::vector<NamedExpression *>>();
   return return_clause;
 }
 
 antlrcpp::Any CypherMainVisitor::visitReturnBody(
     CypherParser::ReturnBodyContext *ctx) {
-  if (ctx->order() || ctx->skip() || ctx->limit()) {
-    // TODO: implement other clauses.
-    throw utils::NotYetImplemented();
+  ReturnBody body;
+  if (ctx->order()) {
+    body.order_by = ctx->order()
+                        ->accept(this)
+                        .as<std::vector<std::pair<Ordering, Expression *>>>();
   }
-  return ctx->returnItems()->accept(this);
+  if (ctx->skip()) {
+    body.skip = static_cast<Expression *>(ctx->skip()->accept(this));
+  }
+  if (ctx->limit()) {
+    body.limit = static_cast<Expression *>(ctx->limit()->accept(this));
+  }
+  body.named_expressions =
+      ctx->returnItems()->accept(this).as<std::vector<NamedExpression *>>();
+  return body;
 }
 
 antlrcpp::Any CypherMainVisitor::visitReturnItems(
     CypherParser::ReturnItemsContext *ctx) {
   if (ctx->getTokens(kReturnAllTokenId).size()) {
-    // TODO: implement other clauses.
+    // TODO: implement *
     throw utils::NotYetImplemented();
   }
   std::vector<NamedExpression *> named_expressions;
@@ -191,6 +198,21 @@ antlrcpp::Any CypherMainVisitor::visitReturnItem(
   return named_expr;
 }
 
+antlrcpp::Any CypherMainVisitor::visitOrder(CypherParser::OrderContext *ctx) {
+  std::vector<std::pair<Ordering, Expression *>> order_by;
+  for (auto *sort_item : ctx->sortItem()) {
+    order_by.push_back(sort_item->accept(this));
+  }
+  return order_by;
+}
+
+antlrcpp::Any CypherMainVisitor::visitSortItem(
+    CypherParser::SortItemContext *ctx) {
+  return std::pair<Ordering, Expression *>(
+      ctx->DESC() || ctx->DESCENDING() ? Ordering::DESC : Ordering::ASC,
+      ctx->expression()->accept(this));
+}
+
 antlrcpp::Any CypherMainVisitor::visitNodePattern(
     CypherParser::NodePatternContext *ctx) {
   auto *node = storage_.Create<NodeAtom>();
@@ -907,16 +929,13 @@ antlrcpp::Any CypherMainVisitor::visitPropertyExpression(
 
 antlrcpp::Any CypherMainVisitor::visitWith(CypherParser::WithContext *ctx) {
   auto *with = storage_.Create<With>();
+  with->body_ = ctx->returnBody()->accept(this);
   if (ctx->DISTINCT()) {
-    // TODO: implement this
-    throw utils::NotYetImplemented();
+    with->body_.distinct = true;
   }
-  with->body_.named_expressions =
-      ctx->returnBody()->accept(this).as<std::vector<NamedExpression *>>();
   if (ctx->where()) {
     with->where_ = ctx->where()->accept(this);
   }
   return with;
 }
 }
-}

src/query/frontend/ast/cypher_main_visitor.hpp

@@ -174,6 +174,16 @@ class CypherMainVisitor : public antlropencypher::CypherBaseVisitor {
    */
   antlrcpp::Any visitReturnItem(CypherParser::ReturnItemContext *ctx) override;
 
+  /**
+   * @return vector<pair<Ordering, Expression*>>
+   */
+  antlrcpp::Any visitOrder(CypherParser::OrderContext *ctx) override;
+
+  /**
+   * @return pair<Ordering, Expression*>
+   */
+  antlrcpp::Any visitSortItem(CypherParser::SortItemContext *ctx) override;
+
   /**
    * @return NodeAtom*
    */

tests/unit/cypher_main_visitor.cpp

@@ -68,6 +68,64 @@ TEST(CypherMainVisitorTest, PropertyLookup) {
             ast_generator.db_accessor_->property("x"));
 }
 
+TEST(CypherMainVisitor, ReturnNoDistinctNoBagSemantics) {
+  AstGenerator ast_generator("RETURN x");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  ASSERT_EQ(return_clause->body_.order_by.size(), 0U);
+  ASSERT_EQ(return_clause->body_.named_expressions.size(), 1U);
+  ASSERT_FALSE(return_clause->body_.limit);
+  ASSERT_FALSE(return_clause->body_.skip);
+  ASSERT_FALSE(return_clause->body_.distinct);
+}
+
+TEST(CypherMainVisitor, ReturnDistinct) {
+  AstGenerator ast_generator("RETURN DISTINCT x");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  ASSERT_TRUE(return_clause->body_.distinct);
+}
+
+TEST(CypherMainVisitor, ReturnLimit) {
+  AstGenerator ast_generator("RETURN x LIMIT 5");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  ASSERT_TRUE(return_clause->body_.limit);
+  auto *literal = dynamic_cast<Literal *>(return_clause->body_.limit);
+  ASSERT_TRUE(literal);
+  ASSERT_EQ(literal->value_.Value<int64_t>(), 5);
+}
+
+TEST(CypherMainVisitor, ReturnSkip) {
+  AstGenerator ast_generator("RETURN x SKIP 5");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  ASSERT_TRUE(return_clause->body_.skip);
+  auto *literal = dynamic_cast<Literal *>(return_clause->body_.skip);
+  ASSERT_TRUE(literal);
+  ASSERT_EQ(literal->value_.Value<int64_t>(), 5);
+}
+
+TEST(CypherMainVisitor, ReturnOrderBy) {
+  AstGenerator ast_generator("RETURN x, y, z ORDER BY z ASC, x, y DESC");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 1U);
+  auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
+  ASSERT_EQ(return_clause->body_.order_by.size(), 3U);
+  std::vector<std::pair<Ordering, std::string>> ordering;
+  for (const auto &sort_item : return_clause->body_.order_by) {
+    auto *identifier = dynamic_cast<Identifier *>(sort_item.second);
+    ordering.emplace_back(sort_item.first, identifier->name_);
+  }
+  ASSERT_THAT(ordering, UnorderedElementsAre(Pair(Ordering::ASC, "z"),
+                                             Pair(Ordering::ASC, "x"),
+                                             Pair(Ordering::DESC, "y")));
+}
+
 TEST(CypherMainVisitorTest, ReturnNamedIdentifier) {
   AstGenerator ast_generator("RETURN var AS var5");
   auto *query = ast_generator.query_;
@@ -262,7 +320,8 @@ TEST(CypherMainVisitorTest, ComparisonOperators) {
   AstGenerator ast_generator("RETURN 2 = 3 != 4 <> 5 < 6 > 7 <= 8 >= 9");
   auto *query = ast_generator.query_;
   auto *return_clause = dynamic_cast<Return *>(query->clauses_[0]);
-  Expression *_operator = return_clause->body_.named_expressions[0]->expression_;
+  Expression *_operator =
+      return_clause->body_.named_expressions[0]->expression_;
   CHECK_COMPARISON(GreaterEqualOperator, 8, 9);
   CHECK_COMPARISON(LessEqualOperator, 7, 8);
   CHECK_COMPARISON(GreaterOperator, 6, 7);
@@ -786,6 +845,10 @@ TEST(CypherMainVisitorTest, With) {
   ASSERT_EQ(query->clauses_.size(), 2U);
   auto *with = dynamic_cast<With *>(query->clauses_[0]);
   ASSERT_TRUE(with);
+  ASSERT_FALSE(with->body_.distinct);
+  ASSERT_FALSE(with->body_.limit);
+  ASSERT_FALSE(with->body_.skip);
+  ASSERT_EQ(with->body_.order_by.size(), 0U);
   ASSERT_FALSE(with->where_);
   ASSERT_EQ(with->body_.named_expressions.size(), 1U);
   auto *named_expr = with->body_.named_expressions[0];
@@ -794,6 +857,34 @@ TEST(CypherMainVisitorTest, With) {
   ASSERT_EQ(identifier->name_, "n");
 }
 
+TEST(CypherMainVisitorTest, WithDistinct) {
+  AstGenerator ast_generator("WITH DISTINCT n AS m RETURN 1");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 2U);
+  auto *with = dynamic_cast<With *>(query->clauses_[0]);
+  ASSERT_TRUE(with->body_.distinct);
+  ASSERT_FALSE(with->where_);
+  ASSERT_EQ(with->body_.named_expressions.size(), 1U);
+  auto *named_expr = with->body_.named_expressions[0];
+  ASSERT_EQ(named_expr->name_, "m");
+  auto *identifier = dynamic_cast<Identifier *>(named_expr->expression_);
+  ASSERT_EQ(identifier->name_, "n");
+}
+
+TEST(CypherMainVisitorTest, WithBag) {
+  AstGenerator ast_generator("WITH n as m ORDER BY m SKIP 1 LIMIT 2 RETURN 1");
+  auto *query = ast_generator.query_;
+  ASSERT_EQ(query->clauses_.size(), 2U);
+  auto *with = dynamic_cast<With *>(query->clauses_[0]);
+  ASSERT_FALSE(with->body_.distinct);
+  ASSERT_FALSE(with->where_);
+  ASSERT_EQ(with->body_.named_expressions.size(), 1U);
+  // No need to check contents of body. That is checked in RETURN clause tests.
+  ASSERT_EQ(with->body_.order_by.size(), 1U);
+  ASSERT_TRUE(with->body_.limit);
+  ASSERT_TRUE(with->body_.skip);
+}
+
 TEST(CypherMainVisitorTest, WithWhere) {
   AstGenerator ast_generator("WITH n AS m WHERE k RETURN 1");
   auto *query = ast_generator.query_;