fix: DH-21400 - Implement findPositionForKey for StaticSymbolTableChunkedOperatorAggregationStateManager (#7572)

This pull request implements the findPositionForKey method for StaticSymbolTableChunkedOperatorAggregationStateManager, replacing a previous UnsupportedOperationException with a working implementation. The implementation uses a lazily-initialized Trove map to map keys to their positions in the state manager. Additionally, two new test methods verify that symbol table-backed DataIndex lookups and filtering work correctly.

Author: lbooker42

engine/table/src/main/java/io/deephaven/engine/table/impl/by/StaticSymbolTableChunkedOperatorAggregationStateManager.java

@@ -3,18 +3,23 @@
 //
 package io.deephaven.engine.table.impl.by;
 
+import gnu.trove.map.hash.TObjectIntHashMap;
+import io.deephaven.base.verify.Assert;
 import io.deephaven.chunk.LongChunk;
 import io.deephaven.chunk.WritableIntChunk;
 import io.deephaven.chunk.WritableLongChunk;
 import io.deephaven.chunk.attributes.Values;
+import io.deephaven.engine.primitive.iterator.CloseableIterator;
 import io.deephaven.engine.rowset.RowSequence;
+import io.deephaven.engine.rowset.RowSequenceFactory;
 import io.deephaven.engine.rowset.chunkattributes.RowKeys;
 import io.deephaven.engine.table.ChunkSource;
 import io.deephaven.engine.table.ColumnSource;
 import io.deephaven.engine.table.Table;
 import io.deephaven.engine.table.impl.SymbolTableToUniqueIdSource;
 import io.deephaven.engine.table.impl.sources.ObjectArraySource;
 import io.deephaven.engine.table.impl.sources.regioned.SymbolTableSource;
+import io.deephaven.engine.table.iterators.ChunkedObjectColumnIterator;
 import io.deephaven.util.QueryConstants;
 import io.deephaven.util.SafeCloseable;
 import io.deephaven.util.mutable.MutableInt;
@@ -34,6 +39,8 @@ public class StaticSymbolTableChunkedOperatorAggregationStateManager implements
     private final int[] keyPositions;
     private int nextPosition = 0;
 
+    private volatile TObjectIntHashMap<String> keyToPosition;
+
     StaticSymbolTableChunkedOperatorAggregationStateManager(final ColumnSource<?> keySource, final Table symbolTable) {
         this.symbolTable = symbolTable;
         tableSize = symbolTable.intSize();
@@ -110,6 +117,7 @@ public void add(final SafeCloseable bc, final RowSequence rowSequence, final Col
     private void updateKeyHashTableSources(final WritableLongChunk<RowKeys> symbolTableValues,
             final int firstNewPosition) {
         keyColumn.ensureCapacity(nextPosition);
+        Assert.eqNull(keyToPosition, "keyToPosition");
 
         final ColumnSource<?> symbolColumnSource = symbolTable.getColumnSource(SymbolTableSource.SYMBOL_COLUMN_NAME);
 
@@ -135,8 +143,23 @@ public ColumnSource<?>[] getKeyHashTableSources() {
 
     @Override
     public int findPositionForKey(final Object key) {
-        // shouldn't be able to get here; rollup/treeview will call this when we're 2+ levels deep in out view. since
-        // we're limited to a single keySource, we cannot be more than 1 level deep
-        throw new UnsupportedOperationException("StaticSymbolTable StateManager must be used with a single keySource");
+        // Build the map if it doesn't exist
+        TObjectIntHashMap<String> localKeyToPosition;
+        if ((localKeyToPosition = keyToPosition) == null) {
+            synchronized (this) {
+                if ((localKeyToPosition = keyToPosition) == null) {
+                    final int length = nextPosition;
+                    localKeyToPosition = new TObjectIntHashMap<>(length, 0.75f, UNKNOWN_ROW);
+                    try (final CloseableIterator<String> keyIterator = new ChunkedObjectColumnIterator<>(
+                            keyColumn, RowSequenceFactory.forRange(0, length - 1))) {
+                        for (int ii = 0; ii < length; ii++) {
+                            localKeyToPosition.put(keyIterator.next(), ii);
+                        }
+                    }
+                    keyToPosition = localKeyToPosition;
+                }
+            }
+        }
+        return localKeyToPosition.get(key);
     }
 }

extensions/parquet/table/src/test/java/io/deephaven/parquet/table/TestSymbolTableSource.java

@@ -4,7 +4,10 @@
 package io.deephaven.parquet.table;
 
 import io.deephaven.base.FileUtils;
+import io.deephaven.engine.table.DataIndex;
 import io.deephaven.engine.table.Table;
+import io.deephaven.engine.table.impl.dataindex.TableBackedDataIndex;
+import io.deephaven.engine.table.impl.indexer.DataIndexer;
 import io.deephaven.engine.table.impl.sources.regioned.SymbolTableSource;
 import io.deephaven.engine.testutil.junit4.EngineCleanup;
 import io.deephaven.engine.util.TableTools;
@@ -61,4 +64,66 @@ public void testWriteAndReadSymbols() {
 
         assertTableEquals(expected, syms);
     }
+
+    @Test
+    public void testSymbolTableDataIndexLookup() {
+        final Table t = TableTools.emptyTable(100)
+                .update("TheBestColumn=i==9?(String)null:`S`+(i%10)", "Sentinel=i");
+        final File toWrite = new File(dataDirectory, "table.parquet");
+        ParquetTools.writeTable(t, toWrite.getPath());
+
+        // Make sure we have the expected symbol table (or not)
+        final Table readBack = ParquetTools.readTable(toWrite.getPath());
+        final SymbolTableSource<String> source =
+                (SymbolTableSource<String>) readBack.getColumnSource("TheBestColumn", String.class);
+        Assert.assertTrue(source.hasSymbolTable(readBack.getRowSet()));
+
+        final DataIndex index = DataIndexer.getOrCreateDataIndex(readBack, "TheBestColumn");
+        Assert.assertTrue("index instanceof TableBackedDataIndex", index instanceof TableBackedDataIndex);
+        final DataIndex.RowKeyLookup rkl = index.rowKeyLookup();
+
+        for (int i = 0; i < 9; i++) {
+            final String key = "S" + i;
+            final long rowKey = rkl.apply(key, false);
+            Assert.assertEquals(i, rowKey);
+        }
+        // Assert null lookup is correct.
+        final long rowKey = rkl.apply(null, false);
+        Assert.assertEquals(9, rowKey);
+    }
+
+    /**
+     * This won't fail after 41.0 due to the table-level filtering replacing AbstractColumnSource#match. Used to test
+     * bugfix against 0.40.x.
+     */
+    @Test
+    public void testFilterIndexedSymbolTable() {
+        final Table t = TableTools.emptyTable(100)
+                .update("TheBestColumn=i==9?(String)null:`S`+(i%10)", "Sentinel=i");
+        final File toWrite = new File(dataDirectory, "table.parquet");
+        ParquetTools.writeTable(t, toWrite.getPath());
+
+        // Make sure we have the expected symbol table (or not)
+        final Table readBack = ParquetTools.readTable(toWrite.getPath());
+        final SymbolTableSource<String> source =
+                (SymbolTableSource<String>) readBack.getColumnSource("TheBestColumn", String.class);
+        Assert.assertTrue(source.hasSymbolTable(readBack.getRowSet()));
+
+        final DataIndex index = DataIndexer.getOrCreateDataIndex(readBack, "TheBestColumn");
+        Assert.assertTrue("index instanceof TableBackedDataIndex", index instanceof TableBackedDataIndex);
+        // materialize the index table
+        final Table indexTable = index.table();
+
+        Table filtered;
+        Table expected;
+
+        filtered = readBack.where("TheBestColumn in `S0`");
+        expected = TableTools.emptyTable(10).update("TheBestColumn=`S0`", "Sentinel=i*10");
+        assertTableEquals(expected, filtered);
+
+        // Assert null filtering is correct.
+        filtered = readBack.where("TheBestColumn in null");
+        expected = TableTools.emptyTable(1).update("TheBestColumn=(String)null", "Sentinel=9");
+        assertTableEquals(expected, filtered);
+    }
 }