refactor: avoid async_trait macro for IcebergWriter and provide extra dyn trait for object safety

Cargo.toml

@@ -59,6 +59,7 @@ bytes = "1.5"
 chrono = "0.4.34"
 ctor = "0.2.8"
 derive_builder = "0.20"
+dyn-clone = "1"
 either = "1"
 env_logger = "0.11.0"
 fnv = "1"

crates/iceberg/Cargo.toml

@@ -58,6 +58,7 @@ bitvec = { workspace = true }
 bytes = { workspace = true }
 chrono = { workspace = true }
 derive_builder = { workspace = true }
+dyn-clone = { workspace = true }
 fnv = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }

crates/iceberg/src/writer/base_writer/data_file_writer.rs

@@ -42,7 +42,6 @@ impl<B: FileWriterBuilder> DataFileWriterBuilder<B> {
     }
 }
 
-#[async_trait::async_trait]
 impl<B: FileWriterBuilder> IcebergWriterBuilder for DataFileWriterBuilder<B> {
     type R = DataFileWriter<B>;
 
@@ -61,13 +60,12 @@ pub struct DataFileWriter<B: FileWriterBuilder> {
     partition_value: Struct,
 }
 
-#[async_trait::async_trait]
 impl<B: FileWriterBuilder> IcebergWriter for DataFileWriter<B> {
     async fn write(&mut self, batch: RecordBatch) -> Result<()> {
         self.inner_writer.as_mut().unwrap().write(&batch).await
     }
 
-    async fn close(&mut self) -> Result<Vec<DataFile>> {
+    async fn close(mut self) -> Result<Vec<DataFile>> {
         let writer = self.inner_writer.take().unwrap();
         Ok(writer
             .close()
@@ -128,7 +126,7 @@ mod test {
             location_gen,
             file_name_gen,
         );
-        let mut data_file_writer = DataFileWriterBuilder::new(pw, None).build().await?;
+        let data_file_writer = DataFileWriterBuilder::new(pw, None).build().await?;
 
         let data_file = data_file_writer.close().await.unwrap();
         assert_eq!(data_file.len(), 1);

crates/iceberg/src/writer/base_writer/equality_delete_writer.rs

@@ -107,7 +107,6 @@ impl EqualityDeleteWriterConfig {
     }
 }
 
-#[async_trait::async_trait]
 impl<B: FileWriterBuilder> IcebergWriterBuilder for EqualityDeleteFileWriterBuilder<B> {
     type R = EqualityDeleteFileWriter<B>;
 
@@ -130,7 +129,6 @@ pub struct EqualityDeleteFileWriter<B: FileWriterBuilder> {
     partition_value: Struct,
 }
 
-#[async_trait::async_trait]
 impl<B: FileWriterBuilder> IcebergWriter for EqualityDeleteFileWriter<B> {
     async fn write(&mut self, batch: RecordBatch) -> Result<()> {
         let batch = self.projector.project_bacth(batch)?;
@@ -144,7 +142,7 @@ impl<B: FileWriterBuilder> IcebergWriter for EqualityDeleteFileWriter<B> {
         }
     }
 
-    async fn close(&mut self) -> Result<Vec<DataFile>> {
+    async fn close(mut self) -> Result<Vec<DataFile>> {
         if let Some(writer) = self.inner_writer.take() {
             Ok(writer
                 .close()
@@ -342,11 +340,11 @@ mod test {
         ])) as ArrayRef;
         let col3 = Arc::new({
             let list_parts = arrow_array::ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
-              Some(
-                  vec![Some(1),]
-              );
-              1024
-          ])
+                Some(
+                    vec![Some(1), ]
+                );
+                1024
+            ])
             .into_parts();
             arrow_array::ListArray::new(
                 if let DataType::List(field) = arrow_schema.fields.get(3).unwrap().data_type() {

crates/iceberg/src/writer/mod.rs

@@ -48,6 +48,8 @@
 pub mod base_writer;
 pub mod file_writer;
 
+use std::future::Future;
+
 use arrow_array::RecordBatch;
 
 use crate::spec::DataFile;
@@ -57,28 +59,128 @@ type DefaultInput = RecordBatch;
 type DefaultOutput = Vec<DataFile>;
 
 /// The builder for iceberg writer.
-#[async_trait::async_trait]
 pub trait IcebergWriterBuilder<I = DefaultInput, O = DefaultOutput>:
     Send + Clone + 'static
 {
     /// The associated writer type.
     type R: IcebergWriter<I, O>;
     /// Build the iceberg writer.
-    async fn build(self) -> Result<Self::R>;
+    fn build(self) -> impl Future<Output = Result<Self::R>> + Send;
 }
 
 /// The iceberg writer used to write data to iceberg table.
-#[async_trait::async_trait]
 pub trait IcebergWriter<I = DefaultInput, O = DefaultOutput>: Send + 'static {
     /// Write data to iceberg table.
-    async fn write(&mut self, input: I) -> Result<()>;
+    fn write(&mut self, input: I) -> impl Future<Output = Result<()>> + Send + '_;
     /// Close the writer and return the written data files.
     /// If close failed, the data written before maybe be lost. User may need to recreate the writer and rewrite the data again.
     /// # NOTE
     /// After close, regardless of success or failure, the writer should never be used again, otherwise the writer will panic.
-    async fn close(&mut self) -> Result<O>;
+    fn close(self) -> impl Future<Output = Result<O>> + Send;
 }
 
+mod dyn_trait {
+    use dyn_clone::{clone_trait_object, DynClone};
+
+    use super::Result;
+    use crate::writer::{DefaultInput, DefaultOutput, IcebergWriter, IcebergWriterBuilder};
+
+    #[async_trait::async_trait]
+    pub trait DynIcebergWriterBuilder<I, O>: Send + DynClone + 'static {
+        async fn dyn_build(self: Box<Self>) -> Result<BoxedIcebergWriter<I, O>>;
+    }
+
+    clone_trait_object!(<I, O> DynIcebergWriterBuilder<I, O>);
+
+    #[async_trait::async_trait]
+    impl<I: 'static + Send, O: 'static + Send, B: IcebergWriterBuilder<I, O>>
+        DynIcebergWriterBuilder<I, O> for B
+    {
+        async fn dyn_build(self: Box<Self>) -> Result<BoxedIcebergWriter<I, O>> {
+            Ok(<Self as IcebergWriterBuilder<I, O>>::build(*self)
+                .await?
+                .boxed())
+        }
+    }
+
+    /// Type alias for `Box<dyn DynIcebergWriterBuilder>`
+    pub type BoxedIcebergWriterBuilder<I = DefaultInput, O = DefaultOutput> =
+        Box<dyn DynIcebergWriterBuilder<I, O>>;
+
+    impl<I: Send + 'static, O: Send + 'static> IcebergWriterBuilder<I, O>
+        for BoxedIcebergWriterBuilder<I, O>
+    {
+        type R = BoxedIcebergWriter<I, O>;
+
+        async fn build(self) -> Result<Self::R> {
+            self.dyn_build().await
+        }
+    }
+
+    /// Extension methods for `IcebergWriterBuilder`
+    pub trait IcebergWriterBuilderDynExt<I: Send + 'static, O: Send + 'static>:
+        IcebergWriterBuilder<I, O> + Sized
+    {
+        /// Create a type erased `IcebergWriterBuilder` wrapped with `Box`.
+        fn boxed(self) -> BoxedIcebergWriterBuilder<I, O> {
+            Box::new(self) as _
+        }
+    }
+
+    impl<B, I: Send + 'static, O: Send + 'static> IcebergWriterBuilderDynExt<I, O> for B where B: IcebergWriterBuilder<I, O>
+    {}
+
+    /// The dyn iceberg writer used to write data to iceberg table.
+    #[async_trait::async_trait]
+    pub trait DynIcebergWriter<I, O>: Send + 'static {
+        /// `write` of trait `IcebergWriter`
+        async fn dyn_write(&mut self, input: I) -> Result<()>;
+        /// `close` of trait `IcebergWriter`
+        async fn dyn_close(self: Box<Self>) -> Result<O>;
+    }
+
+    #[async_trait::async_trait]
+    impl<I: 'static + Send, O: 'static + Send, W: IcebergWriter<I, O>> DynIcebergWriter<I, O> for W {
+        async fn dyn_write(&mut self, input: I) -> Result<()> {
+            self.write(input).await
+        }
+
+        async fn dyn_close(self: Box<Self>) -> Result<O> {
+            (*self).close().await
+        }
+    }
+
+    /// Type alias for `Box<dyn DynIcebergWriter>`
+    pub type BoxedIcebergWriter<I = DefaultInput, O = DefaultOutput> =
+        Box<dyn DynIcebergWriter<I, O>>;
+
+    impl<I: 'static + Send, O: 'static + Send> IcebergWriter<I, O> for BoxedIcebergWriter<I, O> {
+        async fn write(&mut self, input: I) -> Result<()> {
+            (**self).dyn_write(input).await
+        }
+
+        async fn close(self) -> Result<O> {
+            self.dyn_close().await
+        }
+    }
+
+    /// Extension methods for `IcebergWriter`
+    pub trait IcebergWriterDynExt<I: Send + 'static, O: Send + 'static>:
+        IcebergWriter<I, O> + Sized
+    {
+        /// Create a type erased `IcebergWriter` wrapped with `Box`.
+        fn boxed(self) -> BoxedIcebergWriter<I, O> {
+            Box::new(self) as _
+        }
+    }
+
+    impl<I: Send + 'static, O: Send + 'static, W: IcebergWriter<I, O>> IcebergWriterDynExt<I, O> for W {}
+}
+
+pub use dyn_trait::{
+    BoxedIcebergWriter, BoxedIcebergWriterBuilder, IcebergWriterBuilderDynExt, IcebergWriterDynExt,
+};
+
 /// The current file status of iceberg writer. It implement for the writer which write a single
 /// file.
 pub trait CurrentFileStatus {
@@ -92,21 +194,47 @@ pub trait CurrentFileStatus {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use arrow_array::RecordBatch;
     use arrow_schema::Schema;
     use arrow_select::concat::concat_batches;
     use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+    use parquet::file::properties::WriterProperties;
 
-    use super::IcebergWriter;
-    use crate::io::FileIO;
+    use super::{
+        IcebergWriter, IcebergWriterBuilder, IcebergWriterBuilderDynExt, IcebergWriterDynExt,
+    };
+    use crate::io::{FileIO, FileIOBuilder};
     use crate::spec::{DataFile, DataFileFormat};
+    use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+    use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
+    use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
+    use crate::writer::file_writer::ParquetWriterBuilder;
 
     // This function is used to guarantee the trait can be used as a object safe trait.
-    async fn _guarantee_object_safe(mut w: Box<dyn IcebergWriter>) {
-        let _ = w
+    async fn _guarantee_dyn_trait(builder: impl IcebergWriterBuilder) {
+        fn ensure_writer_builder<WB: IcebergWriterBuilder>(builder: WB) -> WB {
+            builder
+        }
+
+        fn ensure_writer<W: IcebergWriter>(writer: W) -> W {
+            writer
+        }
+
+        let writer = ensure_writer(builder.clone().build().await.unwrap());
+        let mut boxed_writer = ensure_writer(writer.boxed());
+        let _ = boxed_writer
+            .write(RecordBatch::new_empty(Schema::empty().into()))
+            .await;
+        let _ = boxed_writer.close().await;
+        let boxed_builder = ensure_writer_builder(builder.boxed());
+        let mut boxed_writer = ensure_writer(boxed_builder.clone().build().await.unwrap());
+
+        let _ = boxed_writer
             .write(RecordBatch::new_empty(Schema::empty().into()))
             .await;
-        let _ = w.close().await;
+        let _ = boxed_writer.close().await;
     }
 
     // This function check:
@@ -131,4 +259,25 @@ mod tests {
         let res = concat_batches(&batch.schema(), &batches).unwrap();
         assert_eq!(*batch, res);
     }
+
+    #[tokio::test]
+    async fn test_build_box_writer() {
+        let temp_dir = tempfile::TempDir::new().unwrap();
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+        let location_gen =
+            MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
+        let file_name_gen =
+            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
+
+        let pw = ParquetWriterBuilder::new(
+            WriterProperties::builder().build(),
+            Arc::new(crate::spec::Schema::builder().build().unwrap()),
+            file_io.clone(),
+            location_gen,
+            file_name_gen,
+        );
+        let data_file_builder = DataFileWriterBuilder::new(pw, None).boxed();
+
+        let _writer = data_file_builder.build().await.unwrap();
+    }
 }