Cap SortPreservingMerge statistics by fetch (apache#23359)

## Which issue does this PR close?

- Closes none.

## Rationale for this change

`SortPreservingMergeExec` supports `fetch`, but its statistics currently
pass through child statistics unchanged. This can make cost-based
optimizers overestimate top-k merge outputs. For example, a
`SortPreservingMergeExec(fetch=1)` side can still look as large as its
input to `JoinSelection`, causing hash join build/probe choices to miss
the small top-k side.

Other fetch-aware operators such as `SortExec`,
`CoalescePartitionsExec`, and limit execs already cap their statistics
by fetch. This PR aligns `SortPreservingMergeExec` with that behavior.

## What changes are included in this PR?

- Cap `SortPreservingMergeExec` statistics with
`Statistics::with_fetch(self.fetch, 0, 1)`.
- Preserve no-op statistics for `fetch = None` and `skip = 0`.
- Avoid scaling byte-size estimates by `0.0` when input row count is
unknown; use absent byte-size stats instead.
- Add regression tests for SPM fetch statistics and for `JoinSelection`
swapping an SPM(fetch=1) side to the hash join build side.

## Are these changes tested?

Yes.

- `cargo test -p datafusion-common test_with_fetch`
- `cargo test -p datafusion-physical-plan sort_preserving_merge`
- `cargo test -p datafusion join_selection --test core_integration`
- `cargo check -p datafusion --test core_integration`
- `cargo fmt --all -- --check`
- `cargo clippy --all-targets --all-features -- -D warnings`

## Are there any user-facing changes?

No user-facing API changes. This improves optimizer statistics for
fetch-limited sort-preserving merge plans.

Author: discord9

datafusion/common/src/stats.rs

@@ -545,6 +545,10 @@ impl Statistics {
         skip: usize,
         n_partitions: usize,
     ) -> Result<Self> {
+        if fetch.is_none() && skip == 0 {
+            return Ok(self);
+        }
+
         let fetch_val = fetch.unwrap_or(usize::MAX);
 
         // Get the ratio of rows after / rows before on a per-partition basis
@@ -598,30 +602,30 @@ impl Statistics {
                 ..
             } => check_num_rows(fetch.and_then(|v| v.checked_mul(n_partitions)), false),
         };
-        let ratio: f64 = match (num_rows_before, self.num_rows) {
+        let ratio: Option<f64> = match (num_rows_before, self.num_rows) {
             (
                 Precision::Exact(nr_before) | Precision::Inexact(nr_before),
                 Precision::Exact(nr_after) | Precision::Inexact(nr_after),
             ) => {
                 if nr_before == 0 {
-                    0.0
+                    Some(0.0)
                 } else {
-                    nr_after as f64 / nr_before as f64
+                    Some(nr_after as f64 / nr_before as f64)
                 }
             }
-            _ => 0.0,
+            _ => None,
         };
         self.column_statistics = self
             .column_statistics
             .into_iter()
             .map(|cs| {
                 let mut cs = cs.to_inexact();
                 // Scale byte_size by the row ratio
-                cs.byte_size = match cs.byte_size {
-                    Precision::Exact(n) | Precision::Inexact(n) => {
+                cs.byte_size = match (cs.byte_size, ratio) {
+                    (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => {
                         Precision::Inexact((n as f64 * ratio) as usize)
                     }
-                    Precision::Absent => Precision::Absent,
+                    _ => Precision::Absent,
                 };
                 // NDV can never exceed the number of rows
                 if let Some(&rows) = self.num_rows.get_value() {
@@ -643,11 +647,11 @@ impl Statistics {
             Some(sum) => Precision::Inexact(sum),
             None => {
                 // Fall back to scaling original total_byte_size if not all columns have byte_size
-                match &self.total_byte_size {
-                    Precision::Exact(n) | Precision::Inexact(n) => {
+                match (&self.total_byte_size, ratio) {
+                    (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => {
                         Precision::Inexact((*n as f64 * ratio) as usize)
                     }
-                    Precision::Absent => Precision::Absent,
+                    _ => Precision::Absent,
                 }
             }
         };
@@ -2376,6 +2380,38 @@ mod tests {
         assert_eq!(result.total_byte_size, Precision::Exact(800));
     }
 
+    #[test]
+    fn test_with_fetch_no_limit_preserves_absent_num_rows() {
+        let original_stats = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Exact(800),
+            column_statistics: vec![col_stats_i64(10)],
+        };
+
+        let result = original_stats.clone().with_fetch(None, 0, 1).unwrap();
+
+        assert_eq!(result, original_stats);
+    }
+
+    #[test]
+    fn test_with_fetch_absent_num_rows_does_not_zero_byte_size() {
+        let original_stats = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Exact(800),
+            column_statistics: vec![col_stats_i64(10)],
+        };
+
+        let result = original_stats.with_fetch(Some(1), 0, 1).unwrap();
+
+        assert_eq!(result.num_rows, Precision::Inexact(1));
+        assert_eq!(result.total_byte_size, Precision::Absent);
+        assert_eq!(result.column_statistics[0].byte_size, Precision::Absent);
+        assert_eq!(
+            result.column_statistics[0].distinct_count,
+            Precision::Inexact(1)
+        );
+    }
+
     #[test]
     fn test_with_fetch_with_skip() {
         // Test with both skip and fetch

datafusion/core/tests/physical_optimizer/join_selection.rs

@@ -35,6 +35,7 @@ use datafusion_physical_expr::expressions::col;
 use datafusion_physical_expr::expressions::{BinaryExpr, Column, NegativeExpr};
 use datafusion_physical_expr::intervals::utils::check_support;
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
+use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::join_selection::JoinSelection;
 use datafusion_physical_plan::ExecutionPlanProperties;
@@ -43,6 +44,7 @@ use datafusion_physical_plan::joins::utils::ColumnIndex;
 use datafusion_physical_plan::joins::utils::JoinFilter;
 use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode};
 use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
     StatisticsContext,
@@ -264,6 +266,84 @@ async fn test_join_with_swap() {
     );
 }
 
+#[tokio::test]
+async fn test_join_with_swap_to_sort_preserving_merge_fetch_side() {
+    let (big, _) = create_big_and_small();
+    let top1_input = Arc::new(StatisticsExec::new(
+        big_statistics(),
+        Schema::new(vec![Field::new("top_col", DataType::Int32, false)]),
+    ));
+    let top1 = Arc::new(
+        SortPreservingMergeExec::new(
+            [PhysicalSortExpr::new_default(Arc::new(Column::new(
+                "top_col", 0,
+            )))]
+            .into(),
+            top1_input,
+        )
+        .with_fetch(Some(1)),
+    );
+
+    let join = Arc::new(
+        HashJoinExec::try_new(
+            Arc::clone(&big),
+            top1,
+            vec![(
+                Arc::new(Column::new_with_schema("big_col", &big.schema()).unwrap()),
+                Arc::new(Column::new("top_col", 0)),
+            )],
+            None,
+            &JoinType::Inner,
+            None,
+            PartitionMode::Partitioned,
+            NullEquality::NullEqualsNothing,
+            false,
+        )
+        .unwrap(),
+    );
+
+    let optimized_join = JoinSelection::new()
+        .optimize(join, &ConfigOptions::new())
+        .unwrap();
+    let optimized_join = optimized_join
+        .downcast_ref::<ProjectionExec>()
+        .map(|projection| projection.input())
+        .unwrap_or(&optimized_join);
+    let swapped_join = optimized_join
+        .downcast_ref::<HashJoinExec>()
+        .expect("optimized plan should contain a hash join");
+
+    let left_spm = swapped_join
+        .left()
+        .downcast_ref::<SortPreservingMergeExec>()
+        .expect("SPM fetch side should become the left/build input");
+    assert_eq!(left_spm.fetch(), Some(1));
+    let statistics_context = StatisticsContext::new();
+    assert_eq!(
+        statistics_context
+            .compute(swapped_join.left().as_ref(), &StatisticsArgs::new())
+            .unwrap()
+            .num_rows,
+        Precision::Inexact(1)
+    );
+    let left_byte_size = statistics_context
+        .compute(swapped_join.left().as_ref(), &StatisticsArgs::new())
+        .unwrap()
+        .total_byte_size;
+    let right_byte_size = big_statistics().total_byte_size;
+    assert!(
+        left_byte_size.get_value() < right_byte_size.get_value(),
+        "SPM fetch side should be estimated smaller than the big side"
+    );
+    assert_eq!(
+        statistics_context
+            .compute(swapped_join.right().as_ref(), &StatisticsArgs::new())
+            .unwrap()
+            .num_rows,
+        big_statistics().num_rows
+    );
+}
+
 #[tokio::test]
 async fn test_left_join_no_swap() {
     let (big, small) = create_big_and_small();

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -36,7 +36,7 @@ use datafusion_execution::TaskContext;
 use datafusion_execution::memory_pool::MemoryConsumer;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements};
 
-use crate::execution_plan::{EvaluationType, SchedulingType};
+use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
 use log::{debug, trace};
 
 /// Sort preserving merge execution plan
@@ -396,7 +396,16 @@ impl ExecutionPlan for SortPreservingMergeExec {
         input_stats: &[Arc<Statistics>],
         _args: &StatisticsArgs,
     ) -> Result<Arc<Statistics>> {
-        Ok(Arc::clone(&input_stats[0]))
+        let stats = input_stats[0].as_ref().clone();
+        Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
+    }
+
+    fn cardinality_effect(&self) -> CardinalityEffect {
+        if self.fetch.is_none() {
+            CardinalityEffect::Equal
+        } else {
+            CardinalityEffect::LowerEqual
+        }
     }
 
     fn supports_limit_pushdown(&self) -> bool {
@@ -446,9 +455,12 @@ mod tests {
     use crate::metrics::{MetricValue, Timestamp};
     use crate::repartition::RepartitionExec;
     use crate::sorts::sort::SortExec;
+    use crate::statistics::StatisticsContext;
     use crate::stream::RecordBatchReceiverStream;
     use crate::test::TestMemoryExec;
-    use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero};
+    use crate::test::exec::{
+        BlockingExec, StatisticsExec, assert_strong_count_converges_to_zero,
+    };
     use crate::test::{self, assert_is_pending, make_partition};
     use crate::{collect, common};
 
@@ -458,8 +470,9 @@ mod tests {
     };
     use arrow::compute::SortOptions;
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+    use datafusion_common::stats::Precision;
     use datafusion_common::test_util::batches_to_string;
-    use datafusion_common::{assert_batches_eq, exec_err};
+    use datafusion_common::{ColumnStatistics, assert_batches_eq, exec_err};
     use datafusion_common_runtime::SpawnedTask;
     use datafusion_execution::RecordBatchStream;
     use datafusion_execution::config::SessionConfig;
@@ -524,6 +537,52 @@ mod tests {
         Ok(Arc::new(spm))
     }
 
+    #[test]
+    fn test_fetch_caps_statistics() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+        let input = Arc::new(StatisticsExec::new(
+            Statistics {
+                num_rows: Precision::Exact(1_000),
+                total_byte_size: Precision::Exact(8_000),
+                column_statistics: vec![ColumnStatistics::new_unknown()],
+            },
+            schema.clone(),
+        ));
+        let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into();
+
+        let spm = SortPreservingMergeExec::new(sort, input).with_fetch(Some(1));
+        let statistics =
+            StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?;
+
+        assert_eq!(statistics.num_rows, Precision::Exact(1));
+        assert_eq!(statistics.total_byte_size, Precision::Inexact(8));
+        assert!(matches!(
+            spm.cardinality_effect(),
+            CardinalityEffect::LowerEqual
+        ));
+        Ok(())
+    }
+
+    #[test]
+    fn test_no_fetch_preserves_statistics() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+        let input_stats = Statistics {
+            num_rows: Precision::Absent,
+            total_byte_size: Precision::Exact(8_000),
+            column_statistics: vec![ColumnStatistics::new_unknown()],
+        };
+        let input = Arc::new(StatisticsExec::new(input_stats.clone(), schema.clone()));
+        let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into();
+
+        let spm = SortPreservingMergeExec::new(sort, input);
+        let statistics =
+            StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?;
+
+        assert_eq!(*statistics, input_stats);
+        assert!(matches!(spm.cardinality_effect(), CardinalityEffect::Equal));
+        Ok(())
+    }
+
     /// This test verifies that memory usage stays within limits when the tie breaker is enabled.
     /// Any errors here could indicate unintended changes in tie breaker logic.
     ///