Randomize groupIds block types to check most AddInput cases

ivancea · ivancea · commit 32588b56bf87 · 2025-05-06T18:18:01.000+02:00
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
@@ -12,8 +12,10 @@
 import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.compute.ConstantBooleanExpressionEvaluator;
 import org.elasticsearch.compute.aggregation.blockhash.BlockHash;
+import org.elasticsearch.compute.aggregation.blockhash.WrappedBlockHash;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BlockTypeRandomizer;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
@@ -39,6 +41,7 @@
 import org.elasticsearch.compute.test.TestBlockFactory;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasables;
+import org.elasticsearch.index.analysis.AnalysisRegistry;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.hamcrest.Matcher;
 
@@ -47,6 +50,7 @@
 import java.util.SortedSet;
 import java.util.TreeSet;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.DoubleStream;
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
@@ -104,7 +108,7 @@ private Operator.OperatorFactory simpleWithMode(
         if (randomBoolean()) {
             supplier = chunkGroups(emitChunkSize, supplier);
         }
-        return new HashAggregationOperator.HashAggregationOperatorFactory(
+        return new RandomizingHashAggregationOperatorFactory(
             List.of(new BlockHash.GroupSpec(0, ElementType.LONG)),
             mode,
             List.of(supplier.groupingAggregatorFactory(mode, channels(mode))),
@@ -777,4 +781,76 @@ public String describe() {
         };
     }
 
+    private record RandomizingHashAggregationOperatorFactory(
+        List<BlockHash.GroupSpec> groups,
+        AggregatorMode aggregatorMode,
+        List<GroupingAggregator.Factory> aggregators,
+        int maxPageSize,
+        AnalysisRegistry analysisRegistry
+    ) implements Operator.OperatorFactory {
+
+        @Override
+        public Operator get(DriverContext driverContext) {
+            Supplier<BlockHash> blockHashSupplier = () -> {
+                BlockHash blockHash = groups.stream().anyMatch(BlockHash.GroupSpec::isCategorize)
+                    ? BlockHash.buildCategorizeBlockHash(
+                        groups,
+                        aggregatorMode,
+                        driverContext.blockFactory(),
+                        analysisRegistry,
+                        maxPageSize
+                    )
+                    : BlockHash.build(groups, driverContext.blockFactory(), maxPageSize, false);
+
+                return new WrappedBlockHash(driverContext.blockFactory(), blockHash) {
+                    @Override
+                    public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
+                        blockHash.add(page, new GroupingAggregatorFunction.AddInput() {
+                            @Override
+                            public void add(int positionOffset, IntBlock groupIds) {
+                                IntBlock newGroupIds = aggregatorMode.isInputPartial()
+                                    ? groupIds
+                                    : BlockTypeRandomizer.randomizeBlockType(groupIds);
+                                addInput.add(positionOffset, newGroupIds);
+                            }
+
+                            @Override
+                            public void add(int positionOffset, IntArrayBlock groupIds) {
+                                add(positionOffset, (IntBlock) groupIds);
+                            }
+
+                            @Override
+                            public void add(int positionOffset, IntBigArrayBlock groupIds) {
+                                add(positionOffset, (IntBlock) groupIds);
+                            }
+
+                            @Override
+                            public void add(int positionOffset, IntVector groupIds) {
+                                add(positionOffset, groupIds.asBlock());
+                            }
+
+                            @Override
+                            public void close() {
+                                addInput.close();
+                            }
+                        });
+                    }
+                };
+            };
+
+            return new HashAggregationOperator(aggregators, blockHashSupplier, driverContext);
+        }
+
+        @Override
+        public String describe() {
+            return new HashAggregationOperator.HashAggregationOperatorFactory(
+                groups,
+                aggregatorMode,
+                aggregators,
+                maxPageSize,
+                analysisRegistry
+            ).describe();
+        }
+    }
+
 }
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/WrappedBlockHash.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/WrappedBlockHash.java
@@ -0,0 +1,69 @@
+/*
+ * 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.aggregation.blockhash;
+
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.BitArray;
+import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.ReleasableIterator;
+
+/**
+ * A test BlockHash that wraps another one.
+ * <p>
+ *     Its methods can be overridden to implement custom behaviours or checks.
+ * </p>
+ */
+public class WrappedBlockHash extends BlockHash {
+    protected BlockHash blockHash;
+
+    public WrappedBlockHash(BlockFactory blockFactory, BlockHash blockHash) {
+        super(blockFactory);
+        this.blockHash = blockHash;
+    }
+
+    @Override
+    public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
+        blockHash.add(page, addInput);
+    }
+
+    @Override
+    public ReleasableIterator<IntBlock> lookup(Page page, ByteSizeValue targetBlockSize) {
+        return blockHash.lookup(page, targetBlockSize);
+    }
+
+    @Override
+    public Block[] getKeys() {
+        return blockHash.getKeys();
+    }
+
+    @Override
+    public IntVector nonEmpty() {
+        return blockHash.nonEmpty();
+    }
+
+    @Override
+    public BitArray seenGroupIds(BigArrays bigArrays) {
+        return blockHash.seenGroupIds(bigArrays);
+    }
+
+    @Override
+    public void close() {
+        blockHash.close();
+    }
+
+    @Override
+    public String toString() {
+        return blockHash.toString();
+    }
+}
diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockTypeRandomizer.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockTypeRandomizer.java
@@ -0,0 +1,102 @@
+/*
+ * 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.data;
+
+import org.elasticsearch.compute.test.TestBlockFactory;
+
+import java.util.BitSet;
+
+import static org.elasticsearch.test.ESTestCase.randomIntBetween;
+
+public class BlockTypeRandomizer {
+    private BlockTypeRandomizer() {}
+
+    /**
+     * Returns a block with the same contents, but with a randomized type (Constant, vector, big-array...).
+     * <p>
+     *     The new block uses a non-breaking block builder, and doesn't increment the circuit breaking.
+     *     This is done to avoid randomly using more memory in tests that expect a deterministic memory usage.
+     * </p>
+     */
+    public static IntBlock randomizeBlockType(IntBlock block) {
+        // Just to track the randomization
+        int classCount = 4;
+
+        BlockFactory blockFactory = TestBlockFactory.getNonBreakingInstance();
+
+        //
+        // ConstantNullBlock. It requires all positions to be null
+        //
+        if (randomIntBetween(0, --classCount) == 0 && block.areAllValuesNull()) {
+            if (block instanceof ConstantNullBlock) {
+                return block;
+            }
+            return new ConstantNullBlock(block.getPositionCount(), blockFactory);
+        }
+
+        //
+        // IntVectorBlock. It doesn't allow nulls or multivalues
+        //
+        if (randomIntBetween(0, --classCount) == 0 && block.doesHaveMultivaluedFields() == false && block.mayHaveNulls() == false) {
+            if (block instanceof IntVectorBlock) {
+                return block;
+            }
+
+            int[] values = new int[block.getPositionCount()];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = block.getInt(i);
+            }
+
+            return new IntVectorBlock(new IntArrayVector(values, block.getPositionCount(), blockFactory));
+        }
+
+        // Both IntArrayBlock and IntBigArrayBlock need a nullsBitSet and a firstValueIndexes int[]
+        int[] firstValueIndexes = new int[block.getPositionCount() + 1];
+        BitSet nullsMask = new BitSet(block.getPositionCount());
+        for (int i = 0; i < block.getPositionCount(); i++) {
+            firstValueIndexes[i] = block.getFirstValueIndex(i);
+
+            if (block.isNull(i)) {
+                nullsMask.set(i);
+            }
+        }
+        int totalValues = block.getFirstValueIndex(block.getPositionCount() - 1) + block.getValueCount(block.getPositionCount() - 1);
+        firstValueIndexes[firstValueIndexes.length - 1] = totalValues;
+
+        //
+        // IntArrayBlock
+        //
+        if (randomIntBetween(0, --classCount) == 0) {
+            if (block instanceof IntVectorBlock) {
+                return block;
+            }
+
+            int[] values = new int[totalValues];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = block.getInt(i);
+            }
+
+            return new IntArrayBlock(values, block.getPositionCount(), firstValueIndexes, nullsMask, block.mvOrdering(), blockFactory);
+        }
+        assert classCount == 1;
+
+        //
+        // IntBigArrayBlock
+        //
+        if (block instanceof IntBigArrayBlock) {
+            return block;
+        }
+
+        var intArray = blockFactory.bigArrays().newIntArray(totalValues);
+        for (int i = 0; i < block.getPositionCount(); i++) {
+            intArray.set(i, block.getInt(i));
+        }
+
+        return new IntBigArrayBlock(intArray, block.getPositionCount(), firstValueIndexes, nullsMask, block.mvOrdering(), blockFactory);
+    }
+}
