This is an automated email from the ASF dual-hosted git repository.
xudong963 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new ab0da55912 Fix: fetch is lost in replace_order_preserving_variants
method during EnforceDistribution (#15808)
ab0da55912 is described below
commit ab0da5591247f87ab75489948285fe0af3b97449
Author: xudong.w <[email protected]>
AuthorDate: Wed Apr 23 09:36:34 2025 +0800
Fix: fetch is lost in replace_order_preserving_variants method during
EnforceDistribution (#15808)
---
.../physical_optimizer/enforce_distribution.rs | 45 ++++++++++++++++++++++
.../physical-optimizer/src/enforce_distribution.rs | 7 +++-
2 files changed, 50 insertions(+), 2 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
index 9898f6204e..5e182cb93b 100644
--- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
+++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
@@ -52,6 +52,7 @@ use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::execution_plan::ExecutionPlan;
use datafusion_physical_plan::expressions::col;
use datafusion_physical_plan::filter::FilterExec;
@@ -3471,3 +3472,47 @@ fn optimize_away_unnecessary_repartition2() ->
Result<()> {
Ok(())
}
+
+#[test]
+fn test_replace_order_preserving_variants_with_fetch() -> Result<()> {
+ // Create a base plan
+ let parquet_exec = parquet_exec();
+
+ let sort_expr = PhysicalSortExpr {
+ expr: Arc::new(Column::new("id", 0)),
+ options: SortOptions::default(),
+ };
+
+ let ordering = LexOrdering::new(vec![sort_expr]);
+
+ // Create a SortPreservingMergeExec with fetch=5
+ let spm_exec = Arc::new(
+ SortPreservingMergeExec::new(ordering,
parquet_exec.clone()).with_fetch(Some(5)),
+ );
+
+ // Create distribution context
+ let dist_context = DistributionContext::new(
+ spm_exec,
+ true,
+ vec![DistributionContext::new(parquet_exec, false, vec![])],
+ );
+
+ // Apply the function
+ let result = replace_order_preserving_variants(dist_context)?;
+
+ // Verify the plan was transformed to CoalescePartitionsExec
+ result
+ .plan
+ .as_any()
+ .downcast_ref::<CoalescePartitionsExec>()
+ .expect("Expected CoalescePartitionsExec");
+
+ // Verify fetch was preserved
+ assert_eq!(
+ result.plan.fetch(),
+ Some(5),
+ "Fetch value was not preserved after transformation"
+ );
+
+ Ok(())
+}
diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs
b/datafusion/physical-optimizer/src/enforce_distribution.rs
index b314b43c6a..523762401d 100644
--- a/datafusion/physical-optimizer/src/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/enforce_distribution.rs
@@ -1018,7 +1018,7 @@ fn remove_dist_changing_operators(
/// " RepartitionExec: partitioning=RoundRobinBatch(10),
input_partitions=2",
/// " DataSourceExec: file_groups={2 groups: \[\[x], \[y]]},
projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet",
/// ```
-fn replace_order_preserving_variants(
+pub fn replace_order_preserving_variants(
mut context: DistributionContext,
) -> Result<DistributionContext> {
context.children = context
@@ -1035,7 +1035,10 @@ fn replace_order_preserving_variants(
if is_sort_preserving_merge(&context.plan) {
let child_plan = Arc::clone(&context.children[0].plan);
- context.plan = Arc::new(CoalescePartitionsExec::new(child_plan));
+ // It's safe to unwrap because `CoalescePartitionsExec` supports
`fetch`.
+ context.plan = CoalescePartitionsExec::new(child_plan)
+ .with_fetch(context.plan.fetch())
+ .unwrap();
return Ok(context);
} else if let Some(repartition) =
context.plan.as_any().downcast_ref::<RepartitionExec>()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
