ES|QL categorize options

Author: jan-elastic

docs/reference/query-languages/esql/_snippets/functions/functionNamedParams/categorize.md

@@ -338,6 +338,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_FIXED_INDEX_LIKE = def(9_119_0_00);
     public static final TransportVersion LOOKUP_JOIN_CCS = def(9_120_0_00);
     public static final TransportVersion NODE_USAGE_STATS_FOR_THREAD_POOLS_IN_CLUSTER_INFO = def(9_121_0_00);
+    public static final TransportVersion ESQL_CATEGORIZE_OPTIONS = def(9_122_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

docs/reference/query-languages/esql/_snippets/functions/layout/categorize.md

@@ -128,16 +128,26 @@ public abstract class BlockHash implements Releasable, SeenGroupIds {
     public record TopNDef(int order, boolean asc, boolean nullsFirst, int limit) {}
 
     /**
-     * @param isCategorize Whether this group is a CATEGORIZE() or not.
-     *                     May be changed in the future when more stateful grouping functions are added.
+     * Configuration for a BlockHash group spec that is doing text categorization.
      */
-    public record GroupSpec(int channel, ElementType elementType, boolean isCategorize, @Nullable TopNDef topNDef) {
+    public record CategorizeDef(String analyzer, OutputFormat outputFormat, int similarityThreshold) {
+        public enum OutputFormat {
+            REGEX,
+            TOKENS
+        }
+    }
+
+    public record GroupSpec(int channel, ElementType elementType, @Nullable CategorizeDef categorizeDef, @Nullable TopNDef topNDef) {
         public GroupSpec(int channel, ElementType elementType) {
-            this(channel, elementType, false, null);
+            this(channel, elementType, null, null);
+        }
+
+        public GroupSpec(int channel, ElementType elementType, CategorizeDef categorizeDef) {
+            this(channel, elementType, categorizeDef, null);
         }
 
-        public GroupSpec(int channel, ElementType elementType, boolean isCategorize) {
-            this(channel, elementType, isCategorize, null);
+        public boolean isCategorize() {
+            return categorizeDef != null;
         }
     }
 
@@ -207,7 +217,13 @@ public static BlockHash buildCategorizeBlockHash(
         int emitBatchSize
     ) {
         if (groups.size() == 1) {
-            return new CategorizeBlockHash(blockFactory, groups.get(0).channel, aggregatorMode, analysisRegistry);
+            return new CategorizeBlockHash(
+                blockFactory,
+                groups.get(0).channel,
+                aggregatorMode,
+                groups.get(0).categorizeDef,
+                analysisRegistry
+            );
         } else {
             assert groups.get(0).isCategorize();
             assert groups.subList(1, groups.size()).stream().noneMatch(GroupSpec::isCategorize);

docs/reference/query-languages/esql/_snippets/functions/parameters/categorize.md

@@ -18,7 +18,6 @@
 import org.elasticsearch.common.util.BytesRefHash;
 import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
-import org.elasticsearch.compute.aggregation.SeenGroupIds;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
@@ -47,12 +46,13 @@
  */
 public class CategorizeBlockHash extends BlockHash {
 
-    private static final CategorizationAnalyzerConfig ANALYZER_CONFIG = CategorizationAnalyzerConfig
+    private static final CategorizationAnalyzerConfig DEFAULT_ANALYZER_CONFIG = CategorizationAnalyzerConfig
         .buildStandardEsqlCategorizationAnalyzer();
     private static final int NULL_ORD = 0;
 
     private final int channel;
     private final AggregatorMode aggregatorMode;
+    private final CategorizeDef categorizeDef;
     private final TokenListCategorizer.CloseableTokenListCategorizer categorizer;
     private final CategorizeEvaluator evaluator;
 
@@ -64,28 +64,38 @@ public class CategorizeBlockHash extends BlockHash {
      */
     private boolean seenNull = false;
 
-    CategorizeBlockHash(BlockFactory blockFactory, int channel, AggregatorMode aggregatorMode, AnalysisRegistry analysisRegistry) {
+    CategorizeBlockHash(
+        BlockFactory blockFactory,
+        int channel,
+        AggregatorMode aggregatorMode,
+        CategorizeDef categorizeDef,
+        AnalysisRegistry analysisRegistry
+    ) {
         super(blockFactory);
 
         this.channel = channel;
         this.aggregatorMode = aggregatorMode;
+        this.categorizeDef = categorizeDef;
 
         this.categorizer = new TokenListCategorizer.CloseableTokenListCategorizer(
             new CategorizationBytesRefHash(new BytesRefHash(2048, blockFactory.bigArrays())),
             CategorizationPartOfSpeechDictionary.getInstance(),
-            0.70f
+            categorizeDef.similarityThreshold() / 100.0f
         );
 
         if (aggregatorMode.isInputPartial() == false) {
-            CategorizationAnalyzer analyzer;
+            CategorizationAnalyzer categorizationAnalyzer;
             try {
                 Objects.requireNonNull(analysisRegistry);
-                analyzer = new CategorizationAnalyzer(analysisRegistry, ANALYZER_CONFIG);
-            } catch (Exception e) {
+                CategorizationAnalyzerConfig config = categorizeDef.analyzer() == null
+                    ? DEFAULT_ANALYZER_CONFIG
+                    : new CategorizationAnalyzerConfig.Builder().setAnalyzer(categorizeDef.analyzer()).build();
+                categorizationAnalyzer = new CategorizationAnalyzer(analysisRegistry, config);
+            } catch (IOException e) {
                 categorizer.close();
                 throw new RuntimeException(e);
             }
-            this.evaluator = new CategorizeEvaluator(analyzer);
+            this.evaluator = new CategorizeEvaluator(categorizationAnalyzer);
         } else {
             this.evaluator = null;
         }
@@ -114,7 +124,7 @@ public IntVector nonEmpty() {
 
     @Override
     public BitArray seenGroupIds(BigArrays bigArrays) {
-        return new SeenGroupIds.Range(seenNull ? 0 : 1, Math.toIntExact(categorizer.getCategoryCount() + 1)).seenGroupIds(bigArrays);
+        return new Range(seenNull ? 0 : 1, Math.toIntExact(categorizer.getCategoryCount() + 1)).seenGroupIds(bigArrays);
     }
 
     @Override
@@ -222,7 +232,7 @@ private Block buildFinalBlock() {
             try (BytesRefBlock.Builder result = blockFactory.newBytesRefBlockBuilder(categorizer.getCategoryCount())) {
                 result.appendNull();
                 for (SerializableTokenListCategory category : categorizer.toCategoriesById()) {
-                    scratch.copyChars(category.getRegex());
+                    scratch.copyChars(getKeyString(category));
                     result.appendBytesRef(scratch.get());
                     scratch.clear();
                 }
@@ -232,14 +242,21 @@ private Block buildFinalBlock() {
 
         try (BytesRefVector.Builder result = blockFactory.newBytesRefVectorBuilder(categorizer.getCategoryCount())) {
             for (SerializableTokenListCategory category : categorizer.toCategoriesById()) {
-                scratch.copyChars(category.getRegex());
+                scratch.copyChars(getKeyString(category));
                 result.appendBytesRef(scratch.get());
                 scratch.clear();
             }
             return result.build().asBlock();
         }
     }
 
+    private String getKeyString(SerializableTokenListCategory category) {
+        return switch (categorizeDef.outputFormat()) {
+            case REGEX -> category.getRegex();
+            case TOKENS -> category.getKeyTokensString();
+        };
+    }
+
     /**
      * Similar implementation to an Evaluator.
      */

docs/reference/query-languages/esql/_snippets/functions/types/categorize.md

@@ -68,7 +68,13 @@ public class CategorizePackedValuesBlockHash extends BlockHash {
 
         boolean success = false;
         try {
-            categorizeBlockHash = new CategorizeBlockHash(blockFactory, specs.get(0).channel(), aggregatorMode, analysisRegistry);
+            categorizeBlockHash = new CategorizeBlockHash(
+                blockFactory,
+                specs.get(0).channel(),
+                aggregatorMode,
+                specs.get(0).categorizeDef(),
+                analysisRegistry
+            );
             packedValuesBlockHash = new PackedValuesBlockHash(delegateSpecs, blockFactory, emitBatchSize);
             success = true;
         } finally {

docs/reference/query-languages/esql/images/functions/categorize.svg

@@ -63,6 +63,12 @@
 
 public class CategorizeBlockHashTests extends BlockHashTestCase {
 
+    private static final BlockHash.CategorizeDef CATEGORIZE_DEF = new BlockHash.CategorizeDef(
+        null,
+        BlockHash.CategorizeDef.OutputFormat.REGEX,
+        70
+    );
+
     private AnalysisRegistry analysisRegistry;
 
     @Before
@@ -98,7 +104,7 @@ public void testCategorizeRaw() {
             page = new Page(builder.build());
         }
 
-        try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, analysisRegistry)) {
+        try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, CATEGORIZE_DEF, analysisRegistry)) {
             for (int i = randomInt(2); i < 3; i++) {
                 hash.add(page, new GroupingAggregatorFunction.AddInput() {
                     private void addBlock(int positionOffset, IntBlock groupIds) {
@@ -170,7 +176,7 @@ public void testCategorizeRawMultivalue() {
             page = new Page(builder.build());
         }
 
-        try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, analysisRegistry)) {
+        try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, CATEGORIZE_DEF, analysisRegistry)) {
             for (int i = randomInt(2); i < 3; i++) {
                 hash.add(page, new GroupingAggregatorFunction.AddInput() {
                     private void addBlock(int positionOffset, IntBlock groupIds) {
@@ -259,8 +265,8 @@ public void testCategorizeIntermediate() {
 
         // Fill intermediatePages with the intermediate state from the raw hashes
         try (
-            BlockHash rawHash1 = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.INITIAL, analysisRegistry);
-            BlockHash rawHash2 = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.INITIAL, analysisRegistry);
+            BlockHash rawHash1 = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.INITIAL, CATEGORIZE_DEF, analysisRegistry);
+            BlockHash rawHash2 = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.INITIAL, CATEGORIZE_DEF, analysisRegistry);
         ) {
             rawHash1.add(page1, new GroupingAggregatorFunction.AddInput() {
                 private void addBlock(int positionOffset, IntBlock groupIds) {
@@ -335,7 +341,7 @@ public void close() {
             page2.releaseBlocks();
         }
 
-        try (var intermediateHash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.FINAL, null)) {
+        try (var intermediateHash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.FINAL, CATEGORIZE_DEF, null)) {
             intermediateHash.add(intermediatePage1, new GroupingAggregatorFunction.AddInput() {
                 private void addBlock(int positionOffset, IntBlock groupIds) {
                     List<Integer> values = IntStream.range(0, groupIds.getPositionCount())
@@ -590,7 +596,7 @@ public void testCategorize_withDriver() {
     }
 
     private BlockHash.GroupSpec makeGroupSpec() {
-        return new BlockHash.GroupSpec(0, ElementType.BYTES_REF, true);
+        return new BlockHash.GroupSpec(0, ElementType.BYTES_REF, CATEGORIZE_DEF);
     }
 
     private void assertHashState(CategorizeBlockHash hash, boolean withNull, String... expectedKeys) {