61yao commented on code in PR #9870:
URL: https://github.com/apache/pinot/pull/9870#discussion_r1035483859
##########
pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java:
##########
@@ -79,13 +79,39 @@ public void testQueryWithException(String query, String
exceptionSnippet) {
}
@Test
- public void testQueryGroupByAfterJoinShouldNotDoDataShuffle()
+ public void testQueryGroupByAfterJoinShouldProperlyRewriteShuffle()
throws Exception {
String query = "SELECT a.col1, a.col2, AVG(b.col3) FROM a JOIN b ON a.col1
= b.col2 "
+ " WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col1,
a.col2";
QueryPlan queryPlan = _queryEnvironment.planQuery(query);
- Assert.assertEquals(queryPlan.getQueryStageMap().size(), 5);
- Assert.assertEquals(queryPlan.getStageMetadataMap().size(), 5);
+ assertGroupBySingletonAfterJoin(queryPlan, true);
+
+ // same query with selection list re-odering should also work
+ query = "SELECT a.col2, a.col1, AVG(b.col3) FROM a JOIN b ON a.col1 =
b.col2 "
+ + " WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col2,
a.col1";
+ queryPlan = _queryEnvironment.planQuery(query);
+ assertGroupBySingletonAfterJoin(queryPlan, true);
+
+ // exact same group key should also work
+ query = "SELECT a.col1, a.col2, AVG(b.col3) FROM a JOIN b ON a.col1 =
b.col2 AND a.col2 = b.col2"
+ + " WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col1,
a.col2";
+ queryPlan = _queryEnvironment.planQuery(query);
+ assertGroupBySingletonAfterJoin(queryPlan, true);
+
+ // shrinking group key should not rewrite into singleton
+ query = "SELECT a.col1, AVG(b.col3) FROM a JOIN b ON a.col1 = b.col2 AND
a.col2 = b.col2"
+ + " WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col1";
+ queryPlan = _queryEnvironment.planQuery(query);
+ assertGroupBySingletonAfterJoin(queryPlan, false);
+
+ // mismatched group key should not rewrite into singleton
+ query = "SELECT a.col3, a.col1, AVG(b.col3) FROM a JOIN b ON a.col1 =
b.col2 AND a.col2 = b.col2"
+ + " WHERE a.col3 >= 0 AND a.col2 = 'a' AND b.col3 < 0 GROUP BY a.col1,
a.col3";
+ queryPlan = _queryEnvironment.planQuery(query);
+ assertGroupBySingletonAfterJoin(queryPlan, false);
Review Comment:
Can we add a test case for expanding group key?
query = "SELECT a.col1, a.col2, AVG(b.col3) FROM a JOIN b ON a.col1 =
b.col2 GROUP BY a.col1, a.col2";
--
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: commits-unsubscr...@pinot.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
