Implement BaseTaskWriter

Author: CTTY

crates/iceberg/src/arrow/mod.rs

@@ -34,4 +34,6 @@ mod value;
 
 pub use reader::*;
 pub use value::*;
-pub(crate) mod record_batch_partition_splitter;
+/// Record batch partition splitter for partitioned tables
+pub mod record_batch_partition_splitter;
+pub use record_batch_partition_splitter::*;

crates/iceberg/src/arrow/record_batch_partition_splitter.rs

@@ -51,6 +51,17 @@ pub struct RecordBatchPartitionSplitter {
 // Remove this after partition writer supported.
 #[allow(dead_code)]
 impl RecordBatchPartitionSplitter {
+    /// Create a new RecordBatchPartitionSplitter.
+    ///
+    /// # Arguments
+    ///
+    /// * `input_schema` - The Arrow schema of the input record batches
+    /// * `iceberg_schema` - The Iceberg schema reference
+    /// * `partition_spec` - The partition specification reference
+    ///
+    /// # Returns
+    ///
+    /// Returns a new `RecordBatchPartitionSplitter` instance or an error if initialization fails.
     pub fn new(
         input_schema: ArrowSchemaRef,
         iceberg_schema: SchemaRef,

crates/iceberg/src/writer/task/mod.rs

@@ -21,6 +21,9 @@
 //! for writing data to Iceberg tables with automatic partition handling.
 
 use crate::Result;
+use crate::arrow::RecordBatchPartitionSplitter;
+use crate::spec::{PartitionKey, PartitionSpecRef, SchemaRef, Struct};
+use crate::writer::partitioning::PartitioningWriter;
 use crate::writer::{DefaultInput, DefaultOutput};
 
 /// High-level async trait for writing tasks to Iceberg tables.
@@ -53,3 +56,97 @@ pub trait TaskWriter<I = DefaultInput, O = DefaultOutput>: Send + 'static {
     /// or an error if the close operation fails.
     async fn close(self) -> Result<O>;
 }
+
+/// A high-level writer implementation for writing data to Iceberg tables.
+///
+/// `BaseTaskWriter` handles both partitioned and non-partitioned tables by composing
+/// a [`PartitioningWriter`] with an optional [`RecordBatchPartitionSplitter`].
+///
+/// # Type Parameters
+///
+/// * `W` - The underlying [`PartitioningWriter`] implementation
+pub struct BaseTaskWriter<W: PartitioningWriter> {
+    writer: W,
+    partition_splitter: Option<RecordBatchPartitionSplitter>,
+    schema: SchemaRef,
+    partition_spec: PartitionSpecRef,
+}
+
+impl<W: PartitioningWriter> BaseTaskWriter<W> {
+    /// Create a new BaseTaskWriter.
+    ///
+    /// # Parameters
+    ///
+    /// * `writer` - The underlying [`PartitioningWriter`] implementation
+    /// * `partition_splitter` - Optional partition splitter for partitioned tables.
+    ///   Should be `None` for unpartitioned tables and `Some` for partitioned tables.
+    /// * `schema` - The Iceberg schema reference
+    /// * `partition_spec` - The partition specification reference
+    ///
+    /// # Returns
+    ///
+    /// Returns a new `BaseTaskWriter` instance, or an error if the partition splitter
+    /// configuration is invalid (e.g., missing splitter for a partitioned table).
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    /// - A partitioned table is provided without a partition splitter
+    pub fn new(
+        writer: W,
+        partition_splitter: Option<RecordBatchPartitionSplitter>,
+        schema: SchemaRef,
+        partition_spec: PartitionSpecRef,
+    ) -> Result<Self> {
+        // Validate that partitioned tables have a splitter
+        if !partition_spec.is_unpartitioned() && partition_splitter.is_none() {
+            return Err(crate::Error::new(
+                crate::ErrorKind::DataInvalid,
+                "Partition splitter is required for partitioned tables",
+            ));
+        }
+
+        Ok(Self {
+            writer,
+            partition_splitter,
+            schema,
+            partition_spec,
+        })
+    }
+}
+
+#[async_trait::async_trait]
+impl<W: PartitioningWriter> TaskWriter for BaseTaskWriter<W> {
+    async fn write(&mut self, input: DefaultInput) -> Result<()> {
+        if self.partition_spec.is_unpartitioned() {
+            // Unpartitioned table: create empty PartitionKey and write directly
+            let partition_key = PartitionKey::new(
+                self.partition_spec.as_ref().clone(),
+                self.schema.clone(),
+                Struct::empty(),
+            );
+            self.writer.write(partition_key, input).await?;
+        } else {
+            // Partitioned table: must have a splitter
+            let splitter = self.partition_splitter.as_ref().ok_or_else(|| {
+                crate::Error::new(
+                    crate::ErrorKind::DataInvalid,
+                    "Partition splitter is required for partitioned tables",
+                )
+            })?;
+
+            // Split batch and write each partition
+            let partitioned_batches = splitter.split(&input)?;
+
+            for (partition_key, batch) in partitioned_batches {
+                self.writer.write(partition_key, batch).await?;
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn close(self) -> Result<DefaultOutput> {
+        self.writer.close().await
+    }
+}