diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs index b704a70002d81..a2ede7ec3de52 100644 --- a/datafusion/common/src/stats.rs +++ b/datafusion/common/src/stats.rs @@ -545,6 +545,10 @@ impl Statistics { skip: usize, n_partitions: usize, ) -> Result { + 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,18 +602,18 @@ 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 = 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 @@ -617,11 +621,11 @@ impl Statistics { .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 diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 80e0a3f23e736..f50170b6fe192 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/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::() + .map(|projection| projection.input()) + .unwrap_or(&optimized_join); + let swapped_join = optimized_join + .downcast_ref::() + .expect("optimized plan should contain a hash join"); + + let left_spm = swapped_join + .left() + .downcast_ref::() + .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(); diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 77a7d8f8f2e11..896ff9828030d 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/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> { - 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. ///