|
@@ -21,6 +21,7 @@ import org.elasticsearch.compute.ann.Fixed.Scope;
|
|
|
import java.util.ArrayList;
|
|
|
import java.util.Arrays;
|
|
|
import java.util.List;
|
|
|
+import java.util.Objects;
|
|
|
import java.util.function.Function;
|
|
|
import java.util.stream.Collectors;
|
|
|
|
|
@@ -33,7 +34,6 @@ import javax.lang.model.type.TypeKind;
|
|
|
import javax.lang.model.type.TypeMirror;
|
|
|
import javax.lang.model.util.Elements;
|
|
|
|
|
|
-import static org.elasticsearch.compute.gen.Methods.appendMethod;
|
|
|
import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
|
|
|
import static org.elasticsearch.compute.gen.Methods.getMethod;
|
|
|
import static org.elasticsearch.compute.gen.Types.BLOCK;
|
|
@@ -71,7 +71,7 @@ public class EvaluatorImplementer {
|
|
|
List<TypeMirror> warnExceptions
|
|
|
) {
|
|
|
this.declarationType = (TypeElement) processFunction.getEnclosingElement();
|
|
|
- this.processFunction = new ProcessFunction(elements, types, processFunction, warnExceptions);
|
|
|
+ this.processFunction = new ProcessFunction(types, processFunction, warnExceptions);
|
|
|
|
|
|
this.implementation = ClassName.get(
|
|
|
elements.getPackageOf(declarationType).toString(),
|
|
@@ -99,7 +99,7 @@ public class EvaluatorImplementer {
|
|
|
builder.addType(factory());
|
|
|
|
|
|
builder.addField(SOURCE, "source", Modifier.PRIVATE, Modifier.FINAL);
|
|
|
- processFunction.args.stream().forEach(a -> a.declareField(builder));
|
|
|
+ processFunction.args.forEach(a -> a.declareField(builder));
|
|
|
builder.addField(DRIVER_CONTEXT, "driverContext", Modifier.PRIVATE, Modifier.FINAL);
|
|
|
|
|
|
builder.addField(WARNINGS, "warnings", Modifier.PRIVATE);
|
|
@@ -117,8 +117,8 @@ public class EvaluatorImplementer {
|
|
|
}
|
|
|
builder.addMethod(realEval(false));
|
|
|
}
|
|
|
- builder.addMethod(toStringMethod());
|
|
|
- builder.addMethod(close());
|
|
|
+ builder.addMethod(processFunction.toStringMethod(implementation));
|
|
|
+ builder.addMethod(processFunction.close());
|
|
|
builder.addMethod(warnings());
|
|
|
return builder.build();
|
|
|
}
|
|
@@ -238,7 +238,7 @@ public class EvaluatorImplementer {
|
|
|
String builtPattern;
|
|
|
if (processFunction.builderArg == null) {
|
|
|
builtPattern = vectorize ? "result.$L(p, " + pattern + ")" : "result.$L(" + pattern + ")";
|
|
|
- args.add(0, appendMethod(resultDataType));
|
|
|
+ args.add(0, processFunction.appendMethod());
|
|
|
} else {
|
|
|
builtPattern = pattern.toString();
|
|
|
}
|
|
@@ -290,35 +290,6 @@ public class EvaluatorImplementer {
|
|
|
builder.endControlFlow();
|
|
|
}
|
|
|
|
|
|
- private MethodSpec toStringMethod() {
|
|
|
- MethodSpec.Builder builder = MethodSpec.methodBuilder("toString").addAnnotation(Override.class);
|
|
|
- builder.addModifiers(Modifier.PUBLIC).returns(String.class);
|
|
|
-
|
|
|
- StringBuilder pattern = new StringBuilder();
|
|
|
- List<Object> args = new ArrayList<>();
|
|
|
- pattern.append("return $S");
|
|
|
- args.add(implementation.simpleName() + "[");
|
|
|
- processFunction.args.stream().forEach(a -> a.buildToStringInvocation(pattern, args, args.size() > 2 ? ", " : ""));
|
|
|
- pattern.append(" + $S");
|
|
|
- args.add("]");
|
|
|
- builder.addStatement(pattern.toString(), args.toArray());
|
|
|
- return builder.build();
|
|
|
- }
|
|
|
-
|
|
|
- private MethodSpec close() {
|
|
|
- MethodSpec.Builder builder = MethodSpec.methodBuilder("close").addAnnotation(Override.class);
|
|
|
- builder.addModifiers(Modifier.PUBLIC);
|
|
|
-
|
|
|
- List<String> invocations = processFunction.args.stream().map(ProcessFunctionArg::closeInvocation).filter(s -> s != null).toList();
|
|
|
- if (invocations.isEmpty() == false) {
|
|
|
- builder.addStatement(
|
|
|
- "$T.closeExpectNoException(" + invocations.stream().collect(Collectors.joining(", ")) + ")",
|
|
|
- Types.RELEASABLES
|
|
|
- );
|
|
|
- }
|
|
|
- return builder.build();
|
|
|
- }
|
|
|
-
|
|
|
static MethodSpec warnings() {
|
|
|
MethodSpec.Builder builder = MethodSpec.methodBuilder("warnings");
|
|
|
builder.addModifiers(Modifier.PRIVATE).returns(WARNINGS);
|
|
@@ -343,42 +314,14 @@ public class EvaluatorImplementer {
|
|
|
builder.addField(SOURCE, "source", Modifier.PRIVATE, Modifier.FINAL);
|
|
|
processFunction.args.stream().forEach(a -> a.declareFactoryField(builder));
|
|
|
|
|
|
- builder.addMethod(factoryCtor());
|
|
|
- builder.addMethod(factoryGet());
|
|
|
- builder.addMethod(toStringMethod());
|
|
|
+ builder.addMethod(processFunction.factoryCtor());
|
|
|
+ builder.addMethod(processFunction.factoryGet(implementation));
|
|
|
+ builder.addMethod(processFunction.toStringMethod(implementation));
|
|
|
|
|
|
return builder.build();
|
|
|
}
|
|
|
|
|
|
- private MethodSpec factoryCtor() {
|
|
|
- MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
|
|
|
- builder.addParameter(SOURCE, "source");
|
|
|
- builder.addStatement("this.source = source");
|
|
|
- processFunction.args.stream().forEach(a -> a.implementFactoryCtor(builder));
|
|
|
-
|
|
|
- return builder.build();
|
|
|
- }
|
|
|
-
|
|
|
- private MethodSpec factoryGet() {
|
|
|
- MethodSpec.Builder builder = MethodSpec.methodBuilder("get").addAnnotation(Override.class);
|
|
|
- builder.addModifiers(Modifier.PUBLIC);
|
|
|
- builder.addParameter(DRIVER_CONTEXT, "context");
|
|
|
- builder.returns(implementation);
|
|
|
-
|
|
|
- List<String> args = new ArrayList<>();
|
|
|
- args.add("source");
|
|
|
- for (ProcessFunctionArg arg : processFunction.args) {
|
|
|
- String invocation = arg.factoryInvocation(builder);
|
|
|
- if (invocation != null) {
|
|
|
- args.add(invocation);
|
|
|
- }
|
|
|
- }
|
|
|
- args.add("context");
|
|
|
- builder.addStatement("return new $T($L)", implementation, args.stream().collect(Collectors.joining(", ")));
|
|
|
- return builder.build();
|
|
|
- }
|
|
|
-
|
|
|
- private interface ProcessFunctionArg {
|
|
|
+ interface ProcessFunctionArg {
|
|
|
/**
|
|
|
* Type containing the actual data for a page of values for this field. Usually a
|
|
|
* Block or Vector, but for fixed fields will be the original fixed type.
|
|
@@ -470,7 +413,7 @@ public class EvaluatorImplementer {
|
|
|
String closeInvocation();
|
|
|
}
|
|
|
|
|
|
- private record StandardProcessFunctionArg(TypeName type, String name) implements ProcessFunctionArg {
|
|
|
+ record StandardProcessFunctionArg(TypeName type, String name) implements ProcessFunctionArg {
|
|
|
@Override
|
|
|
public TypeName dataType(boolean blockStyle) {
|
|
|
if (blockStyle) {
|
|
@@ -726,7 +669,7 @@ public class EvaluatorImplementer {
|
|
|
}
|
|
|
}
|
|
|
|
|
|
- private record FixedProcessFunctionArg(TypeName type, String name, boolean includeInToString, Scope scope, boolean releasable)
|
|
|
+ record FixedProcessFunctionArg(TypeName type, String name, boolean includeInToString, Scope scope, boolean releasable)
|
|
|
implements
|
|
|
ProcessFunctionArg {
|
|
|
@Override
|
|
@@ -999,20 +942,15 @@ public class EvaluatorImplementer {
|
|
|
}
|
|
|
}
|
|
|
|
|
|
- private static class ProcessFunction {
|
|
|
- private final ExecutableElement function;
|
|
|
- private final List<ProcessFunctionArg> args;
|
|
|
+ static class ProcessFunction {
|
|
|
+ final ExecutableElement function;
|
|
|
+ final List<ProcessFunctionArg> args;
|
|
|
private final BuilderProcessFunctionArg builderArg;
|
|
|
private final List<TypeMirror> warnExceptions;
|
|
|
|
|
|
private boolean hasBlockType;
|
|
|
|
|
|
- private ProcessFunction(
|
|
|
- Elements elements,
|
|
|
- javax.lang.model.util.Types types,
|
|
|
- ExecutableElement function,
|
|
|
- List<TypeMirror> warnExceptions
|
|
|
- ) {
|
|
|
+ ProcessFunction(javax.lang.model.util.Types types, ExecutableElement function, List<TypeMirror> warnExceptions) {
|
|
|
this.function = function;
|
|
|
args = new ArrayList<>();
|
|
|
BuilderProcessFunctionArg builderArg = null;
|
|
@@ -1063,12 +1001,89 @@ public class EvaluatorImplementer {
|
|
|
this.warnExceptions = warnExceptions;
|
|
|
}
|
|
|
|
|
|
- private ClassName resultDataType(boolean blockStyle) {
|
|
|
+ TypeName returnType() {
|
|
|
+ return TypeName.get(function.getReturnType());
|
|
|
+ }
|
|
|
+
|
|
|
+ ClassName resultDataType(boolean blockStyle) {
|
|
|
if (builderArg != null) {
|
|
|
return builderArg.type.enclosingClassName();
|
|
|
}
|
|
|
boolean useBlockStyle = blockStyle || warnExceptions.isEmpty() == false;
|
|
|
- return useBlockStyle ? blockType(TypeName.get(function.getReturnType())) : vectorType(TypeName.get(function.getReturnType()));
|
|
|
+ return useBlockStyle ? blockType(returnType()) : vectorType(returnType());
|
|
|
+ }
|
|
|
+
|
|
|
+ String appendMethod() {
|
|
|
+ return Methods.appendMethod(returnType());
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ public String toString() {
|
|
|
+ return "ProcessFunction{"
|
|
|
+ + "function="
|
|
|
+ + function
|
|
|
+ + ", args="
|
|
|
+ + args
|
|
|
+ + ", builderArg="
|
|
|
+ + builderArg
|
|
|
+ + ", warnExceptions="
|
|
|
+ + warnExceptions
|
|
|
+ + ", hasBlockType="
|
|
|
+ + hasBlockType
|
|
|
+ + '}';
|
|
|
+ }
|
|
|
+
|
|
|
+ MethodSpec toStringMethod(ClassName implementation) {
|
|
|
+ MethodSpec.Builder builder = MethodSpec.methodBuilder("toString").addAnnotation(Override.class);
|
|
|
+ builder.addModifiers(Modifier.PUBLIC).returns(String.class);
|
|
|
+
|
|
|
+ StringBuilder pattern = new StringBuilder();
|
|
|
+ List<Object> args = new ArrayList<>();
|
|
|
+ pattern.append("return $S");
|
|
|
+ args.add(implementation.simpleName() + "[");
|
|
|
+ this.args.forEach(a -> a.buildToStringInvocation(pattern, args, args.size() > 2 ? ", " : ""));
|
|
|
+ pattern.append(" + $S");
|
|
|
+ args.add("]");
|
|
|
+ builder.addStatement(pattern.toString(), args.toArray());
|
|
|
+ return builder.build();
|
|
|
+ }
|
|
|
+
|
|
|
+ MethodSpec factoryCtor() {
|
|
|
+ MethodSpec.Builder builder = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
|
|
|
+ builder.addParameter(SOURCE, "source");
|
|
|
+ builder.addStatement("this.source = source");
|
|
|
+ args.stream().forEach(a -> a.implementFactoryCtor(builder));
|
|
|
+ return builder.build();
|
|
|
+ }
|
|
|
+
|
|
|
+ MethodSpec factoryGet(ClassName implementation) {
|
|
|
+ MethodSpec.Builder builder = MethodSpec.methodBuilder("get").addAnnotation(Override.class);
|
|
|
+ builder.addModifiers(Modifier.PUBLIC);
|
|
|
+ builder.addParameter(DRIVER_CONTEXT, "context");
|
|
|
+ builder.returns(implementation);
|
|
|
+
|
|
|
+ List<String> args = new ArrayList<>();
|
|
|
+ args.add("source");
|
|
|
+ for (ProcessFunctionArg arg : this.args) {
|
|
|
+ String invocation = arg.factoryInvocation(builder);
|
|
|
+ if (invocation != null) {
|
|
|
+ args.add(invocation);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ args.add("context");
|
|
|
+ builder.addStatement("return new $T($L)", implementation, String.join(", ", args));
|
|
|
+ return builder.build();
|
|
|
+ }
|
|
|
+
|
|
|
+ MethodSpec close() {
|
|
|
+ MethodSpec.Builder builder = MethodSpec.methodBuilder("close").addAnnotation(Override.class);
|
|
|
+ builder.addModifiers(Modifier.PUBLIC);
|
|
|
+
|
|
|
+ List<String> invocations = args.stream().map(ProcessFunctionArg::closeInvocation).filter(Objects::nonNull).toList();
|
|
|
+ if (invocations.isEmpty() == false) {
|
|
|
+ builder.addStatement("$T.closeExpectNoException(" + String.join(", ", invocations) + ")", Types.RELEASABLES);
|
|
|
+ }
|
|
|
+ return builder.build();
|
|
|
}
|
|
|
}
|
|
|
|