Merge pull request JanKaul#260 from JanKaul/detailed-tracing

Make sure to load Manifests concurrently

Author: JanKaul

datafusion_iceberg/src/statistics.rs

@@ -1,56 +1,23 @@
-use std::ops::Deref;
-
 use datafusion::{
     common::stats::Precision,
     physical_plan::{ColumnStatistics, Statistics},
     scalar::ScalarValue,
 };
-use futures::{future, TryFutureExt, TryStreamExt};
+use iceberg_rust::error::Error;
 use iceberg_rust::spec::{
     manifest::{ManifestEntry, Status},
     schema::Schema,
     values::Value,
 };
-use iceberg_rust::{catalog::tabular::Tabular, error::Error, table::Table};
-use itertools::Itertools;
-
-use super::table::DataFusionTable;
-
-impl DataFusionTable {
-    pub(crate) async fn statistics(&self) -> Result<Statistics, Error> {
-        match self.tabular.read().await.deref() {
-            Tabular::Table(table) => table_statistics(table, &self.snapshot_range).await,
-            Tabular::View(_) => Err(Error::NotSupported("Statistics for views".to_string())),
-            Tabular::MaterializedView(mv) => {
-                let table = mv.storage_table().await?;
-                table_statistics(&table, &self.snapshot_range).await
-            }
-        }
-    }
-}
 
-pub(crate) async fn table_statistics(
-    table: &Table,
-    snapshot_range: &(Option<i64>, Option<i64>),
-) -> Result<Statistics, Error> {
-    let schema = &snapshot_range
-        .1
-        .and_then(|snapshot_id| table.metadata().schema(snapshot_id).ok().cloned())
-        .unwrap_or_else(|| table.current_schema(None).unwrap().clone());
-
-    let sequence_number_range = [snapshot_range.0, snapshot_range.1]
-        .iter()
-        .map(|x| x.and_then(|y| table.metadata().sequence_number(y)))
-        .collect_tuple::<(Option<i64>, Option<i64>)>()
-        .unwrap();
-
-    let manifests = table.manifests(snapshot_range.0, snapshot_range.1).await?;
-    let datafiles = table
-        .datafiles(&manifests, None, sequence_number_range)
-        .await?;
+pub(crate) fn statistics_from_datafiles(
+    schema: &Schema,
+    datafiles: &[(String, ManifestEntry)],
+) -> Statistics {
     datafiles
-        .try_filter(|manifest| future::ready(!matches!(manifest.1.status(), Status::Deleted)))
-        .try_fold(
+        .iter()
+        .filter(|(_, manifest)| !matches!(manifest.status(), Status::Deleted))
+        .fold(
             Statistics {
                 num_rows: Precision::Exact(0),
                 total_byte_size: Precision::Exact(0),
@@ -65,14 +32,14 @@ pub(crate) async fn table_statistics(
                     schema.fields().len()
                 ],
             },
-            |acc, manifest| async move {
-                let column_stats = column_statistics(schema, &manifest.1);
-                Ok(Statistics {
+            |acc, (_, manifest)| {
+                let column_stats = column_statistics(schema, manifest);
+                Statistics {
                     num_rows: acc.num_rows.add(&Precision::Exact(
-                        *manifest.1.data_file().record_count() as usize,
+                        *manifest.data_file().record_count() as usize,
                     )),
                     total_byte_size: acc.total_byte_size.add(&Precision::Exact(
-                        *manifest.1.data_file().file_size_in_bytes() as usize,
+                        *manifest.data_file().file_size_in_bytes() as usize,
                     )),
                     column_statistics: acc
                         .column_statistics
@@ -86,11 +53,9 @@ pub(crate) async fn table_statistics(
                             sum_value: acc.sum_value.add(&x.sum_value),
                         })
                         .collect(),
-                })
+                }
             },
         )
-        .map_err(Error::from)
-        .await
 }
 
 fn column_statistics<'a>(

datafusion_iceberg/src/table.rs

@@ -33,8 +33,9 @@ use tokio::sync::{
     mpsc::{self},
     RwLock, RwLockWriteGuard,
 };
-use tracing::instrument;
+use tracing::{instrument, Instrument};
 
+use crate::statistics::statistics_from_datafiles;
 use crate::{
     error::Error as DataFusionIcebergError,
     pruning_statistics::{transform_predicate, PruneDataFiles, PruneManifests},
@@ -73,7 +74,6 @@ use datafusion::{
         projection::ProjectionExec,
         union::UnionExec,
         DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, SendableRecordBatchStream,
-        Statistics,
     },
     prelude::Expr,
     scalar::ScalarValue,
@@ -259,16 +259,11 @@ impl TableProvider for DataFusionTable {
             }
             Tabular::Table(table) => {
                 let schema = self.schema();
-                let statistics = self
-                    .statistics()
-                    .await
-                    .map_err(DataFusionIcebergError::from)?;
                 table_scan(
                     table,
                     &self.snapshot_range,
                     schema,
                     self.config.as_ref(),
-                    statistics,
                     session_state,
                     projection,
                     filters,
@@ -282,16 +277,11 @@ impl TableProvider for DataFusionTable {
                     .await
                     .map_err(DataFusionIcebergError::from)?;
                 let schema = self.schema();
-                let statistics = self
-                    .statistics()
-                    .await
-                    .map_err(DataFusionIcebergError::from)?;
                 table_scan(
                     &table,
                     &self.snapshot_range,
                     schema,
                     self.config.as_ref(),
-                    statistics,
                     session_state,
                     projection,
                     filters,
@@ -350,7 +340,7 @@ fn fake_object_store_url(table_location_url: &str) -> ObjectStoreUrl {
 }
 
 #[allow(clippy::too_many_arguments)]
-#[instrument(name = "datafusion_iceberg::table_scan", level = "debug", skip(arrow_schema, statistics, session, filters), fields(
+#[instrument(name = "datafusion_iceberg::table_scan", level = "debug", skip(arrow_schema, session, filters), fields(
     table_identifier = %table.identifier(),
     snapshot_range = ?snapshot_range,
     projection = ?projection,
@@ -362,7 +352,6 @@ async fn table_scan(
     snapshot_range: &(Option<i64>, Option<i64>),
     arrow_schema: SchemaRef,
     config: Option<&DataFusionTableConfig>,
-    statistics: Statistics,
     session: &SessionState,
     projection: Option<&Vec<usize>>,
     filters: &[Expr],
@@ -445,7 +434,7 @@ async fn table_scan(
         HashMap::new();
 
     // Prune data & delete file and insert them into the according map
-    if let Some(physical_predicate) = physical_predicate.clone() {
+    let statistics = if let Some(physical_predicate) = physical_predicate.clone() {
         let partition_schema = Arc::new(ArrowSchema::new(table_partition_cols.clone()));
         let partition_column_names = partition_fields
             .iter()
@@ -492,15 +481,13 @@ async fn table_scan(
                 .await
                 .map_err(DataFusionIcebergError::from)?
                 .try_collect()
-                .await
                 .map_err(DataFusionIcebergError::from)?
         } else {
             table
                 .datafiles(&manifests, None, sequence_number_range)
                 .await
                 .map_err(DataFusionIcebergError::from)?
                 .try_collect()
-                .await
                 .map_err(DataFusionIcebergError::from)?
         };
 
@@ -513,6 +500,8 @@ async fn table_scan(
             &data_files,
         ))?;
 
+        let statistics = statistics_from_datafiles(&schema, &data_files);
+
         data_files
             .into_iter()
             .zip(files_to_prune.into_iter())
@@ -537,18 +526,21 @@ async fn table_scan(
                     }
                 };
             });
+        statistics
     } else {
         let manifests = table
             .manifests(snapshot_range.0, snapshot_range.1)
             .await
             .map_err(DataFusionIcebergError::from)?;
-        let data_files: Vec<(ManifestPath, ManifestEntry)> = table
+        let data_files: Vec<_> = table
             .datafiles(&manifests, None, sequence_number_range)
             .await
             .map_err(DataFusionIcebergError::from)?
             .try_collect()
-            .await
             .map_err(DataFusionIcebergError::from)?;
+
+        let statistics = statistics_from_datafiles(&schema, &data_files);
+
         data_files.into_iter().for_each(|manifest| {
             if *manifest.1.status() != Status::Deleted {
                 match manifest.1.data_file().content() {
@@ -570,6 +562,7 @@ async fn table_scan(
                 }
             }
         });
+        statistics
     };
 
     let file_source = Arc::new(
@@ -865,6 +858,9 @@ async fn table_scan(
 
         let other_plan = ParquetFormat::default()
             .create_physical_plan(session, file_scan_config)
+            .instrument(tracing::debug_span!(
+                "datafusion_iceberg::create_physical_plan_scan_data_files"
+            ))
             .await?;
 
         plans.push(other_plan);

iceberg-rust/src/table/manifest_list.rs

@@ -13,7 +13,7 @@ use std::{
 use apache_avro::{
     types::Value as AvroValue, Reader as AvroReader, Schema as AvroSchema, Writer as AvroWriter,
 };
-use futures::{future::join_all, TryFutureExt, TryStreamExt};
+use futures::{future::join_all, stream, TryFutureExt, TryStreamExt};
 use iceberg_rust_spec::{
     manifest::{partition_value_schema, DataFile, ManifestEntry, Status},
     manifest_list::{
@@ -255,7 +255,7 @@ pub async fn snapshot_column_bounds(
 
     let primitive_field_ids = schema.primitive_field_ids().collect::<Vec<_>>();
     let n = primitive_field_ids.len();
-    datafiles
+    stream::iter(datafiles)
         .try_fold(None::<Rectangle>, |acc, (_, manifest)| {
             let primitive_field_ids = &primitive_field_ids;
             async move {