Skip to content

Commit 2047c9a

Browse files
JonasKunzJonasKunz
authored
ESQL: Add support for exponential_histogram in code generation (#137459)
1 parent 490a2d5 commit 2047c9a

File tree

61 files changed

+317
-251
lines changed

Some content is hidden

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

61 files changed

+317
-251
lines changed

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

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import javax.lang.model.util.Elements;
3838

3939
import static java.util.stream.Collectors.joining;
40+
import static org.elasticsearch.compute.gen.Methods.getMethod;
4041
import static org.elasticsearch.compute.gen.Methods.optionalStaticMethod;
4142
import static org.elasticsearch.compute.gen.Methods.requireAnyArgs;
4243
import static org.elasticsearch.compute.gen.Methods.requireAnyType;
@@ -53,7 +54,6 @@
5354
import static org.elasticsearch.compute.gen.Types.BLOCK;
5455
import static org.elasticsearch.compute.gen.Types.BLOCK_ARRAY;
5556
import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR;
56-
import static org.elasticsearch.compute.gen.Types.BYTES_REF;
5757
import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
5858
import static org.elasticsearch.compute.gen.Types.ELEMENT_TYPE;
5959
import static org.elasticsearch.compute.gen.Types.INTERMEDIATE_STATE_DESC;
@@ -62,6 +62,8 @@
6262
import static org.elasticsearch.compute.gen.Types.PAGE;
6363
import static org.elasticsearch.compute.gen.Types.WARNINGS;
6464
import static org.elasticsearch.compute.gen.Types.blockType;
65+
import static org.elasticsearch.compute.gen.Types.fromString;
66+
import static org.elasticsearch.compute.gen.Types.scratchType;
6567
import static org.elasticsearch.compute.gen.Types.vectorType;
6668

6769
/**
@@ -85,7 +87,7 @@ public class AggregatorImplementer {
8587

8688
private final AggregationState aggState;
8789
private final List<Argument> aggParams;
88-
private final boolean hasOnlyBlockArguments;
90+
private final boolean tryToUseVectors;
8991

9092
public AggregatorImplementer(
9193
Elements elements,
@@ -119,7 +121,8 @@ public AggregatorImplementer(
119121
return a;
120122
}).filter(a -> a instanceof PositionArgument == false).toList();
121123

122-
this.hasOnlyBlockArguments = this.aggParams.stream().allMatch(a -> a instanceof BlockArgument);
124+
this.tryToUseVectors = aggParams.stream().anyMatch(a -> (a instanceof BlockArgument) == false)
125+
&& aggParams.stream().noneMatch(a -> a.supportsVectorReadAccess() == false);
123126

124127
this.createParameters = init.getParameters()
125128
.stream()
@@ -199,7 +202,7 @@ private TypeSpec type() {
199202
builder.addMethod(addRawInput());
200203
builder.addMethod(addRawInputExploded(true));
201204
builder.addMethod(addRawInputExploded(false));
202-
if (hasOnlyBlockArguments == false) {
205+
if (tryToUseVectors) {
203206
builder.addMethod(addRawVector(false));
204207
builder.addMethod(addRawVector(true));
205208
}
@@ -340,16 +343,18 @@ private MethodSpec addRawInputExploded(boolean hasMask) {
340343
builder.addStatement("$T $L = page.getBlock(channels.get($L))", a.dataType(true), a.blockName(), i);
341344
}
342345

343-
for (Argument a : aggParams) {
344-
String rawBlock = "addRawBlock("
345-
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
346-
+ (hasMask ? ", mask" : "")
347-
+ ")";
346+
if (tryToUseVectors) {
347+
for (Argument a : aggParams) {
348+
String rawBlock = "addRawBlock("
349+
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
350+
+ (hasMask ? ", mask" : "")
351+
+ ")";
348352

349-
a.resolveVectors(builder, rawBlock, "return");
353+
a.resolveVectors(builder, rawBlock, "return");
354+
}
350355
}
351356

352-
builder.addStatement(invokeAddRaw(hasOnlyBlockArguments, hasMask));
357+
builder.addStatement(invokeAddRaw(tryToUseVectors == false, hasMask));
353358
return builder.build();
354359
}
355360

@@ -499,9 +504,9 @@ private MethodSpec.Builder initAddRaw(boolean blockStyle, boolean masked) {
499504
builder.addParameter(BOOLEAN_VECTOR, "mask");
500505
}
501506
for (Argument a : aggParams) {
502-
if (a.isBytesRef()) {
503-
// Add bytes_ref scratch var that will be used for bytes_ref blocks/vectors
504-
builder.addStatement("$T $L = new $T()", BYTES_REF, a.scratchName(), BYTES_REF);
507+
if (a.scratchType() != null) {
508+
// Add scratch var that will be used for some blocks/vectors, e.g. for bytes_ref
509+
builder.addStatement("$T $L = new $T()", a.scratchType(), a.scratchName(), a.scratchType());
505510
}
506511
}
507512
return builder;
@@ -610,8 +615,8 @@ private MethodSpec addIntermediateInput() {
610615
).map(Methods::requireType).toArray(TypeMatcher[]::new)
611616
)
612617
);
613-
if (intermediateState.stream().map(IntermediateStateDesc::elementType).anyMatch(n -> n.equals("BYTES_REF"))) {
614-
builder.addStatement("$T scratch = new $T()", BYTES_REF, BYTES_REF);
618+
for (IntermediateStateDesc interState : intermediateState) {
619+
interState.addScratchDeclaration(builder);
615620
}
616621
builder.addStatement("$T.combineIntermediate(state, " + intermediateStateRowAccess() + ")", declarationType);
617622
}
@@ -706,13 +711,25 @@ public String access(String position) {
706711
if (block) {
707712
return name();
708713
}
709-
String s = name() + "." + vectorAccessorName(elementType()) + "(" + position;
710-
if (elementType().equals("BYTES_REF")) {
711-
s += ", scratch";
714+
String s = name() + ".";
715+
if (vectorType(elementType) != null) {
716+
s += vectorAccessorName(elementType()) + "(" + position;
717+
} else {
718+
s += getMethod(fromString(elementType())) + "(" + name() + ".getFirstValueIndex(" + position + ")";
719+
}
720+
if (scratchType(elementType()) != null) {
721+
s += ", " + name() + "Scratch";
712722
}
713723
return s + ")";
714724
}
715725

726+
public void addScratchDeclaration(MethodSpec.Builder builder) {
727+
ClassName scratchType = scratchType(elementType());
728+
if (scratchType != null) {
729+
builder.addStatement("$T $L = new $T()", scratchType, name() + "Scratch", scratchType);
730+
}
731+
}
732+
716733
public void assignToVariable(MethodSpec.Builder builder, int offset) {
717734
builder.addStatement("Block $L = page.getBlock(channels.get($L))", name + "Uncast", offset);
718735
ClassName blockType = blockType(elementType());
@@ -721,7 +738,7 @@ public void assignToVariable(MethodSpec.Builder builder, int offset) {
721738
builder.addStatement("return");
722739
builder.endControlFlow();
723740
}
724-
if (block) {
741+
if (block || vectorType(elementType) == null) {
725742
builder.addStatement("$T $L = ($T) $L", blockType, name, blockType, name + "Uncast");
726743
} else {
727744
builder.addStatement("$T $L = (($T) $L).asVector()", vectorType(elementType), name, blockType, name + "Uncast");
@@ -732,6 +749,7 @@ public TypeName combineArgType() {
732749
var type = Types.fromString(elementType);
733750
return block ? blockType(type) : type;
734751
}
752+
735753
}
736754

737755
/**

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ public class EvaluatorImplementer {
5151
private final ProcessFunction processFunction;
5252
private final ClassName implementation;
5353
private final boolean processOutputsMultivalued;
54+
private final boolean vectorsUnsupported;
5455
private final boolean allNullsIsNull;
5556

5657
public EvaluatorImplementer(
@@ -69,6 +70,8 @@ public EvaluatorImplementer(
6970
declarationType.getSimpleName() + extraName + "Evaluator"
7071
);
7172
this.processOutputsMultivalued = this.processFunction.hasBlockType;
73+
boolean anyParameterNotSupportingVectors = this.processFunction.args.stream().anyMatch(a -> a.supportsVectorReadAccess() == false);
74+
vectorsUnsupported = processOutputsMultivalued || anyParameterNotSupportingVectors;
7275
this.allNullsIsNull = allNullsIsNull;
7376
}
7477

@@ -101,7 +104,7 @@ private TypeSpec type() {
101104
builder.addMethod(eval());
102105
builder.addMethod(processFunction.baseRamBytesUsed());
103106

104-
if (processOutputsMultivalued) {
107+
if (vectorsUnsupported) {
105108
if (processFunction.args.stream().anyMatch(x -> x instanceof FixedArgument == false)) {
106109
builder.addMethod(realEval(true));
107110
}
@@ -145,7 +148,7 @@ private MethodSpec eval() {
145148
builder.addModifiers(Modifier.PUBLIC).returns(BLOCK).addParameter(PAGE, "page");
146149
processFunction.args.forEach(a -> a.evalToBlock(builder));
147150
String invokeBlockEval = invokeRealEval(true);
148-
if (processOutputsMultivalued) {
151+
if (vectorsUnsupported) {
149152
builder.addStatement(invokeBlockEval);
150153
} else {
151154
processFunction.args.forEach(a -> a.resolveVectors(builder, invokeBlockEval));

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

Lines changed: 28 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@
5151
import static org.elasticsearch.compute.gen.Types.BIG_ARRAYS;
5252
import static org.elasticsearch.compute.gen.Types.BLOCK;
5353
import static org.elasticsearch.compute.gen.Types.BLOCK_ARRAY;
54-
import static org.elasticsearch.compute.gen.Types.BYTES_REF;
5554
import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
5655
import static org.elasticsearch.compute.gen.Types.ELEMENT_TYPE;
5756
import static org.elasticsearch.compute.gen.Types.GROUPING_AGGREGATOR_EVALUATOR_CONTEXT;
@@ -93,6 +92,7 @@ public class GroupingAggregatorImplementer {
9392
private final AggregationState aggState;
9493
private final List<Argument> aggParams;
9594
private final boolean hasOnlyBlockArguments;
95+
private final boolean allArgumentsSupportVectors;
9696

9797
public GroupingAggregatorImplementer(
9898
Elements elements,
@@ -128,6 +128,7 @@ public GroupingAggregatorImplementer(
128128
}).filter(a -> a instanceof PositionArgument == false).toList();
129129

130130
this.hasOnlyBlockArguments = this.aggParams.stream().allMatch(a -> a instanceof BlockArgument);
131+
this.allArgumentsSupportVectors = aggParams.stream().noneMatch(a -> a.supportsVectorReadAccess() == false);
131132

132133
this.createParameters = init.getParameters()
133134
.stream()
@@ -204,7 +205,7 @@ private TypeSpec type() {
204205
builder.addMethod(prepareProcessRawInputPage());
205206
for (ClassName groupIdClass : GROUP_IDS_CLASSES) {
206207
builder.addMethod(addRawInputLoop(groupIdClass, false));
207-
if (hasOnlyBlockArguments == false) {
208+
if (hasOnlyBlockArguments == false && allArgumentsSupportVectors) {
208209
builder.addMethod(addRawInputLoop(groupIdClass, true));
209210
}
210211
builder.addMethod(addIntermediateInput(groupIdClass));
@@ -330,26 +331,31 @@ private MethodSpec prepareProcessRawInputPage() {
330331
builder.addStatement("$T $L = page.getBlock(channels.get($L))", a.dataType(true), a.blockName(), i);
331332
}
332333

333-
for (Argument a : aggParams) {
334-
builder.addStatement(
335-
"$T $L = $L.asVector()",
336-
vectorType(a.elementType()),
337-
(a instanceof BlockArgument) ? (a.name() + "Vector") : a.vectorName(),
338-
a.blockName()
339-
);
340-
builder.beginControlFlow("if ($L == null)", (a instanceof BlockArgument) ? (a.name() + "Vector") : a.vectorName());
341-
{
334+
String groupIdTrackingStatement = "maybeEnableGroupIdTracking(seenGroupIds, "
335+
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
336+
+ ")";
337+
338+
if (allArgumentsSupportVectors) {
339+
340+
for (Argument a : aggParams) {
342341
builder.addStatement(
343-
"maybeEnableGroupIdTracking(seenGroupIds, "
344-
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
345-
+ ")"
342+
"$T $L = $L.asVector()",
343+
vectorType(a.elementType()),
344+
(a instanceof BlockArgument) ? (a.name() + "Vector") : a.vectorName(),
345+
a.blockName()
346346
);
347-
returnAddInput(builder, false);
347+
builder.beginControlFlow("if ($L == null)", (a instanceof BlockArgument) ? (a.name() + "Vector") : a.vectorName());
348+
{
349+
builder.addStatement(groupIdTrackingStatement);
350+
returnAddInput(builder, false);
351+
}
352+
builder.endControlFlow();
348353
}
349-
builder.endControlFlow();
354+
returnAddInput(builder, true);
355+
} else {
356+
builder.addStatement(groupIdTrackingStatement);
357+
returnAddInput(builder, false);
350358
}
351-
352-
returnAddInput(builder, true);
353359
return builder.build();
354360
}
355361

@@ -443,9 +449,9 @@ private MethodSpec addRawInputLoop(TypeName groupsType, boolean valuesAreVector)
443449
);
444450
}
445451
for (Argument a : aggParams) {
446-
if (a.isBytesRef()) {
447-
// Add bytes_ref scratch var that will be used for bytes_ref blocks/vectors
448-
builder.addStatement("$T $L = new $T()", BYTES_REF, a.scratchName(), BYTES_REF);
452+
if (a.scratchType() != null) {
453+
// Add scratch var that will be used for some blocks/vectors, e.g. for bytes_ref
454+
builder.addStatement("$T $L = new $T()", a.scratchType(), a.scratchName(), a.scratchType());
449455
}
450456
}
451457

@@ -645,11 +651,7 @@ private MethodSpec addIntermediateInput(TypeName groupsType) {
645651
.collect(Collectors.joining(", "));
646652
builder.addStatement("$T.combineIntermediate(state, positionOffset, groups, " + states + ")", declarationType);
647653
} else {
648-
if (intermediateState.stream()
649-
.map(AggregatorImplementer.IntermediateStateDesc::elementType)
650-
.anyMatch(n -> n.equals("BYTES_REF"))) {
651-
builder.addStatement("$T scratch = new $T()", BYTES_REF, BYTES_REF);
652-
}
654+
intermediateState.forEach(state -> state.addScratchDeclaration(builder));
653655
builder.beginControlFlow("for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++)");
654656
{
655657
if (groupsIsBlock) {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,9 @@ public static String getMethod(TypeName elementType) {
314314
if (elementType.equals(TypeName.FLOAT)) {
315315
return "getFloat";
316316
}
317+
if (elementType.equals(Types.EXPONENTIAL_HISTOGRAM)) {
318+
return "getExponentialHistogram";
319+
}
317320
throw new IllegalArgumentException("unknown get method for [" + elementType + "]");
318321
}
319322

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

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -52,13 +52,16 @@ public class Types {
5252
public static final ClassName LONG_BLOCK = ClassName.get(DATA_PACKAGE, "LongBlock");
5353
public static final ClassName DOUBLE_BLOCK = ClassName.get(DATA_PACKAGE, "DoubleBlock");
5454
public static final ClassName FLOAT_BLOCK = ClassName.get(DATA_PACKAGE, "FloatBlock");
55+
public static final ClassName EXPONENTIAL_HISTOGRAM_BLOCK = ClassName.get(DATA_PACKAGE, "ExponentialHistogramBlock");
56+
public static final ClassName EXPONENTIAL_HISTOGRAM_SCRATCH = ClassName.get(DATA_PACKAGE, "ExponentialHistogramScratch");
5557

5658
static final ClassName BOOLEAN_BLOCK_BUILDER = BOOLEAN_BLOCK.nestedClass("Builder");
5759
static final ClassName BYTES_REF_BLOCK_BUILDER = BYTES_REF_BLOCK.nestedClass("Builder");
5860
static final ClassName INT_BLOCK_BUILDER = INT_BLOCK.nestedClass("Builder");
5961
static final ClassName LONG_BLOCK_BUILDER = LONG_BLOCK.nestedClass("Builder");
6062
static final ClassName DOUBLE_BLOCK_BUILDER = DOUBLE_BLOCK.nestedClass("Builder");
6163
static final ClassName FLOAT_BLOCK_BUILDER = FLOAT_BLOCK.nestedClass("Builder");
64+
static final ClassName EXPONENTIAL_HISTOGRAM_BLOCK_BUILDER = ClassName.get(DATA_PACKAGE, "ExponentialHistogramBlockBuilder");
6265

6366
static final ClassName ELEMENT_TYPE = ClassName.get(DATA_PACKAGE, "ElementType");
6467

@@ -133,24 +136,32 @@ public class Types {
133136
static final ClassName SOURCE = ClassName.get("org.elasticsearch.xpack.esql.core.tree", "Source");
134137

135138
public static final ClassName BYTES_REF = ClassName.get("org.apache.lucene.util", "BytesRef");
139+
public static final ClassName EXPONENTIAL_HISTOGRAM = ClassName.get("org.elasticsearch.exponentialhistogram", "ExponentialHistogram");
136140

137141
public static final ClassName RELEASABLE = ClassName.get("org.elasticsearch.core", "Releasable");
138142
public static final ClassName RELEASABLES = ClassName.get("org.elasticsearch.core", "Releasables");
139143

140-
private record TypeDef(TypeName type, String alias, ClassName block, ClassName vector) {
144+
private record TypeDef(TypeName type, String alias, ClassName block, ClassName vector, ClassName scratch) {
141145

142-
public static TypeDef of(TypeName type, String alias, String block, String vector) {
143-
return new TypeDef(type, alias, ClassName.get(DATA_PACKAGE, block), ClassName.get(DATA_PACKAGE, vector));
146+
public static TypeDef of(TypeName type, String alias, String block, String vector, ClassName scratch) {
147+
return new TypeDef(
148+
type,
149+
alias,
150+
ClassName.get(DATA_PACKAGE, block),
151+
vector == null ? null : ClassName.get(DATA_PACKAGE, vector),
152+
scratch
153+
);
144154
}
145155
}
146156

147157
private static final Map<String, TypeDef> TYPES = Stream.of(
148-
TypeDef.of(TypeName.BOOLEAN, "BOOLEAN", "BooleanBlock", "BooleanVector"),
149-
TypeDef.of(TypeName.INT, "INT", "IntBlock", "IntVector"),
150-
TypeDef.of(TypeName.LONG, "LONG", "LongBlock", "LongVector"),
151-
TypeDef.of(TypeName.FLOAT, "FLOAT", "FloatBlock", "FloatVector"),
152-
TypeDef.of(TypeName.DOUBLE, "DOUBLE", "DoubleBlock", "DoubleVector"),
153-
TypeDef.of(BYTES_REF, "BYTES_REF", "BytesRefBlock", "BytesRefVector")
158+
TypeDef.of(TypeName.BOOLEAN, "BOOLEAN", "BooleanBlock", "BooleanVector", null),
159+
TypeDef.of(TypeName.INT, "INT", "IntBlock", "IntVector", null),
160+
TypeDef.of(TypeName.LONG, "LONG", "LongBlock", "LongVector", null),
161+
TypeDef.of(TypeName.FLOAT, "FLOAT", "FloatBlock", "FloatVector", null),
162+
TypeDef.of(TypeName.DOUBLE, "DOUBLE", "DoubleBlock", "DoubleVector", null),
163+
TypeDef.of(BYTES_REF, "BYTES_REF", "BytesRefBlock", "BytesRefVector", BYTES_REF),
164+
TypeDef.of(EXPONENTIAL_HISTOGRAM, "EXPONENTIAL_HISTOGRAM", "ExponentialHistogramBlock", null, EXPONENTIAL_HISTOGRAM_SCRATCH)
154165
)
155166
.flatMap(def -> Stream.of(def.type.toString(), def.type + "[]", def.alias).map(alias -> Map.entry(alias, def)))
156167
.collect(toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));
@@ -183,6 +194,14 @@ static ClassName vectorType(String elementType) {
183194
return findRequired(elementType, "vector").vector;
184195
}
185196

197+
public static ClassName scratchType(String elementType) {
198+
TypeDef typeDef = TYPES.get(elementType);
199+
if (typeDef != null) {
200+
return typeDef.scratch;
201+
}
202+
return null;
203+
}
204+
186205
static ClassName builderType(TypeName resultType) {
187206
if (resultType.equals(BOOLEAN_BLOCK)) {
188207
return BOOLEAN_BLOCK_BUILDER;
@@ -220,6 +239,9 @@ static ClassName builderType(TypeName resultType) {
220239
if (resultType.equals(FLOAT_VECTOR)) {
221240
return FLOAT_VECTOR_BUILDER;
222241
}
242+
if (resultType.equals(EXPONENTIAL_HISTOGRAM_BLOCK)) {
243+
return EXPONENTIAL_HISTOGRAM_BLOCK_BUILDER;
244+
}
223245
throw new IllegalArgumentException("unknown builder type for [" + resultType + "]");
224246
}
225247

@@ -261,6 +283,9 @@ public static TypeName elementType(TypeName t) {
261283
if (t.equals(FLOAT_BLOCK) || t.equals(FLOAT_VECTOR) || t.equals(FLOAT_BLOCK_BUILDER)) {
262284
return TypeName.FLOAT;
263285
}
286+
if (t.equals(EXPONENTIAL_HISTOGRAM_BLOCK) || t.equals(EXPONENTIAL_HISTOGRAM_BLOCK_BUILDER)) {
287+
return EXPONENTIAL_HISTOGRAM;
288+
}
264289
throw new IllegalArgumentException("unknown element type for [" + t + "]");
265290
}
266291

0 commit comments

Comments
 (0)