private static StageScheduler getSourcePartitionedScheduler(
StageExecutionPlan plan,
SqlStageExecution stage,
InternalNodeManager nodeManager,
NodeTaskMap nodeTaskMap,
int splitBatchSize)
{
NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig()
.setIncludeCoordinator(false)
.setMaxSplitsPerNode(20)
.setMaxPendingSplitsPerTask(0);
NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap);
PlanNodeId sourceNode = Iterables.getOnlyElement(plan.getSplitSources().keySet());
SplitSource splitSource = Iterables.getOnlyElement(plan.getSplitSources().values());
SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(splitSource.getConnectorId()), stage::getAllTasks);
return newSourcePartitionedSchedulerAsStageScheduler(stage, sourceNode, splitSource, placementPolicy, splitBatchSize);
}