private static SourcePartitionedScheduler getSourcePartitionedScheduler(
StageExecutionPlan plan,
SqlStageExecution stage, NodeManager nodeManager, NodeTaskMap nodeTaskMap,
int splitBatchSize)
{
NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig()
.setIncludeCoordinator(false)
.setMaxSplitsPerNode(20)
.setMaxPendingSplitsPerNodePerTask(0);
NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap);
SplitSource splitSource = plan.getDataSource().get();
SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(splitSource.getDataSourceName()), stage::getAllTasks);
return new SourcePartitionedScheduler(stage, splitSource, placementPolicy, splitBatchSize);
}