diff --git a/iceberg/src/data_source.rs b/iceberg/src/data_source.rs index ed15343dc..6a3666822 100644 --- a/iceberg/src/data_source.rs +++ b/iceberg/src/data_source.rs @@ -1,7 +1,8 @@ use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::{Statistics, exec_datafusion_err}; +use datafusion::common::stats::Precision; +use datafusion::common::{ColumnStatistics, Statistics, exec_datafusion_err}; use datafusion::config::ConfigOptions; use datafusion::datasource::source::DataSource; use datafusion::error::Result; @@ -18,10 +19,19 @@ use datafusion::prelude::Expr; use datafusion_distributed::WorkUnitFeed; use futures::{StreamExt, TryStreamExt}; use iceberg::arrow::ArrowReaderBuilder; +use iceberg::spec::SnapshotRef; use crate::common::{convert_filters_to_predicate, df_err, iceberg_err}; use crate::{IcebergConfig, IcebergWorkUnitFeed}; +/// Snapshot summary keys defined by the Iceberg table spec: +/// https://iceberg.apache.org/spec/#optional-snapshot-summary-fields +/// +/// iceberg-rust defines them privately: +/// https://github.com/apache/iceberg-rust/blob/4168a0b2950dc5f85588e5cb3ab6796e5228b309/crates/iceberg/src/spec/snapshot_summary.rs#L46-L47 +const TOTAL_RECORDS: &str = "total-records"; +const TOTAL_FILE_SIZE: &str = "total-files-size"; + /// Consumes a stream of [iceberg::scan::FileScanTask]s per partition and reads the underlying /// files into an Arrow stream. /// @@ -115,6 +125,7 @@ pub struct IcebergDataSource { partitioning: Partitioning, fetch: Option, metrics: ExecutionPlanMetricsSet, + current_snapshot: Option, iceberg_file_io: iceberg::io::FileIO, iceberg_runtime: iceberg::Runtime, feed: WorkUnitFeed, @@ -148,6 +159,8 @@ impl IcebergDataSource { .collect::>() }); + let current_snapshot = table.metadata().current_snapshot().cloned(); + let predicates = convert_filters_to_predicate(opts.filters); Self { @@ -167,6 +180,7 @@ impl IcebergDataSource { partitioning, sync_manager: Default::default(), }), + current_snapshot, } } } @@ -249,10 +263,7 @@ impl DataSource for IcebergDataSource { } fn partition_statistics(&self, _partition: Option) -> Result> { - // TODO: Implement planning time statistics for this DataSource. - // At this point, we have information about the iceberg::table::Table which we are about - // to read, so maybe there's something we can get from there. - Ok(Arc::new(Statistics::new_unknown(&self.schema))) + stats_from_snapshot(self.current_snapshot.as_ref(), &self.schema) } fn with_fetch(&self, fetch: Option) -> Option> { @@ -297,3 +308,34 @@ impl DataSource for IcebergDataSource { Ok(SortOrderPushdownResult::Unsupported) } } + +/// Getting statistics from the provided snapshot. +fn stats_from_snapshot( + snapshot: Option<&SnapshotRef>, + schema: &SchemaRef, +) -> Result> { + let Some(snap) = snapshot else { + // A table with no current snapshot has never had a commit. It was created, but zero data files were added + return Ok(Arc::new(Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ColumnStatistics::new_unknown(); schema.fields().len()], + })); + }; + let props = &snap.summary().additional_properties; + + let num_rows = props + .get(TOTAL_RECORDS) + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + let total_byte_size = props + .get(TOTAL_FILE_SIZE) + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + + Ok(Arc::new(Statistics { + num_rows: Precision::Exact(num_rows), + total_byte_size: Precision::Exact(total_byte_size), + column_statistics: vec![ColumnStatistics::new_unknown(); schema.fields().len()], + })) +} diff --git a/iceberg/src/test_utils/harness.rs b/iceberg/src/test_utils/harness.rs index c8be724c3..b837b2ca1 100644 --- a/iceberg/src/test_utils/harness.rs +++ b/iceberg/src/test_utils/harness.rs @@ -7,7 +7,7 @@ use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::dataframe::DataFrame; use datafusion::error::Result; use datafusion::execution::SessionStateBuilder; -use datafusion::physical_plan::displayable; +use datafusion::physical_plan::{ExecutionPlan, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; use futures::StreamExt; use futures::stream::BoxStream; @@ -58,6 +58,10 @@ impl IcebergTestHarness { pretty_format_batches(&batches)?.to_string(), )) } + + pub async fn physical_plan(&self, sql: &str) -> Result> { + self.ctx.sql(sql).await?.create_physical_plan().await + } } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/iceberg/tests/statistics.rs b/iceberg/tests/statistics.rs new file mode 100644 index 000000000..87ffa131b --- /dev/null +++ b/iceberg/tests/statistics.rs @@ -0,0 +1,133 @@ +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::common::Statistics; + use datafusion::common::stats::Precision; + use datafusion::datasource::source::DataSourceExec; + use datafusion::error::Result; + use datafusion::physical_plan::{ExecutionPlan, displayable}; + use datafusion_distributed_iceberg::IcebergDataSource; + use datafusion_distributed_iceberg::test_utils::IcebergTestHarness; + + // Values from testdata/iceberg/taxi/metadata/v1.metadata.json snapshot summary. + const TAXI_ROWS: usize = 175_000; + const TAXI_BYTES: usize = 4_480_382; + const TAXI_COLUMNS: usize = 13; + + #[tokio::test] + async fn reports_exact_row_count_and_byte_size_for_full_scan() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let stats = source_statistics(&harness, "SELECT * FROM taxi").await?; + + assert_eq!(stats.num_rows, Precision::Exact(TAXI_ROWS)); + assert_eq!(stats.total_byte_size, Precision::Exact(TAXI_BYTES)); + Ok(()) + } + + #[tokio::test] + async fn column_statistics_match_full_schema() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let stats = source_statistics(&harness, "SELECT * FROM taxi").await?; + + assert_eq!(stats.column_statistics.len(), TAXI_COLUMNS); + Ok(()) + } + + #[tokio::test] + async fn column_statistics_match_projected_schema() -> Result<()> { + // Regression: a column_statistics vec shorter than the output schema + // makes DataFusion panic while propagating statistics upstream. + let harness = IcebergTestHarness::new().await?; + let stats = source_statistics(&harness, "SELECT vendor_id, pickup_date FROM taxi").await?; + + assert_eq!(stats.column_statistics.len(), 2); + assert_eq!(stats.num_rows, Precision::Exact(TAXI_ROWS)); + Ok(()) + } + + #[tokio::test] + async fn statistics_propagate_through_filter() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let plan = harness + .physical_plan("SELECT vendor_id FROM taxi WHERE pickup_date = DATE '2024-01-10'") + .await?; + let stats = plan.partition_statistics(None)?; + + // The filter cannot keep the count exact, but it must not lose it. + assert!(matches!(stats.num_rows, Precision::Inexact(_))); + assert_eq!(stats.column_statistics.len(), 1); + Ok(()) + } + + #[tokio::test] + async fn statistics_propagate_through_projection_and_sort() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let plan = harness + .physical_plan("SELECT vendor_id, trip_distance FROM taxi ORDER BY pickup_at") + .await?; + let stats = plan.partition_statistics(None)?; + + assert_eq!(stats.num_rows, Precision::Exact(TAXI_ROWS)); + assert_eq!(stats.column_statistics.len(), 2); + Ok(()) + } + + #[tokio::test] + async fn explain_shows_statistics_on_the_iceberg_source() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + let plan = harness.physical_plan("SELECT vendor_id FROM taxi").await?; + let display = displayable(plan.as_ref()) + .set_show_statistics(true) + .indent(true) + .to_string(); + + insta::assert_snapshot!(display, @" + CooperativeExec, statistics=[Rows=Exact(175000), Bytes=Exact(4480382), [(Col[0]:)]] + DataSourceExec: format=iceberg, projection=[vendor_id], statistics=[Rows=Exact(175000), Bytes=Exact(4480382), [(Col[0]:)]] + "); + Ok(()) + } + + #[tokio::test] + async fn exact_row_count_lets_count_star_skip_the_scan() -> Result<()> { + // With Precision::Exact(num_rows) the AggregateStatistics optimizer + // rule answers COUNT(*) from metadata without reading any data file. + let harness = IcebergTestHarness::new().await?; + let (plan, batches) = harness.query("SELECT count(*) FROM taxi").await?; + + insta::assert_snapshot!(plan, @" + ProjectionExec: expr=[175000 as count(*)] + PlaceholderRowExec + "); + insta::assert_snapshot!(batches, @" + +----------+ + | count(*) | + +----------+ + | 175000 | + +----------+ + "); + Ok(()) + } + + /// Finds the single Iceberg `DataSourceExec` in the plan and returns the + /// statistics reported by the `IcebergDataSource` itself. + async fn source_statistics(harness: &IcebergTestHarness, sql: &str) -> Result { + let plan = harness.physical_plan(sql).await?; + let exec = find_iceberg_exec(&plan).expect("plan contains an Iceberg DataSourceExec"); + Ok(Arc::unwrap_or_clone(exec.partition_statistics(None)?)) + } + + fn find_iceberg_exec(plan: &Arc) -> Option> { + if let Some(exec) = plan.downcast_ref::() { + if exec + .data_source() + .downcast_ref::() + .is_some() + { + return Some(Arc::new(exec.clone())); + } + } + plan.children().into_iter().find_map(find_iceberg_exec) + } +}