Fix return collection containing both aggregation and group-by

Summary:
Referring to the TCK failure on:
```
MATCH (a {name: 'Andres'})<-[:FATHER]-(child)
RETURN {foo: a.name='Andres', kids: collect(child.name)}
```

In the planner we'd only treat a list|map as a group_by if it contained
no aggregations. That's changed so that if a map contains both aggregations
and non-aggregations, then non-aggregations are treated as individual
group_by expressions.

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: buda, pullbot, teon.banek

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

src/query/plan/rule_based_planner.cpp

@@ -144,32 +144,47 @@ class ReturnBodyContext : public HierarchicalTreeVisitor {
     return true;
   }
 
+ private:
+  template <typename TLiteral, typename TIteratorToExpression>
+  void PostVisitCollectionLiteral(
+      TLiteral &literal, TIteratorToExpression iterator_to_expression) {
+    // If there is an aggregation in the list, and there are group-bys, then we
+    // need to add the group-bys manually. If there are no aggregations, the
+    // whole list will be added as a group-by.
+    std::vector<Expression *> literal_group_by;
+    bool has_aggr = false;
+    auto it = has_aggregation_.end();
+    auto elements_it = literal.elements_.begin();
+    std::advance(it, -literal.elements_.size());
+    while (it != has_aggregation_.end()) {
+      if (*it) {
+        has_aggr = true;
+      } else {
+        literal_group_by.emplace_back(iterator_to_expression(elements_it));
+      }
+      elements_it++;
+      it = has_aggregation_.erase(it);
+    }
+    has_aggregation_.emplace_back(has_aggr);
+    if (has_aggr) {
+      for (auto expression_ptr : literal_group_by)
+        group_by_.emplace_back(expression_ptr);
+    }
+  }
+
+ public:
   bool PostVisit(ListLiteral &list_literal) override {
     DCHECK(list_literal.elements_.size() <= has_aggregation_.size())
         << "Expected has_aggregation_ flags as much as there are list "
            "elements.";
-    bool has_aggr = false;
+    PostVisitCollectionLiteral(list_literal, [](auto it) { return *it; });
-    auto it = has_aggregation_.end();
-    std::advance(it, -list_literal.elements_.size());
-    while (it != has_aggregation_.end()) {
-      has_aggr = has_aggr || *it;
-      it = has_aggregation_.erase(it);
-    }
-    has_aggregation_.emplace_back(has_aggr);
     return true;
   }
 
   bool PostVisit(MapLiteral &map_literal) override {
     DCHECK(map_literal.elements_.size() <= has_aggregation_.size())
         << "Expected has_aggregation_ flags as much as there are map elements.";
-    bool has_aggr = false;
+    PostVisitCollectionLiteral(map_literal, [](auto it) { return it->second; });
-    auto it = has_aggregation_.end();
-    std::advance(it, -map_literal.elements_.size());
-    while (it != has_aggregation_.end()) {
-      has_aggr = has_aggr || *it;
-      it = has_aggregation_.erase(it);
-    }
-    has_aggregation_.emplace_back(has_aggr);
     return true;
   }
 

tests/unit/query_planner.cpp

@@ -486,7 +486,6 @@ TEST(TestLogicalPlanner, OptionalMatchNamedPatternReturn) {
             ExpectProduce());
 }
 
-
 TEST(TestLogicalPlanner, MatchWhereReturn) {
   // Test MATCH (n) WHERE n.property < 42 RETURN n
   AstTreeStorage storage;
@@ -1084,6 +1083,40 @@ TEST(TestLogicalPlanner, ListSliceAggregationReturn) {
   CheckPlan(storage, aggr, ExpectProduce());
 }
 
+TEST(TestLogicalPlanner, ListWithAggregationAndGroupBy) {
+  // Test RETURN [sum(2), 42]
+  AstTreeStorage storage;
+  auto sum = SUM(LITERAL(2));
+  auto group_by_literal = LITERAL(42);
+  QUERY(RETURN(LIST(sum, group_by_literal), AS("result")));
+  auto aggr = ExpectAggregate({sum}, {group_by_literal});
+  CheckPlan(storage, aggr, ExpectProduce());
+}
+
+TEST(TestLogicalPlanner, AggregatonWithListWithAggregationAndGroupBy) {
+  // Test RETURN sum(2), [sum(3), 42]
+  AstTreeStorage storage;
+  auto sum2 = SUM(LITERAL(2));
+  auto sum3 = SUM(LITERAL(3));
+  auto group_by_literal = LITERAL(42);
+  QUERY(RETURN(sum2, AS("sum2"), LIST(sum3, group_by_literal), AS("list")));
+  auto aggr = ExpectAggregate({sum2, sum3}, {group_by_literal});
+  CheckPlan(storage, aggr, ExpectProduce());
+}
+
+TEST(TestLogicalPlanner, MapWithAggregationAndGroupBy) {
+  // Test RETURN {lit: 42, sum: sum(2)}
+  GraphDb db;
+  AstTreeStorage storage;
+  auto sum = SUM(LITERAL(2));
+  auto group_by_literal = LITERAL(42);
+  QUERY(RETURN(MAP({PROPERTY_PAIR("sum"), sum},
+                   {PROPERTY_PAIR("lit"), group_by_literal}),
+               AS("result")));
+  auto aggr = ExpectAggregate({sum}, {group_by_literal});
+  CheckPlan(storage, aggr, ExpectProduce());
+}
+
 TEST(TestLogicalPlanner, CreateIndex) {
   // Test CREATE INDEX ON :Label(property)
   GraphDb db;
@@ -1248,9 +1281,8 @@ TEST(TestLogicalPlanner, WhereIndexedLabelPropertyRange) {
   AstTreeStorage storage;
   auto lit_42 = LITERAL(42);
   auto n_prop = PROPERTY_LOOKUP("n", property);
-  auto check_planned_range = [&label, &property, &dba](const auto &rel_expr,
+  auto check_planned_range = [&label, &property, &dba](
-                                                       auto lower_bound,
+      const auto &rel_expr, auto lower_bound, auto upper_bound) {
-                                                       auto upper_bound) {
     // Shadow the first storage, so that the query is created in this one.
     AstTreeStorage storage;
     QUERY(MATCH(PATTERN(NODE("n", label))), WHERE(rel_expr), RETURN("n"));