Skip to content

[ESQL] Support first/last functions. #136419

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. #136419

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;

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);
}
}

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) {
builder.addComment("This type does not support vectors because all values are multi-valued");
Copy link
Contributor

Choose a reason for hiding this comment

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

Not related with this PR, but I think this should be an assertion? If for any reason code goes through this, it feels like an error.

However, now that you're not calling it, maybe we shouldn't create this method to begin with, if it won't be usable. Can we skip its creation?

If it adds extra complexity, it can be left for another PR

Copy link
Contributor

Choose a reason for hiding this comment

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

Btw, I don't know much about why these functions can't receive single-valued positions; @craigtaverner will know better if removing/asserting will change logics in any way

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");
}

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);

/**
* 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