yashmayya commented on code in PR #19874:
URL: https://github.com/apache/pinot/pull/19874#discussion_r4286627032
##########
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:
Good point. Added a note to the hook doc in 18ab7ac58e: return the input
workers unchanged, in the same order, to keep the exchange local and the input
partitioning.
##########
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:
Added to the comment in 18ab7ac58e: the `OFFLINE` key with no segments is
only a placeholder, because the join reads the dimension table from memory.
##########
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:
I would rather leave this out of this PR. The partition function and class
ids live only in `DispatchablePlanMetadata` during planning, and
`DispatchableSubPlan` does not expose them, so these asserts need new test
plumbing. Partial coverage exists: in the inherit case the exchange is direct,
so `checkPartitionClassAgreement` in `MailboxAssignmentVisitor` fails planning
if the copied class ids disagree with the input. The redistribution branch sets
the same values as `assignWorkersToIntermediateFragment` does for a stage that
does not inherit partitioning.
--
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]
