diff --git a/fluss-rust/bindings/cpp/src/lib.rs b/fluss-rust/bindings/cpp/src/lib.rs index 3b59f950dc..675a28061c 100644 --- a/fluss-rust/bindings/cpp/src/lib.rs +++ b/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, diff --git a/fluss-rust/bindings/python/fluss/__init__.pyi b/fluss-rust/bindings/python/fluss/__init__.pyi index 9e7b8e75fb..8a2c27f1f9 100644 --- a/fluss-rust/bindings/python/fluss/__init__.pyi +++ b/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: diff --git a/fluss-rust/bindings/python/src/table.rs b/fluss-rust/bindings/python/src/table.rs index ca8c0556e1..ac050ff960 100644 --- a/fluss-rust/bindings/python/src/table.rs +++ b/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) -> 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. diff --git a/fluss-rust/crates/fluss/src/client/table/append.rs b/fluss-rust/crates/fluss/src/client/table/append.rs index dd54a5b119..b00c605210 100644 --- a/fluss-rust/crates/fluss/src/client/table/append.rs +++ b/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: ::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>, + bucketing: Box, +} + +impl BucketRouter { + fn encode_key(&self, row: &dyn InternalRow) -> Result { + 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, partition_getter: Option, writer_client: Arc, table_info: Arc, + bucket_router: Option, } 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 { - 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, 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, + bucket_key: Option, + ) -> Result { 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 { + let idx = UInt32Array::from(indices.to_vec()); + let columns = batch + .columns() + .iter() + .map(|column| arrow::compute::take(column, &idx, None)) + .collect::, _>>() + .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, + }) +} diff --git a/fluss-rust/crates/fluss/src/client/table/upsert.rs b/fluss-rust/crates/fluss/src/client/table/upsert.rs index 8a5547a7c7..28dce4ee79 100644 --- a/fluss-rust/crates/fluss/src/client/table/upsert.rs +++ b/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)> { - 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(&self, row: &R) -> Result { - 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)?; diff --git a/fluss-rust/crates/fluss/src/client/write/mod.rs b/fluss-rust/crates/fluss/src/client/write/mod.rs index a65b5d5af1..37bdbffe3f 100644 --- a/fluss-rust/crates/fluss/src/client/write/mod.rs +++ b/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) -> Self { + self.bucket_key = bucket_key; + self + } + #[allow(clippy::too_many_arguments)] pub fn for_upsert( table_info: Arc, @@ -254,6 +260,18 @@ impl WriteResultFuture { }), } } + + pub fn join(handles: Vec) -> Self { + Self { + inner: Box::pin(async move { + for handle in handles { + let result = handle.wait().await?; + handle.result(result)?; + } + Ok(()) + }), + } + } } impl Future for WriteResultFuture { diff --git a/fluss-rust/crates/fluss/src/client/write/writer_client.rs b/fluss-rust/crates/fluss/src/client/write/writer_client.rs index 53e1fd4877..05b2397b98 100644 --- a/fluss-rust/crates/fluss/src/client/write/writer_client.rs +++ b/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, table_path: Arc, - bucket_key: Option<&Bytes>, config: &Config, ) -> Result> { - 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 = ::of(datalake_format.as_ref()); Ok(Arc::new(HashBucketAssigner::new( diff --git a/fluss-rust/crates/fluss/tests/integration/log_table.rs b/fluss-rust/crates/fluss/tests/integration/log_table.rs index 2d88de45af..18d09c6689 100644 --- a/fluss-rust/crates/fluss/tests/integration/log_table.rs +++ b/fluss-rust/crates/fluss/tests/integration/log_table.rs @@ -23,7 +23,7 @@ mod table_test { create_table, dt_array_int, dt_map_string_int, dt_row_seq_label, extract_ids_from_batches, get_shared_cluster, make_int_array, make_string_array, map_dt_basics_columns, poll_until_count, poll_until_nonempty, row_dt_basics_columns, scalar_dt_columns, - wait_for_partitions_ready, wait_for_table_ready, + wait_for_partitions_ready, wait_for_table_buckets_ready, wait_for_table_ready, }; use arrow::array::record_batch; use fluss::client::{EARLIEST_OFFSET, FlussTable, TableScan}; @@ -2181,4 +2181,228 @@ mod table_test { .await .expect("Failed to drop table"); } + + #[tokio::test] + async fn append_hash_distributes_by_declared_bucket_key() { + let cluster = get_shared_cluster(); + let connection = cluster.get_fluss_connection().await; + let admin = connection.get_admin().expect("admin"); + + let table_path = TablePath::new("fluss", "test_log_append_bucket_key_spread"); + let descriptor = TableDescriptor::builder() + .schema( + Schema::builder() + .column("c1", DataTypes::int()) + .column("c2", DataTypes::string()) + .build() + .expect("schema"), + ) + .distributed_by(Some(3), vec!["c1".to_string()]) + .build() + .expect("descriptor"); + create_table(&admin, &table_path, &descriptor).await; + wait_for_table_buckets_ready(&admin, &table_path, &[0, 1, 2]).await; + + let table = connection.get_table(&table_path).await.expect("table"); + let writer = table + .new_append() + .expect("append") + .create_writer() + .expect("writer"); + + let row_count: i32 = 30; + for id in 1..=row_count { + let mut row = GenericRow::new(2); + row.set_field(0, id); + row.set_field(1, "x"); + writer.append(&row).expect("append row"); + } + writer.flush().await.expect("flush"); + + let offsets = admin + .list_offsets(&table_path, &[0, 1, 2], OffsetSpec::Latest) + .await + .expect("list offsets"); + + let total: i64 = offsets.values().sum(); + assert_eq!( + total, row_count as i64, + "every appended row must be persisted, got per-bucket offsets {offsets:?}" + ); + + let non_empty = offsets.values().filter(|&&o| o > 0).count(); + assert!( + non_empty >= 2, + "rows must hash-distribute across buckets, got per-bucket offsets {offsets:?}" + ); + + admin.drop_table(&table_path, false).await.expect("drop"); + } + + #[tokio::test] + async fn append_same_bucket_key_lands_in_one_bucket() { + let cluster = get_shared_cluster(); + let connection = cluster.get_fluss_connection().await; + let admin = connection.get_admin().expect("admin"); + + let table_path = TablePath::new("fluss", "test_log_append_bucket_key_colocate"); + let descriptor = TableDescriptor::builder() + .schema( + Schema::builder() + .column("c1", DataTypes::int()) + .column("c2", DataTypes::string()) + .build() + .expect("schema"), + ) + .distributed_by(Some(3), vec!["c1".to_string()]) + .build() + .expect("descriptor"); + create_table(&admin, &table_path, &descriptor).await; + wait_for_table_buckets_ready(&admin, &table_path, &[0, 1, 2]).await; + + let table = connection.get_table(&table_path).await.expect("table"); + let writer = table + .new_append() + .expect("append") + .create_writer() + .expect("writer"); + + for _ in 0..5 { + let mut row = GenericRow::new(2); + row.set_field(0, 42); + row.set_field(1, "x"); + writer.append(&row).expect("append row"); + } + writer.flush().await.expect("flush"); + + let offsets = admin + .list_offsets(&table_path, &[0, 1, 2], OffsetSpec::Latest) + .await + .expect("list offsets"); + + let non_empty: Vec<_> = offsets.iter().filter(|&(_, &o)| o > 0).collect(); + assert_eq!( + non_empty.len(), + 1, + "all rows with the same bucket key must land in one bucket, got {offsets:?}" + ); + assert_eq!( + *non_empty[0].1, 5, + "that bucket must hold all five rows, got {offsets:?}" + ); + + admin.drop_table(&table_path, false).await.expect("drop"); + } + + #[tokio::test] + async fn append_arrow_batch_splits_mixed_keys_across_buckets() { + let cluster = get_shared_cluster(); + let connection = cluster.get_fluss_connection().await; + let admin = connection.get_admin().expect("admin"); + + let table_path = TablePath::new("fluss", "test_log_append_arrow_split"); + let descriptor = TableDescriptor::builder() + .schema( + Schema::builder() + .column("c1", DataTypes::int()) + .column("c2", DataTypes::string()) + .build() + .expect("schema"), + ) + .distributed_by(Some(3), vec!["c1".to_string()]) + .build() + .expect("descriptor"); + create_table(&admin, &table_path, &descriptor).await; + wait_for_table_buckets_ready(&admin, &table_path, &[0, 1, 2]).await; + + let table = connection.get_table(&table_path).await.expect("table"); + let writer = table + .new_append() + .expect("append") + .create_writer() + .expect("writer"); + + let batch = record_batch!( + ("c1", Int32, [1, 2, 3, 4, 5, 6, 7, 8, 9]), + ("c2", Utf8, ["a", "b", "c", "d", "e", "f", "g", "h", "i"]) + ) + .unwrap(); + writer.append_arrow_batch(batch).expect("append batch"); + writer.flush().await.expect("flush"); + + let offsets = admin + .list_offsets(&table_path, &[0, 1, 2], OffsetSpec::Latest) + .await + .expect("list offsets"); + + let total: i64 = offsets.values().sum(); + assert_eq!(total, 9, "all rows must be persisted, got {offsets:?}"); + let non_empty = offsets.values().filter(|&&o| o > 0).count(); + assert!( + non_empty >= 2, + "a mixed-key batch must be split across buckets, got {offsets:?}" + ); + + admin.drop_table(&table_path, false).await.expect("drop"); + } + + /// An empty first batch must not lock in a sticky assigner and break later keyed appends. + #[tokio::test] + async fn append_empty_batch_first_keeps_key_distribution() { + let cluster = get_shared_cluster(); + let connection = cluster.get_fluss_connection().await; + let admin = connection.get_admin().expect("admin"); + + let table_path = TablePath::new("fluss", "test_log_append_empty_first"); + let descriptor = TableDescriptor::builder() + .schema( + Schema::builder() + .column("c1", DataTypes::int()) + .column("c2", DataTypes::string()) + .build() + .expect("schema"), + ) + .distributed_by(Some(3), vec!["c1".to_string()]) + .build() + .expect("descriptor"); + create_table(&admin, &table_path, &descriptor).await; + wait_for_table_buckets_ready(&admin, &table_path, &[0, 1, 2]).await; + + let table = connection.get_table(&table_path).await.expect("table"); + let writer = table + .new_append() + .expect("append") + .create_writer() + .expect("writer"); + + let empty = record_batch!(("c1", Int32, [1]), ("c2", Utf8, ["x"])) + .unwrap() + .slice(0, 0); + writer + .append_arrow_batch(empty) + .expect("append empty batch"); + + let batch = record_batch!( + ("c1", Int32, [1, 2, 3, 4, 5, 6, 7, 8, 9]), + ("c2", Utf8, ["a", "b", "c", "d", "e", "f", "g", "h", "i"]) + ) + .unwrap(); + writer.append_arrow_batch(batch).expect("append batch"); + writer.flush().await.expect("flush"); + + let offsets = admin + .list_offsets(&table_path, &[0, 1, 2], OffsetSpec::Latest) + .await + .expect("list offsets"); + + let total: i64 = offsets.values().sum(); + assert_eq!(total, 9, "all rows must be persisted, got {offsets:?}"); + let non_empty = offsets.values().filter(|&&o| o > 0).count(); + assert!( + non_empty >= 2, + "keyed rows must still spread across buckets after an empty first batch, got {offsets:?}" + ); + + admin.drop_table(&table_path, false).await.expect("drop"); + } }