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 51496aebf94 Fix planning of lookup joins over an intermediate stage
(#19874)
51496aebf94 is described below
commit 51496aebf94451e09cf58cc4e4b46863ccc12585
Author: Yash Mayya <[email protected]>
AuthorDate: Thu Oct 15 14:55:12 2026 -0700
Fix planning of lookup joins over an intermediate stage (#19874)
---
.../apache/pinot/query/routing/WorkerManager.java | 89 ++++++++++++----
.../pinot/query/routing/WorkerManagerTest.java | 114 +++++++++++++++++++++
.../src/test/resources/queries/LookupJoin.json | 20 ++++
3 files changed, 204 insertions(+), 19 deletions(-)
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 9743ec88f67..3539df3720c 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
@@ -423,30 +423,15 @@ public class WorkerManager {
assignWorkersToNonRootFragment(child, context, leafOnly,
visitedFragmentIds);
}
Map<Integer, DispatchablePlanMetadata> metadataMap =
context.getDispatchablePlanMetadataMap();
- DispatchablePlanMetadata metadata =
metadataMap.get(fragment.getFragmentId());
- boolean isLeaf = isLeafPlan(metadata);
- if (leafOnly != isLeaf) {
+ if (leafOnly != isAssignedInLeafPass(fragment, metadataMap)) {
return;
}
- if (isLeaf) {
+ if (isLeafPlan(metadataMap.get(fragment.getFragmentId()))) {
// TODO: Revisit this logic and see if we can generalize this
// For LOOKUP join, join is leaf stage because there is no exchange
added to the right side of the join. When we
// find a single local exchange child in the leaf stage, assign workers
based on the local exchange child.
if (isLookupJoin(children)) {
- DispatchablePlanMetadata childMetadata =
metadataMap.get(children.get(0).getFragmentId());
- Map<Integer, QueryServerInstance> workerIdToServerInstanceMap =
assignWorkersForLocalExchange(childMetadata);
- metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap);
- metadata.setPartitionFunction(childMetadata.getPartitionFunction());
- // The worker map comes from the child, so the worker ids stand for
the same partition classes as the child's.
- metadata.setPartitionClassIds(childMetadata.getPartitionClassIds());
- // Fake a segments map so that the worker can be correctly identified
as leaf stage
- Map<String, List<String>> segmentsMap =
Map.of(TableType.OFFLINE.name(), List.of());
- Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap =
-
Maps.newHashMapWithExpectedSize(workerIdToServerInstanceMap.size());
- for (Integer workerId : workerIdToServerInstanceMap.keySet()) {
- workerIdToSegmentsMap.put(workerId, segmentsMap);
- }
- metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap);
+ assignWorkersToLookupJoinFragment(fragment, context, leafOnly);
} else {
assignWorkersToLeafFragment(fragment, context);
}
@@ -455,6 +440,69 @@ public class WorkerManager {
}
}
+ /// Returns whether the leaf pass assigns the workers of the given fragment.
A lookup join fragment scans the
+ /// dimension table, so it is a leaf fragment. But it takes its workers from
its input, so when the input is assigned
+ /// in the intermediate pass (e.g. an aggregation below the join), the
lookup join fragment is assigned in that pass
+ /// too.
+ private static boolean isAssignedInLeafPass(PlanFragment fragment,
+ Map<Integer, DispatchablePlanMetadata> metadataMap) {
+ if (!isLeafPlan(metadataMap.get(fragment.getFragmentId()))) {
+ return false;
+ }
+ List<PlanFragment> children = fragment.getChildren();
+ return !isLookupJoin(children) || isAssignedInLeafPass(children.get(0),
metadataMap);
+ }
+
+ /// Assigns the workers of a lookup join fragment from the workers of its
input, and fakes an empty segments map so
+ /// that the workers are identified as leaf stage workers. The join probes
the dimension table in memory, so it must
+ /// run on servers, which hold the dimension table. A leaf input runs on
servers, so the join takes the same workers.
+ /// An intermediate input can run on instances that cannot scan segments (see
+ /// [#getCandidateServersForReplicatedLeaf]), so then the instances come
from that hook.
+ ///
+ /// @param isInputLeaf whether the input of the join is assigned in the leaf
pass
+ private void assignWorkersToLookupJoinFragment(PlanFragment fragment,
DispatchablePlanContext context,
+ boolean isInputLeaf) {
+ Map<Integer, DispatchablePlanMetadata> metadataMap =
context.getDispatchablePlanMetadataMap();
+ DispatchablePlanMetadata metadata =
metadataMap.get(fragment.getFragmentId());
+ DispatchablePlanMetadata childMetadata =
metadataMap.get(fragment.getChildren().get(0).getFragmentId());
+ Map<Integer, QueryServerInstance> workerIdToServerInstanceMap =
assignWorkersForLocalExchange(childMetadata);
+ boolean inheritPartitioning = true;
+ if (!isInputLeaf) {
+ int numInputWorkers = workerIdToServerInstanceMap.size();
+ List<QueryServerInstance> inputWorkers = new
ArrayList<>(numInputWorkers);
+ for (int workerId = 0; workerId < numInputWorkers; workerId++) {
+ inputWorkers.add(workerIdToServerInstanceMap.get(workerId));
+ }
+ List<QueryServerInstance> workers =
getCandidateServersForReplicatedLeaf(context, inputWorkers);
+ if (!workers.equals(inputWorkers)) {
+ // The workers differ from the input's, so the exchange redistributes
the rows and the worker ids no longer
+ // stand for the input's partition classes.
+ workerIdToServerInstanceMap =
Maps.newHashMapWithExpectedSize(workers.size());
+ for (int workerId = 0; workerId < workers.size(); workerId++) {
+ workerIdToServerInstanceMap.put(workerId, workers.get(workerId));
+ }
+ inheritPartitioning = false;
+ }
+ }
+ metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap);
+ if (inheritPartitioning) {
+ metadata.setPartitionFunction(childMetadata.getPartitionFunction());
+ // The worker map comes from the child, so the worker ids stand for the
same partition classes as the child's.
+ metadata.setPartitionClassIds(childMetadata.getPartitionClassIds());
+ } else {
+ metadata.setPartitionFunction(DEFAULT_SHUFFLE_PARTITION_FUNCTION);
+ }
+ // Fake a segments map so that the worker can be correctly identified as
leaf stage. The OFFLINE key with no
+ // segments is only a placeholder: the join reads the dimension table from
memory, not from segments.
+ Map<String, List<String>> segmentsMap = Map.of(TableType.OFFLINE.name(),
List.of());
+ Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap =
+ Maps.newHashMapWithExpectedSize(workerIdToServerInstanceMap.size());
+ for (Integer workerId : workerIdToServerInstanceMap.keySet()) {
+ workerIdToSegmentsMap.put(workerId, segmentsMap);
+ }
+ metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap);
+ }
+
static boolean isLookupJoin(List<PlanFragment> children) {
if (children.size() != 1) {
return false;
@@ -807,7 +855,10 @@ public class WorkerManager {
}
/// Returns the instances to assign to replicated leaf stage children when
there is no local exchange peer. By
- /// default, uses the same candidates as the intermediate stage.
+ /// default, uses the same candidates as the intermediate stage. Also
returns the instances to assign to a lookup join
+ /// stage over an intermediate input, which are by default the workers of
that input. Return that list unchanged, in
+ /// the same order, to keep the exchange local and the input's partitioning.
Any other list makes the exchange
+ /// redistribute the rows.
///
/// Subclasses can override to use different instances for replicated leaf
stages (e.g., when intermediate stages
/// run on non-server instances that cannot scan segments).
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java
index eafaff9cca7..dfc526e6bfd 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java
@@ -33,6 +33,7 @@ import org.apache.pinot.common.request.BrokerRequest;
import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.Function;
import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.core.routing.MockRoutingManagerFactory;
import org.apache.pinot.core.routing.RoutingManager;
import org.apache.pinot.core.routing.RoutingTable;
import org.apache.pinot.core.routing.SegmentsToQuery;
@@ -41,6 +42,7 @@ import
org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo;
import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo;
import org.apache.pinot.core.transport.ServerInstance;
import org.apache.pinot.query.QueryEnvironment;
+import org.apache.pinot.query.planner.physical.DispatchablePlanContext;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
import org.apache.pinot.spi.config.table.TableConfig;
@@ -2298,6 +2300,118 @@ public class WorkerManagerTest {
}
}
+ private static final String LOOKUP_JOINS_OVER_AGGREGATE_QUERY =
+ "SELECT /*+ joinOptions(join_strategy='lookup') */ t.col1, t.cnt,
dim1.col2, dim2.col2 "
+ + "FROM (SELECT col1, COUNT(*) AS cnt FROM fact GROUP BY col1) t "
+ + "JOIN dim1 ON t.col1 = dim1.col1 JOIN dim2 ON t.col1 = dim2.col1";
+
+ /// A lookup join over an intermediate input (the aggregation of a subquery)
is assigned after its input. By default
+ /// it takes the workers of its input, like a lookup join over a leaf input.
The second lookup join takes the workers
+ /// of the first one, so it must be assigned after the first one too.
+ @Test
+ public void testLookupJoinOverIntermediateInputTakesInputWorkers() {
+ DispatchableSubPlan plan =
newLookupJoinQueryEnvironment(null).planQuery(LOOKUP_JOINS_OVER_AGGREGATE_QUERY);
+ DispatchablePlanFragment dim1Join = lookupJoinFragment(plan, "dim1");
+ DispatchablePlanFragment dim2Join = lookupJoinFragment(plan, "dim2");
+ Map<Integer, String> aggregateWorkers =
workerIdToServer(inputFragment(plan, dim1Join));
+ assertEquals(aggregateWorkers.size(), 2);
+ assertEquals(workerIdToServer(dim1Join), aggregateWorkers);
+ assertEquals(workerIdToServer(dim2Join), aggregateWorkers);
+ }
+
+ /// When the intermediate stages run on an instance that cannot scan
segments, like a broker, a lookup join over an
+ /// intermediate input runs on the servers from
[WorkerManager#getCandidateServersForReplicatedLeaf]. Only servers
+ /// hold the dimension tables.
+ @Test
+ public void testLookupJoinOverIntermediateInputUsesReplicatedLeafServers() {
+ QueryServerInstance broker = new QueryServerInstance("Broker_localhost_9",
"localhost", 9, 9);
+ DispatchableSubPlan plan =
newLookupJoinQueryEnvironment(broker).planQuery(LOOKUP_JOINS_OVER_AGGREGATE_QUERY);
+ DispatchablePlanFragment dim1Join = lookupJoinFragment(plan, "dim1");
+ DispatchablePlanFragment dim2Join = lookupJoinFragment(plan, "dim2");
+ assertEquals(workerIdToServer(inputFragment(plan, dim1Join)), Map.of(0,
broker.getInstanceId()));
+ Set<String> factServers = new
HashSet<>(workerIdToServer(leafFragmentForTable(plan, "fact")).values());
+ assertEquals(factServers.size(), 2);
+ assertEquals(new HashSet<>(workerIdToServer(dim1Join).values()),
factServers);
+ assertEquals(workerIdToServer(dim2Join), workerIdToServer(dim1Join));
+ }
+
+ /// A lookup join over a leaf input takes the workers of the leaf, also when
the intermediate stages run on an
+ /// instance that cannot scan segments.
+ @Test
+ public void testLookupJoinOverLeafInputTakesLeafWorkers() {
+ QueryServerInstance broker = new QueryServerInstance("Broker_localhost_9",
"localhost", 9, 9);
+ DispatchableSubPlan plan = newLookupJoinQueryEnvironment(broker).planQuery(
+ "SELECT /*+ joinOptions(join_strategy='lookup') */ fact.col1,
dim1.col2 "
+ + "FROM fact JOIN dim1 ON fact.col1 = dim1.col1");
+ DispatchablePlanFragment dim1Join = lookupJoinFragment(plan, "dim1");
+ DispatchablePlanFragment factLeaf = inputFragment(plan, dim1Join);
+ assertEquals(factLeaf, leafFragmentForTable(plan, "fact"));
+ assertEquals(workerIdToServer(dim1Join), workerIdToServer(factLeaf));
+ }
+
+ /// Builds a QueryEnvironment over the offline table "fact", with one
segment on each of 2 servers, and the offline
+ /// tables "dim1" and "dim2", with one segment on both servers like
dimension tables.
+ ///
+ /// @param intermediateStageWorker when set, every intermediate stage runs
on this instance, see
+ /// [NonServerIntermediateStageWorkerManager]
+ private static QueryEnvironment newLookupJoinQueryEnvironment(@Nullable
QueryServerInstance intermediateStageWorker) {
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(1, 2);
+ factory.registerTable(getSchemaBuilder("fact").build(), "fact_OFFLINE");
+ factory.registerSegment(1, "fact_OFFLINE", "fact_1");
+ factory.registerSegment(2, "fact_OFFLINE", "fact_2");
+ for (String dimTableName : List.of("dim1", "dim2")) {
+ String dimTableNameWithType = dimTableName + "_OFFLINE";
+ factory.registerTable(getSchemaBuilder(dimTableName).build(),
dimTableNameWithType);
+ factory.registerSegment(1, dimTableNameWithType, dimTableName + "_0");
+ factory.registerSegment(2, dimTableNameWithType, dimTableName + "_0");
+ }
+ RoutingManager routingManager = factory.buildRoutingManager(null);
+ WorkerManager workerManager = intermediateStageWorker == null
+ ? new WorkerManager("Broker_localhost", "localhost", 3, routingManager)
+ : new NonServerIntermediateStageWorkerManager(routingManager,
intermediateStageWorker);
+ return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE,
factory.buildTableCache(), workerManager);
+ }
+
+ /// Runs every intermediate stage on one instance that cannot scan segments,
like a broker. So it runs the replicated
+ /// leaf stages on the servers that host the tables.
+ private static class NonServerIntermediateStageWorkerManager extends
WorkerManager {
+ private final QueryServerInstance _intermediateStageWorker;
+
+ NonServerIntermediateStageWorkerManager(RoutingManager routingManager,
+ QueryServerInstance intermediateStageWorker) {
+ super("Broker_localhost", "localhost", 3, routingManager);
+ _intermediateStageWorker = intermediateStageWorker;
+ }
+
+ @Override
+ protected List<QueryServerInstance>
getCandidateServers(DispatchablePlanContext context) {
+ // Mutable, because the caller sorts it
+ return new ArrayList<>(List.of(_intermediateStageWorker));
+ }
+
+ @Override
+ protected List<QueryServerInstance>
getCandidateServersForReplicatedLeaf(DispatchablePlanContext context,
+ List<QueryServerInstance> intermediateStageWorkers) {
+ return getCandidateServersPerTables(context);
+ }
+ }
+
+ /// Returns the lookup join fragment that probes the given dimension table.
It scans the dimension table, has the
+ /// input of the join as its only child, and every worker is a leaf stage
worker.
+ private static DispatchablePlanFragment
lookupJoinFragment(DispatchableSubPlan dispatchableSubPlan,
+ String dimTableName) {
+ DispatchablePlanFragment fragment =
leafFragmentForTable(dispatchableSubPlan, dimTableName);
+ assertEquals(fragment.getPlanFragment().getChildren().size(), 1, "Not a
lookup join fragment: " + dimTableName);
+ assertEquals(fragment.getWorkerIdToSegmentsMap().keySet(),
workerIdToServer(fragment).keySet());
+ return fragment;
+ }
+
+ /// Returns the fragment that feeds the given fragment, which has a single
child.
+ private static DispatchablePlanFragment inputFragment(DispatchableSubPlan
dispatchableSubPlan,
+ DispatchablePlanFragment fragment) {
+ return
dispatchableSubPlan.getQueryStageMap().get(fragment.getPlanFragment().getChildren().get(0).getFragmentId());
+ }
+
/// Returns the only fragment below the reduce stage with neither segments
nor children, i.e. the join stage.
private static DispatchablePlanFragment joinFragment(DispatchableSubPlan
dispatchableSubPlan) {
for (Map.Entry<Integer, DispatchablePlanFragment> entry :
dispatchableSubPlan.getQueryStageMap().entrySet()) {
diff --git a/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
b/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
index 93b0468d7d2..d5ed8dcabd5 100644
--- a/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
+++ b/pinot-query-runtime/src/test/resources/queries/LookupJoin.json
@@ -73,6 +73,26 @@
["gamma", 1]
],
"ignoreLiteMode": true
+ },
+ {
+ "description": "Lookup join over an aggregation - the input of the
join is an intermediate stage, so the join stage is assigned after it",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */ t.dim_key,
t.cnt, {dim_tbl}.name FROM (SELECT dim_key, COUNT(*) AS cnt FROM {fact_tbl}
GROUP BY dim_key) t JOIN {dim_tbl} ON t.dim_key = {dim_tbl}.id",
+ "outputs": [
+ [1, 2, "alpha"],
+ [2, 1, "beta"],
+ [3, 1, "gamma"]
+ ],
+ "ignoreLiteMode": true
+ },
+ {
+ "description": "Two lookup joins over an aggregation - the upper join
takes the workers of the lower one, which is assigned after the aggregation",
+ "sql": "SELECT /*+ joinOptions(join_strategy='lookup') */ t.dim_key,
t.cnt, d1.name, d2.name FROM (SELECT dim_key, COUNT(*) AS cnt FROM {fact_tbl}
GROUP BY dim_key) t JOIN {dim_tbl} d1 ON t.dim_key = d1.id JOIN {dim_tbl} d2 ON
t.dim_key = d2.id",
+ "outputs": [
+ [1, 2, "alpha", "alpha"],
+ [2, 1, "beta", "beta"],
+ [3, 1, "gamma", "gamma"]
+ ],
+ "ignoreLiteMode": true
}
]
},
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
