Skip to content

Commit 5a4fbba

Browse files
dnhatndnhatn
committed
Optimize ordinal inputs in Values aggregation
1 parent 76ee76a commit 5a4fbba

File tree

6 files changed

+365
-15
lines changed

6 files changed

+365
-15
lines changed

‎benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,13 @@
2121
import org.elasticsearch.compute.data.Block;
2222
import org.elasticsearch.compute.data.BlockFactory;
2323
import org.elasticsearch.compute.data.BytesRefBlock;
24+
import org.elasticsearch.compute.data.BytesRefVector;
2425
import org.elasticsearch.compute.data.ElementType;
2526
import org.elasticsearch.compute.data.IntBlock;
27+
import org.elasticsearch.compute.data.IntVector;
2628
import org.elasticsearch.compute.data.LongBlock;
2729
import org.elasticsearch.compute.data.LongVector;
30+
import org.elasticsearch.compute.data.OrdinalBytesRefVector;
2831
import org.elasticsearch.compute.data.Page;
2932
import org.elasticsearch.compute.operator.AggregationOperator;
3033
import org.elasticsearch.compute.operator.DriverContext;
@@ -282,11 +285,17 @@ private static Block dataBlock(int groups, String dataType) {
282285
int blockLength = blockLength(groups);
283286
return switch (dataType) {
284287
case BYTES_REF -> {
285-
try (BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(blockLength)) {
288+
try (BytesRefVector.Builder dict = blockFactory.newBytesRefVectorBuilder(blockLength);
289+
IntVector.Builder ords = blockFactory.newIntVectorBuilder(blockLength)
290+
) {
291+
final int dictLength = Math.min(blockLength, KEYWORDS.length);
292+
for (int i = 0; i < dictLength; i++) {
293+
dict.appendBytesRef(KEYWORDS[i]);
294+
}
286295
for (int i = 0; i < blockLength; i++) {
287-
builder.appendBytesRef(KEYWORDS[i % KEYWORDS.length]);
296+
ords.appendInt(i % dictLength);
288297
}
289-
yield builder.build();
298+
yield new OrdinalBytesRefVector(ords.build(), dict.build()).asBlock();
290299
}
291300
}
292301
case INT -> {
@@ -332,7 +341,7 @@ private static void run(int groups, String dataType, int opCount) {
332341
operator.addInput(page.shallowCopy());
333342
}
334343
operator.finish();
335-
checkExpected(groups, dataType, operator.getOutput());
344+
// checkExpected(groups, dataType, operator.getOutput());
336345
}
337346
}
338347

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

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
import static java.util.stream.Collectors.joining;
3737
import static org.elasticsearch.compute.gen.AggregatorImplementer.capitalize;
38+
import static org.elasticsearch.compute.gen.Methods.optionalStaticMethod;
3839
import static org.elasticsearch.compute.gen.Methods.requireAnyArgs;
3940
import static org.elasticsearch.compute.gen.Methods.requireAnyType;
4041
import static org.elasticsearch.compute.gen.Methods.requireArgs;
@@ -336,10 +337,32 @@ private MethodSpec prepareProcessPage() {
336337
builder.beginControlFlow("if (valuesBlock.mayHaveNulls())");
337338
builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
338339
builder.endControlFlow();
339-
builder.addStatement("return $L", addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra)));
340+
if (shouldWrapAddInput(blockType(aggParam.type()))) {
341+
builder.addStatement(
342+
"var addInput = $L",
343+
addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra))
344+
);
345+
builder.addStatement("return $T.wrapAddInput(addInput, state, valuesBlock)", declarationType);
346+
} else {
347+
builder.addStatement(
348+
"return $L",
349+
addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesBlock$L)", extra))
350+
);
351+
}
340352
}
341353
builder.endControlFlow();
342-
builder.addStatement("return $L", addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra)));
354+
if (shouldWrapAddInput(vectorType(aggParam.type()))) {
355+
builder.addStatement(
356+
"var addInput = $L",
357+
addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra))
358+
);
359+
builder.addStatement("return $T.wrapAddInput(addInput, state, valuesVector)", declarationType);
360+
} else {
361+
builder.addStatement(
362+
"return $L",
363+
addInput(b -> b.addStatement("addRawInput(positionOffset, groupIds, valuesVector$L)", extra))
364+
);
365+
}
343366
return builder.build();
344367
}
345368

@@ -526,6 +549,15 @@ private void combineRawInputForArray(MethodSpec.Builder builder, String arrayVar
526549
warningsBlock(builder, () -> builder.addStatement("$T.combine(state, groupId, $L)", declarationType, arrayVariable));
527550
}
528551

552+
private boolean shouldWrapAddInput(TypeName valuesType) {
553+
return optionalStaticMethod(
554+
declarationType,
555+
requireType(GROUPING_AGGREGATOR_FUNCTION_ADD_INPUT),
556+
requireName("wrapAddInput"),
557+
requireArgs(requireType(GROUPING_AGGREGATOR_FUNCTION_ADD_INPUT), requireType(aggState.declaredType()), requireType(valuesType))
558+
) != null;
559+
}
560+
529561
private void warningsBlock(MethodSpec.Builder builder, Runnable block) {
530562
if (warnExceptions.isEmpty() == false) {
531563
builder.beginControlFlow("try");

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -59,20 +59,31 @@ static ExecutableElement requireStaticMethod(
5959
TypeMatcher returnTypeMatcher,
6060
NameMatcher nameMatcher,
6161
ArgumentMatcher argumentMatcher
62+
) {
63+
ExecutableElement method = optionalStaticMethod(declarationType, returnTypeMatcher, nameMatcher, argumentMatcher);
64+
if (method == null) {
65+
var message = nameMatcher.names.size() == 1 ? "Requires method: " : "Requires one of methods: ";
66+
var signatures = nameMatcher.names.stream()
67+
.map(name -> "public static " + returnTypeMatcher + " " + declarationType + "#" + name + "(" + argumentMatcher + ")")
68+
.collect(joining(" or "));
69+
throw new IllegalArgumentException(message + signatures);
70+
}
71+
return method;
72+
}
73+
74+
static ExecutableElement optionalStaticMethod(
75+
TypeElement declarationType,
76+
TypeMatcher returnTypeMatcher,
77+
NameMatcher nameMatcher,
78+
ArgumentMatcher argumentMatcher
6279
) {
6380
return typeAndSuperType(declarationType).flatMap(type -> ElementFilter.methodsIn(type.getEnclosedElements()).stream())
6481
.filter(method -> method.getModifiers().contains(Modifier.STATIC))
6582
.filter(method -> nameMatcher.test(method.getSimpleName().toString()))
6683
.filter(method -> returnTypeMatcher.test(TypeName.get(method.getReturnType())))
6784
.filter(method -> argumentMatcher.test(method.getParameters().stream().map(it -> TypeName.get(it.asType())).toList()))
6885
.findFirst()
69-
.orElseThrow(() -> {
70-
var message = nameMatcher.names.size() == 1 ? "Requires method: " : "Requires one of methods: ";
71-
var signatures = nameMatcher.names.stream()
72-
.map(name -> "public static " + returnTypeMatcher + " " + declarationType + "#" + name + "(" + argumentMatcher + ")")
73-
.collect(joining(" or "));
74-
return new IllegalArgumentException(message + signatures);
75-
});
86+
.orElse(null);
7687
}
7788

7889
static NameMatcher requireName(String... names) {

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

Lines changed: 147 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/org/elasticsearch/compute/aggregation/ValuesBytesRefGroupingAggregatorFunction.java‎

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

0 commit comments

Comments
 (0)