coderfender commented on code in PR #2721:
URL:
https://github.com/apache/datafusion-ballista/pull/2721#discussion_r4451354851
##########
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:
Now that the guide points custom nodes at `slice_by_global_partition`, its
out-of-bounds error reaches users. As a `DataFusionError::Internal` it renders
with "This issue was likely caused by a bug in DataFusion's code. Please help
us ... file a bug report" (hit this from a resolver-backed node). Could it, and
`single_child`, use `plan_err!` like `checked_restriction` /
`ensure_leaf_runs_whole`?
##########
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:
With this PR a split stage whose `DataSourceExec` wraps a custom
`DataSource` fails, but I don't see a way for an embedder to make it
restrictable: `PartitionSliceable: ExecutionPlan` rules out implementing it on
the `DataSource`, the orphan rule rules out `DataSourceExec`, and the resolver
returns a reference borrowed from the plan, so it can't hand back a wrapper.
Would a resolver that does the restriction itself, e.g. `Fn(&Arc<dyn
ExecutionPlan>, Vec<Arc<dyn ExecutionPlan>>, &[usize]) -> Option<Result<Arc<dyn
ExecutionPlan>>>`, work? It would cover this case as well as custom nodes.
--
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]
