diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index d29c7eb85..152714bb3 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -38,7 +38,7 @@ use arrow_array::{Array, Int32Array, RecordBatch, StringArray, StructArray}; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use paimon::catalog::Identifier; use paimon::io::FileIOBuilder; -use paimon::spec::{DataType, IntType, Schema, TableSchema, VarCharType}; +use paimon::spec::{CommitKind, DataType, IntType, Schema, TableSchema, VarCharType}; use paimon::table::{SnapshotManager, Table}; use crate::error::*; @@ -73,6 +73,19 @@ fn not_null_table_schema() -> TableSchema { 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 +117,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), @@ -1271,6 +1316,101 @@ fn test_commit_rejects_messages_from_different_builder_identity() { } } +#[test] +fn test_commit_rejects_mismatched_write_kind_and_overwrite_mode() { + let path = "memory:/test_commit_write_context"; + let file_io = memory_file_io(); + setup_table_dirs(&file_io, path); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test"), + path.to_string(), + partitioned_postpone_table_schema(), + None, + ); + let handle = unsafe { wrap_table(table) }; + let commit_user = CString::new("write-context-job").unwrap(); + + unsafe { + let standard_wb = + paimon_table_new_write_builder_with_commit_user(handle, commit_user.as_ptr()) + .write_builder; + let standard_tw = paimon_write_builder_new_write(standard_wb).write; + let (array, schema) = export_batch_to_ffi(make_partitioned_write_batch( + vec!["p"], + vec![1], + vec!["standard"], + )); + assert!(paimon_table_write_write_arrow_batch( + standard_tw, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ) + .is_null()); + let standard_messages = paimon_table_write_prepare_commit(standard_tw).messages; + + let fixed_wb = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ) + .write_builder; + assert!(paimon_write_builder_with_overwrite(fixed_wb).is_null()); + let (array, schema) = + export_batch_to_ffi(make_postpone_bucket_plan_batch(vec!["p"], vec![1])); + assert!(paimon_write_builder_with_postpone_bucket_plan( + fixed_wb, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ) + .is_null()); + let fixed_tw = paimon_write_builder_new_write(fixed_wb).write; + let (array, schema) = export_batch_to_ffi(make_partitioned_write_batch( + vec!["p"], + vec![1], + vec!["fixed"], + )); + assert!(paimon_table_write_write_arrow_batch( + fixed_tw, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ) + .is_null()); + let fixed_messages = paimon_table_write_prepare_commit(fixed_tw).messages; + + let error = paimon_commit_messages_merge(standard_messages, fixed_messages); + assert!(!error.is_null()); + assert!(error_message(error).contains("write kind and overwrite mode")); + paimon_error_free(error); + + let fixed_commit = paimon_write_builder_new_commit(fixed_wb).commit; + let error = paimon_table_commit_commit(fixed_commit, standard_messages); + assert!(!error.is_null()); + assert!(error_message(error).contains("different write kind or overwrite mode")); + paimon_error_free(error); + + let standard_commit = paimon_write_builder_new_commit(standard_wb).commit; + let error = paimon_table_commit_overwrite(standard_commit, standard_messages); + assert!(!error.is_null()); + assert!(error_message(error).contains("append messages cannot be committed")); + paimon_error_free(error); + + assert!(crate::runtime() + .block_on(SnapshotManager::new(file_io, path.to_string()).get_latest_snapshot()) + .unwrap() + .is_none()); + + paimon_table_commit_free(standard_commit); + paimon_table_commit_free(fixed_commit); + paimon_commit_messages_free(fixed_messages); + paimon_commit_messages_free(standard_messages); + paimon_table_write_free(fixed_tw); + paimon_write_builder_free(fixed_wb); + paimon_table_write_free(standard_tw); + paimon_write_builder_free(standard_wb); + unwrap_table(handle); + } +} + #[test] fn test_commit_messages_live_until_explicit_free() { const CHILD_ENV: &str = "PAIMON_C_MESSAGES_LIFETIME_CHILD"; @@ -1442,6 +1582,198 @@ 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 normal = paimon_table_new_write_builder(handle); + assert!(normal.error.is_null()); + assert!(matches!( + (&*((*normal.write_builder).inner as *const WriteBuilderState)).kind, + WriteBuilderKind::Standard + )); + paimon_write_builder_free(normal.write_builder); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(handle); + assert!(fixed.error.is_null()); + assert!(matches!( + (&*((*fixed.write_builder).inner as *const WriteBuilderState)).kind, + WriteBuilderKind::PostponeFixed { .. } + )); + let write = paimon_write_builder_new_write(fixed.write_builder); + assert!(write.write.is_null()); + assert!(error_message(write.error).contains("bucket plan is required")); + paimon_error_free(write.error); + paimon_write_builder_free(fixed.write_builder); + + 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] { + assert!(paimon_write_builder_with_overwrite(wb).is_null()); + let (array, schema) = export_batch_to_ffi(make_postpone_bucket_plan_batch( + vec!["p1", "p2"], + vec![3, 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, partitions, ids, names) in [ + (tw1, vec!["p1"], vec![1], vec!["a"]), + ( + tw2, + vec!["p2", "p2", "p2", "p2"], + vec![2, 3, 4, 5], + vec!["b", "c", "d", "e"], + ), + ] { + let (array, schema) = + export_batch_to_ffi(make_partitioned_write_batch(partitions, 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; + assert!(matches!( + (&*((*commit).inner as *const TableCommitState)).commit, + TableCommitKind::PostponeFixed(_) + )); + let error = paimon_table_commit_commit(commit, messages1); + assert!(error.is_null()); + let snapshot = crate::runtime() + .block_on( + SnapshotManager::new(table_ref(handle).file_io().clone(), path.to_string()) + .get_latest_snapshot(), + ) + .unwrap() + .unwrap(); + assert_eq!(snapshot.commit_kind(), &CommitKind::OVERWRITE); + + 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_distributed_postpone_writers_reject_overlapping_bucket_ownership() { + let path = "memory:/test_distributed_postpone_overlapping_ownership"; + 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("overlapping-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![1])); + assert!(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, + ) + .is_null()); + } + + let tw1 = paimon_write_builder_new_write(wb1).write; + let tw2 = paimon_write_builder_new_write(wb2).write; + for (tw, name) in [(tw1, "first"), (tw2, "second")] { + let (array, schema) = + export_batch_to_ffi(make_partitioned_write_batch(vec!["p"], vec![1], vec![name])); + assert!(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, + ) + .is_null()); + } + + let messages1 = paimon_table_write_prepare_commit(tw1).messages; + let messages2 = paimon_table_write_prepare_commit(tw2).messages; + 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()); + assert!(error_message(error).contains("writer ownership conflict for bucket 0")); + paimon_error_free(error); + + 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 +2053,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 +2068,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..7263dd179 100644 --- a/bindings/c/src/types.rs +++ b/bindings/c/src/types.rs @@ -20,7 +20,10 @@ 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, PostponeFixedBucketTableCommit, + PostponeFixedBucketTableWrite, Table, TableCommit, TableWrite, +}; /// C-compatible key-value pair for options. #[repr(C)] @@ -205,28 +208,60 @@ pub struct paimon_arrow_batch { // === Write/Commit opaque types === -/// Internal state for WriteBuilder that stores table, shared commit_user, and overwrite flag. +pub(crate) enum WriteBuilderKind { + Standard, + PostponeFixed { + bucket_plan: Option, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum WriteKind { + Standard, + PostponeFixed, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct WriteContext { + pub kind: WriteKind, + pub overwrite: bool, +} + pub(crate) struct WriteBuilderState { pub table: Table, pub commit_user: String, pub overwrite: bool, + pub kind: WriteBuilderKind, +} + +pub(crate) enum TableWriteKind { + Standard(Box), + PostponeFixed(Box), } pub(crate) struct TableWriteState { - pub write: TableWrite, + pub write: TableWriteKind, + pub context: WriteContext, pub target_schema: Arc, pub table_location: String, pub commit_user: String, } +pub(crate) enum TableCommitKind { + Standard(TableCommit), + PostponeFixed(PostponeFixedBucketTableCommit), +} + pub(crate) struct TableCommitState { - pub commit: TableCommit, + pub commit: TableCommitKind, + pub context: WriteContext, pub table_location: String, pub commit_user: String, } pub(crate) struct CommitMessagesState { pub messages: Vec, + pub context: WriteContext, pub table_location: String, pub commit_user: String, } diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs index 8e9637997..e79223a19 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, + kind: WriteBuilderKind, ) -> paimon_result_write_builder { if let Err(e) = check_non_null(table, "table") { return paimon_result_write_builder { @@ -45,23 +46,14 @@ unsafe fn new_write_builder( }; } let table_ref = &*((*table).inner as *const Table); - let builder = table_ref.new_write_builder(); - let commit_user = match commit_user { - Some(commit_user) => match builder.with_commit_user(commit_user) { - Ok(builder) => builder.commit_user().to_string(), - Err(e) => { - return paimon_result_write_builder { - write_builder: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } + let state = match create_write_builder_state(table_ref, commit_user, kind) { + Ok(state) => state, + Err(error) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error: paimon_error::from_paimon(error), } - }, - None => builder.commit_user().to_string(), - }; - let state = WriteBuilderState { - table: table_ref.clone(), - commit_user, - overwrite: false, + } }; let inner = Box::into_raw(Box::new(state)) as *mut c_void; paimon_result_write_builder { @@ -70,6 +62,55 @@ unsafe fn new_write_builder( } } +fn create_write_builder_state( + table: &Table, + commit_user: Option, + kind: WriteBuilderKind, +) -> paimon::Result { + let commit_user = match &kind { + WriteBuilderKind::Standard => { + let builder = table.new_write_builder(); + match commit_user { + Some(commit_user) => builder + .with_commit_user(commit_user)? + .commit_user() + .to_string(), + None => builder.commit_user().to_string(), + } + } + WriteBuilderKind::PostponeFixed { .. } => { + let builder = table.new_postpone_fixed_bucket_write_builder()?; + match commit_user { + Some(commit_user) => builder + .with_commit_user(commit_user)? + .commit_user() + .to_string(), + None => builder.commit_user().to_string(), + } + } + }; + Ok(WriteBuilderState { + table: table.clone(), + commit_user, + overwrite: false, + kind, + }) +} + +unsafe fn new_write_builder_with_commit_user( + table: *const paimon_table, + commit_user: *const c_char, + kind: WriteBuilderKind, +) -> paimon_result_write_builder { + match validate_cstr(commit_user, "commit_user") { + Ok(commit_user) => new_write_builder(table, Some(commit_user), kind), + Err(error) => paimon_result_write_builder { + write_builder: ptr::null_mut(), + error, + }, + } +} + /// Create a new WriteBuilder from a Table. /// /// The returned WriteBuilder holds a shared `commit_user` (UUID) that will be @@ -82,7 +123,23 @@ 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, WriteBuilderKind::Standard) +} + +/// Create a one-shot fixed-bucket WriteBuilder for a postpone table. +/// A bucket plan must be set before creating a writer. +/// +/// # 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, + WriteBuilderKind::PostponeFixed { bucket_plan: None }, + ) } /// Create a WriteBuilder with a caller-provided stable commit identity. @@ -98,16 +155,25 @@ pub unsafe extern "C" fn paimon_table_new_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)) + new_write_builder_with_commit_user(table, commit_user, WriteBuilderKind::Standard) +} + +/// Create a fixed-bucket WriteBuilder with a stable commit identity. +/// A bucket plan must be set before creating a writer. +/// +/// # 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 { + new_write_builder_with_commit_user( + table, + commit_user, + WriteBuilderKind::PostponeFixed { bucket_plan: None }, + ) } /// Free a paimon_write_builder. @@ -126,8 +192,8 @@ pub unsafe extern "C" fn paimon_write_builder_free(wb: *mut paimon_write_builder /// Enable overwrite mode for the WriteBuilder. /// -/// In overwrite mode, a subsequent `paimon_table_commit_overwrite` will replace -/// the data in the written partitions rather than appending. +/// Fixed-bucket committers carry this mode. Standard callers use +/// `paimon_table_commit_overwrite`. /// /// # Safety /// `wb` must be a valid pointer from `paimon_table_new_write_builder`, or null (returns error). @@ -143,6 +209,48 @@ pub unsafe extern "C" fn paimon_write_builder_with_overwrite( ptr::null_mut() } +/// Set a shared `partition -> total_buckets` plan. +/// Ownership of `array` and `schema` is transferred to this call; the caller +/// must not release the Arrow C Data structs afterward. +/// +/// # Safety +/// `wb` must be an explicit postpone fixed-bucket builder. `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); + let WriteBuilderKind::PostponeFixed { bucket_plan } = &mut state.kind else { + return invalid_input( + "a postpone bucket plan requires an explicit 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) => { + *bucket_plan = Some(plan); + ptr::null_mut() + } + Err(error) => paimon_error::from_paimon(error), + } +} + // ======================= TableWrite =============================== fn invalid_input(message: impl Into) -> *mut paimon_error { @@ -244,47 +352,22 @@ 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()) - { - Ok(b) => b, - Err(e) => { - return paimon_result_table_write { - write: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } - } - }; - - if state.overwrite { - builder = builder.with_overwrite(); - } - - let tw = match builder.new_write() { - Ok(w) => w, - Err(e) => { + let result = create_table_write(state).and_then(|write| { + paimon::arrow::build_target_arrow_schema(state.table.schema().fields()) + .map(|schema| (write, schema)) + }); + let (write, target_schema) = match result { + Ok(result) => result, + Err(error) => { return paimon_result_table_write { write: ptr::null_mut(), - error: paimon_error::from_paimon(e), + error: paimon_error::from_paimon(error), } } }; - - let target_schema = - match paimon::arrow::build_target_arrow_schema(state.table.schema().fields()) { - Ok(schema) => schema, - Err(e) => { - return paimon_result_table_write { - write: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } - } - }; let table_write = TableWriteState { - write: tw, + write, + context: write_context(state), target_schema, table_location: state.table.location().to_string(), commit_user: state.commit_user.clone(), @@ -296,6 +379,51 @@ pub unsafe extern "C" fn paimon_write_builder_new_write( } } +fn create_table_write(state: &WriteBuilderState) -> paimon::Result { + match &state.kind { + WriteBuilderKind::Standard => { + let mut builder = state + .table + .new_write_builder() + .with_commit_user(state.commit_user.clone())?; + if state.overwrite { + builder = builder.with_overwrite(); + } + builder + .new_write() + .map(Box::new) + .map(TableWriteKind::Standard) + } + WriteBuilderKind::PostponeFixed { bucket_plan } => { + let mut builder = state + .table + .new_postpone_fixed_bucket_write_builder() + .and_then(|builder| builder.with_commit_user(state.commit_user.clone()))?; + if let Some(plan) = bucket_plan.clone() { + builder = builder.with_bucket_plan(plan); + } + if state.overwrite { + builder = builder.with_overwrite(); + } + builder + .new_write() + .map(Box::new) + .map(TableWriteKind::PostponeFixed) + } + } +} + +fn write_context(state: &WriteBuilderState) -> WriteContext { + let kind = match &state.kind { + WriteBuilderKind::Standard => WriteKind::Standard, + WriteBuilderKind::PostponeFixed { .. } => WriteKind::PostponeFixed, + }; + WriteContext { + kind, + overwrite: state.overwrite, + } +} + /// Free a paimon_table_write. /// /// Dropping a TableWrite before calling `prepare_commit` discards any @@ -349,7 +477,11 @@ pub unsafe extern "C" fn paimon_table_write_write_arrow_batch( return error; } - match runtime().block_on(table_write.write.write_arrow_batch(&batch)) { + let result = match &mut table_write.write { + TableWriteKind::Standard(write) => runtime().block_on(write.write_arrow_batch(&batch)), + TableWriteKind::PostponeFixed(write) => runtime().block_on(write.write_arrow_batch(&batch)), + }; + match result { Ok(()) => ptr::null_mut(), Err(e) => paimon_error::from_paimon(e), } @@ -358,8 +490,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 @@ -379,10 +512,15 @@ pub unsafe extern "C" fn paimon_table_write_prepare_commit( } let table_write = &mut *((*tw).inner as *mut TableWriteState); - match runtime().block_on(table_write.write.prepare_commit()) { + let result = match &mut table_write.write { + TableWriteKind::Standard(write) => runtime().block_on(write.prepare_commit()), + TableWriteKind::PostponeFixed(write) => runtime().block_on(write.prepare_commit()), + }; + match result { Ok(messages) => { let messages = CommitMessagesState { messages, + context: table_write.context, table_location: table_write.table_location.clone(), commit_user: table_write.commit_user.clone(), }; @@ -401,6 +539,27 @@ pub unsafe extern "C" fn paimon_table_write_prepare_commit( // ======================= TableCommit =============================== +fn create_table_commit(state: &WriteBuilderState) -> paimon::Result { + match &state.kind { + WriteBuilderKind::Standard => state + .table + .new_write_builder() + .with_commit_user(state.commit_user.clone())? + .try_new_commit() + .map(TableCommitKind::Standard), + WriteBuilderKind::PostponeFixed { .. } => { + let mut builder = state + .table + .new_postpone_fixed_bucket_write_builder()? + .with_commit_user(state.commit_user.clone())?; + if state.overwrite { + builder = builder.with_overwrite(); + } + builder.try_new_commit().map(TableCommitKind::PostponeFixed) + } + } +} + /// Create a new TableCommit from the WriteBuilder. /// /// The committer shares the same `commit_user` as the writer, which is @@ -420,32 +579,19 @@ pub unsafe extern "C" fn paimon_write_builder_new_commit( } let state = &*((*wb).inner as *const WriteBuilderState); - let builder = match state - .table - .new_write_builder() - .with_commit_user(state.commit_user.clone()) - { - Ok(b) => b, - Err(e) => { - return paimon_result_table_commit { - commit: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } - } - }; - - let tc = match builder.try_new_commit() { - Ok(c) => c, - Err(e) => { + let commit = match create_table_commit(state) { + Ok(commit) => commit, + Err(error) => { return paimon_result_table_commit { commit: ptr::null_mut(), - error: paimon_error::from_paimon(e), + error: paimon_error::from_paimon(error), } } }; let table_commit = TableCommitState { - commit: tc, + commit, + context: write_context(state), table_location: state.table.location().to_string(), commit_user: state.commit_user.clone(), }; @@ -515,6 +661,12 @@ pub unsafe extern "C" fn paimon_commit_messages_merge( "commit messages can only be merged when table and commit_user both match", ); } + if target.context != source.context { + return invalid_input(format!( + "commit messages can only be merged when write kind and overwrite mode both match (target {:?}, source {:?})", + target.context, source.context + )); + } target.messages.extend(source.messages.clone()); ptr::null_mut() } @@ -536,10 +688,59 @@ fn validate_commit_context( "commit messages were prepared with a different commit_user", )); } + if commit.context != messages.context { + return Err(invalid_input(format!( + "commit messages were prepared with a different write kind or overwrite mode (message {:?}, committer {:?})", + messages.context, commit.context + ))); + } Ok(()) } -/// Commit the given messages in APPEND mode. +unsafe fn commit_with_identifier_impl( + tc: *const paimon_table_commit, + msgs: *mut paimon_commit_messages, + commit_identifier: i64, + filter_committed: bool, +) -> *mut paimon_error { + if let Err(error) = check_non_null(tc, "tc") { + return error; + } + if let Err(error) = check_non_null(msgs, "msgs") { + return error; + } + + let table_commit = &*((*tc).inner as *const TableCommitState); + let messages = &*((*msgs).inner as *const CommitMessagesState); + if let Err(error) = validate_commit_context(table_commit, messages) { + return error; + } + if messages.context.kind == WriteKind::Standard && messages.context.overwrite { + return invalid_input( + "standard overwrite messages must be committed with paimon_table_commit_overwrite", + ); + } + + let messages = messages.messages.clone(); + let result = match (&table_commit.commit, filter_committed) { + (TableCommitKind::Standard(commit), true) => runtime() + .block_on(commit.filter_and_commit_with_identifier(messages, commit_identifier)), + (TableCommitKind::Standard(commit), false) => { + runtime().block_on(commit.commit_with_identifier(messages, commit_identifier)) + } + (TableCommitKind::PostponeFixed(commit), true) => runtime() + .block_on(commit.filter_and_commit_with_identifier(messages, commit_identifier)), + (TableCommitKind::PostponeFixed(commit), false) => { + runtime().block_on(commit.commit_with_identifier(messages, commit_identifier)) + } + }; + match result { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), + } +} + +/// Commit in append mode, or in the configured fixed-bucket overwrite mode. /// /// Empty messages is a no-op success. /// The caller retains ownership of `msgs`; it may retry after an error and @@ -570,27 +771,7 @@ pub unsafe extern "C" fn paimon_table_commit_commit_with_identifier( msgs: *mut paimon_commit_messages, commit_identifier: i64, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tc, "tc") { - return e; - } - if let Err(e) = check_non_null(msgs, "msgs") { - return e; - } - - let table_commit = &*((*tc).inner as *const TableCommitState); - let messages = &*((*msgs).inner as *const CommitMessagesState); - if let Err(error) = validate_commit_context(table_commit, messages) { - return error; - } - - match runtime().block_on( - table_commit - .commit - .commit_with_identifier(messages.messages.clone(), commit_identifier), - ) { - Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), - } + commit_with_identifier_impl(tc, msgs, commit_identifier, false) } /// Filter a previously committed identifier, then commit if it is new. @@ -605,27 +786,7 @@ pub unsafe extern "C" fn paimon_table_commit_filter_and_commit_with_identifier( msgs: *mut paimon_commit_messages, commit_identifier: i64, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tc, "tc") { - return e; - } - if let Err(e) = check_non_null(msgs, "msgs") { - return e; - } - - let table_commit = &*((*tc).inner as *const TableCommitState); - let messages = &*((*msgs).inner as *const CommitMessagesState); - if let Err(error) = validate_commit_context(table_commit, messages) { - return error; - } - - match runtime().block_on( - table_commit - .commit - .filter_and_commit_with_identifier(messages.messages.clone(), commit_identifier), - ) { - Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), - } + commit_with_identifier_impl(tc, msgs, commit_identifier, true) } /// Commit in OVERWRITE mode, replacing data in the written partitions. @@ -675,20 +836,26 @@ unsafe fn paimon_table_commit_overwrite_impl( if let Err(error) = validate_commit_context(table_commit, messages) { return error; } + if !messages.context.overwrite { + return invalid_input( + "append messages cannot be committed with paimon_table_commit_overwrite", + ); + } - let result = match commit_identifier { - Some(commit_identifier) => { - runtime().block_on(table_commit.commit.overwrite_with_identifier( - messages.messages.clone(), - None, - commit_identifier, - )) + let messages = messages.messages.clone(); + let result = match (&table_commit.commit, commit_identifier) { + (TableCommitKind::Standard(commit), Some(commit_identifier)) => { + runtime().block_on(commit.overwrite_with_identifier(messages, None, commit_identifier)) + } + (TableCommitKind::Standard(commit), None) => { + runtime().block_on(commit.overwrite(messages, None)) + } + (TableCommitKind::PostponeFixed(commit), Some(commit_identifier)) => { + runtime().block_on(commit.overwrite_with_identifier(messages, commit_identifier)) + } + (TableCommitKind::PostponeFixed(commit), None) => { + runtime().block_on(commit.overwrite(messages)) } - None => runtime().block_on( - table_commit - .commit - .overwrite(messages.messages.clone(), None), - ), }; match result { Ok(()) => ptr::null_mut(), @@ -729,13 +896,17 @@ unsafe fn paimon_table_commit_truncate_table_impl( let table_commit = &*((*tc).inner as *const TableCommitState); - let result = match commit_identifier { - Some(commit_identifier) => runtime().block_on( - table_commit - .commit - .truncate_table_with_identifier(commit_identifier), - ), - None => runtime().block_on(table_commit.commit.truncate_table()), + let result = match (&table_commit.commit, commit_identifier) { + (TableCommitKind::Standard(commit), Some(commit_identifier)) => { + runtime().block_on(commit.truncate_table_with_identifier(commit_identifier)) + } + (TableCommitKind::Standard(commit), None) => runtime().block_on(commit.truncate_table()), + (TableCommitKind::PostponeFixed(commit), Some(commit_identifier)) => { + runtime().block_on(commit.truncate_table_with_identifier(commit_identifier)) + } + (TableCommitKind::PostponeFixed(commit), None) => { + runtime().block_on(commit.truncate_table()) + } }; match result { Ok(()) => ptr::null_mut(), @@ -771,7 +942,13 @@ pub unsafe extern "C" fn paimon_table_commit_abort( return error; } - match runtime().block_on(table_commit.commit.abort(&messages.messages)) { + let result = match &table_commit.commit { + TableCommitKind::Standard(commit) => runtime().block_on(commit.abort(&messages.messages)), + TableCommitKind::PostponeFixed(commit) => { + runtime().block_on(commit.abort(&messages.messages)) + } + }; + match result { Ok(()) => ptr::null_mut(), Err(e) => paimon_error::from_paimon(e), } @@ -781,12 +958,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/paimon/src/lib.rs b/crates/paimon/src/lib.rs index e372a40ba..a86a95e92 100644 --- a/crates/paimon/src/lib.rs +++ b/crates/paimon/src/lib.rs @@ -50,9 +50,10 @@ pub use catalog::FileSystemCatalog; pub use table::{ CommitMessage, DataEvolutionDeleteWriter, DataEvolutionWriter, DataSplit, DataSplitBuilder, DeletionFile, IncrementalPlan, IncrementalScan, IncrementalScanMode, IncrementalSplit, - PartitionBucket, Plan, RESTEnv, RESTSnapshotCommit, ReadBuilder, RenamingSnapshotCommit, - RowRange, ScanTrace, SnapshotCommit, SnapshotManager, Table, TableCommit, TableRead, TableScan, - TableUpdate, TableWrite, TagManager, WriteBuilder, + PartitionBucket, Plan, PostponeBucketPlan, PostponeFixedBucketTableCommit, + PostponeFixedBucketTableWrite, RESTEnv, RESTSnapshotCommit, ReadBuilder, + RenamingSnapshotCommit, RowRange, ScanTrace, SnapshotCommit, SnapshotManager, Table, + TableCommit, TableRead, TableScan, TableUpdate, TableWrite, TagManager, WriteBuilder, }; pub use table::{ diff --git a/crates/paimon/src/table/commit_message.rs b/crates/paimon/src/table/commit_message.rs index 5a91e36b5..2a1344489 100644 --- a/crates/paimon/src/table/commit_message.rs +++ b/crates/paimon/src/table/commit_message.rs @@ -27,9 +27,11 @@ pub struct CommitMessage { pub partition: Vec, /// Bucket id. pub bucket: i32, + /// Per-partition bucket count for fixed-bucket postpone writes. + 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. + /// Snapshot id from which state-dependent write conflicts should be checked. pub check_from_snapshot: Option, /// New changelog files to be added. pub new_changelog_files: Vec, @@ -46,6 +48,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..8243fde4f 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -72,7 +72,11 @@ mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; mod pk_vector_position_read; mod pk_vector_scan; +mod postpone_bucket_plan; mod postpone_file_writer; +mod postpone_fixed_bucket_router; +mod postpone_fixed_bucket_write; +mod postpone_fixed_bucket_write_builder; mod prepared_files; mod read_builder; pub mod referenced_files; @@ -120,6 +124,11 @@ pub use incremental_scan::{ }; pub use lumina_index_build_builder::LuminaIndexBuildBuilder; pub use partition_stat::PartitionStat; +pub use postpone_bucket_plan::{PostponeBucketPlan, POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}; +pub use postpone_fixed_bucket_write::{ + PostponeFixedBucketTableCommit, PostponeFixedBucketTableWrite, +}; +pub use postpone_fixed_bucket_write_builder::PostponeFixedBucketWriteBuilder; pub use read_builder::ReadBuilder; pub use rest_env::RESTEnv; pub use scan_trace::ScanTrace; @@ -359,6 +368,13 @@ impl Table { WriteBuilder::new(self) } + /// Create a one-shot fixed-bucket builder for a postpone table. + pub fn new_postpone_fixed_bucket_write_builder( + &self, + ) -> Result> { + PostponeFixedBucketWriteBuilder::new(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_bucket_plan.rs b/crates/paimon/src/table/postpone_bucket_plan.rs new file mode 100644 index 000000000..6b6d867e0 --- /dev/null +++ b/crates/paimon/src/table/postpone_bucket_plan.rs @@ -0,0 +1,122 @@ +// 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. + +use crate::spec::batch_to_serialized_bytes; +use crate::table::Table; +use crate::Result; +use arrow_array::{Array, Int32Array, RecordBatch}; +use std::collections::HashMap; +use std::sync::Arc; + +pub const POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD: &str = "total_buckets"; + +pub(crate) fn data_invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +#[derive(Debug, Clone, Default)] +pub struct PostponeBucketPlan { + bucket_counts: Arc, i32>>, +} + +impl PostponeBucketPlan { + 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(data_invalid(format!( + "Postpone bucket plan expected {} partition column(s) plus '{POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}', got {} columns", + partition_count, + batch.num_columns() + ))); + } + + let expected = crate::arrow::build_target_arrow_schema(&partition_fields)?; + let actual = batch.schema(); + if !expected + .fields() + .iter() + .zip(actual.fields()) + .all(|(expected, actual)| { + actual.name() == expected.name() && actual.data_type() == expected.data_type() + }) + { + return Err(data_invalid( + "Postpone bucket plan partition fields do not match the table schema", + )); + } + for (index, field) in expected.fields().iter().enumerate() { + if !field.is_nullable() && batch.column(index).null_count() != 0 { + return Err(data_invalid(format!( + "Postpone bucket plan partition column '{}' is NOT NULL but contains null values", + field.name() + ))); + } + } + + let count_field = actual.field(partition_count); + if count_field.name() != POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD + || count_field.data_type() != &arrow_schema::DataType::Int32 + { + return Err(data_invalid(format!( + "Postpone bucket plan final field must be '{POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}': Int32, got '{}': {:?}", + count_field.name(), + count_field.data_type() + ))); + } + let counts = batch + .column(partition_count) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + data_invalid("Postpone bucket plan total_buckets column is not Int32") + })?; + 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(data_invalid(format!( + "Postpone bucket plan total_buckets is null at row {row}" + ))); + } + let total_buckets = counts.value(row); + if total_buckets <= 0 { + return Err(data_invalid(format!( + "Postpone bucket plan total_buckets must be positive at row {row}, got {total_buckets}" + ))); + } + if let Some(previous) = bucket_counts.insert(partition, total_buckets) { + if previous != total_buckets { + return Err(data_invalid(format!( + "Postpone bucket plan contains conflicting total bucket counts {previous} and {total_buckets} for one partition" + ))); + } + } + } + Ok(Self { + bucket_counts: Arc::new(bucket_counts), + }) + } + + pub(crate) fn total_buckets(&self, partition: &[u8]) -> Option { + self.bucket_counts.get(partition).copied() + } +} diff --git a/crates/paimon/src/table/postpone_fixed_bucket_router.rs b/crates/paimon/src/table/postpone_fixed_bucket_router.rs new file mode 100644 index 000000000..27b9cb88c --- /dev/null +++ b/crates/paimon/src/table/postpone_fixed_bucket_router.rs @@ -0,0 +1,156 @@ +// 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. + +use super::postpone_bucket_plan::data_invalid; +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::table_write::take_rows; +use crate::table::{PostponeBucketPlan, Table}; +use crate::Result; +use arrow_array::RecordBatch; +use std::collections::HashMap; + +#[derive(Debug)] +pub(crate) struct PostponeBucketBatch { + pub(crate) partition: Vec, + pub(crate) bucket: i32, + pub(crate) batch: RecordBatch, +} + +pub(crate) struct PostponeFixedBucketRouter { + fields: Vec, + partition_field_indices: Vec, + bucket_key_indices: Vec, + bucket_function_type: BucketFunctionType, + plan: PostponeBucketPlan, +} + +impl PostponeFixedBucketRouter { + pub(crate) fn new(table: &Table, plan: PostponeBucketPlan) -> Result { + validate_postpone_fixed_bucket_table(table)?; + let schema = table.schema(); + let options = CoreOptions::new(schema.options()); + let bucket_key_indices = field_indices(schema.fields(), &schema.bucket_keys()); + let bucket_key_fields = bucket_key_indices + .iter() + .map(|&index| schema.fields()[index].clone()) + .collect::>(); + let bucket_function_type = options.bucket_function_type()?; + if !bucket_key_fields.is_empty() { + validate_bucket_function(bucket_function_type, &bucket_key_fields)?; + } + Ok(Self { + fields: schema.fields().to_vec(), + partition_field_indices: field_indices(schema.fields(), schema.partition_keys()), + bucket_key_indices, + bucket_function_type, + plan, + }) + } + + pub(crate) fn route(&self, batch: &RecordBatch) -> Result> { + let mut output = Vec::new(); + for (partition, batch) in + partition_batches(batch, &self.partition_field_indices, &self.fields)? + { + let total_buckets = self.plan.total_buckets(&partition).ok_or_else(|| { + data_invalid("Postpone bucket plan does not contain an input partition") + })?; + 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, + &self.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); + } + for (bucket, rows) in groups { + output.push(PostponeBucketBatch { + partition: partition.clone(), + bucket, + batch: take_rows(&batch, &rows)?, + }); + } + } + Ok(output) + } + + pub(crate) fn total_buckets(&self, partition: &[u8]) -> Option { + self.plan.total_buckets(partition) + } +} + +pub(crate) fn validate_postpone_fixed_bucket_table(table: &Table) -> Result<()> { + let schema = table.schema(); + let bucket = CoreOptions::new(schema.options()).bucket(); + if table.is_format_table() || bucket != POSTPONE_BUCKET || schema.primary_keys().is_empty() { + return Err(crate::Error::Unsupported { + message: format!( + "Postpone fixed-bucket writes require a Paimon primary-key table with bucket=-2, but table '{}' has bucket={bucket}", + table.identifier().full_name() + ), + }); + } + if schema + .partition_keys() + .iter() + .any(|key| !schema.primary_keys().contains(key)) + { + return Err(crate::Error::Unsupported { + message: "Postpone fixed-bucket writes do not support cross-partition updates" + .to_string(), + }); + } + Ok(()) +} + +fn field_indices(fields: &[DataField], names: &[String]) -> Vec { + names + .iter() + .filter_map(|name| fields.iter().position(|field| field.name() == name)) + .collect() +} + +fn partition_batches( + batch: &RecordBatch, + partition_field_indices: &[usize], + fields: &[DataField], +) -> Result, RecordBatch)>> { + let partitions = if partition_field_indices.is_empty() { + vec![EMPTY_SERIALIZED_ROW.clone(); batch.num_rows()] + } else { + batch_to_serialized_bytes(batch, partition_field_indices, fields)? + }; + let mut groups: HashMap, Vec> = HashMap::new(); + for (row, partition) in partitions.into_iter().enumerate() { + groups.entry(partition).or_default().push(row); + } + groups + .into_iter() + .map(|(partition, rows)| Ok((partition, take_rows(batch, &rows)?))) + .collect() +} diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write.rs b/crates/paimon/src/table/postpone_fixed_bucket_write.rs new file mode 100644 index 000000000..939bdd862 --- /dev/null +++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs @@ -0,0 +1,814 @@ +// 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. + +use super::postpone_bucket_plan::data_invalid; +use super::postpone_fixed_bucket_router::{ + validate_postpone_fixed_bucket_table, PostponeFixedBucketRouter, +}; +use crate::spec::CoreOptions; +use crate::table::{CommitMessage, PostponeBucketPlan, Table, TableCommit, TableWrite}; +use crate::Result; +use arrow_array::RecordBatch; + +pub struct PostponeFixedBucketTableWrite { + inner: TableWrite, + router: PostponeFixedBucketRouter, + check_from_snapshot: Option, + prepare_started: bool, +} + +impl PostponeFixedBucketTableWrite { + pub(crate) fn new( + table: &Table, + commit_user: String, + plan: PostponeBucketPlan, + overwrite: bool, + ) -> Result { + validate_postpone_fixed_bucket_write(table)?; + let inner = TableWrite::new(table, commit_user)?; + Ok(Self { + inner: if overwrite { + inner.with_overwrite() + } else { + inner + }, + router: PostponeFixedBucketRouter::new(table, plan)?, + check_from_snapshot: None, + prepare_started: false, + }) + } + + pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) -> Result<()> { + self.ensure_writable()?; + let Some(batch) = self.inner.normalize_write_batch(batch)? else { + return Ok(()); + }; + if self.check_from_snapshot.is_none() { + self.check_from_snapshot = Some(self.inner.pin_sequence_snapshot().await?); + } + for routed in self.router.route(&batch)? { + self.inner + .write_partition_bucket_batch(routed.partition, routed.bucket, routed.batch) + .await?; + } + Ok(()) + } + + pub async fn write_arrow(&mut self, batches: &[RecordBatch]) -> Result<()> { + for batch in batches { + self.write_arrow_batch(batch).await?; + } + Ok(()) + } + + pub async fn prepare_commit(&mut self) -> Result> { + self.ensure_writable()?; + self.prepare_started = true; + let mut messages = self.inner.prepare_commit().await?; + for message in &mut messages { + message.total_buckets = self.router.total_buckets(&message.partition); + message.check_from_snapshot = self.check_from_snapshot; + } + Ok(messages) + } + + fn ensure_writable(&self) -> Result<()> { + if self.prepare_started { + return Err(data_invalid("Fixed-bucket postpone TableWrite only supports one prepare_commit call; create a new writer for the next batch")); + } + Ok(()) + } +} + +pub struct PostponeFixedBucketTableCommit { + inner: TableCommit, + overwrite: bool, +} + +impl PostponeFixedBucketTableCommit { + pub(crate) fn new(table: &Table, commit_user: String, overwrite: bool) -> Self { + Self { + inner: TableCommit::new(table.clone(), commit_user), + overwrite, + } + } + + pub async fn commit(&self, messages: Vec) -> Result<()> { + if self.overwrite { + self.inner.overwrite(messages, None).await + } else { + self.inner.commit(messages).await + } + } + + pub async fn commit_with_identifier( + &self, + messages: Vec, + commit_identifier: i64, + ) -> Result<()> { + if self.overwrite { + self.inner + .overwrite_with_identifier(messages, None, commit_identifier) + .await + } else { + self.inner + .commit_with_identifier(messages, commit_identifier) + .await + } + } + + pub async fn filter_and_commit_with_identifier( + &self, + messages: Vec, + commit_identifier: i64, + ) -> Result<()> { + if self.overwrite { + self.inner + .overwrite_with_identifier(messages, None, commit_identifier) + .await + } else { + self.inner + .filter_and_commit_with_identifier(messages, commit_identifier) + .await + } + } + + pub async fn overwrite(&self, messages: Vec) -> Result<()> { + self.inner.overwrite(messages, None).await + } + + pub async fn overwrite_with_identifier( + &self, + messages: Vec, + commit_identifier: i64, + ) -> Result<()> { + self.inner + .overwrite_with_identifier(messages, None, commit_identifier) + .await + } + + pub async fn truncate_table(&self) -> Result<()> { + self.inner.truncate_table().await + } + + pub async fn truncate_table_with_identifier(&self, commit_identifier: i64) -> Result<()> { + self.inner + .truncate_table_with_identifier(commit_identifier) + .await + } + + pub async fn abort(&self, messages: &[CommitMessage]) -> Result<()> { + self.inner.abort(messages).await + } +} + +fn validate_postpone_fixed_bucket_write(table: &Table) -> Result<()> { + validate_postpone_fixed_bucket_table(table)?; + if CoreOptions::new(table.schema().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() + ), + }); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::catalog::Identifier; + use crate::io::FileIO; + use crate::spec::{DataType, IntType, Schema, TableSchema, VarCharType}; + use crate::table::table_write::tests::{ + make_batch, make_partitioned_batch_3col, read_id_value_rows, setup_dirs, test_file_io, + test_postpone_partitioned_table, test_postpone_pk_table, + }; + use crate::table::{ + CommitMessage, PostponeBucketPlan, PostponeFixedBucketWriteBuilder, SnapshotManager, Table, + TableCommit, TableScan, + }; + use arrow_array::{Int32Array, RecordBatch, StringArray}; + use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; + use std::sync::Arc; + + fn make_partition_bucket_plan_batch( + partitions: Vec<&str>, + total_buckets: Vec, + ) -> RecordBatch { + 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() + } + + fn make_partition_bucket_plan( + table: &Table, + partitions: Vec<&str>, + total_buckets: Vec, + ) -> PostponeBucketPlan { + PostponeBucketPlan::from_arrow( + table, + &make_partition_bucket_plan_batch(partitions, total_buckets), + ) + .unwrap() + } + + fn make_unpartitioned_bucket_plan(table: &Table, total_buckets: i32) -> PostponeBucketPlan { + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ArrowField::new( + "total_buckets", + ArrowDataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(vec![total_buckets]))], + ) + .unwrap(); + PostponeBucketPlan::from_arrow(table, &batch).unwrap() + } + + fn cross_partition_postpone_table(file_io: &FileIO, table_path: &str) -> Table { + let schema = Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .partition_keys(["pt"]) + .option("bucket", "-2") + .build() + .unwrap(); + Table::new( + file_io.clone(), + Identifier::new("default", "cross_partition_postpone"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) + } + + async fn write_fixed_batch( + table: &Table, + commit_user: &str, + total_buckets: i32, + batch: &RecordBatch, + ) -> Vec { + let mut write = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user(commit_user) + .unwrap() + .with_bucket_plan(make_unpartitioned_bucket_plan(table, total_buckets)) + .new_write() + .unwrap(); + write.write_arrow_batch(batch).await.unwrap(); + write.prepare_commit().await.unwrap() + } + + async fn prepare_partitioned_fixed_batch<'a>( + table: &'a Table, + commit_user: &str, + plan: PostponeBucketPlan, + partition: &str, + id: i32, + value: i32, + ) -> (PostponeFixedBucketWriteBuilder<'a>, Vec) { + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user(commit_user) + .unwrap() + .with_bucket_plan(plan); + let mut write = builder.new_write().unwrap(); + write + .write_arrow_batch(&make_partitioned_batch_3col( + vec![partition], + vec![id], + vec![value], + )) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + (builder, messages) + } + + #[test] + fn test_postpone_fixed_bucket_rejects_cross_partition_update() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_cross_partition"; + let table = cross_partition_postpone_table(&file_io, table_path); + let error = match table.new_postpone_fixed_bucket_write_builder() { + Ok(_) => panic!("cross-partition postpone fixed-bucket writes should be rejected"), + Err(error) => error, + }; + assert!(error + .to_string() + .contains("do not support cross-partition updates")); + } + + fn assert_total_buckets(messages: &[CommitMessage], total_buckets: i32) { + assert!(!messages.is_empty()); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(total_buckets))); + } + + fn assert_one_shot_error(error: crate::Error) { + assert!( + matches!(error, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supports one prepare_commit call") + && message.contains("create a new writer")) + ); + } + + #[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_postpone_pk_table(&file_io, table_path); + + let first_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("fixed-user-1") + .unwrap() + .with_bucket_plan(make_unpartitioned_bucket_plan(&table, 2)); + let mut write = first_builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch(vec![1, 2, 3, 4], vec![10, 20, 30, 40])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages.iter().all(|message| message.bucket >= 0)); + assert_total_buckets(&messages, 2); + + let stale_messages = + write_fixed_batch(&table, "stale-user", 1, &make_batch(vec![9], vec![90])).await; + assert_total_buckets(&stale_messages, 1); + TableCommit::new(table.clone(), "fixed-user-1".to_string()) + .commit(messages) + .await + .unwrap(); + let error = TableCommit::new(table.clone(), "stale-user".to_string()) + .commit(stale_messages) + .await + .unwrap_err(); + assert!(error.to_string().contains("Fixed-bucket conflict")); + assert_eq!( + read_id_value_rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40)] + ); + + let second_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("fixed-user-2") + .unwrap() + .with_bucket_plan(make_unpartitioned_bucket_plan(&table, 2)); + let mut write = second_builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch(vec![5], vec![50])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert_total_buckets(&messages, 2); + assert_one_shot_error( + write + .write_arrow_batch(&make_batch(vec![6], vec![60])) + .await + .unwrap_err(), + ); + assert_one_shot_error(write.prepare_commit().await.unwrap_err()); + 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!["p1", "p2"], vec![3, 3]); + + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("shared-plan-user") + .unwrap() + .with_bucket_plan(plan); + let mut first = builder.new_write().unwrap(); + let mut second = builder.new_write().unwrap(); + first + .write_arrow_batch(&make_partitioned_batch_3col(vec!["p1"], vec![1], vec![10])) + .await + .unwrap(); + second + .write_arrow_batch(&make_partitioned_batch_3col( + vec!["p2", "p2", "p2", "p2"], + 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_total_buckets(&messages, 3); + builder.new_commit().commit(messages).await.unwrap(); + } + + #[tokio::test] + async fn test_postpone_distributed_writers_reject_overlapping_bucket_ownership() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_overlapping_writer_ownership"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_partitioned_table(&file_io, table_path); + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("overlapping-writers") + .unwrap() + .with_bucket_plan(make_partition_bucket_plan(&table, vec!["p"], vec![1])); + + 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"], vec![1], vec![20])) + .await + .unwrap(); + + let mut messages = first.prepare_commit().await.unwrap(); + messages.extend(second.prepare_commit().await.unwrap()); + let error = TableCommit::new(table.clone(), "overlapping-writers".to_string()) + .commit(messages.clone()) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("writer ownership conflict for bucket 0")); + + let error = builder.new_commit().commit(messages).await.unwrap_err(); + assert!(error + .to_string() + .contains("writer ownership conflict for bucket 0")); + } + + #[tokio::test] + async fn test_postpone_concurrent_commits_reject_overlapping_bucket_ownership() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_concurrent_writer_ownership"; + 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![1]); + let (first_builder, first_messages) = prepare_partitioned_fixed_batch( + &table, + "concurrent-writer-1", + plan.clone(), + "p", + 1, + 10, + ) + .await; + let (second_builder, second_messages) = + prepare_partitioned_fixed_batch(&table, "concurrent-writer-2", plan, "p", 1, 20).await; + assert_eq!(first_messages[0].new_files[0].min_sequence_number, 0); + assert_eq!(second_messages[0].new_files[0].min_sequence_number, 0); + first_builder + .new_commit() + .commit(first_messages) + .await + .unwrap(); + let error = second_builder + .new_commit() + .commit(second_messages) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("writer ownership conflict for bucket 0")); + } + + #[tokio::test] + async fn test_postpone_concurrent_overwrites_reject_overlapping_bucket_ownership() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_concurrent_overwrite_ownership"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + let plan = make_unpartitioned_bucket_plan(&table, 1); + + let first_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("overwrite-writer-1") + .unwrap() + .with_bucket_plan(plan.clone()) + .with_overwrite(); + let second_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("overwrite-writer-2") + .unwrap() + .with_bucket_plan(plan) + .with_overwrite(); + let mut first_write = first_builder.new_write().unwrap(); + let mut second_write = second_builder.new_write().unwrap(); + first_write + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + second_write + .write_arrow_batch(&make_batch(vec![1], vec![20])) + .await + .unwrap(); + let first_messages = first_write.prepare_commit().await.unwrap(); + let second_messages = second_write.prepare_commit().await.unwrap(); + + first_builder + .new_commit() + .commit(first_messages) + .await + .unwrap(); + let error = second_builder + .new_commit() + .commit(second_messages) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("writer ownership conflict for bucket 0")); + } + + #[tokio::test] + async fn test_postpone_concurrent_commits_allow_disjoint_ownership() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_disjoint_writer_ownership"; + 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!["p1", "p2"], vec![1, 1]); + let (first_builder, first_messages) = + prepare_partitioned_fixed_batch(&table, "disjoint-writer-1", plan.clone(), "p1", 1, 10) + .await; + let (second_builder, second_messages) = + prepare_partitioned_fixed_batch(&table, "disjoint-writer-2", plan, "p2", 2, 20).await; + + first_builder + .new_commit() + .commit(first_messages) + .await + .unwrap(); + second_builder + .new_commit() + .commit(second_messages) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_postpone_rejects_negative_bucket_after_fixed_layout() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_rejects_negative_bucket_after_fixed"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + let fixed_messages = + write_fixed_batch(&table, "fixed-writer", 1, &make_batch(vec![1], vec![10])).await; + TableCommit::new(table.clone(), "fixed-writer".to_string()) + .commit(fixed_messages) + .await + .unwrap(); + + let normal_builder = table + .new_write_builder() + .with_commit_user("normal-writer") + .unwrap(); + let mut normal_write = normal_builder.new_write().unwrap(); + normal_write + .write_arrow_batch(&make_batch(vec![1], vec![20])) + .await + .unwrap(); + let messages = normal_write.prepare_commit().await.unwrap(); + assert!(messages.iter().all(|message| message.bucket == -2)); + let error = normal_builder + .new_commit() + .commit(messages) + .await + .unwrap_err(); + assert!(error.to_string().contains("already uses fixed buckets")); + assert_eq!(read_id_value_rows(&table).await, vec![(1, 10)]); + } + + #[tokio::test] + async fn test_postpone_fixed_commit_rejects_concurrent_negative_bucket() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_rejects_concurrent_negative_bucket"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + let fixed_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("fixed-writer") + .unwrap() + .with_bucket_plan(make_unpartitioned_bucket_plan(&table, 1)); + let mut fixed_write = fixed_builder.new_write().unwrap(); + fixed_write + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let fixed_messages = fixed_write.prepare_commit().await.unwrap(); + + let normal_builder = table + .new_write_builder() + .with_commit_user("normal-writer") + .unwrap(); + let mut normal_write = normal_builder.new_write().unwrap(); + normal_write + .write_arrow_batch(&make_batch(vec![1], vec![20])) + .await + .unwrap(); + normal_builder + .new_commit() + .commit(normal_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let error = fixed_builder + .new_commit() + .commit(fixed_messages) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("another commit wrote bucket=-2 files")); + } + + #[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 error = PostponeBucketPlan::from_arrow( + &table, + &make_partition_bucket_plan_batch(vec!["p"], vec![0]), + ) + .unwrap_err(); + assert!(error.to_string().contains("must be positive")); + + let plan = make_partition_bucket_plan(&table, vec!["p"], vec![2]); + let mut write = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_bucket_plan(plan) + .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_bucket_plan(make_partition_bucket_plan(&table, vec!["p"], vec![1])); + 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_bucket_plan(make_partition_bucket_plan(&table, vec!["p"], vec![3])) + .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_total_buckets(&messages, 3); + overwrite_builder + .new_commit() + .commit(messages) + .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_fixed_bucket_delete_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!["-D"])), + ], + ) + .unwrap(); + + let messages = write_fixed_batch(&table, "fixed-rowkind", 1, &batch).await; + + 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.len(), 1); + assert_eq!(messages[0].new_files[0].row_count, 1); + assert_eq!(messages[0].new_files[0].delete_row_count, Some(1)); + } +} diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs new file mode 100644 index 000000000..44905876d --- /dev/null +++ b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs @@ -0,0 +1,159 @@ +// 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. + +use super::postpone_bucket_plan::data_invalid; +use super::postpone_fixed_bucket_router::validate_postpone_fixed_bucket_table; +use crate::table::write_builder::{ensure_table_write_allowed, validate_commit_user}; +use crate::table::{ + PostponeBucketPlan, PostponeFixedBucketTableCommit, PostponeFixedBucketTableWrite, Table, +}; +use crate::Result; +use uuid::Uuid; + +pub struct PostponeFixedBucketWriteBuilder<'a> { + table: &'a Table, + commit_user: String, + overwrite: bool, + bucket_plan: Option, +} + +impl<'a> PostponeFixedBucketWriteBuilder<'a> { + pub(crate) fn new(table: &'a Table) -> Result { + validate_postpone_fixed_bucket_table(table)?; + Ok(Self { + table, + commit_user: Uuid::new_v4().to_string(), + overwrite: false, + bucket_plan: None, + }) + } + + pub fn commit_user(&self) -> &str { + &self.commit_user + } + + pub fn with_commit_user(mut self, commit_user: impl Into) -> Result { + let commit_user = commit_user.into(); + validate_commit_user(&commit_user)?; + self.commit_user = commit_user; + Ok(self) + } + + pub fn with_overwrite(mut self) -> Self { + self.overwrite = true; + self + } + + pub fn with_bucket_plan(mut self, bucket_plan: PostponeBucketPlan) -> Self { + self.bucket_plan = Some(bucket_plan); + self + } + + pub fn new_commit(&self) -> PostponeFixedBucketTableCommit { + PostponeFixedBucketTableCommit::new(self.table, self.commit_user.clone(), self.overwrite) + } + + pub fn try_new_commit(&self) -> Result { + self.table.ensure_not_branch_reference_for_write()?; + Ok(self.new_commit()) + } + + pub fn new_write(&self) -> Result { + ensure_table_write_allowed(self.table)?; + let plan = self + .bucket_plan + .clone() + .ok_or_else(|| data_invalid("A resolved postpone bucket plan is required"))?; + PostponeFixedBucketTableWrite::new( + self.table, + self.commit_user.clone(), + plan, + self.overwrite, + ) + } +} + +#[cfg(test)] +mod tests { + use crate::table::table_write::tests::{ + make_batch, setup_dirs, test_file_io, test_postpone_pk_table, + }; + use crate::table::{PostponeBucketPlan, Table}; + use arrow_array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + fn bucket_plan(table: &Table, total_buckets: i32) -> PostponeBucketPlan { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "total_buckets", + DataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(vec![total_buckets]))], + ) + .unwrap(); + PostponeBucketPlan::from_arrow(table, &batch).unwrap() + } + + #[tokio::test] + async fn test_postpone_fixed_bucket_builder_modes() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_builder"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + let error = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .new_write() + .err() + .unwrap(); + assert!(error.to_string().contains("bucket plan is required")); + + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("explicit-fixed-user") + .unwrap() + .with_bucket_plan(bucket_plan(&table, 1)); + let mut write = builder.new_write().unwrap(); + write + .write_arrow_batch(&make_batch(vec![4], vec![40])) + .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)); + + let dv_table = table.copy_with_options(std::collections::HashMap::from([( + "deletion-vectors.enabled".to_string(), + "true".to_string(), + )])); + let error = dv_table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_bucket_plan(bucket_plan(&dv_table, 1)) + .new_write() + .err() + .unwrap(); + assert!(matches!(error, crate::Error::Unsupported { ref message } + if message.contains("postpone fixed-bucket writes") + && message.contains("deletion-vectors.enabled=true"))); + } +} diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 809e45d9f..a317536c9 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -28,6 +28,7 @@ use crate::spec::{ CoreOptions, DataFileMeta, DataType, Datum, GlobalIndexColumnUpdateAction, IndexManifest, IndexManifestEntry, Manifest, ManifestEntry, ManifestFileMeta, ManifestList, PartitionComputer, PartitionStatistics, Predicate, Snapshot, EMPTY_SERIALIZED_ROW, MANIFEST_ENTRY_SCHEMA, + POSTPONE_BUCKET, }; use crate::table::commit_message::CommitMessage; use crate::table::global_index_build_common::same_extra_field_ids; @@ -50,6 +51,25 @@ type PartitionBucketKey = (Vec, i32); type RowIdRange = (i64, i64); type ExistingRowIdRanges = HashMap>; +fn validate_bucket_ownership(messages: &[CommitMessage]) -> Result<()> { + let mut owners = HashSet::new(); + for message in messages { + if message.total_buckets.is_none() || message.new_files.is_empty() { + continue; + } + if !owners.insert((message.partition.as_slice(), message.bucket)) { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket writer ownership conflict for bucket {}: route all rows for one partition and bucket to a single writer", + message.bucket + ), + source: None, + }); + } + } + Ok(()) +} + /// Table commit logic for Paimon write operations. /// /// Provides atomic commit functionality including append, overwrite and truncate @@ -156,6 +176,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() { return Ok(()); @@ -199,6 +220,7 @@ impl TableCommit { commit_identifier: i64, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() { return Ok(()); @@ -270,6 +292,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() && static_partitions.is_none() { return Ok(()); @@ -298,11 +321,14 @@ impl TableCommit { } } + let check_from_snapshot = Self::min_check_from_snapshot(&commit_messages); + self.try_commit( CommitEntriesPlan::Overwrite { partition_filter, new_entries, new_index_entries, + check_from_snapshot, cached_snapshot: None, cached_entries: Vec::new(), full_scan_count: 0, @@ -531,6 +557,7 @@ impl TableCommit { partition_filter: Some(partition_filter), new_entries: vec![], new_index_entries: vec![], + check_from_snapshot: None, cached_snapshot: None, cached_entries: Vec::new(), full_scan_count: 0, @@ -603,6 +630,7 @@ impl TableCommit { partition_filter: None, new_entries: vec![], new_index_entries: vec![], + check_from_snapshot: None, cached_snapshot: None, cached_entries: Vec::new(), full_scan_count: 0, @@ -1240,7 +1268,17 @@ 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 has_postpone_entries = self.total_buckets == POSTPONE_BUCKET + && entries.iter().any(|entry| { + *entry.kind() == FileKind::Add && entry.bucket() == POSTPONE_BUCKET + }); + let detect_conflicts = has_delete + || check_from_snapshot.is_some() + || has_partition_bucket_counts + || has_postpone_entries; let base_data_files = if detect_conflicts { self.check_deletion_vector_index_only_conflict( latest_snapshot.as_ref(), @@ -1298,12 +1336,17 @@ impl TableCommit { let entries = self .provide_overwrite_entries(plan, latest_snapshot) .await?; - let (partition_filter, new_index_entries) = match plan { + let (partition_filter, new_index_entries, check_from_snapshot) = match plan { CommitEntriesPlan::Overwrite { partition_filter, new_index_entries, + check_from_snapshot, .. - } => (partition_filter.clone(), new_index_entries.clone()), + } => ( + partition_filter.clone(), + new_index_entries.clone(), + *check_from_snapshot, + ), CommitEntriesPlan::Direct { .. } => unreachable!(), }; let base_data_files = self @@ -1312,7 +1355,7 @@ impl TableCommit { retry_state, &entries, &CommitKind::OVERWRITE, - None, + check_from_snapshot, ) .await?; @@ -1854,6 +1897,22 @@ impl TableCommit { ) -> Result<()> { self.check_delete_entries_against_base(base_entries, delta_entries)?; + // Validate delta entries before duplicate files are merged. + self.check_total_bucket_conflicts(delta_entries)?; + self.check_postpone_bucket_mixing(base_entries, delta_entries)?; + + // Check the final layout so overwrite rescaling remains valid. + 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)?; + self.check_fixed_bucket_ownership_conflicts( + latest_snapshot, + delta_entries, + check_from_snapshot, + ) + .await?; + if !self.data_evolution_enabled { return Ok(()); } @@ -1861,14 +1920,133 @@ 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 } + async fn check_fixed_bucket_ownership_conflicts( + &self, + latest_snapshot: Option<&Snapshot>, + entries: &[ManifestEntry], + check_from_snapshot: Option, + ) -> Result<()> { + let Some(check_from_snapshot) = check_from_snapshot else { + return Ok(()); + }; + let Some(latest_snapshot) = latest_snapshot else { + return Ok(()); + }; + if latest_snapshot.id() <= check_from_snapshot { + return Ok(()); + } + + let fixed_entries = entries + .iter() + .filter(|entry| { + *entry.kind() == FileKind::Add + && entry.total_buckets() != self.total_buckets + && entry.total_buckets() > 0 + }) + .collect::>(); + if fixed_entries.is_empty() { + return Ok(()); + } + let owned_buckets = fixed_entries + .iter() + .map(|entry| (entry.partition(), entry.bucket())) + .collect::>(); + let owned_partitions = fixed_entries + .iter() + .map(|entry| entry.partition()) + .collect::>(); + let partition_filter = self.build_entries_partition_filter(&fixed_entries)?; + + for snapshot_id in check_from_snapshot.max(0) + 1..=latest_snapshot.id() { + let snapshot = self.snapshot_manager.get_snapshot(snapshot_id).await?; + let concurrent_entries = self + .read_delta_entries(partition_filter.as_ref(), &snapshot) + .await?; + for entry in concurrent_entries { + if *entry.kind() != FileKind::Add { + continue; + } + if entry.bucket() == POSTPONE_BUCKET && owned_partitions.contains(entry.partition()) + { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket writer conflict: another commit wrote bucket=-2 files for the same partition after snapshot {check_from_snapshot}" + ), + source: None, + }); + } + if owned_buckets.contains(&(entry.partition(), entry.bucket())) { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket writer ownership conflict for bucket {}: another commit wrote the same partition and bucket after snapshot {check_from_snapshot}", + entry.bucket() + ), + source: None, + }); + } + } + } + Ok(()) + } + + fn check_postpone_bucket_mixing( + &self, + base_entries: &[ManifestEntry], + delta_entries: &[ManifestEntry], + ) -> Result<()> { + if self.total_buckets != POSTPONE_BUCKET { + return Ok(()); + } + let fixed_partitions = base_entries + .iter() + .chain(delta_entries) + .filter(|entry| { + *entry.kind() == FileKind::Add && entry.bucket() >= 0 && entry.total_buckets() > 0 + }) + .map(|entry| entry.partition()) + .collect::>(); + if delta_entries.iter().any(|entry| { + *entry.kind() == FileKind::Add + && entry.bucket() == POSTPONE_BUCKET + && fixed_partitions.contains(entry.partition()) + }) { + return Err(crate::Error::DataInvalid { + message: + "Cannot commit bucket=-2 files for a partition that already uses fixed buckets" + .to_string(), + source: None, + }); + } + Ok(()) + } + + 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!( + "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 +2735,8 @@ 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 place replacement ADDs last. + stats.total_buckets = entry.total_buckets(); } Ok(stats_map.into_values().collect()) @@ -2602,7 +2782,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 +2792,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 +2812,7 @@ impl TableCommit { FileKind::Add, msg.partition.clone(), msg.bucket, - self.total_buckets, + msg.total_buckets.unwrap_or(self.total_buckets), file.clone(), 0, ) @@ -2701,6 +2881,7 @@ enum CommitEntriesPlan { partition_filter: Option, new_entries: Vec, new_index_entries: Vec, + check_from_snapshot: Option, cached_snapshot: Option>, cached_entries: Vec, full_scan_count: usize, @@ -2903,7 +3084,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}; @@ -3088,6 +3269,7 @@ mod tests { partition_filter, new_entries, new_index_entries: vec![], + check_from_snapshot: None, cached_snapshot: None, cached_entries: Vec::new(), full_scan_count: 0, @@ -4133,6 +4315,36 @@ mod tests { assert_eq!(snapshot.commit_kind(), &CommitKind::OVERWRITE); // 300 - 100 (delete a) + 50 (add a2) = 250 assert_eq!(snapshot.total_record_count(), Some(250)); + + 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] diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 490de6357..9c0c130bc 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -41,7 +41,7 @@ use crate::table::partition_filter::PartitionFilter; use crate::table::postpone_file_writer::{PostponeFileWriter, PostponeWriteConfig}; use crate::table::prepared_files::PreparedFiles; use crate::table::row_kind_generator::RowKindGenerator; -use crate::table::{SnapshotManager, Table, TableScan}; +use crate::table::{Snapshot, SnapshotManager, Table, TableScan}; use crate::Result; use arrow_array::RecordBatch; use std::collections::{HashMap, HashSet}; @@ -55,6 +55,27 @@ enum FileWriter { Postpone(PostponeFileWriter), } +pub(super) fn take_rows(batch: &RecordBatch, row_indices: &[usize]) -> Result { + if row_indices.len() == batch.num_rows() { + return Ok(batch.clone()); + } + let indices = + arrow_array::UInt32Array::from(row_indices.iter().map(|&i| i as u32).collect::>()); + let columns = batch + .columns() + .iter() + .map(|col| arrow_select::take::take(col.as_ref(), &indices, None)) + .collect::, _>>() + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to take rows: {e}"), + source: None, + })?; + RecordBatch::try_new(batch.schema(), columns).map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to create sub-batch: {e}"), + source: None, + }) +} + impl FileWriter { async fn write(&mut self, batch: &RecordBatch) -> Result<()> { match self { @@ -109,6 +130,7 @@ pub struct TableWrite { changelog_file_format: String, changelog_file_compression: String, partition_seq_cache: HashMap, HashMap>, + sequence_snapshot: Option>, commit_user: String, /// Bucket assignment strategy (fixed, dynamic, or cross-partition). bucket_assigner: BucketAssignerEnum, @@ -367,6 +389,7 @@ impl TableWrite { changelog_file_format, changelog_file_compression, partition_seq_cache: HashMap::new(), + sequence_snapshot: None, commit_user, bucket_assigner, is_overwrite, @@ -384,11 +407,17 @@ impl TableWrite { /// bucket → (max_sequence_number + 1) for each bucket in that partition. async fn scan_partition_sequence_numbers( table: &Table, + sequence_snapshot: Option>, partition_bytes: &[u8], ) -> crate::Result> { - let snapshot_manager = - SnapshotManager::new(table.file_io().clone(), table.location().to_string()); - let latest_snapshot = snapshot_manager.get_latest_snapshot().await?; + let latest_snapshot = match sequence_snapshot { + Some(snapshot) => snapshot, + None => { + let snapshot_manager = + SnapshotManager::new(table.file_io().clone(), table.location().to_string()); + snapshot_manager.get_latest_snapshot().await? + } + }; let mut bucket_seq: HashMap = HashMap::new(); if let Some(snapshot) = latest_snapshot { let partition_filter = Self::build_partition_filter(table, partition_bytes)?; @@ -407,6 +436,17 @@ impl TableWrite { Ok(bucket_seq) } + pub(super) async fn pin_sequence_snapshot(&mut self) -> Result { + let snapshot_manager = SnapshotManager::new( + self.table.file_io().clone(), + self.table.location().to_string(), + ); + let snapshot = snapshot_manager.get_latest_snapshot().await?; + let snapshot_id = snapshot.as_ref().map_or(0, Snapshot::id); + self.sequence_snapshot = Some(snapshot); + Ok(snapshot_id) + } + /// Build a partition filter from serialized partition bytes. /// /// Uses `PartitionSet` for O(1) byte-level matching when partition fields exist. @@ -437,16 +477,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<()> { - self.validate_write_batch_schema(batch)?; - - if batch.num_rows() == 0 { - return Ok(()); - } - - let batch = self.enrich_rowkind_batch(batch)?; - if batch.num_rows() == 0 { + let Some(batch) = self.normalize_write_batch(batch)? else { return Ok(()); - } + }; let grouped = self.divide_by_partition_bucket(&batch).await?; for ((partition_bytes, bucket), sub_batch) in grouped { @@ -456,6 +489,24 @@ impl TableWrite { Ok(()) } + pub(super) fn normalize_write_batch(&self, batch: &RecordBatch) -> Result> { + self.validate_write_batch_schema(batch)?; + if batch.num_rows() == 0 { + return Ok(None); + } + let batch = self.enrich_rowkind_batch(batch)?; + Ok((batch.num_rows() != 0).then_some(batch)) + } + + pub(super) async fn write_partition_bucket_batch( + &mut self, + partition: Vec, + bucket: i32, + batch: RecordBatch, + ) -> Result<()> { + self.write_bucket(partition, bucket, batch).await + } + fn validate_write_batch_schema(&self, batch: &RecordBatch) -> Result<()> { let expected_schema = &self.write_schema; let actual_schema = batch.schema(); @@ -606,7 +657,7 @@ impl TableWrite { || matches!(self.bucket_assigner, BucketAssignerEnum::CrossPartition(_)) || !output.deletes.is_empty(); for (key, row_indices) in groups { - let sub_batch = Self::take_rows(batch, &row_indices)?; + let sub_batch = take_rows(batch, &row_indices)?; let sub_batch = if needs_value_kind && !batch_has_value_kind { Self::add_value_kind_column(&sub_batch, 0)? } else { @@ -624,7 +675,7 @@ impl TableWrite { .push(*row_idx); } for (key, row_indices) in delete_groups { - let sub_batch = Self::take_rows(batch, &row_indices)?; + let sub_batch = take_rows(batch, &row_indices)?; let delete_batch = Self::add_value_kind_column(&sub_batch, 1)?; result.push((key, delete_batch)); } @@ -633,29 +684,6 @@ impl TableWrite { Ok(result) } - /// Extract rows from a batch by indices. - fn take_rows(batch: &RecordBatch, row_indices: &[usize]) -> Result { - if row_indices.len() == batch.num_rows() { - return Ok(batch.clone()); - } - let indices = arrow_array::UInt32Array::from( - row_indices.iter().map(|&i| i as u32).collect::>(), - ); - let columns: Vec> = batch - .columns() - .iter() - .map(|col| arrow_select::take::take(col.as_ref(), &indices, None)) - .collect::, _>>() - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to take rows: {e}"), - source: None, - })?; - RecordBatch::try_new(batch.schema(), columns).map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to create sub-batch: {e}"), - source: None, - }) - } - /// Add a `_VALUE_KIND` column to a batch with the given value for all rows. fn add_value_kind_column(batch: &RecordBatch, value_kind: i8) -> Result { use arrow_array::Int8Array; @@ -711,7 +739,7 @@ impl TableWrite { if keep_rows.is_empty() { return Ok(RecordBatch::new_empty(batch.schema())); } - let filtered = Self::take_rows(batch, &keep_rows)?; + let filtered = take_rows(batch, &keep_rows)?; Self::add_per_row_value_kind_column(&filtered, kinds) } @@ -937,8 +965,11 @@ impl TableWrite { // Lazily scan partition sequence numbers on first writer creation per partition. // Overwrite mode skips this — old data will be replaced, so seq starts at 0. if !self.is_overwrite && !self.partition_seq_cache.contains_key(partition_bytes) { + let table = self.table.clone(); + let sequence_snapshot = self.sequence_snapshot.clone(); let bucket_seq = - Self::scan_partition_sequence_numbers(&self.table, partition_bytes).await?; + Self::scan_partition_sequence_numbers(&table, sequence_snapshot, partition_bytes) + .await?; self.partition_seq_cache .insert(partition_bytes.to_vec(), bucket_seq); } @@ -980,7 +1011,7 @@ impl TableWrite { } #[cfg(test)] -mod tests { +pub(in crate::table) mod tests { use super::*; use crate::arrow::format::create_format_reader; use crate::catalog::Identifier; @@ -1001,7 +1032,7 @@ mod tests { }; use std::sync::Arc; - fn test_file_io() -> FileIO { + pub(in crate::table) fn test_file_io() -> FileIO { FileIOBuilder::new("memory").build().unwrap() } @@ -1069,7 +1100,7 @@ mod tests { TableSchema::new(0, &schema) } - async fn setup_dirs(file_io: &FileIO, table_path: &str) { + pub(in crate::table) async fn setup_dirs(file_io: &FileIO, table_path: &str) { file_io .mkdirs(&format!("{table_path}/snapshot/")) .await @@ -1080,7 +1111,7 @@ mod tests { .unwrap(); } - fn make_batch(ids: Vec, values: Vec) -> RecordBatch { + pub(in crate::table) fn make_batch(ids: Vec, values: Vec) -> RecordBatch { let schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new("id", ArrowDataType::Int32, false), ArrowField::new("value", ArrowDataType::Int32, false), @@ -3659,7 +3690,7 @@ mod tests { TableSchema::new(0, &schema) } - fn test_postpone_pk_table(file_io: &FileIO, table_path: &str) -> Table { + pub(in crate::table) fn test_postpone_pk_table(file_io: &FileIO, table_path: &str) -> Table { Table::new( file_io.clone(), Identifier::new("default", "test_postpone_table"), @@ -3682,7 +3713,10 @@ mod tests { TableSchema::new(0, &schema) } - fn test_postpone_partitioned_table(file_io: &FileIO, table_path: &str) -> Table { + pub(in crate::table) fn test_postpone_partitioned_table( + file_io: &FileIO, + table_path: &str, + ) -> Table { Table::new( file_io.clone(), Identifier::new("default", "test_postpone_table"), @@ -3692,7 +3726,11 @@ mod tests { ) } - fn make_partitioned_batch_3col(pts: Vec<&str>, ids: Vec, values: Vec) -> RecordBatch { + pub(in crate::table) fn make_partitioned_batch_3col( + pts: Vec<&str>, + ids: Vec, + values: Vec, + ) -> RecordBatch { let schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new("pt", ArrowDataType::Utf8, false), ArrowField::new("id", ArrowDataType::Int32, false), @@ -4217,7 +4255,7 @@ mod tests { .unwrap(); } - async fn read_id_value_rows(table: &Table) -> Vec<(i32, i32)> { + pub(in crate::table) async fn read_id_value_rows(table: &Table) -> Vec<(i32, i32)> { let rb = table.new_read_builder(); let plan = rb.new_scan().plan().await.unwrap(); let read = rb.new_read().unwrap(); diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 52ba57754..57747e1ea 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -179,25 +179,7 @@ impl<'a> PaimonWriteBuilder<'a> { /// For primary-key tables, sequence numbers are lazily scanned per partition /// when the first writer for that partition is created. pub fn new_write(&self) -> crate::Result { - self.ensure_main_branch_write()?; - // A table with a time-travel selector reads a pinned snapshot (and may - // carry that snapshot's historical schema), so writing through the - // same copy would be inconsistent with what its reads observe — even - // when the pinned snapshot happens to share the current schema id. - // Java avoids this structurally (write paths use copyWithoutTimeTravel); - // here the same table copy can serve both reads and writes, so reject - // explicitly. Conflicting selectors (`Err`) cannot be valid for writes - // either. Commit-only flows (new_commit) stay untouched. - let selector = - crate::spec::CoreOptions::new(self.table.schema().options()).try_time_travel_selector(); - if !matches!(selector, Ok(None)) { - return Err(crate::Error::Unsupported { - message: - "Cannot write to a table with a time-travel option set \ - (scan.version / scan.timestamp-millis / scan.snapshot-id / scan.tag-name)" - .to_string(), - }); - } + ensure_table_write_allowed(self.table)?; let write = TableWrite::new(self.table, self.commit_user.clone())?; Ok(if self.overwrite { write.with_overwrite() @@ -223,6 +205,21 @@ impl<'a> PaimonWriteBuilder<'a> { } } +pub(super) fn ensure_table_write_allowed(table: &Table) -> crate::Result<()> { + table.ensure_not_branch_reference_for_write()?; + // A time-travel table may carry a historical schema. + let selector = + crate::spec::CoreOptions::new(table.schema().options()).try_time_travel_selector(); + if !matches!(selector, Ok(None)) { + return Err(crate::Error::Unsupported { + message: "Cannot write to a table with a time-travel option set \ + (scan.version / scan.timestamp-millis / scan.snapshot-id / scan.tag-name)" + .to_string(), + }); + } + Ok(()) +} + pub(super) fn validate_commit_user(commit_user: &str) -> crate::Result<()> { let is_invalid = commit_user.is_empty() || commit_user == "." @@ -410,6 +407,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