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 02c41a39e8f Use a local exchange for UNION ALL inputs in the
multi-stage engine (#19330)
02c41a39e8f is described below
commit 02c41a39e8f0b8d41440696cf02f76c6955e5faa
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Aug 27 12:04:52 2026 -0400
Use a local exchange for UNION ALL inputs in the multi-stage engine (#19330)
---
.../pinot/calcite/rel/hint/PinotHintOptions.java | 14 +-
.../rel/rules/PinotRelDistributionTraitRule.java | 21 +++
.../rules/PinotSetOpExchangeNodeInsertRule.java | 21 ++-
.../planner/physical/MailboxAssignmentVisitor.java | 62 +++++++--
.../apache/pinot/query/routing/WorkerManager.java | 57 +++++++-
.../apache/pinot/query/QueryCompilationTest.java | 148 +++++++++++++++++++--
.../pinot/query/QueryPlannerRuleOptionsTest.java | 12 +-
.../physical/MailboxAssignmentVisitorTest.java | 82 +++++++++++-
.../src/test/resources/queries/AggregatePlans.json | 4 +-
.../resources/queries/ExplainPhysicalPlans.json | 26 ++--
.../src/test/resources/queries/SetOpPlans.json | 34 ++---
.../src/test/resources/queries/QueryHints.json | 28 ++++
12 files changed, 428 insertions(+), 81 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
index 4c3b73aa017..36b43b0f734 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/hint/PinotHintOptions.java
@@ -88,8 +88,9 @@ public class PinotHintOptions {
/// Hint options for set operations (UNION / UNION ALL / INTERSECT / EXCEPT).
public static class SetOpHintOptions {
- /// Forces (or disables) a colocated, pre-partitioned exchange on every
input of a set operation
- /// (UNION / UNION ALL / INTERSECT / EXCEPT), so the inputs are processed
in place without a network shuffle.
+ /// Forces (or disables) a colocated, pre-partitioned exchange on every
input of an INTERSECT, EXCEPT or
+ /// distinct UNION, so the inputs are processed in place without a network
shuffle. On a UNION ALL only
+ /// `'false'` has an effect; see below.
///
/// This is opt-in and honored only by the default (V1) query planner; the
V2 physical planner determines
/// colocation on its own. Like the join hint
[JoinHintOptions#IS_COLOCATED_BY_JOIN_KEYS], it trusts the user:
@@ -98,8 +99,13 @@ public class PinotHintOptions {
/// operation matches rows on the entire output row, so forcing `'true'`
is only correct when all inputs are
/// partitioned the same way (same partition function and count) on one or
more of the projected columns, so that
/// rows which are equal across all projected columns are guaranteed to
land on the same worker. Forcing it when
- /// that does not hold silently produces wrong results for INTERSECT,
EXCEPT and distinct UNION; UNION ALL only
- /// concatenates and is unaffected.
+ /// that does not hold silently produces wrong results for INTERSECT,
EXCEPT and distinct UNION.
+ ///
+ /// UNION ALL is different: it only concatenates, so it is correct under
ANY row-to-worker mapping and never
+ /// needs a shuffle. Its inputs always get a local exchange, and `'true'`
is therefore a no-op on a UNION ALL.
+ /// Only `'false'` changes anything there -- it restores the full-row
shuffle. Note that a partitioning assertion
+ /// made *above* a UNION ALL (e.g. `aggOptions` on an aggregation over it)
is never established by the union's own
+ /// exchanges; it must hold on the physical data.
public static final String IS_COLOCATED_BY_SET_OP_KEYS =
"is_colocated_by_set_op_keys";
/// Reads the hint from a hint list. Unlike
[JoinHintOptions#isColocatedByJoinKeys], this takes the hint list
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRelDistributionTraitRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRelDistributionTraitRule.java
index 985937281bd..09877167631 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRelDistributionTraitRule.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotRelDistributionTraitRule.java
@@ -31,6 +31,7 @@ import org.apache.calcite.rel.RelDistributions;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Exchange;
import org.apache.calcite.rel.core.Project;
+import org.apache.calcite.rel.core.SetOp;
import org.apache.calcite.rel.core.TableScan;
import org.apache.calcite.rel.logical.LogicalFilter;
import org.apache.calcite.rel.logical.LogicalJoin;
@@ -161,6 +162,26 @@ public class PinotRelDistributionTraitRule extends
RelOptRule {
if (inputRelDistribution != null) {
return inputRelDistribution;
}
+ } else if (node instanceof SetOp) {
+ // A set operation sits above the exchanges inserted by
PinotSetOpExchangeNodeInsertRule, and an exchange that
+ // genuinely redistributes describes the output: hash on the projected
columns leaves the output hash
+ // distributed on them, broadcast leaves every worker holding all rows
of every branch. That holds for a
+ // pre-partitioned hash exchange too, because
is_colocated_by_set_op_keys asserts that rows equal across all
+ // projected columns already share a worker -- exactly the claimed
distribution. We take the hint at its word
+ // here, just as the exchange itself does.
+ // A local (SINGLETON) exchange, which is what UNION ALL gets, is the
exception: it redistributes nothing, so
+ // the output keeps whatever placement the inputs happened to have and
SINGLETON ("everything on one worker")
+ // is not true of it. Claiming anything there would let a deduplicating
consumer above it -- for example the
+ // aggregate UnionToDistinctRule puts over a distinct UNION -- skip a
shuffle it needs.
+ // All inputs are checked so a future per-branch decision cannot
silently invalidate this.
+ for (RelNode setOpInput : inputs) {
+ RelNode unboxedInput = PinotRuleUtils.unboxRel(setOpInput);
+ if (!(unboxedInput instanceof PinotLogicalExchange)
+ || ((PinotLogicalExchange)
unboxedInput).getDistribution().getType() == RelDistribution.Type.SINGLETON) {
+ return RelDistributions.of(RelDistribution.Type.RANDOM_DISTRIBUTED,
RelDistributions.EMPTY);
+ }
+ }
+ return ((PinotLogicalExchange) input).getDistribution();
}
// TODO: add the rest of the nodes.
return computeCurrentDistribution(node);
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSetOpExchangeNodeInsertRule.java
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSetOpExchangeNodeInsertRule.java
index e0e90f05331..03537068ca1 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSetOpExchangeNodeInsertRule.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSetOpExchangeNodeInsertRule.java
@@ -26,6 +26,7 @@ import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.rel.RelDistributions;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.SetOp;
+import org.apache.calcite.rel.core.Union;
import org.apache.calcite.rel.hint.Hintable;
import org.apache.calcite.tools.RelBuilderFactory;
import org.apache.calcite.util.ImmutableIntList;
@@ -53,13 +54,27 @@ public class PinotSetOpExchangeNodeInsertRule extends
RelOptRule {
SetOp setOp = call.rel(0);
List<RelNode> inputs = setOp.getInputs();
// When the colocation hint is set, force a pre-partitioned (direct,
no-shuffle) exchange on every input; otherwise
- // leave it null so the planner auto-detects pre-partitioning from the
inputs' distribution. See
+ // leave it null so the planner auto-detects pre-partitioning from the
inputs' distribution. On a UNION ALL only
+ // 'false' has an effect, because its inputs already get a local exchange
either way. See
// PinotHintOptions.SetOpHintOptions.IS_COLOCATED_BY_SET_OP_KEYS for the
correctness contract.
Boolean prePartitioned = resolveColocationHint(setOp);
+ // UNION ALL only concatenates its inputs, so any row-to-worker mapping
produces correct results and no
+ // redistribution is required. Use a local (SINGLETON) exchange so the
union stage inherits its inputs' worker
+ // assignment and the rows are handed over in place, with no shuffle and
no network hop when sender and receiver
+ // land on the same server.
+ // The projected columns are still attached as keys. They are unused while
the exchange stays local, but they let
+ // the mailbox layer promote it to a real HASH distribution when the
inputs do not resolve to the same workers --
+ // the same idiom PinotJoinExchangeNodeInsertRule uses for
distribution_type='local'. Keeping keys also means a
+ // UNION ALL input is never a KEYLESS local exchange, so the "local
exchange with parallelism requires keys"
+ // guard continues to protect the colocated semi-join build side untouched.
+ // An explicit is_colocated_by_set_op_keys='false' hint opts out and
restores the full-row shuffle.
+ List<Integer> keys = ImmutableIntList.range(0,
setOp.getRowType().getFieldCount());
+ boolean useLocalExchange = setOp instanceof Union && ((Union) setOp).all
&& !Boolean.FALSE.equals(prePartitioned);
List<RelNode> newInputs = new ArrayList<>(inputs.size());
for (RelNode input : inputs) {
- RelNode exchange = PinotLogicalExchange.create(input,
- RelDistributions.hash(ImmutableIntList.range(0,
setOp.getRowType().getFieldCount())), prePartitioned);
+ RelNode exchange = useLocalExchange
+ ? PinotLogicalExchange.create(input, RelDistributions.SINGLETON,
keys, null)
+ : PinotLogicalExchange.create(input, RelDistributions.hash(keys),
prePartitioned);
newInputs.add(exchange);
}
call.transformTo(setOp.copy(setOp.getTraitSet(), newInputs));
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java
index a1e0363e608..0036d8f8e81 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitor.java
@@ -44,6 +44,9 @@ public class MailboxAssignmentVisitor extends
DefaultPostOrderTraversalVisitor<V
MailboxSendNode sendNode = (MailboxSendNode) node;
// NOTE: Using Integer to avoid boxing
Integer senderStageId = sendNode.getStageId();
+ // Captured before the loop because the branches below rewrite it, and
one send node can serve several
+ // receiver stages (spools) -- reading it per iteration would make the
wiring depend on visit order.
+ RelDistribution.Type originalDistributionType =
sendNode.getDistributionType();
for (Integer receiverStageId : sendNode.getReceiverStageIds()) {
Map<Integer, DispatchablePlanMetadata> metadataMap =
context.getDispatchablePlanMetadataMap();
DispatchablePlanMetadata senderMetadata =
metadataMap.get(senderStageId);
@@ -55,23 +58,37 @@ public class MailboxAssignmentVisitor extends
DefaultPostOrderTraversalVisitor<V
int numSenders = senderServerMap.size();
int numReceivers = receiverServerMap.size();
- if (sendNode.getDistributionType() == RelDistribution.Type.SINGLETON) {
- // NOTE: We use SINGLETON to represent a local exchange. The actual
distribution type is determined by the
- // parallelism: 1-to-1 when sender and receiver have the same
number of workers, otherwise the data is
- // hash distributed to the parallel receivers on each server.
computeDirectExchange handles the
- // co-location assumption and its cross-server fallback.
- if (numSenders != numReceivers) {
- // Local exchange with parallelism: hash distribute to the
parallel receivers, so keys are required.
+ if (originalDistributionType == RelDistribution.Type.SINGLETON) {
+ // NOTE: We use SINGLETON to represent a local exchange. 1-to-1 when
sender and receiver have the same
+ // number of workers, otherwise each sender fans out to the
parallel receivers on its OWN server.
+ if (numSenders == 0) {
+ // A fully pruned sender stage has nothing to wire, and the
parallelism below would divide by zero.
+ connectWorkers(receiverStageId, receiverServerMap,
senderMailboxesMap, numSenders);
+ connectWorkers(senderStageId, senderServerMap,
receiverMailboxesMap, numReceivers);
+ } else if (numSenders == numReceivers) {
+ computeDirectExchange(senderMailboxesMap, receiverMailboxesMap,
senderStageId, receiverStageId,
+ senderServerMap, receiverServerMap, numSenders, 1,
senderMetadata, receiverMetadata);
+ } else {
+ // Redistributing across several receivers needs keys. A keyless
local exchange here is the colocated
+ // dynamic-broadcast semi-join build side, which needs EVERY
receiver to see the whole build side.
// TODO: Support local exchange with parallelism but no key
Preconditions.checkState(!sendNode.getKeys().isEmpty(), "Local
exchange with parallelism requires keys");
sendNode.setDistributionType(RelDistribution.Type.HASH_DISTRIBUTED);
- Preconditions.checkState(numReceivers % numSenders == 0,
- "Number of receivers: %s should be a multiple of number of
senders: %s for local exchange",
- numReceivers, numSenders);
+ // The parallel wiring addresses a whole receiver range at the
range's first host, which only holds when
+ // the receiver map was derived from the sender by
WorkerManager#assignWorkersForLocalExchange. Otherwise
+ // a range spans hosts and blocks would land on the wrong one,
stalling the receiver until the deadline.
+ // So verify co-residency, and fall back to a hash shuffle when it
does not hold -- correct for every
+ // local exchange kind, since HashExchange routes each key
consistently.
+ if (numReceivers % numSenders == 0
+ && isEachRangeCoResident(senderServerMap, receiverServerMap,
numSenders, numReceivers / numSenders)) {
+ computeDirectExchange(senderMailboxesMap, receiverMailboxesMap,
senderStageId, receiverStageId,
+ senderServerMap, receiverServerMap, numSenders, numReceivers
/ numSenders, senderMetadata,
+ receiverMetadata);
+ } else {
+ connectWorkers(receiverStageId, receiverServerMap,
senderMailboxesMap, numSenders);
+ connectWorkers(senderStageId, senderServerMap,
receiverMailboxesMap, numReceivers);
+ }
}
- int parallelism = numReceivers / numSenders;
- computeDirectExchange(senderMailboxesMap, receiverMailboxesMap,
senderStageId, receiverStageId,
- senderServerMap, receiverServerMap, numSenders, parallelism,
senderMetadata, receiverMetadata);
} else if (senderMetadata.isPrePartitioned() &&
isDirectExchangeCompatible(senderMetadata, receiverMetadata)) {
// Direct exchange: the data is already pre-partitioned, so send it
1-to-1 to the worker with the same worker
// id (with parallelism, fan out each sender worker to a contiguous
range of receiver workers). The
@@ -181,6 +198,21 @@ public class MailboxAssignmentVisitor extends
DefaultPostOrderTraversalVisitor<V
Arrays.toString(receiverPartitionClassIds));
}
+ /// Whether every sender's contiguous receiver range sits entirely on that
sender's own server, which is what
+ /// [#computeDirectExchangeWithParallelism] assumes when it addresses a
whole range at a single host.
+ private static boolean isEachRangeCoResident(Map<Integer,
QueryServerInstance> senderServerMap,
+ Map<Integer, QueryServerInstance> receiverServerMap, int numSenders, int
parallelism) {
+ for (int senderWorkerId = 0; senderWorkerId < numSenders;
senderWorkerId++) {
+ QueryServerInstance senderServer = senderServerMap.get(senderWorkerId);
+ for (int i = 0; i < parallelism; i++) {
+ if (!senderServer.equals(receiverServerMap.get(senderWorkerId *
parallelism + i))) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
/// Wires one sender worker of a direct exchange to the contiguous range of
`parallelism` receiver workers it fans out
/// to. See [#computeDirectExchange].
private void computeDirectExchangeWithParallelism(Map<Integer, Map<Integer,
MailboxInfos>> senderMailboxesMap,
@@ -205,7 +237,9 @@ public class MailboxAssignmentVisitor extends
DefaultPostOrderTraversalVisitor<V
DispatchablePlanMetadata receiver) {
int numSenders = sender.getWorkerIdToServerInstanceMap().size();
int numReceivers = receiver.getWorkerIdToServerInstanceMap().size();
- if (numSenders * sender.getPartitionParallelism() != numReceivers) {
+ // numSenders is 0 when every segment of a leaf stage was pruned. Guard
it: with numReceivers also 0 it
+ // passes the multiplication check and then divides by zero computing the
parallelism.
+ if (numSenders == 0 || numSenders * sender.getPartitionParallelism() !=
numReceivers) {
return false;
}
// A sender whose worker ids stand for partition classes may only be wired
1-to-1 to a receiver whose worker ids
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
index 917cbab89fe..6768fc3b306 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java
@@ -38,6 +38,7 @@ import javax.annotation.Nullable;
import org.apache.calcite.rel.RelDistribution;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.collections4.MapUtils;
+import org.apache.commons.lang3.StringUtils;
import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.calcite.rel.rules.ImmutableTableOptions;
import org.apache.pinot.calcite.rel.rules.TableOptions;
@@ -479,6 +480,43 @@ public class WorkerManager {
return false;
}
+ /// A stage adopts the worker assignment of the FIRST of its local exchange
children (every branch of a UNION ALL
+ /// is one); the others then send 1-to-1 by worker id. Only the worker COUNT
has to match, because
+ /// [MailboxAssignmentVisitor] already handles a child whose workers sit on
different servers -- it costs a network
+ /// hop rather than an in-process handover, still far cheaper than the
shuffle it replaces. A child with no worker
+ /// at all (a fully pruned leaf) can never anchor it: inheriting its empty
map would leave this stage with no
+ /// workers and silently drop every live sibling's rows.
+ private static boolean
canInheritWorkerAssignment(List<DispatchablePlanMetadata> children) {
+ Map<Integer, QueryServerInstance> anchor =
+ children.isEmpty() ? null :
children.get(0).getWorkerIdToServerInstanceMap();
+ if (anchor == null || anchor.isEmpty()) {
+ return false;
+ }
+ for (int i = 1; i < children.size(); i++) {
+ Map<Integer, QueryServerInstance> workers =
children.get(i).getWorkerIdToServerInstanceMap();
+ if (workers == null || workers.size() != anchor.size()) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ /// Whether the partition descriptor inherited from the first local exchange
child describes ALL of them. When it
+ /// does not, the stage still keeps their worker assignment -- a UNION ALL
only concatenates, so placement is free
+ /// -- but it must not advertise a descriptor that only some of its rows
satisfy, because exchanges above it read
+ /// it (see [MailboxAssignmentVisitor#isDirectExchangeCompatible]) to decide
whether they may skip a shuffle.
+ private static boolean shareOnePartitioning(List<DispatchablePlanMetadata>
children) {
+ DispatchablePlanMetadata first = children.get(0);
+ for (int i = 1; i < children.size(); i++) {
+ DispatchablePlanMetadata other = children.get(i);
+ if (!Arrays.equals(first.getPartitionClassIds(),
other.getPartitionClassIds())
+ || !StringUtils.equalsIgnoreCase(first.getPartitionFunction(),
other.getPartitionFunction())) {
+ return false;
+ }
+ }
+ return true;
+ }
+
private Map<Integer, QueryServerInstance>
assignWorkersForLocalExchange(DispatchablePlanMetadata childMetadata) {
int partitionParallelism = childMetadata.getPartitionParallelism();
Map<Integer, QueryServerInstance> childWorkerIdToServerInstanceMap =
childMetadata.getWorkerIdToServerInstanceMap();
@@ -554,15 +592,20 @@ public class WorkerManager {
}
// Assign workers for local exchange if there is one
- DispatchablePlanMetadata localExchangeChildMetadata = null;
- Map<Integer, QueryServerInstance> workerIdToServerInstanceMap = null;
+ List<DispatchablePlanMetadata> localExchangeChildren = new
ArrayList<>(children.size());
for (PlanFragment child : children) {
if (isLocalExchange(child, context)) {
- Preconditions.checkState(localExchangeChildMetadata == null, "Found
multiple local exchanges in the children");
- localExchangeChildMetadata = metadataMap.get(child.getFragmentId());
- workerIdToServerInstanceMap =
assignWorkersForLocalExchange(localExchangeChildMetadata);
+ localExchangeChildren.add(metadataMap.get(child.getFragmentId()));
}
}
+ DispatchablePlanMetadata localExchangeChildMetadata = null;
+ Map<Integer, QueryServerInstance> workerIdToServerInstanceMap = null;
+ boolean inheritPartitioning = false;
+ if (canInheritWorkerAssignment(localExchangeChildren)) {
+ localExchangeChildMetadata = localExchangeChildren.get(0);
+ workerIdToServerInstanceMap =
assignWorkersForLocalExchange(localExchangeChildMetadata);
+ inheritPartitioning = shareOnePartitioning(localExchangeChildren);
+ }
// If there is no local exchange, assign workers to the servers hosting
the tables
List<QueryServerInstance> candidateServers = null;
@@ -624,12 +667,12 @@ public class WorkerManager {
// With a local exchange peer the worker map is copied from it, so the
classes come along; without one it comes
// from the candidate servers, whose worker ids are not classes at all.
childMetadata.setPartitionClassIds(
- localExchangeChildMetadata != null ?
localExchangeChildMetadata.getPartitionClassIds() : null);
+ inheritPartitioning ?
localExchangeChildMetadata.getPartitionClassIds() : null);
}
}
metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap);
- if (localExchangeChildMetadata != null) {
+ if (inheritPartitioning) {
metadata.setPartitionFunction(localExchangeChildMetadata.getPartitionFunction());
metadata.setPartitionClassIds(localExchangeChildMetadata.getPartitionClassIds());
} else {
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 66391119edd..6376839d748 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
@@ -22,6 +22,7 @@ import com.google.common.base.Throwables;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -1274,8 +1275,8 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
/// The `setOpOptions(is_colocated_by_set_op_keys='true')` hint forces a
pre-partitioned (direct) exchange on
/// every input of a set operation, avoiding the shuffle. Here the inputs
project `col3`, which is neither
- /// table's partition column, so without the hint the planner would shuffle.
Covers UNION ALL, INTERSECT and EXCEPT,
- /// which all share
[org.apache.pinot.calcite.rel.rules.PinotSetOpExchangeNodeInsertRule].
+ /// table's partition column, so without the hint the planner would shuffle.
UNION ALL is not covered here: it
+ /// gets a local exchange with or without the hint (see
[#testUnionAllUsesLocalExchange]).
@Test(dataProvider = "setOpColocationHintQueries")
public void testSetOpColocationHintForcesPrePartitionedExchange(String
query) {
List<MailboxSendNode> sendNodes =
findSetOpInputSendNodes(_queryEnvironment.planQuery(query));
@@ -1290,7 +1291,6 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
private Object[][] setOpColocationHintQueries() {
String hint = "/*+ setOpOptions(is_colocated_by_set_op_keys='true') */";
return new Object[][]{
- {"SELECT " + hint + " col3 FROM a UNION ALL SELECT col3 FROM b"},
{"SELECT " + hint + " col3 FROM a INTERSECT SELECT col3 FROM b"},
{"SELECT " + hint + " col3 FROM a EXCEPT SELECT col3 FROM b"},
};
@@ -1301,7 +1301,7 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
@Test
public void testSetOpColocationHintViaOuterSelectWrap() {
String query = "SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='true') */ * FROM "
- + "(SELECT col3 FROM a UNION ALL SELECT col3 FROM b)";
+ + "(SELECT col3 FROM a INTERSECT SELECT col3 FROM b)";
for (MailboxSendNode sendNode :
findSetOpInputSendNodes(_queryEnvironment.planQuery(query))) {
assertTrue(sendNode.isPrePartitioned(), "Hint on the wrapping SELECT
should force a pre-partitioned exchange");
}
@@ -1313,24 +1313,140 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
@Test
public void testSetOpColocationHintFirstInputWins() {
String query = "SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='true') */ col3 FROM a "
- + "UNION ALL SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='false') */ col3 FROM a";
+ + "INTERSECT SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='false') */ col3 FROM a";
for (MailboxSendNode sendNode :
findSetOpInputSendNodes(_queryEnvironment.planQuery(query))) {
assertTrue(sendNode.isPrePartitioned(),
"The first input's hint value should win and apply to all inputs
when branches conflict");
}
}
- /// Without the hint and with inputs that are not partitioned by the
projected column, the exchanges below the set
- /// operation must be regular (shuffled) exchanges.
+ /// Without the hint and with inputs that are not partitioned by the
projected column, the exchanges below a
+ /// distinct set operation must be regular (shuffled) exchanges. UNION ALL
is exempt: it only concatenates, so it
+ /// gets a local exchange instead (see [#testUnionAllUsesLocalExchange]).
+ @Test(dataProvider = "shuffledSetOpQueries")
+ public void testDistinctSetOpWithoutHintIsNotPrePartitioned(String query) {
+ for (MailboxSendNode sendNode :
findSetOpInputSendNodes(_queryEnvironment.planQuery(query))) {
+ assertFalse(sendNode.isPrePartitioned(),
+ "Without the hint and matching partitioning, the set op exchanges
should be a full shuffle");
+ }
+ }
+
+ @DataProvider(name = "shuffledSetOpQueries")
+ private Object[][] shuffledSetOpQueries() {
+ return new Object[][]{
+ {"SELECT col3 FROM a INTERSECT SELECT col3 FROM b"},
+ {"SELECT col3 FROM a EXCEPT SELECT col3 FROM b"},
+ };
+ }
+
+ /// UNION ALL only concatenates, so any row-to-worker mapping is correct and
no redistribution is required. Its
+ /// input exchanges are local (SINGLETON) exchanges: the union stage
inherits its inputs' workers and the rows are
+ /// handed over in place. Note SINGLETON is Pinot's marker for a local
exchange, not a gather to one node.
+ /// The projected columns ride along as keys -- unused while the exchange
stays local, but they let the mailbox
+ /// layer promote it to a real hash shuffle when the branches cannot share a
worker assignment.
@Test
- public void testSetOpWithoutHintIsNotPrePartitioned() {
+ public void testUnionAllUsesLocalExchange() {
+ // Both branches read table a, so they resolve to the same workers and the
local exchange survives to the plan.
+ String query = "SELECT col3 FROM a UNION ALL SELECT col3 FROM a";
+ List<MailboxSendNode> sendNodes =
findSetOpInputSendNodes(_queryEnvironment.planQuery(query));
+ assertFalse(sendNodes.isEmpty());
+ for (MailboxSendNode sendNode : sendNodes) {
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.SINGLETON);
+ assertEquals(sendNode.getKeys(), List.of(0),
+ "A local UNION ALL exchange should carry the projected columns as
keys");
+ }
+ }
+
+ /// When the branches do not resolve to the same workers -- table a lives on
two servers and table b on one -- the
+ /// union stage cannot inherit both. It takes the two-worker layout, so
branch a is still wired 1-to-1 and stays
+ /// SINGLETON, while branch b is promoted to a real hash shuffle on the
projected columns. Correct either way for
+ /// a concatenation, and the aligned branch still pays nothing.
+ @Test
+ public void testUnionAllWithMisalignedBranchesShufflesOnlyTheMisalignedOne()
{
+ String query = "SELECT col3 FROM a UNION ALL SELECT col3 FROM b";
+ List<MailboxSendNode> sendNodes =
findSetOpInputSendNodes(_queryEnvironment.planQuery(query));
+ assertEquals(sendNodes.size(), 2);
+ Set<RelDistribution.Type> types = new HashSet<>();
+ for (MailboxSendNode sendNode : sendNodes) {
+ types.add(sendNode.getDistributionType());
+ }
+ assertEquals(types, Set.of(RelDistribution.Type.SINGLETON,
RelDistribution.Type.HASH_DISTRIBUTED),
+ "The aligned branch should stay local and the misaligned one should be
promoted to a hash shuffle");
+ }
+
+ /// Branches of the SAME WIDTH on DIFFERENT servers still stay local. The
union stage adopts one branch's layout,
+ /// and the other sends 1-to-1 by worker id -- correct for a concatenation
whichever server each worker sits on,
+ /// costing a network hop rather than a shuffle. Requiring identical worker
maps here would push this into a full
+ /// shuffle, which is worse than what it replaces.
+ @Test
+ public void testUnionAllWithEqualWidthBranchesOnDifferentServersStaysLocal()
{
+ // Table a lives only on server 1 and table b only on server 2, so both
resolve to one worker on disjoint servers.
+ QueryEnvironment queryEnvironment = getQueryEnvironment(3, 1, 2,
TABLE_SCHEMAS,
+ Map.of("a_REALTIME", List.of("a1")), Map.of("b_REALTIME",
List.of("b1")), null);
String query = "SELECT col3 FROM a UNION ALL SELECT col3 FROM b";
+ List<MailboxSendNode> sendNodes =
findSetOpInputSendNodes(queryEnvironment.planQuery(query));
+ assertEquals(sendNodes.size(), 2);
+ for (MailboxSendNode sendNode : sendNodes) {
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.SINGLETON,
+ "Equal-width branches should stay local even when their workers sit
on different servers");
+ }
+ }
+
+ /// `setOpOptions(is_colocated_by_set_op_keys='false')` opts UNION ALL out
of the local exchange and restores the
+ /// full-row hash shuffle.
+ @Test
+ public void testUnionAllHintFalseRestoresShuffle() {
+ String query = "SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='false') */ col3 FROM a "
+ + "UNION ALL SELECT col3 FROM b";
for (MailboxSendNode sendNode :
findSetOpInputSendNodes(_queryEnvironment.planQuery(query))) {
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.HASH_DISTRIBUTED);
assertFalse(sendNode.isPrePartitioned(),
- "Without the hint and matching partitioning, the set op exchanges
should be a full shuffle");
+ "setOpOptions(is_colocated_by_set_op_keys='false') should restore
the shuffle for UNION ALL");
}
}
+ /// A non-UNION-ALL set operation's input exchanges shuffle on the full row,
so its output is hash distributed on
+ /// all of its columns and a downstream exchange keyed on those columns is
auto-detected as pre-partitioned.
+ /// INTERSECT ALL is used (rather than INTERSECT) so the aggregate above
survives planning: a distinct set op's
+ /// output is unique on the group keys, which lets Calcite remove the
aggregate and its exchange entirely.
+ @Test
+ public void testSetOpOutputIsHashDistributedOnAllColumns() {
+ String query = "SELECT col1, col2, COUNT(*) FROM "
+ + "(SELECT col1, col2 FROM a INTERSECT ALL SELECT col1, col2 FROM b)
GROUP BY col1, col2";
+ MailboxSendNode sendNode =
findSendNodeAboveSetOp(_queryEnvironment.planQuery(query));
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.HASH_DISTRIBUTED);
+ assertTrue(sendNode.isPrePartitioned(),
+ "An exchange keyed on all columns of a shuffled set op should
auto-detect pre-partitioning");
+ }
+
+ /// UNION ALL's output carries no distribution guarantee: a local exchange
does not redistribute anything, so it
+ /// promises nothing about where equal rows land. A downstream exchange
keyed on all of its columns must NOT be
+ /// auto-detected as pre-partitioned -- the deduplicating aggregate above it
needs a real shuffle to be correct.
+ @Test
+ public void testUnionAllOutputIsNotTreatedAsHashDistributed() {
+ String query = "SELECT col1, col2, COUNT(*) FROM "
+ + "(SELECT col1, col2 FROM a UNION ALL SELECT col1, col2 FROM b) GROUP
BY col1, col2";
+ MailboxSendNode sendNode =
findSendNodeAboveSetOp(_queryEnvironment.planQuery(query));
+ assertFalse(sendNode.isPrePartitioned(),
+ "UNION ALL output must not be treated as hash distributed: its inputs
are not shuffled");
+ }
+
+ /// `is_colocated_by_set_op_keys='true'` asserts that rows equal across all
projected columns already share a
+ /// worker, which is exactly the all-column hash distribution. The set op's
output is therefore claimed as hash
+ /// distributed just as it is for genuinely shuffled inputs, so a downstream
exchange keyed on those columns is
+ /// auto-detected as pre-partitioned too and the colocation carries all the
way up. Same query as
+ /// [#testSetOpOutputIsHashDistributedOnAllColumns] plus the hint.
+ @Test
+ public void testHintedColocatedSetOpOutputIsHashDistributed() {
+ String query = "SELECT col1, col2, COUNT(*) FROM "
+ + "(SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='true') */
col1, col2 FROM a "
+ + "INTERSECT ALL SELECT col1, col2 FROM b) GROUP BY col1, col2";
+ MailboxSendNode sendNode =
findSendNodeAboveSetOp(_queryEnvironment.planQuery(query));
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.HASH_DISTRIBUTED);
+ assertTrue(sendNode.isPrePartitioned(),
+ "A hint-forced colocated set op's output should be treated as hash
distributed on all columns");
+ }
+
/// When each input is declared partitioned by the projected column, the
single-column set-op exchange matches the
/// input partitioning, so the planner auto-detects a pre-partitioned
exchange even without the hint. This is the
/// baseline that [#testSetOpColocationHintFalseDisablesAutoDetected]
overrides.
@@ -1358,7 +1474,7 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
// op matches the input partitioning and the planner auto-detects
pre-partitioning.
private static final String AUTO_DETECTED_SET_OP =
"SELECT col2 FROM a /*+ tableOptions(partition_function='hashcode',
partition_key='col2', partition_size='4') */ "
- + "UNION ALL "
+ + "INTERSECT "
+ "SELECT col1 FROM b /*+
tableOptions(partition_function='hashcode', partition_key='col1', "
+ "partition_size='4') */";
@@ -1386,6 +1502,18 @@ public class QueryCompilationTest extends
QueryEnvironmentTestBase {
return sendNodes;
}
+ /// Finds the [MailboxSendNode] at the root of the fragment containing the
(single) set-op node, i.e. the exchange
+ /// that ships the set operation's output to its consumer stage.
+ private MailboxSendNode findSendNodeAboveSetOp(DispatchableSubPlan
dispatchableSubPlan) {
+ for (DispatchablePlanFragment fragment :
dispatchableSubPlan.getQueryStages()) {
+ PlanNode root = fragment.getPlanFragment().getFragmentRoot();
+ if (root instanceof MailboxSendNode && findNodeOfType(root,
SetOpNode.class) != null) {
+ return (MailboxSendNode) root;
+ }
+ }
+ throw new AssertionError("Expected a fragment rooted at a MailboxSendNode
containing a SetOp node");
+ }
+
/// 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 [#colocatedChains]). The
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java
index b9de8291778..5e8b28eeb1e 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java
@@ -182,9 +182,9 @@ public class QueryPlannerRuleOptionsTest extends
QueryEnvironmentTestBase {
+ " PinotLogicalExchange(distribution=[hash[0]])\n"
+ " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n"
+ " LogicalUnion(all=[true])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " LogicalValues(tuples=[[]])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " LogicalProject(col1=[$0])\n"
+ " PinotLogicalTableScan(table=[[default, b]])\n");
//@formatter:on
@@ -629,12 +629,12 @@ public class QueryPlannerRuleOptionsTest extends
QueryEnvironmentTestBase {
+ " PinotLogicalExchange(distribution=[hash[0]])\n"
+ " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n"
+ " LogicalUnion(all=[true])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n"
+ " PinotLogicalExchange(distribution=[hash[0]])\n"
+ " PinotLogicalAggregate(group=[{0}],
aggType=[LEAF])\n"
+ " PinotLogicalTableScan(table=[[default, a]])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n"
+ " PinotLogicalExchange(distribution=[hash[0]])\n"
+ " PinotLogicalAggregate(group=[{0}],
aggType=[LEAF])\n"
@@ -664,10 +664,10 @@ public class QueryPlannerRuleOptionsTest extends
QueryEnvironmentTestBase {
+ " PinotLogicalExchange(distribution=[hash[0]])\n"
+ " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n"
+ " LogicalUnion(all=[true])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " LogicalProject(col1=[$0])\n"
+ " PinotLogicalTableScan(table=[[default, a]])\n"
- + " PinotLogicalExchange(distribution=[hash[0]])\n"
+ + " PinotLogicalExchange(distribution=[single])\n"
+ " LogicalProject(col1=[$0])\n"
+ " PinotLogicalTableScan(table=[[default, b]])\n");
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java
index 33476c72f87..9bc7a7eae1d 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/MailboxAssignmentVisitorTest.java
@@ -19,8 +19,10 @@
package org.apache.pinot.query.planner.physical;
import java.util.ArrayList;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import org.apache.calcite.rel.RelDistribution;
import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.common.utils.DataSchema;
@@ -106,14 +108,69 @@ public class MailboxAssignmentVisitorTest extends
QueryEnvironmentTestBase {
assertEquals(singleMailbox(receiverMailboxes, 1,
SENDER_STAGE).getHostname(), "host_B");
}
- /// An unequal-but-non-multiple worker count (2 senders, 3 receivers) must
be rejected rather than rounding the
- /// parallelism down to 1 and silently dropping the extra receiver.
+ /// A KEYED local exchange (a UNION ALL input carries the projected columns)
whose worker counts do not divide
+ /// evenly is promoted to a real hash shuffle: HashExchange routes every key
consistently across the receivers, so
+ /// this is correct for a concatenation and for a keyed join alike.
+ @Test
+ public void testKeyedSingletonWithUnequalWorkersShuffles() {
+ DispatchablePlanMetadata sender = metadata(Map.of(0, server("A"), 1,
server("B")));
+ DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1,
server("B"), 2, server("C")));
+ MailboxSendNode sendNode = singletonSendNode(List.of(0));
+ process(sendNode, sender, receiver);
+
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.HASH_DISTRIBUTED);
+ // Shuffled: every receiver worker reads from every sender worker.
+ for (int workerId = 0; workerId < 3; workerId++) {
+ assertEquals(expandedWorkerIds(receiver.getWorkerIdToMailboxesMap(),
workerId, SENDER_STAGE), List.of(0, 1));
+ }
+ }
+
+ /// The parallelism path addresses a whole receiver range at the range's
FIRST host, which is only valid when the
+ /// receiver map was derived from the sender. Here the single sender's two
receivers sit on DIFFERENT servers, so
+ /// that assumption does not hold: posting both to host_A would strand the
receiver on host_B until the deadline.
+ /// Co-residency is verified rather than assumed, and the exchange falls
back to a full shuffle.
+ @Test
+ public void testSingletonWithParallelismAcrossHostsShuffles() {
+ DispatchablePlanMetadata sender = metadata(Map.of(0, server("A")));
+ DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1,
server("B")));
+ MailboxSendNode sendNode = singletonSendNode(List.of(0));
+ process(sendNode, sender, receiver);
+
+ assertEquals(sendNode.getDistributionType(),
RelDistribution.Type.HASH_DISTRIBUTED);
+ Set<String> hosts = new HashSet<>();
+ for (MailboxInfo mailboxInfo :
sender.getWorkerIdToMailboxesMap().get(0).get(RECEIVER_STAGE).getMailboxInfos())
{
+ hosts.add(mailboxInfo.getHostname());
+ }
+ assertEquals(hosts, Set.of("host_A", "host_B"), "Each receiver must be
addressed at its own host");
+ }
+
+ /// A KEYLESS local exchange must still fail loudly when the workers do not
line up. This is the colocated
+ /// dynamic-broadcast semi-join build side: every receiver needs the WHOLE
build side to build its filter
+ /// (the non-colocated variant broadcasts for exactly that reason), so
redistributing it would silently
+ /// drop matches. A UNION ALL input never reaches here because it carries
the projected columns as keys.
@Test(expectedExceptions = IllegalStateException.class,
- expectedExceptionsMessageRegExp = ".*multiple of number of senders.*")
- public void testSingletonRejectsNonMultipleReceiverCount() {
+ expectedExceptionsMessageRegExp = ".*requires keys.*")
+ public void testKeylessSingletonWithUnequalWorkersStillFails() {
DispatchablePlanMetadata sender = metadata(Map.of(0, server("A"), 1,
server("B")));
DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1,
server("B"), 2, server("C")));
- process(singletonSendNode(List.of(0)), sender, receiver);
+ process(singletonSendNode(List.of()), sender, receiver);
+ }
+
+ /// A local exchange whose sender stage has no worker at all (a fully pruned
leaf) has nothing to wire 1-to-1, and
+ /// computing the parallelism would divide by zero. Every live receiver must
still get an entry, holding an empty
+ /// mailbox list, or its WorkerMetadata carries a null mailbox map and fails
during dispatch serialization.
+ @Test
+ public void testSingletonWithZeroSendersDoesNotDivideByZero() {
+ DispatchablePlanMetadata sender = metadata(Map.of());
+ DispatchablePlanMetadata receiver = metadata(Map.of(0, server("A"), 1,
server("B")));
+ process(singletonSendNode(List.of()), sender, receiver);
+
+ assertTrue(sender.getWorkerIdToMailboxesMap().isEmpty());
+ for (int workerId = 0; workerId < 2; workerId++) {
+ MailboxInfos mailboxInfos =
receiver.getWorkerIdToMailboxesMap().get(workerId).get(SENDER_STAGE);
+ assertNotNull(mailboxInfos, "Missing entry for worker: " + workerId);
+ assertTrue(mailboxInfos.getMailboxInfos().isEmpty());
+ }
}
/// A SINGLETON local exchange with parallelism (more receivers than
senders) does not assert co-location either: it
@@ -193,6 +250,21 @@ public class MailboxAssignmentVisitorTest extends
QueryEnvironmentTestBase {
assertTrue(receiver.getWorkerIdToMailboxesMap().isEmpty());
}
+ /// A pre-partitioned hash exchange whose sender stage has zero workers (all
its leaf segments were pruned) must not
+ /// be wired as a direct exchange: with the receiver stage also empty (as
when every branch of a UNION ALL is fully
+ /// pruned), the sender and receiver counts trivially "match" and computing
the fan-out parallelism would divide by
+ /// zero. It must fall back to the regular wiring, which is a no-op for
empty stages.
+ @Test
+ public void testPrePartitionedExchangeWithZeroWorkersFallsBackToShuffle() {
+ DispatchablePlanMetadata sender = metadata(Map.of());
+ sender.setPrePartitioned(true);
+ DispatchablePlanMetadata receiver = metadata(Map.of());
+ process(hashSendNode(), sender, receiver);
+
+ assertTrue(sender.getWorkerIdToMailboxesMap().isEmpty());
+ assertTrue(receiver.getWorkerIdToMailboxesMap().isEmpty());
+ }
+
private static QueryServerInstance server(String id) {
return new QueryServerInstance(id, "host_" + id, 1, 1);
}
diff --git a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
index 2f78a35ef93..364a54fda9d 100644
--- a/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
+++ b/pinot-query-planner/src/test/resources/queries/AggregatePlans.json
@@ -143,13 +143,13 @@
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[PERCENTILE($1,
50)], aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1, 2]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col2=[$0], sum_of_runs=[$1], $f2=[50])",
"\n PinotLogicalAggregate(group=[{0}],
agg#0=[PERCENTILE($1, 50)], aggType=[DIRECT])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n LogicalProject(col2=[$1], col3=[$2], $f2=[50])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1, 2]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col2=[$0], sum_of_runs=[$1], $f2=[50])",
"\n PinotLogicalAggregate(group=[{0}],
agg#0=[PERCENTILE($1, 50)], aggType=[DIRECT])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
diff --git
a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
index 5f6c4b84e9d..fbb428bad31 100644
--- a/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
+++ b/pinot-query-planner/src/test/resources/queries/ExplainPhysicalPlans.json
@@ -737,42 +737,42 @@
]
},
{
- "description": "set op (UNION ALL) without the colocation hint
shuffles (hash redistributes) its inputs to all set-op-stage workers",
+ "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",
"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] UNION_ALL",
- "\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 └── [1]@localhost:1|[0] MAIL_RECEIVE(SINGLETON)",
+ "\n ├── [2]@localhost:2|[1]
MAIL_SEND(SINGLETON)->{[1]@localhost:2|[1]} (Subtree Omitted)",
+ "\n └── [2]@localhost:1|[0]
MAIL_SEND(SINGLETON)->{[1]@localhost:1|[0]}",
"\n └── [2]@localhost:1|[0] PROJECT",
"\n └── [2]@localhost:1|[0] TABLE SCAN (a) null",
- "\n └── [1]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)",
- "\n ├── [3]@localhost:2|[1]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n └── [3]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]}",
+ "\n └── [1]@localhost:1|[0] MAIL_RECEIVE(SINGLETON)",
+ "\n ├── [3]@localhost:2|[1]
MAIL_SEND(SINGLETON)->{[1]@localhost:2|[1]} (Subtree Omitted)",
+ "\n └── [3]@localhost:1|[0]
MAIL_SEND(SINGLETON)->{[1]@localhost:1|[0]}",
"\n └── [3]@localhost:1|[0] PROJECT",
"\n └── [3]@localhost:1|[0] TABLE SCAN (a) null",
"\n"
]
},
{
- "description": "set op (UNION ALL) with the
is_colocated_by_set_op_keys hint uses a pre-partitioned (direct, 1-to-1)
exchange on every input to avoid the shuffle",
- "sql": "EXPLAIN IMPLEMENTATION PLAN FOR SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='true') */ col3 FROM a UNION ALL
SELECT col3 FROM a",
+ "description": "set op (UNION ALL) with
is_colocated_by_set_op_keys='false' opts out of the local exchange and shuffles
(hash redistributes) its inputs to all set-op-stage workers",
+ "sql": "EXPLAIN IMPLEMENTATION PLAN FOR SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='false') */ col3 FROM a UNION ALL
SELECT col3 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] UNION_ALL",
"\n └── [1]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)",
- "\n ├── [2]@localhost:2|[1]
MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n └── [2]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[1]@localhost:1|[0]}",
+ "\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 └── [1]@localhost:1|[0] MAIL_RECEIVE(HASH_DISTRIBUTED)",
- "\n ├── [3]@localhost:2|[1]
MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[1]@localhost:2|[1]} (Subtree
Omitted)",
- "\n └── [3]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)[PARTITIONED]->{[1]@localhost:1|[0]}",
+ "\n ├── [3]@localhost:2|[1]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]} (Subtree
Omitted)",
+ "\n └── [3]@localhost:1|[0]
MAIL_SEND(HASH_DISTRIBUTED)->{[1]@localhost:1|[0],[1]@localhost:2|[1]}",
"\n └── [3]@localhost:1|[0] PROJECT",
"\n └── [3]@localhost:1|[0] TABLE SCAN (a) null",
"\n"
diff --git a/pinot-query-planner/src/test/resources/queries/SetOpPlans.json
b/pinot-query-planner/src/test/resources/queries/SetOpPlans.json
index 1aa229eb9af..9396598c9ad 100644
--- a/pinot-query-planner/src/test/resources/queries/SetOpPlans.json
+++ b/pinot-query-planner/src/test/resources/queries/SetOpPlans.json
@@ -7,10 +7,10 @@
"output": [
"Execution Plan",
"\nLogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col1=[$0], col2=[$1])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col1=[$0], col2=[$1])",
"\n PinotLogicalTableScan(table=[[default, b]])",
"\n"
@@ -22,13 +22,13 @@
"output": [
"Execution Plan",
"\nLogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col1=[$0], col2=[$1])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col1=[$0], col2=[$1])",
"\n PinotLogicalTableScan(table=[[default, b]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col1=[$0], col2=[$1])",
"\n PinotLogicalTableScan(table=[[default, c]])",
"\n"
@@ -43,22 +43,22 @@
"\n PinotLogicalExchange(distribution=[hash[0, 1]])",
"\n PinotLogicalAggregate(group=[{0, 1}], aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0, 1}], aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0, 1]])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0,
1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0,
1]])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[LEAF])",
"\n PinotLogicalTableScan(table=[[default,
a]])",
- "\n PinotLogicalExchange(distribution=[hash[0,
1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0,
1]])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[LEAF])",
"\n PinotLogicalTableScan(table=[[default,
b]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0, 1}], aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0, 1]])",
"\n PinotLogicalAggregate(group=[{0, 1}],
aggType=[LEAF])",
@@ -95,12 +95,12 @@
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{1}],
agg#0=[$SUM0($2)], aggType=[LEAF])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{1}],
agg#0=[$SUM0($2)], aggType=[LEAF])",
@@ -117,12 +117,12 @@
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[COUNT($1)],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{1}], agg#0=[COUNT()],
aggType=[LEAF])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[COUNT($1)],
aggType=[FINAL])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{1}], agg#0=[COUNT()],
aggType=[LEAF])",
@@ -138,12 +138,12 @@
"\nPinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[DIRECT])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[DIRECT])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n LogicalProject(col2=[$1], col3=[$2])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[DIRECT])",
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n LogicalProject(col2=[$1], col3=[$2])",
@@ -160,10 +160,10 @@
"\n PinotLogicalExchange(distribution=[hash[0]])",
"\n PinotLogicalAggregate(group=[{0}], agg#0=[$SUM0($1)],
aggType=[LEAF])",
"\n LogicalUnion(all=[true])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col2=[$1], col3=[$2])",
"\n PinotLogicalTableScan(table=[[default, a]])",
- "\n PinotLogicalExchange(distribution=[hash[0, 1]])",
+ "\n PinotLogicalExchange(distribution=[single])",
"\n LogicalProject(col2=[$1], col3=[$2])",
"\n PinotLogicalTableScan(table=[[default, b]])",
"\n"
diff --git a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
index 913b24df0ab..05cd92da01e 100644
--- a/pinot-query-runtime/src/test/resources/queries/QueryHints.json
+++ b/pinot-query-runtime/src/test/resources/queries/QueryHints.json
@@ -167,6 +167,10 @@
"description": "INTERSECT on the partition column num with
is_colocated_by_set_op_keys hint forces a pre-partitioned (direct) exchange;
tbl1 and tbl2 are both partitioned by num so results stay correct",
"sql": "SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='true') */
{tbl1}.num FROM {tbl1} INTERSECT SELECT {tbl2}.num FROM {tbl2}"
},
+ {
+ "description": "GROUP BY above a colocated INTERSECT: the hint asserts
equal rows already share a worker, so the set op's output is claimed hash
distributed on all columns and the aggregate above skips its shuffle too -- it
must still see complete groups",
+ "sql": "SELECT t.num, COUNT(*) FROM (SELECT /*+
setOpOptions(is_colocated_by_set_op_keys='true') */ {tbl1}.num AS num FROM
{tbl1} INTERSECT SELECT {tbl2}.num AS num FROM {tbl2}) AS t GROUP BY t.num"
+ },
{
"description": "EXCEPT on the partition column num with
is_colocated_by_set_op_keys hint forces a pre-partitioned (direct) exchange;
tbl1 and tbl2 are both partitioned by num so results stay correct",
"sql": "SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='true') */
{tbl1}.num FROM {tbl1} EXCEPT SELECT {tbl2}.num FROM {tbl2}"
@@ -191,6 +195,26 @@
{
"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
[...]
+ },
+ {
+ "description": "UNION ALL without any hint uses local (SINGLETON)
exchanges; UNION ALL only concatenates so a direct exchange is always correct",
+ "sql": "SELECT {tbl1}.num FROM {tbl1} UNION ALL SELECT {tbl2}.num FROM
{tbl2}"
+ },
+ {
+ "description": "UNION ALL of heterogeneous non-partition columns
without any hint: the local exchange is correct under any row-to-worker
mapping",
+ "sql": "SELECT {tbl1}.name FROM {tbl1} UNION ALL SELECT {tbl2}.val
FROM {tbl2}"
+ },
+ {
+ "description": "Distinct UNION deduplicates correctly above the local
UNION ALL exchange: the union's output claims no distribution, so the dedup
aggregate shuffles for correct results",
+ "sql": "SELECT {tbl1}.name FROM {tbl1} UNION SELECT {tbl2}.val FROM
{tbl2}"
+ },
+ {
+ "description": "GROUP BY above a hint-free UNION ALL deduplicates
correctly: the union's output claims no distribution, so the aggregate shuffles
for correct results",
+ "sql": "SELECT t.name, COUNT(*) FROM (SELECT {tbl1}.name AS name FROM
{tbl1} UNION ALL SELECT {tbl2}.val AS name FROM {tbl2}) AS t GROUP BY t.name"
+ },
+ {
+ "description": "is_partitioned_by_group_by_keys above a hint-free
self-UNION ALL: both branches route identically, so the assertion holds on the
physical data and the direct (un-shuffled) aggregate sees complete groups",
+ "sql": "SELECT /*+ aggOptions(is_partitioned_by_group_by_keys='true')
*/ t.num, COUNT(*) FROM (SELECT {tbl1}.num AS num FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */ UNION ALL SELECT {tbl1}.num AS num FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */) AS t GROUP BY t.num"
}
]
},
@@ -272,6 +296,10 @@
{
"description": "Forcing is_colocated_by_set_op_keys='true' when the
inputs have mismatched partition counts (tbl1=2, tbl2=4): the planner cannot
form a direct exchange, so it still shuffles and results stay correct",
"sql": "SELECT /*+ setOpOptions(is_colocated_by_set_op_keys='true') */
{tbl1}.num FROM {tbl1} /*+ tableOptions(partition_function='hashcode',
partition_key='num', partition_size='2') */ INTERSECT SELECT {tbl2}.num FROM
{tbl2} /*+ tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */"
+ },
+ {
+ "description": "Hint-free UNION ALL over inputs with mismatched
partition counts (tbl1=2, tbl2=4): the branches cannot share one worker
assignment, so the local exchange degrades to an arbitrary fan-out and results
stay correct",
+ "sql": "SELECT {tbl1}.num FROM {tbl1} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='2') */ UNION ALL SELECT {tbl2}.num FROM {tbl2} /*+
tableOptions(partition_function='hashcode', partition_key='num',
partition_size='4') */"
}
]
},
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
