Fix wrong results for multi-key GROUP BY with transform expressions

pinot-core/src/main/java/org/apache/pinot/core/query/utils/OrderByComparatorFactory.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.core.query.utils;
 
+import com.google.common.base.Preconditions;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashMap;
@@ -70,11 +71,15 @@ public static Comparator<Record> getRecordKeyComparator(List<OrderByExpressionCo
     return (k1, k2) -> valueComparator.compare(k1.getValues(), k2.getValues());
   }
 
-  private static Map<String, Integer> getGroupByExpressionIndexMap(List<ExpressionContext> groupByExpressions) {
-    Map<String, Integer> groupByExpressionIndexMap = new HashMap<>();
+  private static Map<ExpressionContext, Integer> getGroupByExpressionIndexMap(
+      List<ExpressionContext> groupByExpressions) {
+    Map<ExpressionContext, Integer> groupByExpressionIndexMap = new HashMap<>();
     int numGroupByExpressions = groupByExpressions.size();
     for (int i = 0; i < numGroupByExpressions; i++) {
-      groupByExpressionIndexMap.put(groupByExpressions.get(i).getIdentifier(), i);
+      // NOTE: Key on the whole expression, not on getIdentifier(). getIdentifier() is null for anything that is not a
+      //       plain column reference, so keying on it collapses every transform group-by key onto a single null entry
+      //       and makes all ORDER BY expressions resolve to the same column index.
+      groupByExpressionIndexMap.put(groupByExpressions.get(i), i);
     }
     return groupByExpressionIndexMap;
   }
@@ -93,12 +98,14 @@ private static class OrderByExpressionWithIndex {
   /// Add an index for each orderby expression with respect to its position in the group keys
   private static List<OrderByExpressionWithIndex> getGroupKeyOrderByExpressionFromRowOrderByExpressions(
       List<OrderByExpressionContext> rowOrderByExpressions, List<ExpressionContext> groupByExpressions) {
-    Map<String, Integer> groupByExpressionIndexMap = getGroupByExpressionIndexMap(groupByExpressions);
+    Map<ExpressionContext, Integer> groupByExpressionIndexMap = getGroupByExpressionIndexMap(groupByExpressions);
     List<OrderByExpressionWithIndex> result = new ArrayList<>();
     // get index wrt group key for each order by expression
-    rowOrderByExpressions.forEach(expr ->
-        result.add(
-            new OrderByExpressionWithIndex(expr, groupByExpressionIndexMap.get(expr.getExpression().getIdentifier()))));
+    rowOrderByExpressions.forEach(expr -> {
+      Integer index = groupByExpressionIndexMap.get(expr.getExpression());
+      Preconditions.checkState(index != null, "ORDER BY expression: %s is not a GROUP BY key", expr.getExpression());
+      result.add(new OrderByExpressionWithIndex(expr, index));
+    });
     return result;
   }
 

pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java

@@ -22,9 +22,11 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.Set;
 import org.apache.commons.io.FileUtils;
@@ -186,6 +188,47 @@ public void testMultiColumnHybridGroupKeyGenerator() {
     }
   }
 
+  /// When the group limit is reached, the multi-column generator must still look up each row's own key. The test
+  /// data contains no null values, so enabling null handling must not change the generated group keys at all.
+  @Test
+  public void testMultiColumnGroupKeyGeneratorAtGroupsLimitWithNullHandling() {
+    int numGroupsLimit = 5;
+    ExpressionContext[] groupByExpressions = new ExpressionContext[]{
+        ExpressionContext.forIdentifier(INT_COLUMN), ExpressionContext.forIdentifier(LONG_COLUMN)
+    };
+
+    int[] groupKeysWithoutNullHandling = new int[NUM_RECORDS];
+    new NoDictionaryMultiColumnGroupKeyGenerator(_projectOperator, groupByExpressions, numGroupsLimit, false,
+        null).generateKeysForBlock(_valueBlock, groupKeysWithoutNullHandling);
+
+    int[] groupKeysWithNullHandling = new int[NUM_RECORDS];
+    new NoDictionaryMultiColumnGroupKeyGenerator(_projectOperator, groupByExpressions, numGroupsLimit, true,
+        null).generateKeysForBlock(_valueBlock, groupKeysWithNullHandling);
+
+    assertEquals(groupKeysWithNullHandling, groupKeysWithoutNullHandling);
+
+    // Sanity check the absolute values too, derived from the block's own contents rather than from an assumed row
+    // layout: the first 'numGroupsLimit' distinct keys encountered get sequential ids, and every key first seen after
+    // the limit is reached stays invalid forever.
+    int numDocs = _valueBlock.getNumDocs();
+    assertEquals(numDocs, NUM_RECORDS);
+    int[] intValues = _valueBlock.getBlockValueSet(groupByExpressions[0]).getIntValuesSV();
+    long[] longValues = _valueBlock.getBlockValueSet(groupByExpressions[1]).getLongValuesSV();
+    Map<List<Object>, Integer> keyToId = new HashMap<>();
+    int[] expectedGroupKeys = new int[NUM_RECORDS];
+    int nextId = 0;
+    for (int row = 0; row < numDocs; row++) {
+      List<Object> key = List.of(intValues[row], longValues[row]);
+      Integer id = keyToId.get(key);
+      if (id == null) {
+        id = nextId < numGroupsLimit ? nextId++ : GroupKeyGenerator.INVALID_ID;
+        keyToId.put(key, id);
+      }
+      expectedGroupKeys[row] = id;
+    }
+    assertEquals(groupKeysWithNullHandling, expectedGroupKeys);
+  }
+
   private void testGroupKeyGenerator(int[] groupByColumnIndexes) {
     int numGroupByColumns = groupByColumnIndexes.length;
     GroupKeyGenerator groupKeyGenerator;

pinot-core/src/test/java/org/apache/pinot/core/query/utils/OrderByComparatorFactoryTest.java

@@ -20,13 +20,17 @@
 package org.apache.pinot.core.query.utils;
 
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.FunctionContext;
 import org.apache.pinot.common.request.context.OrderByExpressionContext;
+import org.apache.pinot.core.data.table.Record;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
 
 
 public class OrderByComparatorFactoryTest {
@@ -104,4 +108,38 @@ public void testTwoNullsCompareNextColumn() {
 
     assertEquals(extractColumn(_rows, COLUMN2_INDEX), Arrays.asList(1, 2, 3));
   }
+
+  private static ExpressionContext function(String name, String arg) {
+    return ExpressionContext.forFunction(
+        new FunctionContext(FunctionContext.Type.TRANSFORM, name, List.of(ExpressionContext.forIdentifier(arg))));
+  }
+
+  /// getRecordKeyComparator maps each ORDER BY expression to its position in the GROUP BY list. When several
+  /// group-by keys are transform expressions rather than plain identifiers, every expression must still resolve to
+  /// its own column, otherwise SortedRecordsMerger sees unequal groups as equal and merges them.
+  @Test
+  public void testRecordKeyComparatorWithMultipleTransformGroupByKeys() {
+    ExpressionContext key0 = function("datetrunc", "tsColumn");
+    ExpressionContext key1 = function("jsonextractindex", "jsonColumn");
+    List<ExpressionContext> groupByExpressions = List.of(key0, key1);
+    List<OrderByExpressionContext> orderBys =
+        List.of(new OrderByExpressionContext(key0, ASC, NULLS_LAST), new OrderByExpressionContext(key1, ASC,
+            NULLS_LAST));
+
+    Comparator<Record> comparator =
+        OrderByComparatorFactory.getRecordKeyComparator(orderBys, groupByExpressions, false);
+
+    // Same second key, different first key: these are distinct groups and must not compare equal.
+    Record a = new Record(new Object[]{1L, "x", 10.0});
+    Record b = new Record(new Object[]{2L, "x", 20.0});
+    assertTrue(comparator.compare(a, b) < 0, "rows differing only in the first group key compared equal");
+    assertTrue(comparator.compare(b, a) > 0);
+
+    // Same first key, different second key: also distinct groups.
+    Record c = new Record(new Object[]{1L, "y", 30.0});
+    assertTrue(comparator.compare(a, c) < 0, "rows differing only in the second group key compared equal");
+
+    // Identical keys are the same group.
+    assertEquals(comparator.compare(a, new Record(new Object[]{1L, "x", 99.0})), 0);
+  }
 }

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonExtractIndexGroupByRealtimeTest.java

@@ -0,0 +1,62 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.integration.tests.custom;
+
+import java.io.File;
+import java.util.List;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.Test;
+
+
+/// Same as [JsonExtractIndexGroupByTest] but against a realtime table, so the JSON index being read is the mutable
+/// (consuming segment) implementation. The production table in PINOT-489 is realtime.
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class JsonExtractIndexGroupByRealtimeTest extends JsonExtractIndexGroupByTest {
+  private static final String TABLE_NAME = "RTJsonExtractIndexGroupByTest";
+
+  @Override
+  public String getTableName() {
+    return TABLE_NAME;
+  }
+
+  @Override
+  public boolean isRealtimeTable() {
+    return true;
+  }
+
+  @Override
+  protected int getRealtimeSegmentFlushSize() {
+    // Keep everything in consuming (mutable) segments so the mutable JSON index is exercised.
+    return 1_000_000;
+  }
+
+  @Override
+  protected TableConfig createRealtimeTableConfig(File sampleAvroFile) {
+    AvroFileSchemaKafkaAvroMessageDecoder._avroFile = sampleAvroFile;
+    return new TableConfigBuilder(TableType.REALTIME).setTableName(getTableName())
+        .setStreamConfigs(getStreamConfigs())
+        .setTimeColumnName(getTimeColumnName())
+        .setJsonIndexColumns(List.of(PROPERTIES_FIELD))
+        .setNoDictionaryColumns(List.of(PROPERTIES_FIELD))
+        .setNumReplicas(getNumReplicas())
+        .build();
+  }
+}