ESQL: CATEGORIZE as a BlockHash

docs/changelog/114317.yaml

@@ -0,0 +1,5 @@
+pr: 114317
+summary: "ESQL: CATEGORIZE as a `BlockHash`"
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java

@@ -192,7 +192,10 @@ public enum DataType {
      * inside alongside time-series aggregations. These fields are not parsable from the
      * mapping and should be hidden from users.
      */
-    PARTIAL_AGG(builder().esType("partial_agg").unknownSize());
+    PARTIAL_AGG(builder().esType("partial_agg").unknownSize()),
+
+    // NOCOMMIT: if we keep this DataType, what should esType be?
+    CATEGORY(builder().esType("category").unknownSize());
 
     /**
      * Types that are actively being built. These types are not returned

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

@@ -0,0 +1,98 @@
+/*
+ * 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.BytesRefBuilder;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.BitArray;
+import org.elasticsearch.common.util.BytesRefHash;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefVector;
+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.xpack.ml.aggs.categorization.CategorizationBytesRefHash;
+import org.elasticsearch.xpack.ml.aggs.categorization.CategorizationPartOfSpeechDictionary;
+import org.elasticsearch.xpack.ml.aggs.categorization.SerializableTokenListCategory;
+import org.elasticsearch.xpack.ml.aggs.categorization.TokenListCategorizer;
+
+import java.io.IOException;
+
+public abstract class AbstractCategorizeBlockHash extends BlockHash {
+    // TODO: this should probably also take an emitBatchSize
+    private final int channel;
+    private final boolean outputPartial;
+    protected final TokenListCategorizer.CloseableTokenListCategorizer categorizer;
+
+    AbstractCategorizeBlockHash(BlockFactory blockFactory, int channel, boolean outputPartial) {
+        super(blockFactory);
+        this.channel = channel;
+        this.outputPartial = outputPartial;
+        this.categorizer = new TokenListCategorizer.CloseableTokenListCategorizer(
+            new CategorizationBytesRefHash(new BytesRefHash(2048, blockFactory.bigArrays())),
+            CategorizationPartOfSpeechDictionary.getInstance(),
+            0.70f
+        );
+    }
+
+    protected int channel() {
+        return channel;
+    }
+
+    @Override
+    public Block[] getKeys() {
+        return new Block[] { outputPartial ? buildIntermediateBlock() : buildFinalBlock() };
+    }
+
+    @Override
+    public IntVector nonEmpty() {
+        return IntVector.range(0, categorizer.getCategoryCount(), blockFactory);
+    }
+
+    @Override
+    public BitArray seenGroupIds(BigArrays bigArrays) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public final ReleasableIterator<IntBlock> lookup(Page page, ByteSizeValue targetBlockSize) {
+        throw new UnsupportedOperationException();
+    }
+
+    private Block buildIntermediateBlock() {
+        if (categorizer.getCategoryCount() == 0) {
+            return blockFactory.newConstantNullBlock(categorizer.getCategoryCount());
+        }
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            // TODO be more careful here.
+            out.writeVInt(categorizer.getCategoryCount());
+            for (SerializableTokenListCategory category : categorizer.toCategoriesById()) {
+                category.writeTo(out);
+            }
+            return blockFactory.newConstantBytesRefBlockWith(out.bytes().toBytesRef(), categorizer.getCategoryCount());
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private Block buildFinalBlock() {
+        try (BytesRefVector.Builder result = blockFactory.newBytesRefVectorBuilder(categorizer.getCategoryCount())) {
+            BytesRefBuilder scratch = new BytesRefBuilder();
+            for (SerializableTokenListCategory category : categorizer.toCategoriesById()) {
+                scratch.copyChars(category.getRegex());
+                result.appendBytesRef(scratch.get());
+                scratch.clear();
+            }
+            return result.build().asBlock();
+        }
+    }
+}

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

@@ -13,6 +13,7 @@
 import org.elasticsearch.common.util.BytesRefHash;
 import org.elasticsearch.common.util.LongHash;
 import org.elasticsearch.common.util.LongLongHash;
+import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.aggregation.SeenGroupIds;
 import org.elasticsearch.compute.data.Block;
@@ -34,9 +35,7 @@
  * @see LongHash
  * @see BytesRefHash
  */
-public abstract sealed class BlockHash implements Releasable, SeenGroupIds //
-    permits BooleanBlockHash, BytesRefBlockHash, DoubleBlockHash, IntBlockHash, LongBlockHash, BytesRef3BlockHash, //
-    NullBlockHash, PackedValuesBlockHash, BytesRefLongBlockHash, LongLongBlockHash, TimeSeriesBlockHash {
+public abstract class BlockHash implements Releasable, SeenGroupIds {
 
     protected final BlockFactory blockFactory;
 
@@ -94,9 +93,15 @@ public record GroupSpec(int channel, ElementType elementType) {}
      *                                 production until we can. And this lets us continue to compile and
      *                                 test them.
      */
-    public static BlockHash build(List<GroupSpec> groups, BlockFactory blockFactory, int emitBatchSize, boolean allowBrokenOptimizations) {
+    public static BlockHash build(
+        List<GroupSpec> groups,
+        AggregatorMode aggregatorMode,
+        BlockFactory blockFactory,
+        int emitBatchSize,
+        boolean allowBrokenOptimizations
+    ) {
         if (groups.size() == 1) {
-            return newForElementType(groups.get(0).channel(), groups.get(0).elementType(), blockFactory);
+            return newForElementType(groups.get(0).channel(), groups.get(0).elementType(), aggregatorMode, blockFactory);
         }
         if (groups.size() == 3 && groups.stream().allMatch(g -> g.elementType == ElementType.BYTES_REF)) {
             return new BytesRef3BlockHash(blockFactory, groups.get(0).channel, groups.get(1).channel, groups.get(2).channel, emitBatchSize);
@@ -127,14 +132,17 @@ public static BlockHash buildPackedValuesBlockHash(List<GroupSpec> groups, Block
     /**
      * Creates a specialized hash table that maps a {@link Block} of the given input element type to ids.
      */
-    private static BlockHash newForElementType(int channel, ElementType type, BlockFactory blockFactory) {
+    private static BlockHash newForElementType(int channel, ElementType type, AggregatorMode aggregatorMode, BlockFactory blockFactory) {
         return switch (type) {
             case NULL -> new NullBlockHash(channel, blockFactory);
             case BOOLEAN -> new BooleanBlockHash(channel, blockFactory);
             case INT -> new IntBlockHash(channel, blockFactory);
             case LONG -> new LongBlockHash(channel, blockFactory);
             case DOUBLE -> new DoubleBlockHash(channel, blockFactory);
             case BYTES_REF -> new BytesRefBlockHash(channel, blockFactory);
+            case CATEGORY -> aggregatorMode.isInputPartial()
+                ? new CategorizedIntermediateBlockHash(channel, blockFactory, aggregatorMode.isOutputPartial())
+                : new CategorizeRawBlockHash(channel, blockFactory, aggregatorMode.isOutputPartial());
             default -> throw new IllegalArgumentException("unsupported grouping element type [" + type + "]");
         };
     }

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

@@ -0,0 +1,137 @@
+/*
+ * 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.analysis.core.WhitespaceTokenizer;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
+import org.elasticsearch.compute.ann.Fixed;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
+import org.elasticsearch.compute.data.IntBlock;
+import org.elasticsearch.compute.data.IntVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.index.analysis.CharFilterFactory;
+import org.elasticsearch.index.analysis.CustomAnalyzer;
+import org.elasticsearch.index.analysis.TokenFilterFactory;
+import org.elasticsearch.index.analysis.TokenizerFactory;
+import org.elasticsearch.xpack.ml.aggs.categorization.TokenListCategorizer;
+import org.elasticsearch.xpack.ml.job.categorization.CategorizationAnalyzer;
+
+public class CategorizeRawBlockHash extends AbstractCategorizeBlockHash {
+    private final CategorizeEvaluator evaluator;
+
+    CategorizeRawBlockHash(int channel, BlockFactory blockFactory, boolean outputPartial) {
+        super(blockFactory, channel, outputPartial);
+        CategorizationAnalyzer analyzer = new CategorizationAnalyzer(
+            // TODO: should be the same analyzer as used in Production
+            new CustomAnalyzer(
+                TokenizerFactory.newFactory("whitespace", WhitespaceTokenizer::new),
+                new CharFilterFactory[0],
+                new TokenFilterFactory[0]
+            ),
+            true
+        );
+        this.evaluator = new CategorizeEvaluator(analyzer, categorizer, blockFactory);
+    }
+
+    @Override
+    public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
+        try (IntBlock result = (IntBlock) evaluator.eval(page.getBlock(channel()))) {
+            addInput.add(0, result);
+        }
+    }
+
+    @Override
+    public void close() {
+        evaluator.close();
+    }
+
+    /**
+     * NOCOMMIT: Super-duper copy-pasted from the actually generated evaluator; needs cleanup.
+     */
+    public static final class CategorizeEvaluator implements Releasable {
+        private final CategorizationAnalyzer analyzer;
+
+        private final TokenListCategorizer.CloseableTokenListCategorizer categorizer;
+
+        private final BlockFactory blockFactory;
+
+        static int process(
+            BytesRef v,
+            @Fixed(includeInToString = false, build = true) CategorizationAnalyzer analyzer,
+            @Fixed(includeInToString = false, build = true) TokenListCategorizer.CloseableTokenListCategorizer categorizer
+        ) {
+            return categorizer.computeCategory(v.utf8ToString(), analyzer).getId();
+        }
+
+        public CategorizeEvaluator(
+            CategorizationAnalyzer analyzer,
+            TokenListCategorizer.CloseableTokenListCategorizer categorizer,
+            BlockFactory blockFactory
+        ) {
+            this.analyzer = analyzer;
+            this.categorizer = categorizer;
+            this.blockFactory = blockFactory;
+        }
+
+        public Block eval(BytesRefBlock vBlock) {
+            BytesRefVector vVector = vBlock.asVector();
+            if (vVector == null) {
+                return eval(vBlock.getPositionCount(), vBlock);
+            }
+            IntVector vector = eval(vBlock.getPositionCount(), vVector);
+            return vector.asBlock();
+        }
+
+        public IntBlock eval(int positionCount, BytesRefBlock vBlock) {
+            try (IntBlock.Builder result = blockFactory.newIntBlockBuilder(positionCount)) {
+                BytesRef vScratch = new BytesRef();
+                position: for (int p = 0; p < positionCount; p++) {
+                    if (vBlock.isNull(p)) {
+                        result.appendNull();
+                        continue position;
+                    }
+                    if (vBlock.getValueCount(p) != 1) {
+                        if (vBlock.getValueCount(p) > 1) {
+                            // TODO: handle multi-values
+                        }
+                        result.appendNull();
+                        continue position;
+                    }
+                    result.appendInt(process(vBlock.getBytesRef(vBlock.getFirstValueIndex(p), vScratch), this.analyzer, this.categorizer));
+                }
+                return result.build();
+            }
+        }
+
+        public IntVector eval(int positionCount, BytesRefVector vVector) {
+            try (IntVector.FixedBuilder result = blockFactory.newIntVectorFixedBuilder(positionCount)) {
+                BytesRef vScratch = new BytesRef();
+                position: for (int p = 0; p < positionCount; p++) {
+                    result.appendInt(p, process(vVector.getBytesRef(p, vScratch), this.analyzer, this.categorizer));
+                }
+                return result.build();
+            }
+        }
+
+        @Override
+        public String toString() {
+            return "CategorizeEvaluator";
+        }
+
+        @Override
+        public void close() {
+            Releasables.closeExpectNoException(analyzer, categorizer);
+        }
+    }
+}