xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4370427480


##########
pinot-core/src/test/java/org/apache/pinot/core/plan/maker/QueryOverrideWithHintsTest.java:
##########
@@ -189,4 +200,123 @@ public void testNotRewriteExpressionsWithHints() {
     assertEquals(queryContext.getSelectExpressions().get(0).getFunction(),
         
queryContext.getExpressionOverrideHints().keySet().iterator().next().getFunction());
   }
+
+  @DataProvider
+  public Object[][] aggregationRewriteCases() {
+    return new Object[][]{{false, false}, {true, false}, {false, true}, {true, 
true}};
+  }
+
+  @Test(dataProvider = "aggregationRewriteCases")
+  public void testRewriteCompiledAggregation(boolean nullHandlingEnabled, 
boolean filtered) {
+    String filter = filtered ? " FILTER(WHERE DATETRUNC('MONTH', ts, 
'MILLISECONDS') > 0)" : "";
+    QueryContext queryContext = aggregationQuery(
+        "SELECT DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS'))" + filter
+            + " AS active_days FROM myTable", "MONTH", nullHandlingEnabled);
+    InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext, 
_indexSegment);
+
+    ExpressionContext derived = ExpressionContext.forIdentifier("$ts$MONTH");
+    
assertEquals(queryContext.getAggregationFunctions()[0].getInputExpressions(), 
List.of(derived));
+    assertEquals(queryContext.getColumns(), Set.of("$ts$MONTH"));
+    assertEquals(queryContext.getAliasList(), List.of("active_days"));
+    assertEquals(queryContext.isNullHandlingEnabled(), nullHandlingEnabled);
+    assertEquals(queryContext.hasFilteredAggregations(), filtered);
+    var pair = queryContext.getFilteredAggregationFunctions().get(0);
+    if (filtered) {
+      assertEquals(pair.getRight().getPredicate().getLhs(), derived);
+    }
+    FunctionContext function = 
queryContext.getSelectExpressions().get(0).getFunction();
+    if (filtered) {
+      function = function.getArguments().get(0).getFunction();
+    }
+    
assertEquals(queryContext.getFilteredAggregationsIndexMap().get(Pair.of(function,
 pair.getRight())),
+        Integer.valueOf(0));
+
+    AggregationFunction[] functions = queryContext.getAggregationFunctions();
+    InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext, 
_indexSegment);
+    assertSame(queryContext.getAggregationFunctions(), functions);
+  }
+
+  @Test
+  public void testRewritePreservesAggregationSlots() {
+    String sql = "SELECT DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS')) 
AS days, "
+        + "DISTINCTCOUNT(DATETRUNC('MONTH', ts, 'MILLISECONDS')) AS 
duplicate_days FROM myTable "
+        + "HAVING DISTINCTCOUNT(\"$ts$MONTH\") > 0 ORDER BY 
DISTINCTCOUNT(\"$ts$MONTH\")";
+    QueryContext queryContext = aggregationQuery(sql, "MONTH", false);
+    QueryContext brokerContext = aggregationQuery(sql, "MONTH", false);
+    assertEquals(queryContext.getAggregationFunctions().length, 3);
+    InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext, 
_indexSegment);
+    // Converging SELECT and HAVING/ORDER-BY expressions must keep their 
original wire slots.
+    assertEquals(queryContext.getAggregationFunctions().length, 3);
+    assertEquals(queryContext.getFilteredAggregationFunctions().size(), 3);
+    for (AggregationFunction function : 
queryContext.getAggregationFunctions()) {
+      assertEquals(function.getInputExpressions(), 
List.of(ExpressionContext.forIdentifier("$ts$MONTH")));
+    }
+    FunctionContext having = 
queryContext.getHavingFilter().getPredicate().getLhs().getFunction();
+    
assertEquals(queryContext.getFilteredAggregationsIndexMap().get(Pair.of(having, 
null)), Integer.valueOf(2));
+    assertEquals(queryContext.getAliasList(), List.of("days", 
"duplicate_days"));
+    String[] physicalNames = new String[3];
+    String[] logicalNames = new String[3];
+    ColumnDataType[] types = new ColumnDataType[3];
+    for (int i = 0; i < 3; i++) {
+      physicalNames[i] = 
queryContext.getAggregationFunctions()[i].getResultColumnName();
+      logicalNames[i] = 
brokerContext.getAggregationFunctions()[i].getResultColumnName();
+      types[i] = 
queryContext.getAggregationFunctions()[i].getIntermediateResultColumnType();
+    }
+    DataSchema canonical = 
ReducerDataSchemaUtils.canonicalizeDataSchemaForAggregation(brokerContext,
+        new DataSchema(physicalNames, types));
+    assertEquals(canonical.getColumnNames(), logicalNames);
+    assertEquals(canonical.getColumnDataTypes(), types);
+  }
+
+  @Test
+  public void testMissingIndexRetainsCompiledAggregation() {

Review Comment:
   I validated the unchanged SQL against two real Parquet segments with the 
materialized DAY index, using production planning, two in-process server 
groups, DataTable serialization and broker reduction. All 600 measured 
full-query executions across the baseline and both-fix versions matched 
independent expectations for all 165 groups, including values, column names and 
types. The PR description records the protocol, result hash and JFR evidence. 
This is saved local in-process validation rather than a committed cluster-suite 
regression or customer deployment proof; I am keeping this fix scoped to the 
focused regressions and that real-segment check.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to