ES|QL categorize with multiple groupings

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/BlockHash.java

@@ -174,13 +174,15 @@ public static BlockHash buildCategorizeBlockHash(
         List<GroupSpec> groups,
         AggregatorMode aggregatorMode,
         BlockFactory blockFactory,
-        AnalysisRegistry analysisRegistry
+        AnalysisRegistry analysisRegistry,
+        int emitBatchSize
     ) {
-        if (groups.size() != 1) {
-            throw new IllegalArgumentException("only a single CATEGORIZE group can used");
+        if (groups.size() == 1) {
+            return new CategorizeBlockHash(blockFactory, groups.get(0).channel, aggregatorMode, analysisRegistry);
+        } else {
+            assert groups.get(0).isCategorize();
+            return new CategorizePackedValuesBlockHash(groups, blockFactory, aggregatorMode, analysisRegistry, emitBatchSize);
         }
-
-        return new CategorizeBlockHash(blockFactory, groups.get(0).channel, aggregatorMode, analysisRegistry);
     }
 
     /**

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/CategorizeBlockHash.java

@@ -44,7 +44,7 @@
 import java.util.Objects;
 
 /**
- * Base BlockHash implementation for {@code Categorize} grouping function.
+ * BlockHash implementation for {@code Categorize} grouping function.
  */
 public class CategorizeBlockHash extends BlockHash {
 
@@ -95,12 +95,14 @@ public class CategorizeBlockHash extends BlockHash {
         }
     }
 
+    boolean seenNull() {
+        return seenNull;
+    }
+
     @Override
     public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
-        if (aggregatorMode.isInputPartial() == false) {
-            addInitial(page, addInput);
-        } else {
-            addIntermediate(page, addInput);
+        try (IntBlock block = add(page)) {
+            addInput.add(0, block);
         }
     }
 
@@ -129,29 +131,28 @@ public void close() {
         Releasables.close(evaluator, categorizer);
     }
 
+    private IntBlock add(Page page) {
+        return aggregatorMode.isInputPartial() == false ? addInitial(page) : addIntermediate(page);
+    }
+
     /**
      * Adds initial (raw) input to the state.
      */
-    private void addInitial(Page page, GroupingAggregatorFunction.AddInput addInput) {
-        try (IntBlock result = (IntBlock) evaluator.eval(page.getBlock(channel))) {
-            addInput.add(0, result);
-        }
+    IntBlock addInitial(Page page) {
+        return (IntBlock) evaluator.eval(page.getBlock(channel));
     }
 
     /**
      * Adds intermediate state to the state.
      */
-    private void addIntermediate(Page page, GroupingAggregatorFunction.AddInput addInput) {
+    private IntBlock addIntermediate(Page page) {
         if (page.getPositionCount() == 0) {
-            return;
+            return null;
         }
         BytesRefBlock categorizerState = page.getBlock(channel);
         if (categorizerState.areAllValuesNull()) {
             seenNull = true;
-            try (var newIds = blockFactory.newConstantIntVector(NULL_ORD, 1)) {
-                addInput.add(0, newIds);
-            }
-            return;
+            return blockFactory.newConstantIntBlockWith(NULL_ORD, 1);
         }
 
         Map<Integer, Integer> idMap = readIntermediate(categorizerState.getBytesRef(0, new BytesRef()));
@@ -161,9 +162,7 @@ private void addIntermediate(Page page, GroupingAggregatorFunction.AddInput addI
             for (int i = fromId; i < toId; i++) {
                 newIdsBuilder.appendInt(idMap.get(i));
             }
-            try (IntBlock newIds = newIdsBuilder.build()) {
-                addInput.add(0, newIds);
-            }
+            return newIdsBuilder.build();
         }
     }
 
@@ -172,7 +171,7 @@ private void addIntermediate(Page page, GroupingAggregatorFunction.AddInput addI
      *
      * @return a map from the old category id to the new one. The old ids go from 0 to {@code size - 1}.
      */
-    private Map<Integer, Integer> readIntermediate(BytesRef bytes) {
+    Map<Integer, Integer> readIntermediate(BytesRef bytes) {
         Map<Integer, Integer> idMap = new HashMap<>();
         try (StreamInput in = new BytesArray(bytes).streamInput()) {
             if (in.readBoolean()) {
@@ -198,15 +197,19 @@ private Block buildIntermediateBlock() {
         if (categorizer.getCategoryCount() == 0) {
             return blockFactory.newConstantNullBlock(seenNull ? 1 : 0);
         }
+        int positionCount = categorizer.getCategoryCount() + (seenNull ? 1 : 0);
+        // We're returning a block with N positions just because the Page must have all blocks with the same position count!
+        return blockFactory.newConstantBytesRefBlockWith(serializeCategorizer(), positionCount);
+    }
+
+    BytesRef serializeCategorizer() {
         try (BytesStreamOutput out = new BytesStreamOutput()) {
             out.writeBoolean(seenNull);
             out.writeVInt(categorizer.getCategoryCount());
             for (SerializableTokenListCategory category : categorizer.toCategoriesById()) {
                 category.writeTo(out);
             }
-            // We're returning a block with N positions just because the Page must have all blocks with the same position count!
-            int positionCount = categorizer.getCategoryCount() + (seenNull ? 1 : 0);
-            return blockFactory.newConstantBytesRefBlockWith(out.bytes().toBytesRef(), positionCount);
+            return out.bytes().toBytesRef();
         } catch (IOException e) {
             throw new RuntimeException(e);
         }

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/blockhash/CategorizePackedValuesBlockHash.java

@@ -0,0 +1,161 @@
+/*
+ * 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.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.BitArray;
+import org.elasticsearch.compute.aggregation.AggregatorMode;
+import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.ReleasableIterator;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.index.analysis.AnalysisRegistry;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * BlockHash implementation for {@code Categorize} grouping function as first
+ * grouping expression, followed by one or mode other grouping expressions.
+ */
+public class CategorizePackedValuesBlockHash extends BlockHash {
+
+    private final AggregatorMode aggregatorMode;
+    private final CategorizeBlockHash categorizeBlockHash;
+    private final PackedValuesBlockHash packedValuesBlockHash;
+
+    CategorizePackedValuesBlockHash(
+        List<GroupSpec> specs,
+        BlockFactory blockFactory,
+        AggregatorMode aggregatorMode,
+        AnalysisRegistry analysisRegistry,
+        int emitBatchSize
+    ) {
+        super(blockFactory);
+        this.aggregatorMode = aggregatorMode;
+
+        List<GroupSpec> packedValuesBlockHashSpecs = new ArrayList<>(specs);
+        packedValuesBlockHashSpecs.set(0, new GroupSpec(-1, ElementType.INT));
+
+        boolean success = false;
+        try {
+            categorizeBlockHash = new CategorizeBlockHash(blockFactory, specs.get(0).channel(), aggregatorMode, analysisRegistry);
+            packedValuesBlockHash = new PackedValuesBlockHash(packedValuesBlockHashSpecs, blockFactory, emitBatchSize);
+            success = true;
+        } finally {
+            if (success == false) {
+                close();
+            }
+        }
+    }
+
+    @Override
+    public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
+        try (IntBlock categories = getCategories(page)) {
+            packedValuesBlockHash.add(page.appendBlock(categories), addInput);
+        }
+    }
+
+    private IntBlock getCategories(Page page) {
+        if (aggregatorMode.isInputPartial() == false) {
+            return categorizeBlockHash.addInitial(page);
+        } else {
+            BytesRefBlock stateBlock = page.getBlock(0);
+            BytesRef stateBytes = stateBlock.getBytesRef(0, new BytesRef());
+
+            try (StreamInput in = new BytesArray(stateBytes).streamInput()) {
+                BytesRef categorizerState = in.readBytesRef();
+                Map<Integer, Integer> idMap = categorizeBlockHash.readIntermediate(categorizerState);
+                int[] oldIds = in.readIntArray();
+                try (IntBlock.Builder newIds = blockFactory.newIntBlockBuilder(page.getPositionCount())) {
+                    for (int oldId : oldIds) {
+                        newIds.appendInt(idMap.get(oldId));
+                    }
+                    return newIds.build();
+                }
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    @Override
+    public Block[] getKeys() {
+        Block[] keys = packedValuesBlockHash.getKeys();
+        if (aggregatorMode.isOutputPartial() == false) {
+            try (
+                BytesRefBlock regexes = (BytesRefBlock) categorizeBlockHash.getKeys()[0];
+                BytesRefBlock.Builder builder = blockFactory.newBytesRefBlockBuilder(keys[0].getPositionCount())
+            ) {
+                IntVector idsVector = (IntVector) keys[0].asVector();
+                int idsOffset = categorizeBlockHash.seenNull() ? 0 : -1;
+                BytesRef scratch = new BytesRef();
+                for (int i = 0; i < idsVector.getPositionCount(); i++) {
+                    int id = idsVector.getInt(i);
+                    if (id == 0) {
+                        builder.appendNull();
+                    } else {
+                        builder.appendBytesRef(regexes.getBytesRef(id + idsOffset, scratch));
+                    }
+                }
+                keys[0].close();
+                keys[0] = builder.build();
+            }
+        } else {
+            BytesRef state;
+            try (BytesStreamOutput out = new BytesStreamOutput()) {
+                out.writeBytesRef(categorizeBlockHash.serializeCategorizer());
+                IntVector idsVector = (IntVector) keys[0].asVector();
+                int[] idsArray = new int[idsVector.getPositionCount()];
+                for (int i = 0; i < idsVector.getPositionCount(); i++) {
+                    idsArray[i] = idsVector.getInt(i);
+                }
+                out.writeIntArray(idsArray);
+                state = out.bytes().toBytesRef();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            keys[0].close();
+            keys[0] = blockFactory.newConstantBytesRefBlockWith(state, keys[0].getPositionCount());
+        }
+        return keys;
+    }
+
+    @Override
+    public IntVector nonEmpty() {
+        return packedValuesBlockHash.nonEmpty();
+    }
+
+    @Override
+    public BitArray seenGroupIds(BigArrays bigArrays) {
+        return packedValuesBlockHash.seenGroupIds(bigArrays);
+    }
+
+    @Override
+    public final ReleasableIterator<IntBlock> lookup(Page page, ByteSizeValue targetBlockSize) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void close() {
+        Releasables.close(categorizeBlockHash, packedValuesBlockHash);
+    }
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Page.java

@@ -131,6 +131,9 @@ public <B extends Block> B getBlock(int blockIndex) {
         if (blocksReleased) {
             throw new IllegalStateException("can't read released page");
         }
+        if (blockIndex < 0) {
+            blockIndex += blocks.length;
+        }
         @SuppressWarnings("unchecked")
         B block = (B) blocks[blockIndex];
         if (block.isReleased()) {

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/HashAggregationOperator.java

@@ -51,7 +51,13 @@ public Operator get(DriverContext driverContext) {
             if (groups.stream().anyMatch(BlockHash.GroupSpec::isCategorize)) {
                 return new HashAggregationOperator(
                     aggregators,
-                    () -> BlockHash.buildCategorizeBlockHash(groups, aggregatorMode, driverContext.blockFactory(), analysisRegistry),
+                    () -> BlockHash.buildCategorizeBlockHash(
+                        groups,
+                        aggregatorMode,
+                        driverContext.blockFactory(),
+                        analysisRegistry,
+                        maxPageSize
+                    ),
                     driverContext
                 );
             }