Skip to content

Commit 7609780

Browse files
nik9000nik9000
authored
ESQL: Move most evaluators to BlockFactory (#101491)
This replaces most of the calls to build a `Block.Builder` on `ExpressionEvaluator` subclasses with calls to methods on the `BlockFactory`. Mostly it changes the methods on the generated code.
1 parent 1722c0a commit 7609780

File tree

194 files changed

+928
-752
lines changed

Some content is hidden

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

194 files changed

+928
-752
lines changed

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

Lines changed: 24 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import javax.lang.model.util.Elements;
2424

2525
import static org.elasticsearch.compute.gen.Methods.appendMethod;
26+
import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
2627
import static org.elasticsearch.compute.gen.Methods.getMethod;
2728
import static org.elasticsearch.compute.gen.Types.ABSTRACT_CONVERT_FUNCTION_EVALUATOR;
2829
import static org.elasticsearch.compute.gen.Types.BLOCK;
@@ -33,6 +34,7 @@
3334
import static org.elasticsearch.compute.gen.Types.SOURCE;
3435
import static org.elasticsearch.compute.gen.Types.VECTOR;
3536
import static org.elasticsearch.compute.gen.Types.blockType;
37+
import static org.elasticsearch.compute.gen.Types.builderType;
3638
import static org.elasticsearch.compute.gen.Types.vectorType;
3739

3840
public class ConvertEvaluatorImplementer {
@@ -131,35 +133,37 @@ private MethodSpec evalVector() {
131133
builder.nextControlFlow("catch (Exception e)");
132134
{
133135
builder.addStatement("registerException(e)");
134-
builder.addStatement("return Block.constantNullBlock(positionCount, driverContext.blockFactory())");
136+
builder.addStatement("return driverContext.blockFactory().newConstantNullBlock(positionCount)");
135137
}
136138
builder.endControlFlow();
137139
}
138140
builder.endControlFlow();
139141

140-
ClassName returnBlockType = blockType(resultType);
141-
builder.addStatement(
142-
"$T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory())",
143-
returnBlockType,
144-
returnBlockType
142+
ClassName resultBuilderType = builderType(blockType(resultType));
143+
builder.beginControlFlow(
144+
"try ($T builder = driverContext.blockFactory().$L(positionCount))",
145+
resultBuilderType,
146+
buildFromFactory(resultBuilderType)
145147
);
146-
builder.beginControlFlow("for (int p = 0; p < positionCount; p++)");
147148
{
148-
builder.beginControlFlow("try");
149+
builder.beginControlFlow("for (int p = 0; p < positionCount; p++)");
149150
{
150-
builder.addStatement("builder.$L($N)", appendMethod(resultType), evalValueCall("vector", "p", scratchPadName));
151-
}
152-
builder.nextControlFlow("catch (Exception e)");
153-
{
154-
builder.addStatement("registerException(e)");
155-
builder.addStatement("builder.appendNull()");
151+
builder.beginControlFlow("try");
152+
{
153+
builder.addStatement("builder.$L($N)", appendMethod(resultType), evalValueCall("vector", "p", scratchPadName));
154+
}
155+
builder.nextControlFlow("catch (Exception e)");
156+
{
157+
builder.addStatement("registerException(e)");
158+
builder.addStatement("builder.appendNull()");
159+
}
160+
builder.endControlFlow();
156161
}
157162
builder.endControlFlow();
163+
builder.addStatement("return builder.build()");
158164
}
159165
builder.endControlFlow();
160166

161-
builder.addStatement("return builder.build()");
162-
163167
return builder.build();
164168
}
165169

@@ -170,11 +174,11 @@ private MethodSpec evalBlock() {
170174
TypeName blockType = blockType(argumentType);
171175
builder.addStatement("$T block = ($T) b", blockType, blockType);
172176
builder.addStatement("int positionCount = block.getPositionCount()");
173-
TypeName resultBlockType = blockType(resultType);
177+
TypeName resultBuilderType = builderType(blockType(resultType));
174178
builder.beginControlFlow(
175-
"try ($T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory()))",
176-
resultBlockType,
177-
resultBlockType
179+
"try ($T builder = driverContext.blockFactory().$L(positionCount))",
180+
resultBuilderType,
181+
buildFromFactory(resultBuilderType)
178182
);
179183
String scratchPadName = null;
180184
if (argumentType.equals(BYTES_REF)) {

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import javax.lang.model.util.Elements;
3434

3535
import static org.elasticsearch.compute.gen.Methods.appendMethod;
36+
import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
3637
import static org.elasticsearch.compute.gen.Methods.getMethod;
3738
import static org.elasticsearch.compute.gen.Types.BLOCK_REF;
3839
import static org.elasticsearch.compute.gen.Types.BYTES_REF;
@@ -45,6 +46,7 @@
4546
import static org.elasticsearch.compute.gen.Types.SOURCE;
4647
import static org.elasticsearch.compute.gen.Types.WARNINGS;
4748
import static org.elasticsearch.compute.gen.Types.blockType;
49+
import static org.elasticsearch.compute.gen.Types.builderType;
4850
import static org.elasticsearch.compute.gen.Types.vectorType;
4951

5052
public class EvaluatorImplementer {
@@ -158,11 +160,11 @@ private MethodSpec realEval(boolean blockStyle) {
158160
builder.addParameter(a.dataType(blockStyle), a.paramName(blockStyle));
159161
}
160162
});
163+
TypeName builderType = builderType(resultDataType);
161164
builder.beginControlFlow(
162-
"try($T.Builder result = $T.$L(positionCount, driverContext.blockFactory()))",
163-
resultDataType,
164-
resultDataType,
165-
resultDataType.simpleName().endsWith("Vector") ? "newVectorBuilder" : "newBlockBuilder"
165+
"try($T result = driverContext.blockFactory().$L(positionCount))",
166+
builderType,
167+
buildFromFactory(builderType)
166168
);
167169
{
168170
processFunction.args.stream().forEach(a -> a.createScratch(builder));

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

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,28 @@
1919
import javax.lang.model.util.ElementFilter;
2020

2121
import static org.elasticsearch.compute.gen.Types.BOOLEAN_BLOCK;
22+
import static org.elasticsearch.compute.gen.Types.BOOLEAN_BLOCK_BUILDER;
2223
import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR;
24+
import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR_BUILDER;
25+
import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR_FIXED_BUILDER;
2326
import static org.elasticsearch.compute.gen.Types.BYTES_REF_BLOCK;
27+
import static org.elasticsearch.compute.gen.Types.BYTES_REF_BLOCK_BUILDER;
28+
import static org.elasticsearch.compute.gen.Types.BYTES_REF_VECTOR_BUILDER;
2429
import static org.elasticsearch.compute.gen.Types.DOUBLE_BLOCK;
30+
import static org.elasticsearch.compute.gen.Types.DOUBLE_BLOCK_BUILDER;
2531
import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR;
32+
import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR_BUILDER;
33+
import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR_FIXED_BUILDER;
2634
import static org.elasticsearch.compute.gen.Types.INT_BLOCK;
35+
import static org.elasticsearch.compute.gen.Types.INT_BLOCK_BUILDER;
2736
import static org.elasticsearch.compute.gen.Types.INT_VECTOR;
37+
import static org.elasticsearch.compute.gen.Types.INT_VECTOR_BUILDER;
38+
import static org.elasticsearch.compute.gen.Types.INT_VECTOR_FIXED_BUILDER;
2839
import static org.elasticsearch.compute.gen.Types.LONG_BLOCK;
40+
import static org.elasticsearch.compute.gen.Types.LONG_BLOCK_BUILDER;
2941
import static org.elasticsearch.compute.gen.Types.LONG_VECTOR;
42+
import static org.elasticsearch.compute.gen.Types.LONG_VECTOR_BUILDER;
43+
import static org.elasticsearch.compute.gen.Types.LONG_VECTOR_FIXED_BUILDER;
3044

3145
/**
3246
* Finds declared methods for the code generator.
@@ -95,6 +109,56 @@ static String appendMethod(TypeName t) {
95109
throw new IllegalArgumentException("unknown append method for [" + t + "]");
96110
}
97111

112+
/**
113+
* Returns the name of the method used to build {@code t} instances
114+
* from a {@code BlockFactory}.
115+
*/
116+
static String buildFromFactory(TypeName t) {
117+
if (t.equals(BOOLEAN_BLOCK_BUILDER)) {
118+
return "newBooleanBlockBuilder";
119+
}
120+
if (t.equals(BOOLEAN_VECTOR_FIXED_BUILDER)) {
121+
return "newBooleanVectorFixedBuilder";
122+
}
123+
if (t.equals(BOOLEAN_VECTOR_BUILDER)) {
124+
return "newBooleanVectorBuilder";
125+
}
126+
if (t.equals(BYTES_REF_BLOCK_BUILDER)) {
127+
return "newBytesRefBlockBuilder";
128+
}
129+
if (t.equals(BYTES_REF_VECTOR_BUILDER)) {
130+
return "newBytesRefVectorBuilder";
131+
}
132+
if (t.equals(INT_BLOCK_BUILDER)) {
133+
return "newIntBlockBuilder";
134+
}
135+
if (t.equals(INT_VECTOR_FIXED_BUILDER)) {
136+
return "newIntVectorFixedBuilder";
137+
}
138+
if (t.equals(INT_VECTOR_BUILDER)) {
139+
return "newIntVectorBuilder";
140+
}
141+
if (t.equals(LONG_BLOCK_BUILDER)) {
142+
return "newLongBlockBuilder";
143+
}
144+
if (t.equals(LONG_VECTOR_FIXED_BUILDER)) {
145+
return "newLongVectorFixedBuilder";
146+
}
147+
if (t.equals(LONG_VECTOR_BUILDER)) {
148+
return "newLongVectorBuilder";
149+
}
150+
if (t.equals(DOUBLE_BLOCK_BUILDER)) {
151+
return "newDoubleBlockBuilder";
152+
}
153+
if (t.equals(DOUBLE_VECTOR_BUILDER)) {
154+
return "newDoubleVectorBuilder";
155+
}
156+
if (t.equals(DOUBLE_VECTOR_FIXED_BUILDER)) {
157+
return "newDoubleVectorFixedBuilder";
158+
}
159+
throw new IllegalArgumentException("unknown build method for [" + t + "]");
160+
}
161+
98162
/**
99163
* Returns the name of the method used to get {@code valueType} instances
100164
* from vectors or blocks.

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

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939
import static org.elasticsearch.compute.gen.Types.SOURCE;
4040
import static org.elasticsearch.compute.gen.Types.WARNINGS;
4141
import static org.elasticsearch.compute.gen.Types.blockType;
42+
import static org.elasticsearch.compute.gen.Types.builderType;
43+
import static org.elasticsearch.compute.gen.Types.vectorFixedBuilderType;
4244
import static org.elasticsearch.compute.gen.Types.vectorType;
4345

4446
public class MvEvaluatorImplementer {
@@ -196,28 +198,19 @@ private MethodSpec evalShell(
196198
builder.beginControlFlow("try (ref)");
197199
builder.addStatement("$T v = ($T) ref.block()", blockType, blockType);
198200
builder.addStatement("int positionCount = v.getPositionCount()");
201+
TypeName builderType;
199202
if (nullable) {
200-
TypeName resultBlockType = blockType(resultType);
201-
builder.beginControlFlow(
202-
"try ($T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory()))",
203-
resultBlockType,
204-
resultBlockType
205-
);
203+
builderType = builderType(blockType(resultType));
206204
} else if (resultType.equals(BYTES_REF)) {
207-
TypeName resultVectorType = vectorType(resultType);
208-
builder.beginControlFlow(
209-
"try ($T.Builder builder = $T.newVectorBuilder(positionCount, driverContext.blockFactory()))",
210-
resultVectorType,
211-
resultVectorType
212-
);
205+
builderType = builderType(vectorType(resultType));
213206
} else {
214-
TypeName resultVectorType = vectorType(resultType);
215-
builder.beginControlFlow(
216-
"try ($T.FixedBuilder builder = $T.newVectorFixedBuilder(positionCount, driverContext.blockFactory()))",
217-
resultVectorType,
218-
resultVectorType
219-
);
207+
builderType = vectorFixedBuilderType(resultType);
220208
}
209+
builder.beginControlFlow(
210+
"try ($T builder = driverContext.blockFactory().$L(positionCount))",
211+
builderType,
212+
Methods.buildFromFactory(builderType)
213+
);
221214

222215
if (false == workType.equals(fieldType) && workType.isPrimitive() == false) {
223216
builder.addStatement("$T work = new $T()", workType, workType);

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,17 @@ public class Types {
5858
static final ClassName LONG_VECTOR = ClassName.get(DATA_PACKAGE, "LongVector");
5959
static final ClassName DOUBLE_VECTOR = ClassName.get(DATA_PACKAGE, "DoubleVector");
6060

61+
static final ClassName BOOLEAN_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "BooleanVector", "Builder");
62+
static final ClassName BYTES_REF_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "BytesRefVector", "Builder");
63+
static final ClassName INT_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "IntVector", "Builder");
64+
static final ClassName LONG_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "LongVector", "Builder");
65+
static final ClassName DOUBLE_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "DoubleVector", "Builder");
66+
67+
static final ClassName BOOLEAN_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "BooleanVector", "FixedBuilder");
68+
static final ClassName INT_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "IntVector", "FixedBuilder");
69+
static final ClassName LONG_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "LongVector", "FixedBuilder");
70+
static final ClassName DOUBLE_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "DoubleVector", "FixedBuilder");
71+
6172
static final ClassName BOOLEAN_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "BooleanArrayVector");
6273
static final ClassName BYTES_REF_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "BytesRefArrayVector");
6374
static final ClassName INT_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "IntArrayVector");
@@ -198,6 +209,56 @@ static ClassName vectorType(String elementType) {
198209
throw new IllegalArgumentException("unknown vector type for [" + elementType + "]");
199210
}
200211

212+
static ClassName builderType(TypeName resultType) {
213+
if (resultType.equals(BOOLEAN_BLOCK)) {
214+
return BOOLEAN_BLOCK_BUILDER;
215+
}
216+
if (resultType.equals(BOOLEAN_VECTOR)) {
217+
return BOOLEAN_VECTOR_BUILDER;
218+
}
219+
if (resultType.equals(BYTES_REF_BLOCK)) {
220+
return BYTES_REF_BLOCK_BUILDER;
221+
}
222+
if (resultType.equals(BYTES_REF_VECTOR)) {
223+
return BYTES_REF_VECTOR_BUILDER;
224+
}
225+
if (resultType.equals(INT_BLOCK)) {
226+
return INT_BLOCK_BUILDER;
227+
}
228+
if (resultType.equals(INT_VECTOR)) {
229+
return INT_VECTOR_BUILDER;
230+
}
231+
if (resultType.equals(LONG_BLOCK)) {
232+
return LONG_BLOCK_BUILDER;
233+
}
234+
if (resultType.equals(LONG_VECTOR)) {
235+
return LONG_VECTOR_BUILDER;
236+
}
237+
if (resultType.equals(DOUBLE_BLOCK)) {
238+
return DOUBLE_BLOCK_BUILDER;
239+
}
240+
if (resultType.equals(DOUBLE_VECTOR)) {
241+
return DOUBLE_VECTOR_BUILDER;
242+
}
243+
throw new IllegalArgumentException("unknown builder type for [" + resultType + "]");
244+
}
245+
246+
static ClassName vectorFixedBuilderType(TypeName elementType) {
247+
if (elementType.equals(TypeName.BOOLEAN)) {
248+
return BOOLEAN_VECTOR_FIXED_BUILDER;
249+
}
250+
if (elementType.equals(TypeName.INT)) {
251+
return INT_VECTOR_FIXED_BUILDER;
252+
}
253+
if (elementType.equals(TypeName.LONG)) {
254+
return LONG_VECTOR_FIXED_BUILDER;
255+
}
256+
if (elementType.equals(TypeName.DOUBLE)) {
257+
return DOUBLE_VECTOR_FIXED_BUILDER;
258+
}
259+
throw new IllegalArgumentException("unknown vector fixed builder type for [" + elementType + "]");
260+
}
261+
201262
static ClassName arrayVectorType(TypeName elementType) {
202263
if (elementType.equals(TypeName.BOOLEAN)) {
203264
return BOOLEAN_ARRAY_VECTOR;

‎x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,9 @@ public BooleanBlock.Builder newBooleanBlockBuilder(int estimatedSize) {
109109
return new BooleanBlockBuilder(estimatedSize, this);
110110
}
111111

112+
/**
113+
* Build a {@link BooleanVector.FixedBuilder} that never grows.
114+
*/
112115
public BooleanVector.FixedBuilder newBooleanVectorFixedBuilder(int size) {
113116
return new BooleanVectorFixedBuilder(size, this);
114117
}
@@ -172,6 +175,9 @@ public IntVector.Builder newIntVectorBuilder(int estimatedSize) {
172175
return new IntVectorBuilder(estimatedSize, this);
173176
}
174177

178+
/**
179+
* Build a {@link IntVector.FixedBuilder} that never grows.
180+
*/
175181
public IntVector.FixedBuilder newIntVectorFixedBuilder(int size) {
176182
return new IntVectorFixedBuilder(size, this);
177183
}
@@ -236,7 +242,10 @@ public LongVector.Builder newLongVectorBuilder(int estimatedSize) {
236242
return new LongVectorBuilder(estimatedSize, this);
237243
}
238244

239-
LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
245+
/**
246+
* Build a {@link LongVector.FixedBuilder} that never grows.
247+
*/
248+
public LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
240249
return new LongVectorFixedBuilder(size, this);
241250
}
242251

@@ -286,7 +295,10 @@ public DoubleVector.Builder newDoubleVectorBuilder(int estimatedSize) {
286295
return new DoubleVectorBuilder(estimatedSize, this);
287296
}
288297

289-
DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
298+
/**
299+
* Build a {@link DoubleVector.FixedBuilder} that never grows.
300+
*/
301+
public DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
290302
return new DoubleVectorFixedBuilder(size, this);
291303
}
292304

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/MockBlockFactory.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,8 @@ public LongVector.Builder newLongVectorBuilder(int estimatedSize) {
245245
return b;
246246
}
247247

248-
LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
248+
@Override
249+
public LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
249250
var b = super.newLongVectorFixedBuilder(size);
250251
track(b, trackDetail());
251252
return b;
@@ -286,7 +287,8 @@ public DoubleVector.Builder newDoubleVectorBuilder(int estimatedSize) {
286287
return b;
287288
}
288289

289-
DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
290+
@Override
291+
public DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
290292
var b = super.newDoubleVectorFixedBuilder(size);
291293
track(b, trackDetail());
292294
return b;

0 commit comments

Comments
 (0)