shauryachats commented on code in PR #16204:
URL: https://github.com/apache/pinot/pull/16204#discussion_r2167818457
##########
pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LiteModeWorkerAssignmentRule.java:
##########
@@ -128,6 +129,49 @@ static String stripIdPrefixFromWorker(String worker) {
return worker.split("@")[1];
}
+ /**
+ * Infers Exchange to be added on top of the leaf stage.
+ */
+ private PhysicalExchange inferPDDForLeafExchange(PRelNode leafStageRoot,
List<String> liteModeWorkers) {
+ RelCollation collation =
leafStageRoot.unwrap().getTraitSet().getCollation();
+ PinotDataDistribution pdd;
+ if (collation != null && !collation.getFieldCollations().isEmpty()) {
+ // If the leaf stage root has a collation trait, then we will use a
sorted receive in the exchange, so we can
+ // add the collation to the PDD.
+ pdd = new PinotDataDistribution(
+ RelDistribution.Type.SINGLETON, liteModeWorkers,
liteModeWorkers.hashCode(), null, collation);
+ } else {
+ pdd = new PinotDataDistribution(
+ RelDistribution.Type.SINGLETON, liteModeWorkers,
liteModeWorkers.hashCode(), null, null);
+ }
+ return new PhysicalExchange(nodeId(), leafStageRoot, pdd,
Collections.emptyList(),
+ ExchangeStrategy.SINGLETON_EXCHANGE, collation,
PinotExecStrategyTrait.getDefaultExecStrategy());
+ }
+
+ /**
+ * Infers distribution for the current node based on its inputs and
node-type. Can also add collation to the PDD
+ * automatically (e.g. if the current node is a Sort or the input is sorted
and this node does not drop collation).
+ */
+ private static PinotDataDistribution inferPDD(PRelNode currentNode,
List<PRelNode> newInputs,
+ List<String> liteModeWorkers) {
+ if (currentNode instanceof Sort) {
+ Sort sort = (Sort) currentNode.unwrap();
+ return new PinotDataDistribution(RelDistribution.Type.SINGLETON,
liteModeWorkers,
+ liteModeWorkers.hashCode(), null, sort.getCollation());
+ }
+ PinotDataDistribution currentNodePDD;
+ if (newInputs.isEmpty()) {
+ // Can happen for Values node.
+ currentNodePDD = new
PinotDataDistribution(RelDistribution.Type.SINGLETON, liteModeWorkers,
Review Comment:
nit: We can just `return` here.
--
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: commits-unsubscr...@pinot.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
