Fix: fetch is missing in EnforceSorting optimizer (two places) (#15822)

* Fix: fetch is missing in EnforceSort

* add ut test_parallelize_sort_preserves_fetch

* add ut: test_plan_with_order_preserving_variants_preserves_fetch

* update

* address comments

Author: xudong963

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -3440,3 +3440,38 @@ fn test_handles_multiple_orthogonal_sorts() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_parallelize_sort_preserves_fetch() -> Result<()> {
+    // Create a schema
+    let schema = create_test_schema3()?;
+    let parquet_exec = parquet_exec(&schema);
+    let coalesced = Arc::new(CoalescePartitionsExec::new(parquet_exec.clone()));
+    let top_coalesced = CoalescePartitionsExec::new(coalesced.clone())
+        .with_fetch(Some(10))
+        .unwrap();
+
+    let requirements = PlanWithCorrespondingCoalescePartitions::new(
+        top_coalesced.clone(),
+        true,
+        vec![PlanWithCorrespondingCoalescePartitions::new(
+            coalesced,
+            true,
+            vec![PlanWithCorrespondingCoalescePartitions::new(
+                parquet_exec,
+                false,
+                vec![],
+            )],
+        )],
+    );
+
+    let res = parallelize_sorts(requirements)?;
+
+    // Verify fetch was preserved
+    assert_eq!(
+        res.data.plan.fetch(),
+        Some(10),
+        "Fetch value was not preserved after transformation"
+    );
+    Ok(())
+}

datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -18,7 +18,8 @@
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    check_integrity, sort_preserving_merge_exec, stream_exec_ordered_with_projection,
+    check_integrity, create_test_schema3, sort_preserving_merge_exec,
+    stream_exec_ordered_with_projection,
 };
 
 use datafusion::prelude::SessionContext;
@@ -40,13 +41,14 @@ use datafusion_physical_plan::{
 };
 use datafusion::datasource::source::DataSourceExec;
 use datafusion_common::tree_node::{TransformedResult, TreeNode};
-use datafusion_common::Result;
+use datafusion_common::{assert_contains, Result};
 use datafusion_expr::{JoinType, Operator};
 use datafusion_physical_expr::expressions::{self, col, Column};
 use datafusion_physical_expr::PhysicalSortExpr;
-use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
+use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{plan_with_order_preserving_variants, replace_with_order_preserving_variants, OrderPreservationContext};
 use datafusion_common::config::ConfigOptions;
 
+use crate::physical_optimizer::enforce_sorting::parquet_exec_sorted;
 use object_store::memory::InMemory;
 use object_store::ObjectStore;
 use rstest::rstest;
@@ -1259,3 +1261,52 @@ fn memory_exec_sorted(
         ))
     })
 }
+
+#[test]
+fn test_plan_with_order_preserving_variants_preserves_fetch() -> Result<()> {
+    // Create a schema
+    let schema = create_test_schema3()?;
+    let parquet_sort_exprs = vec![crate::physical_optimizer::test_utils::sort_expr(
+        "a", &schema,
+    )];
+    let parquet_exec = parquet_exec_sorted(&schema, parquet_sort_exprs);
+    let coalesced = CoalescePartitionsExec::new(parquet_exec.clone())
+        .with_fetch(Some(10))
+        .unwrap();
+
+    // Test sort's fetch is greater than coalesce fetch, return error because it's not reasonable
+    let requirements = OrderPreservationContext::new(
+        coalesced.clone(),
+        false,
+        vec![OrderPreservationContext::new(
+            parquet_exec.clone(),
+            false,
+            vec![],
+        )],
+    );
+    let res = plan_with_order_preserving_variants(requirements, false, true, Some(15));
+    assert_contains!(res.unwrap_err().to_string(), "CoalescePartitionsExec fetch [10] should be greater than or equal to SortExec fetch [15]");
+
+    // Test sort is without fetch, expected to get the fetch value from the coalesced
+    let requirements = OrderPreservationContext::new(
+        coalesced.clone(),
+        false,
+        vec![OrderPreservationContext::new(
+            parquet_exec.clone(),
+            false,
+            vec![],
+        )],
+    );
+    let res = plan_with_order_preserving_variants(requirements, false, true, None)?;
+    assert_eq!(res.plan.fetch(), Some(10),);
+
+    // Test sort's fetch is less than coalesces fetch, expected to get the fetch value from the sort
+    let requirements = OrderPreservationContext::new(
+        coalesced,
+        false,
+        vec![OrderPreservationContext::new(parquet_exec, false, vec![])],
+    );
+    let res = plan_with_order_preserving_variants(requirements, false, true, Some(5))?;
+    assert_eq!(res.plan.fetch(), Some(5),);
+    Ok(())
+}

datafusion/physical-optimizer/src/enforce_sorting/mod.rs

@@ -400,6 +400,7 @@ pub fn parallelize_sorts(
             ),
         ))
     } else if is_coalesce_partitions(&requirements.plan) {
+        let fetch = requirements.plan.fetch();
         // There is an unnecessary `CoalescePartitionsExec` in the plan.
         // This will handle the recursive `CoalescePartitionsExec` plans.
         requirements = remove_bottleneck_in_subplan(requirements)?;
@@ -408,7 +409,10 @@ pub fn parallelize_sorts(
 
         Ok(Transformed::yes(
             PlanWithCorrespondingCoalescePartitions::new(
-                Arc::new(CoalescePartitionsExec::new(Arc::clone(&requirements.plan))),
+                // Safe to unwrap, because `CoalescePartitionsExec` has a fetch
+                CoalescePartitionsExec::new(Arc::clone(&requirements.plan))
+                    .with_fetch(fetch)
+                    .unwrap(),
                 false,
                 vec![requirements],
             ),

datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs

@@ -27,7 +27,7 @@ use crate::utils::{
 
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::Transformed;
-use datafusion_common::Result;
+use datafusion_common::{internal_err, Result};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::execution_plan::EmissionType;
@@ -93,7 +93,7 @@ pub fn update_order_preservation_ctx_children_data(opc: &mut OrderPreservationCo
 /// inside `sort_input` with their order-preserving variants. This will
 /// generate an alternative plan, which will be accepted or rejected later on
 /// depending on whether it helps us remove a `SortExec`.
-fn plan_with_order_preserving_variants(
+pub fn plan_with_order_preserving_variants(
     mut sort_input: OrderPreservationContext,
     // Flag indicating that it is desirable to replace `RepartitionExec`s with
     // `SortPreservingRepartitionExec`s:
@@ -138,6 +138,19 @@ fn plan_with_order_preserving_variants(
     } else if is_coalesce_partitions(&sort_input.plan) && is_spm_better {
         let child = &sort_input.children[0].plan;
         if let Some(ordering) = child.output_ordering() {
+            let mut fetch = fetch;
+            if let Some(coalesce_fetch) = sort_input.plan.fetch() {
+                if let Some(sort_fetch) = fetch {
+                    if coalesce_fetch < sort_fetch {
+                        return internal_err!(
+                                "CoalescePartitionsExec fetch [{:?}] should be greater than or equal to SortExec fetch [{:?}]", coalesce_fetch, sort_fetch
+                            );
+                    }
+                } else {
+                    // If the sort node does not have a fetch, we need to keep the coalesce node's fetch.
+                    fetch = Some(coalesce_fetch);
+                }
+            };
             // When the input of a `CoalescePartitionsExec` has an ordering,
             // replace it with a `SortPreservingMergeExec` if appropriate:
             let spm = SortPreservingMergeExec::new(ordering.clone(), Arc::clone(child))