diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs index d29c7eb85..15cb661a7 100644 --- a/bindings/c/src/tests.rs +++ b/bindings/c/src/tests.rs @@ -73,6 +73,17 @@ fn not_null_table_schema() -> TableSchema { TableSchema::new(0, &schema) } +fn postpone_table_schema() -> TableSchema { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .build() + .unwrap(); + TableSchema::new(0, &schema) +} + 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 })) @@ -950,6 +961,47 @@ fn test_write_new_builder_and_free() { } } +#[test] +fn test_postpone_fixed_bucket_builder_is_explicit() { + let path = "memory:/test_explicit_postpone_fixed_bucket_builder"; + let file_io = memory_file_io(); + setup_table_dirs(&file_io, path); + let table = Table::new( + file_io, + Identifier::new("default", "test"), + path.to_string(), + postpone_table_schema(), + None, + ); + let handle = unsafe { wrap_table(table) }; + + unsafe { + let normal = paimon_table_new_write_builder(handle); + assert!(normal.error.is_null()); + let normal_state = &*((*normal.write_builder).inner as *const WriteBuilderState); + assert!(!normal_state.postpone_fixed_bucket); + paimon_write_builder_free(normal.write_builder); + + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder(handle); + assert!(fixed.error.is_null()); + let fixed_state = &*((*fixed.write_builder).inner as *const WriteBuilderState); + assert!(fixed_state.postpone_fixed_bucket); + paimon_write_builder_free(fixed.write_builder); + + let commit_user = CString::new("fixed-user").unwrap(); + let fixed = paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + handle, + commit_user.as_ptr(), + ); + assert!(fixed.error.is_null()); + let fixed_state = &*((*fixed.write_builder).inner as *const WriteBuilderState); + assert!(fixed_state.postpone_fixed_bucket); + assert_eq!(fixed_state.commit_user, "fixed-user"); + paimon_write_builder_free(fixed.write_builder); + unwrap_table(handle); + } +} + #[test] fn test_write_commit_read_roundtrip() { let path = "memory:/test_write_roundtrip"; @@ -1721,6 +1773,11 @@ fn test_null_pointer_handling() { assert!(result.write_builder.is_null()); paimon_error_free(result.error); + let result = paimon_table_new_postpone_fixed_bucket_write_builder(ptr::null()); + assert!(!result.error.is_null()); + assert!(result.write_builder.is_null()); + paimon_error_free(result.error); + let result = paimon_write_builder_new_write(ptr::null()); assert!(!result.error.is_null()); assert!(result.write.is_null()); diff --git a/bindings/c/src/types.rs b/bindings/c/src/types.rs index 1dfd9c3f2..d792a8e05 100644 --- a/bindings/c/src/types.rs +++ b/bindings/c/src/types.rs @@ -210,6 +210,7 @@ pub(crate) struct WriteBuilderState { pub table: Table, pub commit_user: String, pub overwrite: bool, + pub postpone_fixed_bucket: bool, } pub(crate) struct TableWriteState { diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs index 8e9637997..5a7518d3c 100644 --- a/bindings/c/src/write.rs +++ b/bindings/c/src/write.rs @@ -37,6 +37,7 @@ use crate::types::*; unsafe fn new_write_builder( table: *const paimon_table, commit_user: Option, + postpone_fixed_bucket: bool, ) -> paimon_result_write_builder { if let Err(e) = check_non_null(table, "table") { return paimon_result_write_builder { @@ -45,7 +46,19 @@ unsafe fn new_write_builder( }; } let table_ref = &*((*table).inner as *const Table); - let builder = table_ref.new_write_builder(); + let builder = if postpone_fixed_bucket { + match table_ref.new_postpone_fixed_bucket_write_builder() { + Ok(builder) => builder, + Err(e) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error: paimon_error::from_paimon(e), + } + } + } + } else { + table_ref.new_write_builder() + }; let commit_user = match commit_user { Some(commit_user) => match builder.with_commit_user(commit_user) { Ok(builder) => builder.commit_user().to_string(), @@ -62,6 +75,7 @@ unsafe fn new_write_builder( table: table_ref.clone(), commit_user, overwrite: false, + postpone_fixed_bucket, }; let inner = Box::into_raw(Box::new(state)) as *mut c_void; paimon_result_write_builder { @@ -82,7 +96,22 @@ unsafe fn new_write_builder( pub unsafe extern "C" fn paimon_table_new_write_builder( table: *const paimon_table, ) -> paimon_result_write_builder { - new_write_builder(table, None) + new_write_builder(table, None, false) +} + +/// Create an explicit one-shot fixed-bucket WriteBuilder for a postpone table. +/// +/// Normal write builders retain `bucket = -2` postpone semantics. Use this +/// entry point when the batch must be routed to immediately visible real +/// buckets. +/// +/// # Safety +/// `table` must be a valid table pointer, or null (returns error). +#[no_mangle] +pub unsafe extern "C" fn paimon_table_new_postpone_fixed_bucket_write_builder( + table: *const paimon_table, +) -> paimon_result_write_builder { + new_write_builder(table, None, true) } /// Create a WriteBuilder with a caller-provided stable commit identity. @@ -107,7 +136,29 @@ pub unsafe extern "C" fn paimon_table_new_write_builder_with_commit_user( } } }; - new_write_builder(table, Some(commit_user)) + new_write_builder(table, Some(commit_user), false) +} + +/// Create an explicit postpone fixed-bucket WriteBuilder with a stable commit identity. +/// +/// # Safety +/// `table` must be a valid table pointer. `commit_user` must be a valid UTF-8 +/// C string and a safe file-name segment. +#[no_mangle] +pub unsafe extern "C" fn paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user( + table: *const paimon_table, + commit_user: *const c_char, +) -> paimon_result_write_builder { + let commit_user = match validate_cstr(commit_user, "commit_user") { + Ok(commit_user) => commit_user, + Err(error) => { + return paimon_result_write_builder { + write_builder: ptr::null_mut(), + error, + } + } + }; + new_write_builder(table, Some(commit_user), true) } /// Free a paimon_write_builder. @@ -245,19 +296,21 @@ pub unsafe extern "C" fn paimon_write_builder_new_write( } let state = &*((*wb).inner as *const WriteBuilderState); - let mut builder = match state - .table - .new_write_builder() - .with_commit_user(state.commit_user.clone()) - { - Ok(b) => b, - Err(e) => { - return paimon_result_table_write { - write: ptr::null_mut(), - error: paimon_error::from_paimon(e), - } - } + let builder = if state.postpone_fixed_bucket { + state.table.new_postpone_fixed_bucket_write_builder() + } else { + Ok(state.table.new_write_builder()) }; + let mut builder = + match builder.and_then(|builder| builder.with_commit_user(state.commit_user.clone())) { + Ok(b) => b, + Err(e) => { + return paimon_result_table_write { + write: ptr::null_mut(), + error: paimon_error::from_paimon(e), + } + } + }; if state.overwrite { builder = builder.with_overwrite(); @@ -358,8 +411,9 @@ pub unsafe extern "C" fn paimon_table_write_write_arrow_batch( /// Close file writers and produce CommitMessages. /// /// Consumes the open file writers (they are flushed and closed). After this -/// call, the TableWrite can be reused — `write_arrow_batch` may be called -/// again to start a new round of writes. +/// call, the TableWrite can normally be reused — `write_arrow_batch` may be +/// called again to start a new round of writes. Fixed-bucket postpone batch +/// writers are one-shot; create a new TableWrite for the next batch. /// /// The returned `paimon_commit_messages` must be passed to a /// `paimon_table_commit_*` function and then freed with @@ -781,8 +835,12 @@ pub unsafe extern "C" fn paimon_table_commit_abort( const _: unsafe extern "C" fn(*const paimon_table) -> paimon_result_write_builder = paimon_table_new_write_builder; +const _: unsafe extern "C" fn(*const paimon_table) -> paimon_result_write_builder = + paimon_table_new_postpone_fixed_bucket_write_builder; const _: unsafe extern "C" fn(*const paimon_table, *const c_char) -> paimon_result_write_builder = paimon_table_new_write_builder_with_commit_user; +const _: unsafe extern "C" fn(*const paimon_table, *const c_char) -> paimon_result_write_builder = + paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user; const _: unsafe extern "C" fn(*const paimon_write_builder) -> paimon_result_table_write = paimon_write_builder_new_write; const _: unsafe extern "C" fn(*const paimon_write_builder) -> paimon_result_table_commit = diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index e26c6bc34..a88e01154 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -121,6 +121,13 @@ const DEFAULT_PARQUET_ROW_GROUP_PARALLELISM: usize = 8; const DEFAULT_PARQUET_ROW_GROUP_MAX_INFLIGHT_BYTES: i64 = 256 * 1024 * 1024; const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str = "dynamic-bucket.target-row-num"; const DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM: i64 = 200_000; +const POSTPONE_BATCH_WRITE_FIXED_BUCKET_OPTION: &str = "postpone.batch-write-fixed-bucket"; +const POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION: &str = + "postpone.batch-write-fixed-bucket.max-parallelism"; +const POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION: &str = "postpone.target-row-num-per-bucket"; +const POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION: &str = "postpone.target-size-per-bucket"; +const DEFAULT_POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM: i32 = 2048; +const DEFAULT_POSTPONE_TARGET_SIZE_PER_BUCKET: i64 = 1024 * 1024 * 1024; const DEFAULT_GLOBAL_INDEX_ROW_COUNT_PER_SHARD: i64 = 100_000; const DEFAULT_GLOBAL_INDEX_THREAD_NUM: i64 = 32; const DEFAULT_GLOBAL_INDEX_FALLBACK_SCAN_MAX_SIZE: i64 = 256 * 1024 * 1024; @@ -1148,6 +1155,75 @@ impl<'a> CoreOptions<'a> { .unwrap_or(DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM) } + /// Whether high-level batch integrations should select the explicit + /// postpone fixed-bucket writer. Direct `new_write_builder` calls retain + /// normal postpone semantics regardless of this option. + pub fn postpone_batch_write_fixed_bucket(&self) -> bool { + self.options + .get(POSTPONE_BATCH_WRITE_FIXED_BUCKET_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(true) + } + + pub fn postpone_batch_write_fixed_bucket_max_parallelism(&self) -> crate::Result { + let value = self + .options + .get(POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION) + .map(|value| value.parse::()) + .transpose() + .map_err(|error| crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION}' must be a positive integer" + ), + source: Some(Box::new(error)), + })? + .unwrap_or(DEFAULT_POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM); + if value <= 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM_OPTION}' must be positive, got: {value}" + ), + source: None, + }); + } + Ok(value) + } + + pub fn postpone_target_row_num_per_bucket(&self) -> crate::Result> { + let value = self.parse_i64_option(POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION)?; + if value.is_some_and(|value| value <= 0) { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_ROW_NUM_PER_BUCKET_OPTION}' must be positive, got: {}", + value.unwrap() + ), + source: None, + }); + } + Ok(value) + } + + pub fn postpone_target_size_per_bucket(&self) -> crate::Result { + let value = match self.options.get(POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION) { + Some(raw) => parse_memory_size(raw).ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION}' must be a valid positive memory size, got: {raw}" + ), + source: None, + })?, + None => DEFAULT_POSTPONE_TARGET_SIZE_PER_BUCKET, + }; + if value <= 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{POSTPONE_TARGET_SIZE_PER_BUCKET_OPTION}' must be positive, got: {value}" + ), + source: None, + }); + } + Ok(value) + } + /// When true, blob field reads return serialized BlobDescriptor bytes /// instead of actual blob bytes. Default is false. pub fn blob_as_descriptor(&self) -> bool { diff --git a/crates/paimon/src/table/commit_message.rs b/crates/paimon/src/table/commit_message.rs index 5a91e36b5..7332c2c47 100644 --- a/crates/paimon/src/table/commit_message.rs +++ b/crates/paimon/src/table/commit_message.rs @@ -27,6 +27,9 @@ pub struct CommitMessage { pub partition: Vec, /// Bucket id. pub bucket: i32, + /// Per-partition bucket count. Set by fixed-bucket postpone batch writes; + /// ordinary writes use the table-level bucket option. + pub total_buckets: Option, /// New data files to be added. pub new_files: Vec, /// Snapshot id from which row-id/column conflicts should be checked. @@ -46,6 +49,7 @@ impl CommitMessage { Self { partition, bucket, + total_buckets: None, new_files, check_from_snapshot: None, new_changelog_files: Vec::new(), diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 4078e3a23..db064488d 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -72,6 +72,7 @@ mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; mod pk_vector_position_read; mod pk_vector_scan; +mod postpone_bucket; mod postpone_file_writer; mod prepared_files; mod read_builder; @@ -359,6 +360,14 @@ impl Table { WriteBuilder::new(self) } + /// Create an explicit one-shot fixed-bucket writer for a postpone table. + /// + /// Normal write builders retain postpone-bucket semantics. Callers which + /// want immediately visible real buckets must opt in through this builder. + pub fn new_postpone_fixed_bucket_write_builder(&self) -> Result> { + WriteBuilder::new_postpone_fixed_bucket(self) + } + /// Create a copy of this table with extra options merged into the schema. /// /// This never switches the schema version; it corresponds to Java diff --git a/crates/paimon/src/table/postpone_bucket.rs b/crates/paimon/src/table/postpone_bucket.rs new file mode 100644 index 000000000..3669354f0 --- /dev/null +++ b/crates/paimon/src/table/postpone_bucket.rs @@ -0,0 +1,436 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Helpers shared by fixed-bucket postpone batch planning. + +use crate::spec::{BinaryRow, DataField, DataType, IntType}; +use arrow_array::{ + Array, ArrayRef, BinaryArray, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, + MapArray, RecordBatch, StringArray, StringViewArray, StructArray, +}; + +/// Sum the Java `BinaryRow.getSizeInBytes()` value for every row in a batch. +/// +/// Arrow allocation size is not a stable planning metric: offsets, views, and +/// slicing can make the same logical rows occupy different Arrow memory. This +/// estimator follows Paimon's internal BinaryRow/BinaryArray layouts without +/// materializing a second copy of the batch. +pub(crate) fn binary_row_batch_size( + batch: &RecordBatch, + fields: &[DataField], +) -> crate::Result { + if batch.num_columns() != fields.len() { + return Err(crate::Error::DataInvalid { + message: format!( + "BinaryRow size planning expected {} columns, got {}", + fields.len(), + batch.num_columns() + ), + source: None, + }); + } + + let mut total = 0_u128; + for row in 0..batch.num_rows() { + total = total.saturating_add(row_size(batch.columns(), row, fields)?); + } + Ok(total.min(i64::MAX as u128) as i64) +} + +fn row_size(arrays: &[ArrayRef], row: usize, fields: &[DataField]) -> crate::Result { + let mut size = BinaryRow::cal_fix_part_size_in_bytes(fields.len() as i32) as u128; + for (array, field) in arrays.iter().zip(fields) { + size = size.saturating_add(variable_size(array, row, field.data_type())?); + } + Ok(size) +} + +fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + // Java's BinaryRowWriter reserves these fixed variable regions even for a + // null top-level field. + match data_type { + DataType::Decimal(decimal) if decimal.precision() > 18 => return Ok(16), + DataType::Timestamp(timestamp) if timestamp.precision() > 3 => return Ok(8), + DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() > 3 => return Ok(8), + _ => {} + } + + if array.is_null(row) { + return Ok(0); + } + + match data_type { + DataType::Char(_) | DataType::VarChar(_) => { + let len = string_len(array, row, data_type)?; + Ok(binary_size(len)) + } + DataType::Binary(_) | DataType::VarBinary(_) | DataType::Blob(_) => { + let len = binary_len(array, row, data_type)?; + Ok(binary_size(len)) + } + DataType::Variant(_) => variant_size(array, row), + DataType::Array(array_type) => { + if let Some(array) = array.as_any().downcast_ref::() { + let offsets = array.value_offsets(); + Ok(round_to_word(binary_array_size( + array.values(), + offsets[row] as usize, + offsets[row + 1] as usize, + array_type.element_type(), + )?)) + } else if let Some(array) = array.as_any().downcast_ref::() { + let offsets = array.value_offsets(); + Ok(round_to_word(binary_array_size( + array.values(), + offsets[row] as usize, + offsets[row + 1] as usize, + array_type.element_type(), + )?)) + } else { + Err(type_mismatch("ListArray", data_type)) + } + } + DataType::Map(map_type) => { + let map = downcast::(array, "MapArray", data_type)?; + let offsets = map.value_offsets(); + let entries = map.entries(); + map_size( + entries, + offsets[row] as usize, + offsets[row + 1] as usize, + map_type.key_type(), + map_type.value_type(), + ) + } + DataType::Multiset(multiset_type) => { + let map = downcast::(array, "MapArray", data_type)?; + let offsets = map.value_offsets(); + let entries = map.entries(); + map_size( + entries, + offsets[row] as usize, + offsets[row + 1] as usize, + multiset_type.element_type(), + &DataType::Int(IntType::new()), + ) + } + DataType::Row(row_type) => { + let struct_array = downcast::(array, "StructArray", data_type)?; + Ok(round_to_word(row_size( + struct_array.columns(), + row, + row_type.fields(), + )?)) + } + DataType::Vector(vector_type) => { + let element_bytes = u128::from(vector_type.length()) + .saturating_mul(primitive_width(vector_type.element_type())?); + Ok(round_to_word(4_u128.saturating_add(element_bytes))) + } + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Float(_) + | DataType::Double(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) => Ok(0), + } +} + +fn binary_array_size( + values: &ArrayRef, + start: usize, + end: usize, + element_type: &DataType, +) -> crate::Result { + if end < start || end > values.len() { + return Err(crate::Error::DataInvalid { + message: format!( + "Invalid nested array range [{start}, {end}) for {} values", + values.len() + ), + source: None, + }); + } + let count = end - start; + let header = 4_u128.saturating_add((count as u128).div_ceil(32).saturating_mul(4)); + let fixed = (count as u128).saturating_mul(fixed_width(element_type)?); + let mut size = round_to_word(header.saturating_add(fixed)); + for row in start..end { + if !values.is_null(row) { + size = size.saturating_add(variable_size(values, row, element_type)?); + } + } + Ok(size) +} + +fn map_size( + entries: &StructArray, + start: usize, + end: usize, + key_type: &DataType, + value_type: &DataType, +) -> crate::Result { + if entries.num_columns() != 2 { + return Err(crate::Error::DataInvalid { + message: format!( + "BinaryMap size planning expected 2 entry columns, got {}", + entries.num_columns() + ), + source: None, + }); + } + let keys = binary_array_size(entries.column(0), start, end, key_type)?; + let values = binary_array_size(entries.column(1), start, end, value_type)?; + Ok(round_to_word( + 4_u128.saturating_add(keys).saturating_add(values), + )) +} + +fn variant_size(array: &ArrayRef, row: usize) -> crate::Result { + let variant = downcast::( + array, + "StructArray", + &DataType::Variant(crate::spec::VariantType::new()), + )?; + if variant.num_columns() != 2 { + return Err(crate::Error::DataInvalid { + message: format!( + "Variant size planning expected 2 child columns, got {}", + variant.num_columns() + ), + source: None, + }); + } + let value_len = binary_len( + variant.column(0), + row, + &DataType::Variant(crate::spec::VariantType::new()), + )?; + let metadata_len = binary_len( + variant.column(1), + row, + &DataType::Variant(crate::spec::VariantType::new()), + )?; + Ok(round_to_word( + 4_u128 + .saturating_add(value_len as u128) + .saturating_add(metadata_len as u128), + )) +} + +fn string_len(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else { + Err(type_mismatch("StringArray", data_type)) + } +} + +fn binary_len(array: &ArrayRef, row: usize, data_type: &DataType) -> crate::Result { + if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else if let Some(array) = array.as_any().downcast_ref::() { + Ok(array.value(row).len()) + } else { + Err(type_mismatch("BinaryArray", data_type)) + } +} + +fn fixed_width(data_type: &DataType) -> crate::Result { + Ok(match data_type { + DataType::Boolean(_) | DataType::TinyInt(_) => 1, + DataType::SmallInt(_) => 2, + DataType::Int(_) | DataType::Float(_) | DataType::Date(_) | DataType::Time(_) => 4, + DataType::BigInt(_) + | DataType::Double(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Binary(_) + | DataType::VarBinary(_) + | DataType::Blob(_) + | DataType::Variant(_) + | DataType::Timestamp(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Decimal(_) + | DataType::Array(_) + | DataType::Map(_) + | DataType::Multiset(_) + | DataType::Row(_) + | DataType::Vector(_) => 8, + }) +} + +fn primitive_width(data_type: &DataType) -> crate::Result { + match data_type { + DataType::Boolean(_) | DataType::TinyInt(_) => Ok(1), + DataType::SmallInt(_) => Ok(2), + DataType::Int(_) | DataType::Float(_) => Ok(4), + DataType::BigInt(_) | DataType::Double(_) => Ok(8), + other => Err(crate::Error::DataInvalid { + message: format!("Unsupported vector element type for size planning: {other:?}"), + source: None, + }), + } +} + +fn binary_size(len: usize) -> u128 { + if len <= 7 { + 0 + } else { + round_to_word(len as u128) + } +} + +fn round_to_word(size: u128) -> u128 { + size.saturating_add(7) / 8 * 8 +} + +fn downcast<'a, T: 'static>( + array: &'a ArrayRef, + expected: &str, + data_type: &DataType, +) -> crate::Result<&'a T> { + array + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch(expected, data_type)) +} + +fn type_mismatch(expected: &str, data_type: &DataType) -> crate::Error { + crate::Error::DataInvalid { + message: format!("BinaryRow size planning expected {expected} for {data_type:?}"), + source: None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{ArrayType, IntType, LocalZonedTimestampType, TimestampType, VarCharType}; + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, ListArray, TimestampMicrosecondArray}; + use arrow_schema::{DataType as ArrowDataType, Field, Schema, TimeUnit}; + use std::sync::Arc; + + #[test] + fn test_java_binary_row_size_differs_from_arrow_buffers() { + let row_count = 1_000; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("left", ArrowDataType::Utf8, false), + Field::new("right", ArrowDataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..row_count)), + Arc::new(StringArray::from(vec!["a"; row_count as usize])), + Arc::new(StringArray::from(vec!["b"; row_count as usize])), + ], + ) + .unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "left".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + DataField::new( + 2, + "right".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 32_000); + assert_ne!(batch.get_array_memory_size() as i64, 32_000); + } + + #[test] + fn test_nested_and_local_zoned_timestamp_size() { + let array = + ListArray::from_iter_primitive::(vec![Some(vec![Some(1), Some(2)])]); + let timestamp = + TimestampMicrosecondArray::from(vec![Some(1_234_567_i64)]).with_timezone("UTC"); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("items", array.data_type().clone(), true), + Field::new( + "event_time", + ArrowDataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(array), + Arc::new(timestamp), + ], + ) + .unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "items".to_string(), + DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), + ), + DataField::new( + 2, + "event_time".to_string(), + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(6).unwrap()), + ), + ]; + + // 32-byte row fixed part + 16-byte BinaryArray + 8-byte non-compact timestamp. + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 56); + } + + #[test] + fn test_non_compact_null_timestamp_reserves_space() { + let schema = Arc::new(Schema::new(vec![Field::new( + "event_time", + ArrowDataType::Timestamp(TimeUnit::Microsecond, None), + true, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(TimestampMicrosecondArray::from(vec![None]))], + ) + .unwrap(); + let fields = vec![DataField::new( + 0, + "event_time".to_string(), + DataType::Timestamp(TimestampType::new(6).unwrap()), + )]; + + assert_eq!(binary_row_batch_size(&batch, &fields).unwrap(), 24); + } +} diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 809e45d9f..6bf7b917e 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -1240,7 +1240,11 @@ impl TableCommit { } else { CommitKind::APPEND }; - let detect_conflicts = has_delete || check_from_snapshot.is_some(); + let has_partition_bucket_counts = entries + .iter() + .any(|entry| entry.total_buckets() != self.total_buckets); + let detect_conflicts = + has_delete || check_from_snapshot.is_some() || has_partition_bucket_counts; let base_data_files = if detect_conflicts { self.check_deletion_vector_index_only_conflict( latest_snapshot.as_ref(), @@ -1853,6 +1857,7 @@ impl TableCommit { check_from_snapshot: Option, ) -> Result<()> { self.check_delete_entries_against_base(base_entries, delta_entries)?; + self.check_total_bucket_conflicts(base_entries, delta_entries)?; if !self.data_evolution_enabled { return Ok(()); @@ -1869,6 +1874,32 @@ impl TableCommit { .await } + fn check_total_bucket_conflicts( + &self, + base_entries: &[ManifestEntry], + delta_entries: &[ManifestEntry], + ) -> Result<()> { + let mut bucket_counts: HashMap, i32> = HashMap::new(); + for entry in base_entries.iter().chain(delta_entries) { + if entry.bucket() < 0 || entry.total_buckets() <= 0 { + continue; + } + let partition = entry.partition().to_vec(); + if let Some(previous) = bucket_counts.insert(partition, entry.total_buckets()) { + if previous != entry.total_buckets() { + return Err(crate::Error::DataInvalid { + message: format!( + "Postpone fixed-bucket conflict: one partition uses different total bucket counts {previous} and {}", + entry.total_buckets() + ), + source: None, + }); + } + } + } + Ok(()) + } + fn check_deletion_vector_index_only_conflict( &self, latest_snapshot: Option<&Snapshot>, @@ -2557,6 +2588,10 @@ impl TableCommit { stats.file_size_in_bytes += sign * file.file_size; stats.file_count += sign; stats.last_file_creation_time = stats.last_file_creation_time.max(file_creation_time); + // Overwrite entries are ordered DELETE-old then ADD-new. Match + // Java PartitionEntry.merge by retaining the replacement entry's + // bucket count instead of the first value seen for the partition. + stats.total_buckets = entry.total_buckets(); } Ok(stats_map.into_values().collect()) @@ -2602,7 +2637,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 +2647,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 +2667,7 @@ impl TableCommit { FileKind::Add, msg.partition.clone(), msg.bucket, - self.total_buckets, + msg.total_buckets.unwrap_or(self.total_buckets), file.clone(), 0, ) @@ -2903,7 +2938,7 @@ mod tests { use crate::spec::stats::BinaryTableStats; use crate::spec::{ BinaryRowBuilder, DataFileMeta, DeletionVectorMeta, GlobalIndexMeta, IndexFileMeta, - ManifestList, TableSchema, + ManifestList, TableSchema, POSTPONE_BUCKET, }; use chrono::{DateTime, Utc}; @@ -4135,6 +4170,45 @@ mod tests { assert_eq!(snapshot.total_record_count(), Some(250)); } + #[test] + fn test_partition_statistics_keep_replacement_bucket_count() { + let file_io = test_file_io(); + let commit = setup_partitioned_commit( + &file_io, + "memory:/test_partition_statistics_replacement_bucket_count", + ); + let partition = partition_bytes("a"); + + for (old_buckets, new_buckets) in [(-2, 4), (4, 8)] { + let entries = vec![ + ManifestEntry::new( + FileKind::Delete, + partition.clone(), + if old_buckets == POSTPONE_BUCKET { + POSTPONE_BUCKET + } else { + 0 + }, + old_buckets, + test_data_file("old.parquet", 100), + 2, + ), + ManifestEntry::new( + FileKind::Add, + partition.clone(), + 0, + new_buckets, + test_data_file("new.parquet", 50), + 2, + ), + ]; + + let statistics = commit.generate_partition_statistics(&entries).unwrap(); + assert_eq!(statistics.len(), 1); + assert_eq!(statistics[0].total_buckets, new_buckets); + } + } + #[tokio::test] async fn test_overwrite_cache_reuses_when_append_misses_target_partition() { let file_io = test_file_io(); diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 490de6357..432d75aa9 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -21,23 +21,24 @@ //! and [pypaimon FileStoreWrite](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/write/file_store_write.py) use crate::arrow::build_target_arrow_schema; -use crate::spec::PartitionComputer; +use crate::spec::{batch_to_serialized_bytes, PartitionComputer}; use crate::spec::{ - first_row_supports_changelog_producer, BinaryRow, ChangelogProducer, CoreOptions, DataField, - DataType, MergeEngine, RowKindFilter, EMPTY_SERIALIZED_ROW, POSTPONE_BUCKET, - VALUE_KIND_FIELD_NAME, + first_row_supports_changelog_producer, BinaryRow, BucketFunctionType, ChangelogProducer, + CoreOptions, DataField, DataType, MergeEngine, RowKindFilter, EMPTY_SERIALIZED_ROW, + POSTPONE_BUCKET, VALUE_KIND_FIELD_NAME, }; use crate::table::bucket_assigner::{BucketAssignerEnum, PartitionBucketKey}; use crate::table::bucket_assigner_constant::ConstantBucketAssigner; use crate::table::bucket_assigner_cross::CrossPartitionAssigner; use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner; use crate::table::bucket_assigner_fixed::FixedBucketAssigner; -use crate::table::bucket_function::validate_bucket_function; +use crate::table::bucket_function::{batch_bucket_ids, validate_bucket_function}; use crate::table::commit_message::CommitMessage; use crate::table::data_file_writer::DataFileWriter; use crate::table::dedicated_format_file_writer::AppendDedicatedFormatFileWriter; use crate::table::kv_file_writer::{KeyValueFileWriter, KeyValueWriteConfig}; use crate::table::partition_filter::PartitionFilter; +use crate::table::postpone_bucket::binary_row_batch_size; use crate::table::postpone_file_writer::{PostponeFileWriter, PostponeWriteConfig}; use crate::table::prepared_files::PreparedFiles; use crate::table::row_kind_generator::RowKindGenerator; @@ -55,6 +56,28 @@ enum FileWriter { Postpone(PostponeFileWriter), } +/// Planning state for batch writes to postpone-bucket tables. Partitions with +/// an existing real-bucket count are written incrementally, while only new +/// partitions are buffered until their bucket count can be inferred. +struct PostponeFixedBucketState { + partition_field_indices: Vec, + bucket_key_indices: Vec, + bucket_function_type: BucketFunctionType, + max_parallelism: i32, + target_rows_per_bucket: Option, + target_size_per_bucket: i64, + metadata_loaded: bool, + known_bucket_counts: HashMap, i32>, + postpone_row_counts: HashMap, i64>, + buffered_batches: HashMap, Vec>, + /// Bucket counts used by the current prepare-commit round. + bucket_counts: HashMap, i32>, + /// Fixed-bucket planning covers one complete batch. Reusing the writer + /// before its prepared messages are committed could re-plan new partitions + /// against a stale snapshot, so this mode is deliberately one-shot. + prepare_started: bool, +} + impl FileWriter { async fn write(&mut self, batch: &RecordBatch) -> Result<()> { match self { @@ -126,10 +149,26 @@ pub struct TableWrite { has_dedicated_vector_fields: bool, row_kind_generator: Option, row_kind_filter: Option, + postpone_fixed_bucket: Option, } impl TableWrite { pub(crate) fn new(table: &Table, commit_user: String) -> crate::Result { + Self::new_inner(table, commit_user, false) + } + + pub(crate) fn new_postpone_fixed_bucket( + table: &Table, + commit_user: String, + ) -> crate::Result { + Self::new_inner(table, commit_user, true) + } + + fn new_inner( + table: &Table, + commit_user: String, + use_postpone_fixed_bucket: bool, + ) -> crate::Result { let is_overwrite = false; let schema = table.schema(); let write_schema = build_target_arrow_schema(schema.fields())?; @@ -295,6 +334,49 @@ impl TableWrite { let target_bucket_row_number = core_options.dynamic_bucket_target_row_num(); let bucket_function_type = core_options.bucket_function_type()?; + let postpone_fixed_bucket = if use_postpone_fixed_bucket { + if total_buckets != POSTPONE_BUCKET || !has_primary_keys { + return Err(crate::Error::Unsupported { + message: format!( + "Postpone fixed-bucket writes require a primary-key table with bucket=-2, but table '{}' has bucket={total_buckets}", + table.identifier().full_name() + ), + }); + } + if core_options.deletion_vectors_enabled() { + return Err(crate::Error::Unsupported { + message: format!( + "Table '{}' cannot use postpone fixed-bucket writes with deletion-vectors.enabled=true because deletion-vector scans skip the level-0 files produced by batch writers; use the normal postpone writer or disable deletion vectors", + table.identifier().full_name() + ), + }); + } + let bucket_key_fields: Vec = bucket_key_indices + .iter() + .map(|&idx| fields[idx].clone()) + .collect(); + if !bucket_key_fields.is_empty() { + validate_bucket_function(bucket_function_type, &bucket_key_fields)?; + } + Some(PostponeFixedBucketState { + partition_field_indices: partition_field_indices.clone(), + bucket_key_indices: bucket_key_indices.clone(), + bucket_function_type, + max_parallelism: core_options + .postpone_batch_write_fixed_bucket_max_parallelism()?, + target_rows_per_bucket: core_options.postpone_target_row_num_per_bucket()?, + target_size_per_bucket: core_options.postpone_target_size_per_bucket()?, + metadata_loaded: false, + known_bucket_counts: HashMap::new(), + postpone_row_counts: HashMap::new(), + buffered_batches: HashMap::new(), + bucket_counts: HashMap::new(), + prepare_started: false, + }) + } else { + None + }; + let bucket_assigner = if is_dynamic_cross_partition { BucketAssignerEnum::CrossPartition(Box::new(CrossPartitionAssigner::new( table.clone(), @@ -377,6 +459,7 @@ impl TableWrite { has_dedicated_vector_fields, row_kind_generator, row_kind_filter, + postpone_fixed_bucket, }) } @@ -437,6 +520,13 @@ impl TableWrite { /// Write an Arrow RecordBatch. Rows are routed to the correct partition and bucket. pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) -> Result<()> { + if self + .postpone_fixed_bucket + .as_ref() + .is_some_and(|state| state.prepare_started) + { + return Err(Self::postpone_fixed_bucket_one_shot_error()); + } self.validate_write_batch_schema(batch)?; if batch.num_rows() == 0 { @@ -448,6 +538,11 @@ impl TableWrite { return Ok(()); } + if self.postpone_fixed_bucket.is_some() { + self.write_postpone_fixed_batch(&batch).await?; + return Ok(()); + } + let grouped = self.divide_by_partition_bucket(&batch).await?; for ((partition_bytes, bucket), sub_batch) in grouped { self.write_bucket(partition_bytes, bucket, sub_batch) @@ -758,6 +853,255 @@ impl TableWrite { }) } + async fn ensure_postpone_bucket_metadata(&mut self) -> Result<()> { + let should_load = self + .postpone_fixed_bucket + .as_ref() + .is_some_and(|state| !state.metadata_loaded); + if !should_load { + return Ok(()); + } + + let (known_bucket_counts, postpone_row_counts) = + self.load_postpone_bucket_metadata().await?; + let state = self.postpone_fixed_bucket.as_mut().unwrap(); + state.known_bucket_counts = known_bucket_counts; + state.postpone_row_counts = postpone_row_counts; + state.metadata_loaded = true; + Ok(()) + } + + async fn write_postpone_fixed_batch(&mut self, batch: &RecordBatch) -> Result<()> { + self.ensure_postpone_bucket_metadata().await?; + + let (partition_field_indices, bucket_key_indices, bucket_function_type) = { + let state = self.postpone_fixed_bucket.as_ref().unwrap(); + ( + state.partition_field_indices.clone(), + state.bucket_key_indices.clone(), + state.bucket_function_type, + ) + }; + 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, + self.table.schema().fields(), + )? + }; + + let mut groups: HashMap, Vec> = HashMap::new(); + for (row, partition) in partitions.into_iter().enumerate() { + groups.entry(partition).or_default().push(row); + } + for (partition, rows) in groups { + let sub_batch = Self::take_rows(batch, &rows)?; + let known_bucket_count = self + .postpone_fixed_bucket + .as_ref() + .unwrap() + .known_bucket_counts + .get(&partition) + .copied(); + if let Some(total_buckets) = known_bucket_count { + self.postpone_fixed_bucket + .as_mut() + .unwrap() + .bucket_counts + .insert(partition.clone(), total_buckets); + self.route_postpone_fixed_batch( + partition, + sub_batch, + total_buckets, + &bucket_key_indices, + bucket_function_type, + ) + .await?; + } else { + self.postpone_fixed_bucket + .as_mut() + .unwrap() + .buffered_batches + .entry(partition) + .or_default() + .push(sub_batch); + } + } + Ok(()) + } + + async fn load_postpone_bucket_metadata( + &mut self, + ) -> Result<(HashMap, i32>, HashMap, i64>)> { + let mut known_bucket_counts = HashMap::new(); + let mut postpone_row_counts = HashMap::new(); + let snapshot_manager = SnapshotManager::new( + self.table.file_io().clone(), + self.table.location().to_string(), + ); + let Some(snapshot) = snapshot_manager.get_latest_snapshot().await? else { + return Ok((known_bucket_counts, postpone_row_counts)); + }; + + let scan = + TableScan::new(&self.table, None, vec![], None, None, None).with_scan_all_files(); + for entry in scan.plan_manifest_entries(&snapshot).await? { + let partition = entry.partition().to_vec(); + if entry.bucket() == POSTPONE_BUCKET { + let rows = postpone_row_counts.entry(partition).or_insert(0_i64); + *rows = rows.saturating_add(entry.file().row_count); + } else if entry.bucket() >= 0 && entry.total_buckets() > 0 { + if let Some(previous) = + known_bucket_counts.insert(partition.clone(), entry.total_buckets()) + { + if previous != entry.total_buckets() { + return Err(crate::Error::DataInvalid { + message: format!( + "Partition has inconsistent total bucket counts: {previous} and {}", + entry.total_buckets() + ), + source: None, + }); + } + } + } + } + Ok((known_bucket_counts, postpone_row_counts)) + } + + fn infer_postpone_bucket_count( + input_rows: i64, + input_size: i64, + postpone_rows: i64, + target_rows_per_bucket: Option, + target_size_per_bucket: i64, + max_parallelism: i32, + ) -> i32 { + let buckets = if let Some(target_rows) = target_rows_per_bucket { + let total_rows = input_rows.saturating_add(postpone_rows); + total_rows.saturating_add(target_rows - 1) / target_rows + } else { + let estimated_size = if postpone_rows > 0 && input_rows > 0 { + let numerator = i128::from(input_size) + .saturating_mul(i128::from(input_rows.saturating_add(postpone_rows))); + let estimate = (numerator + i128::from(input_rows - 1)) / i128::from(input_rows); + estimate.min(i128::from(i64::MAX)) as i64 + } else { + input_size + }; + estimated_size.saturating_add(target_size_per_bucket - 1) / target_size_per_bucket + }; + buckets.max(1).min(i64::from(max_parallelism)) as i32 + } + + async fn flush_postpone_fixed_batches(&mut self) -> Result<()> { + let Some(state) = self.postpone_fixed_bucket.as_ref() else { + return Ok(()); + }; + if state.buffered_batches.is_empty() { + return Ok(()); + } + + let ( + buffered_batches, + bucket_key_indices, + bucket_function_type, + target_rows_per_bucket, + target_size_per_bucket, + max_parallelism, + postpone_row_counts, + ) = { + let state = self.postpone_fixed_bucket.as_mut().unwrap(); + ( + std::mem::take(&mut state.buffered_batches), + state.bucket_key_indices.clone(), + state.bucket_function_type, + state.target_rows_per_bucket, + state.target_size_per_bucket, + state.max_parallelism, + state.postpone_row_counts.clone(), + ) + }; + + for (partition, batches) in buffered_batches { + let input_rows = batches.iter().fold(0_i64, |rows, batch| { + rows.saturating_add(batch.num_rows() as i64) + }); + let input_size = + batches.iter().try_fold(0_i64, |size, batch| { + Ok::<_, crate::Error>(size.saturating_add(binary_row_batch_size( + batch, + self.table.schema().fields(), + )?)) + })?; + let postpone_rows = if self.is_overwrite { + 0 + } else { + postpone_row_counts.get(&partition).copied().unwrap_or(0) + }; + let total_buckets = Self::infer_postpone_bucket_count( + input_rows, + input_size, + postpone_rows, + target_rows_per_bucket, + target_size_per_bucket, + max_parallelism, + ); + { + let state = self.postpone_fixed_bucket.as_mut().unwrap(); + state + .known_bucket_counts + .insert(partition.clone(), total_buckets); + state.bucket_counts.insert(partition.clone(), total_buckets); + } + + for batch in batches { + self.route_postpone_fixed_batch( + partition.clone(), + batch, + total_buckets, + &bucket_key_indices, + bucket_function_type, + ) + .await?; + } + } + Ok(()) + } + + async fn route_postpone_fixed_batch( + &mut self, + partition: Vec, + batch: RecordBatch, + total_buckets: i32, + bucket_key_indices: &[usize], + bucket_function_type: BucketFunctionType, + ) -> Result<()> { + let buckets = if total_buckets <= 1 || bucket_key_indices.is_empty() { + vec![0; batch.num_rows()] + } else { + batch_bucket_ids( + &batch, + bucket_key_indices, + self.table.schema().fields(), + 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 { + let sub_batch = Self::take_rows(&batch, &rows)?; + self.write_bucket(partition.clone(), bucket, sub_batch) + .await?; + } + Ok(()) + } + /// Write a batch directly to the writer for the given (partition, bucket). async fn write_bucket( &mut self, @@ -782,8 +1126,20 @@ impl TableWrite { } /// Close all writers and collect CommitMessages for use with TableCommit. - /// Writers are cleared after this call, allowing the TableWrite to be reused. + /// Writers are cleared after this call, allowing the TableWrite to be reused, + /// except for fixed-bucket postpone batch writes, which are one-shot. pub async fn prepare_commit(&mut self) -> Result> { + if let Some(state) = self.postpone_fixed_bucket.as_mut() { + if state.prepare_started { + return Err(Self::postpone_fixed_bucket_one_shot_error()); + } + // Mark the one-shot operation before flushing. A failed prepare may + // already have consumed buffered batches or closed file writers and + // therefore cannot be retried safely on the same writer. + state.prepare_started = true; + } + self.flush_postpone_fixed_batches().await?; + let writers: Vec<(PartitionBucketKey, FileWriter)> = self.partition_writers.drain().collect(); @@ -814,6 +1170,10 @@ impl TableWrite { || !index_files.is_empty() { let mut msg = CommitMessage::new(partition_bytes, bucket, files.data_files); + msg.total_buckets = self + .postpone_fixed_bucket + .as_ref() + .and_then(|state| state.bucket_counts.get(&msg.partition).copied()); msg.new_changelog_files = files.changelog_files; msg.new_index_files = index_files; messages.push(msg); @@ -828,9 +1188,22 @@ impl TableWrite { messages.push(msg); } } + if let Some(state) = self.postpone_fixed_bucket.as_mut() { + state.metadata_loaded = false; + state.known_bucket_counts.clear(); + state.postpone_row_counts.clear(); + state.bucket_counts.clear(); + } Ok(messages) } + fn postpone_fixed_bucket_one_shot_error() -> crate::Error { + crate::Error::DataInvalid { + message: "Fixed-bucket postpone TableWrite only supports one prepare_commit call; create a new writer for the next batch".to_string(), + source: None, + } + } + async fn create_writer(&mut self, partition_bytes: Vec, bucket: i32) -> Result<()> { let partition_path = self.resolve_partition_path(&partition_bytes)?; @@ -3669,6 +4042,25 @@ mod tests { ) } + fn test_fixed_postpone_pk_table(file_io: &FileIO, table_path: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.target-row-num-per-bucket", "2") + .option("postpone.batch-write-fixed-bucket.max-parallelism", "8") + .build() + .unwrap(); + Table::new( + file_io.clone(), + Identifier::new("default", "test_fixed_postpone_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ) + } + fn test_postpone_partitioned_schema() -> TableSchema { let schema = Schema::builder() .column("pt", DataType::VarChar(VarCharType::string_type())) @@ -3737,6 +4129,220 @@ mod tests { assert_eq!(snapshot.total_record_count(), Some(3)); } + #[tokio::test] + async fn test_postpone_batch_write_uses_visible_fixed_buckets() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_write"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-user-1".to_string()).unwrap(); + write + .write_arrow_batch(&make_batch(vec![1, 2, 3, 4], vec![10, 20, 30, 40])) + .await + .unwrap(); + let state = write.postpone_fixed_bucket.as_ref().unwrap(); + assert_eq!(state.buffered_batches.len(), 1); + assert!(write.partition_writers.is_empty()); + let messages = write.prepare_commit().await.unwrap(); + assert!(!messages.is_empty()); + assert!(messages.iter().all(|message| message.bucket >= 0)); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + TableCommit::new(table.clone(), "fixed-user-1".to_string()) + .commit(messages) + .await + .unwrap(); + assert_eq!( + read_id_value_rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40)] + ); + + // A later small append reuses the partition's existing bucket count + // instead of inferring a new count from its own input size. + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-user-2".to_string()).unwrap(); + write + .write_arrow_batch(&make_batch(vec![5], vec![50])) + .await + .unwrap(); + // The real-bucket count is loaded on the first write, so existing + // partitions stream into file writers instead of retaining Arrow + // batches until prepare_commit. + let state = write.postpone_fixed_bucket.as_ref().unwrap(); + assert!(state.buffered_batches.is_empty()); + assert!(!write.partition_writers.is_empty()); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + TableCommit::new(table.clone(), "fixed-user-2".to_string()) + .commit(messages) + .await + .unwrap(); + assert_eq!( + read_id_value_rows(&table).await, + vec![(1, 10), (2, 20), (3, 30), (4, 40), (5, 50)] + ); + } + + #[tokio::test] + async fn test_postpone_bucket_count_uses_java_binary_row_size() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_java_binary_row_size"; + setup_dirs(&file_io, table_path).await; + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("left", DataType::VarChar(VarCharType::string_type())) + .column("right", DataType::VarChar(VarCharType::string_type())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("postpone.target-size-per-bucket", "20 kb") + .option("postpone.batch-write-fixed-bucket.max-parallelism", "8") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_java_binary_row_size"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let row_count = 1_000; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("left", ArrowDataType::Utf8, false), + ArrowField::new("right", ArrowDataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from_iter_values(0..row_count)), + Arc::new(StringArray::from(vec!["a"; row_count as usize])), + Arc::new(StringArray::from(vec!["b"; row_count as usize])), + ], + ) + .unwrap(); + + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-size-user".to_string()).unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + + // Java BinaryRows are 32,000 bytes, so a 20 KiB target plans two + // buckets. Arrow buffer sizing would incorrectly plan one. + assert!(!messages.is_empty()); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(2))); + } + + #[tokio::test] + async fn test_postpone_fixed_bucket_batch_write_is_one_shot() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_one_shot"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + let mut write = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-one-shot".to_string()).unwrap(); + + write + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert!(messages + .iter() + .all(|message| message.total_buckets == Some(1))); + + let write_error = write + .write_arrow_batch(&make_batch(vec![2, 3, 4, 5], vec![20, 30, 40, 50])) + .await + .unwrap_err(); + assert!( + matches!(write_error, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supports one prepare_commit call") + && message.contains("create a new writer")) + ); + + let prepare_error = write.prepare_commit().await.unwrap_err(); + assert!( + matches!(prepare_error, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supports one prepare_commit call") + && message.contains("create a new writer")) + ); + } + + #[test] + fn test_postpone_fixed_bucket_rejects_deletion_vectors() { + let file_io = test_file_io(); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "-2") + .option("deletion-vectors.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_postpone_dv"), + "memory:/test_postpone_dv".to_string(), + TableSchema::new(0, &schema), + None, + ); + + let error = TableWrite::new_postpone_fixed_bucket(&table, "test-user".to_string()) + .err() + .expect("fixed-bucket postpone writes must reject deletion vectors"); + assert!(matches!(error, crate::Error::Unsupported { ref message } + if message.contains("postpone fixed-bucket writes") + && message.contains("deletion-vectors.enabled=true"))); + } + + #[tokio::test] + async fn test_postpone_batch_write_rejects_conflicting_bucket_counts() { + let file_io = test_file_io(); + let table_path = "memory:/test_postpone_fixed_bucket_conflict"; + setup_dirs(&file_io, table_path).await; + let table = test_fixed_postpone_pk_table(&file_io, table_path); + + // Both writers plan against the empty table. Their differently sized + // inputs produce different bucket counts for the same partition. + let mut first = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-conflict-1".to_string()).unwrap(); + first + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let first_messages = first.prepare_commit().await.unwrap(); + assert!(first_messages + .iter() + .all(|message| message.total_buckets == Some(1))); + + let mut second = + TableWrite::new_postpone_fixed_bucket(&table, "fixed-conflict-2".to_string()).unwrap(); + second + .write_arrow_batch(&make_batch(vec![2, 3, 4, 5], vec![20, 30, 40, 50])) + .await + .unwrap(); + let second_messages = second.prepare_commit().await.unwrap(); + assert!(second_messages + .iter() + .all(|message| message.total_buckets == Some(2))); + + TableCommit::new(table.clone(), "fixed-conflict-1".to_string()) + .commit(first_messages) + .await + .unwrap(); + let error = TableCommit::new(table, "fixed-conflict-2".to_string()) + .commit(second_messages) + .await + .unwrap_err(); + assert!(error.to_string().contains("Postpone fixed-bucket conflict")); + } + #[tokio::test] async fn test_postpone_write_empty_batch() { let file_io = test_file_io(); diff --git a/crates/paimon/src/table/write_builder.rs b/crates/paimon/src/table/write_builder.rs index 52ba57754..d98348972 100644 --- a/crates/paimon/src/table/write_builder.rs +++ b/crates/paimon/src/table/write_builder.rs @@ -43,6 +43,23 @@ impl<'a> WriteBuilder<'a> { } } + /// Create a builder for one-shot fixed-bucket writes to a postpone table. + /// + /// Unlike [`WriteBuilder::new`], writers created by this builder plan real + /// buckets for `bucket = -2` tables. Keeping this mode explicit prevents a + /// normal batch writer from silently changing postpone-write semantics. + pub fn new_postpone_fixed_bucket(table: &'a Table) -> crate::Result { + if table.is_format_table() { + return Err(crate::Error::Unsupported { + message: "Postpone fixed-bucket writes are only supported for Paimon tables" + .to_string(), + }); + } + Ok(Self(WriteBuilderKind::Paimon( + PaimonWriteBuilder::new(table).with_postpone_fixed_bucket(), + ))) + } + /// Get the commit user shared by writers and committers created by this builder. pub fn commit_user(&self) -> &str { match &self.0 { @@ -120,6 +137,7 @@ struct PaimonWriteBuilder<'a> { table: &'a Table, commit_user: String, overwrite: bool, + postpone_fixed_bucket: bool, } impl<'a> PaimonWriteBuilder<'a> { @@ -128,9 +146,15 @@ impl<'a> PaimonWriteBuilder<'a> { table, commit_user: Uuid::new_v4().to_string(), overwrite: false, + postpone_fixed_bucket: false, } } + fn with_postpone_fixed_bucket(mut self) -> Self { + self.postpone_fixed_bucket = true; + self + } + /// Get the commit user shared by writers and committers created by this builder. /// /// This value is persisted in snapshot metadata and used for duplicate @@ -198,7 +222,11 @@ impl<'a> PaimonWriteBuilder<'a> { .to_string(), }); } - let write = TableWrite::new(self.table, self.commit_user.clone())?; + let write = if self.postpone_fixed_bucket { + TableWrite::new_postpone_fixed_bucket(self.table, self.commit_user.clone())? + } else { + TableWrite::new(self.table, self.commit_user.clone())? + }; Ok(if self.overwrite { write.with_overwrite() } else { @@ -430,6 +458,39 @@ mod tests { assert_eq!(snapshot.commit_user(), "my-commit-user"); } + #[tokio::test] + async fn test_postpone_fixed_bucket_builder_is_explicit() { + let file_io = test_file_io(); + let table_path = "memory:/test_explicit_postpone_fixed_bucket_builder"; + setup_dirs(&file_io, table_path).await; + let table = test_postpone_pk_table(&file_io, table_path); + + let mut normal = table.new_write_builder().new_write().unwrap(); + normal + .write_arrow_batch(&make_batch(vec![1], vec![10])) + .await + .unwrap(); + let normal_messages = normal.prepare_commit().await.unwrap(); + assert_eq!(normal_messages.len(), 1); + assert_eq!(normal_messages[0].bucket, POSTPONE_BUCKET); + assert_eq!(normal_messages[0].total_buckets, None); + + let builder = table + .new_postpone_fixed_bucket_write_builder() + .unwrap() + .with_commit_user("explicit-fixed-user") + .unwrap(); + let mut fixed = builder.new_write().unwrap(); + fixed + .write_arrow_batch(&make_batch(vec![2], vec![20])) + .await + .unwrap(); + let fixed_messages = fixed.prepare_commit().await.unwrap(); + assert_eq!(fixed_messages.len(), 1); + assert_eq!(fixed_messages[0].bucket, 0); + assert_eq!(fixed_messages[0].total_buckets, Some(1)); + } + #[tokio::test] async fn test_branch_reference_rejects_write_and_index_builders() { let table = as_main_branch_reference(test_postpone_pk_table(