fix: Static aggregations when a data index uses different names for the key columns should remap names correctly (#6360)

Fixes #6359

Author: rcaudy

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<Ch
 
     @NotNull
     private static QueryTable staticIndexedAggregation(
+            final QueryTable input,
             final BasicDataIndex dataIndex,
             final String[] keyNames,
             final AggregationContext ac) {
@@ -1628,10 +1628,22 @@ private static QueryTable staticIndexedAggregation(
 
         final Map<String, ColumnSource<?>> 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<String, ColumnSource<?>> inputNamesToSources = input.getColumnSourceMap();
+        final Map<ColumnSource<?>, 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<Object> lookupKeyToRowKey =
-                    DataIndexUtils.buildRowKeyMappingFunction(indexKeyTable, keyNames);
+                    DataIndexUtils.buildRowKeyMappingFunction(indexKeyTable, dataIndexKeyNames);
             return key -> (int) lookupKeyToRowKey.applyAsLong(key);
         });
 

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