diff --git a/engine/table/src/main/java/io/deephaven/engine/table/impl/by/ChunkedOperatorAggregationHelper.java b/engine/table/src/main/java/io/deephaven/engine/table/impl/by/ChunkedOperatorAggregationHelper.java index 568511176ac..6df7b5a9e6f 100644 --- a/engine/table/src/main/java/io/deephaven/engine/table/impl/by/ChunkedOperatorAggregationHelper.java +++ b/engine/table/src/main/java/io/deephaven/engine/table/impl/by/ChunkedOperatorAggregationHelper.java @@ -163,14 +163,13 @@ private static QueryTable aggregation( final PermuteKernel[] permuteKernels = ac.makePermuteKernels(); if (dataIndex != null && initialKeys == null && !input.isRefreshing()) { - return staticIndexedAggregation(dataIndex, keyNames, ac); + return staticIndexedAggregation(input, dataIndex, keyNames, ac); } final Table symbolTable; final boolean useSymbolTable; if (!input.isRefreshing() && control.considerSymbolTables(input, dataIndex != null, keySources)) { Assert.eq(keySources.length, "keySources.length", 1); - symbolTable = ((SymbolTableSource) keySources[0]).getStaticSymbolTable(input.getRowSet(), control.useSymbolTableLookupCaching()); useSymbolTable = control.useSymbolTables(input.size(), symbolTable.size()); @@ -1619,6 +1618,7 @@ private static void modifySlots(RowSetBuilderRandom modifiedBuilder, IntChunk> resultColumnSourceMap = new LinkedHashMap<>(); - // Add the index key columns directly to the result table. - final Table indexKeyTable = indexTable.flatten().select(keyNames); - for (final String keyName : keyNames) { - resultColumnSourceMap.put(keyName, indexKeyTable.getColumnSource(keyName)); + // Add the index table's key columns directly to the result table (after flattening and selecting). + final Table indexKeyTable = indexTable.flatten().select(dataIndex.keyColumnNames().toArray(String[]::new)); + // Index key names corresponding to the input key names in the order of the input key names + final String[] dataIndexKeyNames = new String[keyNames.length]; + final Map> inputNamesToSources = input.getColumnSourceMap(); + final Map, String> indexedSourcesToNames = dataIndex.keyColumnNamesByIndexedColumn(); + for (int ki = 0; ki < keyNames.length; ki++) { + final String keyName = keyNames[ki]; + // Note that we must be careful to find the index key column sources by indexed (input) column source, not + // by input key name, as the input key names may not match the indexed column names. + // We build dataIndexKeyNames as part of this mapping, so we can use them for the row lookup remapping. + final ColumnSource keyColumnSource = inputNamesToSources.get(keyName); + final String dataIndexKeyName = indexedSourcesToNames.get(keyColumnSource); + dataIndexKeyNames[ki] = dataIndexKeyName; + final ColumnSource indexKeyColumnSource = indexKeyTable.getColumnSource(dataIndexKeyName); + resultColumnSourceMap.put(keyName, indexKeyColumnSource); } ac.getResultColumns(resultColumnSourceMap); @@ -1648,7 +1660,7 @@ private static QueryTable staticIndexedAggregation( // Create a lookup function from the index table ac.supplyRowLookup(() -> { final ToLongFunction lookupKeyToRowKey = - DataIndexUtils.buildRowKeyMappingFunction(indexKeyTable, keyNames); + DataIndexUtils.buildRowKeyMappingFunction(indexKeyTable, dataIndexKeyNames); return key -> (int) lookupKeyToRowKey.applyAsLong(key); }); diff --git a/engine/table/src/test/java/io/deephaven/engine/table/impl/QueryTableAggregationTest.java b/engine/table/src/test/java/io/deephaven/engine/table/impl/QueryTableAggregationTest.java index a9aa3a9c392..296d3e257e3 100644 --- a/engine/table/src/test/java/io/deephaven/engine/table/impl/QueryTableAggregationTest.java +++ b/engine/table/src/test/java/io/deephaven/engine/table/impl/QueryTableAggregationTest.java @@ -198,13 +198,26 @@ public int initialHashTableSize(@NotNull final Table table) { @Test public void testStaticGroupedByWithChunks() { - final Table input1 = emptyTable(10000).update("A=Integer.toString(i % 5)", "B=i / 5"); + final Table input1 = emptyTable(10000).update("A=Integer.toString(i % 5)", "B=i / 5", "C=ii"); DataIndexer.getOrCreateDataIndex(input1, "A"); DataIndexer.getOrCreateDataIndex(input1, "B"); + DataIndexer.getOrCreateDataIndex(input1, "A", "B"); individualStaticByTest(input1, null, "A"); individualStaticByTest(input1, null, "B"); + individualStaticByTest(input1, null, "A", "B"); + + // Test scenario where the key columns have different order than in the index + individualStaticByTest(input1, null, "B", "A"); + + // Test scenarios where the key columns have different names than in the index + individualStaticByTest(input1.renameColumns("D=A", "E=B"), null, "D"); + individualStaticByTest(input1.renameColumns("D=A", "E=B"), null, "E"); + individualStaticByTest(input1.renameColumns("D=A", "E=B"), null, "D", "E"); + + // Test scenario where the key columns have different names and order than in the index + individualStaticByTest(input1.renameColumns("D=A", "E=B"), null, "E", "D"); } @Test