Skip to content

Commit 5b32d11

Browse files
codebirdmatriv
codebird
authored and
matriv
committed
SQL: Fix issues with GROUP BY queries (#41964)
Translate to an agg query even if only literals are selected, so that the correct number of rows is returned (number of buckets). Fix issue with key only in GROUP BY (not in select) and WHERE clause: Resolve aggregates and groupings based on the child plan which holds the info info for all the fields of the underlying table. Fixes: #41951 Fixes: #41413 (cherry picked from commit 45b8580)
1 parent b729157 commit 5b32d11

File tree

5 files changed

+97
-9
lines changed

5 files changed

+97
-9
lines changed

‎x-pack/plugin/sql/qa/src/main/resources/agg.sql-spec

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,26 @@ countDistinctAlias
144144
SELECT COUNT(DISTINCT hire_date) AS count FROM test_emp;
145145
countDistinctAndCountSimpleWithAlias
146146
SELECT COUNT(*) cnt, COUNT(DISTINCT first_name) as names, gender FROM test_emp GROUP BY gender ORDER BY gender;
147+
aliasedCountWithFunctionFilterAndGroupBy
148+
SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;
149+
countWithFunctionFilterAndGroupBy
150+
SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;
151+
aliasedCountWithMultiFunctionFilterAndGroupBy
152+
SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender ORDER BY gender;
153+
countWithMultiFunctionFilterAndGroupBy
154+
SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender ORDER BY gender;
155+
aliasedCountWithFunctionFilterAndMultiGroupBy
156+
SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender, salary ORDER BY gender;
157+
countWithFunctionFilterAndMultiGroupBy
158+
SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 GROUP BY gender, salary ORDER BY gender;
159+
aliasedCountWithMultiFunctionFilterAndMultiGroupBy
160+
SELECT COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender, salary ORDER BY gender;
161+
countWithMultiFunctionFilterAndMultiGroupBy
162+
SELECT COUNT(*) FROM test_emp WHERE ABS(salary) > 0 AND YEAR(birth_date) > 1980 GROUP BY gender, salary ORDER BY gender;
163+
aliasedCountLiteralColumnWithFunctionFilterAndMultiGroupBy
164+
SELECT 1, gender as g, COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY g, salary ORDER BY gender;
165+
aliasedCountLiteralColumnWithFunctionFilterAndMultiGroupByWithFunction
166+
SELECT 1, gender as g, COUNT(*) as c FROM test_emp WHERE ABS(salary) > 0 GROUP BY g, YEAR(birth_date) ORDER BY gender, YEAR(birth_date);
147167

148168
aggCountAliasAndWhereClauseMultiGroupBy
149169
SELECT gender g, languages l, COUNT(*) c FROM "test_emp" WHERE emp_no < 10020 GROUP BY gender, languages ORDER BY gender, languages;
@@ -563,6 +583,22 @@ SELECT MIN(salary) min, MAX(salary) max, gender g, languages l, COUNT(*) c FROM
563583
// group by with literal
564584
implicitGroupByWithLiteral
565585
SELECT 10, MAX("salary") FROM test_emp;
586+
literalWithGroupBy
587+
SELECT 1 FROM test_emp GROUP BY gender;
588+
literalsWithGroupBy
589+
SELECT 1, 2 FROM test_emp GROUP BY gender;
590+
aliasedLiteralWithGroupBy
591+
SELECT 1 AS s FROM test_emp GROUP BY gender;
592+
aliasedLiteralsWithGroupBy
593+
SELECT 1 AS s, 2 FROM test_emp GROUP BY gender;
594+
literalsWithMultipleGroupBy
595+
SELECT 1, 2 FROM test_emp GROUP BY gender, salary;
596+
divisionLiteralsAdditionWithMultipleGroupBy
597+
SELECT 144 / 12 AS division, 1, 2 AS x, 1 + 2 AS addition FROM test_emp GROUP BY gender, salary;
598+
aliasedLiteralsWithMultipleGroupBy
599+
SELECT 1 as s, 2 FROM test_emp GROUP BY gender, salary;
600+
aliasedLiteralsWithMultipleGroupByWithFunction
601+
SELECT 1 as s, 2 FROM test_emp GROUP BY gender, YEAR(birth_date);
566602
implicitGroupByWithLiterals
567603
SELECT 10, 'foo', MAX("salary"), 20, 'bar' FROM test_emp;
568604
groupByWithLiteral

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,7 @@ else if (plan instanceof Aggregate) {
341341
List<Expression> groupings = a.groupings();
342342
List<Expression> newGroupings = new ArrayList<>();
343343
AttributeMap<Expression> resolved = Expressions.aliases(a.aggregates());
344+
344345
boolean changed = false;
345346
for (Expression grouping : groupings) {
346347
if (grouping instanceof UnresolvedAttribute) {
@@ -618,7 +619,7 @@ protected LogicalPlan rule(LogicalPlan plan) {
618619
for (Order or : o.order()) {
619620
maybeResolved.add(or.resolved() ? or : tryResolveExpression(or, child));
620621
}
621-
622+
622623
Stream<Order> referencesStream = maybeResolved.stream()
623624
.filter(Expression::resolved);
624625

@@ -629,7 +630,7 @@ protected LogicalPlan rule(LogicalPlan plan) {
629630
// a + 1 in SELECT is actually Alias("a + 1", a + 1) and translates to ReferenceAttribute
630631
// in the output. However it won't match the unnamed a + 1 despite being the same expression
631632
// so explicitly compare the source
632-
633+
633634
// if there's a match, remove the item from the reference stream
634635
if (Expressions.hasReferenceAttribute(child.outputSet())) {
635636
final Map<Attribute, Expression> collectRefs = new LinkedHashMap<>();
@@ -720,6 +721,18 @@ protected LogicalPlan rule(LogicalPlan plan) {
720721
}
721722
}
722723

724+
// Try to resolve aggregates and groupings based on the child plan
725+
if (plan instanceof Aggregate) {
726+
Aggregate a = (Aggregate) plan;
727+
LogicalPlan child = a.child();
728+
List<Expression> newGroupings = new ArrayList<>(a.groupings().size());
729+
a.groupings().forEach(e -> newGroupings.add(tryResolveExpression(e, child)));
730+
List<NamedExpression> newAggregates = new ArrayList<>(a.aggregates().size());
731+
a.aggregates().forEach(e -> newAggregates.add(tryResolveExpression(e, child)));
732+
if (newAggregates.equals(a.aggregates()) == false || newGroupings.equals(a.groupings()) == false) {
733+
return new Aggregate(a.source(), child, newGroupings, newAggregates);
734+
}
735+
}
723736
return plan;
724737
}
725738

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1130,7 +1130,8 @@ protected LogicalPlan rule(LogicalPlan plan) {
11301130

11311131
plan.forEachDown(a -> {
11321132
List<Object> values = extractConstants(a.aggregates());
1133-
if (values.size() == a.aggregates().size() && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
1133+
if (values.size() == a.aggregates().size() && a.groupings().isEmpty()
1134+
&& isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
11341135
optimizedPlan.set(new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray())));
11351136
}
11361137
}, Aggregate.class);

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -393,20 +393,21 @@ protected PhysicalPlan rule(AggregateExec a) {
393393
}
394394
return a;
395395
}
396-
396+
397397
static EsQueryExec fold(AggregateExec a, EsQueryExec exec) {
398-
398+
399399
QueryContainer queryC = exec.queryContainer();
400-
400+
401401
// track aliases defined in the SELECT and used inside GROUP BY
402402
// SELECT x AS a ... GROUP BY a
403403
Map<Attribute, Expression> aliasMap = new LinkedHashMap<>();
404+
String id = null;
404405
for (NamedExpression ne : a.aggregates()) {
405406
if (ne instanceof Alias) {
406407
aliasMap.put(ne.toAttribute(), ((Alias) ne).child());
407408
}
408409
}
409-
410+
410411
if (aliasMap.isEmpty() == false) {
411412
Map<Attribute, Expression> newAliases = new LinkedHashMap<>(queryC.aliases());
412413
newAliases.putAll(aliasMap);
@@ -451,7 +452,7 @@ static EsQueryExec fold(AggregateExec a, EsQueryExec exec) {
451452
target = ((Alias) ne).child();
452453
}
453454

454-
String id = Expressions.id(target);
455+
id = Expressions.id(target);
455456

456457
// literal
457458
if (target.foldable()) {
@@ -587,7 +588,14 @@ else if (target.foldable()) {
587588
}
588589
}
589590
}
590-
591+
// If we're only selecting literals, we have to still execute the aggregation to create
592+
// the correct grouping buckets, in order to return the appropriate number of rows
593+
if (a.aggregates().stream().allMatch(e -> e.anyMatch(Expression::foldable))) {
594+
for (Expression grouping : a.groupings()) {
595+
GroupByKey matchingGroup = groupingContext.groupFor(grouping);
596+
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, false), id);
597+
}
598+
}
591599
return new EsQueryExec(exec.source(), exec.index(), a.output(), queryC);
592600
}
593601

‎x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryTranslatorTests.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
import org.elasticsearch.xpack.ql.expression.Expression;
1717
import org.elasticsearch.xpack.ql.expression.FieldAttribute;
1818
import org.elasticsearch.xpack.ql.expression.Literal;
19+
import org.elasticsearch.xpack.ql.expression.function.aggregate.Count;
1920
import org.elasticsearch.xpack.ql.expression.gen.script.ScriptTemplate;
21+
import org.elasticsearch.xpack.ql.expression.predicate.operator.comparison.GreaterThan;
2022
import org.elasticsearch.xpack.ql.index.EsIndex;
2123
import org.elasticsearch.xpack.ql.index.IndexResolution;
2224
import org.elasticsearch.xpack.ql.plan.logical.Aggregate;
@@ -130,6 +132,34 @@ public void testTermEqualityAnalyzer() {
130132
assertEquals("value", tq.value());
131133
}
132134

135+
public void testAliasAndGroupByResolution(){
136+
LogicalPlan p = plan("SELECT COUNT(*) AS c FROM test WHERE ABS(int) > 0 GROUP BY int");
137+
assertTrue(p instanceof Aggregate);
138+
Aggregate a = (Aggregate) p;
139+
LogicalPlan pc = ((Aggregate) p).child();
140+
assertTrue(pc instanceof Filter);
141+
Expression condition = ((Filter) pc).condition();
142+
assertEquals("GREATERTHAN", ((GreaterThan) condition).functionName());
143+
List<Expression> groupings = a.groupings();
144+
assertTrue(groupings.get(0).resolved());
145+
assertEquals("c", a.aggregates().get(0).name());
146+
assertEquals("COUNT", ((Count) ((Alias) a.aggregates().get(0)).child()).functionName());
147+
}
148+
public void testLiteralWithGroupBy(){
149+
LogicalPlan p = plan("SELECT 1 as t, 2 FROM test GROUP BY int");
150+
assertTrue(p instanceof Aggregate);
151+
Aggregate a = (Aggregate) p;
152+
List<Expression> groupings = a.groupings();
153+
assertEquals(1, groupings.size());
154+
assertTrue(groupings.get(0).resolved());
155+
assertTrue(groupings.get(0) instanceof FieldAttribute);
156+
assertEquals(2, a.aggregates().size());
157+
assertEquals("t", a.aggregates().get(0).name());
158+
assertTrue(((Alias) a.aggregates().get(0)).child() instanceof Literal);
159+
assertEquals("1", ((Alias) a.aggregates().get(0)).child().toString());
160+
assertEquals("2", ((Alias) a.aggregates().get(1)).child().toString());
161+
}
162+
133163
public void testTermEqualityNotAnalyzed() {
134164
LogicalPlan p = plan("SELECT some.string FROM test WHERE int = 5");
135165
assertTrue(p instanceof Project);

0 commit comments

Comments
 (0)