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]