andygrove opened a new issue, #2679: URL: https://github.com/apache/datafusion-ballista/issues/2679
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.** When the scheduler splits a stage between tasks, `restrict_plan_to_partitions` ([task_builder.rs#L65](https://github.com/apache/datafusion-ballista/blob/98d724b3bd9b603211f444d992b0bcca13c75bca/ballista/scheduler/src/state/task_builder.rs#L65)) rewrites the stage plan so that each task sees only its own partitions. It can only restrict two kinds of node: - the leaves that `select_output_partitions` knows ([task_builder.rs#L239-L253](https://github.com/apache/datafusion-ballista/blob/98d724b3bd9b603211f444d992b0bcca13c75bca/ballista/scheduler/src/state/task_builder.rs#L239-L253)): `ShuffleReaderExec`, `RangeShuffleReaderExec`, `DataSourceExec` over files or memory, `PlaceholderRowExec` and `EmptyExec`; - the `PartitionSliceable` operators returned by `as_partition_sliceable` ([plan_algebra.rs#L95-L123](https://github.com/apache/datafusion-ballista/blob/98d724b3bd9b603211f444d992b0bcca13c75bca/ballista/core/src/execution_plans/plan_algebra.rs#L95-L123)), which is a hand-maintained list of two Ballista operators. A downstream crate can implement `PartitionSliceable`, but it can't add its node to that list, and the trait only covers operators with a single child. Any other node that holds data per partition is left as it is. Operators only see task-local partition indices, so a task that runs partitions 4 and 5 of such a node executes its partitions 0 and 1. The stage then returns some partitions twice and others not at all, with no error. The comment on `select_output_partitions` describes this failure mode. Comet Native Executors (apache/datafusion-comet#7288) run each stage as a `CometFragmentExec`, which needs two things per partition. One is the partition's own Comet plan, which names its files. The other is the partition's index in the stage, which Spark uses to seed `rand`, `monotonically_increasing_id` and `spark_partition_id`. The prototype (apache/datafusion-comet#7289) works around the missing hook by giving every fragment an extra `MemorySourceConfig` child with one row per partition, holding that partition's index and plan, because that is a leaf the scheduler already restricts. **Describe the solution you'd like** A way for a downstream crate to make its own nodes restrictable: 1. Generalize `PartitionSliceable` to nodes with any number of children, including leaves. For example: `fn restrict_to_partitions(&self, children: Vec<Arc<dyn ExecutionPlan>>, partitions: &[usize]) -> Result<Arc<dyn ExecutionPlan>>`, where `children` are already restricted. 2. Look nodes up through something a downstream crate can extend, consulted after the built-in list. For example, a resolver on `SchedulerConfig` next to the existing `override_*` hooks, mapping an `Arc<dyn ExecutionPlan>` to `Option<&dyn PartitionSliceable>`. 3. Fail the job when a stage that is split between tasks contains a leaf that can't be restricted, instead of running it and duplicating its output. Longer term, the `ExecutionPlan::select_partitions(&[usize])` method that the comment on `select_output_partitions` suggests for DataFusion would make the list unnecessary. **Describe alternatives you've considered** - The descriptor child described above. It works, but every downstream node needs a fake leaf, and nothing catches a node that forgets one. - One partition per task doesn't avoid the problem: the task's plan is still restricted to that one partition, so an unrestricted node runs its partition 0 in every task. **Additional context** Declaring a child `SinglePartition` in `input_distribution_requirements` already keeps it whole ([task_builder.rs#L191](https://github.com/apache/datafusion-ballista/blob/98d724b3bd9b603211f444d992b0bcca13c75bca/ballista/scheduler/src/state/task_builder.rs#L191)). That covers inputs that every partition reads in full, such as a broadcast build side, so they don't need a separate hook. Part of apache/datafusion-comet#7288. -- 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]
