From 54611b5468efa7574c261ed2cbdb1fe70a387df7 Mon Sep 17 00:00:00 2001 From: ChrisHegarty Date: Thu, 14 Sep 2023 11:56:42 +0100 Subject: [PATCH] Remove default driver context --- .../compute/operator/DriverContext.java | 8 --- .../compute/operator/EvalOperator.java | 2 +- .../compute/operator/FilterOperator.java | 2 +- .../operator/ThrowingDriverContext.java | 64 +++++++++++++++++++ .../elasticsearch/compute/OperatorTests.java | 23 +++++-- .../AggregatorFunctionTestCase.java | 12 ++-- ...untDistinctIntAggregatorFunctionTests.java | 2 +- ...ntDistinctLongAggregatorFunctionTests.java | 2 +- .../GroupingAggregatorFunctionTestCase.java | 28 ++++---- .../SumDoubleAggregatorFunctionTests.java | 10 +-- .../SumIntAggregatorFunctionTests.java | 2 +- .../SumLongAggregatorFunctionTests.java | 4 +- .../lucene/LuceneSourceOperatorTests.java | 2 +- .../lucene/LuceneTopNSourceOperatorTests.java | 2 +- .../ValuesSourceReaderOperatorTests.java | 6 +- .../compute/operator/AnyOperatorTestCase.java | 12 +++- .../compute/operator/AsyncOperatorTests.java | 14 +++- .../compute/operator/DriverContextTests.java | 37 ++++++----- .../operator/ForkingOperatorTestCase.java | 12 ++-- .../compute/operator/LimitOperatorTests.java | 4 +- .../compute/operator/OperatorTestCase.java | 8 +-- .../compute/operator/RowOperatorTests.java | 7 +- .../exchange/ExchangeServiceTests.java | 16 ++++- .../operator/topn/TopNOperatorTests.java | 18 +++--- .../xpack/esql/lookup/EnrichLookupIT.java | 11 +++- .../esql/enrich/EnrichLookupService.java | 12 +++- .../evaluator/mapper/EvaluatorMapper.java | 4 +- .../esql/plugin/TransportEsqlQueryAction.java | 2 +- .../function/AbstractFunctionTestCase.java | 20 ++++-- .../scalar/conditional/CaseTests.java | 7 +- .../function/scalar/math/RoundTests.java | 7 +- .../AbstractMultivalueFunctionTestCase.java | 3 +- .../scalar/multivalue/MvConcatTests.java | 2 +- .../function/scalar/nulls/CoalesceTests.java | 3 +- .../function/scalar/string/ConcatTests.java | 5 +- .../function/scalar/string/LeftTests.java | 3 +- .../function/scalar/string/RightTests.java | 3 +- .../function/scalar/string/SplitTests.java | 3 +- .../scalar/string/SubstringTests.java | 5 +- .../AbstractBinaryOperatorTestCase.java | 3 +- .../operator/arithmetic/NegTests.java | 3 +- .../xpack/esql/planner/EvalMapperTests.java | 13 +++- 42 files changed, 270 insertions(+), 136 deletions(-) create mode 100644 x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ThrowingDriverContext.java diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverContext.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverContext.java index 4e95e582769b5..db60b45f4516c 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverContext.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverContext.java @@ -35,9 +35,6 @@ */ public class DriverContext { - /** A default driver context. The returned bigArrays is non recycling. */ - public static DriverContext DEFAULT = new DriverContext(BigArrays.NON_RECYCLING_INSTANCE); - // Working set. Only the thread executing the driver will update this set. Set workingSet = Collections.newSetFromMap(new IdentityHashMap<>()); @@ -45,11 +42,6 @@ public class DriverContext { private final BigArrays bigArrays; - // For testing - public DriverContext() { - this(BigArrays.NON_RECYCLING_INSTANCE); - } - public DriverContext(BigArrays bigArrays) { Objects.requireNonNull(bigArrays); this.bigArrays = bigArrays; 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 7202f05b5562a..221c65f8a4ce3 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 @@ -25,7 +25,7 @@ public Operator get(DriverContext driverContext) { @Override public String describe() { - return "EvalOperator[evaluator=" + evaluator.get(DriverContext.DEFAULT) + "]"; + return "EvalOperator[evaluator=" + evaluator.get(new ThrowingDriverContext()) + "]"; } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterOperator.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterOperator.java index 20864373e8016..d3e7d6aa3a658 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterOperator.java +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/FilterOperator.java @@ -27,7 +27,7 @@ public Operator get(DriverContext driverContext) { @Override public String describe() { - return "FilterOperator[evaluator=" + evaluatorSupplier.get(DriverContext.DEFAULT) + "]"; + return "FilterOperator[evaluator=" + evaluatorSupplier.get(new ThrowingDriverContext()) + "]"; } } diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ThrowingDriverContext.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ThrowingDriverContext.java new file mode 100644 index 0000000000000..d985d7649ee38 --- /dev/null +++ b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/ThrowingDriverContext.java @@ -0,0 +1,64 @@ +/* + * 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.compute.operator; + +import org.elasticsearch.common.util.BigArrays; +import org.elasticsearch.common.util.ByteArray; +import org.elasticsearch.common.util.DoubleArray; +import org.elasticsearch.common.util.FloatArray; +import org.elasticsearch.common.util.IntArray; +import org.elasticsearch.common.util.LongArray; +import org.elasticsearch.core.Releasable; + +public class ThrowingDriverContext extends DriverContext { + public ThrowingDriverContext() { + super(new ThrowingBigArrays()); + } + + @Override + public BigArrays bigArrays() { + throw new AssertionError("should not reach here"); + } + + @Override + public boolean addReleasable(Releasable releasable) { + throw new AssertionError("should not reach here"); + } + + static class ThrowingBigArrays extends BigArrays { + + ThrowingBigArrays() { + super(null, null, "fake"); + } + + @Override + public ByteArray newByteArray(long size, boolean clearOnResize) { + throw new AssertionError("should not reach here"); + } + + @Override + public IntArray newIntArray(long size, boolean clearOnResize) { + throw new AssertionError("should not reach here"); + } + + @Override + public LongArray newLongArray(long size, boolean clearOnResize) { + throw new AssertionError("should not reach here"); + } + + @Override + public FloatArray newFloatArray(long size, boolean clearOnResize) { + throw new AssertionError("should not reach here"); + } + + @Override + public DoubleArray newDoubleArray(long size, boolean clearOnResize) { + throw new AssertionError("should not reach here"); + } + } +} diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java index 135877e4f5405..d7ec9bcaf99ee 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/OperatorTests.java @@ -27,6 +27,7 @@ import org.apache.lucene.tests.index.RandomIndexWriter; import org.apache.lucene.tests.store.BaseDirectoryWrapper; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.breaker.CircuitBreakingException; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.MockBigArrays; @@ -112,7 +113,7 @@ public void testQueryOperator() throws IOException { assertTrue("duplicated docId=" + docId, actualDocIds.add(docId)); } }); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); drivers.add(new Driver(driverContext, factory.get(driverContext), List.of(), docCollector, () -> {})); } OperatorTestCase.runDriver(drivers); @@ -144,9 +145,10 @@ public void testQueryOperator() throws IOException { } } + // @Repeat(iterations = 1) public void testGroupingWithOrdinals() throws Exception { final String gField = "g"; - final int numDocs = between(100, 10000); + final int numDocs = 2856; // between(100, 10000); final Map expectedCounts = new HashMap<>(); int keyLength = randomIntBetween(1, 10); try (BaseDirectoryWrapper dir = newDirectory(); RandomIndexWriter writer = new RandomIndexWriter(random(), dir)) { @@ -210,7 +212,7 @@ public String toString() { }; try (DirectoryReader reader = writer.getReader()) { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); Driver driver = new Driver( driverContext, @@ -258,14 +260,18 @@ public String toString() { LongBlock counts = page.getBlock(1); for (int i = 0; i < keys.getPositionCount(); i++) { BytesRef spare = new BytesRef(); - actualCounts.put(keys.getBytesRef(i, spare), counts.getLong(i)); + keys.getBytesRef(i, spare); + actualCounts.put(BytesRef.deepCopyOf(spare), counts.getLong(i)); } + // System.out.println("HEGO: keys.getPositionCount=" + keys.getPositionCount()); + // Releasables.close(keys); }), () -> {} ); OperatorTestCase.runDriver(driver); assertThat(actualCounts, equalTo(expectedCounts)); assertDriverContext(driverContext); + org.elasticsearch.common.util.MockBigArrays.ensureAllArraysAreReleased(); } } } @@ -276,7 +282,7 @@ public void testLimitOperator() { var values = randomList(positions, positions, ESTestCase::randomLong); var results = new ArrayList(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( var driver = new Driver( driverContext, @@ -388,6 +394,13 @@ private BigArrays bigArrays() { return new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService()); } + /** + * A {@link DriverContext} that won't throw {@link CircuitBreakingException}. + */ + protected final DriverContext driverContext() { + return new DriverContext(bigArrays()); + } + public static void assertDriverContext(DriverContext driverContext) { assertTrue(driverContext.isFinished()); assertThat(driverContext.getSnapshot().releasables(), empty()); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java index e2f1c606a4c25..a4b6c8b965962 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java @@ -92,7 +92,7 @@ public final void testIgnoresNulls() { int end = between(1_000, 100_000); List results = new ArrayList<>(); List input = CannedSourceOperator.collectPages(simpleInput(end)); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( @@ -110,14 +110,14 @@ public final void testIgnoresNulls() { public final void testMultivalued() { int end = between(1_000, 100_000); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages(new PositionMergingSourceOperator(simpleInput(end))); assertSimpleOutput(input, drive(simple(BigArrays.NON_RECYCLING_INSTANCE).get(driverContext), input.iterator())); } public final void testMultivaluedWithNulls() { int end = between(1_000, 100_000); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages( new NullInsertingSourceOperator(new PositionMergingSourceOperator(simpleInput(end))) ); @@ -125,7 +125,7 @@ public final void testMultivaluedWithNulls() { } public final void testEmptyInput() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), List.of().iterator()); assertThat(results, hasSize(1)); @@ -133,7 +133,7 @@ public final void testEmptyInput() { } public final void testEmptyInputInitialFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = drive( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), @@ -147,7 +147,7 @@ public final void testEmptyInputInitialFinal() { } public final void testEmptyInputInitialIntermediateFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = drive( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java index e559dc4effccb..974046469e518 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java @@ -62,7 +62,7 @@ protected void assertOutputFromEmpty(Block b) { } public void testRejectsDouble() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( driverContext, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java index 57b90fb844f54..04cbe0ed53236 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java @@ -63,7 +63,7 @@ protected void assertOutputFromEmpty(Block b) { } public void testRejectsDouble() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( driverContext, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java index 002790b3735d2..23015d066810a 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java @@ -145,7 +145,7 @@ protected ByteSizeValue smallEnoughToCircuitBreak() { } public final void testNullGroupsAndValues() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(new NullInsertingSourceOperator(simpleInput(end))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -153,7 +153,7 @@ public final void testNullGroupsAndValues() { } public final void testNullGroups() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(nullGroups(simpleInput(end))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -182,7 +182,7 @@ protected void appendNull(ElementType elementType, Block.Builder builder, int bl } public final void testNullValues() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(nullValues(simpleInput(end))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -190,7 +190,7 @@ public final void testNullValues() { } public final void testNullValuesInitialIntermediateFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(nullValues(simpleInput(end))); List results = drive( @@ -218,7 +218,7 @@ protected void appendNull(ElementType elementType, Block.Builder builder, int bl } public final void testMultivalued() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(1_000, 100_000); List input = CannedSourceOperator.collectPages(mergeValues(simpleInput(end))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -226,7 +226,7 @@ public final void testMultivalued() { } public final void testMulitvaluedNullGroupsAndValues() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(new NullInsertingSourceOperator(mergeValues(simpleInput(end)))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -234,7 +234,7 @@ public final void testMulitvaluedNullGroupsAndValues() { } public final void testMulitvaluedNullGroup() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(nullGroups(mergeValues(simpleInput(end)))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -242,7 +242,7 @@ public final void testMulitvaluedNullGroup() { } public final void testMulitvaluedNullValues() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); int end = between(50, 60); List input = CannedSourceOperator.collectPages(nullValues(mergeValues(simpleInput(end)))); List results = drive(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext), input.iterator()); @@ -250,12 +250,12 @@ public final void testMulitvaluedNullValues() { } public final void testNullOnly() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullOnly(List.of(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext))); } public final void testNullOnlyInputInitialFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullOnly( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), @@ -265,7 +265,7 @@ public final void testNullOnlyInputInitialFinal() { } public final void testNullOnlyInputInitialIntermediateFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullOnly( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), @@ -294,12 +294,12 @@ private void assertNullOnly(List operators) { } public final void testNullSome() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullSome(List.of(simple(nonBreakingBigArrays().withCircuitBreaking()).get(driverContext))); } public final void testNullSomeInitialFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullSome( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), @@ -309,7 +309,7 @@ public final void testNullSomeInitialFinal() { } public final void testNullSomeInitialIntermediateFinal() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); assertNullSome( List.of( simpleWithMode(nonBreakingBigArrays().withCircuitBreaking(), AggregatorMode.INITIAL).get(driverContext), diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunctionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunctionTests.java index 909b582bec732..767f9a2d5c25b 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunctionTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumDoubleAggregatorFunctionTests.java @@ -49,7 +49,7 @@ protected void assertSimpleOutput(List input, Block result) { } public void testOverflowSucceeds() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = new ArrayList<>(); try ( Driver d = new Driver( @@ -67,7 +67,7 @@ public void testOverflowSucceeds() { } public void testSummationAccuracy() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = new ArrayList<>(); try ( Driver d = new Driver( @@ -96,7 +96,7 @@ public void testSummationAccuracy() { : randomDoubleBetween(Double.MIN_VALUE, Double.MAX_VALUE, true); sum += values[i]; } - driverContext = new DriverContext(); + driverContext = driverContext(); try ( Driver d = new Driver( driverContext, @@ -118,7 +118,7 @@ public void testSummationAccuracy() { for (int i = 0; i < n; i++) { largeValues[i] = Double.MAX_VALUE; } - driverContext = new DriverContext(); + driverContext = driverContext(); try ( Driver d = new Driver( driverContext, @@ -137,7 +137,7 @@ public void testSummationAccuracy() { for (int i = 0; i < n; i++) { largeValues[i] = -Double.MAX_VALUE; } - driverContext = new DriverContext(); + driverContext = driverContext(); try ( Driver d = new Driver( driverContext, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumIntAggregatorFunctionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumIntAggregatorFunctionTests.java index d9e073ace9b6e..552b0d2d8836f 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumIntAggregatorFunctionTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumIntAggregatorFunctionTests.java @@ -49,7 +49,7 @@ protected void assertSimpleOutput(List input, Block result) { } public void testRejectsDouble() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( driverContext, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumLongAggregatorFunctionTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumLongAggregatorFunctionTests.java index 25e3d62ae9ed8..21880eb6b1a3e 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumLongAggregatorFunctionTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/SumLongAggregatorFunctionTests.java @@ -49,7 +49,7 @@ public void assertSimpleOutput(List input, Block result) { } public void testOverflowFails() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( driverContext, @@ -65,7 +65,7 @@ public void testOverflowFails() { } public void testRejectsDouble() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver d = new Driver( driverContext, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java index f3eef4ea45f90..bbafc8ed753cc 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneSourceOperatorTests.java @@ -143,7 +143,7 @@ public void testEmpty() { } private void testSimple(int size, int limit) { - DriverContext ctx = new DriverContext(); + DriverContext ctx = driverContext(); LuceneSourceOperator.Factory factory = simple(nonBreakingBigArrays(), DataPartitioning.SHARD, size, limit); Operator.OperatorFactory readS = ValuesSourceReaderOperatorTests.factory( reader, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java index 7abf042fa851f..54853abd0cecb 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneTopNSourceOperatorTests.java @@ -148,7 +148,7 @@ public void testEmpty() { } private void testSimple(int size, int limit) { - DriverContext ctx = new DriverContext(); + DriverContext ctx = driverContext(); LuceneTopNSourceOperator.Factory factory = simple(nonBreakingBigArrays(), DataPartitioning.SHARD, size, limit); Operator.OperatorFactory readS = ValuesSourceReaderOperatorTests.factory( reader, diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java index c2c8c9e05c064..64edcaa43d89b 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperatorTests.java @@ -160,7 +160,7 @@ protected SourceOperator simpleInput(int size) { randomPageSize(), LuceneOperator.NO_LIMIT ); - return luceneFactory.get(new DriverContext()); + return luceneFactory.get(driverContext()); } @Override @@ -226,7 +226,7 @@ public void testLoadAllInOnePageShuffled() { } private void loadSimpleAndAssert(List input) { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List results = new ArrayList<>(); List operators = List.of( factory( @@ -390,7 +390,7 @@ public void testValuesSourceReaderOperatorWithNulls() throws IOException { reader = w.getReader(); } - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); var luceneFactory = new LuceneSourceOperator.Factory( List.of(mockSearchContext(reader)), ctx -> new MatchAllDocsQuery(), diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java index e70160041047e..5edaa8d8da340 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AnyOperatorTestCase.java @@ -60,8 +60,7 @@ public final void testSimpleDescription() { Operator.OperatorFactory factory = simple(nonBreakingBigArrays()); String description = factory.describe(); assertThat(description, equalTo(expectedDescriptionOfSimple())); - DriverContext driverContext = new DriverContext(); - try (Operator op = factory.get(driverContext)) { + try (Operator op = factory.get(driverContext())) { if (op instanceof GroupingAggregatorFunction) { assertThat(description, matchesPattern(GROUPING_AGG_FUNCTION_DESCRIBE_PATTERN)); } else { @@ -74,7 +73,7 @@ public final void testSimpleDescription() { * Makes sure the description of {@link #simple} matches the {@link #expectedDescriptionOfSimple}. */ public final void testSimpleToString() { - try (Operator operator = simple(nonBreakingBigArrays()).get(new DriverContext())) { + try (Operator operator = simple(nonBreakingBigArrays()).get(driverContext())) { assertThat(operator.toString(), equalTo(expectedToStringOfSimple())); } } @@ -85,4 +84,11 @@ public final void testSimpleToString() { protected final BigArrays nonBreakingBigArrays() { return new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking(); } + + /** + * A {@link DriverContext} with a nonBreakingBigArrays. + */ + protected final DriverContext driverContext() { + return new DriverContext(nonBreakingBigArrays()); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AsyncOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AsyncOperatorTests.java index 48911c208bdc9..2c566aa46c413 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AsyncOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/AsyncOperatorTests.java @@ -13,6 +13,8 @@ import org.elasticsearch.action.support.ListenableActionFuture; import org.elasticsearch.action.support.PlainActionFuture; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.common.util.concurrent.EsExecutors; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BytesRefBlock; @@ -20,6 +22,7 @@ import org.elasticsearch.compute.data.LongVector; import org.elasticsearch.compute.data.Page; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.FixedExecutorBuilder; import org.elasticsearch.threadpool.TestThreadPool; @@ -116,7 +119,7 @@ public void close() { }); PlainActionFuture future = new PlainActionFuture<>(); Driver driver = new Driver( - new DriverContext(), + driverContext(), sourceOperator, List.of(asyncOperator), outputOperator, @@ -205,4 +208,13 @@ protected void doRun() { threadPool.schedule(command, delay, threadPool.executor(ESQL_TEST_EXECUTOR)); } } + + /** + * A {@link DriverContext} with a nonBreakingBigArrays. + */ + DriverContext driverContext() { + return new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverContextTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverContextTests.java index 8e322d6a80b99..dcf56c09efe05 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverContextTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverContextTests.java @@ -41,12 +41,10 @@ public class DriverContextTests extends ESTestCase { - final BigArrays bigArrays = new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()); - private static final String ESQL_TEST_EXECUTOR = "esql_test_executor"; public void testEmptyFinished() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = new AssertingDriverContext(); driverContext.finish(); assertTrue(driverContext.isFinished()); var snapshot = driverContext.getSnapshot(); @@ -54,7 +52,7 @@ public void testEmptyFinished() { } public void testAddByIdentity() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = new AssertingDriverContext(); ReleasablePoint point1 = new ReleasablePoint(1, 2); ReleasablePoint point2 = new ReleasablePoint(1, 2); assertThat(point1, equalTo(point2)); @@ -68,9 +66,11 @@ public void testAddByIdentity() { } public void testAddFinish() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = new AssertingDriverContext(); int count = randomInt(128); - Set releasables = IntStream.range(0, count).mapToObj(i -> randomReleasable()).collect(toIdentitySet()); + Set releasables = IntStream.range(0, count) + .mapToObj(i -> randomReleasable(driverContext.bigArrays())) + .collect(toIdentitySet()); assertThat(releasables, hasSize(count)); releasables.forEach(driverContext::addReleasable); @@ -84,7 +84,7 @@ public void testAddFinish() { } public void testRemoveAbsent() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = new AssertingDriverContext(); boolean removed = driverContext.removeReleasable(new NoOpReleasable()); assertThat(removed, equalTo(false)); driverContext.finish(); @@ -94,9 +94,11 @@ public void testRemoveAbsent() { } public void testAddRemoveFinish() { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = new AssertingDriverContext(); int count = randomInt(128); - Set releasables = IntStream.range(0, count).mapToObj(i -> randomReleasable()).collect(toIdentitySet()); + Set releasables = IntStream.range(0, count) + .mapToObj(i -> randomReleasable(driverContext.bigArrays())) + .collect(toIdentitySet()); assertThat(releasables, hasSize(count)); releasables.forEach(driverContext::addReleasable); @@ -112,9 +114,7 @@ public void testMultiThreaded() throws Exception { ExecutorService executor = threadPool.executor(ESQL_TEST_EXECUTOR); int tasks = randomIntBetween(4, 32); - List testDrivers = IntStream.range(0, tasks) - .mapToObj(i -> new TestDriver(new AssertingDriverContext(), randomInt(128), bigArrays)) - .toList(); + List testDrivers = IntStream.range(0, tasks).mapToObj(DriverContextTests::newTestDriver).toList(); List> futures = executor.invokeAll(testDrivers, 1, TimeUnit.MINUTES); assertThat(futures, hasSize(tasks)); for (var fut : futures) { @@ -135,9 +135,18 @@ public void testMultiThreaded() throws Exception { finishedReleasables.stream().flatMap(Set::stream).forEach(Releasable::close); } + static TestDriver newTestDriver(int unused) { + var driverContext = new AssertingDriverContext(); + return new TestDriver(driverContext, randomInt(128), driverContext.bigArrays()); + } + static class AssertingDriverContext extends DriverContext { volatile Thread thread; + AssertingDriverContext() { + super(new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService())); + } + @Override public boolean addReleasable(Releasable releasable) { checkThread(); @@ -219,10 +228,6 @@ static Set randomNFromCollection(Set input, int n) { return result; } - Releasable randomReleasable() { - return randomReleasable(bigArrays); - } - static Releasable randomReleasable(BigArrays bigArrays) { return switch (randomInt(3)) { case 0 -> new NoOpReleasable(); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ForkingOperatorTestCase.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ForkingOperatorTestCase.java index 5024c28d86a91..1c12fbf4bcd52 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ForkingOperatorTestCase.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/ForkingOperatorTestCase.java @@ -55,7 +55,7 @@ protected final Operator.OperatorFactory simple(BigArrays bigArrays) { public final void testInitialFinal() { BigArrays bigArrays = nonBreakingBigArrays(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages(simpleInput(between(1_000, 100_000))); List results = new ArrayList<>(); @@ -79,7 +79,7 @@ public final void testInitialFinal() { public final void testManyInitialFinal() { BigArrays bigArrays = nonBreakingBigArrays(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages(simpleInput(between(1_000, 100_000))); List partials = oneDriverPerPage(input, () -> List.of(simpleWithMode(bigArrays, AggregatorMode.INITIAL).get(driverContext))); List results = new ArrayList<>(); @@ -100,7 +100,7 @@ public final void testManyInitialFinal() { public final void testInitialIntermediateFinal() { BigArrays bigArrays = nonBreakingBigArrays(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages(simpleInput(between(1_000, 100_000))); List results = new ArrayList<>(); @@ -126,7 +126,7 @@ public final void testInitialIntermediateFinal() { @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/99160") public final void testManyInitialManyPartialFinal() { BigArrays bigArrays = nonBreakingBigArrays(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List input = CannedSourceOperator.collectPages(simpleInput(between(1_000, 100_000))); List partials = oneDriverPerPage(input, () -> List.of(simpleWithMode(bigArrays, AggregatorMode.INITIAL).get(driverContext))); @@ -217,7 +217,7 @@ List createDriversForInput(BigArrays bigArrays, List input, List

drivers = new ArrayList<>(); for (List pages : splitInput) { - DriverContext driver1Context = new DriverContext(); + DriverContext driver1Context = driverContext(); drivers.add( new Driver( driver1Context, @@ -234,7 +234,7 @@ List createDriversForInput(BigArrays bigArrays, List input, List

oneDriverPerPageList(Iterator> source, Sup List in = source.next(); try ( Driver d = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(in.iterator()), operators.get(), new PageConsumerOperator(result::add), @@ -131,7 +131,7 @@ protected final List oneDriverPerPageList(Iterator> source, Sup private void assertSimple(BigArrays bigArrays, int size) { List input = CannedSourceOperator.collectPages(simpleInput(size)); - List results = drive(simple(bigArrays.withCircuitBreaking()).get(new DriverContext()), input.iterator()); + List results = drive(simple(bigArrays.withCircuitBreaking()).get(driverContext()), input.iterator()); assertSimpleOutput(input, results); } @@ -143,7 +143,7 @@ protected final List drive(List operators, Iterator input) List results = new ArrayList<>(); try ( Driver d = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(input), operators, new PageConsumerOperator(results::add), @@ -166,7 +166,7 @@ public static void runDriver(List drivers) { drivers.add( new Driver( "dummy-session", - new DriverContext(), + new DriverContext(BigArrays.NON_RECYCLING_INSTANCE), () -> "dummy-driver", new SequenceLongBlockSourceOperator(LongStream.range(0, between(1, 100)), between(1, 100)), List.of(), diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/RowOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/RowOperatorTests.java index ac7bc2f7e4ad1..bb2713e105b93 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/RowOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/RowOperatorTests.java @@ -8,12 +8,15 @@ package org.elasticsearch.compute.operator; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BooleanBlock; import org.elasticsearch.compute.data.BytesRefBlock; import org.elasticsearch.compute.data.DoubleBlock; import org.elasticsearch.compute.data.IntBlock; import org.elasticsearch.compute.data.LongBlock; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.test.ESTestCase; import java.util.Arrays; @@ -22,7 +25,9 @@ import static org.hamcrest.Matchers.equalTo; public class RowOperatorTests extends ESTestCase { - final DriverContext driverContext = new DriverContext(); + final DriverContext driverContext = new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); public void testBoolean() { RowOperator.RowOperatorFactory factory = new RowOperator.RowOperatorFactory(List.of(false)); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java index 97c43920f54d0..40187abb3dae1 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/exchange/ExchangeServiceTests.java @@ -16,6 +16,8 @@ import org.elasticsearch.cluster.node.VersionInformation; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.common.util.concurrent.ConcurrentCollections; import org.elasticsearch.common.util.concurrent.EsExecutors; import org.elasticsearch.compute.data.Block; @@ -28,6 +30,7 @@ import org.elasticsearch.compute.operator.SinkOperator; import org.elasticsearch.compute.operator.SourceOperator; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.tasks.Task; import org.elasticsearch.tasks.TaskCancellationService; import org.elasticsearch.test.ESTestCase; @@ -267,14 +270,14 @@ void runConcurrentTest( for (int i = 0; i < numSinks; i++) { String description = "sink-" + i; ExchangeSinkOperator sinkOperator = new ExchangeSinkOperator(exchangeSink.get(), Function.identity()); - DriverContext dc = new DriverContext(); + DriverContext dc = driverContext(); Driver d = new Driver("test-session:1", dc, () -> description, seqNoGenerator.get(dc), List.of(), sinkOperator, () -> {}); drivers.add(d); } for (int i = 0; i < numSources; i++) { String description = "source-" + i; ExchangeSourceOperator sourceOperator = new ExchangeSourceOperator(exchangeSource.get()); - DriverContext dc = new DriverContext(); + DriverContext dc = driverContext(); Driver d = new Driver("test-session:2", dc, () -> description, sourceOperator, List.of(), seqNoCollector.get(dc), () -> {}); drivers.add(d); } @@ -451,4 +454,13 @@ public void sendResponse(Exception exception) throws IOException { in.sendResponse(exception); } } + + /** + * A {@link DriverContext} with a BigArrays that does not circuit break. + */ + DriverContext driverContext() { + return new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); + } } diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java index fee7b4d336270..f7ead4912d1be 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java +++ b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/topn/TopNOperatorTests.java @@ -183,7 +183,7 @@ public void testRamBytesUsed() { List.of(DEFAULT_UNSORTABLE), List.of(new TopNOperator.SortOrder(0, true, false)), pageSize - ).get(new DriverContext()); + ).get(driverContext()); long actualEmpty = RamUsageTester.ramUsed(op) - RamUsageTester.ramUsed(LONG) - RamUsageTester.ramUsed(DEFAULT_UNSORTABLE); assertThat(op.ramBytesUsed(), both(greaterThan(actualEmpty - underCount)).and(lessThan(actualEmpty))); // But when we fill it then we're quite close @@ -452,7 +452,7 @@ public void testCollectAllValues() { } List> actualTop = new ArrayList<>(); - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); try ( Driver driver = new Driver( driverContext, @@ -536,7 +536,7 @@ public void testCollectAllValues_RandomMultiValues() { expectedTop.add(eTop); } - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List> actualTop = new ArrayList<>(); try ( Driver driver = new Driver( @@ -569,7 +569,7 @@ private List> topNTwoColumns( List encoder, List sortOrders ) { - DriverContext driverContext = new DriverContext(); + DriverContext driverContext = driverContext(); List> outputValues = new ArrayList<>(); try ( Driver driver = new Driver( @@ -611,7 +611,7 @@ public void testTopNManyDescriptionAndToString() { + sorts + "]]"; assertThat(factory.describe(), equalTo("TopNOperator[count=10" + tail)); - try (Operator operator = factory.get(new DriverContext())) { + try (Operator operator = factory.get(driverContext())) { assertThat(operator.toString(), equalTo("TopNOperator[count=0/10" + tail)); } } @@ -831,7 +831,7 @@ private void assertSortingOnMV( int topCount = randomIntBetween(1, values.size()); try ( Driver driver = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(List.of(page).iterator()), List.of(new TopNOperator(topCount, List.of(blockType), List.of(encoder), List.of(sortOrders), randomPageSize())), new PageConsumerOperator(p -> readInto(actualValues, p)), @@ -965,7 +965,7 @@ public void testIPSortingSingleValue() throws UnknownHostException { List> actual = new ArrayList<>(); try ( Driver driver = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(List.of(new Page(builder.build())).iterator()), List.of( new TopNOperator( @@ -1088,7 +1088,7 @@ private void assertIPSortingOnMultiValues( List> actual = new ArrayList<>(); try ( Driver driver = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(List.of(new Page(builder.build())).iterator()), List.of( new TopNOperator( @@ -1169,7 +1169,7 @@ public void testZeroByte() { List> actual = new ArrayList<>(); try ( Driver driver = new Driver( - new DriverContext(), + driverContext(), new CannedSourceOperator(List.of(new Page(blocks.toArray(Block[]::new))).iterator()), List.of( new TopNOperator( diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/lookup/EnrichLookupIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/lookup/EnrichLookupIT.java index 86d082f1051ab..3829ed3ac3198 100644 --- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/lookup/EnrichLookupIT.java +++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/lookup/EnrichLookupIT.java @@ -14,6 +14,8 @@ import org.elasticsearch.cluster.node.DiscoveryNode; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.time.DateFormatter; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BytesRefBlock; import org.elasticsearch.compute.data.ElementType; @@ -25,6 +27,7 @@ import org.elasticsearch.compute.operator.OutputOperator; import org.elasticsearch.compute.operator.SourceOperator; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.tasks.CancellableTask; import org.elasticsearch.test.hamcrest.ElasticsearchAssertions; import org.elasticsearch.transport.TransportService; @@ -145,7 +148,7 @@ protected void start(Driver driver, ActionListener listener) { Driver.start(executor, driver, between(1, 1000), listener); } }; - Driver driver = new Driver(new DriverContext(), sourceOperator, List.of(enrichOperator), outputOperator, () -> {}); + Driver driver = new Driver(driverContext(), sourceOperator, List.of(enrichOperator), outputOperator, () -> {}); PlainActionFuture future = new PlainActionFuture<>(); runner.runToCompletion(List.of(driver), future); future.actionGet(TimeValue.timeValueSeconds(30)); @@ -224,4 +227,10 @@ public void testRandom() { public void testMultipleMatches() { } + + static DriverContext driverContext() { + return new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); + } } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupService.java index 4d783b9a1012c..42c705c8d9a76 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupService.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichLookupService.java @@ -22,6 +22,7 @@ import org.elasticsearch.cluster.service.ClusterService; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.util.BigArrays; import org.elasticsearch.common.util.concurrent.ThreadContext; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.ElementType; @@ -96,12 +97,19 @@ public class EnrichLookupService { private final SearchService searchService; private final TransportService transportService; private final Executor executor; + private final BigArrays bigArrays; - public EnrichLookupService(ClusterService clusterService, SearchService searchService, TransportService transportService) { + public EnrichLookupService( + ClusterService clusterService, + SearchService searchService, + TransportService transportService, + BigArrays bigArrays + ) { this.clusterService = clusterService; this.searchService = searchService; this.transportService = transportService; this.executor = transportService.getThreadPool().executor(EsqlPlugin.ESQL_THREAD_POOL_NAME); + this.bigArrays = bigArrays; transportService.registerRequestHandler(LOOKUP_ACTION_NAME, this.executor, LookupRequest::new, new TransportHandler()); } @@ -200,7 +208,7 @@ private void doLookup( OutputOperator outputOperator = new OutputOperator(List.of(), Function.identity(), result::set); Driver driver = new Driver( "enrich-lookup:" + sessionId, - new DriverContext(), + new DriverContext(bigArrays), () -> lookupDescription(sessionId, shardId, matchType, matchField, extractFields, inputPage.getPositionCount()), queryOperator, intermediateOperators, diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/EvaluatorMapper.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/EvaluatorMapper.java index aa4d9235bdb40..a518dd36e3e9e 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/EvaluatorMapper.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/mapper/EvaluatorMapper.java @@ -8,8 +8,8 @@ package org.elasticsearch.xpack.esql.evaluator.mapper; import org.elasticsearch.compute.data.Page; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; +import org.elasticsearch.compute.operator.ThrowingDriverContext; import org.elasticsearch.xpack.ql.expression.Expression; import java.util.function.Function; @@ -31,7 +31,7 @@ public interface EvaluatorMapper { */ default Object fold() { return toJavaObject( - toEvaluator(e -> driverContext -> p -> fromArrayRow(e.fold())[0]).get(DriverContext.DEFAULT).eval(new Page(1)), + toEvaluator(e -> driverContext -> p -> fromArrayRow(e.fold())[0]).get(new ThrowingDriverContext()).eval(new Page(1)), 0 ); } diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java index 614277e9d7216..6d859ef857bc7 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java @@ -65,7 +65,7 @@ public TransportEsqlQueryAction( this.requestExecutor = threadPool.executor(EsqlPlugin.ESQL_THREAD_POOL_NAME); exchangeService.registerTransportHandler(transportService); this.exchangeService = exchangeService; - this.enrichLookupService = new EnrichLookupService(clusterService, searchService, transportService); + this.enrichLookupService = new EnrichLookupService(clusterService, searchService, transportService, bigArrays); this.computeService = new ComputeService( searchService, transportService, diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java index 80aeb25d49d4a..017034eba9c64 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java @@ -10,6 +10,8 @@ import org.apache.lucene.document.InetAddressPoint; import org.apache.lucene.sandbox.document.HalfFloatPoint; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.BlockUtils; import org.elasticsearch.compute.data.Page; @@ -17,6 +19,7 @@ import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator; import org.elasticsearch.core.PathUtils; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.logging.LogManager; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.esql.evaluator.EvalMapper; @@ -180,7 +183,7 @@ public final void testEvaluate() { expression = new FoldNull().rule(expression); assertThat(expression.dataType(), equalTo(testCase.expectedType)); // TODO should we convert unsigned_long into BigDecimal so it's easier to assert? - Object result = toJavaObject(evaluator(expression).get(new DriverContext()).eval(row(testCase.getDataValues())), 0); + Object result = toJavaObject(evaluator(expression).get(driverContext()).eval(row(testCase.getDataValues())), 0); assertThat(result, not(equalTo(Double.NaN))); assertThat(result, not(equalTo(Double.POSITIVE_INFINITY))); assertThat(result, not(equalTo(Double.NEGATIVE_INFINITY))); @@ -194,7 +197,7 @@ public final void testSimpleWithNulls() { // TODO replace this with nulls insert assumeTrue("nothing to do if a type error", testCase.getExpectedTypeError() == null); assumeTrue("All test data types must be representable in order to build fields", testCase.allTypesAreRepresentable()); List simpleData = testCase.getDataValues(); - EvalOperator.ExpressionEvaluator eval = evaluator(buildFieldExpression(testCase)).get(new DriverContext()); + EvalOperator.ExpressionEvaluator eval = evaluator(buildFieldExpression(testCase)).get(driverContext()); Block[] orig = BlockUtils.fromListRow(simpleData); for (int i = 0; i < orig.length; i++) { List data = new ArrayList<>(); @@ -231,7 +234,7 @@ public final void testEvaluateInManyThreads() throws ExecutionException, Interru Page page = row(simpleData); futures.add(exec.submit(() -> { - EvalOperator.ExpressionEvaluator eval = evalSupplier.get(new DriverContext()); + EvalOperator.ExpressionEvaluator eval = evalSupplier.get(driverContext()); for (int c = 0; c < count; c++) { assertThat(toJavaObject(eval.eval(page), 0), testCase.getMatcher()); } @@ -249,7 +252,7 @@ public final void testEvaluatorToString() { assumeTrue("nothing to do if a type error", testCase.getExpectedTypeError() == null); assumeTrue("All test data types must be representable in order to build fields", testCase.allTypesAreRepresentable()); var supplier = evaluator(buildFieldExpression(testCase)); - var ev = supplier.get(new DriverContext()); + var ev = supplier.get(driverContext()); assertThat(ev.toString(), equalTo(testCase.evaluatorToString)); } @@ -608,4 +611,13 @@ private static void writeToTempDir(String subdir, String str, String extension) Path file = dir.resolve(functionName() + "." + extension); Files.writeString(file, str); } + + /** + * A {@link DriverContext} with a BigArrays that does not circuit break. + */ + protected DriverContext driverContext() { + return new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); + } } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java index 79138679e2414..68cea4ea873a5 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/CaseTests.java @@ -14,7 +14,6 @@ import org.elasticsearch.compute.data.Block; import org.elasticsearch.compute.data.IntBlock; import org.elasticsearch.compute.data.Page; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; @@ -89,9 +88,7 @@ protected Expression build(Source source, List args) { public void testEvalCase() { testCase( caseExpr -> toJavaObject( - caseExpr.toEvaluator(child -> evaluator(child)) - .get(new DriverContext()) - .eval(new Page(IntBlock.newConstantBlockWith(0, 1))), + caseExpr.toEvaluator(child -> evaluator(child)).get(driverContext()).eval(new Page(IntBlock.newConstantBlockWith(0, 1))), 0 ) ); @@ -157,7 +154,7 @@ public void testCaseIsLazy() { }; } return evaluator(child); - }).get(new DriverContext()).eval(new Page(IntBlock.newConstantBlockWith(0, 1))), 0)); + }).get(driverContext()).eval(new Page(IntBlock.newConstantBlockWith(0, 1))), 0)); } private static Case caseExpr(Object... args) { diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTests.java index cbfb0d6a579fe..853fe44d12ec9 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTests.java @@ -10,7 +10,6 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.ql.expression.Expression; @@ -86,14 +85,14 @@ public void testExamples() { private Object process(Number val) { return toJavaObject( - evaluator(new Round(Source.EMPTY, field("val", typeOf(val)), null)).get(new DriverContext()).eval(row(List.of(val))), + evaluator(new Round(Source.EMPTY, field("val", typeOf(val)), null)).get(driverContext()).eval(row(List.of(val))), 0 ); } private Object process(Number val, int decimals) { return toJavaObject( - evaluator(new Round(Source.EMPTY, field("val", typeOf(val)), field("decimals", DataTypes.INTEGER))).get(new DriverContext()) + evaluator(new Round(Source.EMPTY, field("val", typeOf(val)), field("decimals", DataTypes.INTEGER))).get(driverContext()) .eval(row(List.of(val, decimals))), 0 ); @@ -119,7 +118,7 @@ protected DataType expectedType(List argTypes) { public void testNoDecimalsToString() { assertThat( - evaluator(new Round(Source.EMPTY, field("val", DataTypes.DOUBLE), null)).get(new DriverContext()).toString(), + evaluator(new Round(Source.EMPTY, field("val", DataTypes.DOUBLE), null)).get(driverContext()).toString(), equalTo("RoundDoubleNoDecimalsEvaluator[val=Attribute[channel=0]]") ); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/AbstractMultivalueFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/AbstractMultivalueFunctionTestCase.java index a300dbb383211..714112b2db543 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/AbstractMultivalueFunctionTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/AbstractMultivalueFunctionTestCase.java @@ -13,7 +13,6 @@ import org.elasticsearch.compute.data.ElementType; import org.elasticsearch.compute.data.Page; import org.elasticsearch.compute.data.Vector; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; @@ -445,7 +444,7 @@ private void testBlock(boolean insertNulls) { builder.copyFrom(oneRowBlock, 0, 1); } Block input = builder.build(); - Block result = evaluator(buildFieldExpression(testCase)).get(new DriverContext()).eval(new Page(input)); + Block result = evaluator(buildFieldExpression(testCase)).get(driverContext()).eval(new Page(input)); assertThat(result.getPositionCount(), equalTo(result.getPositionCount())); for (int p = 0; p < input.getPositionCount(); p++) { diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvConcatTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvConcatTests.java index f6082af0e142e..e7670c9840b91 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvConcatTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvConcatTests.java @@ -72,7 +72,7 @@ public void testNull() { BytesRef bar = new BytesRef("bar"); BytesRef delim = new BytesRef(";"); Expression expression = buildFieldExpression(testCase); - DriverContext dvrCtx = new DriverContext(); + DriverContext dvrCtx = driverContext(); assertThat(toJavaObject(evaluator(expression).get(dvrCtx).eval(row(Arrays.asList(Arrays.asList(foo, bar), null))), 0), nullValue()); assertThat(toJavaObject(evaluator(expression).get(dvrCtx).eval(row(Arrays.asList(foo, null))), 0), nullValue()); 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 b2345e85336d4..a70fbf45ab4fd 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 @@ -11,7 +11,6 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import org.elasticsearch.compute.data.Block; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.evaluator.EvalMapper; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; @@ -88,7 +87,7 @@ public void testCoalesceIsLazy() { return dvrCtx -> page -> { throw new AssertionError("shouldn't be called"); }; } return EvalMapper.toEvaluator(child, layout); - }).get(new DriverContext()).eval(row(testCase.getDataValues())), 0), testCase.getMatcher()); + }).get(driverContext()).eval(row(testCase.getDataValues())), 0), testCase.getMatcher()); } public void testCoalesceNullabilityIsUnknown() { diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java index 7d5c18a5e3fd6..39d328747199d 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/ConcatTests.java @@ -11,7 +11,6 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import org.apache.lucene.util.BytesRef; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.ql.expression.Expression; @@ -99,7 +98,7 @@ public void testMany() { field("a", DataTypes.KEYWORD), IntStream.range(1, 5).mapToObj(i -> field(Integer.toString(i), DataTypes.KEYWORD)).toList() ) - ).get(new DriverContext()).eval(row(simpleData)), + ).get(driverContext()).eval(row(simpleData)), 0 ), equalTo(new BytesRef("cats and dogs")) @@ -121,7 +120,7 @@ public void testSomeConstant() { field("c", DataTypes.KEYWORD) ) ) - ).get(new DriverContext()).eval(row(simpleData)), + ).get(driverContext()).eval(row(simpleData)), 0 ), equalTo(new BytesRef("cats and dogs")) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/LeftTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/LeftTests.java index 77807e6463324..3d9e8d677f3e4 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/LeftTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/LeftTests.java @@ -12,7 +12,6 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.ql.expression.Expression; @@ -199,7 +198,7 @@ public void testUnicode() { private String process(String str, int length) { Block result = evaluator( new Left(Source.EMPTY, field("str", DataTypes.KEYWORD), new Literal(Source.EMPTY, length, DataTypes.INTEGER)) - ).get(new DriverContext()).eval(row(List.of(new BytesRef(str)))); + ).get(driverContext()).eval(row(List.of(new BytesRef(str)))); if (null == result) { return null; } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/RightTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/RightTests.java index 39222386a0cb0..ca9d1ef2dc1ee 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/RightTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/RightTests.java @@ -12,7 +12,6 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.ql.expression.Expression; @@ -201,7 +200,7 @@ public void testUnicode() { private String process(String str, int length) { Block result = evaluator( new Right(Source.EMPTY, field("str", DataTypes.KEYWORD), new Literal(Source.EMPTY, length, DataTypes.INTEGER)) - ).get(new DriverContext()).eval(row(List.of(new BytesRef(str)))); + ).get(driverContext()).eval(row(List.of(new BytesRef(str)))); if (null == result) { return null; } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SplitTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SplitTests.java index fc426be21e3f6..27b8ed722f963 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SplitTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SplitTests.java @@ -13,7 +13,6 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.BytesRefBlock; import org.elasticsearch.compute.data.Page; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; @@ -86,7 +85,7 @@ protected Expression build(Source source, List args) { public void testConstantDelimiter() { EvalOperator.ExpressionEvaluator eval = evaluator( new Split(Source.EMPTY, field("str", DataTypes.KEYWORD), new Literal(Source.EMPTY, new BytesRef(":"), DataTypes.KEYWORD)) - ).get(new DriverContext()); + ).get(driverContext()); /* * 58 is ascii for : and appears in the toString below. We don't convert the delimiter to a * string because we aren't really sure it's printable. It could be a tab or a bell or some diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SubstringTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SubstringTests.java index 19113ed65ffbe..5730b93aecd8d 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SubstringTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/string/SubstringTests.java @@ -12,7 +12,6 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.compute.data.Block; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.ql.expression.Expression; @@ -68,7 +67,7 @@ public Matcher resultsMatcher(List typedData public void testNoLengthToString() { assertThat( evaluator(new Substring(Source.EMPTY, field("str", DataTypes.KEYWORD), field("start", DataTypes.INTEGER), null)).get( - new DriverContext() + driverContext() ).toString(), equalTo("SubstringNoLengthEvaluator[str=Attribute[channel=0], start=Attribute[channel=1]]") ); @@ -137,7 +136,7 @@ private String process(String str, int start, Integer length) { new Literal(Source.EMPTY, start, DataTypes.INTEGER), length == null ? null : new Literal(Source.EMPTY, length, DataTypes.INTEGER) ) - ).get(new DriverContext()).eval(row(List.of(new BytesRef(str)))); + ).get(driverContext()).eval(row(List.of(new BytesRef(str)))); return result == null ? null : ((BytesRef) toJavaObject(result, 0)).utf8ToString(); } diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/AbstractBinaryOperatorTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/AbstractBinaryOperatorTestCase.java index 0ac08b61ec39e..3d5e7820677e4 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/AbstractBinaryOperatorTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/AbstractBinaryOperatorTestCase.java @@ -7,7 +7,6 @@ package org.elasticsearch.xpack.esql.expression.predicate.operator; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.analysis.Verifier; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; @@ -93,7 +92,7 @@ public final void testApplyToAllTypes() { Source src = new Source(Location.EMPTY, lhsType.typeName() + " " + rhsType.typeName()); if (isRepresentable(lhsType) && isRepresentable(rhsType)) { op = build(src, field("lhs", lhsType), field("rhs", rhsType)); - result = toJavaObject(evaluator(op).get(new DriverContext()).eval(row(List.of(lhs.value(), rhs.value()))), 0); + result = toJavaObject(evaluator(op).get(driverContext()).eval(row(List.of(lhs.value(), rhs.value()))), 0); } else { op = build(src, lhs, rhs); result = op.fold(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/NegTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/NegTests.java index f4de880bcd2b0..0138160ebd0fc 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/NegTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/NegTests.java @@ -10,7 +10,6 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; -import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.expression.function.scalar.AbstractScalarFunctionTestCase; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; @@ -172,7 +171,7 @@ public void testEdgeCases() { private Object process(Object val) { if (testCase.allTypesAreRepresentable()) { Neg neg = new Neg(Source.EMPTY, field("val", typeOf(val))); - return toJavaObject(evaluator(neg).get(new DriverContext()).eval(row(List.of(val))), 0); + return toJavaObject(evaluator(neg).get(driverContext()).eval(row(List.of(val))), 0); } else { // just fold if type is not representable Neg neg = new Neg(Source.EMPTY, new Literal(Source.EMPTY, val, typeOf(val))); return neg.fold(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/EvalMapperTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/EvalMapperTests.java index 34e6670862249..7956892c34645 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/EvalMapperTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/planner/EvalMapperTests.java @@ -10,8 +10,11 @@ import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import org.apache.lucene.util.BytesRef; +import org.elasticsearch.common.util.MockBigArrays; +import org.elasticsearch.common.util.PageCacheRecycler; import org.elasticsearch.compute.operator.DriverContext; import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.indices.breaker.NoneCircuitBreakerService; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xpack.esql.SerializationTestUtils; import org.elasticsearch.xpack.esql.evaluator.EvalMapper; @@ -128,8 +131,8 @@ public void testEvaluatorSuppliers() { Layout layout = lb.build(); var supplier = EvalMapper.toEvaluator(expression, layout); - EvalOperator.ExpressionEvaluator evaluator1 = supplier.get(new DriverContext()); - EvalOperator.ExpressionEvaluator evaluator2 = supplier.get(new DriverContext()); + EvalOperator.ExpressionEvaluator evaluator1 = supplier.get(driverContext()); + EvalOperator.ExpressionEvaluator evaluator2 = supplier.get(driverContext()); assertNotNull(evaluator1); assertNotNull(evaluator2); assertTrue(evaluator1 != evaluator2); @@ -143,4 +146,10 @@ public void testExpressionSerialization() { private static FieldAttribute field(String name, DataType type) { return new FieldAttribute(Source.EMPTY, name, new EsField(name, type, Collections.emptyMap(), false)); } + + static DriverContext driverContext() { + return new DriverContext( + new MockBigArrays(PageCacheRecycler.NON_RECYCLING_INSTANCE, new NoneCircuitBreakerService()).withCircuitBreaking() + ); + } }