NoahKusaba commented on code in PR #2721:
URL:
https://github.com/apache/datafusion-ballista/pull/2721#discussion_r4521080875
##########
ballista/scheduler/src/state/task_builder.rs:
##########
@@ -795,4 +864,204 @@ mod tests {
assert_eq!(reader.partition[0][0].partition_id.partition_id, 1);
assert_eq!(reader.partition[1][0].partition_id.partition_id, 3);
}
+
+ // --- Nodes the scheduler does not know: `PayloadExec` stands in for one
---
+
+ use crate::test_utils::payload::{PayloadExec, payload_resolver,
payloads_of};
+ use
ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver;
+ use std::sync::Mutex;
+
+ fn partition_count(plan: &Arc<dyn ExecutionPlan>) -> usize {
+ plan.properties().output_partitioning().partition_count()
+ }
+
+ /// The file of each group of a scan (`group_file_counts` only counts
them).
+ fn group_files(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
+ file_groups(plan)
+ .iter()
+ .map(|group| group.files()[0].object_meta.location.to_string())
+ .collect()
+ }
+
+ #[test]
+ fn resolver_restricts_a_node_the_scheduler_does_not_know() {
+ let resolver = payload_resolver();
+ let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]);
+
+ let restricted =
+ restrict_plan_to_partitions(leaf, &[1, 3],
Some(&resolver)).unwrap();
+
+ assert_eq!(payloads_of(&restricted), ["p1", "p3"]);
+ assert_eq!(partition_count(&restricted), 2);
+ }
+
+ /// A node that cannot implement `PartitionSliceable`, such as a
+ /// `DataSourceExec` over a source of your own, is restricted by the
+ /// resolver building a replacement.
+ #[test]
+ fn resolver_restricts_a_node_without_partition_sliceable() {
+ let resolver: PartitionSliceableResolver =
+ Arc::new(|plan, children, partitions| {
+ assert!(children.is_empty());
+ plan.is::<PayloadExec>().then(|| {
+ Ok(Arc::new(
+
EmptyExec::new(plan.schema()).with_partitions(partitions.len()),
+ ) as Arc<dyn ExecutionPlan>)
+ })
+ });
+ let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]);
+
+ let restricted =
+ restrict_plan_to_partitions(leaf, &[1, 3],
Some(&resolver)).unwrap();
+
+ assert!(restricted.is::<EmptyExec>());
+ assert_eq!(partition_count(&restricted), 2);
+ }
+
+ #[test]
+ fn a_leaf_nobody_can_restrict_is_refused_unless_the_task_runs_all_of_it() {
+ let leaf = PayloadExec::leaf(&["p0", "p1"]);
+
+ let kept = restrict_plan_to_partitions(leaf.clone(), &[0, 1],
None).unwrap();
+ assert!(
+ Arc::ptr_eq(&kept, &leaf),
+ "an unrestricted leaf is left alone"
+ );
+
+ // A one-partition leaf is right for the task that owns it.
+ let single = PayloadExec::leaf(&["p0"]);
+ assert!(restrict_plan_to_partitions(single, &[0], None).is_ok());
+
+ for partitions in [&[0][..], &[1], &[1, 0], &[]] {
+ assert!(
+ restrict_plan_to_partitions(leaf.clone(), partitions,
None).is_err(),
+ "{partitions:?} is not all of the leaf, in order"
+ );
+ }
+
+ let error = restrict_plan_to_partitions(leaf, &[1], None)
+ .unwrap_err()
+ .to_string();
+ assert!(error.contains("PayloadExec"), "{error}");
Review Comment:
Nit: comparing the full error message instead of using `contains` would
catch wording regressions like the typo in `ensure_leaf_runs_whole`.
--
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]
