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 beb891d2282 Keep partitioned workers for window and WITHIN GROUP sorts
over hash exchanges (#19954)
beb891d2282 is described below
commit beb891d2282a3ceea4fd749166c3bbce0f4339ab
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Oct 21 11:48:30 2026 -0700
Keep partitioned workers for window and WITHIN GROUP sorts over hash
exchanges (#19954)
---
.../planner/physical/DispatchablePlanVisitor.java | 21 ++++-
.../apache/pinot/query/QueryCompilationTest.java | 91 +++++++++++++++++++++
.../physical/DispatchablePlanVisitorTest.java | 93 ++++++++++++++++++++++
.../resources/queries/ExplainPhysicalPlans.json | 18 +++++
4 files changed, 220 insertions(+), 3 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
index 5cd307bbebb..0594a689535 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
@@ -22,6 +22,7 @@ import com.google.common.base.Preconditions;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Set;
+import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.core.JoinRelType;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.common.config.provider.TableCache;
@@ -152,12 +153,26 @@ public class DispatchablePlanVisitor implements
PlanNodeVisitor<Void, Dispatchab
public Void visitSort(SortNode node, DispatchablePlanContext context) {
node.getInputs().get(0).visit(this, context);
DispatchablePlanMetadata dispatchablePlanMetadata =
getOrCreateDispatchablePlanMetadata(node, context);
- // Final sort (receives from sort exchange) needs singleton worker
- boolean isFinalSort = node.getInputs().get(0) instanceof
MailboxReceiveNode;
- dispatchablePlanMetadata.setRequireSingleton(isFinalSort);
+ dispatchablePlanMetadata.setRequireSingleton(isFinalSort(node));
return null;
}
+ /// A final sort orders (or trims with fetch/offset) every row its stage
receives, so it needs a singleton worker,
+ /// which also rotates the hotspot server across queries. A Sort without
fetch/offset over a hash-partitioned exchange
+ /// orders each partition only - it feeds a `PARTITION BY ... ORDER BY`
window or a `WITHIN GROUP` aggregate - so it
+ /// keeps the partitioned worker assignment.
+ private static boolean isFinalSort(SortNode node) {
+ if (!(node.getInputs().get(0) instanceof MailboxReceiveNode)) {
+ return false;
+ }
+ if (node.getFetch() >= 0 || node.getOffset() >= 0) {
+ return true;
+ }
+ MailboxReceiveNode receiveNode = (MailboxReceiveNode)
node.getInputs().get(0);
+ return receiveNode.getDistributionType() !=
RelDistribution.Type.HASH_DISTRIBUTED
+ || receiveNode.getKeys().isEmpty();
+ }
+
@Override
public Void visitTableScan(TableScanNode node, DispatchablePlanContext
context) {
DispatchablePlanMetadata dispatchablePlanMetadata =
getOrCreateDispatchablePlanMetadata(node, context);
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 c91508d5d41..63eec96d1dc 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
@@ -45,6 +45,7 @@ import
org.apache.pinot.query.planner.plannode.MailboxSendNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.ProjectNode;
import org.apache.pinot.query.planner.plannode.SetOpNode;
+import org.apache.pinot.query.planner.plannode.SortNode;
import org.apache.pinot.query.planner.plannode.WindowNode;
import org.apache.pinot.query.routing.QueryServerInstance;
import org.apache.pinot.spi.exception.QueryErrorCode;
@@ -1189,6 +1190,96 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
return null;
}
+ /// A Sort without fetch/offset over a hash-partitioned exchange orders each
partition only. It feeds a
+ /// `PARTITION BY ... ORDER BY` window or a `WITHIN GROUP` aggregate, so its
stage must keep the partitioned worker
+ /// assignment (one worker per server, times `stageParallelism`) instead of
the single worker of a final ORDER BY.
+ @Test(dataProvider = "partitionedSortQueries")
+ public void testSortOverPartitionedExchangeKeepsPartitionedWorkers(String
query, int expectedNumWorkers) {
+ Map<SortNode, DispatchablePlanFragment> sortStages =
findSortOverReceiveStages(_queryEnvironment.planQuery(query));
+ assertEquals(sortStages.size(), 1, "Expected one Sort over a mailbox
receive");
+ Map.Entry<SortNode, DispatchablePlanFragment> entry =
sortStages.entrySet().iterator().next();
+ assertTrue(isOverPartitionedReceive(entry.getKey()), "Expected the Sort to
read a hash-partitioned exchange");
+ DispatchablePlanFragment sortStage = entry.getValue();
+ assertEquals(sortStage.getServerInstanceToWorkerIdMap().size(), 2, "Sort
stage should run on every server");
+ assertEquals(sortStage.getWorkerMetadataList().size(), expectedNumWorkers);
+ }
+
+ @DataProvider(name = "partitionedSortQueries")
+ public Object[][] partitionedSortQueries() {
+ return new Object[][]{
+ {"SELECT col1, ROW_NUMBER() OVER (PARTITION BY col1 ORDER BY col3
DESC) FROM a", 2},
+ {"SELECT col1, SUM(col3) OVER (PARTITION BY col1 ORDER BY col3) FROM
a", 2},
+ {"SET stageParallelism=2; SELECT col1, ROW_NUMBER() OVER (PARTITION BY
col1 ORDER BY col3 DESC) FROM a", 4},
+ {"SELECT col1, LISTAGG(col2, ',') WITHIN GROUP (ORDER BY col3) FROM a
GROUP BY col1", 2},
+ {"SET stageParallelism=2; SELECT col1, LISTAGG(col2, ',') WITHIN GROUP
(ORDER BY col3) FROM a GROUP BY col1",
+ 4}
+ };
+ }
+
+ /// A Sort that trims (fetch/offset) or reads a keyless exchange establishes
a total order, so its stage must run on
+ /// a single worker even when `stageParallelism` asks for more.
+ @Test(dataProvider = "globalSortQueries")
+ public void testSortOverGlobalExchangeUsesSingletonWorker(String query) {
+ Map<SortNode, DispatchablePlanFragment> sortStages =
findSortOverReceiveStages(_queryEnvironment.planQuery(query));
+ assertEquals(sortStages.size(), 1, "Expected one Sort over a mailbox
receive");
+ Map.Entry<SortNode, DispatchablePlanFragment> entry =
sortStages.entrySet().iterator().next();
+ assertFalse(isOverPartitionedReceive(entry.getKey()), "Expected the Sort
to read a keyless exchange");
+ assertEquals(entry.getValue().getWorkerMetadataList().size(), 1, "Global
sort stage should use a single worker");
+ }
+
+ @DataProvider(name = "globalSortQueries")
+ public Object[][] globalSortQueries() {
+ return new Object[][]{
+ {"SET stageParallelism=2; SELECT * FROM a ORDER BY col1"},
+ {"SET stageParallelism=2; SELECT * FROM a ORDER BY col1 LIMIT 10"},
+ {"SET stageParallelism=2; SELECT col1, ROW_NUMBER() OVER (ORDER BY
col3) FROM a"},
+ {"SET stageParallelism=2; SELECT LISTAGG(col2, ',') WITHIN GROUP
(ORDER BY col3) FROM a"}
+ };
+ }
+
+ /// A partitioned window under a final ORDER BY: the window's Sort keeps the
partitioned workers, and the final Sort
+ /// still runs on a single worker.
+ @Test
+ public void testPartitionedWindowUnderFinalOrderBy() {
+ String query = "SET stageParallelism=2; SELECT col1, ROW_NUMBER() OVER
(PARTITION BY col1 ORDER BY col3) AS rn "
+ + "FROM a ORDER BY rn";
+ Map<SortNode, DispatchablePlanFragment> sortStages =
findSortOverReceiveStages(_queryEnvironment.planQuery(query));
+ assertEquals(sortStages.size(), 2, "Expected the window Sort and the final
Sort");
+ for (Map.Entry<SortNode, DispatchablePlanFragment> entry :
sortStages.entrySet()) {
+ int expectedNumWorkers = isOverPartitionedReceive(entry.getKey()) ? 4 :
1;
+ assertEquals(entry.getValue().getWorkerMetadataList().size(),
expectedNumWorkers);
+ }
+ }
+
+ /// Maps each Sort that reads straight from a mailbox to the stage that runs
it.
+ private static Map<SortNode, DispatchablePlanFragment>
findSortOverReceiveStages(
+ DispatchableSubPlan dispatchableSubPlan) {
+ Map<SortNode, DispatchablePlanFragment> sortStages = new HashMap<>();
+ for (DispatchablePlanFragment fragment :
dispatchableSubPlan.getQueryStages()) {
+ List<SortNode> sortNodes = new ArrayList<>();
+ collectSortsOverReceive(fragment.getPlanFragment().getFragmentRoot(),
sortNodes);
+ for (SortNode sortNode : sortNodes) {
+ sortStages.put(sortNode, fragment);
+ }
+ }
+ return sortStages;
+ }
+
+ private static void collectSortsOverReceive(PlanNode node, List<SortNode>
sortNodes) {
+ if (node instanceof SortNode && node.getInputs().get(0) instanceof
MailboxReceiveNode) {
+ sortNodes.add((SortNode) node);
+ }
+ for (PlanNode input : node.getInputs()) {
+ collectSortsOverReceive(input, sortNodes);
+ }
+ }
+
+ private static boolean isOverPartitionedReceive(SortNode sortNode) {
+ MailboxReceiveNode receiveNode = (MailboxReceiveNode)
sortNode.getInputs().get(0);
+ return receiveNode.getDistributionType() ==
RelDistribution.Type.HASH_DISTRIBUTED
+ && !receiveNode.getKeys().isEmpty();
+ }
+
/// Tests that FULL OUTER JOIN with only non-equi conditions uses a
singleton worker for the join stage, to ensure
/// correctness of unmatched row tracking. With multiple workers, the
broadcast right table would be on each worker
/// but each worker only sees a subset of left rows, leading to incorrect
unmatched-right-rows output.
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitorTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitorTest.java
new file mode 100644
index 00000000000..ca11e21a862
--- /dev/null
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitorTest.java
@@ -0,0 +1,93 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.query.planner.physical;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.apache.calcite.rel.RelDistribution;
+import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.runtime.PairList;
+import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.query.QueryEnvironment;
+import org.apache.pinot.query.context.PlannerContext;
+import org.apache.pinot.query.planner.plannode.MailboxReceiveNode;
+import org.apache.pinot.query.planner.plannode.MailboxSendNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.planner.plannode.SortNode;
+import org.apache.pinot.query.planner.plannode.ValueNode;
+import org.mockito.Mockito;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+public class DispatchablePlanVisitorTest {
+ private static final int SORT_STAGE = 1;
+ private static final int SENDER_STAGE = 2;
+ private static final DataSchema DATA_SCHEMA =
+ new DataSchema(new String[]{"col"}, new
ColumnDataType[]{ColumnDataType.INT});
+ private static final List<RelFieldCollation> COLLATIONS = List.of(new
RelFieldCollation(0));
+
+ /// Only a Sort over a hash-partitioned exchange that keeps every row orders
each partition independently. Every
+ /// other Sort over a mailbox establishes or trims a total order and needs a
singleton worker.
+ @Test(dataProvider = "sortsOverReceive")
+ public void testSortOverReceiveRequiresSingleton(RelDistribution.Type
distributionType, List<Integer> keys,
+ int fetch, int offset, boolean expectedSingleton) {
+ MailboxSendNode sender = new MailboxSendNode(SENDER_STAGE, DATA_SCHEMA,
+ List.of(new ValueNode(SENDER_STAGE, DATA_SCHEMA,
PlanNode.NodeHint.EMPTY, List.of(), List.of())), SORT_STAGE,
+ PinotRelExchangeType.STREAMING, distributionType, keys, false,
COLLATIONS, false, "absHashCode");
+ MailboxReceiveNode receive = new MailboxReceiveNode(SORT_STAGE,
DATA_SCHEMA, SENDER_STAGE,
+ PinotRelExchangeType.STREAMING, distributionType, keys, COLLATIONS,
false, false, sender);
+ SortNode sort =
+ new SortNode(SORT_STAGE, DATA_SCHEMA, PlanNode.NodeHint.EMPTY,
List.of(receive), COLLATIONS, fetch, offset);
+
+ PlannerContext plannerContext = Mockito.mock(PlannerContext.class);
+ Mockito.when(plannerContext.getOptions()).thenReturn(Map.of());
+
Mockito.when(plannerContext.getEnvConfig()).thenReturn(Mockito.mock(QueryEnvironment.Config.class));
+ DispatchablePlanContext context =
+ new DispatchablePlanContext(null, 0, plannerContext, PairList.of(0,
"col"), Set.of());
+ sort.visit(new DispatchablePlanVisitor(null, null), context);
+
+
assertEquals(context.getDispatchablePlanMetadataMap().get(SORT_STAGE).isRequiresSingletonInstance(),
+ expectedSingleton);
+ }
+
+ @DataProvider(name = "sortsOverReceive")
+ public Object[][] sortsOverReceive() {
+ return new Object[][]{
+ // Window / WITHIN GROUP input: partitioned, keeps every row
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), -1, -1, false},
+ // Trimming a partitioned exchange still needs every row in one place
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), 10, -1, true},
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), -1, 5, true},
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(0), 10, 0, true},
+ // Final ORDER BY: keyless exchange
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(), -1, -1, true},
+ {RelDistribution.Type.HASH_DISTRIBUTED, List.of(), 10, 0, true},
+ // Exchanges that do not partition by key
+ {RelDistribution.Type.RANDOM_DISTRIBUTED, List.of(), -1, -1, true},
+ {RelDistribution.Type.BROADCAST_DISTRIBUTED, List.of(), -1, -1, true},
+ {RelDistribution.Type.SINGLETON, List.of(), -1, -1, true}
+ };
+ }
+}
diff --git
a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
index fbb428bad31..691815bc475 100644
--- a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
+++ b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
@@ -736,6 +736,24 @@
"\n"
]
},
+ {
+ "description": "window function (PARTITION BY ... ORDER BY) sorts each
partition on every window-stage worker; the Sort over the hash exchange is not
a final sort, so the stage is not collapsed to a singleton",
+ "sql": "EXPLAIN IMPLEMENTATION PLAN FOR SELECT col1, ROW_NUMBER() OVER
(PARTITION BY col1 ORDER BY col3 DESC) FROM a",
+ "output": [
+ "[0]@localhost:3|[0] MAIL_RECEIVE(BROADCAST_DISTRIBUTED)",
+ "\n├── [1]@localhost:2|[1]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@localhost:3|[0]} (Subtree Omitted)",
+ "\n└── [1]@localhost:1|[0]
MAIL_SEND(BROADCAST_DISTRIBUTED)->{[0]@localhost:3|[0]}",
+ "\n └── [1]@localhost:1|[0] PROJECT",
+ "\n └── [1]@localhost:1|[0] WINDOW",
+ "\n └── [1]@localhost:1|[0] SORT",
+ "\n └── [1]@localhost:1|[0]
MAIL_RECEIVE(HASH_DISTRIBUTED)",
+ "\n ├── [2]@localhost:2|[1]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
+ "\n └── [2]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]}",
+ "\n └── [2]@localhost:1|[0] PROJECT",
+ "\n └── [2]@localhost:1|[0] TABLE SCAN
(a) null",
+ "\n"
+ ]
+ },
{
"description": "set op (UNION ALL) uses a local (SINGLETON) exchange
on every input: it only concatenates, so no redistribution is needed and each
sender is wired 1-to-1 to the union worker on its own server",
"sql": "EXPLAIN IMPLEMENTATION PLAN FOR SELECT col3 FROM a UNION ALL
SELECT col3 FROM a",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
