ES|QL categorize options

Author: jan-elastic

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/AggregatorBenchmark.java

@@ -191,7 +191,7 @@ private static Operator operator(DriverContext driverContext, String grouping, S
                 new BlockHash.GroupSpec(2, ElementType.BYTES_REF)
             );
             case TOP_N_LONGS -> List.of(
-                new BlockHash.GroupSpec(0, ElementType.LONG, false, new BlockHash.TopNDef(0, true, true, TOP_N_LIMIT))
+                new BlockHash.GroupSpec(0, ElementType.LONG, null, new BlockHash.TopNDef(0, true, true, TOP_N_LIMIT))
             );
             default -> throw new IllegalArgumentException("unsupported grouping [" + grouping + "]");
         };

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 {