diff --git a/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs index 7d45c171b7e8..a7f0c88884e1 100644 --- a/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs @@ -178,7 +178,7 @@ fn get_updated_plan( let child = plan.children()[0].clone(); plan = Arc::new( RepartitionExec::try_new(child, plan.output_partitioning())? - .with_preserve_order(), + .with_preserve_order(true), ) as _ } // When the input of a `CoalescePartitionsExec` has an ordering, replace it diff --git a/datafusion/core/src/physical_plan/repartition/mod.rs b/datafusion/core/src/physical_plan/repartition/mod.rs index c47c9926819b..99b72a1b4064 100644 --- a/datafusion/core/src/physical_plan/repartition/mod.rs +++ b/datafusion/core/src/physical_plan/repartition/mod.rs @@ -419,11 +419,9 @@ impl ExecutionPlan for RepartitionExec { self: Arc, children: Vec>, ) -> Result> { - let mut repartition = - RepartitionExec::try_new(children[0].clone(), self.partitioning.clone())?; - if self.preserve_order { - repartition = repartition.with_preserve_order(); - } + let repartition = + RepartitionExec::try_new(children[0].clone(), self.partitioning.clone())? + .with_preserve_order(self.preserve_order); Ok(Arc::new(repartition)) } @@ -625,8 +623,8 @@ impl RepartitionExec { } /// Set Order preserving flag - pub fn with_preserve_order(mut self) -> Self { - self.preserve_order = true; + pub fn with_preserve_order(mut self, preserve_order: bool) -> Self { + self.preserve_order = preserve_order; self } diff --git a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs index d0ce58f6e617..6b3a633f3c16 100644 --- a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs @@ -140,7 +140,7 @@ mod sp_repartition_fuzz_tests { Arc::new( RepartitionExec::try_new(input, Partitioning::RoundRobinBatch(2)) .unwrap() - .with_preserve_order(), + .with_preserve_order(true), ) } @@ -159,7 +159,7 @@ mod sp_repartition_fuzz_tests { Arc::new( RepartitionExec::try_new(input, Partitioning::Hash(hash_expr, 2)) .unwrap() - .with_preserve_order(), + .with_preserve_order(true), ) }