Skip to content

Commit 4f95c6b

Browse files
nik9000nik9000
authored
ESQL: Allow agg tests to process many columns (elastic#132358)
Allows aggregation function tests to process the input `Page` instead of the first `Block`. This should allow us to write tests for elastic#108385.
1 parent 2a11826 commit 4f95c6b

File tree

52 files changed

+149
-107
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

52 files changed

+149
-107
lines changed

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AbstractTopBytesRefAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.elasticsearch.compute.data.Block;
1212
import org.elasticsearch.compute.data.BlockFactory;
1313
import org.elasticsearch.compute.data.BlockUtils;
14+
import org.elasticsearch.compute.data.Page;
1415
import org.elasticsearch.compute.operator.SequenceBytesRefBlockSourceOperator;
1516
import org.elasticsearch.compute.operator.SourceOperator;
1617

@@ -30,8 +31,8 @@ protected final SourceOperator simpleInput(BlockFactory blockFactory, int size)
3031
protected abstract BytesRef randomValue();
3132

3233
@Override
33-
public final void assertSimpleOutput(List<Block> input, Block result) {
34-
Object[] values = input.stream().flatMap(AggregatorFunctionTestCase::allBytesRefs).sorted().limit(LIMIT).toArray(Object[]::new);
34+
public final void assertSimpleOutput(List<Page> input, Block result) {
35+
Object[] values = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).sorted().limit(LIMIT).toArray(Object[]::new);
3536
assertThat((List<?>) BlockUtils.toJavaObject(result, 0), contains(values));
3637
}
3738
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/AggregatorFunctionTestCase.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,12 @@ protected final int aggregatorIntermediateBlockCount() {
5858

5959
protected abstract String expectedDescriptionOfAggregator();
6060

61-
protected abstract void assertSimpleOutput(List<Block> input, Block result);
61+
/**
62+
* Assert that the result is correct given the input.
63+
* @param input the input pages build by {@link #simpleInput}
64+
* @param result the result of running {@link #aggregatorFunction()}
65+
*/
66+
protected abstract void assertSimpleOutput(List<Page> input, Block result);
6267

6368
@Override
6469
protected Operator.OperatorFactory simpleWithMode(SimpleOptions options, AggregatorMode mode) {
@@ -99,7 +104,7 @@ protected final void assertSimpleOutput(List<Page> input, List<Page> results) {
99104
assertThat(results.get(0).getPositionCount(), equalTo(1));
100105

101106
Block result = results.get(0).getBlock(0);
102-
assertSimpleOutput(input.stream().map(p -> p.<Block>getBlock(0)).toList(), result);
107+
assertSimpleOutput(input, result);
103108
}
104109

105110
public final void testIgnoresNulls() {

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.elasticsearch.compute.data.Block;
1111
import org.elasticsearch.compute.data.BlockFactory;
1212
import org.elasticsearch.compute.data.LongBlock;
13+
import org.elasticsearch.compute.data.Page;
1314
import org.elasticsearch.compute.operator.SourceOperator;
1415
import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
1516

@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
3637
}
3738

3839
@Override
39-
protected void assertSimpleOutput(List<Block> input, Block result) {
40-
long count = input.stream().flatMapToLong(b -> allLongs(b)).count();
40+
protected void assertSimpleOutput(List<Page> input, Block result) {
41+
long count = input.stream().flatMapToLong(p -> allLongs(p.getBlock(0))).count();
4142
assertThat(((LongBlock) result).getLong(0), equalTo(count));
4243
}
4344

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctBooleanAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.elasticsearch.compute.data.Block;
1111
import org.elasticsearch.compute.data.BlockFactory;
1212
import org.elasticsearch.compute.data.LongBlock;
13+
import org.elasticsearch.compute.data.Page;
1314
import org.elasticsearch.compute.operator.SequenceBooleanBlockSourceOperator;
1415
import org.elasticsearch.compute.operator.SourceOperator;
1516

@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
3637
}
3738

3839
@Override
39-
protected void assertSimpleOutput(List<Block> input, Block result) {
40-
long expected = input.stream().flatMap(b -> allBooleans(b)).distinct().count();
40+
protected void assertSimpleOutput(List<Page> input, Block result) {
41+
long expected = input.stream().flatMap(p -> allBooleans(p.getBlock(0))).distinct().count();
4142
long count = ((LongBlock) result).getLong(0);
4243
assertThat(count, equalTo(expected));
4344
}

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctBytesRefAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.elasticsearch.compute.data.Block;
1212
import org.elasticsearch.compute.data.BlockFactory;
1313
import org.elasticsearch.compute.data.LongBlock;
14+
import org.elasticsearch.compute.data.Page;
1415
import org.elasticsearch.compute.operator.BytesRefBlockSourceOperator;
1516
import org.elasticsearch.compute.operator.SourceOperator;
1617

@@ -42,8 +43,8 @@ protected String expectedDescriptionOfAggregator() {
4243
}
4344

4445
@Override
45-
protected void assertSimpleOutput(List<Block> input, Block result) {
46-
long expected = input.stream().flatMap(b -> allBytesRefs(b)).distinct().count();
46+
protected void assertSimpleOutput(List<Page> input, Block result) {
47+
long expected = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).distinct().count();
4748
long count = ((LongBlock) result).getLong(0);
4849
// HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param
4950
// https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations-metrics-cardinality-aggregation.html

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctDoubleAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.elasticsearch.compute.data.Block;
1111
import org.elasticsearch.compute.data.BlockFactory;
1212
import org.elasticsearch.compute.data.LongBlock;
13+
import org.elasticsearch.compute.data.Page;
1314
import org.elasticsearch.compute.operator.SequenceDoubleBlockSourceOperator;
1415
import org.elasticsearch.compute.operator.SourceOperator;
1516
import org.elasticsearch.test.ESTestCase;
@@ -38,8 +39,8 @@ protected String expectedDescriptionOfAggregator() {
3839
}
3940

4041
@Override
41-
protected void assertSimpleOutput(List<Block> input, Block result) {
42-
long expected = input.stream().flatMapToDouble(b -> allDoubles(b)).distinct().count();
42+
protected void assertSimpleOutput(List<Page> input, Block result) {
43+
long expected = input.stream().flatMapToDouble(p -> allDoubles(p.getBlock(0))).distinct().count();
4344

4445
long count = ((LongBlock) result).getLong(0);
4546
// HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctFloatAggregatorFunctionTests.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.elasticsearch.compute.data.Block;
1111
import org.elasticsearch.compute.data.BlockFactory;
1212
import org.elasticsearch.compute.data.LongBlock;
13+
import org.elasticsearch.compute.data.Page;
1314
import org.elasticsearch.compute.operator.SequenceFloatBlockSourceOperator;
1415
import org.elasticsearch.compute.operator.SourceOperator;
1516
import org.elasticsearch.test.ESTestCase;
@@ -38,8 +39,8 @@ protected String expectedDescriptionOfAggregator() {
3839
}
3940

4041
@Override
41-
protected void assertSimpleOutput(List<Block> input, Block result) {
42-
long expected = input.stream().flatMap(AggregatorFunctionTestCase::allFloats).distinct().count();
42+
protected void assertSimpleOutput(List<Page> input, Block result) {
43+
long expected = input.stream().flatMap(p -> allFloats(p.getBlock(0))).distinct().count();
4344

4445
long count = ((LongBlock) result).getLong(0);
4546
// HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctIntAggregatorFunctionTests.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@ protected String expectedDescriptionOfAggregator() {
4545
}
4646

4747
@Override
48-
protected void assertSimpleOutput(List<Block> input, Block result) {
49-
long expected = input.stream().flatMapToInt(b -> allInts(b)).distinct().count();
48+
protected void assertSimpleOutput(List<Page> input, Block result) {
49+
long expected = input.stream().flatMapToInt(p -> allInts(p.getBlock(0))).distinct().count();
5050

5151
long count = ((LongBlock) result).getLong(0);
5252
// HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/CountDistinctLongAggregatorFunctionTests.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ protected String expectedDescriptionOfAggregator() {
4646
}
4747

4848
@Override
49-
protected void assertSimpleOutput(List<Block> input, Block result) {
50-
long expected = input.stream().flatMapToLong(b -> allLongs(b)).distinct().count();
49+
protected void assertSimpleOutput(List<Page> input, Block result) {
50+
long expected = input.stream().flatMapToLong(p -> allLongs(p.getBlock(0))).distinct().count();
5151
long count = ((LongBlock) result).getLong(0);
5252

5353
// HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/FilteredAggregatorFunctionTests.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@ protected String expectedToStringOfSimpleAggregator() {
4646
}
4747

4848
@Override
49-
protected void assertSimpleOutput(List<Block> input, Block result) {
49+
protected void assertSimpleOutput(List<Page> input, Block result) {
5050
long sum = 0;
51-
for (Block block : input) {
52-
IntBlock ints = (IntBlock) block;
51+
for (Page page : input) {
52+
IntBlock ints = page.getBlock(0);
5353
for (int p = 0; p < ints.getPositionCount(); p++) {
5454
/*
5555
* Perform the sum on the values *only* if any of the

0 commit comments

Comments
 (0)