Unit test for CategorizedIntermediateBlockHash.

Author: jan-elastic

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

@@ -9,9 +9,12 @@
 
 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.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;
 import org.elasticsearch.xpack.ml.aggs.categorization.SerializableTokenListCategory;
@@ -44,33 +47,47 @@ protected int channel() {
     @Override
     public Block[] getKeys() {
         if (outputPartial) {
+            return new Block[] { buildIntermediateBlock() };
             // NOCOMMIT load partial
-            Block state = null;
-            Block keys; // NOCOMMIT do we even need to send the keys? it's just going to be 0 to the length of state
+            // Block state = null;
+            // Block keys; // NOCOMMIT do we even need to send the keys? it's just going to be 0 to the length of state
             // return new Block[] {new CompositeBlock()};
-            return null;
+            // return null;
         }
 
         // NOCOMMIT load final
         return new Block[0];
     }
 
+    @Override
+    public IntVector nonEmpty() {
+        // TODO
+        return null;
+    }
+
+    @Override
+    public BitArray seenGroupIds(BigArrays bigArrays) {
+        // TODO
+        return null;
+    }
+
     @Override
     public final ReleasableIterator<IntBlock> lookup(Page page, ByteSizeValue targetBlockSize) {
         throw new UnsupportedOperationException();
     }
 
-    private Block buildIntermediateBlock(BlockFactory blockFactory, int positionCount) {
+    // visible for testing
+    Block buildIntermediateBlock() {
         if (categorizer.getCategoryCount() == 0) {
-            return blockFactory.newConstantNullBlock(positionCount);
+            return blockFactory.newConstantNullBlock(1);
         }
         try (BytesStreamOutput out = new BytesStreamOutput()) {
             // TODO be more careful here.
             out.writeVInt(categorizer.getCategoryCount());
             for (SerializableTokenListCategory category : categorizer.toCategories(categorizer.getCategoryCount())) {
                 category.writeTo(out);
             }
-            return blockFactory.newConstantBytesRefBlockWith(out.bytes().toBytesRef(), positionCount);
+            return blockFactory.newConstantBytesRefBlockWith(out.bytes().toBytesRef(), 1);
         } catch (IOException e) {
             throw new RuntimeException(e);
         }

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

@@ -7,10 +7,7 @@
 
 package org.elasticsearch.compute.aggregation.blockhash;
 
-import org.apache.lucene.analysis.TokenStream;
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.util.BigArrays;
-import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.ann.Fixed;
 import org.elasticsearch.compute.data.Block;
@@ -25,8 +22,6 @@
 import org.elasticsearch.xpack.ml.aggs.categorization.TokenListCategorizer;
 import org.elasticsearch.xpack.ml.job.categorization.CategorizationAnalyzer;
 
-import java.io.IOException;
-
 public class CategorizeRawBlockHash extends AbstractCategorizeBlockHash {
     private final CategorizeEvaluator evaluator;
 
@@ -47,18 +42,6 @@ public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
         addInput.add(0, result);
     }
 
-    @Override
-    public IntVector nonEmpty() {
-        // TODO
-        return null;
-    }
-
-    @Override
-    public BitArray seenGroupIds(BigArrays bigArrays) {
-        // TODO
-        return null;
-    }
-
     @Override
     public void close() {
         evaluator.close();
@@ -79,12 +62,7 @@ static int process(
             @Fixed(includeInToString = false, build = true) CategorizationAnalyzer analyzer,
             @Fixed(includeInToString = false, build = true) TokenListCategorizer.CloseableTokenListCategorizer categorizer
         ) {
-            String s = v.utf8ToString();
-            try (TokenStream ts = analyzer.tokenStream("text", s)) {
-                return categorizer.computeCategory(ts, s.length(), 1).getId();
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
+            return categorizer.computeCategory(v.utf8ToString(), analyzer).getId();
         }
 
         public CategorizeEvaluator(

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

@@ -10,14 +10,10 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.util.BigArrays;
-import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BytesRefBlock;
-import org.elasticsearch.compute.data.CompositeBlock;
 import org.elasticsearch.compute.data.IntBlock;
-import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.xpack.ml.aggs.categorization.SerializableTokenListCategory;
 import org.elasticsearch.xpack.ml.aggs.categorization.TokenListCategorizer;
@@ -32,27 +28,24 @@ public class CategorizedIntermediateBlockHash extends AbstractCategorizeBlockHas
 
     CategorizedIntermediateBlockHash(
         BlockFactory blockFactory,
+        int channel,
         boolean outputPartial,
-        TokenListCategorizer.CloseableTokenListCategorizer categorizer,
-        IntBlockHash hash,
-        int channel
+        TokenListCategorizer.CloseableTokenListCategorizer categorizer
     ) {
         super(blockFactory, channel, outputPartial, categorizer);
-        this.hash = hash;
+        this.hash = new IntBlockHash(channel, blockFactory);
     }
 
     public void add(Page page, GroupingAggregatorFunction.AddInput addInput) {
-        CompositeBlock block = page.getBlock(channel());
-        BytesRefBlock groupingState = block.getBlock(0);
-        BytesRefBlock groups = block.getBlock(0);
+        BytesRefBlock categorizerState = page.getBlock(channel());
         Map<Integer, Integer> idMap;
-        if (groupingState.areAllValuesNull() == false) {
-            idMap = readIntermediate(groupingState.getBytesRef(0, new BytesRef()));
+        if (categorizerState.areAllValuesNull() == false) {
+            idMap = readIntermediate(categorizerState.getBytesRef(0, new BytesRef()));
         } else {
             idMap = Collections.emptyMap();
         }
-        try (IntBlock.Builder newIdsBuilder = blockFactory.newIntBlockBuilder(groups.getTotalValueCount())) {
-            for (int i = 0; i < groups.getTotalValueCount(); i++) {
+        try (IntBlock.Builder newIdsBuilder = blockFactory.newIntBlockBuilder(idMap.size())) {
+            for (int i = 0; i < idMap.size(); i++) {
                 newIdsBuilder.appendInt(idMap.get(i));
             }
             IntBlock newIds = newIdsBuilder.build();
@@ -76,18 +69,9 @@ private Map<Integer, Integer> readIntermediate(BytesRef bytes) {
         }
     }
 
-    @Override
-    public IntVector nonEmpty() {
-        return hash.nonEmpty();
-    }
-
-    @Override
-    public BitArray seenGroupIds(BigArrays bigArrays) {
-        return hash.seenGroupIds(bigArrays);
-    }
-
     @Override
     public void close() {
-
+        categorizer.close();
+        hash.close();
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java

@@ -10,14 +10,8 @@
 import com.carrotsearch.randomizedtesting.annotations.Name;
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import org.apache.lucene.analysis.core.WhitespaceTokenizer;
 import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.unit.ByteSizeValue;
-import org.elasticsearch.common.util.BigArrays;
-import org.elasticsearch.common.util.BytesRefHash;
-import org.elasticsearch.common.util.MockBigArrays;
-import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BooleanBlock;
@@ -27,22 +21,11 @@
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.IntVector;
 import org.elasticsearch.compute.data.LongBlock;
-import org.elasticsearch.compute.data.MockBlockFactory;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.data.TestBlockFactory;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.ReleasableIterator;
 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.indices.breaker.CircuitBreakerService;
-import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.xpack.ml.aggs.categorization.CategorizationBytesRefHash;
-import org.elasticsearch.xpack.ml.aggs.categorization.CategorizationPartOfSpeechDictionary;
-import org.elasticsearch.xpack.ml.aggs.categorization.TokenListCategorizer;
-import org.elasticsearch.xpack.ml.job.categorization.CategorizationAnalyzer;
 import org.junit.After;
 
 import java.util.ArrayList;
@@ -61,8 +44,6 @@
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.startsWith;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 public class BlockHashTests extends BlockHashTestCase {