kosiew commented on code in PR #26716:
URL: https://github.com/apache/datafusion/pull/26716#discussion_r4550812411
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -2030,4 +2109,190 @@ fn update_children(mut dist_context:
DistributionContext) -> Result<Distribution
Ok(dist_context)
}
-// See tests in datafusion/core/tests/physical_optimizer
+// Plan-level tests live in datafusion/core/tests/physical_optimizer.
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ const TARGET: usize = 8;
+ const THRESHOLD: usize = 4;
+
+ fn base_config() -> ChildPolicyConfig {
+ ChildPolicyConfig {
+ target_partitions: TARGET,
+ subset_satisfaction_threshold: THRESHOLD,
+ preserve_file_partitions: 0,
+ enable_round_robin: true,
+ is_partitioned_join: false,
+ }
+ }
+
+ fn status(requirement: Distribution) -> RepartitionRequirementStatus {
+ RepartitionRequirementStatus {
+ requirement,
+ roundrobin_beneficial: true,
+ roundrobin_beneficial_stats: true,
+ hash_necessary: true,
+ }
+ }
+
+ fn hash_on(names: &[&str]) -> Distribution {
+ Distribution::KeyPartitioned(
+ names
+ .iter()
+ .enumerate()
+ .map(|(i, name)| Arc::new(Column::new(name, i)) as _)
+ .collect(),
+ )
+ }
+
+ fn policy(
+ config: &ChildPolicyConfig,
+ current_partitions: usize,
+ requirement: Distribution,
+ ) -> ChildPolicy {
+ ChildPolicy::new(config, current_partitions, &status(requirement), |_|
false)
+ }
+
+ #[test]
+ fn subset_satisfaction_needs_threshold() {
+ let config = base_config();
+ assert!(
+ !policy(&config, THRESHOLD - 1, hash_on(&["a"]))
+ .allow_subset_satisfy_partitioning
+ );
+ assert!(
+ policy(&config, THRESHOLD,
hash_on(&["a"])).allow_subset_satisfy_partitioning
+ );
+ }
+
+ #[test]
+ fn partitioned_join_never_allows_subset_satisfaction() {
+ let config = ChildPolicyConfig {
+ is_partitioned_join: true,
+ ..base_config()
+ };
+ assert!(
+ !policy(&config, TARGET,
hash_on(&["a"])).allow_subset_satisfy_partitioning
+ );
+ }
+
Review Comment:
Optional coverage suggestion: could we test the grouping-ID policy with both
`KeyPartitioned` and deprecated `HashPartitioned`? `ChildPolicy::new` supports
both, but `hash_on` only constructs `KeyPartitioned`. A small table-driven test
would help protect the legacy case during migration.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1647,74 +1786,24 @@ pub fn ensure_distribution_with_stats(
repartition_status_flags.into_iter()
)
.map(
- |(
- child_idx,
- mut child,
- required_input_ordering,
- maintains,
- RepartitionRequirementStatus {
+ |(child_idx, mut child, required_input_ordering, maintains, status)| {
+ let policy = ChildPolicy::new(
Review Comment:
Optional documentation suggestion: could we clarify that the policy
intentionally captures the partition count before source repartitioning, while
hash eligibility uses the updated count afterward? A short comment would make
this snapshot boundary easier to understand without changing the existing
behavior.
--
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]
