diff --git a/bindings/c/src/result.rs b/bindings/c/src/result.rs index a833de5e9..94f87c2de 100644 --- a/bindings/c/src/result.rs +++ b/bindings/c/src/result.rs @@ -109,3 +109,27 @@ pub struct paimon_result_prepare_commit { pub messages: *mut paimon_commit_messages, pub error: *mut paimon_error, } + +#[repr(C)] +pub struct paimon_result_postpone_fixed_bucket_write_builder { + pub write_builder: *mut paimon_postpone_fixed_bucket_write_builder, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_postpone_fixed_bucket_table_write { + pub write: *mut paimon_postpone_fixed_bucket_table_write, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_postpone_fixed_bucket_table_commit { + pub commit: *mut paimon_postpone_fixed_bucket_table_commit, + pub error: *mut paimon_error, +} + +#[repr(C)] +pub struct paimon_result_postpone_fixed_bucket_prepare_commit { + pub messages: *mut paimon_postpone_fixed_bucket_commit_messages, + pub error: *mut paimon_error, +} diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index d29c7eb85..e7da43068 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,75 @@ fn test_commit_rejects_messages_from_different_builder_identity() { } } +#[test] +fn test_fixed_commit_rejects_mismatched_overwrite_mode() { + let path = "memory:/test_fixed_commit_overwrite_mode"; + 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("fixed-overwrite-mode-job").unwrap(); + + unsafe { + let append_wb = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ) + .write_builder; + let (array, schema) = + export_batch_to_ffi(make_postpone_bucket_plan_batch(vec!["p"], vec![1])); + assert!(paimon_postpone_fixed_bucket_write_builder_with_bucket_plan( + append_wb, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ) + .is_null()); + let tw = paimon_postpone_fixed_bucket_write_builder_new_write(append_wb).write; + let (array, schema) = export_batch_to_ffi(make_partitioned_write_batch( + vec!["p"], + vec![1], + vec!["append"], + )); + assert!(paimon_postpone_fixed_bucket_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 messages = paimon_postpone_fixed_bucket_table_write_prepare_commit(tw).messages; + + let overwrite_wb = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ) + .write_builder; + assert!(paimon_postpone_fixed_bucket_write_builder_with_overwrite(overwrite_wb).is_null()); + let commit = paimon_postpone_fixed_bucket_write_builder_new_commit(overwrite_wb).commit; + let error = paimon_postpone_fixed_bucket_table_commit_commit(commit, messages); + assert!(!error.is_null()); + assert!(error_message(error).contains("different overwrite mode")); + paimon_error_free(error); + + assert!(crate::runtime() + .block_on(SnapshotManager::new(file_io, path.to_string()).get_latest_snapshot()) + .unwrap() + .is_none()); + + paimon_postpone_fixed_bucket_table_commit_free(commit); + paimon_postpone_fixed_bucket_write_builder_free(overwrite_wb); + paimon_postpone_fixed_bucket_commit_messages_free(messages); + paimon_postpone_fixed_bucket_table_write_free(tw); + paimon_postpone_fixed_bucket_write_builder_free(append_wb); + unwrap_table(handle); + } +} + #[test] fn test_commit_messages_live_until_explicit_free() { const CHILD_ENV: &str = "PAIMON_C_MESSAGES_LIFETIME_CHILD"; @@ -1442,6 +1556,234 @@ fn test_commit_messages_merge_preserves_all_writer_files() { } } +#[test] +fn test_postpone_bucket_plan_arrow_ownership_on_errors() { + let path = "memory:/test_postpone_bucket_plan_arrow_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) }; + + unsafe { + let (mut array, mut schema) = + export_batch_to_ffi(make_postpone_bucket_plan_batch(vec!["p"], vec![1])); + let error = paimon_postpone_fixed_bucket_write_builder_with_bucket_plan( + ptr::null_mut(), + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ); + assert!(!error.is_null()); + assert!(!array.is_released()); + assert!(schema.release.is_some()); + paimon_error_free(error); + ManuallyDrop::drop(array.as_mut()); + ManuallyDrop::drop(schema.as_mut()); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(handle).write_builder; + let (array, schema) = + export_batch_to_ffi(make_postpone_bucket_plan_batch(vec!["p"], vec![0])); + let error = paimon_postpone_fixed_bucket_write_builder_with_bucket_plan( + fixed, + (&**array) as *const FFI_ArrowArray as *mut c_void, + (&**schema) as *const FFI_ArrowSchema as *mut c_void, + ); + assert!(!error.is_null()); + assert!(array.is_released()); + assert!(schema.release.is_none()); + paimon_error_free(error); + + paimon_postpone_fixed_bucket_write_builder_free(fixed); + unwrap_table(handle); + } +} + +#[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!(!normal.write_builder.is_null()); + paimon_write_builder_free(normal.write_builder); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(handle); + assert!(fixed.error.is_null()); + assert!(!fixed.write_builder.is_null()); + let write = paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_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_postpone_fixed_bucket_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_postpone_fixed_bucket_write_builder_with_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_postpone_fixed_bucket_write_builder_new_write(wb1).write; + let tw2 = paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_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_postpone_fixed_bucket_table_write_prepare_commit(tw1).messages; + let messages2 = paimon_postpone_fixed_bucket_table_write_prepare_commit(tw2).messages; + for messages in [messages1, messages2] { + let state = &*((*messages).inner as *const PostponeFixedBucketCommitMessagesState); + assert!(!state.messages.is_empty()); + assert!(state + .messages + .iter() + .all(|message| message.total_buckets == Some(3))); + } + let error = paimon_postpone_fixed_bucket_commit_messages_merge(messages1, messages2); + assert!(error.is_null()); + let commit = paimon_postpone_fixed_bucket_write_builder_new_commit(wb1).commit; + let error = paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_table_commit_free(commit); + paimon_postpone_fixed_bucket_commit_messages_free(messages2); + paimon_postpone_fixed_bucket_commit_messages_free(messages1); + paimon_postpone_fixed_bucket_table_write_free(tw2); + paimon_postpone_fixed_bucket_table_write_free(tw1); + paimon_postpone_fixed_bucket_write_builder_free(wb2); + paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_write_builder_with_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_postpone_fixed_bucket_write_builder_new_write(wb1).write; + let tw2 = paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_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_postpone_fixed_bucket_table_write_prepare_commit(tw1).messages; + let messages2 = paimon_postpone_fixed_bucket_table_write_prepare_commit(tw2).messages; + let error = paimon_postpone_fixed_bucket_commit_messages_merge(messages1, messages2); + assert!(error.is_null()); + let commit = paimon_postpone_fixed_bucket_write_builder_new_commit(wb1).commit; + let error = paimon_postpone_fixed_bucket_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_postpone_fixed_bucket_table_commit_free(commit); + paimon_postpone_fixed_bucket_commit_messages_free(messages2); + paimon_postpone_fixed_bucket_commit_messages_free(messages1); + paimon_postpone_fixed_bucket_table_write_free(tw2); + paimon_postpone_fixed_bucket_table_write_free(tw1); + paimon_postpone_fixed_bucket_write_builder_free(wb2); + paimon_postpone_fixed_bucket_write_builder_free(wb1); + unwrap_table(handle); + } +} #[test] fn test_write_multiple_batches() { let path = "memory:/test_write_multi_batch"; @@ -1721,34 +2063,78 @@ 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()); paimon_error_free(result.error); + let result = paimon_postpone_fixed_bucket_write_builder_new_write(ptr::null()); + assert!(!result.error.is_null()); + assert!(result.write.is_null()); + paimon_error_free(result.error); + let result = paimon_write_builder_new_commit(ptr::null()); assert!(!result.error.is_null()); assert!(result.commit.is_null()); paimon_error_free(result.error); + let result = paimon_postpone_fixed_bucket_write_builder_new_commit(ptr::null()); + assert!(!result.error.is_null()); + assert!(result.commit.is_null()); + paimon_error_free(result.error); + + let err = paimon_postpone_fixed_bucket_write_builder_with_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()); paimon_error_free(err); + let err = paimon_postpone_fixed_bucket_table_write_write_arrow_batch( + ptr::null_mut(), + ptr::null_mut(), + ptr::null_mut(), + ); + assert!(!err.is_null()); + paimon_error_free(err); + let result = paimon_table_write_prepare_commit(ptr::null_mut()); assert!(!result.error.is_null()); assert!(result.messages.is_null()); paimon_error_free(result.error); + let result = paimon_postpone_fixed_bucket_table_write_prepare_commit(ptr::null_mut()); + assert!(!result.error.is_null()); + assert!(result.messages.is_null()); + paimon_error_free(result.error); + let err = paimon_table_commit_commit(ptr::null(), ptr::null_mut()); assert!(!err.is_null()); paimon_error_free(err); + let err = paimon_postpone_fixed_bucket_table_commit_commit(ptr::null(), ptr::null_mut()); + assert!(!err.is_null()); + paimon_error_free(err); + paimon_write_builder_free(ptr::null_mut()); paimon_table_write_free(ptr::null_mut()); paimon_table_commit_free(ptr::null_mut()); paimon_commit_messages_free(ptr::null_mut()); + paimon_postpone_fixed_bucket_write_builder_free(ptr::null_mut()); + paimon_postpone_fixed_bucket_table_write_free(ptr::null_mut()); + paimon_postpone_fixed_bucket_table_commit_free(ptr::null_mut()); + paimon_postpone_fixed_bucket_commit_messages_free(ptr::null_mut()); } } diff --git a/bindings/c/src/types.rs b/bindings/c/src/types.rs index 1dfd9c3f2..98ada0359 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,15 +208,30 @@ pub struct paimon_arrow_batch { // === Write/Commit opaque types === -/// Internal state for WriteBuilder that stores table, shared commit_user, and overwrite flag. pub(crate) struct WriteBuilderState { pub table: Table, pub commit_user: String, pub overwrite: bool, } +pub(crate) struct PostponeFixedBucketWriteBuilderState { + pub table: Table, + pub commit_user: String, + pub overwrite: bool, + pub bucket_plan: Option, +} + pub(crate) struct TableWriteState { - pub write: TableWrite, + pub write: Box, + pub overwrite: bool, + pub target_schema: Arc, + pub table_location: String, + pub commit_user: String, +} + +pub(crate) struct PostponeFixedBucketTableWriteState { + pub write: Box, + pub overwrite: bool, pub target_schema: Arc, pub table_location: String, pub commit_user: String, @@ -221,12 +239,28 @@ pub(crate) struct TableWriteState { pub(crate) struct TableCommitState { pub commit: TableCommit, + pub overwrite: bool, + pub table_location: String, + pub commit_user: String, +} + +pub(crate) struct PostponeFixedBucketTableCommitState { + pub commit: PostponeFixedBucketTableCommit, + pub overwrite: bool, pub table_location: String, pub commit_user: String, } pub(crate) struct CommitMessagesState { pub messages: Vec, + pub overwrite: bool, + pub table_location: String, + pub commit_user: String, +} + +pub(crate) struct PostponeFixedBucketCommitMessagesState { + pub messages: Vec, + pub overwrite: bool, pub table_location: String, pub commit_user: String, } @@ -251,3 +285,23 @@ pub struct paimon_table_commit { pub struct paimon_commit_messages { pub inner: *mut c_void, } + +#[repr(C)] +pub struct paimon_postpone_fixed_bucket_write_builder { + pub inner: *mut c_void, +} + +#[repr(C)] +pub struct paimon_postpone_fixed_bucket_table_write { + pub inner: *mut c_void, +} + +#[repr(C)] +pub struct paimon_postpone_fixed_bucket_table_commit { + pub inner: *mut c_void, +} + +#[repr(C)] +pub struct paimon_postpone_fixed_bucket_commit_messages { + pub inner: *mut c_void, +} diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs index 8e9637997..c6e9afa7d 100644 --- a/bindings/c/src/write.rs +++ b/bindings/c/src/write.rs @@ -22,12 +22,15 @@ 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::{ - paimon_result_prepare_commit, paimon_result_table_commit, paimon_result_table_write, - paimon_result_write_builder, + paimon_result_postpone_fixed_bucket_prepare_commit, + paimon_result_postpone_fixed_bucket_table_commit, + paimon_result_postpone_fixed_bucket_table_write, + paimon_result_postpone_fixed_bucket_write_builder, paimon_result_prepare_commit, + paimon_result_table_commit, paimon_result_table_write, paimon_result_write_builder, }; use crate::runtime; use crate::types::*; @@ -46,30 +49,105 @@ 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) => builder + .with_commit_user(commit_user) + .map(|builder| builder.commit_user().to_string()), + None => Ok(builder.commit_user().to_string()), + }; + let state = match commit_user { + Ok(commit_user) => WriteBuilderState { + table: table_ref.clone(), + commit_user, + overwrite: false, + }, + Err(error) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error: paimon_error::from_paimon(error), + } + } + }; + let inner = Box::into_raw(Box::new(state)) as *mut c_void; + paimon_result_write_builder { + write_builder: Box::into_raw(Box::new(paimon_write_builder { inner })), + error: ptr::null_mut(), + } +} + +unsafe fn new_postpone_fixed_bucket_write_builder( + table: *const paimon_table, + commit_user: Option, +) -> paimon_result_postpone_fixed_bucket_write_builder { + if let Err(error) = check_non_null(table, "table") { + return paimon_result_postpone_fixed_bucket_write_builder { + write_builder: ptr::null_mut(), + error, + }; + } + let table_ref = &*((*table).inner as *const Table); + let builder = match table_ref.new_postpone_fixed_bucket_write_builder() { + Ok(builder) => builder, + Err(error) => { + return paimon_result_postpone_fixed_bucket_write_builder { + write_builder: ptr::null_mut(), + error: paimon_error::from_paimon(error), + } + } + }; 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 { + Err(error) => { + return paimon_result_postpone_fixed_bucket_write_builder { write_builder: ptr::null_mut(), - error: paimon_error::from_paimon(e), + error: paimon_error::from_paimon(error), } } }, None => builder.commit_user().to_string(), }; - let state = WriteBuilderState { + let state = PostponeFixedBucketWriteBuilderState { table: table_ref.clone(), commit_user, overwrite: false, + bucket_plan: None, }; let inner = Box::into_raw(Box::new(state)) as *mut c_void; - paimon_result_write_builder { - write_builder: Box::into_raw(Box::new(paimon_write_builder { inner })), + paimon_result_postpone_fixed_bucket_write_builder { + write_builder: Box::into_raw(Box::new(paimon_postpone_fixed_bucket_write_builder { + inner, + })), error: ptr::null_mut(), } } +unsafe fn new_write_builder_with_commit_user( + table: *const paimon_table, + commit_user: *const c_char, +) -> paimon_result_write_builder { + match validate_cstr(commit_user, "commit_user") { + Ok(commit_user) => new_write_builder(table, Some(commit_user)), + Err(error) => paimon_result_write_builder { + write_builder: ptr::null_mut(), + error, + }, + } +} + +unsafe fn new_postpone_fixed_bucket_write_builder_with_commit_user( + table: *const paimon_table, + commit_user: *const c_char, +) -> paimon_result_postpone_fixed_bucket_write_builder { + match validate_cstr(commit_user, "commit_user") { + Ok(commit_user) => new_postpone_fixed_bucket_write_builder(table, Some(commit_user)), + Err(error) => paimon_result_postpone_fixed_bucket_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 @@ -85,6 +163,18 @@ pub unsafe extern "C" fn paimon_table_new_write_builder( new_write_builder(table, None) } +/// 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_postpone_fixed_bucket_write_builder { + new_postpone_fixed_bucket_write_builder(table, None) +} + /// Create a WriteBuilder with a caller-provided stable commit identity. /// /// Writers whose messages are merged into one logical commit must use the @@ -98,16 +188,21 @@ 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) +} + +/// 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_postpone_fixed_bucket_write_builder { + new_postpone_fixed_bucket_write_builder_with_commit_user(table, commit_user) } /// Free a paimon_write_builder. @@ -126,9 +221,6 @@ 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. -/// /// # Safety /// `wb` must be a valid pointer from `paimon_table_new_write_builder`, or null (returns error). #[no_mangle] @@ -143,6 +235,79 @@ pub unsafe extern "C" fn paimon_write_builder_with_overwrite( ptr::null_mut() } +/// Free a postpone fixed-bucket write builder. +/// +/// # Safety +/// Only call with a builder returned from +/// `paimon_table_new_postpone_fixed_bucket_write_builder`. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_free( + wb: *mut paimon_postpone_fixed_bucket_write_builder, +) { + if !wb.is_null() { + let wrapper = Box::from_raw(wb); + if !wrapper.inner.is_null() { + drop(Box::from_raw( + wrapper.inner as *mut PostponeFixedBucketWriteBuilderState, + )); + } + } +} + +/// Enable overwrite mode for a postpone fixed-bucket write operation. +/// +/// # Safety +/// `wb` must be a valid fixed-bucket builder, or null (returns error). +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_with_overwrite( + wb: *mut paimon_postpone_fixed_bucket_write_builder, +) -> *mut paimon_error { + if let Err(error) = check_non_null(wb, "wb") { + return error; + } + let state = &mut *((*wb).inner as *mut PostponeFixedBucketWriteBuilderState); + state.overwrite = true; + ptr::null_mut() +} + +/// Set a shared `partition -> total_buckets` plan. +/// The caller retains ownership when pointer or builder validation fails. Once +/// Arrow import starts, this call consumes both structs even if plan validation +/// returns an error. +/// +/// # Safety +/// `wb` must be a valid postpone fixed-bucket builder. `array` and +/// `schema` must point to initialized Arrow C Data structs. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_with_bucket_plan( + wb: *mut paimon_postpone_fixed_bucket_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 PostponeFixedBucketWriteBuilderState); + + let batch = match import_record_batch(array, schema) { + Ok(batch) => batch, + Err(error) => return error, + }; + match PostponeBucketPlan::from_arrow(&state.table, &batch) { + Ok(plan) => { + state.bucket_plan = Some(plan); + ptr::null_mut() + } + Err(error) => paimon_error::from_paimon(error), + } +} + // ======================= TableWrite =============================== fn invalid_input(message: impl Into) -> *mut paimon_error { @@ -227,82 +392,127 @@ unsafe fn import_record_batch( .map_err(|e| invalid_input(format!("Failed to construct Arrow record batch: {e}"))) } -/// Create a new TableWrite from the WriteBuilder. -/// -/// The returned TableWrite accumulates Arrow batches until `prepare_commit` is called. +/// Create a standard TableWrite from a standard WriteBuilder. /// /// # Safety -/// `wb` must be a valid pointer from `paimon_table_new_write_builder`, or null (returns error). +/// wb must be a valid standard builder, or null (returns error). #[no_mangle] pub unsafe extern "C" fn paimon_write_builder_new_write( wb: *const paimon_write_builder, ) -> paimon_result_table_write { - if let Err(e) = check_non_null(wb, "wb") { + if let Err(error) = check_non_null(wb, "wb") { return paimon_result_table_write { write: ptr::null_mut(), - error: e, + error, }; } 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) => { + Ok(builder) => builder, + Err(error) => { return paimon_result_table_write { write: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } + error: paimon_error::from_paimon(error), + }; } }; - if state.overwrite { builder = builder.with_overwrite(); } - - let tw = match builder.new_write() { - Ok(w) => w, - Err(e) => { + let result = builder.new_write().and_then(|write| { + paimon::arrow::build_target_arrow_schema(state.table.schema().fields()) + .map(|schema| (Box::new(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, + let inner = Box::into_raw(Box::new(TableWriteState { + write, + overwrite: state.overwrite, target_schema, table_location: state.table.location().to_string(), commit_user: state.commit_user.clone(), - }; - let inner = Box::into_raw(Box::new(table_write)) as *mut c_void; + })) as *mut c_void; paimon_result_table_write { write: Box::into_raw(Box::new(paimon_table_write { inner })), error: ptr::null_mut(), } } -/// Free a paimon_table_write. +/// Create a postpone fixed-bucket TableWrite. /// -/// Dropping a TableWrite before calling `prepare_commit` discards any -/// uncommitted data. +/// # Safety +/// wb must be a valid fixed-bucket builder, or null (returns error). +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_new_write( + wb: *const paimon_postpone_fixed_bucket_write_builder, +) -> paimon_result_postpone_fixed_bucket_table_write { + if let Err(error) = check_non_null(wb, "wb") { + return paimon_result_postpone_fixed_bucket_table_write { + write: ptr::null_mut(), + error, + }; + } + let state = &*((*wb).inner as *const PostponeFixedBucketWriteBuilderState); + let mut builder = match state + .table + .new_postpone_fixed_bucket_write_builder() + .and_then(|builder| builder.with_commit_user(state.commit_user.clone())) + { + Ok(builder) => builder, + Err(error) => { + return paimon_result_postpone_fixed_bucket_table_write { + write: ptr::null_mut(), + error: paimon_error::from_paimon(error), + }; + } + }; + if let Some(plan) = state.bucket_plan.clone() { + builder = builder.with_bucket_plan(plan); + } + if state.overwrite { + builder = builder.with_overwrite(); + } + let result = builder.new_write().and_then(|write| { + paimon::arrow::build_target_arrow_schema(state.table.schema().fields()) + .map(|schema| (Box::new(write), schema)) + }); + let (write, target_schema) = match result { + Ok(result) => result, + Err(error) => { + return paimon_result_postpone_fixed_bucket_table_write { + write: ptr::null_mut(), + error: paimon_error::from_paimon(error), + }; + } + }; + let inner = Box::into_raw(Box::new(PostponeFixedBucketTableWriteState { + write, + overwrite: state.overwrite, + target_schema, + table_location: state.table.location().to_string(), + commit_user: state.commit_user.clone(), + })) as *mut c_void; + paimon_result_postpone_fixed_bucket_table_write { + write: Box::into_raw(Box::new(paimon_postpone_fixed_bucket_table_write { inner })), + error: ptr::null_mut(), + } +} + +/// Free a standard TableWrite. /// /// # Safety -/// Only call with a write returned from `paimon_write_builder_new_write`. +/// Only call with a write returned from paimon_write_builder_new_write. #[no_mangle] pub unsafe extern "C" fn paimon_table_write_free(tw: *mut paimon_table_write) { if !tw.is_null() { @@ -313,153 +523,247 @@ pub unsafe extern "C" fn paimon_table_write_free(tw: *mut paimon_table_write) { } } -/// Write a single Arrow record batch into the table's writers. -/// -/// The Arrow data is imported via the Arrow C Data Interface. `array` and -/// `schema` must point to valid `ArrowArray` and `ArrowSchema` structs -/// filled by the caller. Ownership is transferred — the caller must not -/// release the structs after this call. +/// Free a postpone fixed-bucket TableWrite. /// /// # Safety -/// `tw` must be a valid pointer from `paimon_write_builder_new_write`, or null (returns error). -/// `array` and `schema` must be valid pointers to initialized ArrowArray / -/// ArrowSchema structs, or null (returns error). +/// Only call with a write returned from +/// paimon_postpone_fixed_bucket_write_builder_new_write. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_write_free( + tw: *mut paimon_postpone_fixed_bucket_table_write, +) { + if !tw.is_null() { + let wrapper = Box::from_raw(tw); + if !wrapper.inner.is_null() { + drop(Box::from_raw( + wrapper.inner as *mut PostponeFixedBucketTableWriteState, + )); + } + } +} + +unsafe fn import_write_batch( + array: *mut c_void, + schema: *mut c_void, + target_schema: &ArrowSchema, +) -> Result { + check_non_null(array, "array")?; + check_non_null(schema, "schema")?; + let batch = import_record_batch(array, schema)?; + validate_batch_schema(&batch, target_schema)?; + Ok(batch) +} + +/// Write one Arrow record batch with a standard TableWrite. +/// +/// Ownership of array and schema is transferred once Arrow import starts. #[no_mangle] pub unsafe extern "C" fn paimon_table_write_write_arrow_batch( tw: *mut paimon_table_write, array: *mut c_void, schema: *mut c_void, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tw, "tw") { - return e; - } - if let Err(e) = check_non_null(array, "array") { - return e; - } - if let Err(e) = check_non_null(schema, "schema") { - return e; + if let Err(error) = check_non_null(tw, "tw") { + return error; } - let table_write = &mut *((*tw).inner as *mut TableWriteState); - let batch = match import_record_batch(array, schema) { + let batch = match import_write_batch(array, schema, &table_write.target_schema) { Ok(batch) => batch, Err(error) => return error, }; - if let Err(error) = validate_batch_schema(&batch, &table_write.target_schema) { - return error; + match runtime().block_on(table_write.write.write_arrow_batch(&batch)) { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), } +} +/// Write one Arrow record batch with a postpone fixed-bucket TableWrite. +/// +/// Ownership of array and schema is transferred once Arrow import starts. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_write_write_arrow_batch( + tw: *mut paimon_postpone_fixed_bucket_table_write, + array: *mut c_void, + schema: *mut c_void, +) -> *mut paimon_error { + if let Err(error) = check_non_null(tw, "tw") { + return error; + } + let table_write = &mut *((*tw).inner as *mut PostponeFixedBucketTableWriteState); + let batch = match import_write_batch(array, schema, &table_write.target_schema) { + Ok(batch) => batch, + Err(error) => return error, + }; match runtime().block_on(table_write.write.write_arrow_batch(&batch)) { Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), + Err(error) => paimon_error::from_paimon(error), } } -/// 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. -/// -/// The returned `paimon_commit_messages` must be passed to a -/// `paimon_table_commit_*` function and then freed with -/// `paimon_commit_messages_free`. +/// Prepare standard commit messages. /// -/// # Safety -/// `tw` must be a valid pointer from `paimon_write_builder_new_write`, or null (returns error). +/// The returned handle remains owned by the caller. #[no_mangle] pub unsafe extern "C" fn paimon_table_write_prepare_commit( tw: *mut paimon_table_write, ) -> paimon_result_prepare_commit { - if let Err(e) = check_non_null(tw, "tw") { + if let Err(error) = check_non_null(tw, "tw") { return paimon_result_prepare_commit { messages: ptr::null_mut(), - error: e, + error, }; } let table_write = &mut *((*tw).inner as *mut TableWriteState); - match runtime().block_on(table_write.write.prepare_commit()) { Ok(messages) => { - let messages = CommitMessagesState { + let inner = Box::into_raw(Box::new(CommitMessagesState { messages, + overwrite: table_write.overwrite, table_location: table_write.table_location.clone(), commit_user: table_write.commit_user.clone(), - }; - let inner = Box::into_raw(Box::new(messages)) as *mut c_void; + })) as *mut c_void; paimon_result_prepare_commit { messages: Box::into_raw(Box::new(paimon_commit_messages { inner })), error: ptr::null_mut(), } } - Err(e) => paimon_result_prepare_commit { + Err(error) => paimon_result_prepare_commit { messages: ptr::null_mut(), - error: paimon_error::from_paimon(e), + error: paimon_error::from_paimon(error), + }, + } +} + +/// Prepare postpone fixed-bucket commit messages. +/// +/// The returned handle remains owned by the caller. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_write_prepare_commit( + tw: *mut paimon_postpone_fixed_bucket_table_write, +) -> paimon_result_postpone_fixed_bucket_prepare_commit { + if let Err(error) = check_non_null(tw, "tw") { + return paimon_result_postpone_fixed_bucket_prepare_commit { + messages: ptr::null_mut(), + error, + }; + } + let table_write = &mut *((*tw).inner as *mut PostponeFixedBucketTableWriteState); + match runtime().block_on(table_write.write.prepare_commit()) { + Ok(messages) => { + let inner = Box::into_raw(Box::new(PostponeFixedBucketCommitMessagesState { + messages, + overwrite: table_write.overwrite, + table_location: table_write.table_location.clone(), + commit_user: table_write.commit_user.clone(), + })) as *mut c_void; + paimon_result_postpone_fixed_bucket_prepare_commit { + messages: Box::into_raw(Box::new(paimon_postpone_fixed_bucket_commit_messages { + inner, + })), + error: ptr::null_mut(), + } + } + Err(error) => paimon_result_postpone_fixed_bucket_prepare_commit { + messages: ptr::null_mut(), + error: paimon_error::from_paimon(error), }, } } // ======================= TableCommit =============================== -/// Create a new TableCommit from the WriteBuilder. -/// -/// The committer shares the same `commit_user` as the writer, which is -/// required for duplicate-commit detection. -/// -/// # Safety -/// `wb` must be a valid pointer from `paimon_table_new_write_builder`, or null (returns error). +/// Create a standard TableCommit from a standard WriteBuilder. #[no_mangle] pub unsafe extern "C" fn paimon_write_builder_new_commit( wb: *const paimon_write_builder, ) -> paimon_result_table_commit { - if let Err(e) = check_non_null(wb, "wb") { + if let Err(error) = check_non_null(wb, "wb") { return paimon_result_table_commit { commit: ptr::null_mut(), - error: e, + error, }; } let state = &*((*wb).inner as *const WriteBuilderState); - - let builder = match state + let commit = match state .table .new_write_builder() .with_commit_user(state.commit_user.clone()) + .and_then(|builder| builder.try_new_commit()) { - Ok(b) => b, - Err(e) => { + 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 inner = Box::into_raw(Box::new(TableCommitState { + commit, + overwrite: state.overwrite, + table_location: state.table.location().to_string(), + commit_user: state.commit_user.clone(), + })) as *mut c_void; + paimon_result_table_commit { + commit: Box::into_raw(Box::new(paimon_table_commit { inner })), + error: ptr::null_mut(), + } +} - let tc = match builder.try_new_commit() { - Ok(c) => c, - Err(e) => { - return paimon_result_table_commit { +/// Create a postpone fixed-bucket TableCommit. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_new_commit( + wb: *const paimon_postpone_fixed_bucket_write_builder, +) -> paimon_result_postpone_fixed_bucket_table_commit { + if let Err(error) = check_non_null(wb, "wb") { + return paimon_result_postpone_fixed_bucket_table_commit { + commit: ptr::null_mut(), + error, + }; + } + let state = &*((*wb).inner as *const PostponeFixedBucketWriteBuilderState); + let builder = match state + .table + .new_postpone_fixed_bucket_write_builder() + .and_then(|builder| builder.with_commit_user(state.commit_user.clone())) + { + Ok(builder) => builder, + Err(error) => { + return paimon_result_postpone_fixed_bucket_table_commit { commit: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } + error: paimon_error::from_paimon(error), + }; } }; - - let table_commit = TableCommitState { - commit: tc, + let builder = if state.overwrite { + builder.with_overwrite() + } else { + builder + }; + let commit = match builder.try_new_commit() { + Ok(commit) => commit, + Err(error) => { + return paimon_result_postpone_fixed_bucket_table_commit { + commit: ptr::null_mut(), + error: paimon_error::from_paimon(error), + }; + } + }; + let inner = Box::into_raw(Box::new(PostponeFixedBucketTableCommitState { + commit, + overwrite: state.overwrite, table_location: state.table.location().to_string(), commit_user: state.commit_user.clone(), - }; - let inner = Box::into_raw(Box::new(table_commit)) as *mut c_void; - paimon_result_table_commit { - commit: Box::into_raw(Box::new(paimon_table_commit { inner })), + })) as *mut c_void; + paimon_result_postpone_fixed_bucket_table_commit { + commit: Box::into_raw(Box::new(paimon_postpone_fixed_bucket_table_commit { + inner, + })), error: ptr::null_mut(), } } -/// Free a paimon_table_commit. -/// -/// # Safety -/// Only call with a commit returned from `paimon_write_builder_new_commit`. +/// Free a standard TableCommit. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_free(tc: *mut paimon_table_commit) { if !tc.is_null() { @@ -470,12 +774,24 @@ pub unsafe extern "C" fn paimon_table_commit_free(tc: *mut paimon_table_commit) } } +/// Free a postpone fixed-bucket TableCommit. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_free( + tc: *mut paimon_postpone_fixed_bucket_table_commit, +) { + if !tc.is_null() { + let wrapper = Box::from_raw(tc); + if !wrapper.inner.is_null() { + drop(Box::from_raw( + wrapper.inner as *mut PostponeFixedBucketTableCommitState, + )); + } + } +} + // ======================= CommitMessages =============================== -/// Free a paimon_commit_messages. -/// -/// # Safety -/// Only call with messages returned from `paimon_table_write_prepare_commit`. +/// Free standard commit messages. #[no_mangle] pub unsafe extern "C" fn paimon_commit_messages_free(msgs: *mut paimon_commit_messages) { if !msgs.is_null() { @@ -486,13 +802,43 @@ pub unsafe extern "C" fn paimon_commit_messages_free(msgs: *mut paimon_commit_me } } -/// Merge `source` messages into `target` for one logical commit. -/// -/// Both handles retain ownership and must be freed separately. They must have -/// been prepared for the same table and `commit_user`. -/// -/// # Safety -/// `target` and `source` must be distinct valid commit-message handles. +/// Free postpone fixed-bucket commit messages. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_commit_messages_free( + msgs: *mut paimon_postpone_fixed_bucket_commit_messages, +) { + if !msgs.is_null() { + let wrapper = Box::from_raw(msgs); + if !wrapper.inner.is_null() { + drop(Box::from_raw( + wrapper.inner as *mut PostponeFixedBucketCommitMessagesState, + )); + } + } +} + +fn validate_message_context( + target_table: &str, + target_user: &str, + target_overwrite: bool, + source_table: &str, + source_user: &str, + source_overwrite: bool, +) -> Result<(), *mut paimon_error> { + if target_table != source_table || target_user != source_user { + return Err(invalid_input( + "commit messages can only be merged when table and commit_user both match", + )); + } + if target_overwrite != source_overwrite { + return Err(invalid_input( + "commit messages can only be merged when overwrite modes match", + )); + } + Ok(()) +} + +/// Merge standard commit messages for one logical commit. #[no_mangle] pub unsafe extern "C" fn paimon_commit_messages_merge( target: *mut paimon_commit_messages, @@ -507,13 +853,48 @@ pub unsafe extern "C" fn paimon_commit_messages_merge( if ptr::eq(target, source.cast_mut()) { return invalid_input("target and source commit messages must be distinct handles"); } - let target = &mut *((*target).inner as *mut CommitMessagesState); let source = &*((*source).inner as *const CommitMessagesState); - if target.table_location != source.table_location || target.commit_user != source.commit_user { - return invalid_input( - "commit messages can only be merged when table and commit_user both match", - ); + if let Err(error) = validate_message_context( + &target.table_location, + &target.commit_user, + target.overwrite, + &source.table_location, + &source.commit_user, + source.overwrite, + ) { + return error; + } + target.messages.extend(source.messages.clone()); + ptr::null_mut() +} + +/// Merge postpone fixed-bucket messages for one logical commit. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_commit_messages_merge( + target: *mut paimon_postpone_fixed_bucket_commit_messages, + source: *const paimon_postpone_fixed_bucket_commit_messages, +) -> *mut paimon_error { + if let Err(error) = check_non_null(target, "target") { + return error; + } + if let Err(error) = check_non_null(source, "source") { + return error; + } + if ptr::eq(target, source.cast_mut()) { + return invalid_input("target and source commit messages must be distinct handles"); + } + let target = &mut *((*target).inner as *mut PostponeFixedBucketCommitMessagesState); + let source = &*((*source).inner as *const PostponeFixedBucketCommitMessagesState); + if let Err(error) = validate_message_context( + &target.table_location, + &target.commit_user, + target.overwrite, + &source.table_location, + &source.commit_user, + source.overwrite, + ) { + return error; } target.messages.extend(source.messages.clone()); ptr::null_mut() @@ -522,32 +903,82 @@ pub unsafe extern "C" fn paimon_commit_messages_merge( // ======================= Commit operations =============================== fn validate_commit_context( - commit: &TableCommitState, - messages: &CommitMessagesState, + commit_table: &str, + commit_user: &str, + commit_overwrite: bool, + message_table: &str, + message_user: &str, + message_overwrite: bool, ) -> Result<(), *mut paimon_error> { - if commit.table_location != messages.table_location { + if commit_table != message_table { return Err(invalid_input(format!( "commit messages were prepared for a different table (message table '{}', committer table '{}')", - messages.table_location, commit.table_location + message_table, commit_table ))); } - if commit.commit_user != messages.commit_user { + if commit_user != message_user { return Err(invalid_input( "commit messages were prepared with a different commit_user", )); } + if commit_overwrite != message_overwrite { + return Err(invalid_input( + "commit messages were prepared with a different overwrite mode", + )); + } Ok(()) } -/// Commit the given messages in APPEND mode. -/// -/// Empty messages is a no-op success. -/// The caller retains ownership of `msgs`; it may retry after an error and -/// must eventually release the handle with `paimon_commit_messages_free`. -/// -/// # Safety -/// `tc` must be a valid pointer from `paimon_write_builder_new_commit`, or null (returns error). -/// `msgs` must be a valid pointer from `paimon_table_write_prepare_commit`, or null (returns error). +unsafe fn standard_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.table_location, + &table_commit.commit_user, + table_commit.overwrite, + &messages.table_location, + &messages.commit_user, + messages.overwrite, + ) { + return error; + } + if messages.overwrite { + return invalid_input( + "standard overwrite messages must be committed with paimon_table_commit_overwrite", + ); + } + let messages = messages.messages.clone(); + let result = if filter_committed { + runtime().block_on( + table_commit + .commit + .filter_and_commit_with_identifier(messages, commit_identifier), + ) + } else { + runtime().block_on( + table_commit + .commit + .commit_with_identifier(messages, commit_identifier), + ) + }; + match result { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), + } +} + +/// Commit standard append messages. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_commit( tc: *const paimon_table_commit, @@ -556,153 +987,89 @@ pub unsafe extern "C" fn paimon_table_commit_commit( paimon_table_commit_commit_with_identifier(tc, msgs, i64::MAX) } -/// Commit the given messages with a caller-provided identifier. -/// -/// Identifiers must increase monotonically for a `commit_user`. All messages -/// for one identifier must be merged and submitted in a single call. -/// This operation does not filter previously committed identifiers. Use -/// `paimon_table_commit_filter_and_commit_with_identifier` when retrying an -/// uncertain result. -/// The caller retains ownership of `msgs` and must free it explicitly. +/// Commit standard append messages with an identifier. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_commit_with_identifier( tc: *const paimon_table_commit, 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), - } + standard_commit_with_identifier_impl(tc, msgs, commit_identifier, false) } -/// Filter a previously committed identifier, then commit if it is new. -/// -/// Identifiers must increase monotonically for a `commit_user`. Use this only -/// to retry the same uncertain result after all writer messages for the -/// logical commit have been merged. -/// The caller retains ownership of `msgs` and must free it explicitly. +/// Filter a committed identifier before committing standard append messages. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_filter_and_commit_with_identifier( tc: *const paimon_table_commit, 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), - } + standard_commit_with_identifier_impl(tc, msgs, commit_identifier, true) } -/// Commit in OVERWRITE mode, replacing data in the written partitions. -/// -/// `static_partitions` is currently passed as `None` (overwrite all -/// partitions that were written to). -/// The caller retains ownership of `msgs`; it may retry after an error and -/// must eventually release the handle with `paimon_commit_messages_free`. -/// -/// # Safety -/// `tc` must be a valid pointer from `paimon_write_builder_new_commit`, or null (returns error). -/// `msgs` must be a valid pointer from `paimon_table_write_prepare_commit`, or null (returns error). +/// Commit standard overwrite messages. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_overwrite( tc: *const paimon_table_commit, msgs: *mut paimon_commit_messages, ) -> *mut paimon_error { - paimon_table_commit_overwrite_impl(tc, msgs, None) + standard_overwrite_impl(tc, msgs, None) } -/// Overwrite with a caller-provided stable commit identifier. -/// -/// The caller retains ownership of `msgs` and must free it explicitly. +/// Commit standard overwrite messages with an identifier. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_overwrite_with_identifier( tc: *const paimon_table_commit, msgs: *mut paimon_commit_messages, commit_identifier: i64, ) -> *mut paimon_error { - paimon_table_commit_overwrite_impl(tc, msgs, Some(commit_identifier)) + standard_overwrite_impl(tc, msgs, Some(commit_identifier)) } -unsafe fn paimon_table_commit_overwrite_impl( +unsafe fn standard_overwrite_impl( tc: *const paimon_table_commit, msgs: *mut paimon_commit_messages, commit_identifier: Option, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tc, "tc") { - return e; + if let Err(error) = check_non_null(tc, "tc") { + return error; } - if let Err(e) = check_non_null(msgs, "msgs") { - return e; + 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) { + if let Err(error) = validate_commit_context( + &table_commit.table_location, + &table_commit.commit_user, + table_commit.overwrite, + &messages.table_location, + &messages.commit_user, + messages.overwrite, + ) { return error; } - + if !messages.overwrite { + return invalid_input( + "append messages cannot be committed with paimon_table_commit_overwrite", + ); + } + let messages = messages.messages.clone(); let result = match commit_identifier { - Some(commit_identifier) => { - runtime().block_on(table_commit.commit.overwrite_with_identifier( - messages.messages.clone(), - None, - commit_identifier, - )) - } - None => runtime().block_on( + Some(commit_identifier) => runtime().block_on( table_commit .commit - .overwrite(messages.messages.clone(), None), + .overwrite_with_identifier(messages, None, commit_identifier), ), + None => runtime().block_on(table_commit.commit.overwrite(messages, None)), }; match result { Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), + Err(error) => paimon_error::from_paimon(error), } } -/// Truncate the entire table — removes all data. -/// -/// This is an OVERWRITE with zero new files. The table's latest snapshot -/// will have no data. -/// -/// # Safety -/// `tc` must be a valid pointer from `paimon_write_builder_new_commit`, or null (returns error). +/// Truncate a table with a standard TableCommit. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_truncate_table( tc: *const paimon_table_commit, @@ -710,7 +1077,7 @@ pub unsafe extern "C" fn paimon_table_commit_truncate_table( paimon_table_commit_truncate_table_impl(tc, None) } -/// Truncate the table with a caller-provided stable commit identifier. +/// Truncate a table with a stable identifier. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_truncate_table_with_identifier( tc: *const paimon_table_commit, @@ -723,12 +1090,10 @@ unsafe fn paimon_table_commit_truncate_table_impl( tc: *const paimon_table_commit, commit_identifier: Option, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tc, "tc") { - return e; + if let Err(error) = check_non_null(tc, "tc") { + return error; } - let table_commit = &*((*tc).inner as *const TableCommitState); - let result = match commit_identifier { Some(commit_identifier) => runtime().block_on( table_commit @@ -739,41 +1104,179 @@ unsafe fn paimon_table_commit_truncate_table_impl( }; match result { Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), + Err(error) => paimon_error::from_paimon(error), } } -/// Abort a prepared commit, cleaning up written data files. -/// -/// This is a best-effort cleanup — it attempts to delete new data, changelog, -/// and index files produced by the writer. Storage deletion errors are ignored -/// so cleanup does not mask an earlier write or commit failure. -/// The caller retains ownership of `msgs` and must free it explicitly. -/// -/// # Safety -/// `tc` must be a valid pointer from `paimon_write_builder_new_commit`, or null (returns error). -/// `msgs` must be a valid pointer from `paimon_table_write_prepare_commit`, or null (returns error). +/// Abort standard commit messages. #[no_mangle] pub unsafe extern "C" fn paimon_table_commit_abort( tc: *const paimon_table_commit, msgs: *mut paimon_commit_messages, ) -> *mut paimon_error { - if let Err(e) = check_non_null(tc, "tc") { - return e; + if let Err(error) = check_non_null(tc, "tc") { + return error; } - if let Err(e) = check_non_null(msgs, "msgs") { - return e; + 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) { + if let Err(error) = validate_commit_context( + &table_commit.table_location, + &table_commit.commit_user, + table_commit.overwrite, + &messages.table_location, + &messages.commit_user, + messages.overwrite, + ) { + return error; + } + match runtime().block_on(table_commit.commit.abort(&messages.messages)) { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), + } +} + +unsafe fn fixed_commit_with_identifier_impl( + tc: *const paimon_postpone_fixed_bucket_table_commit, + msgs: *mut paimon_postpone_fixed_bucket_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 PostponeFixedBucketTableCommitState); + let messages = &*((*msgs).inner as *const PostponeFixedBucketCommitMessagesState); + if let Err(error) = validate_commit_context( + &table_commit.table_location, + &table_commit.commit_user, + table_commit.overwrite, + &messages.table_location, + &messages.commit_user, + messages.overwrite, + ) { + return error; + } + let messages = messages.messages.clone(); + let result = if filter_committed { + runtime().block_on( + table_commit + .commit + .filter_and_commit_with_identifier(messages, commit_identifier), + ) + } else { + runtime().block_on( + table_commit + .commit + .commit_with_identifier(messages, commit_identifier), + ) + }; + match result { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), + } +} + +/// Commit postpone fixed-bucket messages using the builder's mode. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_commit( + tc: *const paimon_postpone_fixed_bucket_table_commit, + msgs: *mut paimon_postpone_fixed_bucket_commit_messages, +) -> *mut paimon_error { + paimon_postpone_fixed_bucket_table_commit_commit_with_identifier(tc, msgs, i64::MAX) +} + +/// Commit postpone fixed-bucket messages with an identifier. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_commit_with_identifier( + tc: *const paimon_postpone_fixed_bucket_table_commit, + msgs: *mut paimon_postpone_fixed_bucket_commit_messages, + commit_identifier: i64, +) -> *mut paimon_error { + fixed_commit_with_identifier_impl(tc, msgs, commit_identifier, false) +} + +/// Filter a committed identifier before committing fixed-bucket messages. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_filter_and_commit_with_identifier( + tc: *const paimon_postpone_fixed_bucket_table_commit, + msgs: *mut paimon_postpone_fixed_bucket_commit_messages, + commit_identifier: i64, +) -> *mut paimon_error { + fixed_commit_with_identifier_impl(tc, msgs, commit_identifier, true) +} + +/// Truncate a table with a postpone fixed-bucket TableCommit. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_truncate_table( + tc: *const paimon_postpone_fixed_bucket_table_commit, +) -> *mut paimon_error { + fixed_truncate_table_impl(tc, None) +} + +/// Truncate a table with a stable identifier. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_truncate_table_with_identifier( + tc: *const paimon_postpone_fixed_bucket_table_commit, + commit_identifier: i64, +) -> *mut paimon_error { + fixed_truncate_table_impl(tc, Some(commit_identifier)) +} + +unsafe fn fixed_truncate_table_impl( + tc: *const paimon_postpone_fixed_bucket_table_commit, + commit_identifier: Option, +) -> *mut paimon_error { + if let Err(error) = check_non_null(tc, "tc") { return error; } + let table_commit = &*((*tc).inner as *const PostponeFixedBucketTableCommitState); + 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()), + }; + match result { + Ok(()) => ptr::null_mut(), + Err(error) => paimon_error::from_paimon(error), + } +} +/// Abort postpone fixed-bucket commit messages. +#[no_mangle] +pub unsafe extern "C" fn paimon_postpone_fixed_bucket_table_commit_abort( + tc: *const paimon_postpone_fixed_bucket_table_commit, + msgs: *mut paimon_postpone_fixed_bucket_commit_messages, +) -> *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 PostponeFixedBucketTableCommitState); + let messages = &*((*msgs).inner as *const PostponeFixedBucketCommitMessagesState); + if let Err(error) = validate_commit_context( + &table_commit.table_location, + &table_commit.commit_user, + table_commit.overwrite, + &messages.table_location, + &messages.commit_user, + messages.overwrite, + ) { + return error; + } match runtime().block_on(table_commit.commit.abort(&messages.messages)) { Ok(()) => ptr::null_mut(), - Err(e) => paimon_error::from_paimon(e), + Err(error) => paimon_error::from_paimon(error), } } @@ -781,37 +1284,87 @@ 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_postpone_fixed_bucket_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_postpone_fixed_bucket_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_postpone_fixed_bucket_write_builder, +) -> paimon_result_postpone_fixed_bucket_table_write = + paimon_postpone_fixed_bucket_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( + *const paimon_postpone_fixed_bucket_write_builder, +) -> paimon_result_postpone_fixed_bucket_table_commit = + paimon_postpone_fixed_bucket_write_builder_new_commit; +const _: unsafe extern "C" fn( + *mut paimon_postpone_fixed_bucket_write_builder, + *mut c_void, + *mut c_void, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_write_builder_with_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( + *mut paimon_postpone_fixed_bucket_table_write, +) -> paimon_result_postpone_fixed_bucket_prepare_commit = + paimon_postpone_fixed_bucket_table_write_prepare_commit; const _: unsafe extern "C" fn( *mut paimon_commit_messages, *const paimon_commit_messages, ) -> *mut paimon_error = paimon_commit_messages_merge; +const _: unsafe extern "C" fn( + *mut paimon_postpone_fixed_bucket_commit_messages, + *const paimon_postpone_fixed_bucket_commit_messages, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_commit_messages_merge; const _: unsafe extern "C" fn( *mut paimon_table_write, *mut c_void, *mut c_void, ) -> *mut paimon_error = paimon_table_write_write_arrow_batch; +const _: unsafe extern "C" fn( + *mut paimon_postpone_fixed_bucket_table_write, + *mut c_void, + *mut c_void, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_table_write_write_arrow_batch; const _: unsafe extern "C" fn( *const paimon_table_commit, *mut paimon_commit_messages, ) -> *mut paimon_error = paimon_table_commit_commit; +const _: unsafe extern "C" fn( + *const paimon_postpone_fixed_bucket_table_commit, + *mut paimon_postpone_fixed_bucket_commit_messages, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_table_commit_commit; const _: unsafe extern "C" fn( *const paimon_table_commit, *mut paimon_commit_messages, i64, ) -> *mut paimon_error = paimon_table_commit_commit_with_identifier; +const _: unsafe extern "C" fn( + *const paimon_postpone_fixed_bucket_table_commit, + *mut paimon_postpone_fixed_bucket_commit_messages, + i64, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_table_commit_commit_with_identifier; const _: unsafe extern "C" fn( *const paimon_table_commit, *mut paimon_commit_messages, i64, ) -> *mut paimon_error = paimon_table_commit_filter_and_commit_with_identifier; +const _: unsafe extern "C" fn( + *const paimon_postpone_fixed_bucket_table_commit, + *mut paimon_postpone_fixed_bucket_commit_messages, + i64, +) -> *mut paimon_error = + paimon_postpone_fixed_bucket_table_commit_filter_and_commit_with_identifier; const _: unsafe extern "C" fn( *const paimon_table_commit, *mut paimon_commit_messages, @@ -829,3 +1382,7 @@ const _: unsafe extern "C" fn( *const paimon_table_commit, *mut paimon_commit_messages, ) -> *mut paimon_error = paimon_table_commit_abort; +const _: unsafe extern "C" fn( + *const paimon_postpone_fixed_bucket_table_commit, + *mut paimon_postpone_fixed_bucket_commit_messages, +) -> *mut paimon_error = paimon_postpone_fixed_bucket_table_commit_abort; 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..55afbf643 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, @@ -39,6 +41,7 @@ pub struct CommitMessage { pub deleted_index_files: Vec, /// Files to be deleted (copy-on-write rewrite: old files replaced by new_files). pub deleted_files: Vec, + fixed_bucket_overwrite: bool, } impl CommitMessage { @@ -46,12 +49,22 @@ impl CommitMessage { Self { partition, bucket, + total_buckets: None, new_files, check_from_snapshot: None, new_changelog_files: Vec::new(), new_index_files: Vec::new(), deleted_index_files: Vec::new(), deleted_files: Vec::new(), + fixed_bucket_overwrite: false, } } + + pub(crate) fn mark_fixed_bucket_overwrite(&mut self) { + self.fixed_bucket_overwrite = true; + } + + pub(crate) fn is_fixed_bucket_overwrite(&self) -> bool { + self.fixed_bucket_overwrite + } } 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..488e3a3f1 --- /dev/null +++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs @@ -0,0 +1,1006 @@ +// 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, + overwrite: bool, + 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)?, + overwrite, + 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); + if self.overwrite { + message.mark_fixed_bucket_overwrite(); + } + 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("another commit wrote the same partition")); + } + + #[tokio::test] + async fn test_postpone_overwrite_rejects_concurrent_write_to_same_partition() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_overwrite_partition_conflict"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + let plan = make_unpartitioned_bucket_plan(&table, 2); + + let overwrite_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("overwrite-writer") + .unwrap() + .with_bucket_plan(plan.clone()) + .with_overwrite(); + let mut overwrite_write = overwrite_builder.new_write().unwrap(); + overwrite_write + .write_arrow_batch(&make_batch( + (0..32).collect(), + (0..32).map(|value| value * 10).collect(), + )) + .await + .unwrap(); + let overwrite_messages = overwrite_write + .prepare_commit() + .await + .unwrap() + .into_iter() + .filter(|message| message.bucket == 0) + .collect::>(); + assert_eq!(overwrite_messages.len(), 1); + + let append_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("append-writer") + .unwrap() + .with_bucket_plan(plan); + let mut append_write = append_builder.new_write().unwrap(); + append_write + .write_arrow_batch(&make_batch( + (100..132).collect(), + (100..132).map(|value| value * 10).collect(), + )) + .await + .unwrap(); + let append_messages = append_write + .prepare_commit() + .await + .unwrap() + .into_iter() + .filter(|message| message.bucket == 1) + .collect::>(); + assert_eq!(append_messages.len(), 1); + append_builder + .new_commit() + .commit(append_messages) + .await + .unwrap(); + + let error = overwrite_builder + .new_commit() + .commit(overwrite_messages) + .await + .unwrap_err(); + assert!(error + .to_string() + .contains("another commit wrote the same partition")); + } + + #[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("Cannot mix bucket=-2 files")); + assert_eq!(read_id_value_rows(&table).await, vec![(1, 10)]); + } + + #[tokio::test] + async fn test_postpone_fixed_append_rejects_existing_negative_bucket() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_rejects_existing_negative_bucket"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + 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, 2, 3, 4], vec![10, 20, 30, 40])) + .await + .unwrap(); + normal_builder + .new_commit() + .commit(normal_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let fixed_messages = + write_fixed_batch(&table, "fixed-append", 1, &make_batch(vec![1], vec![100])).await; + let error = TableCommit::new(table.clone(), "fixed-append".to_string()) + .commit(fixed_messages) + .await + .unwrap_err(); + assert!(error.to_string().contains("Cannot mix bucket=-2 files")); + + let overwrite_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("fixed-overwrite") + .unwrap() + .with_bucket_plan(make_unpartitioned_bucket_plan(&table, 1)) + .with_overwrite(); + let mut overwrite_write = overwrite_builder.new_write().unwrap(); + overwrite_write + .write_arrow_batch(&make_batch(vec![1], vec![100])) + .await + .unwrap(); + overwrite_builder + .new_commit() + .commit(overwrite_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + assert_eq!(read_id_value_rows(&table).await, vec![(1, 100)]); + } + + #[tokio::test] + async fn test_postpone_commit_rejects_writer_mode_mismatch() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_writer_mode_mismatch"; + 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 initial_builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_bucket_plan(plan.clone()); + let mut initial_write = initial_builder.new_write().unwrap(); + initial_write + .write_arrow_batch(&make_batch(vec![1, 2], vec![10, 20])) + .await + .unwrap(); + initial_builder + .new_commit() + .commit(initial_write.prepare_commit().await.unwrap()) + .await + .unwrap(); + + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_bucket_plan(plan); + let append_commit = builder.new_commit(); + let overwrite_builder = builder.with_overwrite(); + let mut overwrite_write = overwrite_builder.new_write().unwrap(); + overwrite_write + .write_arrow_batch(&make_batch(vec![1], vec![100])) + .await + .unwrap(); + let messages = overwrite_write.prepare_commit().await.unwrap(); + + let error = append_commit.commit(messages.clone()).await.unwrap_err(); + assert!(error.to_string().contains("submitted as append")); + assert_eq!(read_id_value_rows(&table).await, vec![(1, 10), (2, 20)]); + + overwrite_builder + .new_commit() + .commit(messages) + .await + .unwrap(); + assert_eq!(read_id_value_rows(&table).await, vec![(1, 100)]); + } + + #[tokio::test] + async fn test_postpone_commit_rejects_invalid_fixed_bucket_range() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_invalid_fixed_bucket_range"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + let messages = + write_fixed_batch(&table, "fixed-writer", 1, &make_batch(vec![1], vec![10])).await; + + for (bucket, total_buckets) in [(1, 1), (-1, 1), (0, 0)] { + let mut invalid = messages.clone(); + invalid[0].bucket = bucket; + invalid[0].total_buckets = Some(total_buckets); + let error = TableCommit::new(table.clone(), "fixed-writer".to_string()) + .commit(invalid) + .await + .unwrap_err(); + assert!(error.to_string().contains("Invalid fixed bucket")); + } + } + + #[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("Cannot mix 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..4330aeb9b 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,53 @@ 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 { + let Some(total_buckets) = message.total_buckets else { + continue; + }; + if total_buckets <= 0 || message.bucket < 0 || message.bucket >= total_buckets { + return Err(crate::Error::DataInvalid { + message: format!( + "Invalid fixed bucket {} for total_buckets={total_buckets}; expected a bucket in [0, {total_buckets})", + message.bucket + ), + source: None, + }); + } + if 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(()) +} + +fn validate_fixed_bucket_commit_mode(messages: &[CommitMessage], overwrite: bool) -> Result<()> { + if messages.iter().any(|message| { + message.total_buckets.is_some() && message.is_fixed_bucket_overwrite() != overwrite + }) { + return Err(crate::Error::DataInvalid { + message: format!( + "Commit messages were prepared for {} but submitted as {}", + if overwrite { "append" } else { "overwrite" }, + if overwrite { "overwrite" } else { "append" }, + ), + source: None, + }); + } + Ok(()) +} + /// Table commit logic for Paimon write operations. /// /// Provides atomic commit functionality including append, overwrite and truncate @@ -156,6 +204,8 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_fixed_bucket_commit_mode(&commit_messages, false)?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() { return Ok(()); @@ -199,6 +249,8 @@ impl TableCommit { commit_identifier: i64, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_fixed_bucket_commit_mode(&commit_messages, false)?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() { return Ok(()); @@ -270,6 +322,8 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + validate_fixed_bucket_commit_mode(&commit_messages, true)?; + validate_bucket_ownership(&commit_messages)?; if commit_messages.is_empty() && static_partitions.is_none() { return Ok(()); @@ -298,11 +352,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 +588,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 +661,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 +1299,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 +1367,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 +1386,7 @@ impl TableCommit { retry_state, &entries, &CommitKind::OVERWRITE, - None, + check_from_snapshot, ) .await?; @@ -1854,6 +1928,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)?; + // Check the final layout so overwrite migration and rescaling remain 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_postpone_bucket_mixing(&merged_entries)?; + self.check_fixed_bucket_ownership_conflicts( + latest_snapshot, + delta_entries, + commit_kind, + check_from_snapshot, + ) + .await?; + if !self.data_evolution_enabled { return Ok(()); } @@ -1861,14 +1951,138 @@ 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], + commit_kind: &CommitKind, + 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 commit_kind == &CommitKind::OVERWRITE + && owned_partitions.contains(entry.partition()) + { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket overwrite conflict: another commit wrote the same partition after snapshot {check_from_snapshot}" + ), + source: None, + }); + } + 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, active_entries: &[ManifestEntry]) -> Result<()> { + if self.total_buckets != POSTPONE_BUCKET { + return Ok(()); + } + let fixed_partitions = active_entries + .iter() + .filter(|entry| { + *entry.kind() == FileKind::Add && entry.bucket() >= 0 && entry.total_buckets() > 0 + }) + .map(|entry| entry.partition()) + .collect::>(); + if active_entries.iter().any(|entry| { + *entry.kind() == FileKind::Add + && entry.bucket() == POSTPONE_BUCKET + && fixed_partitions.contains(entry.partition()) + }) { + return Err(crate::Error::DataInvalid { + message: "Cannot mix bucket=-2 files and fixed buckets in the same partition; use a complete overwrite to migrate the partition" + .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 +2771,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 +2818,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 +2828,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 +2848,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 +2917,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 +3120,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 +3305,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, @@ -3185,6 +3403,38 @@ mod tests { assert_eq!(entries[0].file().file_name, "data-0.parquet"); } + #[tokio::test] + async fn test_commit_rejects_writer_mode_mismatch() { + let file_io = test_file_io(); + let table_path = "memory:/test_commit_writer_mode_mismatch"; + setup_dirs(&file_io, table_path).await; + let commit = setup_commit(&file_io, table_path); + + let mut append = CommitMessage::new(vec![], 0, vec![test_data_file("append.parquet", 1)]); + append.total_buckets = Some(1); + let error = commit.overwrite(vec![append], None).await.unwrap_err(); + assert!(error.to_string().contains("submitted as overwrite")); + + let mut overwrite = + CommitMessage::new(vec![], 0, vec![test_data_file("overwrite.parquet", 1)]); + overwrite.total_buckets = Some(1); + overwrite.mark_fixed_bucket_overwrite(); + let error = commit.commit(vec![overwrite]).await.unwrap_err(); + assert!(error.to_string().contains("submitted as append")); + + assert!(latest_snapshot(&file_io, table_path).await.is_none()); + } + + #[test] + fn test_fixed_bucket_commit_mode_preserves_non_fixed_messages() { + let append = CommitMessage::new(vec![], 0, vec![]); + assert!(validate_fixed_bucket_commit_mode(&[append], true).is_ok()); + + let mut overwrite = CommitMessage::new(vec![], 0, vec![]); + overwrite.mark_fixed_bucket_overwrite(); + assert!(validate_fixed_bucket_commit_mode(&[overwrite], false).is_ok()); + } + #[tokio::test] async fn test_commit_with_identifier_writes_snapshot_identifier() { let file_io = test_file_io(); @@ -4133,6 +4383,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..a792e783f 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -153,8 +153,7 @@ impl<'a> PaimonWriteBuilder<'a> { } /// Mark writers created by this builder as overwrite-aware. - /// - /// The commit kind remains explicit at the commit call site. + /// Their messages must be submitted through [`TableCommit::overwrite`]. pub fn with_overwrite(mut self) -> Self { self.overwrite = true; self @@ -179,25 +178,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 +204,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 +406,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 @@ -506,7 +503,7 @@ mod tests { "Overwrite-aware writer must not produce input changelog files" ); - wb.new_commit().commit(messages).await.unwrap(); + wb.new_commit().overwrite(messages, None).await.unwrap(); let snapshot_manager = crate::table::SnapshotManager::new(file_io.clone(), table_path.to_string()); @@ -515,7 +512,7 @@ mod tests { .await .unwrap() .unwrap(); - assert_eq!(snapshot.commit_kind(), &CommitKind::APPEND); + assert_eq!(snapshot.commit_kind(), &CommitKind::OVERWRITE); } #[test]