feat[iceberg]: table statistics suite

iceberg/src/codec.rs

@@ -95,6 +95,8 @@ impl PhysicalExtensionCodec for IcebergCodec {
             partitioning,
             fetch,
             metrics: Default::default(),
+            column_stats: None,
+            table_snapshot: None,
             iceberg_file_io,
             iceberg_runtime: self.iceberg_runtime.clone(),
             feed,

iceberg/src/config.rs

@@ -14,6 +14,8 @@ extensions_options! {
         pub row_group_filtering_enabled: bool, default = true
         /// Whether to apply row-level selections while reading Parquet files.
         pub row_selection_enabled: bool, default = false
+        /// Whether to include column statistics read during planning
+        pub column_stats_enabled: bool, default = false
     }
 }
 

iceberg/src/data_source.rs

@@ -1,11 +1,12 @@
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::stats::Precision;
 use datafusion::common::{ColumnStatistics, Statistics};
 use datafusion::config::ConfigOptions;
 use datafusion::datasource::source::DataSource;
-use datafusion::error::Result;
+use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_expr::projection::ProjectionExprs;
 use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
@@ -16,11 +17,17 @@ use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSe
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayFormatType, SortOrderPushdownResult};
 use datafusion::prelude::Expr;
+use datafusion::scalar::ScalarValue;
 use datafusion_distributed::WorkUnitFeed;
 use futures::{StreamExt, TryStreamExt};
 use iceberg::arrow::ArrowReaderBuilder;
 use iceberg::io::FileIO;
-use iceberg::spec::SnapshotRef;
+use iceberg::puffin::APACHE_DATASKETCHES_THETA_V1;
+use iceberg::spec::{
+    DataFile, Datum, Manifest, ManifestContentType, ManifestList, PrimitiveLiteral, PrimitiveType,
+    SnapshotRef,
+};
+use iceberg::table::Table;
 
 use crate::common::{convert_filters_to_predicate, df_err, iceberg_err};
 use crate::work_unit_wire::FileScanTaskDecoder;
@@ -127,6 +134,8 @@ pub struct IcebergDataSource {
     pub(crate) partitioning: Partitioning,
     pub(crate) fetch: Option<usize>,
     pub(crate) metrics: ExecutionPlanMetricsSet,
+    pub(crate) column_stats: Option<Vec<ColumnStatistics>>,
+    pub(crate) table_snapshot: Option<SnapshotRef>,
     pub(crate) iceberg_file_io: FileIO,
     pub(crate) iceberg_runtime: iceberg::Runtime,
     pub(crate) feed: WorkUnitFeed<IcebergWorkUnitFeed>,
@@ -148,7 +157,7 @@ impl IcebergDataSource {
         table: iceberg::table::Table,
         schema: SchemaRef,
         partitioning: Partitioning,
-        opts: IcebergDataSourceOptions,
+        opts: IcebergDataSourceOptions<'_>,
     ) -> Self {
         let output_schema = match opts.projection {
             None => schema.clone(),
@@ -159,7 +168,12 @@ impl IcebergDataSource {
                 .map(|p| schema.field(*p).name().clone())
                 .collect::<Vec<String>>()
         });
-
+        // Necessary for time-travel queries
+        let table_snapshot = match opts.snapshot_id {
+            Some(snapshot_id) => table.metadata().snapshot_by_id(snapshot_id),
+            None => table.metadata().current_snapshot(),
+        }
+        .cloned();
         let predicates = convert_filters_to_predicate(opts.filters);
 
         Self {
@@ -179,8 +193,32 @@ impl IcebergDataSource {
                 partitioning,
                 sync_manager: Default::default(),
             }),
+            table_snapshot,
+            column_stats: None,
         }
     }
+
+    /// Creating an instance with per column statistics calculation, including:
+    /// - null_count, min_value, max_value, byte_size
+    pub(crate) async fn with_column_statistics(
+        mut self,
+        table: Table,
+        projection: Option<&Vec<usize>>,
+    ) -> Result<Self> {
+        let schema = match &self.table_snapshot {
+            Some(snap) => snap.schema(table.metadata()).map_err(df_err)?,
+            // empty table
+            None => table.metadata().current_schema().clone(),
+        };
+        let fields = schema.as_struct().fields().to_vec();
+        let field_ids: Vec<i32> = match projection {
+            Some(projection) => projection.iter().map(|&idx| fields[idx].id).collect(),
+            None => fields.iter().map(|f| f.id).collect(),
+        };
+        self.column_stats =
+            Some(compute_column_stats(table, field_ids, self.table_snapshot.clone()).await?);
+        Ok(self)
+    }
 }
 
 impl IcebergDataSource {
@@ -259,13 +297,16 @@ impl DataSource for IcebergDataSource {
     }
 
     fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
-        let Some(feed) = self.feed.inner() else {
+        if self.feed.inner().is_none() {
             return Ok(Arc::new(Statistics::new_unknown(&self.schema)));
-        };
-        stats_from_snapshot(
-            feed.iceberg_table.metadata().current_snapshot(),
-            &self.schema,
-        )
+        }
+        let mut stats = stats_from_snapshot(self.table_snapshot.clone(), &self.schema)?;
+
+        if let Some(col_stats) = &self.column_stats {
+            stats.column_statistics = col_stats.clone();
+        }
+
+        Ok(Arc::new(stats))
     }
 
     fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn DataSource>> {
@@ -311,35 +352,227 @@ impl DataSource for IcebergDataSource {
     }
 }
 
-/// Getting statistics from the provided snapshot.
-fn stats_from_snapshot(
-    snapshot: Option<&SnapshotRef>,
-    schema: &SchemaRef,
-) -> Result<Arc<Statistics>> {
+/// Getting stats out of snapshot's additional properties (no I/O overhead)
+fn stats_from_snapshot(snapshot: Option<SnapshotRef>, schema: &SchemaRef) -> Result<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 {
+        return Ok(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(|value| value.parse().ok())
+        .and_then(|v| v.parse().ok())
         .map(Precision::Exact)
         .unwrap_or(Precision::Absent);
     let total_byte_size = props
         .get(TOTAL_FILE_SIZE)
-        .and_then(|value| value.parse().ok())
+        .and_then(|v| v.parse().ok())
         .map(Precision::Exact)
         .unwrap_or(Precision::Absent);
 
-    Ok(Arc::new(Statistics {
+    Ok(Statistics {
         num_rows,
         total_byte_size,
         column_statistics: vec![ColumnStatistics::new_unknown(); schema.fields().len()],
-    }))
+    })
+}
+
+/// Getting number of distinct values for columns in a table
+///
+/// Note: https://iceberg.apache.org/puffin-spec/#apache-datasketches-theta-v1-blob-type
+fn ndv_from_metadata(table: &Table, snapshot_id: i64) -> HashMap<i32, usize> {
+    let mut ndvs = HashMap::new();
+    let Some(stats) = table.metadata().statistics_for_snapshot(snapshot_id) else {
+        return ndvs;
+    };
+    for blob in &stats.blob_metadata {
+        // Getting sketch estimates of ndv in every column
+        if blob.r#type != APACHE_DATASKETCHES_THETA_V1 {
+            continue;
+        }
+
+        let &[field_id] = &blob.fields[..] else {
+            continue;
+        };
+        if let Some(ndv) = blob.properties.get("ndv").and_then(|v| v.parse().ok()) {
+            ndvs.insert(field_id, ndv);
+        }
+    }
+    ndvs
+}
+
+/// Reading table statistics from data-files concurrently
+pub async fn compute_column_stats(
+    table: Table,
+    fields_ids: Vec<i32>,
+    snapshot: Option<SnapshotRef>,
+) -> Result<Vec<ColumnStatistics>> {
+    match snapshot {
+        Some(actual_snapshot) => {
+            let metadata = table.metadata();
+            let ml_bytes = table
+                .file_io()
+                .new_input(actual_snapshot.manifest_list())
+                .map_err(df_err)?
+                .read()
+                .await
+                .map_err(df_err)?;
+            let manifest_list = Arc::new(
+                ManifestList::parse_with_version(&ml_bytes, metadata.format_version())
+                    .map_err(df_err)?,
+            );
+
+            // If a table has delete files (i.e MOR) - we should account for that fact later and change the counts to `inexact`
+            let has_deletes = manifest_list
+                .entries()
+                .iter()
+                .any(|f| f.content == ManifestContentType::Deletes);
+            // Collecting all of the needed paths before spawning
+            let manifest_paths: Vec<_> = manifest_list
+                .entries()
+                .iter()
+                .filter(|mf| mf.content == ManifestContentType::Data)
+                .map(|mf| mf.manifest_path.clone())
+                .collect();
+            let mut join_set = tokio::task::JoinSet::new();
+
+            for path in manifest_paths {
+                let table = table.clone();
+                let fields_ids = fields_ids.clone();
+
+                join_set.spawn(async move {
+                    let manifest =
+                        Manifest::parse_avro(&table.file_io().new_input(&path)?.read().await?)?;
+
+                    let mut col_stats: Vec<Option<ColumnStatistics>> = vec![None; fields_ids.len()];
+
+                    for entry in manifest.entries().iter().filter(|e| e.is_alive()) {
+                        let df = entry.data_file();
+
+                        for (i, &id) in fields_ids.iter().enumerate() {
+                            let next = data_file_col_stats(df, id);
+                            col_stats[i] = Some(merge_col_stats(col_stats[i].take(), next));
+                        }
+                    }
+
+                    Ok::<_, iceberg::Error>(col_stats)
+                });
+            }
+
+            let mut merged: Vec<Option<ColumnStatistics>> = vec![None; fields_ids.len()];
+            while let Some(result) = join_set.join_next().await {
+                let manifest_stats = result
+                    .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))?
+                    .map_err(df_err)?;
+                for (acc, next) in merged.iter_mut().zip(manifest_stats) {
+                    if let Some(next) = next {
+                        *acc = Some(merge_col_stats(acc.take(), next));
+                    }
+                }
+            }
+
+            let mut merged: Vec<ColumnStatistics> = merged
+                .into_iter()
+                .map(|cs| cs.unwrap_or_else(ColumnStatistics::new_unknown))
+                .collect();
+
+            let ndvs = ndv_from_metadata(&table, actual_snapshot.snapshot_id());
+            for (cs, id) in merged.iter_mut().zip(fields_ids.iter()) {
+                if let Some(&ndv) = ndvs.get(id) {
+                    cs.distinct_count = Precision::Inexact(ndv);
+                }
+            }
+
+            // TODO: probably we can do something about the delete files?
+            // However it wouldn't be free of cost in terms of performance
+            if has_deletes {
+                for cs in &mut merged {
+                    cs.null_count = cs.null_count.to_inexact();
+                }
+            }
+
+            Ok(merged)
+        }
+        None => Ok(vec![ColumnStatistics::new_unknown(); fields_ids.len()]),
+    }
+}
+
+/// Merging table's column statistics incrementally
+fn merge_col_stats(acc: Option<ColumnStatistics>, next: ColumnStatistics) -> ColumnStatistics {
+    match acc {
+        None => next,
+        Some(acc) => ColumnStatistics {
+            null_count: acc.null_count.add(&next.null_count),
+            min_value: acc.min_value.min(&next.min_value),
+            max_value: acc.max_value.max(&next.max_value),
+            byte_size: acc.byte_size.add(&next.byte_size),
+            sum_value: Precision::Absent,
+            distinct_count: Precision::Absent,
+        },
+    }
+}
+
+fn data_file_col_stats(df: &DataFile, id: i32) -> ColumnStatistics {
+    ColumnStatistics {
+        null_count: df
+            .null_value_counts()
+            .get(&id)
+            .map(|n| Precision::Exact(*n as usize))
+            .unwrap_or(Precision::Absent),
+        min_value: df
+            .lower_bounds()
+            .get(&id)
+            .and_then(datum_to_scalar)
+            .map(Precision::Inexact)
+            .unwrap_or(Precision::Absent),
+        max_value: df
+            .upper_bounds()
+            .get(&id)
+            .and_then(datum_to_scalar)
+            .map(Precision::Inexact)
+            .unwrap_or(Precision::Absent),
+        byte_size: df
+            .column_sizes()
+            .get(&id)
+            .map(|n| Precision::Inexact(*n as usize))
+            .unwrap_or(Precision::Absent),
+        sum_value: Precision::Absent,
+        distinct_count: Precision::Absent,
+    }
+}
+
+/// Conversion function of iceberg's Datum
+fn datum_to_scalar(d: &Datum) -> Option<ScalarValue> {
+    match (d.data_type(), d.literal()) {
+        (PrimitiveType::Boolean, PrimitiveLiteral::Boolean(v)) => {
+            Some(ScalarValue::Boolean(Some(*v)))
+        }
+        (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => Some(ScalarValue::Int32(Some(*v))),
+        (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => Some(ScalarValue::Int64(Some(*v))),
+        (PrimitiveType::Float, PrimitiveLiteral::Float(v)) => {
+            Some(ScalarValue::Float32(Some(v.into_inner())))
+        }
+        (PrimitiveType::Double, PrimitiveLiteral::Double(v)) => {
+            Some(ScalarValue::Float64(Some(v.into_inner())))
+        }
+        (PrimitiveType::String, PrimitiveLiteral::String(s)) => {
+            Some(ScalarValue::Utf8(Some(s.clone())))
+        }
+        (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => Some(ScalarValue::Date32(Some(*v))),
+        (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => {
+            Some(ScalarValue::TimestampMicrosecond(Some(*v), None))
+        }
+        (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => Some(
+            ScalarValue::TimestampMicrosecond(Some(*v), Some("UTC".into())),
+        ),
+        (PrimitiveType::Decimal { precision, scale }, PrimitiveLiteral::Int128(v)) => Some(
+            ScalarValue::Decimal128(Some(*v), *precision as u8, *scale as i8),
+        ),
+        _ => None,
+    }
 }