This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] 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 5b65e70cbe Cap SortPreservingMerge statistics by fetch (#23359)
5b65e70cbe is described below
commit 5b65e70cbeb9e7a3a069703ab9f49b44d31f85fc
Author: discord9 <[email protected]>
AuthorDate: Tue Jul 21 05:18:48 2026 +0800
Cap SortPreservingMerge statistics by fetch (#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.
---
datafusion/common/src/stats.rs | 56 ++++++++++++---
.../tests/physical_optimizer/join_selection.rs | 80 ++++++++++++++++++++++
.../src/sorts/sort_preserving_merge.rs | 67 ++++++++++++++++--
3 files changed, 189 insertions(+), 14 deletions(-)
diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs
index b704a70002..a2ede7ec3d 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<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,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<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
@@ -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 2db9f18f31..cca54909a1 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,
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();
diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
index 693f9789b2..a25f3a2862 100644
--- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
+++ b/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.
///
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
