xiangfu0 commented on code in PR #19874:
URL: https://github.com/apache/pinot/pull/19874#discussion_r4278317660
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -455,6 +440,68 @@ private void assignWorkersToNonRootFragment(PlanFragment
fragment, DispatchableP
}
}
+ /// 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)) {
Review Comment:
Nice that the default hook returns the same list, so `.equals` is trivially
true and colocation is preserved. One thing to note: this comparison is
order-sensitive, so a subclass override that returns the same servers in a
different order would be treated as "changed" and forced into a shuffle. That's
a perf pessimization only, never a correctness issue, but it might be worth a
one-line note on the hook that returning the identical/same-order list
preserves colocation.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -455,6 +440,68 @@ private void assignWorkersToNonRootFragment(PlanFragment
fragment, DispatchableP
}
}
+ /// 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);
Review Comment:
Good that this mirrors the existing intermediate-fragment path (set the
shuffle function, leave `partitionClassIds` unset so it defaults to null).
Since the lookup-join metadata is fresh here, no stale class ids can leak
through. Clear and consistent.
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java:
##########
@@ -455,6 +440,68 @@ private void assignWorkersToNonRootFragment(PlanFragment
fragment, DispatchableP
}
}
+ /// 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
+ Map<String, List<String>> segmentsMap = Map.of(TableType.OFFLINE.name(),
List.of());
Review Comment:
This faked segments map is carried over unchanged from the old code, so no
concern for this PR. Minor readability note for future readers: the `OFFLINE`
key and empty list are just a placeholder so the fragment counts as a leaf
stage, not real segment data. The existing comment covers the "why" but not
that the key/value are placeholders.
##########
pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java:
##########
@@ -2298,6 +2300,118 @@ ColocatedTableSpec unavailableSegments(List<String>
unavailableSegments) {
}
}
+ 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));
Review Comment:
Solid coverage of worker placement across all four scenarios. One optional
addition: these assert placement but not the partitioning contract directly. A
quick assert that the redistribution case yields
`DEFAULT_SHUFFLE_PARTITION_FUNCTION` with null class ids (and the inherit case
copies the child's) would pin that branch explicitly. Not blocking.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
