Merge pull request JanKaul#258 from JanKaul/write-manifests-concurrently

Write manifests concurrently

Author: JanKaul

iceberg-rust/src/table/manifest.rs

@@ -17,6 +17,7 @@
 
 use std::{
     collections::HashSet,
+    future::Future,
     io::Read,
     iter::{repeat, Map, Repeat, Zip},
     sync::Arc,
@@ -26,6 +27,7 @@ use apache_avro::{
     to_value, types::Value as AvroValue, Reader as AvroReader, Schema as AvroSchema,
     Writer as AvroWriter,
 };
+use futures::TryFutureExt;
 use iceberg_rust_spec::{
     manifest::{Content, ManifestEntry, ManifestEntryV1, ManifestEntryV2, Status},
     manifest_list::{self, FieldSummary, ManifestListEntry},
@@ -636,6 +638,49 @@ impl<'schema, 'metadata> ManifestWriter<'schema, 'metadata> {
             .await?;
         Ok(self.manifest)
     }
+
+    /// Finishes writing the manifest file concurrently.
+    ///
+    /// This method completes the manifest writing process by finalizing the writer
+    /// and returning both the manifest list entry and a future for the actual file upload.
+    /// The upload operation can be awaited separately, allowing for concurrent processing
+    /// of multiple manifest writes.
+    ///
+    /// # Arguments
+    ///
+    /// * `object_store` - The object store implementation used to persist the manifest file
+    ///
+    /// # Returns
+    ///
+    /// Returns a tuple containing:
+    /// - `ManifestListEntry`: The completed manifest entry with updated metadata
+    /// - `impl Future<Output = Result<PutResult, Error>>`: A future that performs the actual file upload
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - The writer cannot be finalized
+    /// - There are issues preparing the upload operation
+    pub(crate) fn finish_concurrently(
+        mut self,
+        object_store: Arc<dyn ObjectStore>,
+    ) -> Result<(ManifestListEntry, impl Future<Output = Result<(), Error>>), Error> {
+        let manifest_bytes = self.writer.into_inner()?;
+
+        let manifest_length: i64 = manifest_bytes.len() as i64;
+
+        self.manifest.manifest_length += manifest_length;
+
+        let path = strip_prefix(&self.manifest.manifest_path).as_str().into();
+        let future = async move {
+            object_store
+                .put(&path, manifest_bytes.into())
+                .map_ok(|_| ())
+                .map_err(Error::from)
+                .await
+        };
+        Ok((self.manifest, future))
+    }
 }
 
 #[allow(clippy::type_complexity)]

iceberg-rust/src/table/manifest_list.rs

@@ -4,6 +4,7 @@
 
 use std::{
     collections::{HashMap, HashSet},
+    future::Future,
     io::{Cursor, Read},
     iter::{repeat, Map, Repeat, Zip},
     sync::Arc,
@@ -12,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, TryStreamExt};
+use futures::{future::join_all, TryFutureExt, TryStreamExt};
 use iceberg_rust_spec::{
     manifest::{partition_value_schema, DataFile, ManifestEntry, Status},
     manifest_list::{
@@ -700,6 +701,24 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
         .await
     }
 
+    #[inline]
+    pub(crate) async fn append_concurrently(
+        &mut self,
+        data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
+        snapshot_id: i64,
+        object_store: Arc<dyn ObjectStore>,
+        content: Content,
+    ) -> Result<impl Future<Output = Result<(), Error>>, Error> {
+        self.append_filtered_concurrently(
+            data_files,
+            snapshot_id,
+            None::<fn(&Result<ManifestEntry, Error>) -> bool>,
+            object_store,
+            content,
+        )
+        .await
+    }
+
     /// Appends data files to a single manifest with optional filtering and finalizes the manifest list.
     ///
     /// This method extends the basic `append` functionality by providing the ability to
@@ -754,6 +773,7 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
     ///     object_store,
     /// ).await?;
     /// ```
+    #[inline]
     pub(crate) async fn append_filtered(
         &mut self,
         data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
@@ -762,6 +782,20 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
         object_store: Arc<dyn ObjectStore>,
         content: Content,
     ) -> Result<(), Error> {
+        self.append_filtered_concurrently(data_files, snapshot_id, filter, object_store, content)
+            .await?
+            .await?;
+        Ok(())
+    }
+
+    pub(crate) async fn append_filtered_concurrently(
+        &mut self,
+        data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
+        snapshot_id: i64,
+        filter: Option<impl Fn(&Result<ManifestEntry, Error>) -> bool>,
+        object_store: Arc<dyn ObjectStore>,
+        content: Content,
+    ) -> Result<impl Future<Output = Result<(), Error>>, Error> {
         let selected_manifest = match content {
             Content::Data => self.selected_data_manifest.take(),
             Content::Deletes => self.selected_delete_manifest.take(),
@@ -820,11 +854,11 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
             manifest_writer.append(manifest_entry?)?;
         }
 
-        let manifest = manifest_writer.finish(object_store.clone()).await?;
+        let (manifest, future) = manifest_writer.finish_concurrently(object_store.clone())?;
 
         self.writer.append_ser(manifest)?;
 
-        Ok(())
+        Ok(future)
     }
 
     /// Appends data files by splitting them across multiple manifests and finalizes the manifest list.
@@ -875,15 +909,15 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
     ///     object_store,
     /// ).await?;
     /// ```
-    pub(crate) async fn append_multiple(
+    pub(crate) async fn append_multiple_concurrently(
         &mut self,
         data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
         snapshot_id: i64,
         n_splits: u32,
         object_store: Arc<dyn ObjectStore>,
         content: Content,
-    ) -> Result<(), Error> {
-        self.append_multiple_filtered(
+    ) -> Result<impl Future<Output = Result<(), Error>>, Error> {
+        self.append_multiple_filtered_concurrently(
             data_files,
             snapshot_id,
             n_splits,
@@ -955,6 +989,7 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
     ///     object_store,
     /// ).await?;
     /// ```
+    #[inline]
     pub(crate) async fn append_multiple_filtered(
         &mut self,
         data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
@@ -964,6 +999,28 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
         object_store: Arc<dyn ObjectStore>,
         content: Content,
     ) -> Result<(), Error> {
+        self.append_multiple_filtered_concurrently(
+            data_files,
+            snapshot_id,
+            n_splits,
+            filter,
+            object_store,
+            content,
+        )
+        .await?
+        .await?;
+        Ok(())
+    }
+
+    pub(crate) async fn append_multiple_filtered_concurrently(
+        &mut self,
+        data_files: impl Iterator<Item = Result<ManifestEntry, Error>>,
+        snapshot_id: i64,
+        n_splits: u32,
+        filter: Option<impl Fn(&Result<ManifestEntry, Error>) -> bool>,
+        object_store: Arc<dyn ObjectStore>,
+        content: Content,
+    ) -> Result<impl Future<Output = Result<(), Error>>, Error> {
         let partition_fields = self
             .table_metadata
             .current_partition_fields(self.branch.as_deref())?;
@@ -1032,7 +1089,7 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
             split_datafiles(data_files, bounds, &partition_column_names, n_splits)?
         };
 
-        let manifest_futures = splits
+        let (manifests, manifest_futures) = splits
             .into_iter()
             .map(|entries| {
                 let manifest_location = self.next_manifest_location();
@@ -1050,17 +1107,17 @@ impl<'schema, 'metadata> ManifestListWriter<'schema, 'metadata> {
                     manifest_writer.append(manifest_entry)?;
                 }
 
-                Ok::<_, Error>(manifest_writer.finish(object_store.clone()))
+                manifest_writer.finish_concurrently(object_store.clone())
             })
-            .collect::<Result<Vec<_>, _>>()?;
-
-        let manifests = futures::future::try_join_all(manifest_futures).await?;
+            .collect::<Result<(Vec<_>, Vec<_>), _>>()?;
 
         for manifest in manifests {
             self.writer.append_ser(manifest)?;
         }
 
-        Ok(())
+        let future = futures::future::try_join_all(manifest_futures).map_ok(|_| ());
+
+        Ok(future)
     }
 
     pub(crate) async fn finish(

iceberg-rust/src/table/transaction/operation.rs

@@ -3,9 +3,12 @@
 */
 
 use std::collections::HashSet;
+use std::future::Future;
+use std::pin::Pin;
 use std::{collections::HashMap, sync::Arc};
 
 use bytes::Bytes;
+use futures::future;
 use iceberg_rust_spec::manifest_list::{
     manifest_list_schema_v1, manifest_list_schema_v2, Content, ManifestListEntry,
 };
@@ -327,6 +330,10 @@ impl Operation {
 
                 // Write manifest files
                 // Split manifest file if limit is exceeded
+                #[allow(clippy::type_complexity)]
+                let mut futures: Vec<
+                    Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>,
+                > = Vec::new();
                 for (content, files, n_files) in [
                     (Content::Data, Either::Left(new_datafile_iter), n_data_files),
                     (
@@ -339,26 +346,37 @@ impl Operation {
                         let n_splits = manifest_list_writer.n_splits(n_files, content);
 
                         if n_splits == 0 {
-                            manifest_list_writer
-                                .append(files, snapshot_id, object_store.clone(), content)
+                            let future = manifest_list_writer
+                                .append_concurrently(
+                                    files,
+                                    snapshot_id,
+                                    object_store.clone(),
+                                    content,
+                                )
                                 .await?;
+                            futures.push(Box::pin(future));
                         } else {
-                            manifest_list_writer
-                                .append_multiple(
+                            let future = manifest_list_writer
+                                .append_multiple_concurrently(
                                     files,
                                     snapshot_id,
                                     n_splits,
                                     object_store.clone(),
                                     content,
                                 )
                                 .await?;
+                            futures.push(Box::pin(future));
                         }
                     }
                 }
 
-                let new_manifest_list_location = manifest_list_writer
-                    .finish(snapshot_id, object_store)
-                    .await?;
+                let manifest_future = future::try_join_all(futures);
+
+                let (_, new_manifest_list_location) = future::try_join(
+                    manifest_future,
+                    manifest_list_writer.finish(snapshot_id, object_store),
+                )
+                .await?;
 
                 let snapshot_operation = match (n_data_files, n_delete_files) {
                     (0, 0) => return Ok((None, Vec::new())),