[rust] Hash-distribute log-table appends by declared bucket key

fluss-rust/bindings/cpp/src/lib.rs

@@ -543,6 +543,8 @@ mod ffi {
         // AppendWriter
         unsafe fn delete_append_writer(writer: *mut AppendWriter);
         fn append(self: &mut AppendWriter, row: &GenericRowInner) -> FfiPtrResult;
+        // Partition (if partitioned) comes from the first row, so all rows must
+        // share one partition; rows are distributed across buckets by key.
         fn append_arrow_batch(
             self: &mut AppendWriter,
             array_ptr: usize,

fluss-rust/bindings/python/fluss/__init__.pyi

@@ -721,8 +721,21 @@ class AppendWriter:
             Use flush() to ensure all queued records are sent and acknowledged.
         """
         ...
-    def write_arrow(self, table: pa.Table) -> None: ...
-    def write_arrow_batch(self, batch: pa.RecordBatch) -> WriteResultHandle: ...
+    def write_arrow(self, table: pa.Table) -> None:
+        """Write a PyArrow Table (one batch at a time).
+
+        For a partitioned table, every batch must contain rows of a single
+        partition (see write_arrow_batch()).
+        """
+        ...
+    def write_arrow_batch(self, batch: pa.RecordBatch) -> WriteResultHandle:
+        """Write a PyArrow RecordBatch.
+
+        For a partitioned table the partition is derived from the first row, so
+        all rows must belong to the same partition. Rows are distributed across
+        buckets by their bucket key automatically.
+        """
+        ...
     def write_pandas(self, df: pd.DataFrame) -> None: ...
     async def flush(self) -> None: ...
     async def __aenter__(self) -> AppendWriter:

fluss-rust/bindings/python/src/table.rs

@@ -1032,7 +1032,10 @@ pub struct AppendWriter {
 
 #[pymethods]
 impl AppendWriter {
-    /// Write Arrow table data (fire-and-forget, use flush() to ensure delivery)
+    /// Write Arrow table data (fire-and-forget, use flush() to ensure delivery).
+    ///
+    /// For a partitioned table, every batch must contain rows of a single partition
+    /// (see `write_arrow_batch`).
     pub fn write_arrow(&self, py: Python, table: Py<PyAny>) -> PyResult<()> {
         // Convert Arrow Table to batches and write each batch
         let batches = table.call_method0(py, "to_batches")?;
@@ -1047,6 +1050,10 @@ impl AppendWriter {
 
     /// Write Arrow batch data.
     ///
+    /// For a partitioned table the partition is derived from the first row, so all
+    /// rows must belong to the same partition. Rows are distributed across buckets
+    /// by their bucket key automatically.
+    ///
     /// Returns:
     ///     WriteResultHandle that can be ignored (fire-and-forget) or
     ///     awaited via `handle.wait()` for server acknowledgment.

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

@@ -15,13 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::bucketing::BucketingFunction;
 use crate::client::table::partition_getter::{PartitionGetter, get_physical_path};
 use crate::client::{WriteRecord, WriteResultFuture, WriterClient};
-use crate::error::Error::IllegalArgument;
+use crate::error::Error::{IllegalArgument, UnexpectedError};
 use crate::error::Result;
 use crate::metadata::{PhysicalTablePath, TableInfo, TablePath};
+use crate::row::encode::{KeyEncoder, KeyEncoderFactory};
 use crate::row::{ColumnarRow, InternalRow};
-use arrow::array::RecordBatch;
+use arrow::array::{RecordBatch, UInt32Array};
+use bytes::Bytes;
+use parking_lot::Mutex;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 pub struct TableAppend {
@@ -53,20 +58,54 @@ impl TableAppend {
             None
         };
 
+        let bucket_router = if self.table_info.has_bucket_key() {
+            let data_lake_format = self.table_info.table_config.get_datalake_format()?;
+            let encoder = KeyEncoderFactory::of_bucket_key_encoder(
+                self.table_info.row_type(),
+                self.table_info.get_bucket_keys(),
+                &data_lake_format,
+            )?;
+            Some(BucketRouter {
+                encoder: Mutex::new(encoder),
+                bucketing: <dyn BucketingFunction>::of(data_lake_format.as_ref()),
+            })
+        } else {
+            None
+        };
+
         Ok(AppendWriter {
             table_path: Arc::clone(&self.table_path),
             partition_getter,
             writer_client: self.writer_client.clone(),
             table_info: Arc::clone(&self.table_info),
+            bucket_router,
         })
     }
 }
 
+struct BucketRouter {
+    encoder: Mutex<Box<dyn KeyEncoder>>,
+    bucketing: Box<dyn BucketingFunction>,
+}
+
+impl BucketRouter {
+    fn encode_key(&self, row: &dyn InternalRow) -> Result<Bytes> {
+        self.encoder.lock().encode_key(row)
+    }
+
+    fn bucket_of(&self, row: &dyn InternalRow, num_buckets: i32) -> Result<(i32, Bytes)> {
+        let key = self.encode_key(row)?;
+        let bucket = self.bucketing.bucketing(&key, num_buckets)?;
+        Ok((bucket, key))
+    }
+}
+
 pub struct AppendWriter {
     table_path: Arc<TablePath>,
     partition_getter: Option<PartitionGetter>,
     writer_client: Arc<WriterClient>,
     table_info: Arc<TableInfo>,
+    bucket_router: Option<BucketRouter>,
 }
 
 impl AppendWriter {
@@ -103,12 +142,17 @@ impl AppendWriter {
             self.partition_getter.as_ref(),
             row,
         )?);
+        let bucket_key = match &self.bucket_router {
+            Some(router) => Some(router.encode_key(row)?),
+            None => None,
+        };
         let record = WriteRecord::for_append(
             Arc::clone(&self.table_info),
             physical_table_path,
             self.table_info.schema_id,
             row,
-        );
+        )
+        .with_bucket_key(bucket_key);
         let result_handle = self.writer_client.send(&record)?;
         Ok(WriteResultFuture::new(result_handle))
     }
@@ -118,14 +162,19 @@ impl AppendWriter {
     /// This method returns a [`WriteResultFuture`] immediately after queueing the write,
     /// enabling fire-and-forget semantics for efficient batching.
     ///
-    /// For partitioned tables, the partition is derived from the **first row** of the batch.
-    /// Callers must ensure all rows in the batch belong to the same partition.
+    /// For a partitioned table the partition is derived from the **first row**, so all rows
+    /// in the batch must belong to the same partition. Rows are distributed across buckets by
+    /// their bucket key automatically.
     ///
     /// # Returns
     /// A [`WriteResultFuture`] that can be awaited to wait for server acknowledgment,
     /// or dropped for fire-and-forget behavior (use `flush()` to ensure delivery).
     pub fn append_arrow_batch(&self, batch: RecordBatch) -> Result<WriteResultFuture> {
-        let physical_table_path = if self.partition_getter.is_some() && batch.num_rows() > 0 {
+        if batch.num_rows() == 0 {
+            // Nothing to write; also avoids a keyless send to a bucket-key table.
+            return Ok(WriteResultFuture::join(Vec::new()));
+        }
+        let physical_table_path = if self.partition_getter.is_some() {
             let first_row = ColumnarRow::new(
                 Arc::new(batch.clone()),
                 Arc::new(self.table_info.row_type.clone()),
@@ -141,12 +190,60 @@ impl AppendWriter {
             Arc::new(PhysicalTablePath::of(Arc::clone(&self.table_path)))
         };
 
+        let Some(router) = self.bucket_router.as_ref() else {
+            return self.send_arrow_batch(batch, physical_table_path, None);
+        };
+
+        // Group rows by bucket, keeping one key per bucket (it hashes back there).
+        let num_buckets = self.table_info.get_num_buckets();
+        let batch_arc = Arc::new(batch.clone());
+        let row_type = Arc::new(self.table_info.row_type.clone());
+        let mut groups: HashMap<i32, (Vec<u32>, Bytes)> = HashMap::new();
+        // Reuse one row view; ColumnarRow::new re-downcasts every column.
+        let mut row = ColumnarRow::new(Arc::clone(&batch_arc), Arc::clone(&row_type), 0, None)?;
+        for i in 0..batch.num_rows() {
+            row.set_row_id(i);
+            let (bucket, key) = router.bucket_of(&row, num_buckets)?;
+            groups
+                .entry(bucket)
+                .or_insert_with(|| (Vec::new(), key))
+                .0
+                .push(i as u32);
+        }
+
+        if groups.len() == 1 {
+            let (_, (_, rep_key)) = groups.into_iter().next().unwrap();
+            return self.send_arrow_batch(batch, physical_table_path, Some(rep_key));
+        }
+
+        let mut handles = Vec::with_capacity(groups.len());
+        for (_bucket, (indices, rep_key)) in groups {
+            let sub_batch = take_rows(&batch, &indices)?;
+            let record = WriteRecord::for_append_record_batch(
+                Arc::clone(&self.table_info),
+                Arc::clone(&physical_table_path),
+                self.table_info.schema_id,
+                sub_batch,
+            )
+            .with_bucket_key(Some(rep_key));
+            handles.push(self.writer_client.send(&record)?);
+        }
+        Ok(WriteResultFuture::join(handles))
+    }
+
+    fn send_arrow_batch(
+        &self,
+        batch: RecordBatch,
+        physical_table_path: Arc<PhysicalTablePath>,
+        bucket_key: Option<Bytes>,
+    ) -> Result<WriteResultFuture> {
         let record = WriteRecord::for_append_record_batch(
             Arc::clone(&self.table_info),
             physical_table_path,
             self.table_info.schema_id,
             batch,
-        );
+        )
+        .with_bucket_key(bucket_key);
         let result_handle = self.writer_client.send(&record)?;
         Ok(WriteResultFuture::new(result_handle))
     }
@@ -155,3 +252,20 @@ impl AppendWriter {
         self.writer_client.flush().await
     }
 }
+
+fn take_rows(batch: &RecordBatch, indices: &[u32]) -> Result<RecordBatch> {
+    let idx = UInt32Array::from(indices.to_vec());
+    let columns = batch
+        .columns()
+        .iter()
+        .map(|column| arrow::compute::take(column, &idx, None))
+        .collect::<std::result::Result<Vec<_>, _>>()
+        .map_err(|e| UnexpectedError {
+            message: format!("Failed to split arrow batch by bucket: {e}"),
+            source: None,
+        })?;
+    RecordBatch::try_new(batch.schema(), columns).map_err(|e| UnexpectedError {
+        message: format!("Failed to rebuild split arrow batch: {e}"),
+        source: None,
+    })
+}

fluss-rust/crates/fluss/src/client/table/upsert.rs

@@ -16,17 +16,18 @@
 // under the License.
 
 use crate::client::{RowBytes, WriteFormat, WriteRecord, WriteResultFuture, WriterClient};
-use crate::error::Error::{IllegalArgument, UnexpectedError};
+use crate::error::Error::IllegalArgument;
 use crate::error::Result;
 use crate::metadata::{RowType, TableInfo, TablePath};
 use crate::row::InternalRow;
 use crate::row::encode::{KeyEncoder, KeyEncoderFactory, RowEncoder, RowEncoderFactory};
 use crate::row::field_getter::FieldGetter;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 use crate::client::table::partition_getter::{PartitionGetter, get_physical_path};
 use bitvec::prelude::bitvec;
 use bytes::Bytes;
+use parking_lot::Mutex;
 
 #[allow(dead_code)]
 pub struct TableUpsert {
@@ -307,34 +308,16 @@ impl UpsertWriter {
     }
 
     fn get_keys(&self, row: &dyn InternalRow) -> Result<(Bytes, Option<Bytes>)> {
-        let key = self
-            .primary_key_encoder
-            .lock()
-            .map_err(|e| UnexpectedError {
-                message: format!("primary_key_encoder lock poisoned: {e}"),
-                source: None,
-            })?
-            .encode_key(row)?;
+        let key = self.primary_key_encoder.lock().encode_key(row)?;
         let bucket_key = match &self.bucket_key_encoder {
-            Some(encoder) => Some(
-                encoder
-                    .lock()
-                    .map_err(|e| UnexpectedError {
-                        message: format!("bucket_key_encoder lock poisoned: {e}"),
-                        source: None,
-                    })?
-                    .encode_key(row)?,
-            ),
+            Some(encoder) => Some(encoder.lock().encode_key(row)?),
             None => Some(key.clone()),
         };
         Ok((key, bucket_key))
     }
 
     fn encode_row<R: InternalRow>(&self, row: &R) -> Result<Bytes> {
-        let mut encoder = self.row_encoder.lock().map_err(|e| UnexpectedError {
-            message: format!("row_encoder lock poisoned: {e}"),
-            source: None,
-        })?;
+        let mut encoder = self.row_encoder.lock();
         encoder.start_new_row()?;
         for (pos, field_getter) in self.field_getters.iter().enumerate() {
             let datum = field_getter.get_field(row)?;

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

@@ -162,6 +162,12 @@ impl<'a> WriteRecord<'a> {
         }
     }
 
+    /// Sets the bucket key used to hash-assign this record to a bucket.
+    pub fn with_bucket_key(mut self, bucket_key: Option<Bytes>) -> Self {
+        self.bucket_key = bucket_key;
+        self
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn for_upsert(
         table_info: Arc<TableInfo>,
@@ -254,6 +260,18 @@ impl WriteResultFuture {
             }),
         }
     }
+
+    pub fn join(handles: Vec<ResultHandle>) -> Self {
+        Self {
+            inner: Box::pin(async move {
+                for handle in handles {
+                    let result = handle.wait().await?;
+                    handle.result(result)?;
+                }
+                Ok(())
+            }),
+        }
+    }
 }
 
 impl Future for WriteResultFuture {

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

@@ -155,12 +155,8 @@ impl WriterClient {
             if let Some(assigner) = self.bucket_assigners.get(table_path) {
                 assigner.clone()
             } else {
-                let assigner = Self::create_bucket_assigner(
-                    table_info,
-                    Arc::clone(table_path),
-                    bucket_key,
-                    &self.config,
-                )?;
+                let assigner =
+                    Self::create_bucket_assigner(table_info, Arc::clone(table_path), &self.config)?;
                 self.bucket_assigners
                     .insert(Arc::clone(table_path), Arc::clone(&assigner));
                 assigner
@@ -232,10 +228,10 @@ impl WriterClient {
     pub fn create_bucket_assigner(
         table_info: &Arc<TableInfo>,
         table_path: Arc<PhysicalTablePath>,
-        bucket_key: Option<&Bytes>,
         config: &Config,
     ) -> Result<Arc<dyn BucketAssigner>> {
-        if bucket_key.is_some() {
+        // Decide from the table's bucket key, not an individual record's key.
+        if table_info.has_bucket_key() {
             let datalake_format = table_info.get_table_config().get_datalake_format()?;
             let function = <dyn BucketingFunction>::of(datalake_format.as_ref());
             Ok(Arc::new(HashBucketAssigner::new(