[feat] Support append arrow record batch (#34)

Author: luoyuxia

.github/workflows/ci.yml

@@ -91,7 +91,7 @@ jobs:
         # only run IT in linux since no docker in macos by default
         run: |
           if [ "$RUNNER_OS" == "Linux" ]; then
-            cargo test --features integration_tests --all-targets --workspace
+            RUST_TEST_THREADS=1 cargo test --features integration_tests --all-targets --workspace  -- --nocapture
           fi
         env:
           RUST_LOG: DEBUG

crates/fluss/src/client/table/append.rs

@@ -16,12 +16,12 @@
 // under the License.
 
 use crate::client::{WriteRecord, WriterClient};
+use crate::error::Result;
 use crate::metadata::{TableInfo, TablePath};
 use crate::row::GenericRow;
+use arrow::array::RecordBatch;
 use std::sync::Arc;
 
-use crate::error::Result;
-
 #[allow(dead_code)]
 pub struct TableAppend {
     table_path: TablePath,
@@ -63,6 +63,13 @@ impl AppendWriter {
         result_handle.result(result)
     }
 
+    pub async fn append_arrow_batch(&self, batch: RecordBatch) -> Result<()> {
+        let record = WriteRecord::new_record_batch(self.table_path.clone(), batch);
+        let result_handle = self.writer_client.send(&record).await?;
+        let result = result_handle.wait().await?;
+        result_handle.result(result)
+    }
+
     pub async fn flush(&self) -> Result<()> {
         self.writer_client.flush().await
     }

crates/fluss/src/client/write/accumulator.rs

@@ -17,7 +17,7 @@
 
 use crate::client::write::batch::WriteBatch::ArrowLog;
 use crate::client::write::batch::{ArrowLogWriteBatch, WriteBatch};
-use crate::client::{ResultHandle, WriteRecord};
+use crate::client::{Record, ResultHandle, WriteRecord};
 use crate::cluster::{BucketLocation, Cluster, ServerNode};
 use crate::config::Config;
 use crate::error::Result;
@@ -105,6 +105,7 @@ impl RecordAccumulator {
             row_type,
             bucket_id,
             current_time_ms(),
+            matches!(record.row, Record::RecordBatch(_)),
         ));
 
         let batch_id = batch.batch_id();
@@ -159,7 +160,6 @@ impl RecordAccumulator {
                 true, false, true,
             ));
         }
-
         self.append_new_batch(cluster, record, bucket_id, &mut dq_guard)
     }
 

crates/fluss/src/client/write/batch.rs

@@ -18,11 +18,10 @@
 use crate::BucketId;
 use crate::client::broadcast::{BatchWriteResult, BroadcastOnce};
 use crate::client::{ResultHandle, WriteRecord};
-use crate::metadata::{DataType, TablePath};
-use std::cmp::max;
-
 use crate::error::Result;
+use crate::metadata::{DataType, TablePath};
 use crate::record::MemoryLogRecordsArrowBuilder;
+use std::cmp::max;
 
 #[allow(dead_code)]
 pub struct InnerWriteBatch {
@@ -140,12 +139,16 @@ impl ArrowLogWriteBatch {
         row_type: &DataType,
         bucket_id: BucketId,
         create_ms: i64,
+        to_append_record_batch: bool,
     ) -> Self {
         let base = InnerWriteBatch::new(batch_id, table_path, create_ms, bucket_id);
-
         Self {
             write_batch: base,
-            arrow_builder: MemoryLogRecordsArrowBuilder::new(schema_id, row_type),
+            arrow_builder: MemoryLogRecordsArrowBuilder::new(
+                schema_id,
+                row_type,
+                to_append_record_batch,
+            ),
         }
     }
 
@@ -157,8 +160,13 @@ impl ArrowLogWriteBatch {
         if self.arrow_builder.is_closed() || self.arrow_builder.is_full() {
             Ok(None)
         } else {
-            self.arrow_builder.append(&write_record.row)?;
-            Ok(Some(ResultHandle::new(self.write_batch.results.receiver())))
+            // append successfully
+            if self.arrow_builder.append(write_record)? {
+                Ok(Some(ResultHandle::new(self.write_batch.results.receiver())))
+            } else {
+                // append fail
+                Ok(None)
+            }
         }
     }
 

crates/fluss/src/client/write/mod.rs

@@ -23,6 +23,7 @@ use crate::error::Error;
 use crate::metadata::TablePath;
 use crate::row::GenericRow;
 pub use accumulator::*;
+use arrow::array::RecordBatch;
 use std::sync::Arc;
 
 pub(crate) mod broadcast;
@@ -34,13 +35,28 @@ mod writer_client;
 pub use writer_client::WriterClient;
 
 pub struct WriteRecord<'a> {
-    pub row: GenericRow<'a>,
+    pub row: Record<'a>,
     pub table_path: Arc<TablePath>,
 }
 
+pub enum Record<'a> {
+    Row(GenericRow<'a>),
+    RecordBatch(Arc<RecordBatch>),
+}
+
 impl<'a> WriteRecord<'a> {
     pub fn new(table_path: Arc<TablePath>, row: GenericRow<'a>) -> Self {
-        Self { row, table_path }
+        Self {
+            row: Record::Row(row),
+            table_path,
+        }
+    }
+
+    pub fn new_record_batch(table_path: Arc<TablePath>, row: RecordBatch) -> Self {
+        Self {
+            row: Record::RecordBatch(Arc::new(row)),
+            table_path,
+        }
     }
 }
 

crates/fluss/src/client/write/sender.rs

@@ -122,7 +122,6 @@ impl Sender {
         collated: &HashMap<i32, Vec<Arc<ReadyWriteBatch>>>,
     ) -> Result<()> {
         for (leader_id, batches) in collated {
-            println!("send request batch");
             self.send_write_request(*leader_id, self.ack, batches)
                 .await?;
         }

crates/fluss/src/client/write/writer_client.rs

@@ -90,20 +90,12 @@ impl WriterClient {
         let table_path = &record.table_path;
         let cluster = self.metadata.get_cluster();
 
-        let bucket_assigner = {
-            if let Some(assigner) = self.bucket_assigners.get(table_path) {
-                assigner.clone()
-            } else {
-                let assigner = Arc::new(Self::create_bucket_assigner(table_path.as_ref()));
-                self.bucket_assigners
-                    .insert(table_path.as_ref().clone(), assigner.clone());
-                assigner
-            }
-        };
+        let (bucket_assigner, bucket_id) = self.assign_bucket(table_path);
 
-        let bucket_id = bucket_assigner.assign_bucket(None, &cluster);
-
-        let mut result = self.accumulate.append(record, 1, &cluster, true).await?;
+        let mut result = self
+            .accumulate
+            .append(record, bucket_id, &cluster, true)
+            .await?;
 
         if result.abort_record_for_new_batch {
             let prev_bucket_id = bucket_id;
@@ -121,6 +113,21 @@ impl WriterClient {
 
         Ok(result.result_handle.expect("result_handle should exist"))
     }
+    fn assign_bucket(&self, table_path: &Arc<TablePath>) -> (Arc<Box<dyn BucketAssigner>>, i32) {
+        let cluster = self.metadata.get_cluster();
+        let bucket_assigner = {
+            if let Some(assigner) = self.bucket_assigners.get(table_path) {
+                assigner.clone()
+            } else {
+                let assigner = Arc::new(Self::create_bucket_assigner(table_path.as_ref()));
+                self.bucket_assigners
+                    .insert(table_path.as_ref().clone(), assigner.clone());
+                assigner
+            }
+        };
+        let bucket_id = bucket_assigner.assign_bucket(None, &cluster);
+        (bucket_assigner, bucket_id)
+    }
 
     pub async fn close(self) -> Result<()> {
         self.shutdown_tx