use iterator instead of stream

Author: JanKaul

datafusion_iceberg/src/statistics.rs

@@ -5,7 +5,7 @@ use datafusion::{
     physical_plan::{ColumnStatistics, Statistics},
     scalar::ScalarValue,
 };
-use futures::{future, TryFutureExt, TryStreamExt};
+use futures::{future, stream, TryFutureExt, TryStreamExt};
 use iceberg_rust::spec::{
     manifest::{ManifestEntry, Status},
     schema::Schema,
@@ -48,7 +48,7 @@ pub(crate) async fn table_statistics(
     let datafiles = table
         .datafiles(&manifests, None, sequence_number_range)
         .await?;
-    datafiles
+    stream::iter(datafiles)
         .try_filter(|manifest| future::ready(!matches!(manifest.1.status(), Status::Deleted)))
         .try_fold(
             Statistics {

datafusion_iceberg/src/table.rs

@@ -33,7 +33,7 @@ use tokio::sync::{
     mpsc::{self},
     RwLock, RwLockWriteGuard,
 };
-use tracing::{instrument, Instrument};
+use tracing::instrument;
 
 use crate::{
     error::Error as DataFusionIcebergError,
@@ -492,21 +492,13 @@ async fn table_scan(
                 .await
                 .map_err(DataFusionIcebergError::from)?
                 .try_collect()
-                .instrument(tracing::debug_span!(
-                    "datafusion_iceberg::collect_datafiles"
-                ))
-                .await
                 .map_err(DataFusionIcebergError::from)?
         } else {
             table
                 .datafiles(&manifests, None, sequence_number_range)
                 .await
                 .map_err(DataFusionIcebergError::from)?
                 .try_collect()
-                .instrument(tracing::debug_span!(
-                    "datafusion_iceberg::collect_datafiles"
-                ))
-                .await
                 .map_err(DataFusionIcebergError::from)?
         };
 
@@ -553,10 +545,6 @@ async fn table_scan(
             .await
             .map_err(DataFusionIcebergError::from)?
             .try_collect()
-            .instrument(tracing::debug_span!(
-                "datafusion_iceberg::collect_datafiles"
-            ))
-            .await
             .map_err(DataFusionIcebergError::from)?;
         data_files.into_iter().for_each(|manifest| {
             if *manifest.1.status() != Status::Deleted {

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 {

iceberg-rust/src/table/mod.rs

@@ -12,13 +12,13 @@
 
 use std::{io::Cursor, sync::Arc};
 
-use futures::future::{self, try_join_all};
+use futures::future::try_join_all;
 use itertools::Itertools;
 use manifest::ManifestReader;
 use manifest_list::read_snapshot;
 use object_store::{path::Path, ObjectStore};
 
-use futures::{stream, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use futures::{stream, StreamExt, TryFutureExt, TryStreamExt};
 use iceberg_rust_spec::util::{self};
 use iceberg_rust_spec::{
     spec::{
@@ -262,7 +262,8 @@ impl Table {
         manifests: &'a [ManifestListEntry],
         filter: Option<Vec<bool>>,
         sequence_number_range: (Option<i64>, Option<i64>),
-    ) -> Result<impl Stream<Item = Result<(ManifestPath, ManifestEntry), Error>> + 'a, Error> {
+    ) -> Result<impl Iterator<Item = Result<(ManifestPath, ManifestEntry), Error>> + 'a, Error>
+    {
         datafiles(
             self.object_store(),
             manifests,
@@ -279,7 +280,7 @@ impl Table {
     ) -> Result<bool, Error> {
         let manifests = self.manifests(start, end).await?;
         let datafiles = self.datafiles(&manifests, None, (None, None)).await?;
-        datafiles
+        stream::iter(datafiles)
             .try_any(|entry| async move { !matches!(entry.1.data_file().content(), Content::Data) })
             .await
     }
@@ -311,7 +312,7 @@ async fn datafiles(
     manifests: &'_ [ManifestListEntry],
     filter: Option<Vec<bool>>,
     sequence_number_range: (Option<i64>, Option<i64>),
-) -> Result<impl Stream<Item = Result<(ManifestPath, ManifestEntry), Error>> + '_, Error> {
+) -> Result<impl Iterator<Item = Result<(ManifestPath, ManifestEntry), Error>> + '_, Error> {
     // filter manifest files according to filter vector
     let iter: Box<dyn Iterator<Item = &ManifestListEntry> + Send + Sync> = match filter {
         Some(predicate) => {
@@ -345,31 +346,34 @@ async fn datafiles(
 
     let results = try_join_all(futures).await?;
 
-    Ok(stream::iter(results).flat_map(move |result| {
+    Ok(results.into_iter().flat_map(move |result| {
         let (bytes, path, sequence_number) = result;
 
         let reader = ManifestReader::new(bytes).unwrap();
-        stream::iter(reader).try_filter_map(move |mut x| {
-            future::ready({
-                let sequence_number = if let Some(sequence_number) = x.sequence_number() {
-                    *sequence_number
-                } else {
-                    *x.sequence_number_mut() = Some(sequence_number);
-                    sequence_number
-                };
+        reader.filter_map(move |x| {
+            let mut x = match x {
+                Ok(entry) => entry,
+                Err(_) => return None,
+            };
 
-                let filter = match sequence_number_range {
-                    (Some(start), Some(end)) => start < sequence_number && sequence_number <= end,
-                    (Some(start), None) => start < sequence_number,
-                    (None, Some(end)) => sequence_number <= end,
-                    _ => true,
-                };
-                if filter {
-                    Ok(Some((path.to_owned(), x)))
-                } else {
-                    Ok(None)
-                }
-            })
+            let sequence_number = if let Some(sequence_number) = x.sequence_number() {
+                *sequence_number
+            } else {
+                *x.sequence_number_mut() = Some(sequence_number);
+                sequence_number
+            };
+
+            let filter = match sequence_number_range {
+                (Some(start), Some(end)) => start < sequence_number && sequence_number <= end,
+                (Some(start), None) => start < sequence_number,
+                (None, Some(end)) => sequence_number <= end,
+                _ => true,
+            };
+            if filter {
+                Some(Ok((path.to_owned(), x)))
+            } else {
+                None
+            }
         })
     }))
 }
@@ -390,7 +394,7 @@ pub(crate) async fn delete_all_table_files(
     let snapshots = &metadata.snapshots;
 
     // stream::iter(datafiles.into_iter())
-    datafiles
+    stream::iter(datafiles)
         .try_for_each_concurrent(None, |datafile| {
             let object_store = object_store.clone();
             async move {