feat[iceberg]: runtime statistics

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<usize>,
     metrics: ExecutionPlanMetricsSet,
+    current_snapshot: Option<SnapshotRef>,
     iceberg_file_io: iceberg::io::FileIO,
     iceberg_runtime: iceberg::Runtime,
     feed: WorkUnitFeed<IcebergWorkUnitFeed>,
@@ -148,6 +159,8 @@ impl IcebergDataSource {
                 .collect::<Vec<String>>()
         });
 
+        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<usize>) -> Result<Arc<Statistics>> {
-        // 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<usize>) -> Option<Arc<dyn DataSource>> {
@@ -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<Arc<Statistics>> {
+    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()],
+    }))
+}

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<Arc<dyn ExecutionPlan>> {
+        self.ctx.sql(sql).await?.create_physical_plan().await
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

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<Statistics> {
+        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<dyn ExecutionPlan>) -> Option<Arc<DataSourceExec>> {
+        if let Some(exec) = plan.downcast_ref::<DataSourceExec>() {
+            if exec
+                .data_source()
+                .downcast_ref::<IcebergDataSource>()
+                .is_some()
+            {
+                return Some(Arc::new(exec.clone()));
+            }
+        }
+        plan.children().into_iter().find_map(find_iceberg_exec)
+    }
+}