Skip to content

Commit 686a11a

Browse files
authored
ESQL: Ungrouped agg implementations for FIRST/LAST (#132513)
Actually a few things: 1. Add generation for ungroups FIRST/LAST 2. Change the description of FIRST/LAST to be more readable 3. Allow that description to be specified instead of generated 4. Add grouped and ungrouped tests for FIRST/LAST 5. Grouped tests weren't passing with `null` groups. Always enabled group id tracking to fix it. 6. Reworked how grouped aggs trigger group id tracking for all grouped aggs to make them a little more consistent. Relates to #108385
1 parent 216753a commit 686a11a

File tree

131 files changed

+4505
-346
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

131 files changed

+4505
-346
lines changed

x-pack/plugin/esql/compute/build.gradle

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ def addOccurrence(props, Occurrence) {
8686
newProps["Occurrence"] = Occurrence
8787
newProps["First"] = Occurrence == "First" ? "true" : ""
8888
newProps["Last"] = Occurrence == "Last" ? "true" : ""
89+
newProps["occurrence"] = Occurrence.toLowerCase(Locale.ROOT)
8990
return newProps
9091
}
9192

@@ -469,6 +470,27 @@ tasks.named('stringTemplates').configure {
469470
it.inputFile = stateInputFile
470471
it.outputFile = "org/elasticsearch/compute/aggregation/DoubleState.java"
471472
}
473+
474+
/*
475+
* Generates pairwise states. We generate the ones that we need at the moment,
476+
* but add more if you need more.
477+
*/
478+
File twoStateInputFile = file("src/main/java/org/elasticsearch/compute/aggregation/X-2State.java.st")
479+
[longProperties].forEach { v1 ->
480+
[intProperties, longProperties, floatProperties, doubleProperties].forEach { v2 ->
481+
{
482+
var properties = [:]
483+
v1.forEach { k, v -> properties["v1_" + k] = v}
484+
v2.forEach { k, v -> properties["v2_" + k] = v}
485+
template {
486+
it.properties = properties
487+
it.inputFile = twoStateInputFile
488+
it.outputFile = "org/elasticsearch/compute/aggregation/${v1.Type}${v2.Type}State.java"
489+
}
490+
}
491+
}
492+
}
493+
472494
File fallibleStateInputFile = new File("${projectDir}/src/main/java/org/elasticsearch/compute/aggregation/X-FallibleState.java.st")
473495
template {
474496
it.properties = booleanProperties

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/AggregatorFunctionSupplierImplementer.java

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,15 @@
2525
import java.util.stream.Collectors;
2626
import java.util.stream.Stream;
2727

28+
import javax.lang.model.element.ExecutableElement;
2829
import javax.lang.model.element.Modifier;
2930
import javax.lang.model.element.TypeElement;
3031
import javax.lang.model.util.Elements;
3132

33+
import static org.elasticsearch.compute.gen.Methods.optionalStaticMethod;
34+
import static org.elasticsearch.compute.gen.Methods.requireArgs;
35+
import static org.elasticsearch.compute.gen.Methods.requireName;
36+
import static org.elasticsearch.compute.gen.Methods.requireType;
3237
import static org.elasticsearch.compute.gen.Types.AGGREGATOR_FUNCTION_SUPPLIER;
3338
import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
3439
import static org.elasticsearch.compute.gen.Types.LIST_AGG_FUNC_DESC;
@@ -210,17 +215,24 @@ private MethodSpec describe() {
210215
MethodSpec.Builder builder = MethodSpec.methodBuilder("describe").returns(String.class);
211216
builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
212217

213-
String name = declarationType.getSimpleName().toString();
214-
name = name.replace("BytesRef", "Byte"); // The hack expects one word types so let's make BytesRef into Byte
215-
String[] parts = name.split("(?=\\p{Upper})");
216-
if (false == parts[parts.length - 1].equals("Aggregator") || parts.length < 3) {
217-
throw new IllegalArgumentException("Can't generate description for " + declarationType.getSimpleName());
218+
ExecutableElement describe = optionalStaticMethod(declarationType, requireType(STRING), requireName("describe"), requireArgs());
219+
if (describe == null) {
220+
String name = declarationType.getSimpleName().toString();
221+
name = name.replace("BytesRef", "Byte"); // The hack expects one word types so let's make BytesRef into Byte
222+
String[] parts = name.split("(?=\\p{Upper})");
223+
if (false == parts[parts.length - 1].equals("Aggregator") || parts.length < 3) {
224+
throw new IllegalArgumentException("Can't generate description for " + declarationType.getSimpleName());
225+
}
226+
227+
String operation = Arrays.stream(parts, 0, parts.length - 2)
228+
.map(s -> s.toLowerCase(Locale.ROOT))
229+
.collect(Collectors.joining("_"));
230+
String type = parts[parts.length - 2];
231+
232+
builder.addStatement("return $S", operation + " of " + type.toLowerCase(Locale.ROOT) + "s");
233+
} else {
234+
builder.addStatement("return $T.$L()", declarationType, "describe");
218235
}
219-
220-
String operation = Arrays.stream(parts, 0, parts.length - 2).map(s -> s.toLowerCase(Locale.ROOT)).collect(Collectors.joining("_"));
221-
String type = parts[parts.length - 2];
222-
223-
builder.addStatement("return $S", operation + " of " + type.toLowerCase(Locale.ROOT) + "s");
224236
return builder.build();
225237
}
226238
}

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/GroupingAggregatorImplementer.java

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,7 @@ private TypeSpec type() {
198198
builder.addMethod(addRawInputLoop(groupIdClass, true));
199199
builder.addMethod(addIntermediateInput(groupIdClass));
200200
}
201+
builder.addMethod(maybeEnableGroupIdTracking());
201202
builder.addMethod(selectedMayContainUnseenGroups());
202203
builder.addMethod(evaluateIntermediate());
203204
builder.addMethod(evaluateFinal());
@@ -321,9 +322,11 @@ private MethodSpec prepareProcessRawInputPage() {
321322
builder.addStatement("$T $L = $L.asVector()", vectorType(p.type()), p.vectorName(), p.blockName());
322323
builder.beginControlFlow("if ($L == null)", p.vectorName());
323324
{
324-
builder.beginControlFlow("if ($L.mayHaveNulls())", p.blockName());
325-
builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
326-
builder.endControlFlow();
325+
builder.addStatement(
326+
"maybeEnableGroupIdTracking(seenGroupIds, "
327+
+ aggParams.stream().map(AggregationParameter::blockName).collect(joining(", "))
328+
+ ")"
329+
);
327330
returnAddInput(builder, false);
328331
}
329332
builder.endControlFlow();
@@ -351,6 +354,23 @@ private void returnAddInput(MethodSpec.Builder builder, boolean valuesAreVector)
351354
}
352355
}
353356

357+
private MethodSpec maybeEnableGroupIdTracking() {
358+
MethodSpec.Builder builder = MethodSpec.methodBuilder("maybeEnableGroupIdTracking");
359+
builder.addModifiers(Modifier.PRIVATE).returns(TypeName.VOID);
360+
builder.addParameter(SEEN_GROUP_IDS, "seenGroupIds");
361+
for (AggregationParameter p : aggParams) {
362+
builder.addParameter(blockType(p.type()), p.blockName());
363+
}
364+
365+
for (AggregationParameter p : aggParams) {
366+
builder.beginControlFlow("if ($L.mayHaveNulls())", p.blockName());
367+
builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
368+
builder.endControlFlow();
369+
}
370+
371+
return builder.build();
372+
}
373+
354374
/**
355375
* Generate an {@code AddInput} implementation. That's a collection path optimized for the input data.
356376
*/

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/FirstDoubleByTimestampAggregator.java

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/FirstFloatByTimestampAggregator.java

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/FirstIntByTimestampAggregator.java

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/aggregation/FirstLongByTimestampAggregator.java

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)