This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new f216ed25ac2 Add tests for chaining colocated joins, window functions
and set operations without shuffle (#18997)
f216ed25ac2 is described below
commit f216ed25ac2319fd3654febe062c9b097b2f74b3
Author: Yash Mayya <[email protected]>
AuthorDate: Fri Jul 17 11:47:51 2026 -0700
Add tests for chaining colocated joins, window functions and set operations
without shuffle (#18997)
---
.../apache/pinot/query/QueryCompilationTest.java | 78 ++++++++++++++++++++++
.../src/test/resources/queries/QueryHints.json | 8 +++
2 files changed, 86 insertions(+)
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
index 123bb5cafb3..715cd421774 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
@@ -1161,6 +1161,84 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
return sendNodes;
}
+ /// When colocation hints are applied to a chain of operations that are all
keyed on the same (partition) column, the
+ /// whole chain executes without a data shuffle: every hash-distributed
exchange in the plan is pre-partitioned. This
+ /// covers combinations of colocated joins, window functions and set
operations (see {@link #colocatedChains}). The
+ /// join key is {@code a.col2}/{@code b.col1} (each table's partition
column), so the join is colocated, and the
+ /// window and set op keep that same key.
+ @Test(dataProvider = "colocatedChains")
+ public void testColocatedOperationsChainWithoutShuffle(String description,
String query) {
+ List<MailboxSendNode> hashSends =
findHashSendNodes(_queryEnvironment.planQuery(query));
+ assertFalse(hashSends.isEmpty(), "Expected at least one hash exchange in
the chain: " + description);
+ for (MailboxSendNode sendNode : hashSends) {
+ assertTrue(sendNode.isPrePartitioned(),
+ "Colocated chain '" + description + "' should have no shuffle, but a
hash exchange was not pre-partitioned");
+ }
+ }
+
+ @DataProvider(name = "colocatedChains")
+ private Object[][] colocatedChains() {
+ // joinOptions + windowOptions together (attach to the join and window in
a single SELECT).
+ String joinAndWindowHint = "/*+
joinOptions(is_colocated_by_join_keys='true'), "
+ + "windowOptions(is_partitioned_by_window_keys='true') */";
+ String joinHint = "/*+ joinOptions(is_colocated_by_join_keys='true') */";
+ String windowHint = "/*+
windowOptions(is_partitioned_by_window_keys='true') */";
+ String setOpHint = "/*+ setOpOptions(is_colocated_by_set_op_keys='true')
*/";
+ String join = "a JOIN b ON a.col2 = b.col1";
+ return new Object[][]{
+ {"join -> window",
+ "SELECT " + joinAndWindowHint + " a.col2, SUM(a.col3) OVER
(PARTITION BY a.col2) FROM " + join},
+ {"set op of colocated joins", "SELECT " + setOpHint + " * FROM (SELECT
" + joinHint + " a.col2 FROM " + join
+ + " UNION ALL SELECT " + joinHint + " a.col2 FROM " + join + ")"},
+ {"set op of colocated windows", "SELECT " + setOpHint + " * FROM
(SELECT " + windowHint
+ + " col2, SUM(col3) OVER (PARTITION BY col2) FROM a UNION ALL
SELECT " + windowHint
+ + " col2, SUM(col3) OVER (PARTITION BY col2) FROM a)"},
+ {"join -> window -> set op", "SELECT " + setOpHint + " * FROM (SELECT
" + joinAndWindowHint
+ + " a.col2, SUM(a.col3) OVER (PARTITION BY a.col2) FROM " + join +
" UNION ALL SELECT " + joinAndWindowHint
+ + " a.col2, SUM(a.col3) OVER (PARTITION BY a.col2) FROM " + join +
")"},
+ };
+ }
+
+ /// Baseline for {@link #testColocatedOperationsChainWithoutShuffle}:
without the colocation hints the same
+ /// join -> window chain still shuffles (at least one hash exchange is not
pre-partitioned), proving the hints are
+ /// what eliminate the shuffles.
+ @Test
+ public void testChainWithoutColocationHintsStillShuffles() {
+ String query = "SELECT a.col2, SUM(a.col3) OVER (PARTITION BY a.col2) FROM
a JOIN b ON a.col2 = b.col1";
+ List<MailboxSendNode> hashSends =
findHashSendNodes(_queryEnvironment.planQuery(query));
+ assertTrue(hashSends.stream().anyMatch(sendNode ->
!sendNode.isPrePartitioned()),
+ "Without the colocation hints, the join -> window chain should contain
at least one shuffled hash exchange");
+ }
+
+ /// Set-op-specific baseline for the "set op of colocated joins" chain: with
+ /// {@code is_colocated_by_set_op_keys='false'} the set-op exchange is
forced back to a shuffle (while the per-branch
+ /// joins stay colocated), proving the set-op hint is what keeps the set-op
level of the chain colocated.
+ @Test
+ public void testSetOpChainWithoutSetOpHintShuffles() {
+ String join = "a JOIN b ON a.col2 = b.col1";
+ String joinHint = "/*+ joinOptions(is_colocated_by_join_keys='true') */";
+ String query = "SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='false') */ * FROM (SELECT " + joinHint
+ + " a.col2 FROM " + join + " UNION ALL SELECT " + joinHint + " a.col2
FROM " + join + ")";
+ List<MailboxSendNode> hashSends =
findHashSendNodes(_queryEnvironment.planQuery(query));
+ assertTrue(hashSends.stream().anyMatch(sendNode ->
!sendNode.isPrePartitioned()),
+ "With is_colocated_by_set_op_keys='false', the set-op level of the
chain should shuffle");
+ }
+
+ /// Collects the {@link MailboxSendNode} at the root of every
hash-distributed stage in the plan (i.e. every
+ /// inter-stage hash exchange). A pre-partitioned send is a direct,
no-shuffle exchange; a non-pre-partitioned one is
+ /// a full shuffle.
+ private List<MailboxSendNode> findHashSendNodes(DispatchableSubPlan
dispatchableSubPlan) {
+ List<MailboxSendNode> sendNodes = new ArrayList<>();
+ for (DispatchablePlanFragment fragment :
dispatchableSubPlan.getQueryStages()) {
+ PlanNode root = fragment.getPlanFragment().getFragmentRoot();
+ if (root instanceof MailboxSendNode
+ && ((MailboxSendNode) root).getDistributionType() ==
RelDistribution.Type.HASH_DISTRIBUTED) {
+ sendNodes.add((MailboxSendNode) root);
+ }
+ }
+ return sendNodes;
+ }
+
/**
* Finds the DispatchablePlanFragment containing a JoinNode (non-leaf,
non-root stage).
*/
diff --git a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
index c8f872bd3ee..913b24df0ab 100644
--- a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
+++ b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
@@ -183,6 +183,14 @@
{
"description": "INTERSECT with is_colocated_by_set_op_keys='false'
disables auto-detected pre-partitioning and falls back to a shuffle; results
stay correct",
"sql": "SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='false')
*/ {tbl1}.num FROM {tbl1} /*+ tableOptions(partition_function='hashcode',
partition_key='num', partition_size='4') */ INTERSECT SELECT {tbl2}.num FROM
{tbl2} /*+ tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */"
+ },
+ {
+ "description": "Chained colocated operations stay correct: a colocated
JOIN on num feeding a window partitioned by num (the no-shuffle plan shape is
asserted in QueryCompilationTest)",
+ "sql": "SELECT /*+ joinOptions(is_colocated_by_join_keys='true'),
windowOptions(is_partitioned_by_window_keys='true') */ {tbl1}.num, COUNT(*)
OVER (PARTITION BY {tbl1}.num) FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */ JOIN {tbl2} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */ ON {tbl1}.num = {tbl2}.num"
+ },
+ {
+ "description": "Chained colocated operations stay correct: a UNION ALL
of two colocated JOINs on num, set-op hint on the outer wrap and join hint per
branch (the no-shuffle plan shape is asserted in QueryCompilationTest)",
+ "sql": "SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='true') */
* FROM (SELECT /*+ joinOptions(is_colocated_by_join_keys='true') */ {tbl1}.num
FROM {tbl1} /*+ tableOptions(partition_function='hashcode',
partition_key='num', partition_size='4') */ JOIN {tbl2} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */ ON {tbl1}.num = {tbl2}.num UNION ALL SELECT /*+
joinOptions(is_colocated_by_join_keys='true') */ {tbl1}.num FROM {tbl1} /*+ ta
[...]
}
]
},
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
