From edb6c0314107c96957d56c6216dddb6a7d200339 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Mon, 13 Jan 2025 15:15:53 -0500 Subject: [PATCH 1/8] ESQL: Speed COALESCE --- .../compute/operator/EvalBenchmark.java | 80 +++++++++++++++++- .../function/scalar/nulls/Coalesce.java | 82 ++++++++++++++++--- .../function/scalar/nulls/CoalesceTests.java | 60 ++++++++++++++ 3 files changed, 208 insertions(+), 14 deletions(-) diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java index 9aab4a3e3210f..6c49311e1148d 100644 --- a/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java +++ b/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java @@ -37,6 +37,7 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc; import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs; import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin; +import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce; import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike; import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add; import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals; @@ -96,6 +97,9 @@ public class EvalBenchmark { "add_double", "case_1_eager", "case_1_lazy", + "coalesce_1_noop", + "coalesce_1_eager", + "coalesce_1_lazy", "date_trunc", "equal_to_const", "long_equal_to_long", @@ -140,7 +144,32 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) { lhs = new Add(Source.EMPTY, lhs, new Literal(Source.EMPTY, 1L, DataType.LONG)); rhs = new Add(Source.EMPTY, rhs, new Literal(Source.EMPTY, 1L, DataType.LONG)); } - yield EvalMapper.toEvaluator(new Case(Source.EMPTY, condition, List.of(lhs, rhs)), layout(f1, f2)).get(driverContext); + EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator( + new Case(Source.EMPTY, condition, List.of(lhs, rhs)), + layout(f1, f2) + ).get(driverContext); + String desc = operation.endsWith("lazy") ? "CaseLazyEvaluator" : "CaseEagerEvaluator"; + if (evaluator.toString().contains(desc) == false) { + throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]"); + } + yield evaluator; + } + case "coalesce_1_noop", "coalesce_1_eager", "coalesce_1_lazy" -> { + FieldAttribute f1 = longField(); + FieldAttribute f2 = longField(); + Expression lhs = f1; + if (operation.endsWith("lazy")) { + lhs = new Add(Source.EMPTY, lhs, new Literal(Source.EMPTY, 1L, DataType.LONG)); + } + EvalOperator.ExpressionEvaluator evaluator = EvalMapper.toEvaluator( + new Coalesce(Source.EMPTY, lhs, List.of(f2)), + layout(f1, f2) + ).get(driverContext); + String desc = operation.endsWith("lazy") ? "CoalesceEvaluator" : "CoalesceEvaluator"; // NOCOMMIT eager/lazy here + if (evaluator.toString().contains(desc) == false) { + throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]"); + } + yield evaluator; } case "date_trunc" -> { FieldAttribute timestamp = new FieldAttribute( @@ -255,6 +284,38 @@ private static void checkExpected(String operation, Page actual) { } } } + case "coalesce_1_noop" -> { + LongVector f1 = actual.getBlock(0).asVector(); + LongVector result = actual.getBlock(2).asVector(); + for (int i = 0; i < BLOCK_LENGTH; i++) { + long expected = f1.getLong(i); + if (result.getLong(i) != expected) { + throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]"); + } + } + } + case "coalesce_1_eager" -> { + LongBlock f1 = actual.getBlock(0); + LongVector f2 = actual.getBlock(1).asVector(); + LongVector result = actual.getBlock(2).asVector(); + for (int i = 0; i < BLOCK_LENGTH; i++) { + long expected = i % 5 == 0 ? f2.getLong(i) : f1.getLong(f1.getFirstValueIndex(i)); + if (result.getLong(i) != expected) { + throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]"); + } + } + } + case "coalesce_1_lazy" -> { + LongBlock f1 = actual.getBlock(0); + LongVector f2 = actual.getBlock(1).asVector(); + LongVector result = actual.getBlock(2).asVector(); + for (int i = 0; i < BLOCK_LENGTH; i++) { + long expected = i % 5 == 0 ? f2.getLong(i) : f1.getLong(f1.getFirstValueIndex(i)) + 1; + if (result.getLong(i) != expected) { + throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]"); + } + } + } case "date_trunc" -> { LongVector v = actual.getBlock(1).asVector(); long oneDay = TimeValue.timeValueHours(24).millis(); @@ -299,7 +360,7 @@ private static void checkExpected(String operation, Page actual) { } } } - default -> throw new UnsupportedOperationException(); + default -> throw new UnsupportedOperationException(operation); } } @@ -319,7 +380,7 @@ private static Page page(String operation) { } yield new Page(builder.build()); } - case "case_1_eager", "case_1_lazy" -> { + case "case_1_eager", "case_1_lazy", "coalesce_1_noop" -> { var f1 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); var f2 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); for (int i = 0; i < BLOCK_LENGTH; i++) { @@ -328,6 +389,19 @@ private static Page page(String operation) { } yield new Page(f1.build(), f2.build()); } + case "coalesce_1_eager", "coalesce_1_lazy" -> { + var f1 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); + var f2 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); + for (int i = 0; i < BLOCK_LENGTH; i++) { + if (i % 5 == 0) { + f1.appendNull(); + } else { + f1.appendLong(i); + } + f2.appendLong(-i); + } + yield new Page(f1.build(), f2.build()); + } case "long_equal_to_long" -> { var lhs = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); var rhs = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index 52686430ca5b5..0de302ff9766b 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -11,6 +11,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; @@ -35,6 +36,8 @@ import java.io.IOException; import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -217,26 +220,83 @@ private record CoalesceEvaluator(DriverContext driverContext, ElementType result EvalOperator.ExpressionEvaluator { @Override public Block eval(Page page) { - /* - * We have to evaluate lazily so any errors or warnings that would be - * produced by the right hand side are avoided. And so if anything - * on the right hand side is slow we skip it. - * - * And it'd be good if that lazy evaluation were fast. But this - * implementation isn't. It's fairly simple - running position at - * a time - but it's not at all fast. - */ + return blockAtATime(page); + } + + /** + * Evaluate COALESCE block-at-a-time for as long as we can, then shift to + * position-at-a-time. + *

+ * Block-at-a-time evaluation is the "normal" way to run the compute engine, + * just calling {@link ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain ExpressionEvaluator#eval} and: + *

+ *
    + *
  • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
    • + *
  • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
    • + *
  • If this is the last evaluator we just return it. COALESCE done.
    • + *
  • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a block-at-a-time evaluator. + *
    • + *
+ */ + private Block blockAtATime(Page page) { + int lastFullBlockIdx = 0; + while (true) { + Block lastFullBlock = evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + try { + // The result has some nulls and some non-nulls. + return positionAtATime(page, lastFullBlockIdx, lastFullBlock); + } finally { + lastFullBlock.close(); + } + } + } + + /** + * Evaluate COALESCE position-at-a-time. We have a block that contains + * some nulls and some non-nulls and we have some remaining evaluators. + * For each position we either: + *
    + *
  • Take the non-null values from the {@code lastFullBlock}
    • + *
  • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
    • + *
+ *

+ * It's important that we're evaluating position-by-position because + * the evaluators may produce warnings, and we really don't want those + * warnings to sneak into the output if they are for values that + * aren't needed. This lazy evaluation is not fast. Not at all. But + * it's very important not to leak the warnings. + *

+ */ + private Block positionAtATime(Page page, int firstToEvaluate, Block lastFullBlock) { int positionCount = page.getPositionCount(); try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } int[] positions = new int[] { p }; Page limited = new Page( 1, IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) ); try (Releasable ignored = limited::releaseBlocks) { - for (EvalOperator.ExpressionEvaluator eval : evaluators) { - try (Block block = eval.eval(limited)) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (Block block = evaluators.get(e).eval(limited)) { if (false == block.isNull(0)) { result.copyFrom(block, 0, 1); continue position; diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java index 797c99992815e..3363972285a7a 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java @@ -12,8 +12,15 @@ import org.elasticsearch.common.network.NetworkAddress; import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BlockUtils; +import org.elasticsearch.compute.data.ElementType; 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.index.mapper.TestBlock; +import org.elasticsearch.xpack.esql.EsqlTestUtils; +import org.elasticsearch.xpack.esql.TestBlockFactory; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; import org.elasticsearch.xpack.esql.core.expression.Literal; @@ -28,17 +35,24 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.VaragsTestCaseBuilder; import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.SpatialRelatesFunctionTestCase; import org.elasticsearch.xpack.esql.planner.Layout; +import org.elasticsearch.xpack.esql.planner.PlannerUtils; import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter; import java.time.ZonedDateTime; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.IntStream; import static org.elasticsearch.compute.data.BlockUtils.toJavaObject; import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral; +import static org.hamcrest.Matchers.either; import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.sameInstance; public class CoalesceTests extends AbstractScalarFunctionTestCase { public CoalesceTests(@Name("TestCase") Supplier testCaseSupplier) { @@ -228,4 +242,50 @@ public void testCoalesceNotNullable() { // Known not to be nullable because it contains a non-null literal assertThat(exp.nullable(), equalTo(Nullability.FALSE)); } + + /** + * Inserts random non-null garbage around the expected data + * and runs + */ + public void testEvaluateWithGarbage() { + DriverContext context = driverContext(); + Expression expression = randomBoolean() ? buildDeepCopyOfFieldExpression(testCase) : buildFieldExpression(testCase); + int positions = between(2, 1024); + List data = testCase.getData(); + Page onePositionPage = row(testCase.getDataValues()); + Block[] manyPositionsBlocks = new Block[Math.toIntExact(data.stream().filter(d -> d.isForceLiteral() == false).count())]; + int realPosition = between(0, positions - 1); + try { + int b = 0; + for (TestCaseSupplier.TypedData d : data) { + ElementType elementType = PlannerUtils.toElementType(d.type()); + try (Block.Builder builder = elementType.newBlockBuilder(positions, context.blockFactory())) { + for (int p = 0; p < positions; p++) { + if (p == realPosition) { + builder.copyFrom(onePositionPage.getBlock(b), 0, 1); + } else { + builder.copyFrom( + BlockUtils.constantBlock(TestBlockFactory.getNonBreakingInstance(), randomLiteral(d.type()).value(), 1), + 0, + 1 + ); + } + } + manyPositionsBlocks[b] = builder.build(); + } + b++; + } + try ( + EvalOperator.ExpressionEvaluator eval = evaluator(expression).get(context); + Block block = eval.eval(new Page(positions, manyPositionsBlocks)) + ) { + assertThat(block.getPositionCount(), is(positions)); + assertThat(toJavaObjectUnsignedLongAware(block, realPosition), testCase.getMatcher()); + assertThat("evaluates to tracked block", block.blockFactory(), sameInstance(context.blockFactory())); + } + } finally { + Releasables.close(onePositionPage::releaseBlocks, Releasables.wrap(manyPositionsBlocks)); + } + } + } From 8bd0989eace9f814d98627c02d79726db948ea8f Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Mon, 13 Jan 2025 17:48:34 -0500 Subject: [PATCH 2/8] WIP --- benchmarks/README.md | 5 +- .../compute/operator/EvalBenchmark.java | 20 +- .../compute/data/BooleanBlockBuilder.java | 6 - .../compute/data/BytesRefBlockBuilder.java | 6 - .../compute/data/DoubleBlockBuilder.java | 6 - .../compute/data/FloatBlockBuilder.java | 6 - .../compute/data/IntBlockBuilder.java | 6 - .../compute/data/X-BlockBuilder.java.st | 2 + .../function/scalar/nulls/Coalesce.java | 186 +++++++++++++----- .../function/scalar/nulls/CoalesceTests.java | 14 +- 10 files changed, 156 insertions(+), 101 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index d7b324acfef81..687577bd9672c 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -126,9 +126,12 @@ exit Grab the async profiler from https://github.com/jvm-profiling-tools/async-profiler and run `prof async` like so: ``` -gradlew -p benchmarks/ run --args 'LongKeyedBucketOrdsBenchmark.multiBucket -prof "async:libPath=/home/nik9000/Downloads/tmp/async-profiler-1.8.3-linux-x64/build/libasyncProfiler.so;dir=/tmp/prof;output=flamegraph"' +gradlew -p benchmarks/ run --args 'LongKeyedBucketOrdsBenchmark.multiBucket -prof "async:libPath=/home/nik9000/Downloads/async-profiler-3.0-29ee888-linux-x64/lib/libasyncProfiler.so;dir=/tmp/prof;output=flamegraph"' ``` +Note: As of January 2025 the latest release of async profiler doesn't work + with out JDK but the nightly is fine. + If you are on Mac, this'll warn you that you downloaded the shared library from the internet. You'll need to go to settings and allow it to run. diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java index 6c49311e1148d..0b4aac138a92a 100644 --- a/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java +++ b/benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java @@ -97,9 +97,9 @@ public class EvalBenchmark { "add_double", "case_1_eager", "case_1_lazy", - "coalesce_1_noop", - "coalesce_1_eager", - "coalesce_1_lazy", + "coalesce_2_noop", + "coalesce_2_eager", + "coalesce_2_lazy", "date_trunc", "equal_to_const", "long_equal_to_long", @@ -154,7 +154,7 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) { } yield evaluator; } - case "coalesce_1_noop", "coalesce_1_eager", "coalesce_1_lazy" -> { + case "coalesce_2_noop", "coalesce_2_eager", "coalesce_2_lazy" -> { FieldAttribute f1 = longField(); FieldAttribute f2 = longField(); Expression lhs = f1; @@ -165,7 +165,7 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) { new Coalesce(Source.EMPTY, lhs, List.of(f2)), layout(f1, f2) ).get(driverContext); - String desc = operation.endsWith("lazy") ? "CoalesceEvaluator" : "CoalesceEvaluator"; // NOCOMMIT eager/lazy here + String desc = operation.endsWith("lazy") ? "CoalesceLazyEvaluator" : "CoalesceEagerEvaluator"; if (evaluator.toString().contains(desc) == false) { throw new IllegalArgumentException("Evaluator was [" + evaluator + "] but expected one containing [" + desc + "]"); } @@ -284,7 +284,7 @@ private static void checkExpected(String operation, Page actual) { } } } - case "coalesce_1_noop" -> { + case "coalesce_2_noop" -> { LongVector f1 = actual.getBlock(0).asVector(); LongVector result = actual.getBlock(2).asVector(); for (int i = 0; i < BLOCK_LENGTH; i++) { @@ -294,7 +294,7 @@ private static void checkExpected(String operation, Page actual) { } } } - case "coalesce_1_eager" -> { + case "coalesce_2_eager" -> { LongBlock f1 = actual.getBlock(0); LongVector f2 = actual.getBlock(1).asVector(); LongVector result = actual.getBlock(2).asVector(); @@ -305,7 +305,7 @@ private static void checkExpected(String operation, Page actual) { } } } - case "coalesce_1_lazy" -> { + case "coalesce_2_lazy" -> { LongBlock f1 = actual.getBlock(0); LongVector f2 = actual.getBlock(1).asVector(); LongVector result = actual.getBlock(2).asVector(); @@ -380,7 +380,7 @@ private static Page page(String operation) { } yield new Page(builder.build()); } - case "case_1_eager", "case_1_lazy", "coalesce_1_noop" -> { + case "case_1_eager", "case_1_lazy", "coalesce_2_noop" -> { var f1 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); var f2 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); for (int i = 0; i < BLOCK_LENGTH; i++) { @@ -389,7 +389,7 @@ private static Page page(String operation) { } yield new Page(f1.build(), f2.build()); } - case "coalesce_1_eager", "coalesce_1_lazy" -> { + case "coalesce_2_eager", "coalesce_2_lazy" -> { var f1 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); var f2 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH); for (int i = 0; i < BLOCK_LENGTH; i++) { diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java index 32627a0e0d36b..e42429195013d 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java @@ -73,12 +73,6 @@ public BooleanBlockBuilder endPositionEntry() { @Override public BooleanBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { - if (block.areAllValuesNull()) { - for (int p = beginInclusive; p < endExclusive; p++) { - appendNull(); - } - return this; - } return copyFrom((BooleanBlock) block, beginInclusive, endExclusive); } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java index 6232cbdd2717c..871dbdb0fdff9 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java @@ -76,12 +76,6 @@ protected void writeNullValue() { @Override public BytesRefBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { - if (block.areAllValuesNull()) { - for (int p = beginInclusive; p < endExclusive; p++) { - appendNull(); - } - return this; - } return copyFrom((BytesRefBlock) block, beginInclusive, endExclusive); } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java index 5921c2daa9f92..a5c4811d54d4b 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java @@ -73,12 +73,6 @@ public DoubleBlockBuilder endPositionEntry() { @Override public DoubleBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { - if (block.areAllValuesNull()) { - for (int p = beginInclusive; p < endExclusive; p++) { - appendNull(); - } - return this; - } return copyFrom((DoubleBlock) block, beginInclusive, endExclusive); } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java index 9c1e7aba49a21..6e110f4d8d5fd 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java @@ -73,12 +73,6 @@ public FloatBlockBuilder endPositionEntry() { @Override public FloatBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { - if (block.areAllValuesNull()) { - for (int p = beginInclusive; p < endExclusive; p++) { - appendNull(); - } - return this; - } return copyFrom((FloatBlock) block, beginInclusive, endExclusive); } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java index 85f943004de29..c2ce7348954b5 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java @@ -73,12 +73,6 @@ public IntBlockBuilder endPositionEntry() { @Override public IntBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { - if (block.areAllValuesNull()) { - for (int p = beginInclusive; p < endExclusive; p++) { - appendNull(); - } - return this; - } return copyFrom((IntBlock) block, beginInclusive, endExclusive); } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st index 8397a0f5274f1..b61a8f40d3b05 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st @@ -120,6 +120,8 @@ $endif$ return copyFrom(($Type$Block) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index 0de302ff9766b..010518cf3d03b 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -11,7 +11,6 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.compute.data.Block; -import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; @@ -36,15 +35,14 @@ import java.io.IOException; import java.util.List; -import java.util.Set; -import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; import static org.elasticsearch.xpack.esql.core.type.DataType.NULL; /** - * Function returning the first non-null value. + * Function returning the first non-null value. {@code COALESCE} runs as though + * it were lazily evaluating each position in each incoming {@link Block}. */ public class Coalesce extends EsqlScalarFunction implements OptionalArgument { public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Coalesce", Coalesce::new); @@ -198,10 +196,27 @@ public boolean foldable() { @Override public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) { List childEvaluators = children().stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + PlannerUtils.toElementType(dataType()), + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEvaluator( + return new CoalesceLazyEvaluator( context, PlannerUtils.toElementType(dataType()), childEvaluators.stream().map(x -> x.get(context)).toList() @@ -210,24 +225,33 @@ public ExpressionEvaluator get(DriverContext context) { @Override public String toString() { - return "CoalesceEvaluator[values=" + childEvaluators + ']'; + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; } }; } - private record CoalesceEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) - implements - EvalOperator.ExpressionEvaluator { + protected abstract static sealed class AbstractCoalesceEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceEagerEvaluator, CoalesceLazyEvaluator { + protected final DriverContext driverContext; + protected final ElementType resultType; + protected final List evaluators; + + protected AbstractCoalesceEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { + this.driverContext = driverContext; + this.resultType = resultType; + this.evaluators = evaluators; + } + @Override - public Block eval(Page page) { - return blockAtATime(page); + public final Block eval(Page page) { + return entireBlock(page); } /** - * Evaluate COALESCE block-at-a-time for as long as we can, then shift to - * position-at-a-time. + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. *

- * Block-at-a-time evaluation is the "normal" way to run the compute engine, + * Entire Block evaluation is the "normal" way to run the compute engine, * just calling {@link ExpressionEvaluator#eval}. It's much faster so we try * that first. For each evaluator, we {@linkplain ExpressionEvaluator#eval} and: *

@@ -237,51 +261,115 @@ public Block eval(Page page) { *
  • If this is the last evaluator we just return it. COALESCE done.
    • *
  • * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop - * into a block-at-a-time evaluator. + * into a per position evaluator. *
    • * */ - private Block blockAtATime(Page page) { + private Block entireBlock(Page page) { int lastFullBlockIdx = 0; while (true) { Block lastFullBlock = evaluators.get(lastFullBlockIdx++).eval(page); if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { return lastFullBlock; } - if (lastFullBlock.areAllValuesNull()) { - // Result is all nulls and isn't the last result so we don't need any of it. - lastFullBlock.close(); - continue; - } - try { - // The result has some nulls and some non-nulls. - return positionAtATime(page, lastFullBlockIdx, lastFullBlock); - } finally { - lastFullBlock.close(); - } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); } } /** - * Evaluate COALESCE position-at-a-time. We have a block that contains - * some nulls and some non-nulls and we have some remaining evaluators. - * For each position we either: - *
      - *
    • Take the non-null values from the {@code lastFullBlock}
      • - *
    • - * Evaluator the remaining evaluators one at a time, keeping - * the first non-null value. - *
      • - *
    + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. *

    - * It's important that we're evaluating position-by-position because - * the evaluators may produce warnings, and we really don't want those - * warnings to sneak into the output if they are for values that - * aren't needed. This lazy evaluation is not fast. Not at all. But - * it's very important not to leak the warnings. + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. *

    */ - private Block positionAtATime(Page page, int firstToEvaluate, Block lastFullBlock) { + protected abstract Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link ExpressionEvaluator.Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + private static final class CoalesceEagerEvaluator extends AbstractCoalesceEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { + super(driverContext, resultType, evaluators); + } + + @Override + protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + Block[] flatten = new Block[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { + position: for (int p = 0; p < positionCount; p++) { + for (Block f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p, p + 1); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + private static final class CoalesceLazyEvaluator extends AbstractCoalesceEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { + super(driverContext, resultType, evaluators); + } + + @Override + protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { int positionCount = page.getPositionCount(); try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { position: for (int p = 0; p < positionCount; p++) { @@ -307,17 +395,9 @@ private Block positionAtATime(Page page, int firstToEvaluate, Block lastFullBloc } } return result.build(); + } finally { + lastFullBlock.close(); } } - - @Override - public String toString() { - return "CoalesceEvaluator[values=" + evaluators + ']'; - } - - @Override - public void close() { - Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); - } } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java index 3363972285a7a..7baa382d625ac 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java @@ -62,7 +62,7 @@ public CoalesceTests(@Name("TestCase") Supplier testC @ParametersFactory public static Iterable parameters() { List noNullsSuppliers = new ArrayList<>(); - VaragsTestCaseBuilder builder = new VaragsTestCaseBuilder(type -> "Coalesce"); + VaragsTestCaseBuilder builder = new VaragsTestCaseBuilder(type -> "CoalesceEager"); builder.expectString(strings -> strings.filter(v -> v != null).findFirst()); builder.expectLong(longs -> longs.filter(v -> v != null).findFirst()); builder.expectInt(ints -> ints.filter(v -> v != null).findFirst()); @@ -77,7 +77,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(first, DataType.IP, "first"), new TestCaseSupplier.TypedData(second, DataType.IP, "second") ), - "CoalesceEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.IP, equalTo(first == null ? second : first) ); @@ -92,7 +92,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(first, DataType.VERSION, "first"), new TestCaseSupplier.TypedData(second, DataType.VERSION, "second") ), - "CoalesceEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.VERSION, equalTo(first == null ? second : first) ); @@ -105,7 +105,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(firstDate, DataType.DATETIME, "first"), new TestCaseSupplier.TypedData(secondDate, DataType.DATETIME, "second") ), - "CoalesceEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.DATETIME, equalTo(firstDate == null ? secondDate : firstDate) ); @@ -118,12 +118,13 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(firstDate, DataType.DATE_NANOS, "first"), new TestCaseSupplier.TypedData(secondDate, DataType.DATE_NANOS, "second") ), - "CoalesceEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.DATE_NANOS, equalTo(firstDate == null ? secondDate : firstDate) ); })); + List suppliers = new ArrayList<>(noNullsSuppliers); for (TestCaseSupplier s : noNullsSuppliers) { for (int nullUpTo = 1; nullUpTo < s.types().size(); nullUpTo++) { @@ -180,7 +181,7 @@ protected static void addSpatialCombinations(List suppliers) { TestCaseSupplier.testCaseSupplier( leftDataSupplier, rightDataSupplier, - (l, r) -> equalTo("CoalesceEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]"), + (l, r) -> equalTo("CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]"), dataType, (l, r) -> l ) @@ -287,5 +288,4 @@ public void testEvaluateWithGarbage() { Releasables.close(onePositionPage::releaseBlocks, Releasables.wrap(manyPositionsBlocks)); } } - } From 759d05236933094745e7c77d6953c4029fb40647 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 14 Jan 2025 05:52:22 -0500 Subject: [PATCH 3/8] Gotta make it nicer --- .../compute/data/BooleanBlockBuilder.java | 8 + .../compute/data/BytesRefBlockBuilder.java | 8 + .../compute/data/DoubleBlockBuilder.java | 8 + .../compute/data/FloatBlockBuilder.java | 8 + .../compute/data/IntBlockBuilder.java | 8 + .../compute/data/LongBlockBuilder.java | 2 + .../function/scalar/nulls/Coalesce.java | 146 +++++++++--------- 7 files changed, 115 insertions(+), 73 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java index e42429195013d..d38c29b864661 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java @@ -73,9 +73,17 @@ public BooleanBlockBuilder endPositionEntry() { @Override public BooleanBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { + if (block.areAllValuesNull()) { + for (int p = beginInclusive; p < endExclusive; p++) { + appendNull(); + } + return this; + } return copyFrom((BooleanBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java index 871dbdb0fdff9..d0f930f35d947 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java @@ -76,9 +76,17 @@ protected void writeNullValue() { @Override public BytesRefBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { + if (block.areAllValuesNull()) { + for (int p = beginInclusive; p < endExclusive; p++) { + appendNull(); + } + return this; + } return copyFrom((BytesRefBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java index a5c4811d54d4b..d95cf79c314d1 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java @@ -73,9 +73,17 @@ public DoubleBlockBuilder endPositionEntry() { @Override public DoubleBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { + if (block.areAllValuesNull()) { + for (int p = beginInclusive; p < endExclusive; p++) { + appendNull(); + } + return this; + } return copyFrom((DoubleBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java index 6e110f4d8d5fd..891af2ff8922d 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java @@ -73,9 +73,17 @@ public FloatBlockBuilder endPositionEntry() { @Override public FloatBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { + if (block.areAllValuesNull()) { + for (int p = beginInclusive; p < endExclusive; p++) { + appendNull(); + } + return this; + } return copyFrom((FloatBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java index c2ce7348954b5..3c56d7722a943 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java @@ -73,9 +73,17 @@ public IntBlockBuilder endPositionEntry() { @Override public IntBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusive) { + if (block.areAllValuesNull()) { + for (int p = beginInclusive; p < endExclusive; p++) { + appendNull(); + } + return this; + } return copyFrom((IntBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java index d24ae214da63a..c592191ffb384 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java @@ -82,6 +82,8 @@ public LongBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusi return copyFrom((LongBlock) block, beginInclusive, endExclusive); } + // NOCOMMIT it's slow to check if all values are null for single position copies. + /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index 010518cf3d03b..b52224543f5b6 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -309,94 +309,94 @@ public final String toString() { public final void close() { Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); } - } - /** - * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. - * First we evaluate all remaining evaluators, and then we pluck the first non-null - * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the - * evaluators make them so we only use it for evaluators that are - * {@link ExpressionEvaluator.Factory#eagerEvalSafeInLazy safe} to evaluate eagerly - * in a lazy environment. - */ - private static final class CoalesceEagerEvaluator extends AbstractCoalesceEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { - super(driverContext, resultType, evaluators); - } + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link ExpressionEvaluator.Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + private static final class CoalesceEagerEvaluator extends AbstractCoalesceEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { + super(driverContext, resultType, evaluators); + } - @Override - protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { - int positionCount = page.getPositionCount(); - Block[] flatten = new Block[evaluators.size() - firstToEvaluate + 1]; - try { - flatten[0] = lastFullBlock; - for (int f = 1; f < flatten.length; f++) { - flatten[f] = evaluators.get(firstToEvaluate + f - 1).eval(page); - } - try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { - position: for (int p = 0; p < positionCount; p++) { - for (Block f : flatten) { - if (false == f.isNull(p)) { - result.copyFrom(f, p, p + 1); - continue position; + @Override + protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + Block[] flatten = new Block[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { + position: for (int p = 0; p < positionCount; p++) { + for (Block f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p, p + 1); + continue position; + } } + result.appendNull(); } - result.appendNull(); + return result.build(); } - return result.build(); + } finally { + Releasables.close(flatten); } - } finally { - Releasables.close(flatten); } } - } - /** - * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. - * For each position we either: - *
      - *
    • Take the non-null values from the {@code lastFullBlock}
      • - *
    • - * Evaluator the remaining evaluators one at a time, keeping - * the first non-null value. - *
      • - *
    - */ - private static final class CoalesceLazyEvaluator extends AbstractCoalesceEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { - super(driverContext, resultType, evaluators); - } + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + private static final class CoalesceLazyEvaluator extends AbstractCoalesceEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { + super(driverContext, resultType, evaluators); + } - @Override - protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { - int positionCount = page.getPositionCount(); - try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { - position: for (int p = 0; p < positionCount; p++) { - if (lastFullBlock.isNull(p) == false) { - result.copyFrom(lastFullBlock, p, p + 1); - continue; - } - int[] positions = new int[] { p }; - Page limited = new Page( - 1, - IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) - ); - try (Releasable ignored = limited::releaseBlocks) { - for (int e = firstToEvaluate; e < evaluators.size(); e++) { - try (Block block = evaluators.get(e).eval(limited)) { - if (false == block.isNull(0)) { - result.copyFrom(block, 0, 1); - continue position; + @Override + protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (Block block = evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0, 1); + continue position; + } } } + result.appendNull(); } - result.appendNull(); } + return result.build(); + } finally { + lastFullBlock.close(); } - return result.build(); - } finally { - lastFullBlock.close(); } } } From f28f6af9daf730090a27ec38f5af56a5844ca7ad Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 14 Jan 2025 10:24:15 -0500 Subject: [PATCH 4/8] COALESCE woo --- x-pack/plugin/esql/build.gradle | 27 ++ .../compute/data/BooleanBlock.java | 5 + .../compute/data/BooleanBlockBuilder.java | 51 ++-- .../compute/data/BytesRefBlock.java | 5 + .../compute/data/BytesRefBlockBuilder.java | 50 ++-- .../compute/data/DoubleBlock.java | 5 + .../compute/data/DoubleBlockBuilder.java | 51 ++-- .../compute/data/FloatBlock.java | 5 + .../compute/data/FloatBlockBuilder.java | 51 ++-- .../elasticsearch/compute/data/IntBlock.java | 5 + .../compute/data/IntBlockBuilder.java | 51 ++-- .../elasticsearch/compute/data/LongBlock.java | 5 + .../compute/data/LongBlockBuilder.java | 51 ++-- .../org/elasticsearch/compute/data/Block.java | 4 + .../compute/data/X-Block.java.st | 5 + .../compute/data/X-BlockBuilder.java.st | 55 ++-- .../data/BlockBuilderCopyFromTests.java | 12 +- .../nulls/CoalesceBooleanEvaluator.java | 232 +++++++++++++++++ .../nulls/CoalesceBytesRefEvaluator.java | 234 +++++++++++++++++ .../scalar/nulls/CoalesceDoubleEvaluator.java | 232 +++++++++++++++++ .../scalar/nulls/CoalesceIntEvaluator.java | 232 +++++++++++++++++ .../scalar/nulls/CoalesceLongEvaluator.java | 232 +++++++++++++++++ .../function/scalar/conditional/Case.java | 2 + .../function/scalar/nulls/Coalesce.java | 236 +++-------------- .../scalar/nulls/X-CoalesceEvaluator.java.st | 238 ++++++++++++++++++ 25 files changed, 1750 insertions(+), 326 deletions(-) create mode 100644 x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st diff --git a/x-pack/plugin/esql/build.gradle b/x-pack/plugin/esql/build.gradle index 9cbcd8d49b586..51a83a749fbe3 100644 --- a/x-pack/plugin/esql/build.gradle +++ b/x-pack/plugin/esql/build.gradle @@ -332,4 +332,31 @@ tasks.named('stringTemplates').configure { it.inputFile = inInputFile it.outputFile = "org/elasticsearch/xpack/esql/expression/predicate/operator/comparison/InBytesRefEvaluator.java" } + + File coalesceInputFile = file("src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st") + template { + it.properties = booleanProperties + it.inputFile = coalesceInputFile + it.outputFile = "org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java" + } + template { + it.properties = intProperties + it.inputFile = coalesceInputFile + it.outputFile = "org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java" + } + template { + it.properties = longProperties + it.inputFile = coalesceInputFile + it.outputFile = "org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java" + } + template { + it.properties = doubleProperties + it.inputFile = coalesceInputFile + it.outputFile = "org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java" + } + template { + it.properties = bytesRefProperties + it.inputFile = coalesceInputFile + it.outputFile = "org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java" + } } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java index 5d2d6c97a11f1..0faa1c5104ce3 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java @@ -223,6 +223,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.BooleanBuilder permi */ Builder copyFrom(BooleanBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(BooleanBlock block, int position); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java index d38c29b864661..f9319791ad76a 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.data; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.BitArray; @@ -82,12 +83,14 @@ public BooleanBlockBuilder copyFrom(Block block, int beginInclusive, int endExcl return copyFrom((BooleanBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(BooleanBlock, int)}. + *

    */ + @Override public BooleanBlockBuilder copyFrom(BooleanBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -103,21 +106,7 @@ public BooleanBlockBuilder copyFrom(BooleanBlock block, int beginInclusive, int private void copyFromBlock(BooleanBlock block, int beginInclusive, int endExclusive) { for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendBoolean(block.getBoolean(i++)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p); } } @@ -127,6 +116,34 @@ private void copyFromVector(BooleanVector vector, int beginInclusive, int endExc } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public BooleanBlockBuilder copyFrom(BooleanBlock block, int position) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendBoolean(block.getBoolean(i++)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendBoolean(block.getBoolean(i++)); + } + endPositionEntry(); + return this; + } + @Override public BooleanBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java index 6fe45f33a7df6..5b236b74cb6b3 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java @@ -228,6 +228,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.BytesRefBuilder perm */ Builder copyFrom(BytesRefBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(BytesRefBlock block, int position, BytesRef scratch); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java index d0f930f35d947..b5b549bb6d6b1 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java @@ -85,12 +85,14 @@ public BytesRefBlockBuilder copyFrom(Block block, int beginInclusive, int endExc return copyFrom((BytesRefBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(BytesRefBlock, int, BytesRef scratch)}. + *

    */ + @Override public BytesRefBlockBuilder copyFrom(BytesRefBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -107,21 +109,7 @@ public BytesRefBlockBuilder copyFrom(BytesRefBlock block, int beginInclusive, in private void copyFromBlock(BytesRefBlock block, int beginInclusive, int endExclusive) { BytesRef scratch = new BytesRef(); for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendBytesRef(block.getBytesRef(i++, scratch)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p, scratch); } } @@ -132,6 +120,34 @@ private void copyFromVector(BytesRefVector vector, int beginInclusive, int endEx } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public BytesRefBlockBuilder copyFrom(BytesRefBlock block, int position, BytesRef scratch) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendBytesRef(block.getBytesRef(i++, scratch)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendBytesRef(block.getBytesRef(i++, scratch)); + } + endPositionEntry(); + return this; + } + @Override public BytesRefBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java index 395ccd412fabb..4cc89d7439dd7 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java @@ -217,6 +217,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.DoubleBuilder permit */ Builder copyFrom(DoubleBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(DoubleBlock block, int position); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java index d95cf79c314d1..e7eaea166a9a6 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.data; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.DoubleArray; @@ -82,12 +83,14 @@ public DoubleBlockBuilder copyFrom(Block block, int beginInclusive, int endExclu return copyFrom((DoubleBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(DoubleBlock, int)}. + *

    */ + @Override public DoubleBlockBuilder copyFrom(DoubleBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -103,21 +106,7 @@ public DoubleBlockBuilder copyFrom(DoubleBlock block, int beginInclusive, int en private void copyFromBlock(DoubleBlock block, int beginInclusive, int endExclusive) { for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendDouble(block.getDouble(i++)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p); } } @@ -127,6 +116,34 @@ private void copyFromVector(DoubleVector vector, int beginInclusive, int endExcl } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public DoubleBlockBuilder copyFrom(DoubleBlock block, int position) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendDouble(block.getDouble(i++)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendDouble(block.getDouble(i++)); + } + endPositionEntry(); + return this; + } + @Override public DoubleBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java index 633c9f309901a..6888c1badde9c 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java @@ -216,6 +216,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.FloatBuilder permits */ Builder copyFrom(FloatBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(FloatBlock block, int position); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java index 891af2ff8922d..a2e1e17aacd5e 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.data; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.FloatArray; @@ -82,12 +83,14 @@ public FloatBlockBuilder copyFrom(Block block, int beginInclusive, int endExclus return copyFrom((FloatBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(FloatBlock, int)}. + *

    */ + @Override public FloatBlockBuilder copyFrom(FloatBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -103,21 +106,7 @@ public FloatBlockBuilder copyFrom(FloatBlock block, int beginInclusive, int endE private void copyFromBlock(FloatBlock block, int beginInclusive, int endExclusive) { for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendFloat(block.getFloat(i++)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p); } } @@ -127,6 +116,34 @@ private void copyFromVector(FloatVector vector, int beginInclusive, int endExclu } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public FloatBlockBuilder copyFrom(FloatBlock block, int position) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendFloat(block.getFloat(i++)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendFloat(block.getFloat(i++)); + } + endPositionEntry(); + return this; + } + @Override public FloatBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java index 7c77d9965391e..025c1f0d964fe 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java @@ -216,6 +216,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.IntBuilder permits I */ Builder copyFrom(IntBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(IntBlock block, int position); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java index 3c56d7722a943..d08e55b1ca0ae 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.data; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.IntArray; @@ -82,12 +83,14 @@ public IntBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusiv return copyFrom((IntBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(IntBlock, int)}. + *

    */ + @Override public IntBlockBuilder copyFrom(IntBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -103,21 +106,7 @@ public IntBlockBuilder copyFrom(IntBlock block, int beginInclusive, int endExclu private void copyFromBlock(IntBlock block, int beginInclusive, int endExclusive) { for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendInt(block.getInt(i++)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p); } } @@ -127,6 +116,34 @@ private void copyFromVector(IntVector vector, int beginInclusive, int endExclusi } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public IntBlockBuilder copyFrom(IntBlock block, int position) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendInt(block.getInt(i++)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendInt(block.getInt(i++)); + } + endPositionEntry(); + return this; + } + @Override public IntBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java index 6c88da8860ca7..f71da8fad860f 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java @@ -217,6 +217,11 @@ sealed interface Builder extends Block.Builder, BlockLoader.LongBuilder permits */ Builder copyFrom(LongBlock block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom(LongBlock block, int position); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java index c592191ffb384..1717a55c0a391 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java @@ -7,6 +7,7 @@ package org.elasticsearch.compute.data; +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.LongArray; @@ -82,12 +83,14 @@ public LongBlockBuilder copyFrom(Block block, int beginInclusive, int endExclusi return copyFrom((LongBlock) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom(LongBlock, int)}. + *

    */ + @Override public LongBlockBuilder copyFrom(LongBlock block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -103,21 +106,7 @@ public LongBlockBuilder copyFrom(LongBlock block, int beginInclusive, int endExc private void copyFromBlock(LongBlock block, int beginInclusive, int endExclusive) { for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { - appendLong(block.getLong(i++)); - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p); } } @@ -127,6 +116,34 @@ private void copyFromVector(LongVector vector, int beginInclusive, int endExclus } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public LongBlockBuilder copyFrom(LongBlock block, int position) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + appendLong(block.getLong(i++)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + appendLong(block.getLong(i++)); + } + endPositionEntry(); + return this; + } + @Override public LongBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java index edf54a829deba..41f7e85dfc08f 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java @@ -280,6 +280,10 @@ interface Builder extends BlockLoader.Builder, Releasable { /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single position copies see {@link IntBlockBuilder#copyFrom(IntBlock, int)}, + * {@link LongBlockBuilder#copyFrom(LongBlock, int)}, etc. + *

    */ Builder copyFrom(Block block, int beginInclusive, int endExclusive); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st index 67e4ac4bb334f..326f7f4d05bb8 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st @@ -288,6 +288,11 @@ $endif$ */ Builder copyFrom($Type$Block block, int beginInclusive, int endExclusive); + /** + * Copy the values in {@code block} at {@code position}. + */ + Builder copyFrom($Type$Block block, int position$if(BytesRef)$, BytesRef scratch$endif$); + @Override Builder appendNull(); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st index b61a8f40d3b05..3064bd6a2725d 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st @@ -15,6 +15,7 @@ import org.elasticsearch.common.util.BytesRefArray; import org.elasticsearch.core.Releasables; $else$ +import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.RamUsageEstimator; import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.util.$Array$; @@ -120,12 +121,14 @@ $endif$ return copyFrom(($Type$Block) block, beginInclusive, endExclusive); } - // NOCOMMIT it's slow to check if all values are null for single position copies. - /** * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. + *

    + * For single-position copies see {@link #copyFrom($Type$Block, int$if(BytesRef)$, BytesRef scratch$endif$)}. + *

    */ + @Override public $Type$BlockBuilder copyFrom($Type$Block block, int beginInclusive, int endExclusive) { if (endExclusive > block.getPositionCount()) { throw new IllegalArgumentException("can't copy past the end [" + endExclusive + " > " + block.getPositionCount() + "]"); @@ -144,25 +147,7 @@ $if(BytesRef)$ BytesRef scratch = new BytesRef(); $endif$ for (int p = beginInclusive; p < endExclusive; p++) { - if (block.isNull(p)) { - appendNull(); - continue; - } - int count = block.getValueCount(p); - if (count > 1) { - beginPositionEntry(); - } - int i = block.getFirstValueIndex(p); - for (int v = 0; v < count; v++) { -$if(BytesRef)$ - appendBytesRef(block.getBytesRef(i++, scratch)); -$else$ - append$Type$(block.get$Type$(i++)); -$endif$ - } - if (count > 1) { - endPositionEntry(); - } + copyFrom(block, p$if(BytesRef)$, scratch$endif$); } } @@ -179,6 +164,34 @@ $endif$ } } + /** + * Copy the values in {@code block} at {@code position}. + *

    + * Note that there isn't a version of this method on {@link Block.Builder} that takes + * {@link Block}. That'd be quite slow, running position by position. And it's important + * to know if you are copying {@link BytesRef}s so you can have the scratch. + *

    + */ + @Override + public $Type$BlockBuilder copyFrom($Type$Block block, int position$if(BytesRef)$, BytesRef scratch$endif$) { + if (block.isNull(position)) { + appendNull(); + return this; + } + int count = block.getValueCount(position); + int i = block.getFirstValueIndex(position); + if (count == 1) { + append$Type$(block.get$Type$(i++$if(BytesRef)$, scratch$endif$)); + return this; + } + beginPositionEntry(); + for (int v = 0; v < count; v++) { + append$Type$(block.get$Type$(i++$if(BytesRef)$, scratch$endif$)); + } + endPositionEntry(); + return this; + } + @Override public $Type$BlockBuilder mvOrdering(Block.MvOrdering mvOrdering) { this.mvOrdering = mvOrdering; diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java index dc12a78954c5e..b89fc2b0232db 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java @@ -10,6 +10,7 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import org.apache.lucene.util.BytesRef; import org.elasticsearch.test.ESTestCase; import java.util.ArrayList; @@ -89,7 +90,16 @@ private void assertEvens(Block block) { Block.Builder builder = elementType.newBlockBuilder(block.getPositionCount() / 2, blockFactory); List> expected = new ArrayList<>(); for (int i = 0; i < block.getPositionCount(); i += 2) { - builder.copyFrom(block, i, i + 1); + switch (elementType) { + case BOOLEAN -> ((BooleanBlockBuilder) builder).copyFrom((BooleanBlock) block, i); + case BYTES_REF -> ((BytesRefBlockBuilder) builder).copyFrom((BytesRefBlock) block, i, new BytesRef()); + case DOUBLE -> ((DoubleBlockBuilder) builder).copyFrom((DoubleBlock) block, i); + case FLOAT -> ((FloatBlockBuilder) builder).copyFrom((FloatBlock) block, i); + case INT -> ((IntBlockBuilder) builder).copyFrom((IntBlock) block, i); + case LONG -> ((LongBlockBuilder) builder).copyFrom((LongBlock) block, i); + default -> throw new IllegalArgumentException(); + } + expected.add(BasicBlockTests.valuesAtPositions(block, i, i + 1).get(0)); } assertBlockValues(builder.build(), expected); diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java new file mode 100644 index 0000000000000..2196ea959d0ee --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java @@ -0,0 +1,232 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BooleanBlock; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class CoalesceBooleanEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceBooleanEvaluator.CoalesceEagerEvaluator, // + CoalesceBooleanEvaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected CoalesceBooleanEvaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final BooleanBlock eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private BooleanBlock entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + BooleanBlock lastFullBlock = (BooleanBlock) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract BooleanBlock perPosition(Page page, BooleanBlock lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends CoalesceBooleanEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected BooleanBlock perPosition(Page page, BooleanBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + BooleanBlock[] flatten = new BooleanBlock[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = (BooleanBlock) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (BooleanBlock.Builder result = driverContext.blockFactory().newBooleanBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for (BooleanBlock f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends CoalesceBooleanEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected BooleanBlock perPosition(Page page, BooleanBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + try (BooleanBlock.Builder result = driverContext.blockFactory().newBooleanBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (BooleanBlock block = (BooleanBlock) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java new file mode 100644 index 0000000000000..682b05111cc07 --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java @@ -0,0 +1,234 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class CoalesceBytesRefEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceBytesRefEvaluator.CoalesceEagerEvaluator, // + CoalesceBytesRefEvaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected CoalesceBytesRefEvaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final BytesRefBlock eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private BytesRefBlock entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + BytesRefBlock lastFullBlock = (BytesRefBlock) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract BytesRefBlock perPosition(Page page, BytesRefBlock lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends CoalesceBytesRefEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected BytesRefBlock perPosition(Page page, BytesRefBlock lastFullBlock, int firstToEvaluate) { + BytesRef scratch = new BytesRef(); + int positionCount = page.getPositionCount(); + BytesRefBlock[] flatten = new BytesRefBlock[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = (BytesRefBlock) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (BytesRefBlock.Builder result = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for (BytesRefBlock f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p, scratch); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends CoalesceBytesRefEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected BytesRefBlock perPosition(Page page, BytesRefBlock lastFullBlock, int firstToEvaluate) { + BytesRef scratch = new BytesRef(); + int positionCount = page.getPositionCount(); + try (BytesRefBlock.Builder result = driverContext.blockFactory().newBytesRefBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (BytesRefBlock block = (BytesRefBlock) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0, scratch); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java new file mode 100644 index 0000000000000..708ee324adfd0 --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java @@ -0,0 +1,232 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.DoubleBlock; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class CoalesceDoubleEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceDoubleEvaluator.CoalesceEagerEvaluator, // + CoalesceDoubleEvaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected CoalesceDoubleEvaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final DoubleBlock eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private DoubleBlock entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + DoubleBlock lastFullBlock = (DoubleBlock) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract DoubleBlock perPosition(Page page, DoubleBlock lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends CoalesceDoubleEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected DoubleBlock perPosition(Page page, DoubleBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + DoubleBlock[] flatten = new DoubleBlock[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = (DoubleBlock) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (DoubleBlock.Builder result = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for (DoubleBlock f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends CoalesceDoubleEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected DoubleBlock perPosition(Page page, DoubleBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + try (DoubleBlock.Builder result = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (DoubleBlock block = (DoubleBlock) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java new file mode 100644 index 0000000000000..071f9d883022f --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java @@ -0,0 +1,232 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.IntBlock; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class CoalesceIntEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceIntEvaluator.CoalesceEagerEvaluator, // + CoalesceIntEvaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected CoalesceIntEvaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final IntBlock eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private IntBlock entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + IntBlock lastFullBlock = (IntBlock) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract IntBlock perPosition(Page page, IntBlock lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends CoalesceIntEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected IntBlock perPosition(Page page, IntBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + IntBlock[] flatten = new IntBlock[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = (IntBlock) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (IntBlock.Builder result = driverContext.blockFactory().newIntBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for (IntBlock f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends CoalesceIntEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected IntBlock perPosition(Page page, IntBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + try (IntBlock.Builder result = driverContext.blockFactory().newIntBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (IntBlock block = (IntBlock) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java new file mode 100644 index 0000000000000..07159a9c0f698 --- /dev/null +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java @@ -0,0 +1,232 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class CoalesceLongEvaluator implements EvalOperator.ExpressionEvaluator permits + CoalesceLongEvaluator.CoalesceEagerEvaluator, // + CoalesceLongEvaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected CoalesceLongEvaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final LongBlock eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private LongBlock entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + LongBlock lastFullBlock = (LongBlock) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract LongBlock perPosition(Page page, LongBlock lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends CoalesceLongEvaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected LongBlock perPosition(Page page, LongBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + LongBlock[] flatten = new LongBlock[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = (LongBlock) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try (LongBlock.Builder result = driverContext.blockFactory().newLongBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for (LongBlock f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends CoalesceLongEvaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected LongBlock perPosition(Page page, LongBlock lastFullBlock, int firstToEvaluate) { + int positionCount = page.getPositionCount(); + try (LongBlock.Builder result = driverContext.blockFactory().newLongBlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try (LongBlock block = (LongBlock) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java index 824f02ca7ccbb..4623a9a20d901 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java @@ -522,6 +522,8 @@ public Block eval(Page page) { ) { for (int p = 0; p < lhs.getPositionCount(); p++) { if (lhsOrRhs.mask().getBoolean(p)) { + // This is quite slow compared to using the single valued copy and unrolling per-type like coalesce + // *especially* when copying a Block. Vectors are pretty fast this way. builder.copyFrom(lhs, p, p + 1); } else { builder.copyFrom(rhs, p, p + 1); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index b52224543f5b6..8cf830691d19e 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -22,6 +22,7 @@ import org.elasticsearch.xpack.esql.core.expression.Expressions; import org.elasticsearch.xpack.esql.core.expression.Nullability; import org.elasticsearch.xpack.esql.core.expression.TypeResolutions; +import org.elasticsearch.xpack.esql.core.expression.function.Function; import org.elasticsearch.xpack.esql.core.tree.NodeInfo; import org.elasticsearch.xpack.esql.core.tree.Source; import org.elasticsearch.xpack.esql.core.type.DataType; @@ -34,11 +35,26 @@ import org.elasticsearch.xpack.esql.planner.PlannerUtils; import java.io.IOException; +import java.util.EnumMap; import java.util.List; +import java.util.Map; +import java.util.function.BiFunction; import java.util.stream.IntStream; import java.util.stream.Stream; +import static org.elasticsearch.xpack.esql.core.type.DataType.BOOLEAN; +import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_POINT; +import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_SHAPE; +import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME; +import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS; +import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_POINT; +import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_SHAPE; +import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER; +import static org.elasticsearch.xpack.esql.core.type.DataType.IP; +import static org.elasticsearch.xpack.esql.core.type.DataType.KEYWORD; +import static org.elasticsearch.xpack.esql.core.type.DataType.LONG; import static org.elasticsearch.xpack.esql.core.type.DataType.NULL; +import static org.elasticsearch.xpack.esql.core.type.DataType.VERSION; /** * Function returning the first non-null value. {@code COALESCE} runs as though @@ -195,209 +211,25 @@ public boolean foldable() { @Override public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) { - List childEvaluators = children().stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - PlannerUtils.toElementType(dataType()), - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } - - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - PlannerUtils.toElementType(dataType()), - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } - - @Override - public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; - } - }; + return SUPPORTED.get(dataType()).toEvaluator(toEvaluator, children()); } - protected abstract static sealed class AbstractCoalesceEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceEagerEvaluator, CoalesceLazyEvaluator { - protected final DriverContext driverContext; - protected final ElementType resultType; - protected final List evaluators; - - protected AbstractCoalesceEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { - this.driverContext = driverContext; - this.resultType = resultType; - this.evaluators = evaluators; - } - - @Override - public final Block eval(Page page) { - return entireBlock(page); - } - - /** - * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to - * {@link #perPosition} evaluation. - *

    - * Entire Block evaluation is the "normal" way to run the compute engine, - * just calling {@link ExpressionEvaluator#eval}. It's much faster so we try - * that first. For each evaluator, we {@linkplain ExpressionEvaluator#eval} and: - *

    - *
      - *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • - *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • - *
    • If this is the last evaluator we just return it. COALESCE done.
      • - *
    • - * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop - * into a per position evaluator. - *
      • - *
    - */ - private Block entireBlock(Page page) { - int lastFullBlockIdx = 0; - while (true) { - Block lastFullBlock = evaluators.get(lastFullBlockIdx++).eval(page); - if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { - return lastFullBlock; - } - if (lastFullBlock.areAllValuesNull()) { - // Result is all nulls and isn't the last result so we don't need any of it. - lastFullBlock.close(); - continue; - } - // The result has some nulls and some non-nulls. - return perPosition(page, lastFullBlock, lastFullBlockIdx); - } - } - - /** - * Evaluate each position of the incoming {@link Page} for COALESCE - * independently. Our attempt to evaluate entire blocks has yielded - * a block that contains some nulls and some non-nulls and we have - * to fill in the nulls with the results of calling the remaining - * evaluators. - *

    - * This must not return warnings caused by - * evaluating positions for which a previous evaluator returned - * non-null. These are positions that, at least from the perspective - * of a compute engine user, don't have to be - * evaluated. Put another way, this must function as though - * {@code COALESCE} were per-position lazy. It can manage that - * any way it likes. - *

    - */ - protected abstract Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate); - - @Override - public final String toString() { - return getClass().getSimpleName() + "[values=" + evaluators + ']'; - } - - @Override - public final void close() { - Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); - } - - /** - * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. - * First we evaluate all remaining evaluators, and then we pluck the first non-null - * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the - * evaluators make them so we only use it for evaluators that are - * {@link ExpressionEvaluator.Factory#eagerEvalSafeInLazy safe} to evaluate eagerly - * in a lazy environment. - */ - private static final class CoalesceEagerEvaluator extends AbstractCoalesceEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { - super(driverContext, resultType, evaluators); - } - - @Override - protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { - int positionCount = page.getPositionCount(); - Block[] flatten = new Block[evaluators.size() - firstToEvaluate + 1]; - try { - flatten[0] = lastFullBlock; - for (int f = 1; f < flatten.length; f++) { - flatten[f] = evaluators.get(firstToEvaluate + f - 1).eval(page); - } - try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { - position: for (int p = 0; p < positionCount; p++) { - for (Block f : flatten) { - if (false == f.isNull(p)) { - result.copyFrom(f, p, p + 1); - continue position; - } - } - result.appendNull(); - } - return result.build(); - } - } finally { - Releasables.close(flatten); - } - } - } - - /** - * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. - * For each position we either: - *
      - *
    • Take the non-null values from the {@code lastFullBlock}
      • - *
    • - * Evaluator the remaining evaluators one at a time, keeping - * the first non-null value. - *
      • - *
    - */ - private static final class CoalesceLazyEvaluator extends AbstractCoalesceEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, ElementType resultType, List evaluators) { - super(driverContext, resultType, evaluators); - } - - @Override - protected Block perPosition(Page page, Block lastFullBlock, int firstToEvaluate) { - int positionCount = page.getPositionCount(); - try (Block.Builder result = resultType.newBlockBuilder(positionCount, driverContext.blockFactory())) { - position: for (int p = 0; p < positionCount; p++) { - if (lastFullBlock.isNull(p) == false) { - result.copyFrom(lastFullBlock, p, p + 1); - continue; - } - int[] positions = new int[] { p }; - Page limited = new Page( - 1, - IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) - ); - try (Releasable ignored = limited::releaseBlocks) { - for (int e = firstToEvaluate; e < evaluators.size(); e++) { - try (Block block = evaluators.get(e).eval(limited)) { - if (false == block.isNull(0)) { - result.copyFrom(block, 0, 1); - continue position; - } - } - } - result.appendNull(); - } - } - return result.build(); - } finally { - lastFullBlock.close(); - } - } - } + interface BuildCoalesce { + ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator, List children); } + private static final Map SUPPORTED = Map.ofEntries( + // We intend to support all types here. + Map.entry(BOOLEAN, CoalesceBooleanEvaluator::toEvaluator), + Map.entry(CARTESIAN_POINT, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(CARTESIAN_SHAPE, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(DATE_NANOS, CoalesceLongEvaluator::toEvaluator), + Map.entry(DATETIME, CoalesceLongEvaluator::toEvaluator), + Map.entry(GEO_POINT, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(GEO_SHAPE, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(INTEGER, CoalesceIntEvaluator::toEvaluator), + Map.entry(IP, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(KEYWORD, CoalesceBytesRefEvaluator::toEvaluator), + Map.entry(LONG, CoalesceLongEvaluator::toEvaluator), + Map.entry(VERSION, CoalesceBytesRefEvaluator::toEvaluator) + ); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st new file mode 100644 index 0000000000000..c4ee70a15af0f --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st @@ -0,0 +1,238 @@ +/* + * 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.nulls; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.$Type$Block; +import org.elasticsearch.compute.data.ElementType; +import org.elasticsearch.compute.data.Page; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.core.Releasable; +import org.elasticsearch.core.Releasables; +import org.elasticsearch.xpack.esql.core.expression.Expression; +import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link Coalesce}. + * This class is generated. Edit {@code X-InEvaluator.java.st} instead. + */ +abstract sealed class Coalesce$Type$Evaluator implements EvalOperator.ExpressionEvaluator permits + Coalesce$Type$Evaluator.CoalesceEagerEvaluator, // + Coalesce$Type$Evaluator.CoalesceLazyEvaluator { + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceEagerEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; + } + }; + } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLazyEvaluator( + context, + childEvaluators.stream().map(x -> x.get(context)).toList() + ); + } + + @Override + public String toString() { + return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } + + protected final DriverContext driverContext; + protected final List evaluators; + + protected Coalesce$Type$Evaluator(DriverContext driverContext, List evaluators) { + this.driverContext = driverContext; + this.evaluators = evaluators; + } + + @Override + public final $Type$Block eval(Page page) { + return entireBlock(page); + } + + /** + * Evaluate COALESCE for an entire {@link Block} for as long as we can, then shift to + * {@link #perPosition} evaluation. + *

    + * Entire Block evaluation is the "normal" way to run the compute engine, + * just calling {@link EvalOperator.ExpressionEvaluator#eval}. It's much faster so we try + * that first. For each evaluator, we {@linkplain EvalOperator.ExpressionEvaluator#eval} and: + *

    + *
      + *
    • If the {@linkplain Block} doesn't have any nulls we return it. COALESCE done.
      • + *
    • If the {@linkplain Block} is only nulls we skip it and try the next evaluator.
      • + *
    • If this is the last evaluator we just return it. COALESCE done.
      • + *
    • + * Otherwise, the {@linkplain Block} has mixed nulls and non-nulls so we drop + * into a per position evaluator. + *
      • + *
    + */ + private $Type$Block entireBlock(Page page) { + int lastFullBlockIdx = 0; + while (true) { + $Type$Block lastFullBlock = ($Type$Block) evaluators.get(lastFullBlockIdx++).eval(page); + if (lastFullBlockIdx == evaluators.size() || lastFullBlock.asVector() != null) { + return lastFullBlock; + } + if (lastFullBlock.areAllValuesNull()) { + // Result is all nulls and isn't the last result so we don't need any of it. + lastFullBlock.close(); + continue; + } + // The result has some nulls and some non-nulls. + return perPosition(page, lastFullBlock, lastFullBlockIdx); + } + } + + /** + * Evaluate each position of the incoming {@link Page} for COALESCE + * independently. Our attempt to evaluate entire blocks has yielded + * a block that contains some nulls and some non-nulls and we have + * to fill in the nulls with the results of calling the remaining + * evaluators. + *

    + * This must not return warnings caused by + * evaluating positions for which a previous evaluator returned + * non-null. These are positions that, at least from the perspective + * of a compute engine user, don't have to be + * evaluated. Put another way, this must function as though + * {@code COALESCE} were per-position lazy. It can manage that + * any way it likes. + *

    + */ + protected abstract $Type$Block perPosition(Page page, $Type$Block lastFullBlock, int firstToEvaluate); + + @Override + public final String toString() { + return getClass().getSimpleName() + "[values=" + evaluators + ']'; + } + + @Override + public final void close() { + Releasables.closeExpectNoException(() -> Releasables.close(evaluators)); + } + + /** + * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. + * First we evaluate all remaining evaluators, and then we pluck the first non-null + * value from each one. This is much faster than + * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * evaluators make them so we only use it for evaluators that are + * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly + * in a lazy environment. + */ + static final class CoalesceEagerEvaluator extends Coalesce$Type$Evaluator { + CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected $Type$Block perPosition(Page page, $Type$Block lastFullBlock, int firstToEvaluate) { +$if(BytesRef)$ + BytesRef scratch = new BytesRef(); +$endif$ + int positionCount = page.getPositionCount(); + $Type$Block[] flatten = new $Type$Block[evaluators.size() - firstToEvaluate + 1]; + try { + flatten[0] = lastFullBlock; + for (int f = 1; f < flatten.length; f++) { + flatten[f] = ($Type$Block) evaluators.get(firstToEvaluate + f - 1).eval(page); + } + try ($Type$Block.Builder result = driverContext.blockFactory().new$Type$BlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + for ($Type$Block f : flatten) { + if (false == f.isNull(p)) { + result.copyFrom(f, p$if(BytesRef)$, scratch$endif$); + continue position; + } + } + result.appendNull(); + } + return result.build(); + } + } finally { + Releasables.close(flatten); + } + } + } + + /** + * Evaluates {@code COALESCE} lazily per position if entire-block evaluation fails. + * For each position we either: + *
      + *
    • Take the non-null values from the {@code lastFullBlock}
      • + *
    • + * Evaluator the remaining evaluators one at a time, keeping + * the first non-null value. + *
      • + *
    + */ + static final class CoalesceLazyEvaluator extends Coalesce$Type$Evaluator { + CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + super(driverContext, evaluators); + } + + @Override + protected $Type$Block perPosition(Page page, $Type$Block lastFullBlock, int firstToEvaluate) { +$if(BytesRef)$ + BytesRef scratch = new BytesRef(); +$endif$ + int positionCount = page.getPositionCount(); + try ($Type$Block.Builder result = driverContext.blockFactory().new$Type$BlockBuilder(positionCount)) { + position: for (int p = 0; p < positionCount; p++) { + if (lastFullBlock.isNull(p) == false) { + result.copyFrom(lastFullBlock, p, p + 1); + continue; + } + int[] positions = new int[] { p }; + Page limited = new Page( + 1, + IntStream.range(0, page.getBlockCount()).mapToObj(b -> page.getBlock(b).filter(positions)).toArray(Block[]::new) + ); + try (Releasable ignored = limited::releaseBlocks) { + for (int e = firstToEvaluate; e < evaluators.size(); e++) { + try ($Type$Block block = ($Type$Block) evaluators.get(e).eval(limited)) { + if (false == block.isNull(0)) { + result.copyFrom(block, 0$if(BytesRef)$, scratch$endif$); + continue position; + } + } + } + result.appendNull(); + } + } + return result.build(); + } finally { + lastFullBlock.close(); + } + } + } +} From 40794d0dc9311d68fa14d359516c6f4fb70a9104 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 14 Jan 2025 11:18:22 -0500 Subject: [PATCH 5/8] Formatting --- .../nulls/CoalesceBooleanEvaluator.java | 55 ++++++++---------- .../nulls/CoalesceBytesRefEvaluator.java | 54 ++++++++---------- .../scalar/nulls/CoalesceDoubleEvaluator.java | 55 ++++++++---------- .../scalar/nulls/CoalesceIntEvaluator.java | 55 ++++++++---------- .../scalar/nulls/CoalesceLongEvaluator.java | 55 ++++++++---------- .../scalar/nulls/X-CoalesceEvaluator.java.st | 56 +++++++++---------- 6 files changed, 146 insertions(+), 184 deletions(-) diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java index 2196ea959d0ee..97b4cba0d9938 100644 --- a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBooleanEvaluator.java @@ -7,10 +7,8 @@ package org.elasticsearch.xpack.esql.expression.function.scalar.nulls; -import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BooleanBlock; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +26,36 @@ * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class CoalesceBooleanEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceBooleanEvaluator.CoalesceEagerEvaluator, // - CoalesceBooleanEvaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + CoalesceBooleanEvaluator.CoalesceBooleanEagerEvaluator, // + CoalesceBooleanEvaluator.CoalesceBooleanLazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new CoalesceBooleanEagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "CoalesceBooleanEagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceBooleanLazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "CoalesceBooleanLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +137,13 @@ public final void close() { * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link CoalesceBooleanLazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends CoalesceBooleanEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceBooleanEagerEvaluator extends CoalesceBooleanEvaluator { + CoalesceBooleanEagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -192,8 +185,8 @@ protected BooleanBlock perPosition(Page page, BooleanBlock lastFullBlock, int fi * * */ - static final class CoalesceLazyEvaluator extends CoalesceBooleanEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceBooleanLazyEvaluator extends CoalesceBooleanEvaluator { + CoalesceBooleanLazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java index 682b05111cc07..7d6834e765a96 100644 --- a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceBytesRefEvaluator.java @@ -10,7 +10,6 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BytesRefBlock; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +27,36 @@ * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class CoalesceBytesRefEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceBytesRefEvaluator.CoalesceEagerEvaluator, // - CoalesceBytesRefEvaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + CoalesceBytesRefEvaluator.CoalesceBytesRefEagerEvaluator, // + CoalesceBytesRefEvaluator.CoalesceBytesRefLazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new CoalesceBytesRefEagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "CoalesceBytesRefEagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceBytesRefLazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "CoalesceBytesRefLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +138,13 @@ public final void close() { * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link CoalesceBytesRefLazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends CoalesceBytesRefEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceBytesRefEagerEvaluator extends CoalesceBytesRefEvaluator { + CoalesceBytesRefEagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -193,8 +187,8 @@ protected BytesRefBlock perPosition(Page page, BytesRefBlock lastFullBlock, int * * */ - static final class CoalesceLazyEvaluator extends CoalesceBytesRefEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceBytesRefLazyEvaluator extends CoalesceBytesRefEvaluator { + CoalesceBytesRefLazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java index 708ee324adfd0..4c01a961ecbee 100644 --- a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceDoubleEvaluator.java @@ -7,10 +7,8 @@ package org.elasticsearch.xpack.esql.expression.function.scalar.nulls; -import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.DoubleBlock; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +26,36 @@ * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class CoalesceDoubleEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceDoubleEvaluator.CoalesceEagerEvaluator, // - CoalesceDoubleEvaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + CoalesceDoubleEvaluator.CoalesceDoubleEagerEvaluator, // + CoalesceDoubleEvaluator.CoalesceDoubleLazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new CoalesceDoubleEagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "CoalesceDoubleEagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceDoubleLazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "CoalesceDoubleLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +137,13 @@ public final void close() { * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link CoalesceDoubleLazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends CoalesceDoubleEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceDoubleEagerEvaluator extends CoalesceDoubleEvaluator { + CoalesceDoubleEagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -192,8 +185,8 @@ protected DoubleBlock perPosition(Page page, DoubleBlock lastFullBlock, int firs * * */ - static final class CoalesceLazyEvaluator extends CoalesceDoubleEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceDoubleLazyEvaluator extends CoalesceDoubleEvaluator { + CoalesceDoubleLazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java index 071f9d883022f..e90bd4b8e5e35 100644 --- a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceIntEvaluator.java @@ -7,10 +7,8 @@ package org.elasticsearch.xpack.esql.expression.function.scalar.nulls; -import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.IntBlock; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +26,36 @@ * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class CoalesceIntEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceIntEvaluator.CoalesceEagerEvaluator, // - CoalesceIntEvaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + CoalesceIntEvaluator.CoalesceIntEagerEvaluator, // + CoalesceIntEvaluator.CoalesceIntLazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new CoalesceIntEagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "CoalesceIntEagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceIntLazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "CoalesceIntLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +137,13 @@ public final void close() { * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link CoalesceIntLazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends CoalesceIntEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceIntEagerEvaluator extends CoalesceIntEvaluator { + CoalesceIntEagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -192,8 +185,8 @@ protected IntBlock perPosition(Page page, IntBlock lastFullBlock, int firstToEva * * */ - static final class CoalesceLazyEvaluator extends CoalesceIntEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceIntLazyEvaluator extends CoalesceIntEvaluator { + CoalesceIntLazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } diff --git a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java index 07159a9c0f698..53a21ad1198f4 100644 --- a/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java +++ b/x-pack/plugin/esql/src/main/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceLongEvaluator.java @@ -7,10 +7,8 @@ package org.elasticsearch.xpack.esql.expression.function.scalar.nulls; -import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.LongBlock; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +26,36 @@ * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class CoalesceLongEvaluator implements EvalOperator.ExpressionEvaluator permits - CoalesceLongEvaluator.CoalesceEagerEvaluator, // - CoalesceLongEvaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + CoalesceLongEvaluator.CoalesceLongEagerEvaluator, // + CoalesceLongEvaluator.CoalesceLongLazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new CoalesceLongEagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "CoalesceLongEagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new CoalesceLongLazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "CoalesceLongLazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +137,13 @@ public final void close() { * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link CoalesceLongLazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends CoalesceLongEvaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceLongEagerEvaluator extends CoalesceLongEvaluator { + CoalesceLongEagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -192,8 +185,8 @@ protected LongBlock perPosition(Page page, LongBlock lastFullBlock, int firstToE * * */ - static final class CoalesceLazyEvaluator extends CoalesceLongEvaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class CoalesceLongLazyEvaluator extends CoalesceLongEvaluator { + CoalesceLongLazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st index c4ee70a15af0f..33841f03f7803 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/X-CoalesceEvaluator.java.st @@ -7,10 +7,11 @@ package org.elasticsearch.xpack.esql.expression.function.scalar.nulls; +$if(BytesRef)$ import org.apache.lucene.util.BytesRef; +$endif$ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.$Type$Block; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; @@ -28,41 +29,36 @@ import java.util.stream.IntStream; * This class is generated. Edit {@code X-InEvaluator.java.st} instead. */ abstract sealed class Coalesce$Type$Evaluator implements EvalOperator.ExpressionEvaluator permits - Coalesce$Type$Evaluator.CoalesceEagerEvaluator, // - Coalesce$Type$Evaluator.CoalesceLazyEvaluator { - static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { - List childEvaluators = children.stream().map(toEvaluator::apply).toList(); - if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { - return new ExpressionEvaluator.Factory() { - @Override - public ExpressionEvaluator get(DriverContext context) { - return new CoalesceEagerEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); - } + Coalesce$Type$Evaluator.Coalesce$Type$EagerEvaluator, // + Coalesce$Type$Evaluator.Coalesce$Type$LazyEvaluator { - @Override - public String toString() { - return "CoalesceEagerEvaluator[values=" + childEvaluators + ']'; - } - }; - } + static ExpressionEvaluator.Factory toEvaluator(EvaluatorMapper.ToEvaluator toEvaluator, List children) { + List childEvaluators = children.stream().map(toEvaluator::apply).toList(); + if (childEvaluators.stream().allMatch(ExpressionEvaluator.Factory::eagerEvalSafeInLazy)) { return new ExpressionEvaluator.Factory() { @Override public ExpressionEvaluator get(DriverContext context) { - return new CoalesceLazyEvaluator( - context, - childEvaluators.stream().map(x -> x.get(context)).toList() - ); + return new Coalesce$Type$EagerEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); } @Override public String toString() { - return "CoalesceLazyEvaluator[values=" + childEvaluators + ']'; + return "Coalesce$Type$EagerEvaluator[values=" + childEvaluators + ']'; } }; } + return new ExpressionEvaluator.Factory() { + @Override + public ExpressionEvaluator get(DriverContext context) { + return new Coalesce$Type$LazyEvaluator(context, childEvaluators.stream().map(x -> x.get(context)).toList()); + } + + @Override + public String toString() { + return "Coalesce$Type$LazyEvaluator[values=" + childEvaluators + ']'; + } + }; + } protected final DriverContext driverContext; protected final List evaluators; @@ -144,13 +140,13 @@ abstract sealed class Coalesce$Type$Evaluator implements EvalOperator.Expression * Evaluates {@code COALESCE} eagerly per position if entire-block evaluation fails. * First we evaluate all remaining evaluators, and then we pluck the first non-null * value from each one. This is much faster than - * {@link CoalesceLazyEvaluator} but will include spurious warnings if any of the + * {@link Coalesce$Type$LazyEvaluator} but will include spurious warnings if any of the * evaluators make them so we only use it for evaluators that are * {@link Factory#eagerEvalSafeInLazy safe} to evaluate eagerly * in a lazy environment. */ - static final class CoalesceEagerEvaluator extends Coalesce$Type$Evaluator { - CoalesceEagerEvaluator(DriverContext driverContext, List evaluators) { + static final class Coalesce$Type$EagerEvaluator extends Coalesce$Type$Evaluator { + Coalesce$Type$EagerEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } @@ -195,8 +191,8 @@ $endif$ * * */ - static final class CoalesceLazyEvaluator extends Coalesce$Type$Evaluator { - CoalesceLazyEvaluator(DriverContext driverContext, List evaluators) { + static final class Coalesce$Type$LazyEvaluator extends Coalesce$Type$Evaluator { + Coalesce$Type$LazyEvaluator(DriverContext driverContext, List evaluators) { super(driverContext, evaluators); } From c7e1648a704f41b5fac7ac47216f4b3b17789eff Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 14 Jan 2025 12:44:13 -0500 Subject: [PATCH 6/8] Speed up COALESCE significantly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``` before after (operation) Score Error Score Error Units coalesce_2_noop 75.949 ± 3.961 -> 0.010 ± 0.001 ns/op 99.9% coalesce_2_eager 99.299 ± 6.959 -> 4.292 ± 0.227 ns/op 95.7% coalesce_2_lazy 113.118 ± 5.747 -> 26.746 ± 0.954 ns/op 76.4% ``` We tend to advise folks that "COALESCE is faster than CASE", but, as of 8.16.0/#112295 that wasn't the true. I was working with someone a few days ago to port a scripted_metric aggregation to ESQL and we saw COALESCE taking ~60% of the time. That won't do. The trouble is that CASE and COALESCE have to be *lazy*, meaning that operations like: ``` COALESCE(a, 1 / b) ``` should never emit a warning if `a` is not `null`, even if `b` is `0`. In 8.16/#112295 CASE grew an optimization where it could operate non-lazily if it was flagged as "safe". This brings a similar optimization to COALESCE, see it above as "case_2_eager", a 95.7% improvement. It also brings and arguably more important optimization - entire-block execution for COALESCE. The schort version is that, if the first parameter of COALESCE returns no nulls we can return it without doing anything lazily. There are a few more cases, but the upshot is that COALESCE is pretyt much *free* in cases where long strings of results are `null` or not `null`. That's the `coalesce_2_noop` line. Finally, when there mixed null and non-null values we were using a single builder with some fairly inefficient paths. This specializes them per type and skips some slow null-checking where possible. That's the `coalesce_2_lazy` result, a more modest 76.4%. --- .../compute/operator/EvalOperator.java | 5 +++ .../function/scalar/nulls/Coalesce.java | 35 +++++++------------ .../function/scalar/nulls/CoalesceTests.java | 33 +++++++++++++---- 3 files changed, 44 insertions(+), 29 deletions(-) diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java index 349ce7b00ff10..5837a7f07d60b 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java @@ -96,6 +96,11 @@ public Block eval(Page page) { public void close() { } + + @Override + public String toString() { + return "ConstantNull"; + } }; } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index 2cc5c74c51477..6522954a73768 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -11,6 +11,7 @@ import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.Expressions; @@ -28,7 +29,6 @@ import java.io.IOException; import java.util.List; -import java.util.Map; import java.util.stream.Stream; import static org.elasticsearch.xpack.esql.core.type.DataType.BOOLEAN; @@ -36,6 +36,7 @@ import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_SHAPE; import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME; import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS; +import static org.elasticsearch.xpack.esql.core.type.DataType.DOUBLE; import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_POINT; import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_SHAPE; import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER; @@ -200,26 +201,16 @@ public boolean foldable() { @Override public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) { - return SUPPORTED.get(dataType()).toEvaluator(toEvaluator, children()); + return switch (dataType()) { + case BOOLEAN -> CoalesceBooleanEvaluator.toEvaluator(toEvaluator, children()); + case DOUBLE, COUNTER_DOUBLE -> CoalesceDoubleEvaluator.toEvaluator(toEvaluator, children()); + case INTEGER, COUNTER_INTEGER -> CoalesceIntEvaluator.toEvaluator(toEvaluator, children()); + case LONG, DATE_NANOS, DATETIME, COUNTER_LONG, UNSIGNED_LONG -> CoalesceLongEvaluator.toEvaluator(toEvaluator, children()); + case KEYWORD, TEXT, SEMANTIC_TEXT, CARTESIAN_POINT, CARTESIAN_SHAPE, GEO_POINT, GEO_SHAPE, IP, VERSION -> + CoalesceBytesRefEvaluator.toEvaluator(toEvaluator, children()); + case NULL -> EvalOperator.CONSTANT_NULL_FACTORY; + case UNSUPPORTED, SHORT, BYTE, DATE_PERIOD, OBJECT, DOC_DATA_TYPE, SOURCE, TIME_DURATION, FLOAT, HALF_FLOAT, TSID_DATA_TYPE, + SCALED_FLOAT, PARTIAL_AGG -> throw new UnsupportedOperationException("can't be coalesced"); + }; } - - interface BuildCoalesce { - ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator, List children); - } - - private static final Map SUPPORTED = Map.ofEntries( - // We intend to support all types here. - Map.entry(BOOLEAN, CoalesceBooleanEvaluator::toEvaluator), - Map.entry(CARTESIAN_POINT, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(CARTESIAN_SHAPE, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(DATE_NANOS, CoalesceLongEvaluator::toEvaluator), - Map.entry(DATETIME, CoalesceLongEvaluator::toEvaluator), - Map.entry(GEO_POINT, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(GEO_SHAPE, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(INTEGER, CoalesceIntEvaluator::toEvaluator), - Map.entry(IP, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(KEYWORD, CoalesceBytesRefEvaluator::toEvaluator), - Map.entry(LONG, CoalesceLongEvaluator::toEvaluator), - Map.entry(VERSION, CoalesceBytesRefEvaluator::toEvaluator) - ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java index 9322bd6b01973..9be505a9be746 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java @@ -47,6 +47,7 @@ import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.nullValue; import static org.hamcrest.Matchers.sameInstance; public class CoalesceTests extends AbstractScalarFunctionTestCase { @@ -57,7 +58,7 @@ public CoalesceTests(@Name("TestCase") Supplier testC @ParametersFactory public static Iterable parameters() { List noNullsSuppliers = new ArrayList<>(); - VaragsTestCaseBuilder builder = new VaragsTestCaseBuilder(type -> "CoalesceEager"); + VaragsTestCaseBuilder builder = new VaragsTestCaseBuilder(type -> "Coalesce" + type + "Eager"); builder.expectString(strings -> strings.filter(v -> v != null).findFirst()); builder.expectLong(longs -> longs.filter(v -> v != null).findFirst()); builder.expectInt(ints -> ints.filter(v -> v != null).findFirst()); @@ -72,7 +73,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(first, DataType.IP, "first"), new TestCaseSupplier.TypedData(second, DataType.IP, "second") ), - "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceBytesRefEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.IP, equalTo(first == null ? second : first) ); @@ -87,7 +88,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(first, DataType.VERSION, "first"), new TestCaseSupplier.TypedData(second, DataType.VERSION, "second") ), - "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceBytesRefEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.VERSION, equalTo(first == null ? second : first) ); @@ -100,7 +101,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(firstDate, DataType.DATETIME, "first"), new TestCaseSupplier.TypedData(secondDate, DataType.DATETIME, "second") ), - "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceLongEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.DATETIME, equalTo(firstDate == null ? secondDate : firstDate) ); @@ -113,7 +114,7 @@ public static Iterable parameters() { new TestCaseSupplier.TypedData(firstDate, DataType.DATE_NANOS, "first"), new TestCaseSupplier.TypedData(secondDate, DataType.DATE_NANOS, "second") ), - "CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", + "CoalesceLongEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]", DataType.DATE_NANOS, equalTo(firstDate == null ? secondDate : firstDate) ); @@ -137,6 +138,20 @@ public static Iterable parameters() { suppliers.add(new TestCaseSupplier(nullCaseName(s, nullUpTo, true), types, () -> nullCase(s.get(), finalNullUpTo, true))); } } + suppliers.add( + new TestCaseSupplier( + List.of(DataType.NULL, DataType.NULL), + () -> new TestCaseSupplier.TestCase( + List.of( + new TestCaseSupplier.TypedData(null, DataType.NULL, "first"), + new TestCaseSupplier.TypedData(null, DataType.NULL, "second") + ), + "ConstantNull", + DataType.NULL, + nullValue() + ) + ) + ); return parameterSuppliersFromTypedData(suppliers); } @@ -175,7 +190,7 @@ protected static void addSpatialCombinations(List suppliers) { TestCaseSupplier.testCaseSupplier( leftDataSupplier, rightDataSupplier, - (l, r) -> equalTo("CoalesceEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]"), + (l, r) -> equalTo("CoalesceBytesRefEagerEvaluator[values=[Attribute[channel=0], Attribute[channel=1]]]"), dataType, (l, r) -> l ) @@ -243,7 +258,11 @@ public void testCoalesceNotNullable() { sub.add(between(0, sub.size()), randomLiteral(sub.get(sub.size() - 1).dataType())); Coalesce exp = build(Source.EMPTY, sub); // Known not to be nullable because it contains a non-null literal - assertThat(exp.nullable(), equalTo(Nullability.FALSE)); + if (testCase.expectedType() == DataType.NULL) { + assertThat(exp.nullable(), equalTo(Nullability.UNKNOWN)); + } else { + assertThat(exp.nullable(), equalTo(Nullability.FALSE)); + } } /** From 8d81a3e2cedc65f8e90c295577eae8fa9e9d3725 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Thu, 23 Jan 2025 11:07:20 -0500 Subject: [PATCH 7/8] Update --- .gitattributes | 1 + benchmarks/README.md | 2 +- .../compute/data/BooleanBlock.java | 5 +- .../compute/data/BooleanBlockBuilder.java | 5 +- .../compute/data/BytesRefBlock.java | 7 ++- .../compute/data/BytesRefBlockBuilder.java | 7 ++- .../compute/data/DoubleBlock.java | 5 +- .../compute/data/DoubleBlockBuilder.java | 5 +- .../compute/data/FloatBlock.java | 5 +- .../compute/data/FloatBlockBuilder.java | 5 +- .../elasticsearch/compute/data/IntBlock.java | 5 +- .../compute/data/IntBlockBuilder.java | 5 +- .../elasticsearch/compute/data/LongBlock.java | 5 +- .../compute/data/LongBlockBuilder.java | 5 +- .../org/elasticsearch/compute/data/Block.java | 3 +- .../compute/data/X-Block.java.st | 9 ++- .../compute/data/X-BlockBuilder.java.st | 9 ++- .../compute/operator/EvalOperator.java | 5 +- .../data/BlockBuilderCopyFromTests.java | 2 +- .../function/scalar/conditional/Case.java | 5 +- .../function/scalar/nulls/Coalesce.java | 2 +- .../function/scalar/nulls/CoalesceTests.java | 61 ++++++++++++------- 22 files changed, 118 insertions(+), 45 deletions(-) diff --git a/.gitattributes b/.gitattributes index 04881c92ede00..a0f434f16b32b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -11,6 +11,7 @@ x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/*.interp li x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseLexer*.java linguist-generated=true x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseParser*.java linguist-generated=true x-pack/plugin/esql/src/main/generated/** linguist-generated=true +x-pack/plugin/esql/src/main/generated-src/** linguist-generated=true # ESQL functions docs are autogenerated. More information at `docs/reference/esql/functions/README.md` docs/reference/esql/functions/*/** linguist-generated=true diff --git a/benchmarks/README.md b/benchmarks/README.md index 687577bd9672c..0cf95a2e81b9a 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -130,7 +130,7 @@ gradlew -p benchmarks/ run --args 'LongKeyedBucketOrdsBenchmark.multiBucket -pro ``` Note: As of January 2025 the latest release of async profiler doesn't work - with out JDK but the nightly is fine. + with our JDK but the nightly is fine. If you are on Mac, this'll warn you that you downloaded the shared library from the internet. You'll need to go to settings and allow it to run. diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java index 0faa1c5104ce3..b08b80acc6976 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java @@ -224,7 +224,10 @@ sealed interface Builder extends Block.Builder, BlockLoader.BooleanBuilder permi Builder copyFrom(BooleanBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. */ Builder copyFrom(BooleanBlock block, int position); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java index f9319791ad76a..7f4705ddecb27 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java @@ -117,7 +117,10 @@ private void copyFromVector(BooleanVector vector, int beginInclusive, int endExc } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java index 5b236b74cb6b3..6661895722725 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlock.java @@ -229,7 +229,12 @@ sealed interface Builder extends Block.Builder, BlockLoader.BytesRefBuilder perm Builder copyFrom(BytesRefBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. + * @param scratch Scratch string used to prevent allocation. Share this + between many calls to this function. */ Builder copyFrom(BytesRefBlock block, int position, BytesRef scratch); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java index b5b549bb6d6b1..0a2b350780405 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java @@ -121,7 +121,12 @@ private void copyFromVector(BytesRefVector vector, int beginInclusive, int endEx } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. + * @param scratch Scratch string used to prevent allocation. Share this + between many calls to this function. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java index 4cc89d7439dd7..04df6253662a9 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlock.java @@ -218,7 +218,10 @@ sealed interface Builder extends Block.Builder, BlockLoader.DoubleBuilder permit Builder copyFrom(DoubleBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. */ Builder copyFrom(DoubleBlock block, int position); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java index e7eaea166a9a6..8ecc9b91e0ffe 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java @@ -117,7 +117,10 @@ private void copyFromVector(DoubleVector vector, int beginInclusive, int endExcl } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java index 6888c1badde9c..0679e38b63219 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlock.java @@ -217,7 +217,10 @@ sealed interface Builder extends Block.Builder, BlockLoader.FloatBuilder permits Builder copyFrom(FloatBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. */ Builder copyFrom(FloatBlock block, int position); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java index a2e1e17aacd5e..8504912adc057 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/FloatBlockBuilder.java @@ -117,7 +117,10 @@ private void copyFromVector(FloatVector vector, int beginInclusive, int endExclu } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java index 025c1f0d964fe..6af61695929df 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlock.java @@ -217,7 +217,10 @@ sealed interface Builder extends Block.Builder, BlockLoader.IntBuilder permits I Builder copyFrom(IntBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. */ Builder copyFrom(IntBlock block, int position); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java index d08e55b1ca0ae..31449b6f1cd72 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java @@ -117,7 +117,10 @@ private void copyFromVector(IntVector vector, int beginInclusive, int endExclusi } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java index f71da8fad860f..090efd9a31579 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlock.java @@ -218,7 +218,10 @@ sealed interface Builder extends Block.Builder, BlockLoader.LongBuilder permits Builder copyFrom(LongBlock block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. */ Builder copyFrom(LongBlock block, int position); diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java index 1717a55c0a391..bf25347edd989 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java @@ -117,7 +117,10 @@ private void copyFromVector(LongVector vector, int beginInclusive, int endExclus } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java index 41f7e85dfc08f..de87c08f7ceb1 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java @@ -281,7 +281,8 @@ interface Builder extends BlockLoader.Builder, Releasable { * Copy the values in {@code block} from {@code beginInclusive} to * {@code endExclusive} into this builder. *

    - * For single position copies see {@link IntBlockBuilder#copyFrom(IntBlock, int)}, + * For single position copies use the faster + * {@link IntBlockBuilder#copyFrom(IntBlock, int)}, * {@link LongBlockBuilder#copyFrom(LongBlock, int)}, etc. *

    */ diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st index 326f7f4d05bb8..6c1616c370721 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st @@ -289,7 +289,14 @@ $endif$ Builder copyFrom($Type$Block block, int beginInclusive, int endExclusive); /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. +$if(BytesRef)$ + * @param scratch Scratch string used to prevent allocation. Share this + between many calls to this function. +$endif$ */ Builder copyFrom($Type$Block block, int position$if(BytesRef)$, BytesRef scratch$endif$); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st index 3064bd6a2725d..d60e1de179d20 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st @@ -165,7 +165,14 @@ $endif$ } /** - * Copy the values in {@code block} at {@code position}. + * Copy the values in {@code block} at {@code position}. If this position + * has a single value, this'll copy a single value. If this positions has + * many values, it'll copy all of them. If this is {@code null}, then it'll + * copy the {@code null}. +$if(BytesRef)$ + * @param scratch Scratch string used to prevent allocation. Share this + between many calls to this function. +$endif$ *

    * Note that there isn't a version of this method on {@link Block.Builder} that takes * {@link Block}. That'd be quite slow, running position by position. And it's important diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java index 5837a7f07d60b..2573baf78b16a 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/EvalOperator.java @@ -99,14 +99,15 @@ public void close() { @Override public String toString() { - return "ConstantNull"; + return CONSTANT_NULL_NAME; } }; } @Override public String toString() { - return "ConstantNull"; + return CONSTANT_NULL_NAME; } }; + private static final String CONSTANT_NULL_NAME = "ConstantNull"; } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java index 3e1d9bce50a56..1d3c8df914bc6 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderCopyFromTests.java @@ -100,7 +100,7 @@ private void assertEvens(Block block) { case FLOAT -> ((FloatBlockBuilder) builder).copyFrom((FloatBlock) block, i); case INT -> ((IntBlockBuilder) builder).copyFrom((IntBlock) block, i); case LONG -> ((LongBlockBuilder) builder).copyFrom((LongBlock) block, i); - default -> throw new IllegalArgumentException(); + default -> throw new IllegalArgumentException("unsupported type: " + elementType); } expected.add(valuesAtPositions(block, i, i + 1).get(0)); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java index 5990102a66d9e..04da04e1b3927 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java @@ -524,8 +524,9 @@ public Block eval(Page page) { ) { for (int p = 0; p < lhs.getPositionCount(); p++) { if (lhsOrRhs.mask().getBoolean(p)) { - // This is quite slow compared to using the single valued copy and unrolling per-type like coalesce - // *especially* when copying a Block. Vectors are pretty fast this way. + // TODO Copy the per-type specialization that COALESCE has. + // There's also a slowdown because copying from a block checks to see if there are any nulls and that's slow. + // Vectors do not, so this still shows as fairly fast. But not as fast as the per-type unrolling. builder.copyFrom(lhs, p, p + 1); } else { builder.copyFrom(rhs, p, p + 1); diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java index 6522954a73768..611c7a456864a 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/Coalesce.java @@ -210,7 +210,7 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) { CoalesceBytesRefEvaluator.toEvaluator(toEvaluator, children()); case NULL -> EvalOperator.CONSTANT_NULL_FACTORY; case UNSUPPORTED, SHORT, BYTE, DATE_PERIOD, OBJECT, DOC_DATA_TYPE, SOURCE, TIME_DURATION, FLOAT, HALF_FLOAT, TSID_DATA_TYPE, - SCALED_FLOAT, PARTIAL_AGG -> throw new UnsupportedOperationException("can't be coalesced"); + SCALED_FLOAT, PARTIAL_AGG -> throw new UnsupportedOperationException(dataType() + " can't be coalesced"); }; } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java index 9be505a9be746..99b21928c9a67 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java @@ -12,13 +12,14 @@ import org.elasticsearch.common.network.NetworkAddress; import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.BlockUtils; import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.compute.test.TestBlockFactory; import org.elasticsearch.core.Releasables; -import org.elasticsearch.xpack.esql.TestBlockFactory; import org.elasticsearch.xpack.esql.core.expression.Expression; import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; import org.elasticsearch.xpack.esql.core.expression.FoldContext; @@ -266,8 +267,12 @@ public void testCoalesceNotNullable() { } /** - * Inserts random non-null garbage around the expected data - * and runs + * Inserts random non-null garbage around the expected data and runs COALESCE. + *

    + * This is important for catching the case where your value is null, but the rest of the block + * isn't null. An off-by-one error in the evaluators can break this in a way that the standard + * tests weren't catching and this does. + *

    */ public void testEvaluateWithGarbage() { DriverContext context = driverContext(); @@ -275,38 +280,48 @@ public void testEvaluateWithGarbage() { int positions = between(2, 1024); List data = testCase.getData(); Page onePositionPage = row(testCase.getDataValues()); - Block[] manyPositionsBlocks = new Block[Math.toIntExact(data.stream().filter(d -> d.isForceLiteral() == false).count())]; + Block[] blocks = new Block[Math.toIntExact(data.stream().filter(d -> d.isForceLiteral() == false).count())]; int realPosition = between(0, positions - 1); try { - int b = 0; + int blocksIndex = 0; for (TestCaseSupplier.TypedData d : data) { - ElementType elementType = PlannerUtils.toElementType(d.type()); - try (Block.Builder builder = elementType.newBlockBuilder(positions, context.blockFactory())) { - for (int p = 0; p < positions; p++) { - if (p == realPosition) { - builder.copyFrom(onePositionPage.getBlock(b), 0, 1); - } else { - builder.copyFrom( - BlockUtils.constantBlock(TestBlockFactory.getNonBreakingInstance(), randomLiteral(d.type()).value(), 1), - 0, - 1 - ); - } - } - manyPositionsBlocks[b] = builder.build(); - } - b++; + blocks[blocksIndex] = blockWithRandomGarbage( + context.blockFactory(), + d.type(), + onePositionPage.getBlock(blocksIndex), + positions, + realPosition + ); + blocksIndex++; } try ( EvalOperator.ExpressionEvaluator eval = evaluator(expression).get(context); - Block block = eval.eval(new Page(positions, manyPositionsBlocks)) + Block block = eval.eval(new Page(positions, blocks)) ) { assertThat(block.getPositionCount(), is(positions)); assertThat(toJavaObjectUnsignedLongAware(block, realPosition), testCase.getMatcher()); assertThat("evaluates to tracked block", block.blockFactory(), sameInstance(context.blockFactory())); } } finally { - Releasables.close(onePositionPage::releaseBlocks, Releasables.wrap(manyPositionsBlocks)); + Releasables.close(onePositionPage::releaseBlocks, Releasables.wrap(blocks)); + } + } + + private Block blockWithRandomGarbage( + BlockFactory blockFactory, + DataType type, + Block singlePositionBlock, + int totalPositions, + int insertLocation + ) { + try (Block.Builder builder = PlannerUtils.toElementType(type).newBlockBuilder(totalPositions, blockFactory)) { + for (int p = 0; p < totalPositions; p++) { + Block copyFrom = p == insertLocation + ? singlePositionBlock + : BlockUtils.constantBlock(TestBlockFactory.getNonBreakingInstance(), randomLiteral(type).value(), 1); + builder.copyFrom(copyFrom, 0, 1); + } + return builder.build(); } } } From c93198d33db4048db22d784969edf8d12646a2e9 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Thu, 23 Jan 2025 11:16:54 -0500 Subject: [PATCH 8/8] Fix --- .../esql/expression/function/scalar/nulls/CoalesceTests.java | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java index 99b21928c9a67..1235a175294af 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/nulls/CoalesceTests.java @@ -14,7 +14,6 @@ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockFactory; import org.elasticsearch.compute.data.BlockUtils; -import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator;