Skip to content

Commit 81d6700

Browse files
nik9000gmjehovich
authored andcommitted
ESQL: More code gen standardization (elastic#134935)
Makes a few changes to code gen that should be a noop but makes it easier to make further code gen changes. 1. Extract the argument infrastructure from `EvaluatorImplementer` so we can share it more easily with aggs. Makes a whole package for it! 2. Read values into a local before calling the in `ExpressionEvaluator`. This mimics how aggregations work more closely and will make it easier to move aggregations to the same argument infrastructure. It's the bulk of this change but should be the same at runtime. 3. Moves parsing of aggregation parameters to the infrastructure we built for `ExpressionEvaluator`. In the future move aggs over to this infrastructure entirely but that's too much for this change. Relates to elastic#108385
1 parent c6318da commit 81d6700

File tree

245 files changed

+2090
-1253
lines changed

Some content is hidden

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

245 files changed

+2090
-1253
lines changed

x-pack/plugin/esql/compute/gen/src/main/java/module-info.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
requires java.compiler;
1616

1717
exports org.elasticsearch.compute.gen;
18+
exports org.elasticsearch.compute.gen.argument;
1819

1920
provides javax.annotation.processing.Processor with AggregatorProcessor, ConsumeProcessor, EvaluatorProcessor;
2021
}

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

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@
1818
import org.elasticsearch.compute.ann.Aggregator;
1919
import org.elasticsearch.compute.ann.IntermediateState;
2020
import org.elasticsearch.compute.gen.Methods.TypeMatcher;
21+
import org.elasticsearch.compute.gen.argument.Argument;
22+
import org.elasticsearch.compute.gen.argument.ArrayArgument;
23+
import org.elasticsearch.compute.gen.argument.StandardArgument;
2124

2225
import java.util.ArrayList;
2326
import java.util.Arrays;
@@ -30,8 +33,6 @@
3033
import javax.lang.model.element.ExecutableElement;
3134
import javax.lang.model.element.Modifier;
3235
import javax.lang.model.element.TypeElement;
33-
import javax.lang.model.element.VariableElement;
34-
import javax.lang.model.type.TypeKind;
3536
import javax.lang.model.type.TypeMirror;
3637
import javax.lang.model.util.Elements;
3738

@@ -87,6 +88,7 @@ public class AggregatorImplementer {
8788

8889
public AggregatorImplementer(
8990
Elements elements,
91+
javax.lang.model.util.Types types,
9092
TypeElement declarationType,
9193
IntermediateState[] interStateAnno,
9294
List<TypeMirror> warnExceptions
@@ -108,7 +110,14 @@ public AggregatorImplementer(
108110
requireName("combine"),
109111
requireArgsStartsWith(requireType(aggState.declaredType()), requireAnyType("<aggregation input column type>"))
110112
);
111-
this.aggParams = combine.getParameters().stream().skip(1).map(AggregationParameter::create).toList();
113+
this.aggParams = combine.getParameters().stream().skip(1).map(v -> {
114+
Argument a = Argument.fromParameter(types, v);
115+
return switch (a) {
116+
case StandardArgument sa -> new AggregationParameter(sa.name(), sa.type(), false);
117+
case ArrayArgument aa -> new AggregationParameter(aa.name(), aa.componentType(), true);
118+
default -> throw new IllegalArgumentException("unsupported argument [" + a + "]");
119+
};
120+
}).toList();
112121

113122
this.createParameters = init.getParameters()
114123
.stream()
@@ -768,14 +777,6 @@ private static String primitiveStateStoreClassname(TypeName declaredType, boolea
768777
}
769778

770779
public record AggregationParameter(String name, TypeName type, boolean isArray) {
771-
public static AggregationParameter create(VariableElement v) {
772-
return new AggregationParameter(
773-
v.getSimpleName().toString(),
774-
TypeName.get(v.asType()),
775-
Objects.equals(v.asType().getKind(), TypeKind.ARRAY)
776-
);
777-
}
778-
779780
public String blockName() {
780781
return name + "Block";
781782
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,13 @@ public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnv
8787
);
8888
if (aggClass.getAnnotation(Aggregator.class) != null) {
8989
IntermediateState[] intermediateState = aggClass.getAnnotation(Aggregator.class).value();
90-
implementer = new AggregatorImplementer(env.getElementUtils(), aggClass, intermediateState, warnExceptionsTypes);
90+
implementer = new AggregatorImplementer(
91+
env.getElementUtils(),
92+
env.getTypeUtils(),
93+
aggClass,
94+
intermediateState,
95+
warnExceptionsTypes
96+
);
9197
write(aggClass, "aggregator", implementer.sourceFile(), env);
9298
}
9399
GroupingAggregatorImplementer groupingAggregatorImplementer = null;
@@ -98,6 +104,7 @@ public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnv
98104
}
99105
groupingAggregatorImplementer = new GroupingAggregatorImplementer(
100106
env.getElementUtils(),
107+
env.getTypeUtils(),
101108
aggClass,
102109
intermediateState,
103110
warnExceptionsTypes

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

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@
1313
import com.squareup.javapoet.TypeName;
1414
import com.squareup.javapoet.TypeSpec;
1515

16+
import org.elasticsearch.compute.gen.argument.Argument;
17+
import org.elasticsearch.compute.gen.argument.FixedArgument;
18+
import org.elasticsearch.compute.gen.argument.StandardArgument;
19+
1620
import java.util.ArrayList;
1721
import java.util.List;
1822
import java.util.stream.Collectors;
@@ -61,14 +65,14 @@ public ConvertEvaluatorImplementer(
6165
this.processFunction = new EvaluatorImplementer.ProcessFunction(types, processFunction, warnExceptions);
6266
this.canProcessOrdinals = warnExceptions.isEmpty()
6367
&& this.processFunction.returnType().equals(BYTES_REF)
64-
&& this.processFunction.args.getFirst() instanceof EvaluatorImplementer.StandardProcessFunctionArg s
68+
&& this.processFunction.args.getFirst() instanceof StandardArgument s
6569
&& s.type().equals(BYTES_REF);
6670

67-
if (this.processFunction.args.getFirst() instanceof EvaluatorImplementer.StandardProcessFunctionArg == false) {
71+
if (this.processFunction.args.getFirst() instanceof StandardArgument == false) {
6872
throw new IllegalArgumentException("first argument must be the field to process");
6973
}
7074
for (int a = 1; a < this.processFunction.args.size(); a++) {
71-
if (this.processFunction.args.get(a) instanceof EvaluatorImplementer.FixedProcessFunctionArg == false) {
75+
if (this.processFunction.args.get(a) instanceof FixedArgument == false) {
7276
throw new IllegalArgumentException("fixed function args supported after the first");
7377
// TODO support more function types when we need them
7478
}
@@ -101,7 +105,7 @@ private TypeSpec type() {
101105
builder.superclass(ABSTRACT_CONVERT_FUNCTION_EVALUATOR);
102106
builder.addField(baseRamBytesUsed(implementation));
103107

104-
for (EvaluatorImplementer.ProcessFunctionArg a : processFunction.args) {
108+
for (Argument a : processFunction.args) {
105109
a.declareField(builder);
106110
}
107111
builder.addMethod(ctor());
@@ -124,7 +128,7 @@ private MethodSpec ctor() {
124128
MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
125129
builder.addParameter(SOURCE, "source");
126130
builder.addStatement("super(driverContext, source)");
127-
for (EvaluatorImplementer.ProcessFunctionArg a : processFunction.args) {
131+
for (Argument a : processFunction.args) {
128132
a.implementCtor(builder);
129133
}
130134
builder.addParameter(DRIVER_CONTEXT, "driverContext");
@@ -134,7 +138,7 @@ private MethodSpec ctor() {
134138
private MethodSpec next() {
135139
MethodSpec.Builder builder = MethodSpec.methodBuilder("next").addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
136140
builder.returns(EXPRESSION_EVALUATOR);
137-
builder.addStatement("return $N", ((EvaluatorImplementer.StandardProcessFunctionArg) processFunction.args.getFirst()).name());
141+
builder.addStatement("return $N", ((StandardArgument) processFunction.args.getFirst()).name());
138142
return builder.build();
139143
}
140144

0 commit comments

Comments
 (0)