gene-bordegaray commented on code in PR #23184:
URL: https://github.com/apache/datafusion/pull/23184#discussion_r3534995660


##########
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:
   Hahaha yes I'm trying my hardest to make it more readable jnremently 😭. But 
please if it's harder now let me know 



-- 
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]

Reply via email to