ESQL: Allow agg tests to process many columns

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBytesRefBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -30,8 +31,8 @@ protected final SourceOperator simpleInput(BlockFactory blockFactory, int size)
     protected abstract BytesRef randomValue();
 
     @Override
-    public final void assertSimpleOutput(List<Block> input, Block result) {
-        Object[] values = input.stream().flatMap(AggregatorFunctionTestCase::allBytesRefs).sorted().limit(LIMIT).toArray(Object[]::new);
+    public final void assertSimpleOutput(List<Page> input, Block result) {
+        Object[] values = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).sorted().limit(LIMIT).toArray(Object[]::new);
         assertThat((List<?>) BlockUtils.toJavaObject(result, 0), contains(values));
     }
 }

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

@@ -58,7 +58,12 @@ protected final int aggregatorIntermediateBlockCount() {
 
     protected abstract String expectedDescriptionOfAggregator();
 
-    protected abstract void assertSimpleOutput(List<Block> input, Block result);
+    /**
+     * Assert that the result is correct given the input.
+     * @param input the input pages build by {@link #simpleInput}
+     * @param result the result of running {@link #aggregatorFunction()}
+     */
+    protected abstract void assertSimpleOutput(List<Page> input, Block result);
 
     @Override
     protected Operator.OperatorFactory simpleWithMode(SimpleOptions options, AggregatorMode mode) {
@@ -99,7 +104,7 @@ protected final void assertSimpleOutput(List<Page> input, List<Page> results) {
         assertThat(results.get(0).getPositionCount(), equalTo(1));
 
         Block result = results.get(0).getBlock(0);
-        assertSimpleOutput(input.stream().map(p -> p.<Block>getBlock(0)).toList(), result);
+        assertSimpleOutput(input, result);
     }
 
     public final void testIgnoresNulls() {

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
 
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long count = input.stream().flatMapToLong(b -> allLongs(b)).count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long count = input.stream().flatMapToLong(p -> allLongs(p.getBlock(0))).count();
         assertThat(((LongBlock) result).getLong(0), equalTo(count));
     }
 

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBooleanBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMap(b -> allBooleans(b)).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMap(p -> allBooleans(p.getBlock(0))).distinct().count();
         long count = ((LongBlock) result).getLong(0);
         assertThat(count, equalTo(expected));
     }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.BytesRefBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -42,8 +43,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMap(b -> allBytesRefs(b)).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).distinct().count();
         long count = ((LongBlock) result).getLong(0);
         // HLL is an approximation algorithm and precision depends on the number of values computed and the precision_threshold param
         // 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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceDoubleBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.test.ESTestCase;
@@ -38,8 +39,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMapToDouble(b -> allDoubles(b)).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMapToDouble(p -> allDoubles(p.getBlock(0))).distinct().count();
 
         long count = ((LongBlock) result).getLong(0);
         // 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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceFloatBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.test.ESTestCase;
@@ -38,8 +39,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMap(AggregatorFunctionTestCase::allFloats).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMap(p -> allFloats(p.getBlock(0))).distinct().count();
 
         long count = ((LongBlock) result).getLong(0);
         // 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

@@ -45,8 +45,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMapToInt(b -> allInts(b)).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMapToInt(p -> allInts(p.getBlock(0))).distinct().count();
 
         long count = ((LongBlock) result).getLong(0);
         // 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

@@ -46,8 +46,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
-        long expected = input.stream().flatMapToLong(b -> allLongs(b)).distinct().count();
+    protected void assertSimpleOutput(List<Page> input, Block result) {
+        long expected = input.stream().flatMapToLong(p -> allLongs(p.getBlock(0))).distinct().count();
         long count = ((LongBlock) result).getLong(0);
 
         // 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

@@ -46,10 +46,10 @@ protected String expectedToStringOfSimpleAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
+    protected void assertSimpleOutput(List<Page> input, Block result) {
         long sum = 0;
-        for (Block block : input) {
-            IntBlock ints = (IntBlock) block;
+        for (Page page : input) {
+            IntBlock ints = page.getBlock(0);
             for (int p = 0; p < ints.getPositionCount(); p++) {
                 /*
                  * Perform the sum on the values *only* if any of the

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBooleanBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        Boolean max = input.stream().flatMap(b -> allBooleans(b)).max(Comparator.naturalOrder()).get();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        Boolean max = input.stream().flatMap(p -> allBooleans(p.getBlock(0))).max(Comparator.naturalOrder()).get();
         assertThat(((BooleanBlock) result).getBoolean(0), equalTo(max));
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBytesRefBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -41,8 +42,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        Optional<BytesRef> max = input.stream().flatMap(b -> allBytesRefs(b)).max(Comparator.naturalOrder());
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        Optional<BytesRef> max = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).max(Comparator.naturalOrder());
         if (max.isEmpty()) {
             assertThat(result.isNull(0), equalTo(true));
             return;

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceDoubleBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.test.ESTestCase;
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        double max = input.stream().flatMapToDouble(b -> allDoubles(b)).max().getAsDouble();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        double max = input.stream().flatMapToDouble(p -> allDoubles(p.getBlock(0))).max().getAsDouble();
         assertThat(((DoubleBlock) result).getDouble(0), equalTo(max));
     }
 }

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceFloatBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.test.ESTestCase;
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        Float max = input.stream().flatMap(AggregatorFunctionTestCase::allFloats).max(floatComparator()).get();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        Float max = input.stream().flatMap(p -> allFloats(p.getBlock(0))).max(floatComparator()).get();
         assertThat(((FloatBlock) result).getFloat(0), equalTo(max));
     }
 }

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceIntBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -35,8 +36,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        int max = input.stream().flatMapToInt(b -> allInts(b)).max().getAsInt();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        int max = input.stream().flatMapToInt(p -> allInts(p.getBlock(0))).max().getAsInt();
         assertThat(((IntBlock) result).getInt(0), equalTo(max));
     }
 }

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BlockUtils;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBytesRefBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -42,8 +43,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        Optional<BytesRef> max = input.stream().flatMap(b -> allBytesRefs(b)).max(Comparator.naturalOrder());
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        Optional<BytesRef> max = input.stream().flatMap(p -> allBytesRefs(p.getBlock(0))).max(Comparator.naturalOrder());
         if (max.isEmpty()) {
             assertThat(result.isNull(0), equalTo(true));
             return;

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
 
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        long max = input.stream().flatMapToLong(b -> allLongs(b)).max().getAsLong();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        long max = input.stream().flatMapToLong(p -> allLongs(p.getBlock(0))).max().getAsLong();
         assertThat(((LongBlock) result).getLong(0), equalTo(max));
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceDoubleBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -39,7 +40,7 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
+    protected void assertSimpleOutput(List<Page> input, Block result) {
         assertThat(((DoubleBlock) result).getDouble(0), equalTo(0.8));
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceFloatBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -39,7 +40,7 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
+    protected void assertSimpleOutput(List<Page> input, Block result) {
         assertThat(((DoubleBlock) result).getDouble(0), closeTo(0.8, 0.001d));
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceIntBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -39,7 +40,7 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
+    protected void assertSimpleOutput(List<Page> input, Block result) {
         assertThat(((DoubleBlock) result).getDouble(0), equalTo(23.0));
     }
 }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SourceOperator;
 import org.elasticsearch.compute.test.SequenceLongBlockSourceOperator;
 
@@ -39,7 +40,7 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    protected void assertSimpleOutput(List<Block> input, Block result) {
+    protected void assertSimpleOutput(List<Page> input, Block result) {
         assertThat(((DoubleBlock) result).getDouble(0), equalTo(23.0));
     }
 }

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.SequenceBooleanBlockSourceOperator;
 import org.elasticsearch.compute.operator.SourceOperator;
 
@@ -36,8 +37,8 @@ protected String expectedDescriptionOfAggregator() {
     }
 
     @Override
-    public void assertSimpleOutput(List<Block> input, Block result) {
-        Boolean min = input.stream().flatMap(b -> allBooleans(b)).min(Comparator.naturalOrder()).get();
+    public void assertSimpleOutput(List<Page> input, Block result) {
+        Boolean min = input.stream().flatMap(p -> allBooleans(p.getBlock(0))).min(Comparator.naturalOrder()).get();
         assertThat(((BooleanBlock) result).getBoolean(0), equalTo(min));
     }
 }