Skip to content

Commit e824c30

Browse files
bpinteabpintea
committed
Review comments
1 parent 4f004e6 commit e824c30

File tree

3 files changed

+176
-49
lines changed

3 files changed

+176
-49
lines changed

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/inlinestats.csv-spec‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,3 +1386,94 @@ total_duration:long | day:date | days:date
13861386
17016205 |2023-10-23T13:00:00.000Z|[2023-10-23T12:00:00.000Z, 2023-10-23T13:00:00.000Z]
13871387
;
13881388

1389+
1390+
evalBeforeInlinestatsAndKeepAfter1
1391+
required_capability: inlinestats_v8
1392+
1393+
FROM employees
1394+
| WHERE still_hired == false
1395+
| EVAL sal = salary/1000
1396+
| INLINESTATS totalK = SUM(sal), count=COUNT(*) BY gender
1397+
| KEEP emp_no, still_hired, totalK, count
1398+
| SORT emp_no
1399+
| LIMIT 5
1400+
;
1401+
1402+
emp_no:integer |still_hired:boolean|totalK:long|count:long
1403+
10003 |false |1567 |32
1404+
10006 |false |810 |16
1405+
10009 |false |810 |16
1406+
10010 |false |378 |7
1407+
10012 |false |378 |7
1408+
;
1409+
1410+
evalBeforeInlinestatsAndKeepAfter2
1411+
required_capability: inlinestats_v8
1412+
1413+
FROM employees
1414+
| EVAL salaryK = salary/1000
1415+
| INLINESTATS total = SUM(salaryK), count=COUNT(*) BY gender
1416+
| KEEP emp_no, still_hired, total, count
1417+
| WHERE still_hired == false
1418+
| SORT emp_no
1419+
| LIMIT 5
1420+
;
1421+
1422+
emp_no:integer |still_hired:boolean|total:long|count:long
1423+
10003 |false |2644 |57
1424+
10006 |false |1648 |33
1425+
10009 |false |1648 |33
1426+
10010 |false |482 |10
1427+
10012 |false |482 |10
1428+
;
1429+
1430+
evalBeforeInlinestatsAndKeepAfter3
1431+
required_capability: inlinestats_v8
1432+
1433+
FROM employees
1434+
| EVAL salaryK = salary/1000
1435+
| INLINESTATS total = SUM(salaryK) BY gender
1436+
| KEEP emp_no, still_hired, total
1437+
| SORT emp_no
1438+
| LIMIT 5
1439+
;
1440+
1441+
emp_no:integer |still_hired:boolean|total:long
1442+
10001 |true |2644
1443+
10002 |true |1648
1444+
10003 |false |2644
1445+
10004 |true |2644
1446+
10005 |true |2644
1447+
;
1448+
1449+
evalBeforeInlinestatsAndKeepAfter4
1450+
required_capability: inlinestats_v8
1451+
1452+
FROM employees
1453+
| EVAL salaryK = salary/1000
1454+
| INLINESTATS count = COUNT(*) BY salaryK
1455+
| KEEP emp_no, still_hired, count
1456+
| SORT emp_no
1457+
| LIMIT 5
1458+
;
1459+
1460+
emp_no:integer |still_hired:boolean|count:long
1461+
10001 |true |1
1462+
10002 |true |3
1463+
10003 |false |2
1464+
10004 |true |2
1465+
10005 |true |1
1466+
;
1467+
1468+
evalBeforeInlinestatsAndKeepAfter5
1469+
required_capability: inlinestats_v8
1470+
1471+
ROW salary = 12300, emp_no = 5, gender = "F"
1472+
| EVAL salaryK = salary/1000
1473+
| INLINESTATS sum = SUM(salaryK) BY gender
1474+
| KEEP emp_no
1475+
;
1476+
1477+
emp_no:integer
1478+
5
1479+
;

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java‎

Lines changed: 53 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
5050
Holder<Boolean> forkPresent = new Holder<>(false);
5151

5252
// while going top-to-bottom (upstream)
53-
var pl = plan.transformDown(p -> {
53+
return plan.transformDown(p -> {
5454
// Note: It is NOT required to do anything special for binary plans like JOINs, except INLINESTATS. It is perfectly fine that
5555
// transformDown descends first into the left side, adding all kinds of attributes to the `used` set, and then descends into
5656
// the right side - even though the `used` set will contain stuff only used in the left hand side. That's because any attribute
@@ -78,11 +78,11 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
7878
do {
7979
recheck.set(false);
8080
p = switch (p) {
81-
case Aggregate agg -> pruneColumns(agg, used, inlineJoin);
82-
case InlineJoin inj -> pruneColumns(inj, used, recheck);
83-
case Eval eval -> pruneColumns(eval, used, recheck);
84-
case Project project -> inlineJoin ? pruneColumns(project, used) : p;
85-
case EsRelation esr -> pruneColumns(esr, used);
81+
case Aggregate agg -> pruneColumnsInAggregate(agg, used, inlineJoin);
82+
case InlineJoin inj -> pruneColumnsInInlineJoin(inj, used, recheck);
83+
case Eval eval -> pruneColumnsInEval(eval, used, recheck);
84+
case Project project -> inlineJoin ? pruneColumnsInProject(project, used) : p;
85+
case EsRelation esr -> pruneColumnsInEsRelation(esr, used);
8686
default -> p;
8787
};
8888
} while (recheck.get());
@@ -92,61 +92,63 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
9292
// preserve the state before going to the next node
9393
return p;
9494
});
95-
96-
return pl;
9795
}
9896

99-
private static LogicalPlan pruneColumns(Aggregate aggregate, AttributeSet.Builder used, boolean inlineJoin) {
97+
private static LogicalPlan pruneColumnsInAggregate(Aggregate aggregate, AttributeSet.Builder used, boolean inlineJoin) {
10098
LogicalPlan p = aggregate;
10199

102100
var remaining = pruneUnusedAndAddReferences(aggregate.aggregates(), used);
103101

104-
if (remaining != null) {
105-
if (remaining.isEmpty()) {
106-
if (inlineJoin) {
102+
if (remaining == null) {
103+
return p;
104+
}
105+
106+
if (remaining.isEmpty()) {
107+
if (inlineJoin) {
108+
p = emptyLocalRelation(aggregate);
109+
} else if (aggregate.groupings().isEmpty()) {
110+
// We still need to have a plan that produces 1 row per group.
111+
p = new LocalRelation(
112+
aggregate.source(),
113+
List.of(Expressions.attribute(aggregate.aggregates().getFirst())),
114+
LocalSupplier.of(new Block[] { BlockUtils.constantBlock(PlannerUtils.NON_BREAKING_BLOCK_FACTORY, null, 1) })
115+
);
116+
} else {
117+
// Aggs cannot produce pages with 0 columns, so retain one grouping.
118+
Attribute attribute = Expressions.attribute(aggregate.groupings().getFirst());
119+
NamedExpression firstAggregate = aggregate.aggregates().getFirst();
120+
remaining = List.of(new Alias(firstAggregate.source(), firstAggregate.name(), attribute, firstAggregate.id()));
121+
p = aggregate.with(aggregate.groupings(), remaining);
122+
}
123+
} else {
124+
// not expecting high groups cardinality, nested loops in lists should be fine, no need for a HashSet
125+
if (inlineJoin && aggregate.groupings().containsAll(remaining)) {
126+
// It's an INLINEJOIN and all remaining attributes are groupings, which are already part of the IJ output (from the
127+
// left-hand side).
128+
// TODO: INLINESTATS: revisit condition when adding support for INLINESTATS filters
129+
if (aggregate.child() instanceof StubRelation stub) {
130+
var message = "Aggregate groups references ["
131+
+ remaining
132+
+ "] not in child's (StubRelation) output: ["
133+
+ stub.outputSet()
134+
+ "]";
135+
assert stub.outputSet().containsAll(Expressions.asAttributes(remaining)) : message;
136+
107137
p = emptyLocalRelation(aggregate);
108-
} else if (aggregate.groupings().isEmpty()) {
109-
// We still need to have a plan that produces 1 row per group.
110-
p = new LocalRelation(
111-
aggregate.source(),
112-
List.of(Expressions.attribute(aggregate.aggregates().getFirst())),
113-
LocalSupplier.of(new Block[] { BlockUtils.constantBlock(PlannerUtils.NON_BREAKING_BLOCK_FACTORY, null, 1) })
114-
);
115138
} else {
116-
// Aggs cannot produce pages with 0 columns, so retain one grouping.
117-
Attribute attribute = Expressions.attribute(aggregate.groupings().getFirst());
118-
NamedExpression firstAggregate = aggregate.aggregates().getFirst();
119-
remaining = List.of(new Alias(firstAggregate.source(), firstAggregate.name(), attribute, firstAggregate.id()));
120-
p = aggregate.with(aggregate.groupings(), remaining);
121-
}
122-
} else {
123-
if (inlineJoin && aggregate.groupings().containsAll(remaining)) { // not expecting high groups cardinality
124-
// It's an INLINEJOIN and all remaining attributes are groupings, which are already part of the IJ output (from the
125-
// left-hand side).
126-
if (aggregate.child() instanceof StubRelation stub) {
127-
var message = "Aggregate groups references ["
128-
+ remaining
129-
+ "] not in child's (StubRelation) output: ["
130-
+ stub.outputSet()
131-
+ "]";
132-
assert stub.outputSet().containsAll(Expressions.asAttributes(remaining)) : message;
133-
134-
p = emptyLocalRelation(aggregate);
135-
} else {
136-
// There are no aggregates to compute, just output the groupings; these are already in the IJ output, so only
137-
// restrict the output to what remained.
138-
p = new Project(aggregate.source(), aggregate.child(), remaining);
139-
}
140-
} else { // not an INLINEJOIN or there are actually aggregates to compute
141-
p = aggregate.with(aggregate.groupings(), remaining);
139+
// There are no aggregates to compute, just output the groupings; these are already in the IJ output, so only
140+
// restrict the output to what remained.
141+
p = new Project(aggregate.source(), aggregate.child(), remaining);
142142
}
143+
} else { // not an INLINEJOIN or there are actually aggregates to compute
144+
p = aggregate.with(aggregate.groupings(), remaining);
143145
}
144146
}
145147

146148
return p;
147149
}
148150

149-
private static LogicalPlan pruneColumns(InlineJoin ij, AttributeSet.Builder used, Holder<Boolean> recheck) {
151+
private static LogicalPlan pruneColumnsInInlineJoin(InlineJoin ij, AttributeSet.Builder used, Holder<Boolean> recheck) {
150152
LogicalPlan p = ij;
151153

152154
used.addAll(ij.references());
@@ -162,7 +164,7 @@ private static LogicalPlan pruneColumns(InlineJoin ij, AttributeSet.Builder used
162164
return p;
163165
}
164166

165-
private static LogicalPlan pruneColumns(Eval eval, AttributeSet.Builder used, Holder<Boolean> recheck) {
167+
private static LogicalPlan pruneColumnsInEval(Eval eval, AttributeSet.Builder used, Holder<Boolean> recheck) {
166168
LogicalPlan p = eval;
167169

168170
var remaining = pruneUnusedAndAddReferences(eval.fields(), used);
@@ -179,7 +181,7 @@ private static LogicalPlan pruneColumns(Eval eval, AttributeSet.Builder used, Ho
179181
return p;
180182
}
181183

182-
private static LogicalPlan pruneColumns(Project project, AttributeSet.Builder used) {
184+
private static LogicalPlan pruneColumnsInProject(Project project, AttributeSet.Builder used) {
183185
LogicalPlan p = project;
184186

185187
var remaining = pruneUnusedAndAddReferences(project.projections(), used);
@@ -188,13 +190,15 @@ private static LogicalPlan pruneColumns(Project project, AttributeSet.Builder us
188190
? emptyLocalRelation(project)
189191
: new Project(project.source(), project.child(), remaining);
190192
} else if (project.output().stream().allMatch(FieldAttribute.class::isInstance)) {
193+
// Use empty relation as a marker for a subsequent pass, in case the project is only outputting field attributes (which are
194+
// already part of the INLINEJOIN left-hand side output).
191195
p = emptyLocalRelation(project);
192196
}
193197

194198
return p;
195199
}
196200

197-
private static LogicalPlan pruneColumns(EsRelation esr, AttributeSet.Builder used) {
201+
private static LogicalPlan pruneColumnsInEsRelation(EsRelation esr, AttributeSet.Builder used) {
198202
LogicalPlan p = esr;
199203

200204
if (esr.indexMode() == IndexMode.LOOKUP) {

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5981,6 +5981,38 @@ public void testInlinestatsWithAvg() {
59815981
var stub = as(agg.child(), StubRelation.class);
59825982
}
59835983

5984+
/*
5985+
* EsqlProject[[emp_no{r}#5]]
5986+
* \_Limit[1000[INTEGER],false]
5987+
* \_LocalRelation[[salary{r}#3, emp_no{r}#5, gender{r}#7],
5988+
* org.elasticsearch.xpack.esql.plan.logical.local.CopyingLocalSupplier@9d5b596d]
5989+
*/
5990+
public void testInlinestatsWithRow() {
5991+
var query = """
5992+
ROW salary = 12300, emp_no = 5, gender = "F"
5993+
| EVAL salaryK = salary/1000
5994+
| INLINESTATS sum = SUM(salaryK) BY gender
5995+
| KEEP emp_no
5996+
""";
5997+
if (releaseBuildForInlinestats(query)) {
5998+
return;
5999+
}
6000+
var plan = optimizedPlan(query);
6001+
6002+
var esqlProject = as(plan, EsqlProject.class);
6003+
assertThat(Expressions.names(esqlProject.projections()), is(List.of("emp_no")));
6004+
var limit = asLimit(esqlProject.child(), 1000, false);
6005+
var localRelation = as(limit.child(), LocalRelation.class);
6006+
assertThat(
6007+
localRelation.output(),
6008+
contains(
6009+
new ReferenceAttribute(EMPTY, "salary", INTEGER),
6010+
new ReferenceAttribute(EMPTY, "emp_no", INTEGER),
6011+
new ReferenceAttribute(EMPTY, "gender", KEYWORD)
6012+
)
6013+
);
6014+
}
6015+
59846016
/**
59856017
* Expects
59866018
*

0 commit comments

Comments
 (0)