xiangfu0 commented on code in PR #20035:
URL: https://github.com/apache/pinot/pull/20035#discussion_r4370150197
##########
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() {
+ QueryContext queryContext = aggregationQuery(
+ "SELECT DISTINCTCOUNT(DATETRUNC('DAY', ts, 'MILLISECONDS')) AS
active_days FROM myTable", "DAY", false);
+ AggregationFunction[] functions = queryContext.getAggregationFunctions();
+ List<ExpressionContext> inputs = functions[0].getInputExpressions();
+ InstancePlanMakerImplV2.rewriteQueryContextWithHints(queryContext,
_indexSegment);
+ assertSame(queryContext.getAggregationFunctions(), functions);
+
assertEquals(queryContext.getAggregationFunctions()[0].getInputExpressions(),
inputs);
+ assertEquals(queryContext.getColumns(), Set.of("ts"));
+ assertEquals(queryContext.getAliasList(), List.of("active_days"));
+ }
+
+ @DataProvider
+ public Object[][] instanceRewriteCases() {
+ return new Object[][]{{false, false}, {true, false}, {false, true}, {true,
true}};
+ }
+
+ @Test(dataProvider = "instanceRewriteCases")
+ public void testInstanceRewriteRequiresIndexOnEverySegment(boolean
streaming, boolean allIndexed) {
Review Comment:
Minor coverage thought: the agg-recompilation tests all use DISTINCTCOUNT in
SELECT, so no case exercises recompilation when the hinted expression sits in
GROUP BY (e.g. `SELECT SUM(DATETRUNC('MONTH', ts, 'MILLISECONDS')) ... GROUP BY
DATETRUNC('MONTH', ts, 'MILLISECONDS')`). Would be nice to add a data-provider
row asserting both the group-by identifier and getInputExpressions() resolve to
$ts$MONTH. Not blocking.
##########
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:
All assertions here check QueryContext state, which is a strong proxy, but
there's no end-to-end test that a query over real segment data with a timestamp
index returns correct aggregation results. If you want extra confidence, a
focused case in the custom-data cluster suite would close the loop. Optional.
--
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]
