Plan Distinct

Summary:
Support WITH/RETURN DISTINCT in test macros.
Test planning Distinct.
Implement OutputSymbols in Distinct operator.

Reviewers: florijan, mislav.bradac

Reviewed By: mislav.bradac

Subscribers: pullbot

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

src/query/plan/operator.cpp

@@ -1587,6 +1587,11 @@ std::unique_ptr<Cursor> Distinct::MakeCursor(GraphDbAccessor &db) {
   return std::make_unique<DistinctCursor>(*this, db);
 }
 
+std::vector<Symbol> Distinct::OutputSymbols(const SymbolTable &symbol_table) {
+  // Propagate this to potential Produce.
+  return input_->OutputSymbols(symbol_table);
+}
+
 Distinct::DistinctCursor::DistinctCursor(Distinct &self, GraphDbAccessor &db)
     : self_(self), input_cursor_(self.input_->MakeCursor(db)) {}
 

src/query/plan/operator.hpp

@@ -1311,6 +1311,7 @@ class Distinct : public LogicalOperator {
            const std::vector<Symbol> &value_symbols);
   void Accept(LogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(GraphDbAccessor &db) override;
+  std::vector<Symbol> OutputSymbols(const SymbolTable &) override;
 
  private:
   const std::shared_ptr<LogicalOperator> input_;

src/query/plan/planner.cpp

@@ -364,10 +364,6 @@ class ReturnBodyContext : public TreeVisitorBase {
 
 auto GenReturnBody(LogicalOperator *input_op, bool advance_command,
                    const ReturnBodyContext &body, bool accumulate = false) {
-  if (body.distinct()) {
-    // TODO: Plan with distinct, when operator available.
-    throw utils::NotYetImplemented();
-  }
   std::vector<Symbol> used_symbols(body.used_symbols().begin(),
                                    body.used_symbols().end());
   auto last_op = input_op;
@@ -390,6 +386,12 @@ auto GenReturnBody(LogicalOperator *input_op, bool advance_command,
     last_op = new Filter(std::shared_ptr<LogicalOperator>(last_op),
                          body.where()->expression_);
   }
+  // Distinct in ReturnBody only makes Produce values unique, so plan after it.
+  // Hopefully, it is more efficient to have Filter before Distinct.
+  if (body.distinct()) {
+    last_op = new Distinct(std::shared_ptr<LogicalOperator>(last_op),
+                           body.output_symbols());
+  }
   // Like Where, OrderBy can read from symbols established by named expressions
   // in Produce, so it must come after it.
   if (!body.order_by().empty()) {
@@ -538,7 +540,8 @@ std::unique_ptr<LogicalOperator> MakeLogicalPlan(
                                   unwind->named_expression_->expression_,
                                   symbol_table.at(*unwind->named_expression_));
     } else {
-      throw utils::NotYetImplemented();
+      throw utils::NotYetImplemented(
+          "Encountered a clause which cannot be converted to operator(s)");
     }
   }
   return std::unique_ptr<LogicalOperator>(input_op);

tests/unit/query_common.hpp

@@ -233,8 +233,9 @@ void FillReturnBody(ReturnBody &body, NamedExpression *named_expr, T... rest) {
 ///
 /// @sa GetWith
 template <class... T>
-auto GetReturn(AstTreeStorage &storage, T... exprs) {
+auto GetReturn(AstTreeStorage &storage, bool distinct, T... exprs) {
   auto ret = storage.Create<Return>();
+  ret->body_.distinct = distinct;
   FillReturnBody(ret->body_, exprs...);
   return ret;
 }
@@ -246,8 +247,9 @@ auto GetReturn(AstTreeStorage &storage, T... exprs) {
 ///
 /// @sa GetReturn
 template <class... T>
-auto GetWith(AstTreeStorage &storage, T... exprs) {
+auto GetWith(AstTreeStorage &storage, bool distinct, T... exprs) {
   auto with = storage.Create<With>();
+  with->body_.distinct = distinct;
   FillReturnBody(with->body_, exprs...);
   return with;
 }
@@ -378,8 +380,12 @@ auto GetMerge(AstTreeStorage &storage, Pattern *pattern, OnMatch on_match,
 // 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<query::NamedExpression>((name))
-#define RETURN(...) query::test_common::GetReturn(storage, __VA_ARGS__)
-#define WITH(...) query::test_common::GetWith(storage, __VA_ARGS__)
+#define RETURN(...) query::test_common::GetReturn(storage, false, __VA_ARGS__)
+#define WITH(...) query::test_common::GetWith(storage, false, __VA_ARGS__)
+#define RETURN_DISTINCT(...) \
+  query::test_common::GetReturn(storage, true, __VA_ARGS__)
+#define WITH_DISTINCT(...) \
+  query::test_common::GetWith(storage, true, __VA_ARGS__)
 #define UNWIND(...) query::test_common::GetUnwind(storage, __VA_ARGS__)
 #define ORDER_BY(...) query::test_common::GetOrderBy(__VA_ARGS__)
 #define SKIP(expr) \

tests/unit/query_planner.cpp

@@ -73,6 +73,7 @@ class PlanChecker : public LogicalOperatorVisitor {
     return false;
   }
   void Visit(Unwind &op) override { CheckOp(op); }
+  void Visit(Distinct &op) override { CheckOp(op); }
 
   std::list<BaseOpChecker *> checkers_;
 
@@ -119,6 +120,7 @@ using ExpectSkip = OpChecker<Skip>;
 using ExpectLimit = OpChecker<Limit>;
 using ExpectOrderBy = OpChecker<OrderBy>;
 using ExpectUnwind = OpChecker<Unwind>;
+using ExpectDistinct = OpChecker<Distinct>;
 
 class ExpectAccumulate : public OpChecker<Accumulate> {
  public:
@@ -680,4 +682,32 @@ TEST(TestLogicalPlanner, MatchUnwindReturn) {
   CheckPlan(*query, ExpectScanAll(), ExpectUnwind(), ExpectProduce());
 }
 
+TEST(TestLogicalPlanner, ReturnDistinctOrderBySkipLimit) {
+  // Test RETURN DISTINCT 1 ORDER BY 1 SKIP 1 LIMIT 1
+  AstTreeStorage storage;
+  auto query = QUERY(RETURN_DISTINCT(LITERAL(1), AS("1"), ORDER_BY(LITERAL(1)),
+                                     SKIP(LITERAL(1)), LIMIT(LITERAL(1))));
+  CheckPlan(*query, ExpectProduce(), ExpectDistinct(), ExpectOrderBy(),
+            ExpectSkip(), ExpectLimit());
+}
+
+TEST(TestLogicalPlanner, CreateWithDistinctSumWhereReturn) {
+  // Test CREATE (n) WITH DISTINCT SUM(n.prop) AS s WHERE s < 42 RETURN s
+  Dbms dbms;
+  auto dba = dbms.active();
+  auto prop = dba->property("prop");
+  AstTreeStorage storage;
+  auto node_n = NODE("n");
+  auto sum = SUM(PROPERTY_LOOKUP("n", prop));
+  auto query =
+      QUERY(CREATE(PATTERN(node_n)), WITH_DISTINCT(sum, AS("s")),
+            WHERE(LESS(IDENT("s"), LITERAL(42))), RETURN(IDENT("s"), AS("s")));
+  auto symbol_table = MakeSymbolTable(*query);
+  auto acc = ExpectAccumulate({symbol_table.at(*node_n->identifier_)});
+  auto aggr = ExpectAggregate({sum}, {});
+  auto plan = MakeLogicalPlan(*query, symbol_table);
+  CheckPlan(*plan, symbol_table, ExpectCreateNode(), acc, aggr, ExpectProduce(),
+            ExpectFilter(), ExpectDistinct(), ExpectProduce());
+}
+
 }  // namespace