Skip to content

[ESQL] Support first/last functions. #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[ESQL] Support first/last functions. #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 56 additions & 53 deletions ...n/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/AggregatorImplementer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ public class AggregatorImplementer {
private final List<IntermediateStateDesc> intermediateState;

private final AggregationState aggState;
private final List<AggregationParameter> aggParams;
private final List<Argument> aggParams;

public AggregatorImplementer(
Elements elements,
Expand All @@ -111,15 +111,13 @@ public AggregatorImplementer(
requireName("combine"),
requireArgsStartsWith(requireType(aggState.declaredType()), requireAnyType("<aggregation input column type>"))
);
this.aggParams = combine.getParameters().stream().skip(1).flatMap(v -> {
this.aggParams = combine.getParameters().stream().skip(1).map(v -> {
Argument a = Argument.fromParameter(types, v);
return switch (a) {
case StandardArgument sa -> Stream.of(new AggregationParameter(sa.name(), sa.type(), false));
case BlockArgument ba -> Stream.of(new AggregationParameter(ba.name(), Types.elementType(ba.type()), true));
case PositionArgument pa -> Stream.of();
default -> throw new IllegalArgumentException("unsupported argument [" + declarationType + "][" + a + "]");
};
}).toList();
if ((a instanceof StandardArgument || a instanceof BlockArgument || a instanceof PositionArgument) == false) {
throw new IllegalArgumentException("unsupported argument [" + declarationType + "][" + a + "]");
}
return a;
}).filter(a -> a instanceof PositionArgument == false).toList();

this.createParameters = init.getParameters()
.stream()
Expand Down Expand Up @@ -334,33 +332,36 @@ private MethodSpec addRawInputExploded(boolean hasMask) {
}

for (int i = 0; i < aggParams.size(); i++) {
AggregationParameter p = aggParams.get(i);
builder.addStatement("$T $L = page.getBlock(channels.get($L))", blockType(p.type()), p.blockName(), i);
Argument a = aggParams.get(i);
builder.addStatement("$T $L = page.getBlock(channels.get($L))", a.dataType(true), a.blockName(), i);
}
for (AggregationParameter p : aggParams) {
builder.addStatement("$T $L = $L.asVector()", vectorType(p.type()), p.vectorName(), p.blockName());
builder.beginControlFlow("if ($L == null)", p.vectorName());

boolean isBlockArgument = aggParams.getFirst() instanceof BlockArgument;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: Replace the invalid use of aggParams.getFirst() with a safe indexed access guarded by an emptiness check (use aggParams.get(0) when non-empty). [possible bug]

Suggested change
boolean isBlockArgument = aggParams.getFirst() instanceof BlockArgument;
boolean isBlockArgument = aggParams.isEmpty() == false && aggParams.get(0) instanceof BlockArgument;
Why Change? ⭐

The original code called getFirst() on a List, which is not a method on java.util.List and would not compile.
Replacing it with a guarded indexed access (aggParams.isEmpty() == false && aggParams.get(0) instanceof BlockArgument)
uses standard List APIs (isEmpty and get). This prevents IndexOutOfBoundsException by checking emptiness before get(0).
The surrounding code previously iterated over aggParams or used its size, so using a guard here preserves semantics while fixing
the compile error. The change is syntactically valid Java and uses only symbols present in the file (aggParams, BlockArgument).

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Array Bounds Risk

Direct getFirst() call without size validation can throw NoSuchElementException if aggParams list is empty. Runtime exception causes aggregation function generation failure preventing system operation.

Standards
  • ISO-IEC-25010-Reliability-Fault-Tolerance
  • DbC-Precondition-Violation


for (Argument a : aggParams) {
String rawBlock = "addRawBlock("
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
+ (hasMask ? ", mask" : "")
+ ")";

if (isBlockArgument == false) {
a.resolveVectors(builder, rawBlock, "return");
} else {
builder.addStatement(rawBlock);
}
Comment on lines +341 to +351
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Redundant Block Processing

String concatenation for rawBlock statement occurs inside loop for each argument, creating identical strings repeatedly. This creates unnecessary string allocation overhead proportional to argument count.

Standards
  • ISO-IEC-25010-Performance-Efficiency-Resource-Utilization
  • Optimization-Pattern-Loop-Invariant-Code-Motion

Comment on lines +339 to +351
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Block Processing Logic

Loop determines block processing mode based only on first parameter but applies same rawBlock statement construction inside loop for each argument. Logic creates redundant addRawBlock statements when isBlockArgument is true since same statement executes multiple times within the loop.

Standards
  • Algorithm-Correctness-Loop-Logic
  • Logic-Verification-Control-Flow

}
Comment on lines +341 to +352
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There are a couple of issues in this loop.

  1. The rawBlock string is constructed inside the for loop, which is inefficient as it will be identical in every iteration. It should be defined once before this logic.
  2. More critically, if isBlockArgument is true, builder.addStatement(rawBlock) is executed for every argument in aggParams. This generates the addRawBlock(...) call multiple times, which is incorrect. It should only be generated once.

The logic should be restructured to address both of these points.

Suggested change
for (Argument a : aggParams) {
String rawBlock = "addRawBlock("
+ aggParams.stream().map(arg -> arg.blockName()).collect(joining(", "))
+ (hasMask ? ", mask" : "")
+ ")";
if (isBlockArgument == false) {
a.resolveVectors(builder, rawBlock, "return");
} else {
builder.addStatement(rawBlock);
}
}
String rawBlock = "addRawBlock("
+ aggParams.stream().map(Argument::blockName).collect(joining(", "))
+ (hasMask ? ", mask" : "")
+ ")";
if (isBlockArgument) {
builder.addStatement(rawBlock);
} else {
for (Argument a : aggParams) {
a.resolveVectors(builder, rawBlock, "return");
}
}


if (isBlockArgument == false) {
builder.addStatement(
"addRawBlock("
+ aggParams.stream().map(AggregationParameter::blockName).collect(joining(", "))
+ (hasMask ? ", mask" : "")
+ ")"
"addRawVector(" + aggParams.stream().map(a -> a.vectorName()).collect(joining(", ")) + (hasMask ? ", mask" : "") + ")"
);
builder.addStatement("return");
builder.endControlFlow();
}
builder.addStatement(
"addRawVector("
+ aggParams.stream().map(AggregationParameter::vectorName).collect(joining(", "))
+ (hasMask ? ", mask" : "")
+ ")"
);
return builder.build();
}

private MethodSpec addRawVector(boolean masked) {
MethodSpec.Builder builder = initAddRaw(true, masked);
if (aggParams.getFirst().isArray()) {
if (aggParams.getFirst() instanceof BlockArgument) {
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: Replace the aggParams.getFirst() check with a guarded indexed access (use aggParams.get(0) only when the list is not empty) to avoid compile errors. [possible bug]

Suggested change
if (aggParams.getFirst() instanceof BlockArgument) {
if (aggParams.isEmpty() == false && aggParams.get(0) instanceof BlockArgument) {
Why Change? ⭐

This replaces an invalid List.getFirst() invocation with a safe check that first ensures the list is non-empty,
then accesses the first element via get(0). The code uses standard List methods (isEmpty and get), so it will compile.
The guard avoids an IndexOutOfBoundsException in the unlikely case aggParams is empty. All referenced symbols exist
in the file and the change preserves the original intent of checking whether the first argument is a BlockArgument.

builder.addComment("This type does not support vectors because all values are multi-valued");
return builder.build();
}
Expand All @@ -383,8 +384,8 @@ private MethodSpec addRawVector(boolean masked) {
if (masked) {
builder.beginControlFlow("if (mask.getBoolean(valuesPosition) == false)").addStatement("continue").endControlFlow();
}
for (AggregationParameter p : aggParams) {
p.read(builder, true);
for (Argument a : aggParams) {
a.read(builder, a.vectorName(), "valuesPosition");
}
combineRawInput(builder, false);
}
Expand All @@ -406,8 +407,8 @@ private void addRawVectorWithFirst(MethodSpec.Builder builder, boolean firstPass
builder.addStatement("continue");
builder.endControlFlow();
}
for (AggregationParameter p : aggParams) {
p.read(builder, true);
for (Argument a : aggParams) {
a.read(builder, a.vectorName(), "valuesPosition");
}
combineRawInput(builder, firstPass);
builder.addStatement("valuesPosition++");
Expand All @@ -427,14 +428,14 @@ private MethodSpec addRawBlock(boolean masked) {
if (masked) {
builder.beginControlFlow("if (mask.getBoolean(p) == false)").addStatement("continue").endControlFlow();
}
for (AggregationParameter p : aggParams) {
builder.addStatement("int $LValueCount = $L.getValueCount(p)", p.name(), p.blockName());
builder.beginControlFlow("if ($LValueCount == 0)", p.name());
for (Argument a : aggParams) {
builder.addStatement("int $LValueCount = $L.getValueCount(p)", a.name(), a.blockName());
builder.beginControlFlow("if ($LValueCount == 0)", a.name());
builder.addStatement("continue");
builder.endControlFlow();
}

if (aggParams.getFirst().isArray()) {
if (aggParams.getFirst() instanceof BlockArgument) {
if (aggParams.size() > 1) {
throw new IllegalArgumentException("array mode not supported for multiple args");
}
Expand All @@ -446,18 +447,18 @@ private MethodSpec addRawBlock(boolean masked) {
if (first == null && aggState.hasSeen()) {
builder.addStatement("state.seen(true)");
}
for (AggregationParameter p : aggParams) {
builder.addStatement("int $L = $L.getFirstValueIndex(p)", p.startName(), p.blockName());
builder.addStatement("int $L = $L + $LValueCount", p.endName(), p.startName(), p.name());
for (Argument a : aggParams) {
builder.addStatement("int $L = $L.getFirstValueIndex(p)", a.startName(), a.blockName());
builder.addStatement("int $L = $L + $LValueCount", a.endName(), a.startName(), a.name());
builder.beginControlFlow(
"for (int $L = $L; $L < $L; $L++)",
p.offsetName(),
p.startName(),
p.offsetName(),
p.endName(),
p.offsetName()
a.offsetName(),
a.startName(),
a.offsetName(),
a.endName(),
a.offsetName()
);
p.read(builder, false);
a.read(builder, a.blockName(), a.offsetName());
}
if (first != null) {
builder.addComment("Check seen in every iteration to save on complexity in the Block path");
Expand All @@ -474,7 +475,7 @@ private MethodSpec addRawBlock(boolean masked) {
} else {
combineRawInput(builder, false);
}
for (AggregationParameter p : aggParams) {
for (Argument a : aggParams) {
builder.endControlFlow();
}
}
Expand All @@ -486,19 +487,21 @@ private MethodSpec addRawBlock(boolean masked) {
private MethodSpec.Builder initAddRaw(boolean valuesAreVector, boolean masked) {
MethodSpec.Builder builder = MethodSpec.methodBuilder(valuesAreVector ? "addRawVector" : "addRawBlock");
builder.addModifiers(Modifier.PRIVATE);
for (AggregationParameter p : aggParams) {
for (Argument a : aggParams) {
boolean isBlockArgument = a instanceof BlockArgument;
TypeName typeName = isBlockArgument ? Types.elementType(a.type()) : a.type();
builder.addParameter(
valuesAreVector ? vectorType(p.type) : blockType(p.type),
valuesAreVector ? p.vectorName() : p.blockName()
valuesAreVector ? vectorType(typeName) : blockType(typeName),
valuesAreVector ? a.vectorName() : a.blockName()
);
}
if (masked) {
builder.addParameter(BOOLEAN_VECTOR, "mask");
}
for (AggregationParameter p : aggParams) {
if (p.isBytesRef()) {
for (Argument a : aggParams) {
if (a.isBytesRef()) {
// Add bytes_ref scratch var that will be used for bytes_ref blocks/vectors
builder.addStatement("$T $L = new $T()", BYTES_REF, p.scratchName(), BYTES_REF);
builder.addStatement("$T $L = new $T()", BYTES_REF, a.scratchName(), BYTES_REF);
}
}
return builder;
Expand Down Expand Up @@ -527,9 +530,9 @@ private void invokeCombineRawInput(TypeName returnType, MethodSpec.Builder build
} else {
throw new IllegalArgumentException("combine must return void or a primitive");
}
for (AggregationParameter p : aggParams) {
for (Argument a : aggParams) {
pattern.append(", $L");
params.add(p.valueName());
params.add(a.valueName());
}
if (returnType.isPrimitive()) {
pattern.append(")");
Expand Down
46 changes: 45 additions & 1 deletion ...lugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/argument/Argument.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,14 @@
import org.elasticsearch.compute.gen.Types;

import java.util.List;
import java.util.Objects;

import javax.lang.model.element.VariableElement;
import javax.lang.model.type.ArrayType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;

import static org.elasticsearch.compute.gen.Types.BYTES_REF;
import static org.elasticsearch.compute.gen.argument.StandardArgument.isBlockType;

/**
Expand Down Expand Up @@ -64,6 +66,46 @@ static Argument fromParameter(javax.lang.model.util.Types types, VariableElement
return new StandardArgument(type, name);
}

default String blockName() {
return paramName(true);
}

default String vectorName() {
return paramName(false);
}

default String valueName() {
return name() + "Value";
}

default String startName() {
return name() + "Start";
}

default String endName() {
return name() + "End";
}

default String offsetName() {
return name() + "Offset";
}

default String scratchName() {
if (isBytesRef() == false) {
throw new IllegalStateException("can't build scratch for non-BytesRef");
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: Include the argument name in the exception message from scratchName() to make failures easier to diagnose. [enhancement]

Suggested change
throw new IllegalStateException("can't build scratch for non-BytesRef");
throw new IllegalStateException("can't build scratch for non-BytesRef for parameter: " + name());
Why Change? ⭐

The proposed change only enhances the exception message by concatenating the existing name() value. In an interface default method it is valid to call the abstract name() method implemented by concrete classes, so this is syntactically correct and will compile. The change does not alter control flow or state, only the text of the IllegalStateException; it therefore cannot introduce runtime bugs beyond the original exception and actually aids debugging by including the parameter name. Assumes implementations provide a reasonable name() and that calling it in the error path is acceptable.

}

return name() + "Scratch";
}

default boolean isBytesRef() {
return Objects.equals(type(), BYTES_REF);
}

String name();

TypeName type();

/**
* 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.
Expand Down Expand Up @@ -121,7 +163,7 @@ static Argument fromParameter(javax.lang.model.util.Types types, VariableElement
* call the block flavored evaluator if this is a block. Noop if the
* parameter is {@link Fixed}.
*/
void resolveVectors(MethodSpec.Builder builder, String invokeBlockEval);
void resolveVectors(MethodSpec.Builder builder, String... invokeBlockEval);

/**
* Create any scratch structures needed by {@code eval}.
Expand All @@ -144,6 +186,8 @@ static Argument fromParameter(javax.lang.model.util.Types types, VariableElement
*/
void read(MethodSpec.Builder builder, boolean blockStyle);

void read(MethodSpec.Builder builder, String accessor, String firstParam);
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: Provide a default implementation for the newly added read(builder, accessor, firstParam) overload to avoid breaking existing implementers of this interface; make it throw an explicit UnsupportedOperationException until concrete implementations are updated. [possible bug]

Suggested change
void read(MethodSpec.Builder builder, String accessor, String firstParam);
default void read(MethodSpec.Builder builder, String accessor, String firstParam) {
throw new UnsupportedOperationException("read(builder, accessor, firstParam) is not implemented for " + getClass().getName());
}
Why Change? ⭐

Converting the newly added abstract overload to a default method that throws UnsupportedOperationException is syntactically valid for a Java interface and will compile. Default methods can reference instance methods such as getClass(), so the diagnostic message will include the concrete implementor's class name. This preserves binary compatibility for existing implementers (they won't be forced to implement the new overload immediately) while clearly signalling unimplemented behavior at runtime. Assumes it's acceptable in this codebase to temporarily provide a throwing default until concrete classes are updated.


/**
* Build the invocation of the process method for this parameter.
*/
Expand Down
29 changes: 17 additions & 12 deletions .../esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/argument/ArrayArgument.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,13 @@
import static org.elasticsearch.compute.gen.Types.blockType;
import static org.elasticsearch.compute.gen.Types.vectorType;

public record ArrayArgument(TypeName componentType, String name) implements Argument {
public record ArrayArgument(TypeName type, String name) implements Argument {
@Override
public TypeName dataType(boolean blockStyle) {
if (blockStyle) {
return ArrayTypeName.of(blockType(componentType));
return ArrayTypeName.of(blockType(type));
}
return ArrayTypeName.of(vectorType(componentType));
return ArrayTypeName.of(vectorType(type));
}

@Override
Expand Down Expand Up @@ -76,7 +76,7 @@ public String factoryInvocation(MethodSpec.Builder factoryMethodBuilder) {

@Override
public void evalToBlock(MethodSpec.Builder builder) {
TypeName blockType = blockType(componentType);
TypeName blockType = blockType(type);
builder.addStatement("$T[] $LBlocks = new $T[$L.length]", blockType, name, blockType, name);
builder.beginControlFlow("try ($T $LRelease = $T.wrap($LBlocks))", RELEASABLE, name, RELEASABLES, name);
builder.beginControlFlow("for (int i = 0; i < $LBlocks.length; i++)", name);
Expand All @@ -92,20 +92,20 @@ public void closeEvalToBlock(MethodSpec.Builder builder) {
}

@Override
public void resolveVectors(MethodSpec.Builder builder, String invokeBlockEval) {
TypeName vectorType = vectorType(componentType);
public void resolveVectors(MethodSpec.Builder builder, String... invokeBlockEval) {
TypeName vectorType = vectorType(type);
builder.addStatement("$T[] $LVectors = new $T[$L.length]", vectorType, name, vectorType, name);
builder.beginControlFlow("for (int i = 0; i < $LBlocks.length; i++)", name);
builder.addStatement("$LVectors[i] = $LBlocks[i].asVector()", name, name);
builder.beginControlFlow("if ($LVectors[i] == null)", name).addStatement(invokeBlockEval).endControlFlow();
builder.beginControlFlow("if ($LVectors[i] == null)", name).addStatement(invokeBlockEval[0]).endControlFlow();
builder.endControlFlow();
}

@Override
public void createScratch(MethodSpec.Builder builder) {
builder.addStatement("$T[] $LValues = new $T[$L.length]", componentType, name, componentType, name);
if (componentType.equals(BYTES_REF)) {
builder.addStatement("$T[] $LScratch = new $T[$L.length]", componentType, name, componentType, name);
builder.addStatement("$T[] $LValues = new $T[$L.length]", type, name, type, name);
if (type.equals(BYTES_REF)) {
builder.addStatement("$T[] $LScratch = new $T[$L.length]", type, name, type, name);
builder.beginControlFlow("for (int i = 0; i < $L.length; i++)", name);
builder.addStatement("$LScratch[i] = new $T()", name, BYTES_REF);
builder.endControlFlow();
Expand Down Expand Up @@ -135,14 +135,19 @@ public void read(MethodSpec.Builder builder, boolean blockStyle) {
} else {
lookupVar = "p";
}
if (componentType.equals(BYTES_REF)) {
if (type.equals(BYTES_REF)) {
builder.addStatement("$LValues[i] = $L[i].getBytesRef($L, $LScratch[i])", name, paramName(blockStyle), lookupVar, name);
} else {
builder.addStatement("$LValues[i] = $L[i].$L($L)", name, paramName(blockStyle), getMethod(componentType), lookupVar);
builder.addStatement("$LValues[i] = $L[i].$L($L)", name, paramName(blockStyle), getMethod(type), lookupVar);
}
builder.endControlFlow();
}

@Override
public void read(MethodSpec.Builder builder, String accessor, String firstParam) {
// nothing to do
}

@Override
public void buildInvocation(StringBuilder pattern, List<Object> args, boolean blockStyle) {
pattern.append("$LValues");
Expand Down
7 changes: 6 additions & 1 deletion .../esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/argument/BlockArgument.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ public void closeEvalToBlock(MethodSpec.Builder builder) {
}

@Override
public void resolveVectors(MethodSpec.Builder builder, String invokeBlockEval) {
public void resolveVectors(MethodSpec.Builder builder, String... invokeBlockEval) {
// nothing to do
}

Expand Down Expand Up @@ -95,6 +95,11 @@ public void read(MethodSpec.Builder builder, boolean blockStyle) {
// nothing to do
}

@Override
public void read(MethodSpec.Builder builder, String accessor, String firstParam) {
// nothing to do
}

@Override
public void buildInvocation(StringBuilder pattern, List<Object> args, boolean blockStyle) {
pattern.append("$L");
Expand Down
Loading