From 61e28bffd032d80896852bea9713bf56af297676 Mon Sep 17 00:00:00 2001 From: florijan Date: Fri, 24 Nov 2017 11:24:26 +0100 Subject: [PATCH] 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 | 47 ++++++++++++++++++--------- tests/unit/query_planner.cpp | 40 ++++++++++++++++++++--- 2 files changed, 67 insertions(+), 20 deletions(-) diff --git a/src/query/plan/rule_based_planner.cpp b/src/query/plan/rule_based_planner.cpp index 7987ad83e..7205b2c65 100644 --- a/src/query/plan/rule_based_planner.cpp +++ b/src/query/plan/rule_based_planner.cpp @@ -144,32 +144,47 @@ class ReturnBodyContext : public HierarchicalTreeVisitor { return true; } + private: + template + 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 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; - 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); + PostVisitCollectionLiteral(list_literal, [](auto it) { return *it; }); 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; - 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); + PostVisitCollectionLiteral(map_literal, [](auto it) { return it->second; }); return true; } diff --git a/tests/unit/query_planner.cpp b/tests/unit/query_planner.cpp index 96b92e6ce..4f2453c34 100644 --- a/tests/unit/query_planner.cpp +++ b/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 lower_bound, - auto upper_bound) { + auto check_planned_range = [&label, &property, &dba]( + const auto &rel_expr, auto lower_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"));