andygrove commented on code in PR #2721:
URL:
https://github.com/apache/datafusion-ballista/pull/2721#discussion_r4455500791
##########
ballista/core/src/execution_plans/plan_algebra.rs:
##########
@@ -122,6 +137,41 @@ pub fn as_partition_sliceable(
None
}
+/// Looks up the [`PartitionSliceable`] of a node that
[`as_partition_sliceable`]
+/// does not list, so the scheduler can restrict `ExecutionPlan` nodes of your
+/// own. Set it with `SchedulerConfig::with_partition_sliceable_resolver`.
+///
+/// It is asked after the built-in list, and never about a node that is read
+/// whole: one below a collapse, or declared `SinglePartition` by its parent.
It
+/// is called for most nodes of every task's plan as tasks are launched, so
keep
+/// it cheap, typically a `downcast_ref` to your own types.
+pub type PartitionSliceableResolver = Arc<
+ dyn for<'a> Fn(&'a Arc<dyn ExecutionPlan>) -> Option<&'a dyn
PartitionSliceable>
+ + Send
+ + Sync,
+>;
+
+/// Unwraps the only child that a single-input operator's
+/// [`restrict_to_partitions`](PartitionSliceable::restrict_to_partitions)
+/// receives.
+///
+/// # Arguments
+///
+/// * `owner` - operator name, for the error message
+/// * `children` - the children `restrict_to_partitions` received
+pub fn single_child(
+ owner: &str,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let [child] = <[_; 1]>::try_from(children).map_err(|children| {
+ DataFusionError::Internal(format!(
+ "{owner}: expected 1 child to restrict, got {}",
+ children.len()
+ ))
+ })?;
+ Ok(child)
+}
+
/// Take `values[global]` for each global partition index, in task-local
Review Comment:
Good catch. `slice_by_global_partition` and `single_child` now return
`plan_err!`, in 29493d90e.
##########
ballista/core/src/execution_plans/plan_algebra.rs:
##########
@@ -122,6 +137,41 @@ pub fn as_partition_sliceable(
None
}
+/// Looks up the [`PartitionSliceable`] of a node that
[`as_partition_sliceable`]
+/// does not list, so the scheduler can restrict `ExecutionPlan` nodes of your
+/// own. Set it with `SchedulerConfig::with_partition_sliceable_resolver`.
+///
+/// It is asked after the built-in list, and never about a node that is read
+/// whole: one below a collapse, or declared `SinglePartition` by its parent.
It
+/// is called for most nodes of every task's plan as tasks are launched, so
keep
+/// it cheap, typically a `downcast_ref` to your own types.
+pub type PartitionSliceableResolver = Arc<
Review Comment:
You're right, the orphan rule and the borrowed return type leave no way to
restrict a `DataSourceExec` over a custom source. The resolver now takes the
node, its restricted children and the partitions, and returns the restricted
node itself, as you suggested. The walker restricts children once and then asks
the built-in list and the resolver. There's a test where the resolver builds a
replacement for a node without `PartitionSliceable`. In 29493d90e.
--
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]
