gabotechs commented on code in PR #23184:
URL: https://github.com/apache/datafusion/pull/23184#discussion_r3534575830
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1309,75 +1389,91 @@ pub fn ensure_distribution(
}
};
- let streaming_benefit = if child.data {
- preserving_order_enables_streaming(&plan, &child.plan)?
- } else {
- false
- };
+ Ok((child, required_input_ordering, maintains, requirement))
+ },
+ )
+ .collect::<Result<Vec<_>>>()?;
- // There is an ordering requirement of the operator:
- if let Some(required_input_ordering) = required_input_ordering {
- // Either:
- // - Ordering requirement cannot be satisfied by preserving
ordering through repartitions, or
- // - using order preserving variant is not desirable.
- let sort_req = required_input_ordering.into_single();
- let ordering_satisfied = child
- .plan
- .equivalence_properties()
- .ordering_satisfy_requirement(sort_req.clone())?;
-
- if (!ordering_satisfied ||
!order_preserving_variants_desirable)
- && !streaming_benefit
- && child.data
- {
- child = replace_order_preserving_variants(child)?;
- // If ordering requirements were satisfied before
repartitioning,
- // make sure ordering requirements are still satisfied
after.
- if ordering_satisfied {
- // Make sure to satisfy ordering requirement:
- child = add_sort_above_with_check(
- child,
- sort_req,
- plan.downcast_ref::<OutputRequirementExec>()
- .map(|output| output.fetch())
- .unwrap_or(None),
- )?;
- }
- }
- // Stop tracking distribution changing operators
- child.data = false;
- } else {
+ enforce_distribution_relationships(
+ plan.name(),
+ &input_distributions,
+ &mut children,
+ target_partitions,
+ )?;
+
+ let children = children
+ .into_iter()
+ .map(
+ |(mut child, required_input_ordering, maintains, requirement)| {
let streaming_benefit = if child.data {
preserving_order_enables_streaming(&plan, &child.plan)?
} else {
false
};
- // no ordering requirement
- match requirement {
- // Operator requires specific distribution.
- Distribution::SinglePartition
- | Distribution::HashPartitioned(_)
- | Distribution::KeyPartitioned(_) => {
- // If the parent doesn't maintain input order,
preserving
- // ordering is pointless. However, if it does maintain
- // input order, we keep order-preserving variants so
- // ordering can flow through to ancestors that need it.
- if !maintains && !streaming_benefit {
- child = replace_order_preserving_variants(child)?;
+
+ // There is an ordering requirement of the operator:
+ if let Some(required_input_ordering) = required_input_ordering
{
+ // Either:
+ // - Ordering requirement cannot be satisfied by
preserving ordering through repartitions, or
+ // - using order preserving variant is not desirable.
+ let sort_req = required_input_ordering.into_single();
+ let ordering_satisfied = child
+ .plan
+ .equivalence_properties()
+ .ordering_satisfy_requirement(sort_req.clone())?;
+
+ if (!ordering_satisfied ||
!order_preserving_variants_desirable)
+ && !streaming_benefit
+ && child.data
+ {
+ child = replace_order_preserving_variants(child)?;
+ // If ordering requirements were satisfied before
repartitioning,
+ // make sure ordering requirements are still satisfied
after.
+ if ordering_satisfied {
+ // Make sure to satisfy ordering requirement:
Review Comment:
The amount of `if` statement branching happening here is most likely going
to take a huge hit in maintainability.
This was already a difficult optimizer rule to understand, but I fear now
it's worst.
I'll likely need to invest some time in understanding this before I can
contribute any meaningful suggestion. In the mean time, if you see
opportunities for simplyfing this code, that's more than welcome.
##########
datafusion/physical-plan/src/execution_plan.rs:
##########
@@ -164,6 +167,12 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send +
Sync {
check_default_invariants(self, check)
}
+ /// Specifies data distribution requirements for all children, including
+ /// non-default satisfaction policies or cross-child relationships.
+ fn required_input_distributions(&self) -> RequiredInputDistributions {
+ RequiredInputDistributions::new(self.required_input_distribution())
+ }
+
/// Specifies the data distribution requirements for all the
Review Comment:
Is there any chance we can not maintain these two methods?
- `required_input_distributions`
- `required_input_distribution`
Having the two at the same time does not seem like a nice public API design
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1041,6 +1045,92 @@ fn get_repartition_requirement_status(
.collect())
}
+/// Enforce cross-child distribution relationships after each child has already
+/// satisfied its own distribution requirement.
+///
+/// If the candidate children are still not co-partitioned, this adds the
+/// conservative key-based repartitions needed to make partition indexes line
up.
+#[expect(
+ deprecated,
+ reason = "HashPartitioned is accepted during the KeyPartitioned migration"
+)]
+fn enforce_distribution_relationships(
+ plan_name: &str,
+ input_distributions: &RequiredInputDistributions,
+ children: &mut [(
+ DistributionContext,
+ Option<OrderingRequirements>,
+ bool,
+ Distribution,
+ )],
+ target_partitions: usize,
Review Comment:
This type is pretty messed up, I'm surprised that clippy allows this. It
would be nice to avoid these wide unnamed tuple arguments, as reading the
signature is impossible to figure out what's any of that.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -1796,6 +1797,19 @@ impl ExecutionPlan for AggregateExec {
}
}
+ fn required_input_distributions(&self) -> RequiredInputDistributions {
+ let requirements =
+
RequiredInputDistributions::new(self.required_input_distribution());
+ match &self.mode {
+ AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned
+ if !self.group_by.has_grouping_set() =>
+ {
+ requirements.with_range_key_partitioning(0)
Review Comment:
Without looking at the specific implementation, this reads:
"The required input distribution must be range key partitioned"
Which I assume is wrong right?
##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -882,9 +883,31 @@ impl HashJoinExec {
return false;
}
+ if self.mode == PartitionMode::Partitioned
+ && !self.has_partitioned_dynamic_filter_routing()
+ {
+ return false;
+ }
+
true
}
+ fn has_partitioned_dynamic_filter_routing(&self) -> bool {
+ match (
+ self.left.output_partitioning(),
+ self.right.output_partitioning(),
+ ) {
+ (
+ Partitioning::Hash(_, left_partition_count),
+ Partitioning::Hash(_, right_partition_count),
+ ) => left_partition_count == right_partition_count,
+ (left_partitioning, right_partitioning) => {
Review Comment:
I imagine that dynamic filters should work at some point with co-partitioned
joins on a range key right? probably another issue worth adding to your batch?
--
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]
