Skip to content

Commit

Permalink
ESQL: Better tests to AUTO_BUCKET
Browse files Browse the repository at this point in the history
This improves the tests for AUTO_BUCKET marginally, specifically so that
it tests all valid combinations of arguments and generates a correct
types table. This'll combine nicely with elastic#106782 to generate the
signatures that kibana needs for it's editor.
  • Loading branch information
nik9000 committed Apr 8, 2024
1 parent 9edd67f commit 7dc4768
Show file tree
Hide file tree
Showing 7 changed files with 321 additions and 98 deletions.
37 changes: 36 additions & 1 deletion docs/reference/esql/functions/types/auto_bucket.asciidoc
Original file line number Diff line number Diff line change
Expand Up @@ -5,5 +5,40 @@
[%header.monospaced.styled,format=dsv,separator=|]
|===
field | buckets | from | to | result

datetime | integer | datetime | datetime | datetime
datetime | integer | datetime | keyword | datetime
datetime | integer | datetime | text | datetime
datetime | integer | keyword | datetime | datetime
datetime | integer | keyword | keyword | datetime
datetime | integer | keyword | text | datetime
datetime | integer | text | datetime | datetime
datetime | integer | text | keyword | datetime
datetime | integer | text | text | datetime
double | integer | double | double | double
double | integer | double | integer | double
double | integer | double | long | double
double | integer | integer | double | double
double | integer | integer | integer | double
double | integer | integer | long | double
double | integer | long | double | double
double | integer | long | integer | double
double | integer | long | long | double
integer | integer | double | double | double
integer | integer | double | integer | double
integer | integer | double | long | double
integer | integer | integer | double | double
integer | integer | integer | integer | double
integer | integer | integer | long | double
integer | integer | long | double | double
integer | integer | long | integer | double
integer | integer | long | long | double
long | integer | double | double | double
long | integer | double | integer | double
long | integer | double | long | double
long | integer | integer | double | double
long | integer | integer | integer | double
long | integer | integer | long | double
long | integer | long | double | double
long | integer | long | integer | double
long | integer | long | long | double
|===
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ synopsis:keyword
"double asin(number:double|integer|long|unsigned_long)"
"double atan(number:double|integer|long|unsigned_long)"
"double atan2(y_coordinate:double|integer|long|unsigned_long, x_coordinate:double|integer|long|unsigned_long)"
"double|date auto_bucket(field:integer|long|double|date, buckets:integer, from:integer|long|double|date|string, to:integer|long|double|date|string)"
"double|date auto_bucket(field:integer|long|double|date, buckets:integer, from:integer|long|double|date|keyword|text, to:integer|long|double|date|keyword|text)"
"double avg(number:double|integer|long)"
"boolean|cartesian_point|date|double|geo_point|integer|ip|keyword|long|text|unsigned_long|version case(condition:boolean, trueValue...:boolean|cartesian_point|date|double|geo_point|integer|ip|keyword|long|text|unsigned_long|version)"
"double|integer|long|unsigned_long ceil(number:double|integer|long|unsigned_long)"
Expand Down Expand Up @@ -117,7 +117,7 @@ acos |number |"double|integer|long|unsigne
asin |number |"double|integer|long|unsigned_long" |Number between -1 and 1. If `null`, the function returns `null`.
atan |number |"double|integer|long|unsigned_long" |Numeric expression. If `null`, the function returns `null`.
atan2 |[y_coordinate, x_coordinate] |["double|integer|long|unsigned_long", "double|integer|long|unsigned_long"] |[y coordinate. If `null`\, the function returns `null`., x coordinate. If `null`\, the function returns `null`.]
auto_bucket |[field, buckets, from, to] |["integer|long|double|date", integer, "integer|long|double|date|string", "integer|long|double|date|string"] |["", "", "", ""]
auto_bucket |[field, buckets, from, to] |["integer|long|double|date", integer, "integer|long|double|date|keyword|text", "integer|long|double|date|keyword|text"] |["", "", "", ""]
avg |number |"double|integer|long" |[""]
case |[condition, trueValue] |[boolean, "boolean|cartesian_point|date|double|geo_point|integer|ip|keyword|long|text|unsigned_long|version"] |["", ""]
ceil |number |"double|integer|long|unsigned_long" |Numeric expression. If `null`, the function returns `null`.
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,139 @@
// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License
// 2.0; you may not use this file except in compliance with the Elastic License
// 2.0.
package org.elasticsearch.xpack.esql.expression.function.scalar.string;

import java.lang.IllegalArgumentException;
import java.lang.Override;
import java.lang.String;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.compute.data.Block;
import org.elasticsearch.compute.data.BytesRefBlock;
import org.elasticsearch.compute.data.BytesRefVector;
import org.elasticsearch.compute.data.IntBlock;
import org.elasticsearch.compute.data.IntVector;
import org.elasticsearch.compute.data.Page;
import org.elasticsearch.compute.operator.DriverContext;
import org.elasticsearch.compute.operator.EvalOperator;
import org.elasticsearch.core.Releasables;
import org.elasticsearch.xpack.esql.expression.function.Warnings;
import org.elasticsearch.xpack.ql.tree.Source;

/**
* {@link EvalOperator.ExpressionEvaluator} implementation for {@link Locate}.
* This class is generated. Do not edit it.
*/
public final class LocateNoStartEvaluator implements EvalOperator.ExpressionEvaluator {
private final Warnings warnings;

private final EvalOperator.ExpressionEvaluator str;

private final EvalOperator.ExpressionEvaluator substr;

private final DriverContext driverContext;

public LocateNoStartEvaluator(Source source, EvalOperator.ExpressionEvaluator str,
EvalOperator.ExpressionEvaluator substr, DriverContext driverContext) {
this.warnings = new Warnings(source);
this.str = str;
this.substr = substr;
this.driverContext = driverContext;
}

@Override
public Block eval(Page page) {
try (BytesRefBlock strBlock = (BytesRefBlock) str.eval(page)) {
try (BytesRefBlock substrBlock = (BytesRefBlock) substr.eval(page)) {
BytesRefVector strVector = strBlock.asVector();
if (strVector == null) {
return eval(page.getPositionCount(), strBlock, substrBlock);
}
BytesRefVector substrVector = substrBlock.asVector();
if (substrVector == null) {
return eval(page.getPositionCount(), strBlock, substrBlock);
}
return eval(page.getPositionCount(), strVector, substrVector).asBlock();
}
}
}

public IntBlock eval(int positionCount, BytesRefBlock strBlock, BytesRefBlock substrBlock) {
try(IntBlock.Builder result = driverContext.blockFactory().newIntBlockBuilder(positionCount)) {
BytesRef strScratch = new BytesRef();
BytesRef substrScratch = new BytesRef();
position: for (int p = 0; p < positionCount; p++) {
if (strBlock.isNull(p)) {
result.appendNull();
continue position;
}
if (strBlock.getValueCount(p) != 1) {
if (strBlock.getValueCount(p) > 1) {
warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value"));
}
result.appendNull();
continue position;
}
if (substrBlock.isNull(p)) {
result.appendNull();
continue position;
}
if (substrBlock.getValueCount(p) != 1) {
if (substrBlock.getValueCount(p) > 1) {
warnings.registerException(new IllegalArgumentException("single-value function encountered multi-value"));
}
result.appendNull();
continue position;
}
result.appendInt(Locate.process(strBlock.getBytesRef(strBlock.getFirstValueIndex(p), strScratch), substrBlock.getBytesRef(substrBlock.getFirstValueIndex(p), substrScratch)));
}
return result.build();
}
}

public IntVector eval(int positionCount, BytesRefVector strVector, BytesRefVector substrVector) {
try(IntVector.Builder result = driverContext.blockFactory().newIntVectorBuilder(positionCount)) {
BytesRef strScratch = new BytesRef();
BytesRef substrScratch = new BytesRef();
position: for (int p = 0; p < positionCount; p++) {
result.appendInt(Locate.process(strVector.getBytesRef(p, strScratch), substrVector.getBytesRef(p, substrScratch)));
}
return result.build();
}
}

@Override
public String toString() {
return "LocateNoStartEvaluator[" + "str=" + str + ", substr=" + substr + "]";
}

@Override
public void close() {
Releasables.closeExpectNoException(str, substr);
}

static class Factory implements EvalOperator.ExpressionEvaluator.Factory {
private final Source source;

private final EvalOperator.ExpressionEvaluator.Factory str;

private final EvalOperator.ExpressionEvaluator.Factory substr;

public Factory(Source source, EvalOperator.ExpressionEvaluator.Factory str,
EvalOperator.ExpressionEvaluator.Factory substr) {
this.source = source;
this.str = str;
this.substr = substr;
}

@Override
public LocateNoStartEvaluator get(DriverContext context) {
return new LocateNoStartEvaluator(source, str.get(context), substr.get(context), context);
}

@Override
public String toString() {
return "LocateNoStartEvaluator[" + "str=" + str + ", substr=" + substr + "]";
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ public AutoBucket(
Source source,
@Param(name = "field", type = { "integer", "long", "double", "date" }) Expression field,
@Param(name = "buckets", type = { "integer" }) Expression buckets,
@Param(name = "from", type = { "integer", "long", "double", "date", "string" }) Expression from,
@Param(name = "to", type = { "integer", "long", "double", "date", "string" }) Expression to
@Param(name = "from", type = { "integer", "long", "double", "date", "keyword", "text" }) Expression from,
@Param(name = "to", type = { "integer", "long", "double", "date", "keyword", "text" }) Expression to
) {
super(source, List.of(field, buckets, from, to));
this.field = field;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,10 @@ public static ExpressionEvaluator.Factory evaluator(Expression e) {
}
Layout.Builder builder = new Layout.Builder();
buildLayout(builder, e);
assertTrue(e.resolved());
Expression.TypeResolution resolution = e.typeResolved();
if (resolution.unresolved()) {
throw new AssertionError("expected resolved " + resolution.message());
}
return EvalMapper.toEvaluator(e, builder.build());
}

Expand Down Expand Up @@ -263,7 +266,10 @@ private void testEvaluate(boolean readFloating) {
}
return;
}
assertFalse("expected resolved", expression.typeResolved().unresolved());
Expression.TypeResolution resolution = expression.typeResolved();
if (resolution.unresolved()) {
throw new AssertionError("expected resolved " + resolution.message());
}
expression = new FoldNull().rule(expression);
assertThat(expression.dataType(), equalTo(testCase.expectedType()));
logger.info("Result type: " + expression.dataType());
Expand Down Expand Up @@ -651,6 +657,28 @@ public static void testFunctionInfo() {
* on input types like {@link Greatest} or {@link Coalesce}.
*/
protected static List<TestCaseSupplier> anyNullIsNull(boolean entirelyNullPreservesType, List<TestCaseSupplier> testCaseSuppliers) {
return anyNullIsNull(
testCaseSuppliers,
(nullPosition, nullValueDataType, original) -> entirelyNullPreservesType == false
&& nullValueDataType == DataTypes.NULL
&& original.getData().size() == 1 ? DataTypes.NULL : original.expectedType(),
(nullPosition, original) -> original
);
}

public interface ExpectedType {
DataType expectedType(int nullPosition, DataType nullValueDataType, TestCaseSupplier.TestCase original);
}

public interface ExpectedEvaluatorToString {
Matcher<String> evaluatorToString(int nullPosition, Matcher<String> original);
}

protected static List<TestCaseSupplier> anyNullIsNull(
List<TestCaseSupplier> testCaseSuppliers,
ExpectedType expectedType,
ExpectedEvaluatorToString evaluatorToString
) {
typesRequired(testCaseSuppliers);
List<TestCaseSupplier> suppliers = new ArrayList<>(testCaseSuppliers.size());
suppliers.addAll(testCaseSuppliers);
Expand All @@ -673,15 +701,12 @@ protected static List<TestCaseSupplier> anyNullIsNull(boolean entirelyNullPreser
TestCaseSupplier.TestCase oc = original.get();
List<TestCaseSupplier.TypedData> data = IntStream.range(0, oc.getData().size()).mapToObj(i -> {
TestCaseSupplier.TypedData od = oc.getData().get(i);
if (i == finalNullPosition) {
return new TestCaseSupplier.TypedData(null, od.type(), od.name());
}
return od;
return i == finalNullPosition ? od.forceValueToNull() : od;
}).toList();
return new TestCaseSupplier.TestCase(
data,
oc.evaluatorToString(),
oc.expectedType(),
evaluatorToString.evaluatorToString(finalNullPosition, oc.evaluatorToString()),
expectedType.expectedType(finalNullPosition, oc.getData().get(finalNullPosition).type(), oc),
nullValue(),
null,
oc.getExpectedTypeError(),
Expand All @@ -704,7 +729,7 @@ protected static List<TestCaseSupplier> anyNullIsNull(boolean entirelyNullPreser
return new TestCaseSupplier.TestCase(
data,
equalTo("LiteralsEvaluator[lit=null]"),
entirelyNullPreservesType == false && oc.getData().size() == 1 ? DataTypes.NULL : oc.expectedType(),
expectedType.expectedType(finalNullPosition, DataTypes.NULL, oc),
nullValue(),
null,
oc.getExpectedTypeError(),
Expand Down Expand Up @@ -810,9 +835,8 @@ private static Stream<List<DataType>> allPermutations(int argumentCount) {
if (argumentCount == 0) {
return Stream.of(List.of());
}
if (argumentCount > 4) {
// TODO check for a limit 4. is arbitrary.
throw new IllegalArgumentException("would generate too many types");
if (argumentCount > 3) {
throw new IllegalArgumentException("would generate too many combinations");
}
Stream<List<DataType>> stream = representable().map(t -> List.of(t));
for (int i = 1; i < argumentCount; i++) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1325,6 +1325,14 @@ public TypedData forceLiteral() {
return new TypedData(data, type, name, true);
}

/**
* Return a {@link TypedData} that always returns {@code null} for it's
* value without modifying anything else in the supplier.
*/
public TypedData forceValueToNull() {
return new TypedData(null, type, name, forceLiteral);
}

@Override
public String toString() {
if (type == DataTypes.UNSIGNED_LONG && data instanceof Long longData) {
Expand Down
Loading

0 comments on commit 7dc4768

Please sign in to comment.