Cap SortPreservingMerge statistics by fetch

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,
     execution_plan::{Boundedness, EmissionType},
@@ -265,6 +267,86 @@ 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));
+    assert_eq!(
+        swapped_join
+            .left()
+            .statistics_with_args(&StatisticsArgs::new())
+            .unwrap()
+            .num_rows,
+        Precision::Inexact(1)
+    );
+    let left_byte_size = swapped_join
+        .left()
+        .statistics_with_args(&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!(
+        swapped_join
+            .right()
+            .statistics_with_args(&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

@@ -382,7 +382,9 @@ impl ExecutionPlan for SortPreservingMergeExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        args.compute_child_statistics(&self.input, None)
+        let stats =
+            Arc::unwrap_or_clone(args.compute_child_statistics(&self.input, None)?);
+        Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
     }
 
     fn supports_limit_pushdown(&self) -> bool {
@@ -434,7 +436,9 @@ mod tests {
     use crate::sorts::sort::SortExec;
     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};
 
@@ -444,8 +448,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;
@@ -510,6 +515,45 @@ 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 = spm.statistics_with_args(&StatisticsArgs::new())?;
+
+        assert_eq!(statistics.num_rows, Precision::Exact(1));
+        assert_eq!(statistics.total_byte_size, Precision::Inexact(8));
+        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 = spm.statistics_with_args(&StatisticsArgs::new())?;
+
+        assert_eq!(*statistics, input_stats);
+        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.
     ///