You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ad...@apache.org on 2025/09/09 16:01:19 UTC
(datafusion) branch main updated: better preserve statistics when applying limits (#17381)
This is an automated email from the ASF dual-hosted git repository.
adriangb 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 da1e7ae2e5 better preserve statistics when applying limits (#17381)
da1e7ae2e5 is described below
commit da1e7ae2e59291499100753acd749d3a23cc88b5
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Tue Sep 9 11:01:15 2025 -0500
better preserve statistics when applying limits (#17381)
---
datafusion/common/src/stats.rs | 298 ++++++++++++++++++++-
.../core/tests/custom_sources_cases/statistics.rs | 10 +-
.../physical_optimizer/partition_statistics.rs | 19 +-
datafusion/physical-plan/src/coalesce_batches.rs | 9 +-
.../physical-plan/src/coalesce_partitions.rs | 2 +-
datafusion/physical-plan/src/limit.rs | 18 +-
datafusion/physical-plan/src/sorts/sort.rs | 19 +-
.../sqllogictest/test_files/parquet_statistics.slt | 2 +-
8 files changed, 330 insertions(+), 47 deletions(-)
diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs
index e461274abd..f54fa55aec 100644
--- a/datafusion/common/src/stats.rs
+++ b/datafusion/common/src/stats.rs
@@ -22,7 +22,7 @@ use std::fmt::{self, Debug, Display};
use crate::{Result, ScalarValue};
use crate::error::_plan_err;
-use arrow::datatypes::{DataType, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Schema};
/// Represents a value with a degree of certainty. `Precision` is used to
/// propagate information the precision of statistical values.
@@ -391,13 +391,15 @@ impl Statistics {
/// parameter to compute global statistics in a multi-partition setting.
pub fn with_fetch(
mut self,
- schema: SchemaRef,
fetch: Option<usize>,
skip: usize,
n_partitions: usize,
) -> Result<Self> {
let fetch_val = fetch.unwrap_or(usize::MAX);
+ // Get the ratio of rows after / rows before on a per-partition basis
+ let num_rows_before = self.num_rows;
+
self.num_rows = match self {
Statistics {
num_rows: Precision::Exact(nr),
@@ -431,8 +433,7 @@ impl Statistics {
// At this point we know that we were given a `fetch` value
// as the `None` case would go into the branch above. Since
// the input has more rows than `fetch + skip`, the number
- // of rows will be the `fetch`, but we won't be able to
- // predict the other statistics.
+ // of rows will be the `fetch`, other statistics will have to be downgraded to inexact.
check_num_rows(
fetch_val.checked_mul(n_partitions),
// We know that we have an estimate for the number of rows:
@@ -445,8 +446,32 @@ impl Statistics {
..
} => check_num_rows(fetch.and_then(|v| v.checked_mul(n_partitions)), false),
};
- self.column_statistics = Statistics::unknown_column(&schema);
- self.total_byte_size = Precision::Absent;
+ let ratio: 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
+ } else {
+ nr_after as f64 / nr_before as f64
+ }
+ }
+ _ => 0.0,
+ };
+ self.column_statistics = self
+ .column_statistics
+ .into_iter()
+ .map(ColumnStatistics::to_inexact)
+ .collect();
+ // Adjust the total_byte_size for the ratio of rows before and after, also marking it as inexact
+ self.total_byte_size = match &self.total_byte_size {
+ Precision::Exact(n) | Precision::Inexact(n) => {
+ let adjusted = (*n as f64 * ratio) as usize;
+ Precision::Inexact(adjusted)
+ }
+ Precision::Absent => Precision::Absent,
+ };
Ok(self)
}
@@ -1199,4 +1224,265 @@ mod tests {
// Distinct count should be Absent after merge
assert_eq!(col_stats.distinct_count, Precision::Absent);
}
+
+ #[test]
+ fn test_with_fetch_basic_preservation() {
+ // Test that column statistics and byte size are preserved (as inexact) when applying fetch
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(1000),
+ total_byte_size: Precision::Exact(8000),
+ column_statistics: vec![
+ ColumnStatistics {
+ null_count: Precision::Exact(10),
+ max_value: Precision::Exact(ScalarValue::Int32(Some(100))),
+ min_value: Precision::Exact(ScalarValue::Int32(Some(0))),
+ sum_value: Precision::Exact(ScalarValue::Int32(Some(5050))),
+ distinct_count: Precision::Exact(50),
+ },
+ ColumnStatistics {
+ null_count: Precision::Exact(20),
+ max_value: Precision::Exact(ScalarValue::Int64(Some(200))),
+ min_value: Precision::Exact(ScalarValue::Int64(Some(10))),
+ sum_value: Precision::Exact(ScalarValue::Int64(Some(10100))),
+ distinct_count: Precision::Exact(75),
+ },
+ ],
+ };
+
+ // Apply fetch of 100 rows (10% of original)
+ let result = original_stats.clone().with_fetch(Some(100), 0, 1).unwrap();
+
+ // Check num_rows
+ assert_eq!(result.num_rows, Precision::Exact(100));
+
+ // Check total_byte_size is scaled proportionally and marked as inexact
+ // 100/1000 = 0.1, so 8000 * 0.1 = 800
+ assert_eq!(result.total_byte_size, Precision::Inexact(800));
+
+ // Check column statistics are preserved but marked as inexact
+ assert_eq!(result.column_statistics.len(), 2);
+
+ // First column
+ assert_eq!(
+ result.column_statistics[0].null_count,
+ Precision::Inexact(10)
+ );
+ assert_eq!(
+ result.column_statistics[0].max_value,
+ Precision::Inexact(ScalarValue::Int32(Some(100)))
+ );
+ assert_eq!(
+ result.column_statistics[0].min_value,
+ Precision::Inexact(ScalarValue::Int32(Some(0)))
+ );
+ assert_eq!(
+ result.column_statistics[0].sum_value,
+ Precision::Inexact(ScalarValue::Int32(Some(5050)))
+ );
+ assert_eq!(
+ result.column_statistics[0].distinct_count,
+ Precision::Inexact(50)
+ );
+
+ // Second column
+ assert_eq!(
+ result.column_statistics[1].null_count,
+ Precision::Inexact(20)
+ );
+ assert_eq!(
+ result.column_statistics[1].max_value,
+ Precision::Inexact(ScalarValue::Int64(Some(200)))
+ );
+ assert_eq!(
+ result.column_statistics[1].min_value,
+ Precision::Inexact(ScalarValue::Int64(Some(10)))
+ );
+ assert_eq!(
+ result.column_statistics[1].sum_value,
+ Precision::Inexact(ScalarValue::Int64(Some(10100)))
+ );
+ assert_eq!(
+ result.column_statistics[1].distinct_count,
+ Precision::Inexact(75)
+ );
+ }
+
+ #[test]
+ fn test_with_fetch_inexact_input() {
+ // Test that inexact input statistics remain inexact
+ let original_stats = Statistics {
+ num_rows: Precision::Inexact(1000),
+ total_byte_size: Precision::Inexact(8000),
+ column_statistics: vec![ColumnStatistics {
+ null_count: Precision::Inexact(10),
+ max_value: Precision::Inexact(ScalarValue::Int32(Some(100))),
+ min_value: Precision::Inexact(ScalarValue::Int32(Some(0))),
+ sum_value: Precision::Inexact(ScalarValue::Int32(Some(5050))),
+ distinct_count: Precision::Inexact(50),
+ }],
+ };
+
+ let result = original_stats.clone().with_fetch(Some(500), 0, 1).unwrap();
+
+ // Check num_rows is inexact
+ assert_eq!(result.num_rows, Precision::Inexact(500));
+
+ // Check total_byte_size is scaled and inexact
+ // 500/1000 = 0.5, so 8000 * 0.5 = 4000
+ assert_eq!(result.total_byte_size, Precision::Inexact(4000));
+
+ // Column stats remain inexact
+ assert_eq!(
+ result.column_statistics[0].null_count,
+ Precision::Inexact(10)
+ );
+ }
+
+ #[test]
+ fn test_with_fetch_skip_all_rows() {
+ // Test when skip >= num_rows (all rows are skipped)
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(100),
+ total_byte_size: Precision::Exact(800),
+ column_statistics: vec![col_stats_i64(10)],
+ };
+
+ let result = original_stats.clone().with_fetch(Some(50), 100, 1).unwrap();
+
+ assert_eq!(result.num_rows, Precision::Exact(0));
+ // When ratio is 0/100 = 0, byte size should be 0
+ assert_eq!(result.total_byte_size, Precision::Inexact(0));
+ }
+
+ #[test]
+ fn test_with_fetch_no_limit() {
+ // Test when fetch is None and skip is 0 (no limit applied)
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(100),
+ total_byte_size: Precision::Exact(800),
+ column_statistics: vec![col_stats_i64(10)],
+ };
+
+ let result = original_stats.clone().with_fetch(None, 0, 1).unwrap();
+
+ // Stats should be unchanged when no fetch and no skip
+ assert_eq!(result.num_rows, Precision::Exact(100));
+ assert_eq!(result.total_byte_size, Precision::Exact(800));
+ }
+
+ #[test]
+ fn test_with_fetch_with_skip() {
+ // Test with both skip and fetch
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(1000),
+ total_byte_size: Precision::Exact(8000),
+ column_statistics: vec![col_stats_i64(10)],
+ };
+
+ // Skip 200, fetch 300, so we get rows 200-500
+ let result = original_stats
+ .clone()
+ .with_fetch(Some(300), 200, 1)
+ .unwrap();
+
+ assert_eq!(result.num_rows, Precision::Exact(300));
+ // 300/1000 = 0.3, so 8000 * 0.3 = 2400
+ assert_eq!(result.total_byte_size, Precision::Inexact(2400));
+ }
+
+ #[test]
+ fn test_with_fetch_multi_partition() {
+ // Test with multiple partitions
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(1000), // per partition
+ total_byte_size: Precision::Exact(8000),
+ column_statistics: vec![col_stats_i64(10)],
+ };
+
+ // Fetch 100 per partition, 4 partitions = 400 total
+ let result = original_stats.clone().with_fetch(Some(100), 0, 4).unwrap();
+
+ assert_eq!(result.num_rows, Precision::Exact(400));
+ // 400/1000 = 0.4, so 8000 * 0.4 = 3200
+ assert_eq!(result.total_byte_size, Precision::Inexact(3200));
+ }
+
+ #[test]
+ fn test_with_fetch_absent_stats() {
+ // Test with absent statistics
+ let original_stats = Statistics {
+ num_rows: Precision::Absent,
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics {
+ null_count: Precision::Absent,
+ max_value: Precision::Absent,
+ min_value: Precision::Absent,
+ sum_value: Precision::Absent,
+ distinct_count: Precision::Absent,
+ }],
+ };
+
+ let result = original_stats.clone().with_fetch(Some(100), 0, 1).unwrap();
+
+ // With absent input stats, output should be inexact estimate
+ assert_eq!(result.num_rows, Precision::Inexact(100));
+ assert_eq!(result.total_byte_size, Precision::Absent);
+ // Column stats should remain absent
+ assert_eq!(result.column_statistics[0].null_count, Precision::Absent);
+ }
+
+ #[test]
+ fn test_with_fetch_fetch_exceeds_rows() {
+ // Test when fetch is larger than available rows after skip
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(100),
+ total_byte_size: Precision::Exact(800),
+ column_statistics: vec![col_stats_i64(10)],
+ };
+
+ // Skip 50, fetch 100, but only 50 rows remain
+ let result = original_stats.clone().with_fetch(Some(100), 50, 1).unwrap();
+
+ assert_eq!(result.num_rows, Precision::Exact(50));
+ // 50/100 = 0.5, so 800 * 0.5 = 400
+ assert_eq!(result.total_byte_size, Precision::Inexact(400));
+ }
+
+ #[test]
+ fn test_with_fetch_preserves_all_column_stats() {
+ // Comprehensive test that all column statistic fields are preserved
+ let original_col_stats = ColumnStatistics {
+ null_count: Precision::Exact(42),
+ max_value: Precision::Exact(ScalarValue::Int32(Some(999))),
+ min_value: Precision::Exact(ScalarValue::Int32(Some(-100))),
+ sum_value: Precision::Exact(ScalarValue::Int32(Some(123456))),
+ distinct_count: Precision::Exact(789),
+ };
+
+ let original_stats = Statistics {
+ num_rows: Precision::Exact(1000),
+ total_byte_size: Precision::Exact(8000),
+ column_statistics: vec![original_col_stats.clone()],
+ };
+
+ let result = original_stats.with_fetch(Some(250), 0, 1).unwrap();
+
+ let result_col_stats = &result.column_statistics[0];
+
+ // All values should be preserved but marked as inexact
+ assert_eq!(result_col_stats.null_count, Precision::Inexact(42));
+ assert_eq!(
+ result_col_stats.max_value,
+ Precision::Inexact(ScalarValue::Int32(Some(999)))
+ );
+ assert_eq!(
+ result_col_stats.min_value,
+ Precision::Inexact(ScalarValue::Int32(Some(-100)))
+ );
+ assert_eq!(
+ result_col_stats.sum_value,
+ Precision::Inexact(ScalarValue::Int32(Some(123456)))
+ );
+ assert_eq!(result_col_stats.distinct_count, Precision::Inexact(789));
+ }
}
diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs
index f9b0db0e80..403c04f173 100644
--- a/datafusion/core/tests/custom_sources_cases/statistics.rs
+++ b/datafusion/core/tests/custom_sources_cases/statistics.rs
@@ -265,17 +265,19 @@ async fn sql_filter() -> Result<()> {
#[tokio::test]
async fn sql_limit() -> Result<()> {
let (stats, schema) = fully_defined();
- let col_stats = Statistics::unknown_column(&schema);
let ctx = init_ctx(stats.clone(), schema)?;
let df = ctx.sql("SELECT * FROM stats_table LIMIT 5").await.unwrap();
let physical_plan = df.create_physical_plan().await.unwrap();
- // when the limit is smaller than the original number of lines
- // we loose all statistics except the for number of rows which becomes the limit
+ // when the limit is smaller than the original number of lines we mark the statistics as inexact
assert_eq!(
Statistics {
num_rows: Precision::Exact(5),
- column_statistics: col_stats,
+ column_statistics: stats
+ .column_statistics
+ .iter()
+ .map(|c| c.clone().to_inexact())
+ .collect(),
total_byte_size: Precision::Absent
},
physical_plan.partition_statistics(None)?
diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
index 7e9d5bf1b9..2332d79ac9 100644
--- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
@@ -490,11 +490,20 @@ mod test {
.map(|idx| local_limit.partition_statistics(Some(idx)))
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
- let schema = scan.schema();
- let mut expected_statistic_partition = Statistics::new_unknown(&schema);
- expected_statistic_partition.num_rows = Precision::Exact(1);
- assert_eq!(statistics[0], expected_statistic_partition);
- assert_eq!(statistics[1], expected_statistic_partition);
+ let mut expected_0 = statistics[0].clone();
+ expected_0.column_statistics = expected_0
+ .column_statistics
+ .into_iter()
+ .map(|c| c.to_inexact())
+ .collect();
+ let mut expected_1 = statistics[1].clone();
+ expected_1.column_statistics = expected_1
+ .column_statistics
+ .into_iter()
+ .map(|c| c.to_inexact())
+ .collect();
+ assert_eq!(statistics[0], expected_0);
+ assert_eq!(statistics[1], expected_1);
Ok(())
}
diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs
index d98530d28e..397bd9a377 100644
--- a/datafusion/physical-plan/src/coalesce_batches.rs
+++ b/datafusion/physical-plan/src/coalesce_batches.rs
@@ -202,12 +202,9 @@ impl ExecutionPlan for CoalesceBatchesExec {
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
- self.input.partition_statistics(partition)?.with_fetch(
- self.schema(),
- self.fetch,
- 0,
- 1,
- )
+ self.input
+ .partition_statistics(partition)?
+ .with_fetch(self.fetch, 0, 1)
}
fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs
index 976ff70502..323a844ae3 100644
--- a/datafusion/physical-plan/src/coalesce_partitions.rs
+++ b/datafusion/physical-plan/src/coalesce_partitions.rs
@@ -214,7 +214,7 @@ impl ExecutionPlan for CoalescePartitionsExec {
fn partition_statistics(&self, _partition: Option<usize>) -> Result<Statistics> {
self.input
.partition_statistics(None)?
- .with_fetch(self.schema(), self.fetch, 0, 1)
+ .with_fetch(self.fetch, 0, 1)
}
fn supports_limit_pushdown(&self) -> bool {
diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs
index 2224f85cc1..4611c547f2 100644
--- a/datafusion/physical-plan/src/limit.rs
+++ b/datafusion/physical-plan/src/limit.rs
@@ -194,12 +194,9 @@ impl ExecutionPlan for GlobalLimitExec {
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
- self.input.partition_statistics(partition)?.with_fetch(
- self.schema(),
- self.fetch,
- self.skip,
- 1,
- )
+ self.input
+ .partition_statistics(partition)?
+ .with_fetch(self.fetch, self.skip, 1)
}
fn fetch(&self) -> Option<usize> {
@@ -338,12 +335,9 @@ impl ExecutionPlan for LocalLimitExec {
}
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
- self.input.partition_statistics(partition)?.with_fetch(
- self.schema(),
- Some(self.fetch),
- 0,
- 1,
- )
+ self.input
+ .partition_statistics(partition)?
+ .with_fetch(Some(self.fetch), 0, 1)
}
fn fetch(&self) -> Option<usize> {
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index c026e43097..7f47d60c73 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -1291,19 +1291,14 @@ impl ExecutionPlan for SortExec {
fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
if !self.preserve_partitioning() {
- return self.input.partition_statistics(None)?.with_fetch(
- self.schema(),
- self.fetch,
- 0,
- 1,
- );
+ return self
+ .input
+ .partition_statistics(None)?
+ .with_fetch(self.fetch, 0, 1);
}
- self.input.partition_statistics(partition)?.with_fetch(
- self.schema(),
- self.fetch,
- 0,
- 1,
- )
+ self.input
+ .partition_statistics(partition)?
+ .with_fetch(self.fetch, 0, 1)
}
fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
diff --git a/datafusion/sqllogictest/test_files/parquet_statistics.slt b/datafusion/sqllogictest/test_files/parquet_statistics.slt
index f9e8990299..c04235ef4e 100644
--- a/datafusion/sqllogictest/test_files/parquet_statistics.slt
+++ b/datafusion/sqllogictest/test_files/parquet_statistics.slt
@@ -112,7 +112,7 @@ query TT
EXPLAIN SELECT * FROM test_table WHERE column1 = 1;
----
physical_plan
-01)CoalesceBatchesExec: target_batch_size=8192, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:)]]
+01)CoalesceBatchesExec: target_batch_size=8192, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(1)))]]
02)--FilterExec: column1@0 = 1, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)))]]
03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:)]]
04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/1.parquet]]}, projection=[column1], file_type=parquet, predicate=column1@0 = 1, pruning_predicate=column1_null_count@2 != row_count@3 AND column1_min@0 <= 1 AND 1 <= column1_max@1, required_guarantees=[column1 in (1)], statistics=[Rows=Absent, Bytes=Abs [...]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org