diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index d29c7eb85..234f72563 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -73,6 +73,42 @@ fn not_null_table_schema() -> TableSchema { TableSchema::new(0, &schema) } +fn postpone_table_schema() -> TableSchema { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .build() + .unwrap(); + TableSchema::new(0, &schema) +} + +fn legacy_postpone_table_schema() -> TableSchema { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.batch-write-fixed-bucket", "false") + .build() + .unwrap(); + TableSchema::new(0, &schema) +} + +fn partitioned_postpone_table_schema() -> TableSchema { + let schema = Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .primary_key(["pt", "id"]) + .partition_keys(["pt"]) + .option("bucket", "-2") + .build() + .unwrap(); + TableSchema::new(0, &schema) +} + unsafe fn wrap_table(table: Table) -> *mut paimon_table { let inner = Box::into_raw(Box::new(table)) as *mut c_void; Box::into_raw(Box::new(paimon_table { inner })) @@ -104,6 +140,38 @@ fn make_batch(ids: Vec, names: Vec<&str>) -> RecordBatch { .unwrap() } +fn make_partitioned_write_batch(pts: Vec<&str>, ids: Vec, names: Vec<&str>) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("pt", ArrowDataType::Utf8, false), + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(pts)), + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + ], + ) + .unwrap() +} + +fn make_postpone_bucket_plan_batch(partitions: Vec<&str>, counts: Vec) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("pt", ArrowDataType::Utf8, false), + ArrowField::new("total_buckets", ArrowDataType::Int32, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(partitions)), + Arc::new(Int32Array::from(counts)), + ], + ) + .unwrap() +} + fn make_type_mismatch_batch(ids: Vec<&str>, names: Vec<&str>) -> RecordBatch { let schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new("id", ArrowDataType::Utf8, false), @@ -950,6 +1018,93 @@ fn test_write_new_builder_and_free() { } } +#[test] +fn test_postpone_fixed_bucket_builder_respects_option() { + let path = "memory:/test_postpone_fixed_bucket_builder_option"; + let file_io = memory_file_io(); + setup_table_dirs(&file_io, path); + let table = Table::new( + file_io, + Identifier::new("default", "test"), + path.to_string(), + postpone_table_schema(), + None, + ); + let handle = unsafe { wrap_table(table) }; + + unsafe { + let normal = paimon_table_new_write_builder(handle); + assert!(normal.error.is_null()); + let normal_state = &*((*normal.write_builder).inner as *const WriteBuilderState); + assert!(normal_state.postpone_fixed_bucket); + + let write = paimon_write_builder_new_write(normal.write_builder); + assert!(write.error.is_null()); + let (array, schema) = export_batch_to_ffi(make_batch(vec![1], vec!["a"])); + let error = paimon_table_write_write_arrow_batch( + write.write, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ); + assert!(error.is_null()); + let prepared = paimon_table_write_prepare_commit(write.write); + assert!(prepared.error.is_null()); + let messages = &*((*prepared.messages).inner as *const CommitMessagesState); + assert!(messages.messages.iter().all(|message| message.bucket >= 0)); + assert!(messages + .messages + .iter() + .all(|message| message.total_buckets == Some(1))); + paimon_commit_messages_free(prepared.messages); + paimon_table_write_free(write.write); + paimon_write_builder_free(normal.write_builder); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(handle); + assert!(fixed.error.is_null()); + let fixed_state = &*((*fixed.write_builder).inner as *const WriteBuilderState); + assert!(fixed_state.postpone_fixed_bucket); + paimon_write_builder_free(fixed.write_builder); + + let commit_user = CString::new("fixed-user").unwrap(); + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ); + assert!(fixed.error.is_null()); + let fixed_state = &*((*fixed.write_builder).inner as *const WriteBuilderState); + assert!(fixed_state.postpone_fixed_bucket); + assert_eq!(fixed_state.commit_user, "fixed-user"); + paimon_write_builder_free(fixed.write_builder); + unwrap_table(handle); + } + + let legacy_path = "memory:/test_legacy_postpone_builder_option"; + let file_io = memory_file_io(); + setup_table_dirs(&file_io, legacy_path); + let legacy_table = Table::new( + file_io, + Identifier::new("default", "test"), + legacy_path.to_string(), + legacy_postpone_table_schema(), + None, + ); + let legacy_handle = unsafe { wrap_table(legacy_table) }; + unsafe { + let normal = paimon_table_new_write_builder(legacy_handle); + assert!(normal.error.is_null()); + let normal_state = &*((*normal.write_builder).inner as *const WriteBuilderState); + assert!(!normal_state.postpone_fixed_bucket); + paimon_write_builder_free(normal.write_builder); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(legacy_handle); + assert!(fixed.error.is_null()); + let fixed_state = &*((*fixed.write_builder).inner as *const WriteBuilderState); + assert!(fixed_state.postpone_fixed_bucket); + paimon_write_builder_free(fixed.write_builder); + unwrap_table(legacy_handle); + } +} + #[test] fn test_write_commit_read_roundtrip() { let path = "memory:/test_write_roundtrip"; @@ -1442,6 +1597,90 @@ fn test_commit_messages_merge_preserves_all_writer_files() { } } +#[test] +fn test_distributed_postpone_writers_share_bucket_plan() { + let path = "memory:/test_distributed_postpone_bucket_plan"; + let file_io = memory_file_io(); + setup_table_dirs(&file_io, path); + let table = Table::new( + file_io, + Identifier::new("default", "test"), + path.to_string(), + partitioned_postpone_table_schema(), + None, + ); + let handle = unsafe { wrap_table(table) }; + let commit_user = CString::new("distributed-postpone-job").unwrap(); + + unsafe { + let wb1 = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ) + .write_builder; + let wb2 = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ) + .write_builder; + + for wb in [wb1, wb2] { + let (array, schema) = + export_batch_to_ffi(make_postpone_bucket_plan_batch(vec!["p"], vec![3])); + let error = paimon_write_builder_with_postpone_bucket_plan( + wb, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ); + assert!(error.is_null()); + } + + let tw1 = paimon_write_builder_new_write(wb1).write; + let tw2 = paimon_write_builder_new_write(wb2).write; + for (tw, ids, names) in [ + (tw1, vec![1], vec!["a"]), + (tw2, vec![2, 3, 4, 5], vec!["b", "c", "d", "e"]), + ] { + let (array, schema) = export_batch_to_ffi(make_partitioned_write_batch( + vec!["p"; ids.len()], + ids, + names, + )); + let error = paimon_table_write_write_arrow_batch( + tw, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ); + assert!(error.is_null()); + } + + let messages1 = paimon_table_write_prepare_commit(tw1).messages; + let messages2 = paimon_table_write_prepare_commit(tw2).messages; + for messages in [messages1, messages2] { + let state = &*((*messages).inner as *const CommitMessagesState); + assert!(!state.messages.is_empty()); + assert!(state + .messages + .iter() + .all(|message| message.total_buckets == Some(3))); + } + let error = paimon_commit_messages_merge(messages1, messages2); + assert!(error.is_null()); + let commit = paimon_write_builder_new_commit(wb1).commit; + let error = paimon_table_commit_commit(commit, messages1); + assert!(error.is_null()); + + paimon_table_commit_free(commit); + paimon_commit_messages_free(messages2); + paimon_commit_messages_free(messages1); + paimon_table_write_free(tw2); + paimon_table_write_free(tw1); + paimon_write_builder_free(wb2); + paimon_write_builder_free(wb1); + unwrap_table(handle); + } +} + #[test] fn test_write_multiple_batches() { let path = "memory:/test_write_multi_batch"; @@ -1721,6 +1960,11 @@ fn test_null_pointer_handling() { assert!(result.write_builder.is_null()); paimon_error_free(result.error); + let result = paimon_table_new_postpone_fixed_bucket_write_builder(ptr::null()); + assert!(!result.error.is_null()); + assert!(result.write_builder.is_null()); + paimon_error_free(result.error); + let result = paimon_write_builder_new_write(ptr::null()); assert!(!result.error.is_null()); assert!(result.write.is_null()); @@ -1731,6 +1975,14 @@ fn test_null_pointer_handling() { assert!(result.commit.is_null()); paimon_error_free(result.error); + let err = paimon_write_builder_with_postpone_bucket_plan( + ptr::null_mut(), + ptr::null_mut(), + ptr::null_mut(), + ); + assert!(!err.is_null()); + paimon_error_free(err); + let err = paimon_table_write_write_arrow_batch(ptr::null_mut(), ptr::null_mut(), ptr::null_mut()); assert!(!err.is_null()); diff --git a/bindings/c/src/types.rs b/bindings/c/src/types.rs index 1dfd9c3f2..b5a2422c5 100644 --- a/bindings/c/src/types.rs +++ b/bindings/c/src/types.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow_schema::Schema as ArrowSchema; use paimon::spec::{DataField, Predicate}; -use paimon::table::{CommitMessage, Table, TableCommit, TableWrite}; +use paimon::table::{CommitMessage, PostponeBucketPlan, Table, TableCommit, TableWrite}; /// C-compatible key-value pair for options. #[repr(C)] @@ -210,6 +210,8 @@ pub(crate) struct WriteBuilderState { pub table: Table, pub commit_user: String, pub overwrite: bool, + pub postpone_fixed_bucket: bool, + pub postpone_bucket_plan: Option, } pub(crate) struct TableWriteState { diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs index 8e9637997..57f6ab6a8 100644 --- a/bindings/c/src/write.rs +++ b/bindings/c/src/write.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use arrow_array::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; use arrow_array::{Array, RecordBatch, RecordBatchOptions, StructArray}; use arrow_schema::{DataType as ArrowDataType, Schema as ArrowSchema}; -use paimon::table::Table; +use paimon::table::{PostponeBucketPlan, Table}; use crate::error::{check_non_null, paimon_error, validate_cstr, PaimonErrorCode}; use crate::result::{ @@ -37,6 +37,7 @@ use crate::types::*; unsafe fn new_write_builder( table: *const paimon_table, commit_user: Option, + force_postpone_fixed_bucket: bool, ) -> paimon_result_write_builder { if let Err(e) = check_non_null(table, "table") { return paimon_result_write_builder { @@ -45,7 +46,20 @@ unsafe fn new_write_builder( }; } let table_ref = &*((*table).inner as *const Table); - let builder = table_ref.new_write_builder(); + let builder = if force_postpone_fixed_bucket { + match table_ref.new_postpone_fixed_bucket_write_builder() { + Ok(builder) => builder, + Err(e) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error: paimon_error::from_paimon(e), + } + } + } + } else { + table_ref.new_write_builder() + }; + let postpone_fixed_bucket = builder.uses_postpone_fixed_bucket(); let commit_user = match commit_user { Some(commit_user) => match builder.with_commit_user(commit_user) { Ok(builder) => builder.commit_user().to_string(), @@ -62,6 +76,8 @@ unsafe fn new_write_builder( table: table_ref.clone(), commit_user, overwrite: false, + postpone_fixed_bucket, + postpone_bucket_plan: None, }; let inner = Box::into_raw(Box::new(state)) as *mut c_void; paimon_result_write_builder { @@ -82,7 +98,19 @@ unsafe fn new_write_builder( pub unsafe extern "C" fn paimon_table_new_write_builder( table: *const paimon_table, ) -> paimon_result_write_builder { - new_write_builder(table, None) + new_write_builder(table, None, false) +} + +/// Create a WriteBuilder which forces one-shot fixed-bucket writes for a +/// postpone table, even when `postpone.batch-write-fixed-bucket=false`. +/// +/// # Safety +/// `table` must be a valid table pointer, or null (returns error). +#[no_mangle] +pub unsafe extern "C" fn paimon_table_new_postpone_fixed_bucket_write_builder( + table: *const paimon_table, +) -> paimon_result_write_builder { + new_write_builder(table, None, true) } /// Create a WriteBuilder with a caller-provided stable commit identity. @@ -107,7 +135,29 @@ pub unsafe extern "C" fn paimon_table_new_write_builder_with_commit_user( } } }; - new_write_builder(table, Some(commit_user)) + new_write_builder(table, Some(commit_user), false) +} + +/// Create an explicit postpone fixed-bucket WriteBuilder with a stable commit identity. +/// +/// # Safety +/// `table` must be a valid table pointer. `commit_user` must be a valid UTF-8 +/// C string and a safe file-name segment. +#[no_mangle] +pub unsafe extern "C" fn paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + table: *const paimon_table, + commit_user: *const c_char, +) -> paimon_result_write_builder { + let commit_user = match validate_cstr(commit_user, "commit_user") { + Ok(commit_user) => commit_user, + Err(error) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error, + } + } + }; + new_write_builder(table, Some(commit_user), true) } /// Free a paimon_write_builder. @@ -143,6 +193,55 @@ pub unsafe extern "C" fn paimon_write_builder_with_overwrite( ptr::null_mut() } +/// Supply a shared bucket plan for distributed postpone fixed-bucket writers. +/// +/// The Arrow batch must contain the table's partition columns in partition-key +/// order followed by a non-null Int32 column named `total_buckets`. For an +/// unpartitioned table, only the `total_buckets` column is present. Ownership +/// of the Arrow C Data structs is transferred to this function. +/// +/// Every writer whose commit messages will be merged must receive the same +/// plan, and the plan must contain every partition written by those writers. +/// +/// # Safety +/// `wb` must select postpone fixed-bucket mode through table configuration or +/// an explicit fixed-bucket constructor. `array` and `schema` must point to +/// initialized Arrow C Data structs. +#[no_mangle] +pub unsafe extern "C" fn paimon_write_builder_with_postpone_bucket_plan( + wb: *mut paimon_write_builder, + array: *mut c_void, + schema: *mut c_void, +) -> *mut paimon_error { + if let Err(error) = check_non_null(wb, "wb") { + return error; + } + if let Err(error) = check_non_null(array, "array") { + return error; + } + if let Err(error) = check_non_null(schema, "schema") { + return error; + } + let state = &mut *((*wb).inner as *mut WriteBuilderState); + if !state.postpone_fixed_bucket { + return invalid_input( + "a postpone bucket plan requires a postpone fixed-bucket write builder", + ); + } + + let batch = match import_record_batch(array, schema) { + Ok(batch) => batch, + Err(error) => return error, + }; + match PostponeBucketPlan::from_arrow(&state.table, &batch) { + Ok(plan) => { + state.postpone_bucket_plan = Some(plan); + ptr::null_mut() + } + Err(error) => paimon_error::from_paimon(error), + } +} + // ======================= TableWrite =============================== fn invalid_input(message: impl Into) -> *mut paimon_error { @@ -245,11 +344,17 @@ pub unsafe extern "C" fn paimon_write_builder_new_write( } let state = &*((*wb).inner as *const WriteBuilderState); - let mut builder = match state - .table - .new_write_builder() - .with_commit_user(state.commit_user.clone()) - { + let builder = if state.postpone_fixed_bucket { + state.table.new_postpone_fixed_bucket_write_builder() + } else { + Ok(state.table.new_write_builder()) + }; + let mut builder = match builder + .and_then(|builder| builder.with_commit_user(state.commit_user.clone())) + .and_then(|builder| match state.postpone_bucket_plan.clone() { + Some(plan) => builder.with_postpone_bucket_plan(plan), + None => Ok(builder), + }) { Ok(b) => b, Err(e) => { return paimon_result_table_write { @@ -358,8 +463,9 @@ pub unsafe extern "C" fn paimon_table_write_write_arrow_batch( /// Close file writers and produce CommitMessages. /// /// Consumes the open file writers (they are flushed and closed). After this -/// call, the TableWrite can be reused — `write_arrow_batch` may be called -/// again to start a new round of writes. +/// call, the TableWrite can normally be reused — `write_arrow_batch` may be +/// called again to start a new round of writes. Fixed-bucket postpone batch +/// writers are one-shot; create a new TableWrite for the next batch. /// /// The returned `paimon_commit_messages` must be passed to a /// `paimon_table_commit_*` function and then freed with @@ -781,12 +887,21 @@ pub unsafe extern "C" fn paimon_table_commit_abort( const _: unsafe extern "C" fn(*const paimon_table) -> paimon_result_write_builder = paimon_table_new_write_builder; +const _: unsafe extern "C" fn(*const paimon_table) -> paimon_result_write_builder = + paimon_table_new_postpone_fixed_bucket_write_builder; const _: unsafe extern "C" fn(*const paimon_table, *const c_char) -> paimon_result_write_builder = paimon_table_new_write_builder_with_commit_user; +const _: unsafe extern "C" fn(*const paimon_table, *const c_char) -> paimon_result_write_builder = + paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user; const _: unsafe extern "C" fn(*const paimon_write_builder) -> paimon_result_table_write = paimon_write_builder_new_write; const _: unsafe extern "C" fn(*const paimon_write_builder) -> paimon_result_table_commit = paimon_write_builder_new_commit; +const _: unsafe extern "C" fn( + *mut paimon_write_builder, + *mut c_void, + *mut c_void, +) -> *mut paimon_error = paimon_write_builder_with_postpone_bucket_plan; const _: unsafe extern "C" fn(*mut paimon_table_write) -> paimon_result_prepare_commit = paimon_table_write_prepare_commit; const _: unsafe extern "C" fn( diff --git a/crates/integrations/datafusion/tests/pk_tables.rs b/crates/integrations/datafusion/tests/pk_tables.rs index 0f44c8197..917ee229d 100644 --- a/crates/integrations/datafusion/tests/pk_tables.rs +++ b/crates/integrations/datafusion/tests/pk_tables.rs @@ -1941,9 +1941,9 @@ async fn test_pk_first_row_insert_overwrite() { // ======================= Postpone Bucket (bucket = -2) ======================= -/// Postpone bucket files are invisible to normal SELECT but visible via scan_all_files. +/// Batch writes use real fixed buckets by default and are immediately visible. #[tokio::test] -async fn test_postpone_write_invisible_to_select() { +async fn test_postpone_batch_write_uses_visible_fixed_bucket() { let (_tmp, catalog) = create_test_env(); let sql_context = create_sql_context(catalog.clone()).await; sql_context @@ -1970,7 +1970,7 @@ async fn test_postpone_write_invisible_to_select() { .await .unwrap(); - // scan_all_files should find the postpone file + // The default batch path must write a real bucket rather than bucket -2. let table = catalog .get_table(&Identifier::new("test_db", "t_postpone")) .await @@ -1983,11 +1983,68 @@ async fn test_postpone_write_invisible_to_select() { .await .unwrap(); let file_count: usize = plan.splits().iter().map(|s| s.data_files().len()).sum(); - assert_eq!(file_count, 1, "scan_all_files should find 1 postpone file"); + assert_eq!( + file_count, 1, + "scan_all_files should find 1 fixed-bucket file" + ); + assert!(plan.splits().iter().all(|split| split.bucket() >= 0)); - // Normal SELECT should return 0 rows (postpone files are invisible) + // Real buckets are visible to the normal read path immediately. let count = row_count(&sql_context, "SELECT * FROM paimon.test_db.t_postpone").await; - assert_eq!(count, 0, "SELECT should return 0 rows for postpone table"); + assert_eq!(count, 3); +} + +/// The compatibility switch retains legacy invisible bucket -2 writes. +#[tokio::test] +async fn test_postpone_fixed_bucket_write_can_be_disabled() { + let (_tmp, catalog) = create_test_env(); + let sql_context = create_sql_context(catalog.clone()).await; + sql_context + .sql("CREATE SCHEMA paimon.test_db") + .await + .expect("CREATE SCHEMA failed"); + + sql_context + .sql( + "CREATE TABLE paimon.test_db.t_postpone_legacy ( + id INT NOT NULL, value INT, + PRIMARY KEY (id) + ) WITH ( + 'bucket' = '-2', + 'postpone.batch-write-fixed-bucket' = 'false' + )", + ) + .await + .unwrap(); + sql_context + .sql("INSERT INTO paimon.test_db.t_postpone_legacy VALUES (1, 10)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let table = catalog + .get_table(&Identifier::new("test_db", "t_postpone_legacy")) + .await + .unwrap(); + let plan = table + .new_read_builder() + .new_scan() + .with_scan_all_files() + .plan() + .await + .unwrap(); + assert_eq!(plan.splits().len(), 1); + assert_eq!(plan.splits()[0].bucket(), -2); + assert_eq!( + row_count( + &sql_context, + "SELECT * FROM paimon.test_db.t_postpone_legacy", + ) + .await, + 0 + ); } /// INSERT OVERWRITE on a postpone table should replace old files with new ones. @@ -2031,7 +2088,7 @@ async fn test_postpone_insert_overwrite() { .await .unwrap(); let file_count: usize = plan.splits().iter().map(|s| s.data_files().len()).sum(); - assert_eq!(file_count, 1, "After INSERT: 1 postpone file"); + assert_eq!(file_count, 1, "After INSERT: 1 fixed-bucket file"); // INSERT OVERWRITE should replace old file sql_context diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index e26c6bc34..a17303b71 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -121,6 +121,13 @@ const DEFAULT_PARQUET_ROW_GROUP_PARALLELISM: usize = 8; const DEFAULT_PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES: i64 = 256 * 1024 * 1024; const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str = "dynamic-bucket.target-row-num"; const DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM: i64 = 200_000; +const POSTPONE_BATCH_WRITE_FIXED_BUCKET_OPTION: &str = "postpone.batch-write-fixed-bucket"; +const POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION: &str = + "postpone.batch-write-fixed-bucket.max-parallelism"; +const POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION: &str = "postpone.target-row-num-per-bucket"; +const POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION: &str = "postpone.target-size-per-bucket"; +const DEFAULT_POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM: i32 = 2048; +const DEFAULT_POSTPONE_TARGET_SIZE_PER_BUCKET: i64 = 1024 * 1024 * 1024; const DEFAULT_GLOBAL_INDEX_ROW_COUNT_PER_SHARD: i64 = 100_000; const DEFAULT_GLOBAL_INDEX_THREAD_NUM: i64 = 32; const DEFAULT_GLOBAL_INDEX_FALLBACK_SCAN_MAX_SIZE: i64 = 256 * 1024 * 1024; @@ -1148,6 +1155,75 @@ impl<'a> CoreOptions<'a> { .unwrap_or(DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM) } + /// Whether batch writes to postpone-bucket primary-key tables use real + /// fixed buckets. This is enabled by default; set it to false to retain + /// legacy bucket -2 writes. + pub fn postpone_batch_write_fixed_bucket(&self) -> bool { + self.options + .get(POSTPONE_BATCH_WRITE_FIXED_BUCKET_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(true) + } + + pub fn postpone_batch_write_fixed_bucket_max_parallelism(&self) -> crate::Result { + let value = self + .options + .get(POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION) + .map(|value| value.parse::()) + .transpose() + .map_err(|error| crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION}' must be a positive integer" + ), + source: Some(Box::new(error)), + })? + .unwrap_or(DEFAULT_POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM); + if value <= 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION}' must be positive, got: {value}" + ), + source: None, + }); + } + Ok(value) + } + + pub fn postpone_target_row_num_per_bucket(&self) -> crate::Result> { + let value = self.parse_i64_option(POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION)?; + if value.is_some_and(|value| value <= 0) { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION}' must be positive, got: {}", + value.unwrap() + ), + source: None, + }); + } + Ok(value) + } + + pub fn postpone_target_size_per_bucket(&self) -> crate::Result { + let value = match self.options.get(POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION) { + Some(raw) => parse_memory_size(raw).ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION}' must be a valid positive memory size, got: {raw}" + ), + source: None, + })?, + None => DEFAULT_POSTPONE_TARGET_SIZE_PER_BUCKET, + }; + if value <= 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION}' must be positive, got: {value}" + ), + source: None, + }); + } + Ok(value) + } + /// When true, blob field reads return serialized BlobDescriptor bytes /// instead of actual blob bytes. Default is false. pub fn blob_as_descriptor(&self) -> bool { diff --git a/crates/paimon/src/table/commit_message.rs b/crates/paimon/src/table/commit_message.rs index 5a91e36b5..7332c2c47 100644 --- a/crates/paimon/src/table/commit_message.rs +++ b/crates/paimon/src/table/commit_message.rs @@ -27,6 +27,9 @@ pub struct CommitMessage { pub partition: Vec, /// Bucket id. pub bucket: i32, + /// Per-partition bucket count. Set by fixed-bucket postpone batch writes; + /// ordinary writes use the table-level bucket option. + pub total_buckets: Option, /// New data files to be added. pub new_files: Vec, /// Snapshot id from which row-id/column conflicts should be checked. @@ -46,6 +49,7 @@ impl CommitMessage { Self { partition, bucket, + total_buckets: None, new_files, check_from_snapshot: None, new_changelog_files: Vec::new(), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 4078e3a23..1431d117e 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -72,6 +72,8 @@ mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; mod pk_vector_position_read; mod pk_vector_scan; +mod postpone_batch_table_write; +mod postpone_bucket; mod postpone_file_writer; mod prepared_files; mod read_builder; @@ -120,6 +122,9 @@ pub use incremental_scan::{ }; pub use lumina_index_build_builder::LuminaIndexBuildBuilder; pub use partition_stat::PartitionStat; +pub use postpone_batch_table_write::{ + PostponeBucketPlan, POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD, +}; pub use read_builder::ReadBuilder; pub use rest_env::RESTEnv; pub use scan_trace::ScanTrace; @@ -359,6 +364,12 @@ impl Table { WriteBuilder::new(self) } + /// Create a writer which forces one-shot fixed-bucket writes for a + /// postpone table, even when `postpone.batch-write-fixed-bucket=false`. + pub fn new_postpone_fixed_bucket_write_builder(&self) -> Result> { + WriteBuilder::new_postpone_fixed_bucket(self) + } + /// Create a copy of this table with extra options merged into the schema. /// /// This never switches the schema version; it corresponds to Java diff --git a/crates/paimon/src/table/postpone_batch_table_write.rs b/crates/paimon/src/table/postpone_batch_table_write.rs new file mode 100644 index 000000000..a9b2cae07 --- /dev/null +++ b/crates/paimon/src/table/postpone_batch_table_write.rs @@ -0,0 +1,514 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! One-shot fixed-bucket planning for batch writes to postpone tables. +//! +//! This mirrors pypaimon's `PostponeFixedBucketBatchTableWrite`: partitions +//! with an existing real-bucket count stream directly to their writers, while +//! new partitions are buffered until `prepare_commit` can infer one bucket +//! count from the complete batch. + +use crate::spec::{ + batch_to_serialized_bytes, BucketFunctionType, CoreOptions, DataField, EMPTY_SERIALIZED_ROW, + POSTPONE_BUCKET, +}; +use crate::table::bucket_function::{batch_bucket_ids, validate_bucket_function}; +use crate::table::postpone_bucket::binary_row_batch_size; +use crate::table::{SnapshotManager, Table, TableScan}; +use crate::Result; +use arrow_array::{Array, Int32Array, RecordBatch, UInt32Array}; +use std::collections::HashMap; + +/// Column name used by [`PostponeBucketPlan::from_arrow`] for bucket counts. +pub const POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD: &str = "total_buckets"; + +/// A shared fixed-bucket plan for distributed writers. +/// +/// The plan must be computed from global partition statistics and supplied to +/// every writer participating in one logical batch commit. +#[derive(Debug, Clone)] +pub struct PostponeBucketPlan { + bucket_counts: HashMap, i32>, +} + +impl PostponeBucketPlan { + /// Build a plan from an Arrow batch containing the table's partition + /// columns, in partition-key order, followed by a non-null Int32 + /// `total_buckets` column. For an unpartitioned table, the batch contains + /// only `total_buckets` and normally has one row. + pub fn from_arrow(table: &Table, batch: &RecordBatch) -> Result { + let partition_fields = table.schema().partition_fields(); + let partition_count = partition_fields.len(); + if batch.num_columns() != partition_count + 1 { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan expected {} partition column(s) plus '{POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}', got {} columns", + partition_count, + batch.num_columns() + ), + source: None, + }); + } + + let expected_partition_schema = crate::arrow::build_target_arrow_schema(&partition_fields)?; + let batch_schema = batch.schema(); + for (index, expected) in expected_partition_schema.fields().iter().enumerate() { + let actual = batch_schema.field(index); + if actual.name() != expected.name() || actual.data_type() != expected.data_type() { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan partition field mismatch at index {index}: expected '{}': {:?}, got '{}': {:?}", + expected.name(), + expected.data_type(), + actual.name(), + actual.data_type() + ), + source: None, + }); + } + if !expected.is_nullable() && batch.column(index).null_count() != 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan partition column '{}' is NOT NULL but contains null values", + expected.name() + ), + source: None, + }); + } + } + + let count_field = batch_schema.field(partition_count); + if count_field.name() != POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD + || count_field.data_type() != &arrow_schema::DataType::Int32 + { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan final field must be '{POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}': Int32, got '{}': {:?}", + count_field.name(), + count_field.data_type() + ), + source: None, + }); + } + let counts = batch + .column(partition_count) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: "Postpone bucket plan total_buckets column is not Int32".to_string(), + source: None, + })?; + let partition_indices = (0..partition_count).collect::>(); + let partitions = batch_to_serialized_bytes(batch, &partition_indices, &partition_fields)?; + let mut bucket_counts = HashMap::with_capacity(batch.num_rows()); + for (row, partition) in partitions.into_iter().enumerate() { + if counts.is_null(row) { + return Err(crate::Error::DataInvalid { + message: format!("Postpone bucket plan total_buckets is null at row {row}"), + source: None, + }); + } + let total_buckets = counts.value(row); + if total_buckets <= 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan total_buckets must be positive at row {row}, got {total_buckets}" + ), + source: None, + }); + } + if let Some(previous) = bucket_counts.insert(partition, total_buckets) { + if previous != total_buckets { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone bucket plan contains conflicting total bucket counts {previous} and {total_buckets} for one partition" + ), + source: None, + }); + } + } + } + Ok(Self { bucket_counts }) + } + + fn into_bucket_counts(self) -> HashMap, i32> { + self.bucket_counts + } +} + +pub(super) struct PostponeBucketBatch { + pub(super) partition: Vec, + pub(super) bucket: i32, + pub(super) batch: RecordBatch, +} + +/// Planning state for a single fixed-bucket batch write to a postpone table. +pub(super) struct PostponeFixedBucketWriter { + partition_field_indices: Vec, + bucket_key_indices: Vec, + bucket_function_type: BucketFunctionType, + max_parallelism: i32, + target_rows_per_bucket: Option, + target_size_per_bucket: Option, + plan_provided: bool, + metadata_loaded: bool, + known_bucket_counts: HashMap, i32>, + postpone_row_counts: HashMap, i64>, + buffered_batches: HashMap, Vec>, + /// Bucket counts used by this prepare-commit round. + bucket_counts: HashMap, i32>, + prepare_started: bool, +} + +impl PostponeFixedBucketWriter { + pub(super) fn new( + table: &Table, + partition_field_indices: Vec, + bucket_key_indices: Vec, + bucket_function_type: BucketFunctionType, + bucket_plan: Option, + ) -> Result { + let schema = table.schema(); + let options = CoreOptions::new(schema.options()); + let total_buckets = options.bucket(); + if total_buckets != POSTPONE_BUCKET || schema.primary_keys().is_empty() { + return Err(crate::Error::Unsupported { + message: format!( + "Postpone fixed-bucket writes require a primary-key table with bucket=-2, but table '{}' has bucket={total_buckets}", + table.identifier().full_name() + ), + }); + } + if options.deletion_vectors_enabled() { + return Err(crate::Error::Unsupported { + message: format!( + "Table '{}' cannot use postpone fixed-bucket writes with deletion-vectors.enabled=true because deletion-vector scans skip the level-0 files produced by batch writers; use the normal postpone writer or disable deletion vectors", + table.identifier().full_name() + ), + }); + } + + let bucket_key_fields: Vec = bucket_key_indices + .iter() + .map(|&index| schema.fields()[index].clone()) + .collect(); + if !bucket_key_fields.is_empty() { + validate_bucket_function(bucket_function_type, &bucket_key_fields)?; + } + + let target_rows_per_bucket = options.postpone_target_row_num_per_bucket()?; + let target_size_per_bucket = if target_rows_per_bucket.is_none() { + Some(options.postpone_target_size_per_bucket()?) + } else { + None + }; + let plan_provided = bucket_plan.is_some(); + let known_bucket_counts = bucket_plan + .map(PostponeBucketPlan::into_bucket_counts) + .unwrap_or_default(); + + Ok(Self { + partition_field_indices, + bucket_key_indices, + bucket_function_type, + max_parallelism: options.postpone_batch_write_fixed_bucket_max_parallelism()?, + target_rows_per_bucket, + target_size_per_bucket, + plan_provided, + metadata_loaded: plan_provided, + known_bucket_counts, + postpone_row_counts: HashMap::new(), + buffered_batches: HashMap::new(), + bucket_counts: HashMap::new(), + prepare_started: false, + }) + } + + pub(super) fn ensure_writable(&self) -> Result<()> { + if self.prepare_started { + return Err(Self::one_shot_error()); + } + Ok(()) + } + + pub(super) fn start_prepare(&mut self) -> Result<()> { + self.ensure_writable()?; + // A failed prepare may already have consumed buffered batches or + // closed file writers, so the same writer cannot be retried safely. + self.prepare_started = true; + Ok(()) + } + + pub(super) async fn write_batch( + &mut self, + table: &Table, + batch: &RecordBatch, + ) -> Result> { + self.ensure_metadata_loaded(table).await?; + + let partitions = if self.partition_field_indices.is_empty() { + vec![EMPTY_SERIALIZED_ROW.clone(); batch.num_rows()] + } else { + batch_to_serialized_bytes( + batch, + &self.partition_field_indices, + table.schema().fields(), + )? + }; + + let mut groups: HashMap, Vec> = HashMap::new(); + for (row, partition) in partitions.into_iter().enumerate() { + groups.entry(partition).or_default().push(row); + } + + if self.plan_provided + && groups + .keys() + .any(|partition| !self.known_bucket_counts.contains_key(partition)) + { + return Err(crate::Error::DataInvalid { + message: "Postpone bucket plan does not contain an input partition".to_string(), + source: None, + }); + } + + let mut output = Vec::new(); + for (partition, rows) in groups { + let sub_batch = take_rows(batch, &rows)?; + if let Some(total_buckets) = self.known_bucket_counts.get(&partition).copied() { + self.bucket_counts.insert(partition.clone(), total_buckets); + output.extend(self.route_batch(table, partition, sub_batch, total_buckets)?); + } else { + self.buffered_batches + .entry(partition) + .or_default() + .push(sub_batch); + } + } + Ok(output) + } + + pub(super) async fn prepare_batch( + &mut self, + table: &Table, + is_overwrite: bool, + ) -> Result> { + if self.buffered_batches.is_empty() { + return Ok(Vec::new()); + } + + let buffered_batches = std::mem::take(&mut self.buffered_batches); + let mut output = Vec::new(); + for (partition, batches) in buffered_batches { + let input_rows = batches.iter().fold(0_i64, |rows, batch| { + rows.saturating_add(batch.num_rows() as i64) + }); + // Match pypaimon: row-count planning does not inspect row sizes. + // Size planning ignores the trailing internal `_VALUE_KIND` field + // appended by TableWrite after row-kind generation. + let input_size = if self.target_rows_per_bucket.is_none() { + batches.iter().try_fold(0_i64, |size, batch| { + Ok::<_, crate::Error>( + size.saturating_add(binary_row_batch_size(batch, table.schema().fields())?), + ) + })? + } else { + 0 + }; + let postpone_rows = if is_overwrite { + 0 + } else { + self.postpone_row_counts + .get(&partition) + .copied() + .unwrap_or(0) + }; + let total_buckets = infer_bucket_count( + input_rows, + input_size, + postpone_rows, + self.target_rows_per_bucket, + self.target_size_per_bucket, + self.max_parallelism, + ); + self.known_bucket_counts + .insert(partition.clone(), total_buckets); + self.bucket_counts.insert(partition.clone(), total_buckets); + + for batch in batches { + output.extend(self.route_batch(table, partition.clone(), batch, total_buckets)?); + } + } + Ok(output) + } + + pub(super) fn total_buckets(&self, partition: &[u8]) -> Option { + self.bucket_counts.get(partition).copied() + } + + pub(super) fn finish(&mut self) { + self.metadata_loaded = false; + self.known_bucket_counts.clear(); + self.postpone_row_counts.clear(); + self.bucket_counts.clear(); + } + + #[cfg(test)] + pub(super) fn buffered_partition_count(&self) -> usize { + self.buffered_batches.len() + } + + async fn ensure_metadata_loaded(&mut self, table: &Table) -> Result<()> { + if self.metadata_loaded { + return Ok(()); + } + let (known_bucket_counts, postpone_row_counts) = load_bucket_metadata(table).await?; + self.known_bucket_counts = known_bucket_counts; + self.postpone_row_counts = postpone_row_counts; + self.metadata_loaded = true; + Ok(()) + } + + fn route_batch( + &self, + table: &Table, + partition: Vec, + batch: RecordBatch, + total_buckets: i32, + ) -> Result> { + let buckets = if total_buckets <= 1 || self.bucket_key_indices.is_empty() { + vec![0; batch.num_rows()] + } else { + batch_bucket_ids( + &batch, + &self.bucket_key_indices, + table.schema().fields(), + self.bucket_function_type, + total_buckets, + )? + }; + let mut groups: HashMap> = HashMap::new(); + for (row, bucket) in buckets.into_iter().enumerate() { + groups.entry(bucket).or_default().push(row); + } + groups + .into_iter() + .map(|(bucket, rows)| { + Ok(PostponeBucketBatch { + partition: partition.clone(), + bucket, + batch: take_rows(&batch, &rows)?, + }) + }) + .collect() + } + + fn one_shot_error() -> crate::Error { + crate::Error::DataInvalid { + message: "Fixed-bucket postpone TableWrite only supports one prepare_commit call; create a new writer for the next batch".to_string(), + source: None, + } + } +} + +async fn load_bucket_metadata( + table: &Table, +) -> Result<(HashMap, i32>, HashMap, i64>)> { + let mut known_bucket_counts = HashMap::new(); + let mut postpone_row_counts = HashMap::new(); + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + let Some(snapshot) = snapshot_manager.get_latest_snapshot().await? else { + return Ok((known_bucket_counts, postpone_row_counts)); + }; + + let scan = TableScan::new(table, None, vec![], None, None, None).with_scan_all_files(); + for entry in scan.plan_manifest_entries(&snapshot).await? { + let partition = entry.partition().to_vec(); + if entry.bucket() == POSTPONE_BUCKET { + let rows = postpone_row_counts.entry(partition).or_insert(0_i64); + *rows = rows.saturating_add(entry.file().row_count); + } else if entry.bucket() >= 0 && entry.total_buckets() > 0 { + if let Some(previous) = + known_bucket_counts.insert(partition.clone(), entry.total_buckets()) + { + if previous != entry.total_buckets() { + return Err(crate::Error::DataInvalid { + message: format!( + "Partition has inconsistent total bucket counts: {previous} and {}", + entry.total_buckets() + ), + source: None, + }); + } + } + } + } + Ok((known_bucket_counts, postpone_row_counts)) +} + +fn infer_bucket_count( + input_rows: i64, + input_size: i64, + postpone_rows: i64, + target_rows_per_bucket: Option, + target_size_per_bucket: Option, + max_parallelism: i32, +) -> i32 { + let buckets = if let Some(target_rows) = target_rows_per_bucket { + let total_rows = input_rows.saturating_add(postpone_rows); + total_rows.saturating_add(target_rows - 1) / target_rows + } else { + let target_size_per_bucket = target_size_per_bucket + .expect("size target is validated when row-count target is absent"); + let estimated_size = if postpone_rows > 0 && input_rows > 0 { + let numerator = i128::from(input_size) + .saturating_mul(i128::from(input_rows.saturating_add(postpone_rows))); + let estimate = (numerator + i128::from(input_rows - 1)) / i128::from(input_rows); + estimate.min(i128::from(i64::MAX)) as i64 + } else { + input_size + }; + estimated_size.saturating_add(target_size_per_bucket - 1) / target_size_per_bucket + }; + buckets.max(1).min(i64::from(max_parallelism)) as i32 +} + +fn take_rows(batch: &RecordBatch, row_indices: &[usize]) -> Result { + if row_indices.len() == batch.num_rows() { + return Ok(batch.clone()); + } + let indices = UInt32Array::from( + row_indices + .iter() + .map(|&index| index as u32) + .collect::>(), + ); + let columns = batch + .columns() + .iter() + .map(|column| arrow_select::take::take(column.as_ref(), &indices, None)) + .collect::, _>>() + .map_err(|error| crate::Error::DataInvalid { + message: format!("Failed to take rows for postpone bucket planning: {error}"), + source: None, + })?; + RecordBatch::try_new(batch.schema(), columns).map_err(|error| crate::Error::DataInvalid { + message: format!("Failed to create postpone bucket batch: {error}"), + source: None, + }) +} diff --git a/crates/paimon/src/table/postpone_bucket.rs b/crates/paimon/src/table/postpone_bucket.rs new file mode 100644 index 000000000..c0725a25c --- /dev/null +++ b/crates/paimon/src/table/postpone_bucket.rs @@ -0,0 +1,464 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Helpers shared by fixed-bucket postpone batch planning. + +use crate::spec::{BinaryRow, DataField, DataType, IntType, VALUE_KIND_FIELD_NAME}; +use arrow_array::{ + Array, ArrayRef, BinaryArray, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, + MapArray, RecordBatch, StringArray, StringViewArray, StructArray, +}; + +/// Sum the Java `BinaryRow.getSizeInBytes()` value for every row in a batch. +/// +/// Arrow allocation size is not a stable planning metric: offsets, views, and +/// slicing can make the same logical rows occupy different Arrow memory. This +/// estimator follows Paimon's internal BinaryRow/BinaryArray layouts without +/// materializing a second copy of the batch. The trailing internal +/// `_VALUE_KIND` column is not part of the user row and is excluded. +pub(crate) fn binary_row_batch_size( + batch: &RecordBatch, + fields: &[DataField], +) -> crate::Result { + let field_count = fields.len(); + let has_value_kind = batch.num_columns() == field_count + 1 + && batch.schema().field(field_count).name() == VALUE_KIND_FIELD_NAME + && batch.schema().field(field_count).data_type() == &arrow_schema::DataType::Int8; + if batch.num_columns() != field_count && !has_value_kind { + return Err(crate::Error::DataInvalid { + message: format!( + "BinaryRow size planning expected {field_count} user columns with an optional trailing {VALUE_KIND_FIELD_NAME}, got {}", + batch.num_columns() + ), + source: None, + }); + } + + let mut total = 0_u128; + for row in 0..batch.num_rows() { + total = total.saturating_add(row_size(&batch.columns()[..field_count], row, fields)?); + } + Ok(total.min(i64::MAX as u128) as i64) +} + +fn row_size(arrays: &[ArrayRef], row: usize, fields: &[DataField]) -> crate::Result { + let mut size = BinaryRow::cal_fix_part_size_in_bytes(fields.len() as i32) as u128; + for (array, field) in arrays.iter().zip(fields) { + size = size.saturating_add(variable_size(array, row, field.data_type())?); + } + Ok(size) +} + +fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + // Java's BinaryRowWriter reserves these fixed variable regions even for a + // null top-level field. + match data_type { + DataType::Decimal(decimal) if decimal.precision() > 18 => return Ok(16), + DataType::Timestamp(timestamp) if timestamp.precision() > 3 => return Ok(8), + DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() > 3 => return Ok(8), + _ => {} + } + + if array.is_null(row) { + return Ok(0); + } + + match data_type { + DataType::Char(_) | DataType::VarChar(_) => { + let len = string_len(array, row, data_type)?; + Ok(binary_size(len)) + } + DataType::Binary(_) | DataType::VarBinary(_) | DataType::Blob(_) => { + let len = binary_len(array, row, data_type)?; + Ok(binary_size(len)) + } + DataType::Variant(_) => variant_size(array, row), + DataType::Array(array_type) => { + if let Some(array) = array.as_any().downcast_ref::() { + let offsets = array.value_offsets(); + Ok(round_to_word(binary_array_size( + array.values(), + offsets[row] as usize, + offsets[row + 1] as usize, + array_type.element_type(), + )?)) + } else if let Some(array) = array.as_any().downcast_ref::() { + let offsets = array.value_offsets(); + Ok(round_to_word(binary_array_size( + array.values(), + offsets[row] as usize, + offsets[row + 1] as usize, + array_type.element_type(), + )?)) + } else { + Err(type_mismatch("ListArray", data_type)) + } + } + DataType::Map(map_type) => { + let map = downcast::(array, "MapArray", data_type)?; + let offsets = map.value_offsets(); + let entries = map.entries(); + map_size( + entries, + offsets[row] as usize, + offsets[row + 1] as usize, + map_type.key_type(), + map_type.value_type(), + ) + } + DataType::Multiset(multiset_type) => { + let map = downcast::(array, "MapArray", data_type)?; + let offsets = map.value_offsets(); + let entries = map.entries(); + map_size( + entries, + offsets[row] as usize, + offsets[row + 1] as usize, + multiset_type.element_type(), + &DataType::Int(IntType::new()), + ) + } + DataType::Row(row_type) => { + let struct_array = downcast::(array, "StructArray", data_type)?; + Ok(round_to_word(row_size( + struct_array.columns(), + row, + row_type.fields(), + )?)) + } + DataType::Vector(vector_type) => { + let element_bytes = u128::from(vector_type.length()) + .saturating_mul(primitive_width(vector_type.element_type())?); + Ok(round_to_word(4_u128.saturating_add(element_bytes))) + } + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Float(_) + | DataType::Double(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) => Ok(0), + } +} + +fn binary_array_size( + values: &ArrayRef, + start: usize, + end: usize, + element_type: &DataType, +) -> crate::Result { + if end < start || end > values.len() { + return Err(crate::Error::DataInvalid { + message: format!( + "Invalid nested array range [{start}, {end}) for {} values", + values.len() + ), + source: None, + }); + } + let count = end - start; + let header = 4_u128.saturating_add((count as u128).div_ceil(32).saturating_mul(4)); + let fixed = (count as u128).saturating_mul(fixed_width(element_type)?); + let mut size = round_to_word(header.saturating_add(fixed)); + for row in start..end { + if !values.is_null(row) { + size = size.saturating_add(variable_size(values, row, element_type)?); + } + } + Ok(size) +} + +fn map_size( + entries: &StructArray, + start: usize, + end: usize, + key_type: &DataType, + value_type: &DataType, +) -> crate::Result { + if entries.num_columns() != 2 { + return Err(crate::Error::DataInvalid { + message: format!( + "BinaryMap size planning expected 2 entry columns, got {}", + entries.num_columns() + ), + source: None, + }); + } + let keys = binary_array_size(entries.column(0), start, end, key_type)?; + let values = binary_array_size(entries.column(1), start, end, value_type)?; + Ok(round_to_word( + 4_u128.saturating_add(keys).saturating_add(values), + )) +} + +fn variant_size(array: &ArrayRef, row: usize) -> crate::Result { + let variant = downcast::( + array, + "StructArray", + &DataType::Variant(crate::spec::VariantType::new()), + )?; + if variant.num_columns() != 2 { + return Err(crate::Error::DataInvalid { + message: format!( + "Variant size planning expected 2 child columns, got {}", + variant.num_columns() + ), + source: None, + }); + } + let value_len = binary_len( + variant.column(0), + row, + &DataType::Variant(crate::spec::VariantType::new()), + )?; + let metadata_len = binary_len( + variant.column(1), + row, + &DataType::Variant(crate::spec::VariantType::new()), + )?; + Ok(round_to_word( + 4_u128 + .saturating_add(value_len as u128) + .saturating_add(metadata_len as u128), + )) +} + +fn string_len(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else { + Err(type_mismatch("StringArray", data_type)) + } +} + +fn binary_len(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else { + Err(type_mismatch("BinaryArray", data_type)) + } +} + +fn fixed_width(data_type: &DataType) -> crate::Result { + Ok(match data_type { + DataType::Boolean(_) | DataType::TinyInt(_) => 1, + DataType::SmallInt(_) => 2, + DataType::Int(_) | DataType::Float(_) | DataType::Date(_) | DataType::Time(_) => 4, + DataType::BigInt(_) + | DataType::Double(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Binary(_) + | DataType::VarBinary(_) + | DataType::Blob(_) + | DataType::Variant(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) + | DataType::Array(_) + | DataType::Map(_) + | DataType::Multiset(_) + | DataType::Row(_) + | DataType::Vector(_) => 8, + }) +} + +fn primitive_width(data_type: &DataType) -> crate::Result { + match data_type { + DataType::Boolean(_) | DataType::TinyInt(_) => Ok(1), + DataType::SmallInt(_) => Ok(2), + DataType::Int(_) | DataType::Float(_) => Ok(4), + DataType::BigInt(_) | DataType::Double(_) => Ok(8), + other => Err(crate::Error::DataInvalid { + message: format!("Unsupported vector element type for size planning: {other:?}"), + source: None, + }), + } +} + +fn binary_size(len: usize) -> u128 { + if len <= 7 { + 0 + } else { + round_to_word(len as u128) + } +} + +fn round_to_word(size: u128) -> u128 { + size.saturating_add(7) / 8 * 8 +} + +fn downcast<'a, T: 'static>( + array: &'a ArrayRef, + expected: &str, + data_type: &DataType, +) -> crate::Result<&'a T> { + array + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch(expected, data_type)) +} + +fn type_mismatch(expected: &str, data_type: &DataType) -> crate::Error { + crate::Error::DataInvalid { + message: format!("BinaryRow size planning expected {expected} for {data_type:?}"), + source: None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{ArrayType, IntType, LocalZonedTimestampType, TimestampType, VarCharType}; + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, Int8Array, ListArray, TimestampMicrosecondArray}; + use arrow_schema::{DataType as ArrowDataType, Field, Schema, TimeUnit}; + use std::sync::Arc; + + #[test] + fn test_java_binary_row_size_differs_from_arrow_buffers() { + let row_count = 1_000; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("left", ArrowDataType::Utf8, false), + Field::new("right", ArrowDataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..row_count)), + Arc::new(StringArray::from(vec!["a"; row_count as usize])), + Arc::new(StringArray::from(vec!["b"; row_count as usize])), + ], + ) + .unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "left".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + DataField::new( + 2, + "right".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 32_000); + assert_ne!(batch.get_array_memory_size() as i64, 32_000); + } + + #[test] + fn test_internal_value_kind_is_excluded() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("value", ArrowDataType::Int32, false), + Field::new(VALUE_KIND_FIELD_NAME, ArrowDataType::Int8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), + Arc::new(Int8Array::from(vec![0])), + ], + ) + .unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "value".to_string(), DataType::Int(IntType::new())), + ]; + + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 24); + } + + #[test] + fn test_nested_and_local_zoned_timestamp_size() { + let array = + ListArray::from_iter_primitive::(vec![Some(vec![Some(1), Some(2)])]); + let timestamp = + TimestampMicrosecondArray::from(vec![Some(1_234_567_i64)]).with_timezone("UTC"); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("items", array.data_type().clone(), true), + Field::new( + "event_time", + ArrowDataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(array), + Arc::new(timestamp), + ], + ) + .unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "items".to_string(), + DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), + ), + DataField::new( + 2, + "event_time".to_string(), + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(6).unwrap()), + ), + ]; + + // 32-byte row fixed part + 16-byte BinaryArray + 8-byte non-compact timestamp. + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 56); + } + + #[test] + fn test_non_compact_null_timestamp_reserves_space() { + let schema = Arc::new(Schema::new(vec![Field::new( + "event_time", + ArrowDataType::Timestamp(TimeUnit::Microsecond, None), + true, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(TimestampMicrosecondArray::from(vec![None]))], + ) + .unwrap(); + let fields = vec![DataField::new( + 0, + "event_time".to_string(), + DataType::Timestamp(TimestampType::new(6).unwrap()), + )]; + + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 24); + } +} diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 809e45d9f..30179fab1 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -1240,7 +1240,11 @@ impl TableCommit { } else { CommitKind::APPEND }; - let detect_conflicts = has_delete || check_from_snapshot.is_some(); + let has_partition_bucket_counts = entries + .iter() + .any(|entry| entry.total_buckets() != self.total_buckets); + let detect_conflicts = + has_delete || check_from_snapshot.is_some() || has_partition_bucket_counts; let base_data_files = if detect_conflicts { self.check_deletion_vector_index_only_conflict( latest_snapshot.as_ref(), @@ -1854,6 +1858,20 @@ impl TableCommit { ) -> Result<()> { self.check_delete_entries_against_base(base_entries, delta_entries)?; + // Validate every file produced by this commit before merging entries. + // This catches inconsistent duplicate ADDs even when they describe the + // same file and would otherwise collapse during merge_active_entries. + self.check_total_bucket_conflicts(delta_entries)?; + + // Check the final active layout rather than raw base + delta entries. + // In an overwrite, DELETE-old followed by ADD-new is a valid bucket + // rescale; concurrent APPENDs with incompatible counts remain active + // together and are still rejected. + let mut all_entries = base_entries.to_vec(); + all_entries.extend(delta_entries.iter().cloned()); + let merged_entries = merge_active_entries(all_entries); + self.check_total_bucket_conflicts(&merged_entries)?; + if !self.data_evolution_enabled { return Ok(()); } @@ -1861,14 +1879,33 @@ impl TableCommit { let next_row_id = latest_snapshot.and_then(Snapshot::next_row_id); self.check_row_id_existence(base_entries, delta_entries, next_row_id)?; - let mut all_entries = base_entries.to_vec(); - all_entries.extend(delta_entries.iter().cloned()); - let merged_entries = merge_active_entries(all_entries); self.check_row_id_range_conflicts(commit_kind, check_from_snapshot, &merged_entries)?; self.check_row_id_from_snapshot(latest_snapshot, delta_entries, check_from_snapshot) .await } + fn check_total_bucket_conflicts(&self, entries: &[ManifestEntry]) -> Result<()> { + let mut bucket_counts: HashMap, i32> = HashMap::new(); + for entry in entries { + if *entry.kind() != FileKind::Add || entry.bucket() < 0 || entry.total_buckets() <= 0 { + continue; + } + let partition = entry.partition().to_vec(); + if let Some(previous) = bucket_counts.insert(partition, entry.total_buckets()) { + if previous != entry.total_buckets() { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket conflict: one partition uses different total bucket counts {previous} and {}", + entry.total_buckets() + ), + source: None, + }); + } + } + } + Ok(()) + } + fn check_deletion_vector_index_only_conflict( &self, latest_snapshot: Option<&Snapshot>, @@ -2557,6 +2594,10 @@ impl TableCommit { stats.file_size_in_bytes += sign * file.file_size; stats.file_count += sign; stats.last_file_creation_time = stats.last_file_creation_time.max(file_creation_time); + // Overwrite entries are ordered DELETE-old then ADD-new. Match + // Java PartitionEntry.merge by retaining the replacement entry's + // bucket count instead of the first value seen for the partition. + stats.total_buckets = entry.total_buckets(); } Ok(stats_map.into_values().collect()) @@ -2602,7 +2643,7 @@ impl TableCommit { FileKind::Add, msg.partition.clone(), msg.bucket, - self.total_buckets, + msg.total_buckets.unwrap_or(self.total_buckets), file.clone(), 2, ) @@ -2612,7 +2653,7 @@ impl TableCommit { FileKind::Delete, msg.partition.clone(), msg.bucket, - self.total_buckets, + msg.total_buckets.unwrap_or(self.total_buckets), file.clone(), 2, ) @@ -2632,7 +2673,7 @@ impl TableCommit { FileKind::Add, msg.partition.clone(), msg.bucket, - self.total_buckets, + msg.total_buckets.unwrap_or(self.total_buckets), file.clone(), 0, ) @@ -2903,7 +2944,7 @@ mod tests { use crate::spec::stats::BinaryTableStats; use crate::spec::{ BinaryRowBuilder, DataFileMeta, DeletionVectorMeta, GlobalIndexMeta, IndexFileMeta, - ManifestList, TableSchema, + ManifestList, TableSchema, POSTPONE_BUCKET, }; use chrono::{DateTime, Utc}; @@ -4135,6 +4176,45 @@ mod tests { assert_eq!(snapshot.total_record_count(), Some(250)); } + #[test] + fn test_partition_statistics_keep_replacement_bucket_count() { + let file_io = test_file_io(); + let commit = setup_partitioned_commit( + &file_io, + "memory:/test_partition_statistics_replacement_bucket_count", + ); + let partition = partition_bytes("a"); + + for (old_buckets, new_buckets) in [(-2, 4), (4, 8)] { + let entries = vec![ + ManifestEntry::new( + FileKind::Delete, + partition.clone(), + if old_buckets == POSTPONE_BUCKET { + POSTPONE_BUCKET + } else { + 0 + }, + old_buckets, + test_data_file("old.parquet", 100), + 2, + ), + ManifestEntry::new( + FileKind::Add, + partition.clone(), + 0, + new_buckets, + test_data_file("new.parquet", 50), + 2, + ), + ]; + + let statistics = commit.generate_partition_statistics(&entries).unwrap(); + assert_eq!(statistics.len(), 1); + assert_eq!(statistics[0].total_buckets, new_buckets); + } + } + #[tokio::test] async fn test_overwrite_cache_reuses_when_append_misses_target_partition() { let file_io = test_file_io(); diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 490de6357..e2635a106 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -38,6 +38,7 @@ use crate::table::data_file_writer::DataFileWriter; use crate::table::dedicated_format_file_writer::AppendDedicatedFormatFileWriter; use crate::table::kv_file_writer::{KeyValueFileWriter, KeyValueWriteConfig}; use crate::table::partition_filter::PartitionFilter; +use crate::table::postpone_batch_table_write::{PostponeBucketPlan, PostponeFixedBucketWriter}; use crate::table::postpone_file_writer::{PostponeFileWriter, PostponeWriteConfig}; use crate::table::prepared_files::PreparedFiles; use crate::table::row_kind_generator::RowKindGenerator; @@ -126,10 +127,35 @@ pub struct TableWrite { has_dedicated_vector_fields: bool, row_kind_generator: Option, row_kind_filter: Option, + postpone_fixed_bucket: Option, } impl TableWrite { pub(crate) fn new(table: &Table, commit_user: String) -> crate::Result { + Self::new_inner(table, commit_user, false, None) + } + + pub(crate) fn new_postpone_fixed_bucket( + table: &Table, + commit_user: String, + ) -> crate::Result { + Self::new_inner(table, commit_user, true, None) + } + + pub(crate) fn new_postpone_fixed_bucket_with_plan( + table: &Table, + commit_user: String, + bucket_plan: PostponeBucketPlan, + ) -> crate::Result { + Self::new_inner(table, commit_user, true, Some(bucket_plan)) + } + + fn new_inner( + table: &Table, + commit_user: String, + use_postpone_fixed_bucket: bool, + postpone_bucket_plan: Option, + ) -> crate::Result { let is_overwrite = false; let schema = table.schema(); let write_schema = build_target_arrow_schema(schema.fields())?; @@ -295,6 +321,18 @@ impl TableWrite { let target_bucket_row_number = core_options.dynamic_bucket_target_row_num(); let bucket_function_type = core_options.bucket_function_type()?; + let postpone_fixed_bucket = use_postpone_fixed_bucket + .then(|| { + PostponeFixedBucketWriter::new( + table, + partition_field_indices.clone(), + bucket_key_indices.clone(), + bucket_function_type, + postpone_bucket_plan, + ) + }) + .transpose()?; + let bucket_assigner = if is_dynamic_cross_partition { BucketAssignerEnum::CrossPartition(Box::new(CrossPartitionAssigner::new( table.clone(), @@ -377,6 +415,7 @@ impl TableWrite { has_dedicated_vector_fields, row_kind_generator, row_kind_filter, + postpone_fixed_bucket, }) } @@ -437,6 +476,9 @@ impl TableWrite { /// Write an Arrow RecordBatch. Rows are routed to the correct partition and bucket. pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) -> Result<()> { + if let Some(writer) = self.postpone_fixed_bucket.as_ref() { + writer.ensure_writable()?; + } self.validate_write_batch_schema(batch)?; if batch.num_rows() == 0 { @@ -448,6 +490,24 @@ impl TableWrite { return Ok(()); } + if self.postpone_fixed_bucket.is_some() { + let routed = self + .postpone_fixed_bucket + .as_mut() + .unwrap() + .write_batch(&self.table, &batch) + .await?; + for routed_batch in routed { + self.write_bucket( + routed_batch.partition, + routed_batch.bucket, + routed_batch.batch, + ) + .await?; + } + return Ok(()); + } + let grouped = self.divide_by_partition_bucket(&batch).await?; for ((partition_bytes, bucket), sub_batch) in grouped { self.write_bucket(partition_bytes, bucket, sub_batch) @@ -758,7 +818,6 @@ impl TableWrite { }) } - /// Write a batch directly to the writer for the given (partition, bucket). async fn write_bucket( &mut self, partition_bytes: Vec, @@ -782,8 +841,22 @@ impl TableWrite { } /// Close all writers and collect CommitMessages for use with TableCommit. - /// Writers are cleared after this call, allowing the TableWrite to be reused. + /// Writers are cleared after this call, allowing the TableWrite to be reused, + /// except for fixed-bucket postpone batch writes, which are one-shot. pub async fn prepare_commit(&mut self) -> Result> { + if let Some(writer) = self.postpone_fixed_bucket.as_mut() { + writer.start_prepare()?; + let routed = writer.prepare_batch(&self.table, self.is_overwrite).await?; + for routed_batch in routed { + self.write_bucket( + routed_batch.partition, + routed_batch.bucket, + routed_batch.batch, + ) + .await?; + } + } + let writers: Vec<(PartitionBucketKey, FileWriter)> = self.partition_writers.drain().collect(); @@ -814,6 +887,10 @@ impl TableWrite { || !index_files.is_empty() { let mut msg = CommitMessage::new(partition_bytes, bucket, files.data_files); + msg.total_buckets = self + .postpone_fixed_bucket + .as_ref() + .and_then(|writer| writer.total_buckets(&msg.partition)); msg.new_changelog_files = files.changelog_files; msg.new_index_files = index_files; messages.push(msg); @@ -828,6 +905,9 @@ impl TableWrite { messages.push(msg); } } + if let Some(writer) = self.postpone_fixed_bucket.as_mut() { + writer.finish(); + } Ok(messages) } @@ -3669,6 +3749,25 @@ mod tests { ) } + fn test_fixed_postpone_pk_table(file_io: &FileIO, table_path: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.target-row-num-per-bucket", "2") + .option("postpone.batch-write-fixed-bucket.max-parallelism", "8") + .build() + .unwrap(); + Table::new( + file_io.clone(), + Identifier::new("default", "test_fixed_postpone_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) + } + fn test_postpone_partitioned_schema() -> TableSchema { let schema = Schema::builder() .column("pt", DataType::VarChar(VarCharType::string_type())) @@ -3709,6 +3808,46 @@ mod tests { .unwrap() } + fn make_partition_bucket_plan( + table: &Table, + partitions: Vec<&str>, + total_buckets: Vec, + ) -> PostponeBucketPlan { + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("pt", ArrowDataType::Utf8, false), + ArrowField::new("total_buckets", ArrowDataType::Int32, false), + ])), + vec![ + Arc::new(StringArray::from(partitions)), + Arc::new(Int32Array::from(total_buckets)), + ], + ) + .unwrap(); + PostponeBucketPlan::from_arrow(table, &batch).unwrap() + } + + #[test] + fn test_postpone_bucket_plan_rejects_non_positive_counts() { + let file_io = test_file_io(); + let table = + test_postpone_partitioned_table(&file_io, "memory:/test_postpone_invalid_bucket_plan"); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("pt", ArrowDataType::Utf8, false), + ArrowField::new("total_buckets", ArrowDataType::Int32, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["p"])), + Arc::new(Int32Array::from(vec![0])), + ], + ) + .unwrap(); + + let error = PostponeBucketPlan::from_arrow(&table, &batch).unwrap_err(); + assert!(error.to_string().contains("must be positive")); + } + #[tokio::test] async fn test_postpone_write_and_commit() { let file_io = test_file_io(); @@ -3737,6 +3876,429 @@ mod tests { assert_eq!(snapshot.total_record_count(), Some(3)); } + #[tokio::test] + async fn test_postpone_batch_write_uses_visible_fixed_buckets() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_write"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-user-1".to_string()).unwrap(); + write + .write_arrow_batch(&make_batch(vec![1, 2, 3, 4], vec![10, 20, 30, 40])) + .await + .unwrap(); + let state = write.postpone_fixed_bucket.as_ref().unwrap(); + assert_eq!(state.buffered_partition_count(), 1); + assert!(write.partition_writers.is_empty()); + let messages = write.prepare_commit().await.unwrap(); + assert!(!messages.is_empty()); + assert!(messages.iter().all(|message| message.bucket >= 0)); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + TableCommit::new(table.clone(), "fixed-user-1".to_string()) + .commit(messages) + .await + .unwrap(); + assert_eq!( + read_id_value_rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40)] + ); + + // A later small append reuses the partition's existing bucket count + // instead of inferring a new count from its own input size. + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-user-2".to_string()).unwrap(); + write + .write_arrow_batch(&make_batch(vec![5], vec![50])) + .await + .unwrap(); + // The real-bucket count is loaded on the first write, so existing + // partitions stream into file writers instead of retaining Arrow + // batches until prepare_commit. + let state = write.postpone_fixed_bucket.as_ref().unwrap(); + assert_eq!(state.buffered_partition_count(), 0); + assert!(!write.partition_writers.is_empty()); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + TableCommit::new(table.clone(), "fixed-user-2".to_string()) + .commit(messages) + .await + .unwrap(); + assert_eq!( + read_id_value_rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40), (5, 50)] + ); + } + + #[tokio::test] + async fn test_postpone_distributed_writers_share_bucket_plan() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_shared_bucket_plan"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_partitioned_table(&file_io, table_path); + let plan = make_partition_bucket_plan(&table, vec!["p"], vec![3]); + + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("shared-plan-user") + .unwrap() + .with_postpone_bucket_plan(plan) + .unwrap(); + let mut first = builder.new_write().unwrap(); + let mut second = builder.new_write().unwrap(); + first + .write_arrow_batch(&make_partitioned_batch_3col(vec!["p"], vec![1], vec![10])) + .await + .unwrap(); + second + .write_arrow_batch(&make_partitioned_batch_3col( + vec!["p", "p", "p", "p"], + vec![2, 3, 4, 5], + vec![20, 30, 40, 50], + )) + .await + .unwrap(); + + let mut messages = first.prepare_commit().await.unwrap(); + messages.extend(second.prepare_commit().await.unwrap()); + assert!(!messages.is_empty()); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(3))); + builder.new_commit().commit(messages).await.unwrap(); + } + + #[tokio::test] + async fn test_postpone_provided_plan_must_cover_input_partitions() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_incomplete_bucket_plan"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_partitioned_table(&file_io, table_path); + let plan = make_partition_bucket_plan(&table, vec!["p"], vec![2]); + let mut write = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_postpone_bucket_plan(plan) + .unwrap() + .new_write() + .unwrap(); + + let error = write + .write_arrow_batch(&make_partitioned_batch_3col( + vec!["missing"], + vec![1], + vec![10], + )) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("does not contain an input partition")); + } + + #[tokio::test] + async fn test_postpone_overwrite_allows_bucket_rescale() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_overwrite_bucket_rescale"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_partitioned_table(&file_io, table_path); + + let initial_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("initial-layout") + .unwrap() + .with_postpone_bucket_plan(make_partition_bucket_plan(&table, vec!["p"], vec![1])) + .unwrap(); + let mut initial_write = initial_builder.new_write().unwrap(); + initial_write + .write_arrow_batch(&make_partitioned_batch_3col(vec!["p"], vec![1], vec![10])) + .await + .unwrap(); + initial_builder + .new_commit() + .commit(initial_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let overwrite_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("replacement-layout") + .unwrap() + .with_postpone_bucket_plan(make_partition_bucket_plan(&table, vec!["p"], vec![3])) + .unwrap() + .with_overwrite(); + let mut overwrite_write = overwrite_builder.new_write().unwrap(); + overwrite_write + .write_arrow_batch(&make_partitioned_batch_3col(vec!["p"], vec![2], vec![20])) + .await + .unwrap(); + let messages = overwrite_write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(3))); + overwrite_builder + .new_commit() + .overwrite(messages, None) + .await + .unwrap(); + + let snapshot = SnapshotManager::new(file_io, table_path.to_string()) + .get_latest_snapshot() + .await + .unwrap() + .unwrap(); + let entries = TableScan::new(&table, None, vec![], None, None, None) + .with_scan_all_files() + .plan_manifest_entries(&snapshot) + .await + .unwrap(); + assert!(!entries.is_empty()); + assert!(entries.iter().all(|entry| entry.total_buckets() == 3)); + } + + #[tokio::test] + async fn test_postpone_row_target_ignores_invalid_size_target() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_row_target_precedence"; + setup_dirs(&file_io, table_path).await; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.target-row-num-per-bucket", "2") + .option("postpone.target-size-per-bucket", "invalid") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_row_target_precedence"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "row-target".to_string()).unwrap(); + write + .write_arrow_batch(&make_batch(vec![1, 2, 3], vec![10, 20, 30])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + } + + #[tokio::test] + async fn test_postpone_fixed_bucket_write_with_rowkind_field() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_rowkind"; + setup_dirs(&file_io, table_path).await; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .column("op", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("rowkind.field", "op") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_fixed_bucket_rowkind"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("value", ArrowDataType::Int32, false), + ArrowField::new("op", ArrowDataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), + Arc::new(StringArray::from(vec!["+I"])), + ], + ) + .unwrap(); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-rowkind".to_string()).unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].bucket, 0); + assert_eq!(messages[0].total_buckets, Some(1)); + assert_eq!(messages[0].new_files[0].row_count, 1); + } + + #[tokio::test] + async fn test_postpone_bucket_count_uses_java_binary_row_size() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_java_binary_row_size"; + setup_dirs(&file_io, table_path).await; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("left", DataType::VarChar(VarCharType::string_type())) + .column("right", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.target-size-per-bucket", "20 kb") + .option("postpone.batch-write-fixed-bucket.max-parallelism", "8") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_java_binary_row_size"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let row_count = 1_000; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("left", ArrowDataType::Utf8, false), + ArrowField::new("right", ArrowDataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from_iter_values(0..row_count)), + Arc::new(StringArray::from(vec!["a"; row_count as usize])), + Arc::new(StringArray::from(vec!["b"; row_count as usize])), + ], + ) + .unwrap(); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-size-user".to_string()).unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + + // Java BinaryRows are 32,000 bytes, so a 20 KiB target plans two + // buckets. Arrow buffer sizing would incorrectly plan one. + assert!(!messages.is_empty()); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + } + + #[tokio::test] + async fn test_postpone_fixed_bucket_batch_write_is_one_shot() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_one_shot"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-one-shot".to_string()).unwrap(); + + write + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(1))); + + let write_error = write + .write_arrow_batch(&make_batch(vec![2, 3, 4, 5], vec![20, 30, 40, 50])) + .await + .unwrap_err(); + assert!( + matches!(write_error, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supports one prepare_commit call") + && message.contains("create a new writer")) + ); + + let prepare_error = write.prepare_commit().await.unwrap_err(); + assert!( + matches!(prepare_error, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supports one prepare_commit call") + && message.contains("create a new writer")) + ); + } + + #[test] + fn test_postpone_fixed_bucket_rejects_deletion_vectors() { + let file_io = test_file_io(); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("deletion-vectors.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_dv"), + "memory:/test_postpone_dv".to_string(), + TableSchema::new(0, &schema), + None, + ); + + let error = TableWrite::new_postpone_fixed_bucket(&table, "test-user".to_string()) + .err() + .expect("fixed-bucket postpone writes must reject deletion vectors"); + assert!(matches!(error, crate::Error::Unsupported { ref message } + if message.contains("postpone fixed-bucket writes") + && message.contains("deletion-vectors.enabled=true"))); + } + + #[tokio::test] + async fn test_postpone_batch_write_rejects_conflicting_bucket_counts() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_conflict"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + + // Both writers plan against the empty table. Their differently sized + // inputs produce different bucket counts for the same partition. + let mut first = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-conflict-1".to_string()).unwrap(); + first + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let first_messages = first.prepare_commit().await.unwrap(); + assert!(first_messages + .iter() + .all(|message| message.total_buckets == Some(1))); + + let mut second = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-conflict-2".to_string()).unwrap(); + second + .write_arrow_batch(&make_batch(vec![2, 3, 4, 5], vec![20, 30, 40, 50])) + .await + .unwrap(); + let second_messages = second.prepare_commit().await.unwrap(); + assert!(second_messages + .iter() + .all(|message| message.total_buckets == Some(2))); + + TableCommit::new(table.clone(), "fixed-conflict-1".to_string()) + .commit(first_messages) + .await + .unwrap(); + let error = TableCommit::new(table, "fixed-conflict-2".to_string()) + .commit(second_messages) + .await + .unwrap_err(); + assert!(error.to_string().contains("Postpone fixed-bucket conflict")); + } + #[tokio::test] async fn test_postpone_write_empty_batch() { let file_io = test_file_io(); diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 52ba57754..c79afd485 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -20,7 +20,10 @@ //! Reference: [pypaimon WriteBuilder](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/write/write_builder.py) use super::format_write_builder::FormatWriteBuilder; -use crate::table::{DataEvolutionDeleteWriter, Table, TableCommit, TableUpdate, TableWrite}; +use crate::spec::{CoreOptions, POSTPONE_BUCKET}; +use crate::table::{ + DataEvolutionDeleteWriter, PostponeBucketPlan, Table, TableCommit, TableUpdate, TableWrite, +}; use uuid::Uuid; /// Builder for creating table writers and committers. @@ -43,6 +46,20 @@ impl<'a> WriteBuilder<'a> { } } + /// Create a builder which forces one-shot fixed-bucket writes for a + /// postpone table, even when `postpone.batch-write-fixed-bucket=false`. + pub fn new_postpone_fixed_bucket(table: &'a Table) -> crate::Result { + if table.is_format_table() { + return Err(crate::Error::Unsupported { + message: "Postpone fixed-bucket writes are only supported for Paimon tables" + .to_string(), + }); + } + Ok(Self(WriteBuilderKind::Paimon( + PaimonWriteBuilder::new(table).with_postpone_fixed_bucket(), + ))) + } + /// Get the commit user shared by writers and committers created by this builder. pub fn commit_user(&self) -> &str { match &self.0 { @@ -51,6 +68,15 @@ impl<'a> WriteBuilder<'a> { } } + /// Whether writers created by this builder use the postpone fixed-bucket + /// batch path. + pub fn uses_postpone_fixed_bucket(&self) -> bool { + match &self.0 { + WriteBuilderKind::Paimon(builder) => builder.postpone_fixed_bucket, + WriteBuilderKind::Format(_) => false, + } + } + /// Set the commit user shared by writers and committers created by this builder. pub fn with_commit_user(self, commit_user: impl Into) -> crate::Result { match self.0 { @@ -75,6 +101,19 @@ impl<'a> WriteBuilder<'a> { } } + /// Supply a bucket plan shared by all distributed postpone writers in one + /// logical batch. The plan must cover every input partition. + pub fn with_postpone_bucket_plan(self, bucket_plan: PostponeBucketPlan) -> crate::Result { + match self.0 { + WriteBuilderKind::Paimon(builder) => Ok(Self(WriteBuilderKind::Paimon( + builder.with_postpone_bucket_plan(bucket_plan)?, + ))), + WriteBuilderKind::Format(_) => Err(crate::Error::Unsupported { + message: "Postpone bucket plans are only supported for Paimon tables".to_string(), + }), + } + } + /// Create a new TableCommit for committing write results. pub fn new_commit(&self) -> TableCommit { match &self.0 { @@ -120,17 +159,42 @@ struct PaimonWriteBuilder<'a> { table: &'a Table, commit_user: String, overwrite: bool, + postpone_fixed_bucket: bool, + postpone_bucket_plan: Option, } impl<'a> PaimonWriteBuilder<'a> { pub fn new(table: &'a Table) -> Self { + let schema = table.schema(); + let options = CoreOptions::new(schema.options()); + let postpone_fixed_bucket = options.bucket() == POSTPONE_BUCKET + && !schema.primary_keys().is_empty() + && options.postpone_batch_write_fixed_bucket(); Self { table, commit_user: Uuid::new_v4().to_string(), overwrite: false, + postpone_fixed_bucket, + postpone_bucket_plan: None, } } + fn with_postpone_fixed_bucket(mut self) -> Self { + self.postpone_fixed_bucket = true; + self + } + + fn with_postpone_bucket_plan(mut self, bucket_plan: PostponeBucketPlan) -> crate::Result { + if !self.postpone_fixed_bucket { + return Err(crate::Error::Unsupported { + message: "A postpone bucket plan requires an explicit postpone fixed-bucket write builder" + .to_string(), + }); + } + self.postpone_bucket_plan = Some(bucket_plan); + Ok(self) + } + /// Get the commit user shared by writers and committers created by this builder. /// /// This value is persisted in snapshot metadata and used for duplicate @@ -198,7 +262,20 @@ impl<'a> PaimonWriteBuilder<'a> { .to_string(), }); } - let write = TableWrite::new(self.table, self.commit_user.clone())?; + let write = if self.postpone_fixed_bucket { + match self.postpone_bucket_plan.clone() { + Some(bucket_plan) => TableWrite::new_postpone_fixed_bucket_with_plan( + self.table, + self.commit_user.clone(), + bucket_plan, + )?, + None => { + TableWrite::new_postpone_fixed_bucket(self.table, self.commit_user.clone())? + } + } + } else { + TableWrite::new(self.table, self.commit_user.clone())? + }; Ok(if self.overwrite { write.with_overwrite() } else { @@ -251,6 +328,7 @@ mod tests { }; use arrow_array::{Int32Array, Int64Array, RecordBatch, StringArray}; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; + use std::collections::HashMap; use std::sync::Arc; fn test_file_io() -> FileIO { @@ -395,7 +473,13 @@ mod tests { let table_path = "memory:/test_write_builder_commit_user"; setup_dirs(&file_io, table_path).await; - let table = test_postpone_pk_table(&file_io, table_path); + // Keep this test on the legacy postpone writer because its file-name + // contract embeds commit_user; fixed-bucket files use regular names. + let table = + test_postpone_pk_table(&file_io, table_path).copy_with_options(HashMap::from([( + "postpone.batch-write-fixed-bucket".to_string(), + "false".to_string(), + )])); let wb = table .new_write_builder() .with_commit_user("my-commit-user") @@ -410,6 +494,7 @@ mod tests { let messages = write.prepare_commit().await.unwrap(); assert_eq!(messages[0].bucket, POSTPONE_BUCKET); + assert_eq!(messages[0].total_buckets, None); assert!( messages[0].new_files[0] .file_name @@ -430,6 +515,53 @@ mod tests { assert_eq!(snapshot.commit_user(), "my-commit-user"); } + #[tokio::test] + async fn test_postpone_fixed_bucket_builder_respects_option() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_builder_option"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + let mut default_write = table.new_write_builder().new_write().unwrap(); + default_write + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let default_messages = default_write.prepare_commit().await.unwrap(); + assert_eq!(default_messages.len(), 1); + assert_eq!(default_messages[0].bucket, 0); + assert_eq!(default_messages[0].total_buckets, Some(1)); + + let legacy_table = table.copy_with_options(HashMap::from([( + "postpone.batch-write-fixed-bucket".to_string(), + "false".to_string(), + )])); + let mut legacy_write = legacy_table.new_write_builder().new_write().unwrap(); + legacy_write + .write_arrow_batch(&make_batch(vec![2], vec![20])) + .await + .unwrap(); + let legacy_messages = legacy_write.prepare_commit().await.unwrap(); + assert_eq!(legacy_messages.len(), 1); + assert_eq!(legacy_messages[0].bucket, POSTPONE_BUCKET); + assert_eq!(legacy_messages[0].total_buckets, None); + + let builder = legacy_table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("explicit-fixed-user") + .unwrap(); + let mut fixed = builder.new_write().unwrap(); + fixed + .write_arrow_batch(&make_batch(vec![3], vec![30])) + .await + .unwrap(); + let fixed_messages = fixed.prepare_commit().await.unwrap(); + assert_eq!(fixed_messages.len(), 1); + assert_eq!(fixed_messages[0].bucket, 0); + assert_eq!(fixed_messages[0].total_buckets, Some(1)); + } + #[tokio::test] async fn test_branch_reference_rejects_write_and_index_builders() { let table = as_main_branch_reference(test_postpone_pk_table(