From 2d76d3cc7c0c486e7dc5b849c3c2131cbbdb1e9d Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 17:28:50 +0800 Subject: [PATCH 1/4] [core] Preserve append writer configuration per bucket across prepares --- .../python/pypaimon_rust/datafusion.pyi | 2 + bindings/python/src/write.rs | 197 +++++- .../src/table/bucket_assigner_constant.rs | 16 +- .../paimon/src/table/bucket_assigner_fixed.rs | 16 +- crates/paimon/src/table/table_write.rs | 117 +++- crates/paimon/src/table/table_write_append.rs | 283 ++++++++ crates/paimon/src/table/table_write_input.rs | 38 +- .../src/table/table_write_precomputed.rs | 3 + .../tests/filtered_write_validation_test.rs | 70 +- .../paimon/tests/table_partial_write_test.rs | 641 +++++++++++++++++- 10 files changed, 1316 insertions(+), 67 deletions(-) create mode 100644 crates/paimon/src/table/table_write_append.rs diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 4d4dcefb6..0fef05bb9 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -446,6 +446,7 @@ class CommitMessage: ... class BatchTableWrite: + def _row_write_columns(self, routing_batch: pyarrow.RecordBatch) -> List[str]: ... def with_blob_uri_reader_factory( self, factory: Optional[_UriReaderFactory] ) -> "BatchTableWrite": @@ -575,6 +576,7 @@ class StreamTableUpdate: ... class StreamTableWrite: + def _row_write_columns(self, routing_batch: pyarrow.RecordBatch) -> List[str]: ... def with_blob_uri_reader_factory( self, factory: Optional[_UriReaderFactory] ) -> "StreamTableWrite": ... diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index e0dce1757..61130dcbb 100644 --- a/bindings/python/src/write.rs +++ b/bindings/python/src/write.rs @@ -88,6 +88,7 @@ impl WriteContext { ))), table_location: self.table.location().to_string(), commit_user: self.commit_user.clone(), + blob_consumers: Vec::new(), }) } } @@ -346,6 +347,19 @@ impl WriteTarget { } } + fn row_write_columns(&self, routing: &RecordBatch) -> paimon::Result> { + match self { + Self::Table(writer) => Ok(writer + .write_fields_for_row(routing)? + .iter() + .map(|field| field.name().to_owned()) + .collect()), + Self::PostponeFixed(_) => Err(paimon::Error::Unsupported { + message: "_row_write_columns requires a Paimon append table".into(), + }), + } + } + async fn write_arrow( &mut self, batch: &RecordBatch, @@ -464,6 +478,7 @@ impl PyPostponeFixedBucketWriteBuilder { ))), table_location: context.table.location().to_string(), commit_user: context.commit_user.clone(), + blob_consumers: Vec::new(), }, prepared: false, }) @@ -478,12 +493,15 @@ struct WriteState { inner: Option, table_location: String, commit_user: String, + // Core keeps prior consumers in captured partition/bucket settings. Expose + // every owned Python reference to GC, including replaced consumers. + blob_consumers: Vec>, } /// Descriptor serialization and exception transport only; record boundaries, /// callback timing and requested flushes belong to Rust core. struct PythonBlobConsumer { - callback: Py, + callback: Mutex>>, } impl BlobConsumer for PythonBlobConsumer { @@ -496,10 +514,14 @@ impl BlobConsumer for PythonBlobConsumer { Python::attach(|py| { let encoded = descriptor.map(|descriptor| PyBytes::new(py, &descriptor.serialize())); - self.callback - .bind(py) - .call1((field_name, encoded))? - .is_truthy() + let callback = self + .callback + .lock() + .unwrap() + .as_ref() + .map(|callback| callback.clone_ref(py)) + .ok_or_else(|| PyRuntimeError::new_err("Blob consumer has been cleared"))?; + callback.bind(py).call1((field_name, encoded))?.is_truthy() }) }); result.map_err(callback_error) @@ -858,7 +880,9 @@ impl WriteState { if !callback.bind(py).is_callable() { return Err(PyTypeError::new_err("blob_consumer must be callable")); } - Ok(Arc::new(PythonBlobConsumer { callback })) + Ok(Arc::new(PythonBlobConsumer { + callback: Mutex::new(Some(callback)), + })) }) .transpose()?; self.inner @@ -870,9 +894,42 @@ impl WriteState { .map(|consumer| consumer as Arc), ) .map_err(to_py_err)?; + // Release replaced callbacks which no bucket has captured. A snapshot + // retained by core contributes another Arc, without another Py ref. + self.blob_consumers + .retain(|consumer| Arc::strong_count(consumer) > 1); + if let Some(consumer) = consumer { + self.blob_consumers.push(consumer); + } Ok(()) } + fn traverse_callbacks(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + for consumer in &self.blob_consumers { + visit.call(&*consumer.callback.lock().unwrap())?; + } + Ok(()) + } + + fn clear_callbacks(&self) { + let references = self + .blob_consumers + .iter() + .map(|consumer| consumer.callback.lock().unwrap().take()) + .collect::>(); + // Destructors may re-enter Native. Drop Python references outside locks. + drop(references); + } + + fn row_write_columns(&self, routing: &Bound<'_, PyAny>) -> PyResult> { + let routing = crate::arrow_input::write_batch_from_pyarrow(routing)?; + let inner = self + .inner + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?; + inner.row_write_columns(&routing).map_err(to_py_err) + } + fn write_result(&self, result: paimon::Result<()>) -> PyResult<()> { result.map_err(to_callback_error) } @@ -1213,6 +1270,18 @@ impl PyBatchTableUpdate { #[pymethods] impl PyBatchTableWrite { + fn _row_write_columns(&self, routing_batch: &Bound<'_, PyAny>) -> PyResult> { + self.state.row_write_columns(routing_batch) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.state.traverse_callbacks(visit) + } + + fn __clear__(&mut self) { + self.state.clear_callbacks(); + } + #[pyo3(signature = (factory))] fn with_blob_uri_reader_factory( mut slf: PyRefMut<'_, Self>, @@ -1245,6 +1314,7 @@ impl PyBatchTableWrite { if let Some(mut writer) = self.state.inner.take() { py.detach(|| runtime().block_on(writer.close())); } + self.state.clear_callbacks(); } #[pyo3(signature = (batch, bucket=None))] @@ -1279,6 +1349,18 @@ pub struct PyStreamTableWrite { #[pymethods] impl PyStreamTableWrite { + fn _row_write_columns(&self, routing_batch: &Bound<'_, PyAny>) -> PyResult> { + self.state.row_write_columns(routing_batch) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.state.traverse_callbacks(visit) + } + + fn __clear__(&mut self) { + self.state.clear_callbacks(); + } + #[pyo3(signature = (factory))] fn with_blob_uri_reader_factory( mut slf: PyRefMut<'_, Self>, @@ -1311,6 +1393,7 @@ impl PyStreamTableWrite { if let Some(mut writer) = self.state.inner.take() { py.detach(|| runtime().block_on(writer.close())); } + self.state.clear_callbacks(); } #[pyo3(signature = (batch, bucket=None))] @@ -1687,6 +1770,108 @@ mod tests { use paimon::spec::{BigIntType, IntType, Schema, TableSchema, VarCharType}; use pyo3::ffi::c_str; + #[test] + fn replaced_blob_consumers_are_visible_to_writer_cycle_collection() { + Python::attach(|py| { + for batch in [false, true] { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "blob_callback_gc"), + "memory:/blob_callback_gc".into(), + TableSchema::new(0, &schema), + None, + )); + let context = WriteContext::new(table); + let writer = if batch { + Py::new( + py, + PyBatchTableWrite { + state: context.new_write(false).unwrap(), + prepared: false, + }, + ) + .unwrap() + .into_any() + } else { + Py::new( + py, + PyStreamTableWrite { + state: context.new_write(false).unwrap(), + }, + ) + .unwrap() + .into_any() + }; + let locals = PyDict::new(py); + locals.set_item("writer", writer).unwrap(); + py.run( + c_str!( + r#" +import gc +finalized = [] +class Callback: + def __init__(self, writer, label): + self.writer = writer + self.label = label + def __call__(self, field, descriptor): + return False + def __del__(self): + finalized.append(self.label) +writer.with_blob_consumer(Callback(writer, 'first')) +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + // Capture the first callback in core without importing PyArrow + // or building/installing the extension for a Python test. + let input = RecordBatch::try_from_iter([( + "id", + Arc::new(arrow::array::Int32Array::from(vec![1])) as arrow::array::ArrayRef, + )]) + .unwrap(); + let object = locals.get_item("writer").unwrap().unwrap(); + if batch { + let mut writer = object.extract::>().unwrap(); + let inner = writer.state.inner.as_mut().unwrap(); + assert_eq!(inner.row_write_columns(&input).unwrap(), vec!["id"]); + assert!(inner.row_write_columns(&input.slice(0, 0)).is_err()); + py.detach(|| runtime().block_on(inner.write_arrow(&input, None))) + .unwrap(); + } else { + let mut writer = object + .extract::>() + .unwrap(); + let inner = writer.state.inner.as_mut().unwrap(); + assert_eq!(inner.row_write_columns(&input).unwrap(), vec!["id"]); + assert!(inner.row_write_columns(&input.slice(0, 0)).is_err()); + py.detach(|| runtime().block_on(inner.write_arrow(&input, None))) + .unwrap(); + } + drop(object); + py.run( + c_str!( + r#" +writer.with_blob_consumer(Callback(writer, 'second')) +assert finalized == [] +del writer +gc.collect() +assert sorted(finalized) == ['first', 'second'], finalized +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + } + }); + } + #[test] fn merge_table_source_conditions_resolve_read_metadata() { Python::attach(|py| { diff --git a/crates/paimon/src/table/bucket_assigner_constant.rs b/crates/paimon/src/table/bucket_assigner_constant.rs index 843d2464a..b8d6e7eb3 100644 --- a/crates/paimon/src/table/bucket_assigner_constant.rs +++ b/crates/paimon/src/table/bucket_assigner_constant.rs @@ -45,11 +45,9 @@ impl ConstantBucketAssigner { pub fn bucket(&self) -> i32 { self.bucket } -} -impl BucketAssigner for ConstantBucketAssigner { - async fn assign_batch( - &mut self, + pub(super) fn assign_batch_sync( + &self, batch: &RecordBatch, fields: &[DataField], ) -> Result { @@ -67,6 +65,16 @@ impl BucketAssigner for ConstantBucketAssigner { skips: Vec::new(), }) } +} + +impl BucketAssigner for ConstantBucketAssigner { + async fn assign_batch( + &mut self, + batch: &RecordBatch, + fields: &[DataField], + ) -> Result { + self.assign_batch_sync(batch, fields) + } async fn prepare_commit_index( &mut self, diff --git a/crates/paimon/src/table/bucket_assigner_fixed.rs b/crates/paimon/src/table/bucket_assigner_fixed.rs index 684a98a84..042ec0ebd 100644 --- a/crates/paimon/src/table/bucket_assigner_fixed.rs +++ b/crates/paimon/src/table/bucket_assigner_fixed.rs @@ -52,11 +52,9 @@ impl FixedBucketAssigner { total_buckets, } } -} -impl BucketAssigner for FixedBucketAssigner { - async fn assign_batch( - &mut self, + pub(super) fn assign_batch_sync( + &self, batch: &RecordBatch, fields: &[DataField], ) -> Result { @@ -82,6 +80,16 @@ impl BucketAssigner for FixedBucketAssigner { skips: Vec::new(), }) } +} + +impl BucketAssigner for FixedBucketAssigner { + async fn assign_batch( + &mut self, + batch: &RecordBatch, + fields: &[DataField], + ) -> Result { + self.assign_batch_sync(batch, fields) + } async fn prepare_commit_index( &mut self, diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 39842497a..9dfd77503 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -52,11 +52,23 @@ use arrow_array::{Array, ArrayRef, Int8Array, RecordBatch}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; +#[path = "table_write_append.rs"] +mod append; #[path = "table_write_input.rs"] mod input; #[path = "table_write_precomputed.rs"] mod precomputed; +/// Java's append writer captures its factory settings per partition/bucket. +/// Keep only these settings across prepares, not completed file resources. +#[derive(Clone)] +struct AppendWriteConfig { + schema: Arc, + fields: Vec, + blob_consumer: Option>, + blob_uri_reader_factory: Option>, +} + /// Enum to hold either an append-only writer, a key-value writer, or a postpone writer. enum FileWriter { Append(DataFileWriter), @@ -143,6 +155,8 @@ pub struct TableWrite { write_fields: Vec, written: bool, partition_writers: HashMap, + append_write_configs: HashMap, + append_write_type_changed: bool, partition_computer: PartitionComputer, partition_keys: Vec, schema_id: i64, @@ -213,6 +227,8 @@ impl TableWrite { write_fields: schema.fields().to_vec(), written: false, partition_writers: HashMap::new(), + append_write_configs: HashMap::new(), + append_write_type_changed: false, partition_computer: PartitionComputer::new( schema.partition_keys(), schema.fields(), @@ -513,6 +529,8 @@ impl TableWrite { write_fields: fields.to_vec(), written: false, partition_writers: HashMap::new(), + append_write_configs: HashMap::new(), + append_write_type_changed: false, partition_computer, partition_keys, schema_id: schema.id(), @@ -641,15 +659,13 @@ impl TableWrite { /// Full-table inputs may provide routing columns which are not written. /// Selected inputs must include partition columns and, unless a bucket is /// supplied through [`Self::write_arrow`], the columns used for hashing. - /// Configure this before writing any nonempty batch. + /// Changes apply to partition/bucket writers created afterwards. Existing + /// writers retain their selected columns, including across prepares. A full + /// table batch is projected to each writer's captured columns; a selected + /// batch must match the destination writer's names and order, optionally + /// accompanied by routing and configured row-kind fields. pub fn with_write_type(&mut self, write_columns: Vec) -> Result<&mut Self> { self.ensure_active()?; - if self.written { - return Err(crate::Error::DataInvalid { - message: "with_write_type must be called before any data is written".into(), - source: None, - }); - } if self.format_writer.is_some() || !self.table.schema().primary_keys().is_empty() { return Err(crate::Error::Unsupported { message: "with_write_type requires a Paimon append table".into(), @@ -722,27 +738,22 @@ impl TableWrite { )) }; let write_schema = build_target_arrow_schema(&fields)?; + self.append_write_type_changed |= self.written && self.write_schema != write_schema; self.bucket_assigner = assigner; self.write_schema = write_schema; self.write_fields = fields; Ok(self) } - /// Receive dedicated append BLOB payload descriptors. Configure before - /// writing, as in PyPaimon. Java's primary-key and inline BLOB writers do - /// not invoke this consumer. Automatic abort preserves consumer Blob files; + /// Receive dedicated append BLOB payload descriptors. Changes apply to + /// newly created partition/bucket writers, as in Java. Primary-key and inline + /// BLOB writers do not invoke it. Automatic abort preserves consumer Blob files; /// callers can explicitly discard prepared files with TableCommit::abort. pub fn with_blob_consumer( &mut self, consumer: Option>, ) -> Result<&mut Self> { self.ensure_active()?; - if self.written { - return Err(crate::Error::DataInvalid { - message: "with_blob_consumer must be called before any write operation".into(), - source: None, - }); - } if self.format_writer.is_some() { return Err(crate::Error::Unsupported { message: "with_blob_consumer requires a Paimon table".into(), @@ -782,6 +793,9 @@ impl TableWrite { return writer.write(batch).await; } self.ensure_active()?; + if self.has_reconfigured_append_type() { + return self.write_reconfigured_append(batch, None).await; + } let filtered = self.filter_write_input(batch)?; let input = filtered.as_ref().unwrap_or(batch); let Some(batch) = self.normalize_filtered_write_batch(input)? else { @@ -881,6 +895,23 @@ impl TableWrite { // Reject malformed schemas even when every event would be ignored. // Empty arrays check equivalent layouts without examining row values. self.validate_write_batch_schema(&RecordBatch::new_empty(batch.schema()))?; + let projected = self.project_full_write_input(batch)?; + self.filter_write_input_with_schema( + batch, + projected.as_ref().unwrap_or(batch), + &self.write_schema, + ) + } + + fn filter_write_input_with_schema( + &self, + batch: &RecordBatch, + physical: &RecordBatch, + expected_schema: &arrow_schema::Schema, + ) -> Result> { + if self.primary_key_indices.is_empty() && self.row_kind_generator.is_none() { + return Ok(None); + } let kinds = if let Some(generator) = &self.row_kind_generator { generator.generate_batch(batch)? } else if let Some(column) = batch.column_by_name(VALUE_KIND_FIELD_NAME) { @@ -899,9 +930,7 @@ impl TableWrite { // still identify the event; other engines and partition migrations // require every declared top-level NOT NULL field. let delete_required_fields = self.delete_not_null_fields(); - let projected = self.project_full_write_input(batch)?; - let physical = projected.as_ref().unwrap_or(batch); - for (index, field) in self.write_schema.fields().iter().enumerate() { + for (index, field) in expected_schema.fields().iter().enumerate() { if field.is_nullable() || physical.column(index).null_count() == 0 { continue; } @@ -1540,7 +1569,12 @@ impl TableWrite { // Only this failed invocation's outputs: no message has been returned // or submitted. Earlier successful prepare calls belong to the caller. let mut cleanup = messages.to_vec(); - if self.primary_key_indices.is_empty() && self.blob_consumer.is_some() { + if self.primary_key_indices.is_empty() + && self + .append_write_configs + .values() + .any(|config| config.blob_consumer.is_some()) + { // A successful partition may already have exposed its dedicated // Blob files to the consumer. Preserve Java's deletion policy in // this table-wide cleanup too. Explicit TableCommit.abort remains @@ -1561,7 +1595,15 @@ impl TableWrite { let partition_path = self.resolve_partition_path(&partition_bytes)?; let writer = if self.primary_key_indices.is_empty() { - self.create_append_writer(partition_path, bucket)? + let key = (partition_bytes.clone(), bucket); + let config = self + .append_write_configs + .get(&key) + .cloned() + .unwrap_or_else(|| self.append_write_config()); + let writer = self.create_append_writer(partition_path, bucket, &config)?; + self.append_write_configs.insert(key, config); + writer } else if bucket == POSTPONE_BUCKET { self.create_postpone_writer(partition_path, bucket)? } else { @@ -1584,7 +1626,12 @@ impl TableWrite { } /// Create an append-only writer for non-PK tables. - fn create_append_writer(&self, partition_path: String, bucket: i32) -> Result { + fn create_append_writer( + &self, + partition_path: String, + bucket: i32, + config: &AppendWriteConfig, + ) -> Result { if self.has_blob_fields || self.has_dedicated_vector_fields || !self.blob_view_fields.is_empty() @@ -1606,15 +1653,15 @@ impl TableWrite { self.file_format.clone(), self.vector_target_file_size, self.vector_file_format.as_deref(), - &self.write_schema, - &self.write_fields, + &config.schema, + &config.fields, fields, self.table.schema().options(), &self.blob_inline_fields, )? .with_blob_writer_options( - self.blob_consumer.clone(), - self.blob_uri_reader_factory.clone(), + config.blob_consumer.clone(), + config.blob_uri_reader_factory.clone(), )? .with_resources(self.resources.clone()), ))) @@ -1631,21 +1678,22 @@ impl TableWrite { self.file_compression_zstd_level, self.write_buffer_size, self.file_format.clone(), - self.write_fields.clone(), + config.fields.clone(), self.table.schema().options().clone(), Some(0), None, (!super::data_evolution_fields::can_omit_normal_write_cols( self.table.schema().fields(), - &self - .write_fields + &config + .fields .iter() .map(|field| field.name().to_string()) .collect::>(), &CoreOptions::new(self.table.schema().options()), )) .then(|| { - self.write_fields + config + .fields .iter() .map(|field| field.name().to_string()) .collect() @@ -1654,7 +1702,7 @@ impl TableWrite { .with_file_index( self.file_index_options .as_ref() - .and_then(|options| options.project_to_fields(&self.write_fields)) + .and_then(|options| options.project_to_fields(&config.fields)) .map(Arc::new), ) .with_target_file_row_num(self.target_file_row_num) @@ -1800,7 +1848,7 @@ pub(in crate::table) mod tests { use arrow_array::{ArrayRef, LargeBinaryArray}; use std::sync::Mutex; - for has_consumer in [false, true] { + for (has_consumer, replace_consumer) in [(false, false), (true, false), (true, true)] { let io = test_file_io(); let schema = Schema::builder() .column("pt", DataType::Int(IntType::new())) @@ -1842,6 +1890,11 @@ pub(in crate::table) mod tests { ]) .unwrap(); writer.write_arrow_batch(&batch).await.unwrap(); + if replace_consumer { + // Captured callbacks own these payloads even when the factory + // has subsequently disabled its consumer for future writers. + writer.with_blob_consumer(None).unwrap(); + } assert_eq!(writer.partition_writers.len(), 2); let FileWriter::AppendDedicated(partition) = writer.partition_writers.values_mut().next().unwrap() diff --git a/crates/paimon/src/table/table_write_append.rs b/crates/paimon/src/table/table_write_append.rs new file mode 100644 index 000000000..aca44f9c0 --- /dev/null +++ b/crates/paimon/src/table/table_write_append.rs @@ -0,0 +1,283 @@ +// 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. + +//! Append writer factory configuration survives physical-file prepares. + +use super::*; + +impl TableWrite { + /// Resolve the columns captured by the destination append writer without + /// creating files or changing bucket assignment state. Row adapters can + /// encode only these columns, leaving unused payloads untouched. The input + /// must contain exactly one row and the table's routing columns. + pub fn write_fields_for_row(&self, routing_batch: &RecordBatch) -> Result> { + self.ensure_active()?; + if self.format_writer.is_some() || !self.primary_key_indices.is_empty() { + return Err(crate::Error::Unsupported { + message: "write_fields_for_row requires a Paimon append table".into(), + }); + } + if routing_batch.num_rows() != 1 { + return Err(crate::Error::DataInvalid { + message: "write_fields_for_row requires exactly one routing row".into(), + source: None, + }); + } + let names = self.automatic_routing_names(); + let key = if names.is_empty() { + // A zero-column routing batch still represents one row. Avoid a + // schema normalizer which would infer zero rows from no arrays. + let bucket = CoreOptions::new(self.table.schema().options()).bucket(); + ( + EMPTY_SERIALIZED_ROW.clone(), + if bucket == POSTPONE_BUCKET { + POSTPONE_BUCKET + } else { + 0 + }, + ) + } else { + let routing = self.normalize_routing_columns(routing_batch, &names)?; + let mut output = self.assign_append_routing(routing)?; + (output.partition_bytes.remove(0), output.buckets[0]) + }; + Ok(self + .append_write_configs + .get(&key) + .map(|config| &config.fields) + .unwrap_or(&self.write_fields) + .clone()) + } + + pub(super) fn append_write_config(&self) -> AppendWriteConfig { + AppendWriteConfig { + schema: self.write_schema.clone(), + fields: self.write_fields.clone(), + blob_consumer: self.blob_consumer.clone(), + blob_uri_reader_factory: self.blob_uri_reader_factory.clone(), + } + } + + pub(super) fn has_reconfigured_append_type(&self) -> bool { + self.append_write_type_changed + } + + /// Java captures the physical write type per partition/bucket. Arrow's + /// named columns must agree with that captured type, rather than being + /// reinterpreted positionally as Java's InternalRow can be. + pub(super) async fn write_reconfigured_append( + &mut self, + input: &RecordBatch, + bucket: Option, + ) -> Result<()> { + let full_schema = build_target_arrow_schema(self.table.schema().fields())?; + let actual_schema = input.schema(); + let names_match = |expected: &arrow_schema::Schema| { + actual_schema + .fields() + .iter() + .map(|f| f.name()) + .eq(expected.fields().iter().map(|f| f.name())) + }; + let mut routing_names = if bucket.is_some() { + self.partition_keys.clone() + } else { + self.automatic_routing_names() + }; + if let Some(field) = CoreOptions::new(self.table.schema().options()).rowkind_field() { + routing_names.push(field.to_owned()); + } + let routing_transport = |expected: &arrow_schema::Schema| { + is_routing_transport(&actual_schema, expected, &routing_names) + }; + let full_input = names_match(&full_schema); + let candidates = std::iter::once(self.write_schema.as_ref()) + .chain( + self.append_write_configs + .values() + .map(|config| config.schema.as_ref()), + ) + .filter(|schema| names_match(schema) || routing_transport(schema)) + .collect::>(); + let logical_schema = if full_input + || names_match(&self.write_schema) + || routing_transport(&self.write_schema) + { + self.write_schema.clone() + } else if let Some(exact) = candidates + .iter() + .copied() + .find(|schema| names_match(schema)) + { + // A selected Arrow layout explicitly declares its logical fields. + Arc::new(exact.clone()) + } else if candidates.is_empty() { + return Err(crate::Error::DataInvalid { + message: "Input columns must match the table or a selected append write type, with optional routing columns".into(), + source: None, + }); + } else { + // Historical types can differ only in routing/event fields while + // admitting the same transport. Do not choose their logical NULL + // checks by HashMap order: retain the common selected fields here, + // then validate each destination's complete physical type below. + let fields = self + .table + .schema() + .fields() + .iter() + .filter(|field| { + candidates + .iter() + .all(|candidate| candidate.field_with_name(field.name()).is_ok()) + }) + .cloned() + .collect::>(); + build_target_arrow_schema(&fields)? + }; + // Validate every provided transport field, without inspecting values of + // omitted payloads. Row marshalling may add routing/event fields to a + // captured subset, but arbitrary payload extras are not accepted. + let transport_fields = actual_schema + .fields() + .iter() + .map(|field| { + self.table + .schema() + .fields() + .iter() + .find(|candidate| candidate.name() == field.name()) + .cloned() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("Column '{}' is not in table schema", field.name()), + source: None, + }) + }) + .collect::>>()?; + let transport_schema = build_target_arrow_schema(&transport_fields)?; + self.normalize_batch_schema(&input.slice(0, 0), &transport_schema)?; + if input.num_rows() == 0 { + return Ok(()); + } + let logical = project(input, &logical_schema)?; + // The facade resolves events and checks its selected NOT NULL fields + // before routing, as Java TableWriteImpl does. Ignored events must not + // cast routing values or open Blob payloads. + let filtered = self.filter_write_input_with_schema(input, &logical, &logical_schema)?; + let input = filtered.as_ref().unwrap_or(input); + if input.num_rows() == 0 { + return Ok(()); + } + let groups = if let Some(bucket) = bucket { + let partitions = self.precomputed_partition_bytes(input, input)?; + let mut groups: HashMap> = HashMap::new(); + for (row, partition) in partitions.into_iter().enumerate() { + groups.entry((partition, bucket)).or_default().push(row); + } + groups + .into_iter() + .map(|(key, rows)| Ok((key, take_rows(input, &rows)?))) + .collect::>>()? + } else { + let routing = self.normalize_routing_columns(input, &self.automatic_routing_names())?; + self.divide_write_input(input, Some(routing)).await? + }; + let mut prepared = Vec::with_capacity(groups.len()); + for (key, batch) in groups { + let expected = self + .append_write_configs + .get(&key) + .map(|config| &config.schema) + .unwrap_or(&self.write_schema); + let physical = if full_input || routing_transport(expected) { + project(&batch, expected)? + } else { + // Selected inputs retain their named layout. Never reinterpret + // equal-width input positionally for a previously opened writer. + batch + }; + if physical.num_rows() > 0 { + prepared.push((key, self.normalize_batch_schema(&physical, expected)?)); + } + } + if bucket.is_some() && prepared.len() > 1 { + return Err(crate::Error::DataInvalid { + message: "A precomputed bucket group contained multiple partitions".into(), + source: None, + }); + } + // Preflight every destination before creating files or modifying any + // existing writer. A bad subset cannot partially stage a mixed batch. + for ((partition, bucket), batch) in prepared { + self.written = true; + self.write_bucket(partition, bucket, batch).await?; + } + Ok(()) + } +} + +fn project(batch: &RecordBatch, expected: &arrow_schema::Schema) -> Result { + let indices = expected + .fields() + .iter() + .map(|field| { + batch + .schema() + .index_of(field.name()) + .map_err(|error| crate::Error::DataInvalid { + message: format!("Cannot project captured append write type: {error}"), + source: None, + }) + }) + .collect::>>()?; + batch + .project(&indices) + .map_err(|error| crate::Error::DataInvalid { + message: format!("Cannot project captured append write type: {error}"), + source: None, + }) +} + +/// Routing/event fields may accompany a named subset. Equal-width reordered +/// subsets remain invalid Arrow layouts; only additional routing fields permit +/// projection, and every captured payload field must still be present. +fn is_routing_transport( + actual: &arrow_schema::Schema, + expected: &arrow_schema::Schema, + routing_names: &[String], +) -> bool { + if actual.fields().len() <= expected.fields().len() { + return false; + } + let names = actual + .fields() + .iter() + .map(|field| field.name()) + .collect::>(); + names.len() == actual.fields().len() + && expected + .fields() + .iter() + .all(|field| names.contains(field.name())) + && actual.fields().iter().all(|field| { + expected + .fields() + .iter() + .any(|candidate| candidate.name() == field.name()) + || routing_names.contains(field.name()) + }) +} diff --git a/crates/paimon/src/table/table_write_input.rs b/crates/paimon/src/table/table_write_input.rs index b2e99b4ce..5c549a78b 100644 --- a/crates/paimon/src/table/table_write_input.rs +++ b/crates/paimon/src/table/table_write_input.rs @@ -64,7 +64,7 @@ impl TableWrite { }) } - fn normalize_routing_columns( + pub(super) fn normalize_routing_columns( &self, input: &RecordBatch, names: &[String], @@ -113,14 +113,19 @@ impl TableWrite { Ok(RoutingInput { batch, fields }) } - pub(super) fn normalize_automatic_routing_input( - &self, - input: &RecordBatch, - ) -> Result> { + pub(super) fn automatic_routing_names(&self) -> Vec { let mut names = self.partition_keys.clone(); if CoreOptions::new(self.table.schema().options()).bucket() > 1 { names.extend(self.table.schema().bucket_keys()); } + names + } + + pub(super) fn normalize_automatic_routing_input( + &self, + input: &RecordBatch, + ) -> Result> { + let names = self.automatic_routing_names(); if names .iter() .all(|name| self.write_fields.iter().any(|field| field.name() == name)) @@ -141,6 +146,11 @@ impl TableWrite { let Some(routing) = routing else { return self.divide_by_partition_bucket(physical).await; }; + let output = self.assign_append_routing(routing)?; + Self::group_assigned_rows(physical, output) + } + + pub(super) fn assign_append_routing(&self, routing: RoutingInput) -> Result { let indices = |names: &[String]| { names .iter() @@ -155,27 +165,25 @@ impl TableWrite { }; let partitions = indices(&self.partition_keys); let options = CoreOptions::new(self.table.schema().options()); - let mut assigner = if options.bucket() > 1 { - BucketAssignerEnum::Fixed(FixedBucketAssigner::new( + if options.bucket() > 1 { + FixedBucketAssigner::new( partitions, indices(&self.table.schema().bucket_keys()), options.bucket_function_type()?, options.bucket(), - )) + ) + .assign_batch_sync(&routing.batch, &routing.fields) } else { - BucketAssignerEnum::Constant(ConstantBucketAssigner::new( + ConstantBucketAssigner::new( partitions, if options.bucket() == POSTPONE_BUCKET { POSTPONE_BUCKET } else { 0 }, - )) - }; - let output = assigner - .assign_batch(&routing.batch, &routing.fields) - .await?; - Self::group_assigned_rows(physical, output) + ) + .assign_batch_sync(&routing.batch, &routing.fields) + } } pub(super) fn precomputed_partition_bytes( diff --git a/crates/paimon/src/table/table_write_precomputed.rs b/crates/paimon/src/table/table_write_precomputed.rs index 3730d3d0f..3656777f8 100644 --- a/crates/paimon/src/table/table_write_precomputed.rs +++ b/crates/paimon/src/table/table_write_precomputed.rs @@ -63,6 +63,9 @@ impl TableWrite { upper_bound - 1, ))); } + if self.has_reconfigured_append_type() { + return self.write_reconfigured_append(batch, Some(bucket)).await; + } let filtered = self.filter_write_input(batch)?; let input = filtered.as_ref().unwrap_or(batch); let Some(batch) = self.normalize_filtered_write_batch(input)? else { diff --git a/crates/paimon/tests/filtered_write_validation_test.rs b/crates/paimon/tests/filtered_write_validation_test.rs index 4d51a78ed..efced043c 100644 --- a/crates/paimon/tests/filtered_write_validation_test.rs +++ b/crates/paimon/tests/filtered_write_validation_test.rs @@ -364,8 +364,8 @@ async fn completely_filtered_batches_do_not_lock_writer_configuration() { writer.with_blob_consumer(None).unwrap(); writer.with_blob_uri_reader_factory(None).unwrap(); writer.write_arrow_batch(&batch("+I")).await.unwrap(); - assert!(writer.with_write_type(vec!["id".into()]).is_err()); - assert!(writer.with_blob_consumer(None).is_err()); + writer.with_write_type(vec!["id".into()]).unwrap(); + writer.with_blob_consumer(None).unwrap(); assert!(writer.with_blob_uri_reader_factory(None).is_err()); let messages = writer.prepare_commit().await.unwrap(); assert_eq!( @@ -397,3 +397,69 @@ async fn completely_filtered_batches_do_not_lock_writer_configuration() { .iter() .all(|batch| batch.column_by_name("op").unwrap().is_null(0))); } + +#[tokio::test] +async fn reconfigured_append_filters_events_before_nested_value_validation() { + let (value_type, invalid_values) = payload("array", false); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", value_type) + .column("pt", DataType::Int(IntType::new())) + .column("op", DataType::VarChar(VarCharType::string_type())) + .partition_keys(["pt"]) + .option("rowkind.field", "op") + .option("ignore-delete", "true") + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let (io, table) = memory_table( + "memory:/reconfigured-filtered", + TableSchema::new(0, &schema), + ); + setup_dirs(&io, table.location()).await; + persist_table_schema(&io, table.location(), table.schema()).await; + let mut writer = table.new_write_builder().new_write().unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + let input = |ids, values: ArrayRef, partitions, kinds| { + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("value", values), + ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), + ("op", Arc::new(StringArray::from(kinds)) as ArrayRef), + ]) + .unwrap() + }; + writer + .write_arrow_batch(&input( + vec![11], + invalid_values.slice(0, 1), + vec![1], + vec!["+I"], + )) + .await + .unwrap(); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + writer + .write_arrow_batch(&input( + vec![99, 88], + invalid_values, + vec![1, 2], + vec!["-D", "+I"], + )) + .await + .unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + messages + .iter() + .flat_map(|m| &m.new_files) + .map(|f| f.row_count) + .sum::(), + 2 + ); +} diff --git a/crates/paimon/tests/table_partial_write_test.rs b/crates/paimon/tests/table_partial_write_test.rs index fd0c87b32..9c735c62c 100644 --- a/crates/paimon/tests/table_partial_write_test.rs +++ b/crates/paimon/tests/table_partial_write_test.rs @@ -158,8 +158,8 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { .with_write_type(columns.into_iter().map(str::to_string).collect()) .is_err()); } - // Empty batches do not lock the configuration; nonempty ones do, even - // after their files have been prepared. + // Invalid requests preserve the factory configuration. Valid requests + // remain allowed after writing and preparing existing bucket writers. writer .write_arrow_batch(&partial(vec![], vec![])) .await @@ -172,9 +172,9 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { .await .unwrap(); let messages = writer.prepare_commit().await.unwrap(); - assert!(writer + writer .with_write_type(vec!["id".into(), "omitted".into(), "pt".into()]) - .is_err()); + .unwrap(); builder.new_commit().commit(messages).await.unwrap(); assert_eq!(rows(&table).await, vec![vec![Some(1), None, Some(2)]]); } @@ -336,3 +336,636 @@ async fn selected_columns_validate_input_before_staging() { 1 ); } + +fn selected_value(values: Vec, partitions: Vec) -> RecordBatch { + RecordBatch::try_from_iter([ + ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), + ("omitted", Arc::new(Int32Array::from(values)) as ArrayRef), + ]) + .unwrap() +} + +fn full_values(ids: Vec, values: Vec, partitions: Vec) -> RecordBatch { + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), + ("omitted", Arc::new(Int32Array::from(values)) as ArrayRef), + ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), + ]) + .unwrap() +} + +#[tokio::test] +async fn append_write_type_is_captured_per_bucket_across_prepare_rounds() { + for explicit_bucket in [false, true] { + let table = table(!explicit_bucket, explicit_bucket, false).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + let bucket = explicit_bucket.then_some(0); + writer + .write_arrow(&[partial(vec![11], vec![Some(1)])], bucket) + .await + .unwrap(); + // Repeating the same type does not lock an already active writer. + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + writer + .with_write_type(vec!["pt".into(), "omitted".into()]) + .unwrap(); + writer + .write_arrow(&[selected_value(vec![22], vec![2])], bucket) + .await + .unwrap(); + // A selected input cannot be reinterpreted positionally in an old + // bucket. Validation must preserve both its old data and the factory. + assert!(writer + .write_arrow(&[selected_value(vec![999], vec![1])], bucket) + .await + .is_err()); + writer + .write_arrow(&[partial(vec![33], vec![Some(1)])], bucket) + .await + .unwrap(); + let first = writer.prepare_commit().await.unwrap(); + assert!(first + .iter() + .flat_map(|m| &m.new_files) + .any(|f| f.write_cols == Some(vec!["pt".into(), "id".into()]))); + assert!(first + .iter() + .flat_map(|m| &m.new_files) + .any(|f| f.write_cols == Some(vec!["pt".into(), "omitted".into()]))); + builder.new_commit().commit(first).await.unwrap(); + // Finished physical resources are gone, but captured settings remain. + writer + .with_write_type(vec!["pt".into(), "omitted".into()]) + .unwrap(); + writer + .write_arrow(&[partial(vec![44], vec![Some(1)])], bucket) + .await + .unwrap(); + writer + .write_arrow(&[selected_value(vec![55], vec![2])], bucket) + .await + .unwrap(); + writer + .write_arrow(&[selected_value(vec![66], vec![3])], bucket) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + assert_eq!( + rows(&table).await, + vec![ + vec![None, Some(22), Some(2)], + vec![None, Some(55), Some(2)], + vec![None, Some(66), Some(3)], + vec![Some(11), None, Some(1)], + vec![Some(33), None, Some(1)], + vec![Some(44), None, Some(1)], + ] + ); + } +} + +#[tokio::test] +async fn reconfigured_mixed_input_is_preflighted_and_full_input_projects_each_bucket() { + let table = table(true, false, false).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + writer + .write_arrow_batch(&partial(vec![11], vec![Some(1)])) + .await + .unwrap(); + writer + .with_write_type(vec!["pt".into(), "omitted".into()]) + .unwrap(); + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + assert!(writer + .write_arrow_batch(&selected_value(vec![999, 888], vec![1, 2])) + .await + .is_err()); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + writer + .write_arrow_batch(&selected_value(vec![22], vec![2])) + .await + .unwrap(); + writer + .write_arrow_batch(&full_values(vec![33, 44], vec![333, 444], vec![1, 3])) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + // Subsequent valid changes do not alter existing buckets' captured order. + writer + .with_write_type(vec!["id".into(), "pt".into()]) + .unwrap(); + writer + .write_arrow_batch(&full_values(vec![55, 66], vec![555, 666], vec![1, 2])) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + assert_eq!( + rows(&table).await, + vec![ + vec![None, Some(22), Some(2)], + vec![None, Some(444), Some(3)], + vec![None, Some(666), Some(2)], + vec![Some(11), None, Some(1)], + vec![Some(33), None, Some(1)], + vec![Some(55), None, Some(1)], + ] + ); +} + +#[tokio::test] +async fn append_blob_consumer_is_captured_per_partition_across_prepares() { + use paimon::spec::BlobDescriptor; + use std::sync::Mutex; + let table = table(true, false, true).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + let calls = Arc::new(Mutex::new(Vec::new())); + let consumer = |name: &'static str| { + let calls = calls.clone(); + Arc::new(move |_: &str, descriptor: Option<&BlobDescriptor>| { + if descriptor.is_some() { + calls.lock().unwrap().push(name); + } + Ok(false) + }) as Arc + }; + let blob_batch = |partition| { + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef), + ( + "omitted", + Arc::new(arrow_array::LargeBinaryArray::from(vec![ + b"payload".as_slice() + ])) as ArrayRef, + ), + ( + "pt", + Arc::new(Int32Array::from(vec![partition])) as ArrayRef, + ), + ]) + .unwrap() + }; + writer.with_blob_consumer(Some(consumer("first"))).unwrap(); + writer.write_arrow_batch(&blob_batch(1)).await.unwrap(); + writer.with_blob_consumer(Some(consumer("second"))).unwrap(); + writer.write_arrow_batch(&blob_batch(1)).await.unwrap(); + writer.write_arrow_batch(&blob_batch(2)).await.unwrap(); + let first = writer.prepare_commit().await.unwrap(); + builder.new_commit().commit(first).await.unwrap(); + writer.with_blob_consumer(None).unwrap(); + writer.write_arrow_batch(&blob_batch(1)).await.unwrap(); + writer.write_arrow_batch(&blob_batch(2)).await.unwrap(); + writer.write_arrow_batch(&blob_batch(3)).await.unwrap(); + let second = writer.prepare_commit().await.unwrap(); + writer.close().await; + // Closing the writer cannot remove any successfully prepared payload. + builder.new_commit().commit(second).await.unwrap(); + let mut received = calls.lock().unwrap().clone(); + received.sort(); + assert_eq!( + received, + vec!["first", "first", "first", "second", "second"] + ); +} + +#[tokio::test] +async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() { + for fixed in [false, true] { + let table = table(false, fixed, false).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + let routing = |partition| { + RecordBatch::try_from_iter([ + ( + "pt", + Arc::new(Int32Array::from(vec![partition])) as ArrayRef, + ), + ("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef), + ]) + .unwrap() + }; + let names = |fields: Vec| { + fields + .into_iter() + .map(|f| f.name().to_owned()) + .collect::>() + }; + let before = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + assert_eq!( + names(writer.write_fields_for_row(&routing(1)).unwrap()), + vec!["pt", "id"] + ); + assert!(writer.prepare_commit().await.unwrap().is_empty()); + // A query does not create or freeze the partition writer. + writer + .with_write_type(vec!["pt".into(), "omitted".into()]) + .unwrap(); + assert_eq!( + names(writer.write_fields_for_row(&routing(1)).unwrap()), + vec!["pt", "omitted"] + ); + assert_eq!( + table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + writer + .write_arrow_batch(&partial(vec![11], vec![Some(1)])) + .await + .unwrap(); + writer + .with_write_type(vec!["pt".into(), "omitted".into()]) + .unwrap(); + assert_eq!( + names(writer.write_fields_for_row(&routing(1)).unwrap()), + vec!["pt", "id"] + ); + assert_eq!( + names(writer.write_fields_for_row(&routing(2)).unwrap()), + vec!["pt", "omitted"] + ); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + names(writer.write_fields_for_row(&routing(1)).unwrap()), + vec!["pt", "id"] + ); + for invalid in [ + routing(1).slice(0, 0), + partial(vec![1, 2], vec![Some(1), Some(2)]), + ] { + assert!(writer.write_fields_for_row(&invalid).is_err()); + } + let missing_partition = + RecordBatch::try_from_iter([("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef)]) + .unwrap(); + assert!(writer.write_fields_for_row(&missing_partition).is_err()); + // Failed queries leave prepared outputs and future factory settings intact. + assert_eq!( + names(writer.write_fields_for_row(&routing(2)).unwrap()), + vec!["pt", "omitted"] + ); + writer.close().await; + builder.new_commit().commit(messages).await.unwrap(); + assert_eq!(rows(&table).await, vec![vec![Some(11), None, Some(1)]]); + } +} + +#[tokio::test] +async fn unpartitioned_row_field_resolution_accepts_zero_column_routing_row() { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("omitted", DataType::Int(IntType::new())) + .build() + .unwrap(); + let (_, table) = memory_table("memory:/row_fields", TableSchema::new(0, &schema)); + let mut writer = table.new_write_builder().new_write().unwrap(); + let routing = RecordBatch::try_new_with_options( + Arc::new(arrow_schema::Schema::empty()), + vec![], + &arrow_array::RecordBatchOptions::new().with_row_count(Some(1)), + ) + .unwrap(); + assert_eq!(writer.write_fields_for_row(&routing).unwrap().len(), 2); + writer.with_write_type(vec!["id".into()]).unwrap(); + assert_eq!( + writer.write_fields_for_row(&routing).unwrap()[0].name(), + "id" + ); + assert!(writer.prepare_commit().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn captured_row_transport_only_allows_additional_routing_and_event_fields() { + for fixed in [false, true] { + let mut schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("new", DataType::Int(IntType::with_nullable(false))) + .column("unused", DataType::Int(IntType::new())) + .column("pt", DataType::Int(IntType::new())) + .column( + "op", + DataType::VarChar(paimon::spec::VarCharType::string_type()), + ) + .partition_keys(["pt"]) + .option("rowkind.field", "op") + .option("ignore-delete", "true") + .option("data-evolution.enabled", (!fixed).to_string()) + .option("row-tracking.enabled", (!fixed).to_string()); + if fixed { + schema = schema.option("bucket", "4").option("bucket-key", "id"); + } + let (io, table) = memory_table( + "memory:/captured_row_transport", + TableSchema::new(0, &schema.build().unwrap()), + ); + setup_dirs(&io, table.location()).await; + persist_table_schema(&io, table.location(), table.schema()).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.with_write_type(vec!["id".into()]).unwrap(); + let original = RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef), + ("new", Arc::new(Int32Array::from(vec![22])) as ArrayRef), + ("unused", Arc::new(Int32Array::from(vec![33])) as ArrayRef), + ("pt", Arc::new(Int32Array::from(vec![1])) as ArrayRef), + ( + "op", + Arc::new(arrow_array::StringArray::from(vec!["+I"])) as ArrayRef, + ), + ]) + .unwrap(); + writer.write_arrow_batch(&original).await.unwrap(); + writer.with_write_type(vec!["new".into()]).unwrap(); + let old = |partitions: Vec, kinds: Vec<&str>, extra_payload| { + let mut columns = vec![( + "id", + Arc::new(Int32Array::from(vec![11; partitions.len()])) as ArrayRef, + )]; + if extra_payload { + columns.push(( + "unused", + Arc::new(Int32Array::from(vec![99; partitions.len()])) as ArrayRef, + )); + } + columns.push(("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef)); + columns.push(( + "op", + Arc::new(arrow_array::StringArray::from(kinds)) as ArrayRef, + )); + RecordBatch::try_from_iter(columns).unwrap() + }; + let before = io + .list_status_recursive(table.location()) + .await + .unwrap() + .len(); + // The old subset is not valid for the new partition; no group may be staged. + assert!(writer + .write_arrow_batch(&old(vec![1, 2], vec!["+I", "+I"], false)) + .await + .is_err()); + // A non-routing payload cannot be smuggled into a selected layout. + assert!(writer + .write_arrow_batch(&old(vec![1], vec!["+I"], true)) + .await + .is_err()); + assert_eq!( + io.list_status_recursive(table.location()) + .await + .unwrap() + .len(), + before + ); + // Missing the new factory's NOT NULL column is valid for the captured + // old subset. Ignored rows neither create a bucket nor inspect payloads. + writer + .write_arrow_batch(&old(vec![1, 2], vec!["+I", "-D"], false)) + .await + .unwrap(); + let mut columns = vec![]; + if fixed { + columns.push(("id", original.column(0).clone())); + } + columns.extend([ + ("new", original.column(1).clone()), + ("pt", Arc::new(Int32Array::from(vec![2])) as ArrayRef), + ("op", original.column(4).clone()), + ]); + writer + .write_arrow_batch(&RecordBatch::try_from_iter(columns).unwrap()) + .await + .unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + messages + .iter() + .flat_map(|m| &m.new_files) + .map(|f| f.row_count) + .sum::(), + 3 + ); + builder.new_commit().commit(messages).await.unwrap(); + writer + .write_arrow_batch(&old(vec![1], vec!["+I"], false)) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let mut reader = table.new_read_builder(); + reader + .with_projection(&["id", "new", "unused", "pt"]) + .unwrap(); + let plan = reader.new_scan().plan().await.unwrap(); + let batches: Vec = reader + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut actual = Vec::new(); + for batch in batches { + for row in 0..batch.num_rows() { + actual.push( + batch + .columns() + .iter() + .map(|column| { + let array = column.as_any().downcast_ref::().unwrap(); + array.is_valid(row).then(|| array.value(row)) + }) + .collect::>(), + ); + } + } + actual.sort(); + assert_eq!( + actual, + vec![ + vec![None, Some(22), None, Some(2)], + vec![Some(11), None, None, Some(1)], + vec![Some(11), None, None, Some(1)], + vec![Some(11), None, None, Some(1)], + ] + ); + } +} + +#[tokio::test] +async fn historical_row_layouts_do_not_randomly_validate_unwritten_routing_fields() { + for _ in 0..4 { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::with_nullable(false))) + .column("new", DataType::Int(IntType::with_nullable(false))) + .column("pt", DataType::Int(IntType::with_nullable(false))) + .column( + "op", + DataType::VarChar(paimon::spec::VarCharType::string_type()), + ) + .partition_keys(["pt"]) + .option("rowkind.field", "op") + .option("ignore-delete", "true") + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let (io, table) = memory_table("memory:/historical_routing", TableSchema::new(0, &schema)); + setup_dirs(&io, table.location()).await; + persist_table_schema(&io, table.location(), table.schema()).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + let input = |id, partition| { + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![id])) as ArrayRef), + ("new", Arc::new(Int32Array::from(vec![99])) as ArrayRef), + ( + "pt", + Arc::new(Int32Array::from(vec![partition])) as ArrayRef, + ), + ( + "op", + Arc::new(arrow_array::StringArray::from(vec!["+I"])) as ArrayRef, + ), + ]) + .unwrap() + }; + // Omitted routing columns do not impose their physical NOT NULL rule. + writer.with_write_type(vec!["id".into()]).unwrap(); + writer.write_arrow_batch(&input(11, None)).await.unwrap(); + writer + .with_write_type(vec!["id".into(), "pt".into()]) + .unwrap(); + writer.write_arrow_batch(&input(22, Some(1))).await.unwrap(); + writer.with_write_type(vec!["new".into()]).unwrap(); + // id,pt,op can represent either historical type. Actual destinations + // decide whether pt is selected, rather than an arbitrary map entry. + let rows = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from(vec![33, 44, 55, 66])) as ArrayRef, + ), + ( + "pt", + Arc::new(Int32Array::from(vec![None, Some(1), None, Some(1)])) as ArrayRef, + ), + ( + "op", + Arc::new(arrow_array::StringArray::from(vec!["+I", "+I", "-D", "-D"])) as ArrayRef, + ), + ]) + .unwrap(); + let ignored_null_payload = RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![None])) as ArrayRef), + ("pt", Arc::new(Int32Array::from(vec![None])) as ArrayRef), + ( + "op", + Arc::new(arrow_array::StringArray::from(vec!["-D"])) as ArrayRef, + ), + ]) + .unwrap(); + // Logical NOT NULL payload checks still precede ignore-delete. + assert!(writer + .write_arrow_batch(&ignored_null_payload) + .await + .is_err()); + fn require_send(future: T) -> T { + future + } + require_send(writer.write_arrow_batch(&rows)).await.unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let mut reader = table.new_read_builder(); + reader.with_projection(&["id", "pt"]).unwrap(); + let plan = reader.new_scan().plan().await.unwrap(); + let batches: Vec = reader + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut actual = Vec::new(); + for batch in batches { + for row in 0..batch.num_rows() { + actual.push( + batch + .columns() + .iter() + .map(|column| { + let array = column.as_any().downcast_ref::().unwrap(); + array.is_valid(row).then(|| array.value(row)) + }) + .collect::>(), + ); + } + } + actual.sort(); + assert_eq!( + actual, + vec![ + vec![Some(11), None], + vec![Some(22), Some(1)], + vec![Some(33), None], + vec![Some(44), Some(1)] + ] + ); + } +} From e1ff9b5dd9acd41fbe1e80b87c2a22682579836c Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 09:44:50 +0800 Subject: [PATCH 2/4] build: update TLS dependencies for security checks --- Cargo.lock | 8 ++++---- DEPENDENCIES.rust.tsv | 4 ++-- benchmarks/tpcds/DEPENDENCIES.rust.tsv | 4 ++-- bindings/c/DEPENDENCIES.rust.tsv | 4 ++-- bindings/go/DEPENDENCIES.rust.tsv | 4 ++-- bindings/python/DEPENDENCIES.rust.tsv | 4 ++-- crates/integration_tests/DEPENDENCIES.rust.tsv | 4 ++-- crates/integrations/datafusion/DEPENDENCIES.rust.tsv | 4 ++-- crates/paimon-rest-server/DEPENDENCIES.rust.tsv | 4 ++-- crates/paimon/DEPENDENCIES.rust.tsv | 4 ++-- 10 files changed, 22 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d73fbb31d..779c94e79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6115,9 +6115,9 @@ checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "aws-lc-rs", "ring", @@ -8170,9 +8170,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" dependencies = [ "rustls-pki-types", ] diff --git a/DEPENDENCIES.rust.tsv b/DEPENDENCIES.rust.tsv index 2313f8f4c..3a7378bf6 100644 --- a/DEPENDENCIES.rust.tsv +++ b/DEPENDENCIES.rust.tsv @@ -509,7 +509,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -681,7 +681,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X diff --git a/benchmarks/tpcds/DEPENDENCIES.rust.tsv b/benchmarks/tpcds/DEPENDENCIES.rust.tsv index 24915be44..a7bb34556 100644 --- a/benchmarks/tpcds/DEPENDENCIES.rust.tsv +++ b/benchmarks/tpcds/DEPENDENCIES.rust.tsv @@ -333,7 +333,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -433,7 +433,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/c/DEPENDENCIES.rust.tsv b/bindings/c/DEPENDENCIES.rust.tsv index 9dbbc6e93..82566005e 100644 --- a/bindings/c/DEPENDENCIES.rust.tsv +++ b/bindings/c/DEPENDENCIES.rust.tsv @@ -294,7 +294,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -393,7 +393,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/go/DEPENDENCIES.rust.tsv b/bindings/go/DEPENDENCIES.rust.tsv index 9dbbc6e93..82566005e 100644 --- a/bindings/go/DEPENDENCIES.rust.tsv +++ b/bindings/go/DEPENDENCIES.rust.tsv @@ -294,7 +294,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -393,7 +393,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/bindings/python/DEPENDENCIES.rust.tsv b/bindings/python/DEPENDENCIES.rust.tsv index 563899bf6..fd4ef2cc8 100644 --- a/bindings/python/DEPENDENCIES.rust.tsv +++ b/bindings/python/DEPENDENCIES.rust.tsv @@ -433,7 +433,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -565,7 +565,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X diff --git a/crates/integration_tests/DEPENDENCIES.rust.tsv b/crates/integration_tests/DEPENDENCIES.rust.tsv index 4d16dab49..6f7b6affb 100644 --- a/crates/integration_tests/DEPENDENCIES.rust.tsv +++ b/crates/integration_tests/DEPENDENCIES.rust.tsv @@ -264,7 +264,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -357,7 +357,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv index 091e6a231..787cf0b7a 100644 --- a/crates/integrations/datafusion/DEPENDENCIES.rust.tsv +++ b/crates/integrations/datafusion/DEPENDENCIES.rust.tsv @@ -428,7 +428,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -582,7 +582,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi@0.3.9 X X winapi-i686-pc-windows-gnu@0.4.0 X X diff --git a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv index 1ef833b0a..eadd3ccb0 100644 --- a/crates/paimon-rest-server/DEPENDENCIES.rust.tsv +++ b/crates/paimon-rest-server/DEPENDENCIES.rust.tsv @@ -268,7 +268,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -363,7 +363,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X wide@0.7.33 X X X winapi-util@0.1.11 X X windows-core@0.62.2 X X diff --git a/crates/paimon/DEPENDENCIES.rust.tsv b/crates/paimon/DEPENDENCIES.rust.tsv index cceabf938..52f38e7cd 100644 --- a/crates/paimon/DEPENDENCIES.rust.tsv +++ b/crates/paimon/DEPENDENCIES.rust.tsv @@ -427,7 +427,7 @@ rustls-native-certs@0.8.4 X X X rustls-pki-types@1.15.0 X X rustls-platform-verifier@0.7.0 X X rustls-platform-verifier-android@0.1.1 X X -rustls-webpki@0.103.15 X +rustls-webpki@0.103.16 X rustversion@1.0.23 X X ryu@1.0.23 X X safe_arch@0.7.4 X X X @@ -585,7 +585,7 @@ wasm-bindgen-shared@0.2.126 X X wasm-streams@0.5.0 X X web-sys@0.3.103 X X web-time@1.1.0 X X -webpki-root-certs@1.0.9 X +webpki-root-certs@1.0.10 X whoami@2.1.3 X X X wide@0.7.33 X X X winapi@0.3.9 X X From c987478ec8d5212a6c150845576ad76f8fc3e50a Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 09:46:05 +0800 Subject: [PATCH 3/4] fix: align append writer retirement and internal row adapters --- bindings/python/Cargo.toml | 2 +- bindings/python/src/write.rs | 191 ++++++++- bindings/python/tests/test_blob_consumer.py | 39 +- bindings/python/tests/test_write.py | 24 +- crates/paimon/Cargo.toml | 2 + .../paimon/src/table/blob_consumer_tests.rs | 4 +- crates/paimon/src/table/internal.rs | 42 ++ crates/paimon/src/table/mod.rs | 3 + crates/paimon/src/table/table_write.rs | 136 ++++++- crates/paimon/src/table/table_write_append.rs | 62 ++- .../src/table/table_write_lifecycle_tests.rs | 365 ++++++++++++++++++ .../paimon/tests/table_partial_write_test.rs | 29 +- 12 files changed, 835 insertions(+), 64 deletions(-) create mode 100644 crates/paimon/src/table/internal.rs create mode 100644 crates/paimon/src/table/table_write_lifecycle_tests.rs diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml index 733b7b6d3..08063569f 100644 --- a/bindings/python/Cargo.toml +++ b/bindings/python/Cargo.toml @@ -36,7 +36,7 @@ chrono = "0.4.38" datafusion = { workspace = true } datafusion-ffi = { workspace = true } futures = "0.3" -paimon = { path = "../../crates/paimon", features = ["storage-all"] } +paimon = { path = "../../crates/paimon", features = ["storage-all", "internal-api"] } paimon-datafusion = { path = "../../crates/integrations/datafusion", features = ["fulltext"] } paimon-vindex-core = "0.5.0" pyo3 = { version = "0.28", features = ["abi3-py310", "chrono"] } diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index 61130dcbb..dfc33002c 100644 --- a/bindings/python/src/write.rs +++ b/bindings/python/src/write.rs @@ -349,11 +349,12 @@ impl WriteTarget { fn row_write_columns(&self, routing: &RecordBatch) -> paimon::Result> { match self { - Self::Table(writer) => Ok(writer - .write_fields_for_row(routing)? - .iter() - .map(|field| field.name().to_owned()) - .collect()), + Self::Table(writer) => Ok(paimon::table::internal::resolve_row_write_fields( + writer, routing, + )? + .iter() + .map(|field| field.name().to_owned()) + .collect()), Self::PostponeFixed(_) => Err(paimon::Error::Unsupported { message: "_row_write_columns requires a Paimon append table".into(), }), @@ -380,9 +381,15 @@ impl WriteTarget { } } - async fn prepare_commit(&mut self) -> paimon::Result> { + async fn prepare_commit( + &mut self, + commit_identifier: i64, + ) -> paimon::Result> { match self { - Self::Table(writer) => writer.prepare_commit().await, + Self::Table(writer) => { + paimon::table::internal::prepare_commit_with_identifier(writer, commit_identifier) + .await + } Self::PostponeFixed(writer) => writer.prepare_commit().await, } } @@ -896,14 +903,23 @@ impl WriteState { .map_err(to_py_err)?; // Release replaced callbacks which no bucket has captured. A snapshot // retained by core contributes another Arc, without another Py ref. - self.blob_consumers - .retain(|consumer| Arc::strong_count(consumer) > 1); + self.prune_blob_consumers(); if let Some(consumer) = consumer { self.blob_consumers.push(consumer); } Ok(()) } + fn prune_blob_consumers(&mut self) { + let (retained, retired): (Vec<_>, Vec<_>) = std::mem::take(&mut self.blob_consumers) + .into_iter() + .partition(|consumer| Arc::strong_count(consumer) > 1); + self.blob_consumers = retained; + // Python destructors can re-enter Native. No callback mutex is held, + // and the retained references are restored before any destructor runs. + drop(retired); + } + fn traverse_callbacks(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { for consumer in &self.blob_consumers { visit.call(&*consumer.callback.lock().unwrap())?; @@ -957,14 +973,18 @@ impl WriteState { self.write_result(result) } - fn prepare_commit(&mut self, py: Python<'_>) -> PyResult> { + fn prepare_commit( + &mut self, + py: Python<'_>, + commit_identifier: i64, + ) -> PyResult> { let inner = self .inner .as_mut() .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?; - let messages = py - .detach(|| runtime().block_on(inner.prepare_commit())) - .map_err(to_callback_error)?; + let result = py.detach(|| runtime().block_on(inner.prepare_commit(commit_identifier))); + self.prune_blob_consumers(); + let messages = result.map_err(to_callback_error)?; Ok(wrap_messages( messages, &self.table_location, @@ -1334,7 +1354,7 @@ impl PyBatchTableWrite { )); } self.prepared = true; - self.state.prepare_commit(py) + self.state.prepare_commit(py, i64::MAX) } } @@ -1407,15 +1427,15 @@ impl PyStreamTableWrite { } /// Rust currently flushes synchronously and has no background compaction. - /// The identifier accompanies the returned messages in StreamTableCommit.commit. + /// The identifier also bounds retirement of captured bucket configuration. fn prepare_commit( &mut self, py: Python<'_>, wait_compaction: bool, commit_identifier: i64, ) -> PyResult> { - let _ = (wait_compaction, commit_identifier); - self.state.prepare_commit(py) + let _ = wait_compaction; + self.state.prepare_commit(py, commit_identifier) } } @@ -1872,6 +1892,141 @@ assert sorted(finalized) == ['first', 'second'], finalized }); } + #[test] + fn idle_bucket_prepare_releases_only_retired_python_blob_consumers() { + Python::attach(|py| { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("pt", DataType::Int(IntType::new())) + .partition_keys(["pt"]) + .option("bucket", "-1") + .build() + .unwrap(); + let table = Arc::new(Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "retired_callbacks"), + "memory:/retired_callbacks".into(), + TableSchema::new(0, &schema), + None, + )); + let writer = Py::new( + py, + PyStreamTableWrite { + state: WriteContext::new(table.clone()).new_write(false).unwrap(), + }, + ) + .unwrap(); + let locals = PyDict::new(py); + locals.set_item("writer", &writer).unwrap(); + py.run( + c_str!( + r#" +import gc +import weakref +finalized = [] +class Callback: + def __init__(self, label): + self.label = label + def __call__(self, field, descriptor): + return False + def __del__(self): + finalized.append(self.label) +def install(label): + callback = Callback(label) + writer.with_blob_consumer(callback) + return weakref.ref(callback) +first = install('first') +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let write_partition = |partition| { + let input = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(arrow::array::Int32Array::from(vec![1])) as arrow::array::ArrayRef, + ), + ( + "pt", + Arc::new(arrow::array::Int32Array::from(vec![partition])) + as arrow::array::ArrayRef, + ), + ]) + .unwrap(); + let mut borrowed = writer.borrow_mut(py); + let inner = borrowed.state.inner.as_mut().unwrap(); + py.detach(|| runtime().block_on(inner.write_arrow(&input, None))) + .unwrap(); + }; + write_partition(1); + py.run( + c_str!("second = install('second')"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + write_partition(2); + py.run( + c_str!("current = install('current')"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + assert_eq!( + writer + .borrow_mut(py) + .state + .prepare_commit(py, 1) + .unwrap() + .len(), + 2 + ); + let prepared_files = py + .detach(|| { + runtime().block_on(table.file_io().list_status_recursive(table.location())) + }) + .unwrap(); + assert!(!prepared_files.is_empty()); + py.run( + c_str!("gc.collect(); assert first() is not None and second() is not None and current() is not None"), + Some(&locals), Some(&locals), + ).unwrap(); + write_partition(2); + assert_eq!( + writer + .borrow_mut(py) + .state + .prepare_commit(py, 2) + .unwrap() + .len(), + 1 + ); + py.run( + c_str!("gc.collect(); assert first() is None; assert second() is not None and current() is not None; assert finalized == ['first']"), + Some(&locals), Some(&locals), + ).unwrap(); + for file in prepared_files.iter().filter(|file| !file.is_dir) { + assert!(py + .detach(|| runtime().block_on(table.file_io().exists(&file.path))) + .unwrap()); + } + writer.borrow_mut(py).close(py); + py.run( + c_str!("gc.collect(); assert second() is None and current() is None"), + Some(&locals), + Some(&locals), + ) + .unwrap(); + for file in prepared_files.iter().filter(|file| !file.is_dir) { + assert!(py + .detach(|| runtime().block_on(table.file_io().exists(&file.path))) + .unwrap()); + } + }); + } + #[test] fn merge_table_source_conditions_resolve_read_metadata() { Python::attach(|py| { @@ -2393,7 +2548,7 @@ batch = pa.RecordBatch.from_pylist([ } else { result.unwrap(); } - assert!(state.prepare_commit(py).unwrap().is_empty()); + assert!(state.prepare_commit(py, i64::MAX).unwrap().is_empty()); } } }); diff --git a/bindings/python/tests/test_blob_consumer.py b/bindings/python/tests/test_blob_consumer.py index d2e8f068c..d86a2ee82 100644 --- a/bindings/python/tests/test_blob_consumer.py +++ b/bindings/python/tests/test_blob_consumer.py @@ -24,20 +24,25 @@ from pypaimon_rust.datafusion import PaimonCatalog, SQLContext -def _table(tmp_path): +def _table(tmp_path, partitioned=False): context = SQLContext() context.register_catalog('paimon', {'warehouse': str(tmp_path)}) context.sql('CREATE SCHEMA paimon.consumer') - context.sql("""CREATE TABLE paimon.consumer.t (id INT, payload BLOB) WITH ( + partition = ', pt INT) PARTITIONED BY (pt)' if partitioned else ')' + context.sql("""CREATE TABLE paimon.consumer.t (id INT, payload BLOB{} WITH ( 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true', - 'blob.copy-buffer-size' = '2 B')""") + 'blob.copy-buffer-size' = '2 B')""".format(partition)) return context, PaimonCatalog({'warehouse': str(tmp_path)}).get_table('consumer.t') -def _batch(): - return pa.record_batch([pa.array([1, 2, 3], type=pa.int32()), - pa.array([b'payload', None, b''], type=pa.large_binary())], - names=['id', 'payload']) +def _batch(partition=None, first_id=1): + columns = [pa.array(range(first_id, first_id + 3), type=pa.int32()), + pa.array([b'payload', None, b''], type=pa.large_binary())] + names = ['id', 'payload'] + if partition is not None: + columns.append(pa.array([partition] * 3, type=pa.int32())) + names.append('pt') + return pa.record_batch(columns, names=names) def _descriptor_payload(encoded): @@ -54,7 +59,7 @@ def _descriptor_payload(encoded): @pytest.mark.parametrize('stream', [False, True]) def test_blob_consumer_descriptor_bytes_and_nulls(tmp_path, stream): - context, table = _table(tmp_path) + context, table = _table(tmp_path, partitioned=True) builder = table.new_stream_write_builder() if stream else table.new_batch_write_builder() writer = builder.new_write() received = [] @@ -65,18 +70,22 @@ def callback(name, encoded): try: assert writer.with_blob_consumer(callback) is writer - writer.write_arrow(_batch()) - with pytest.raises(ValueError, match='before any write'): - writer.with_blob_consumer(None) + writer.write_arrow(_batch(1)) + assert writer.with_blob_consumer(None) is writer + # The first partition retains its callback, while the new partition + # captures the cleared factory setting. + writer.write_arrow(_batch(1, 4)) + writer.write_arrow(_batch(2, 7)) messages = writer.prepare_commit(True, 1) if stream else writer.prepare_commit() commit = builder.new_commit() commit.commit(1, messages) if stream else commit.commit(messages) - assert [name for name, _ in received] == ['payload'] * 3 + assert [name for name, _ in received] == ['payload'] * 6 assert [_descriptor_payload(encoded) if encoded is not None else None - for _, encoded in received] == [b'payload', None, b''] + for _, encoded in received] == [b'payload', None, b''] * 2 assert pa.Table.from_batches(context.sql( - 'SELECT id, payload FROM paimon.consumer.t ORDER BY id')).to_pydict() == { - 'id': [1, 2, 3], 'payload': [b'payload', None, b'']} + 'SELECT id, payload, pt FROM paimon.consumer.t ORDER BY id')).to_pydict() == { + 'id': list(range(1, 10)), 'payload': [b'payload', None, b''] * 3, + 'pt': [1] * 6 + [2] * 3} finally: writer.close() with pytest.raises(RuntimeError, match='closed'): diff --git a/bindings/python/tests/test_write.py b/bindings/python/tests/test_write.py index a915f3140..289ed4b6c 100644 --- a/bindings/python/tests/test_write.py +++ b/bindings/python/tests/test_write.py @@ -66,27 +66,33 @@ def test_write_commit_read_roundtrip(): @pytest.mark.parametrize("streaming", [False, True]) def test_partial_write_type_roundtrip_and_lifecycle(tmp_path, streaming): - ctx = _make_empty_table(str(tmp_path)) + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": str(tmp_path)}) + ctx.sql("CREATE SCHEMA paimon.wdb") + ctx.sql("CREATE TABLE paimon.wdb.t (id INT, name STRING, pt STRING) PARTITIONED BY (pt)") table = _get_table(str(tmp_path)) builder = table.new_stream_write_builder() if streaming else table.new_batch_write_builder() writer = builder.new_write() - assert writer.with_write_type(["id"]) is writer - writer.write_arrow(pa.record_batch([pa.array([], type=pa.int32())], names=["id"])) - assert writer.with_write_type(["id"]) is writer + assert writer.with_write_type(["id", "pt"]) is writer + empty = pa.record_batch( + [pa.array([], type=pa.int32()), pa.array([], type=pa.string())], names=["id", "pt"]) + writer.write_arrow(empty) + assert writer.with_write_type(["id", "pt"]) is writer for columns in [[], ["missing"], ["id", "id"]]: with pytest.raises(ValueError): writer.with_write_type(columns) - writer.write_arrow(pa.record_batch([pa.array([1, 2], type=pa.int32())], names=["id"])) - with pytest.raises(ValueError, match="before any data"): - writer.with_write_type(["id", "name"]) + schema = pa.schema([("id", pa.int32()), ("name", pa.string()), ("pt", pa.string())]) + writer.write_arrow(pa.record_batch([[1, 2], ["ignored", "ignored"], ["old", "old"]], schema=schema)) + assert writer.with_write_type(["name", "pt"]) is writer + writer.write_arrow(pa.record_batch([[3, 4], ["ignored", "new"], ["old", "new"]], schema=schema)) messages = writer.prepare_commit(True, 7) if streaming else writer.prepare_commit() commit = builder.new_commit() commit.commit(7, messages) if streaming else commit.commit(messages) writer.close() with pytest.raises(RuntimeError, match="closed"): writer.with_write_type(["id"]) - assert pa.Table.from_batches(ctx.sql("SELECT id, name FROM paimon.wdb.t ORDER BY id")).to_pydict() == { - "id": [1, 2], "name": [None, None]} + assert pa.Table.from_batches(ctx.sql("SELECT id, name, pt FROM paimon.wdb.t ORDER BY pt, id")).to_pydict() == { + "id": [None, 1, 2, 3], "name": ["new", None, None, None], "pt": ["new", "old", "old", "old"]} @pytest.mark.parametrize("streaming", [False, True]) diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index 34644e6c1..b680891fa 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -45,6 +45,8 @@ storage-all = [ "storage-oss-cpp", ] fulltext = ["dep:paimon-ftindex-core"] +# Unstable bridge used by the repository's language bindings. +internal-api = [] vortex = ["dep:vortex"] storage-memory = ["opendal/services-memory"] diff --git a/crates/paimon/src/table/blob_consumer_tests.rs b/crates/paimon/src/table/blob_consumer_tests.rs index 9c3ef0a5b..d11b28eec 100644 --- a/crates/paimon/src/table/blob_consumer_tests.rs +++ b/crates/paimon/src/table/blob_consumer_tests.rs @@ -94,7 +94,9 @@ async fn append_consumer_survives_file_rolls_and_checkpoint_recreation() { .sum::(), 2 ); - assert!(writer.with_blob_consumer(None).is_err()); + // Java changes the factory setting while an existing bucket keeps the + // consumer it captured, including across consecutive prepares. + writer.with_blob_consumer(None).unwrap(); } writer.close().await; let events = events.lock().unwrap().clone(); diff --git a/crates/paimon/src/table/internal.rs b/crates/paimon/src/table/internal.rs new file mode 100644 index 000000000..4d17261b7 --- /dev/null +++ b/crates/paimon/src/table/internal.rs @@ -0,0 +1,42 @@ +// 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. + +//! Unstable adapters for the repository's language bindings. +//! These operations are available only with `internal-api`. + +use super::{CommitMessage, TableWrite}; +use crate::spec::DataField; +use arrow_array::RecordBatch; + +/// Inspect only routing metadata before a binding encodes a named row. +/// This does not create a writer or capture a new bucket configuration. +pub fn resolve_row_write_fields( + writer: &TableWrite, + routing_batch: &RecordBatch, +) -> crate::Result> { + writer.resolve_row_write_fields(routing_batch) +} + +/// Forward an identifier already supplied by the batch or stream facade. +pub async fn prepare_commit_with_identifier( + writer: &mut TableWrite, + commit_identifier: i64, +) -> crate::Result> { + writer + .prepare_commit_with_identifier(commit_identifier) + .await +} diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index a8b35eb12..e42b559a8 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -137,6 +137,9 @@ mod query_auth; mod system_table; pub use system_table::SystemRead; +#[cfg(feature = "internal-api")] +#[doc(hidden)] +pub mod internal; mod key_normalization; mod read_builder; mod read_execution; diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 9dfd77503..891c7623a 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -56,9 +56,18 @@ use std::sync::Arc; mod append; #[path = "table_write_input.rs"] mod input; +#[cfg(test)] +#[path = "table_write_lifecycle_tests.rs"] +mod lifecycle_tests; #[path = "table_write_precomputed.rs"] mod precomputed; +#[derive(Clone, Copy)] +enum AppendPrepareIdentifier { + Known(i64), + Unknown, +} + /// Java's append writer captures its factory settings per partition/bucket. /// Keep only these settings across prepares, not completed file resources. #[derive(Clone)] @@ -67,6 +76,7 @@ struct AppendWriteConfig { fields: Vec, blob_consumer: Option>, blob_uri_reader_factory: Option>, + last_modified_prepare: Option, } /// Enum to hold either an append-only writer, a key-value writer, or a postpone writer. @@ -660,7 +670,8 @@ impl TableWrite { /// Selected inputs must include partition columns and, unless a bucket is /// supplied through [`Self::write_arrow`], the columns used for hashing. /// Changes apply to partition/bucket writers created afterwards. Existing - /// writers retain their selected columns, including across prepares. A full + /// writers retain their selected columns across active prepares, until an + /// idle prepare can safely retire them. A full /// table batch is projected to each writer's captured columns; a selected /// batch must match the destination writer's names and order, optionally /// accompanied by routing and configured row-kind fields. @@ -1486,19 +1497,70 @@ impl TableWrite { for (_, mut writer) in self.partition_writers.drain() { writer.abort().await; } + self.append_write_configs.clear(); + self.blob_consumer = None; + self.blob_uri_reader_factory = None; } /// Flush all writers and collect CommitMessages for use with TableCommit. /// Append writers are closed; primary-key writers retain their sequence /// counters after handing off the prepared files, as Java does. + /// Captured append settings survive active prepares. Empty prepares retire + /// idle bucket-unaware writers. This generic operation has no checkpoint + /// identifier, so it conservatively retains fixed-bucket settings because + /// it cannot prove the previous output has committed. /// /// Invalidate the snapshot sequence cache used to seed previously unseen /// buckets. Already-open key-value writers keep their own advancing counter. pub async fn prepare_commit(&mut self) -> Result> { + self.prepare_commit_internal(None).await + } + + #[cfg(any(test, feature = "internal-api"))] + pub(crate) async fn prepare_commit_with_identifier( + &mut self, + commit_identifier: i64, + ) -> Result> { + self.prepare_commit_internal(Some(commit_identifier)).await + } + + async fn prepare_commit_internal( + &mut self, + commit_identifier: Option, + ) -> Result> { if let Some(writer) = self.format_writer.as_mut() { return writer.prepare_commit().await; } self.ensure_active()?; + // Java keeps newly created writers on the first prepare. Later empty + // increments may retire them, but active buckets keep their captured + // factory configuration across successful prepares. + let can_retire = self.append_write_configs.values().any(|config| { + matches!( + config.last_modified_prepare, + Some(AppendPrepareIdentifier::Unknown) + ) || matches!( + config.last_modified_prepare, + Some(AppendPrepareIdentifier::Known(id)) if id != i64::MIN + ) + }); + let fixed_bucket = CoreOptions::new(self.table.schema().options()).bucket() >= 0; + let latest_committed_identifier = if can_retire + && fixed_bucket + && commit_identifier.is_some() + && self.append_write_configs.iter().any(|(key, config)| { + !self.partition_writers.contains_key(key) + && matches!( + config.last_modified_prepare, + Some(AppendPrepareIdentifier::Known(_)) + ) + }) { + // Resolve before draining files: a metadata failure leaves pending + // output available for a retry, with no prepared files lost. + Self::latest_committed_append_identifier(&self.table, &self.commit_user).await? + } else { + None + }; self.partition_seq_cache.clear(); let writers: Vec<(PartitionBucketKey, FileWriter)> = self.partition_writers.drain().collect(); @@ -1560,21 +1622,81 @@ impl TableWrite { messages.push(msg); } } + let modified = messages + .iter() + .map(|message| (message.partition.clone(), message.bucket)) + .collect::>(); + self.append_write_configs.retain(|key, config| { + if modified.contains(key) { + config.last_modified_prepare = Some(match commit_identifier { + Some(id) => AppendPrepareIdentifier::Known(id), + None => AppendPrepareIdentifier::Unknown, + }); + return true; + } + if !can_retire { + return true; + } + if !fixed_bucket { + return false; + } + // The generic core prepare has no checkpoint identifier, so it + // cannot prove a fixed bucket's previous output has committed. + match ( + commit_identifier, + config.last_modified_prepare, + latest_committed_identifier, + ) { + (Some(_), Some(AppendPrepareIdentifier::Known(last)), Some(latest)) => { + last >= latest + } + _ => true, + } + }); Ok(messages) } + async fn latest_committed_append_identifier( + table: &Table, + commit_user: &str, + ) -> Result> { + // FileSystemWriteRestore deliberately reads filesystem snapshots, even + // for catalog-managed tables, and stops at the first expired snapshot. + let manager = table.snapshot_manager().with_rest_env(None); + let Some(mut id) = manager.get_latest_snapshot_id().await? else { + return Ok(None); + }; + while id >= 1 { + let snapshot = match manager.get_snapshot(id).await { + Ok(snapshot) => snapshot, + Err(crate::Error::SnapshotNotExist { .. }) => break, + Err(crate::Error::IoUnexpected { source, .. }) + if source.kind() == opendal::ErrorKind::NotFound => + { + break; + } + Err(error) => return Err(error), + }; + if snapshot.commit_user() == commit_user { + return Ok(Some(snapshot.commit_identifier())); + } + id -= 1; + } + Ok(None) + } + async fn fail_prepare(&mut self, messages: &[CommitMessage]) { + let preserve_blob_files = self.primary_key_indices.is_empty() + && self + .append_write_configs + .values() + .any(|config| config.blob_consumer.is_some()); self.failed = true; self.close().await; // Only this failed invocation's outputs: no message has been returned // or submitted. Earlier successful prepare calls belong to the caller. let mut cleanup = messages.to_vec(); - if self.primary_key_indices.is_empty() - && self - .append_write_configs - .values() - .any(|config| config.blob_consumer.is_some()) - { + if preserve_blob_files { // A successful partition may already have exposed its dedicated // Blob files to the consumer. Preserve Java's deletion policy in // this table-wide cleanup too. Explicit TableCommit.abort remains diff --git a/crates/paimon/src/table/table_write_append.rs b/crates/paimon/src/table/table_write_append.rs index aca44f9c0..9163c7f35 100644 --- a/crates/paimon/src/table/table_write_append.rs +++ b/crates/paimon/src/table/table_write_append.rs @@ -24,16 +24,20 @@ impl TableWrite { /// creating files or changing bucket assignment state. Row adapters can /// encode only these columns, leaving unused payloads untouched. The input /// must contain exactly one row and the table's routing columns. - pub fn write_fields_for_row(&self, routing_batch: &RecordBatch) -> Result> { + #[cfg(any(test, feature = "internal-api"))] + pub(crate) fn resolve_row_write_fields( + &self, + routing_batch: &RecordBatch, + ) -> Result> { self.ensure_active()?; if self.format_writer.is_some() || !self.primary_key_indices.is_empty() { return Err(crate::Error::Unsupported { - message: "write_fields_for_row requires a Paimon append table".into(), + message: "Row write field resolution requires a Paimon append table".into(), }); } if routing_batch.num_rows() != 1 { return Err(crate::Error::DataInvalid { - message: "write_fields_for_row requires exactly one routing row".into(), + message: "Row write field resolution requires exactly one routing row".into(), source: None, }); } @@ -69,6 +73,7 @@ impl TableWrite { fields: self.write_fields.clone(), blob_consumer: self.blob_consumer.clone(), blob_uri_reader_factory: self.blob_uri_reader_factory.clone(), + last_modified_prepare: None, } } @@ -281,3 +286,54 @@ fn is_routing_transport( || routing_names.contains(field.name()) }) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::Identifier; + use crate::io::FileIOBuilder; + use crate::spec::{DataType, IntType, Schema, TableSchema}; + + #[test] + fn row_write_field_resolution_does_not_capture_a_writer() { + let schema = Schema::builder() + .column("first", DataType::Int(IntType::new())) + .column("second", DataType::Int(IntType::new())) + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "row_fields"), + "memory:/row_fields".into(), + TableSchema::new(0, &schema), + None, + ); + let mut writer = TableWrite::new(&table, "row-fields".into()).unwrap(); + let routing = |rows| { + RecordBatch::try_new_with_options( + Arc::new(arrow_schema::Schema::empty()), + vec![], + &arrow_array::RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .unwrap() + }; + assert_eq!( + writer.resolve_row_write_fields(&routing(1)).unwrap().len(), + 2 + ); + writer.with_write_type(vec!["first".into()]).unwrap(); + assert_eq!( + writer.resolve_row_write_fields(&routing(1)).unwrap()[0].name(), + "first" + ); + writer.with_write_type(vec!["second".into()]).unwrap(); + assert_eq!( + writer.resolve_row_write_fields(&routing(1)).unwrap()[0].name(), + "second" + ); + assert!(writer.resolve_row_write_fields(&routing(0)).is_err()); + assert!(writer.resolve_row_write_fields(&routing(2)).is_err()); + assert!(writer.partition_writers.is_empty()); + assert!(writer.append_write_configs.is_empty()); + } +} diff --git a/crates/paimon/src/table/table_write_lifecycle_tests.rs b/crates/paimon/src/table/table_write_lifecycle_tests.rs new file mode 100644 index 000000000..b0ce2da4c --- /dev/null +++ b/crates/paimon/src/table/table_write_lifecycle_tests.rs @@ -0,0 +1,365 @@ +// 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::*; +use crate::catalog::Identifier; +use crate::io::FileIOBuilder; +use crate::spec::{BlobDescriptor, BlobType, CommitKind, IntType, Schema, TableSchema}; +use arrow_array::{Int32Array, LargeBinaryArray}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Mutex; + +async fn table(fixed: bool, blob: bool) -> Table { + let io = FileIOBuilder::new("memory").build().unwrap(); + let mut schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "value", + if blob { + DataType::Blob(BlobType::new()) + } else { + DataType::Int(IntType::new()) + }, + ) + .column("pt", DataType::Int(IntType::new())) + .partition_keys(["pt"]); + if fixed { + schema = schema.option("bucket", "2").option("bucket-key", "id"); + } else { + schema = schema + .option("row-tracking.enabled", "true") + .option("data-evolution.enabled", "true"); + } + let table = Table::new( + io, + Identifier::new("default", "append_lifecycle"), + "memory:/append_lifecycle".into(), + TableSchema::new(0, &schema.build().unwrap()), + None, + ); + table + .file_io() + .new_output(&table.schema_manager().schema_path(0)) + .unwrap() + .write(serde_json::to_vec(table.schema()).unwrap().into()) + .await + .unwrap(); + table +} + +fn batch(partition: i32, blob: bool) -> RecordBatch { + let value: ArrayRef = if blob { + Arc::new(LargeBinaryArray::from(vec![b"payload".as_slice()])) + } else { + Arc::new(Int32Array::from(vec![20])) + }; + RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![10])) as ArrayRef), + ("value", value), + ( + "pt", + Arc::new(Int32Array::from(vec![partition])) as ArrayRef, + ), + ]) + .unwrap() +} + +async fn write(writer: &mut TableWrite, fixed: bool, partition: i32) { + writer + .write_arrow(&[batch(partition, false)], fixed.then_some(0)) + .await + .unwrap(); +} + +fn assert_columns(messages: &[CommitMessage], column: &str) { + assert_eq!(messages.len(), 1); + assert!(!messages[0].new_files.is_empty()); + for file in &messages[0].new_files { + assert_eq!(file.write_cols, Some(vec!["pt".into(), column.into()])); + assert_eq!(file.row_count, 1); + } +} + +async fn snapshot(table: &Table, id: i64, user: &str, identifier: i64) { + let snapshot = Snapshot::builder() + .version(3) + .id(id) + .schema_id(0) + .base_manifest_list(String::new()) + .delta_manifest_list(String::new()) + .commit_user(user.into()) + .commit_identifier(identifier) + .commit_kind(CommitKind::APPEND) + .time_millis(0) + .build(); + let manager = table.snapshot_manager(); + table + .file_io() + .new_output(&manager.snapshot_path(id)) + .unwrap() + .write(serde_json::to_vec(&snapshot).unwrap().into()) + .await + .unwrap(); + manager.write_latest_hint(id).await.unwrap(); +} + +#[tokio::test] +async fn unaware_append_recaptures_type_only_after_an_idle_prepare() { + for idle_prepares in 0..=2 { + let table = table(false, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, false, 1).await; + assert_columns(&writer.prepare_commit().await.unwrap(), "id"); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + for _ in 0..idle_prepares { + assert!(writer.prepare_commit().await.unwrap().is_empty()); + } + write(&mut writer, false, 1).await; + assert_columns( + &writer.prepare_commit().await.unwrap(), + if idle_prepares == 0 { "id" } else { "value" }, + ); + } +} + +#[tokio::test] +async fn empty_prepares_do_not_capture_append_factory_settings() { + let table = table(false, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + for _ in 0..2 { + assert!(writer.prepare_commit().await.unwrap().is_empty()); + } + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + write(&mut writer, false, 1).await; + assert_columns(&writer.prepare_commit().await.unwrap(), "value"); +} + +#[tokio::test] +async fn idle_append_retirement_recaptures_consumers_and_close_releases_them() { + for idle_prepares in 0..=2 { + let table = table(false, true).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + let old_calls = Arc::new(AtomicUsize::new(0)); + let new_calls = Arc::new(AtomicUsize::new(0)); + let consumer = |calls: Arc| { + Arc::new(move |_: &str, descriptor: Option<&BlobDescriptor>| { + if descriptor.is_some() { + calls.fetch_add(1, Ordering::Relaxed); + } + Ok(false) + }) as Arc + }; + let old = consumer(old_calls.clone()); + let old_ref = Arc::downgrade(&old); + writer.with_blob_consumer(Some(old)).unwrap(); + writer.write_arrow_batch(&batch(1, true)).await.unwrap(); + writer.prepare_commit().await.unwrap(); + let new = consumer(new_calls.clone()); + let new_ref = Arc::downgrade(&new); + writer.with_blob_consumer(Some(new)).unwrap(); + for _ in 0..idle_prepares { + assert!(writer.prepare_commit().await.unwrap().is_empty()); + } + assert_eq!(old_ref.upgrade().is_some(), idle_prepares == 0); + writer.write_arrow_batch(&batch(1, true)).await.unwrap(); + writer.prepare_commit().await.unwrap(); + assert_eq!( + old_calls.load(Ordering::Relaxed), + if idle_prepares == 0 { 2 } else { 1 } + ); + assert_eq!( + new_calls.load(Ordering::Relaxed), + usize::from(idle_prepares != 0) + ); + writer.close().await; + assert!(old_ref.upgrade().is_none()); + assert!(new_ref.upgrade().is_none()); + } +} + +#[tokio::test] +async fn fixed_append_waits_for_a_strictly_newer_committed_identifier_of_its_user() { + let table = table(true, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, true, 1).await; + assert_columns( + &writer.prepare_commit_with_identifier(1).await.unwrap(), + "id", + ); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + for (id, user, identifier) in [(1, "owner", 1), (2, "owner", 1), (3, "other", 99)] { + snapshot(&table, id, user, identifier).await; + assert!(writer + .prepare_commit_with_identifier(2) + .await + .unwrap() + .is_empty()); + assert_eq!(writer.append_write_configs.len(), 1); + } + snapshot(&table, 4, "owner", 2).await; + assert!(writer + .prepare_commit_with_identifier(3) + .await + .unwrap() + .is_empty()); + write(&mut writer, true, 1).await; + assert_columns( + &writer.prepare_commit_with_identifier(3).await.unwrap(), + "value", + ); +} + +#[tokio::test] +async fn fixed_append_keeps_unknown_prepare_identifiers() { + for initial_identifier in [None, Some(1)] { + let table = table(true, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, true, 1).await; + writer + .prepare_commit_internal(initial_identifier) + .await + .unwrap(); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + snapshot(&table, 1, "owner", 2).await; + assert!(writer.prepare_commit().await.unwrap().is_empty()); + if initial_identifier.is_none() { + assert!(writer + .prepare_commit_with_identifier(3) + .await + .unwrap() + .is_empty()); + } + write(&mut writer, true, 1).await; + assert_columns(&writer.prepare_commit().await.unwrap(), "id"); + } +} + +#[tokio::test] +async fn fixed_append_stops_at_an_expired_snapshot_gap() { + let table = table(true, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, true, 1).await; + writer.prepare_commit_with_identifier(1).await.unwrap(); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + snapshot(&table, 1, "owner", 2).await; + snapshot(&table, 3, "other", 99).await; + assert!(writer + .prepare_commit_with_identifier(3) + .await + .unwrap() + .is_empty()); + write(&mut writer, true, 1).await; + assert_columns( + &writer.prepare_commit_with_identifier(3).await.unwrap(), + "id", + ); +} + +#[tokio::test] +async fn fixed_append_metadata_errors_leave_pending_output_available() { + let table = table(true, false).await; + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, true, 1).await; + writer.prepare_commit_with_identifier(1).await.unwrap(); + writer + .with_write_type(vec!["pt".into(), "value".into()]) + .unwrap(); + write(&mut writer, true, 2).await; + let manager = table.snapshot_manager(); + table + .file_io() + .new_output(&manager.snapshot_path(1)) + .unwrap() + .write(bytes::Bytes::from_static(b"invalid snapshot JSON")) + .await + .unwrap(); + manager.write_latest_hint(1).await.unwrap(); + assert!(matches!( + writer.prepare_commit_with_identifier(2).await, + Err(crate::Error::DataInvalid { .. }) + )); + snapshot(&table, 1, "owner", 1).await; + assert_columns( + &writer.prepare_commit_with_identifier(2).await.unwrap(), + "value", + ); +} + +#[tokio::test] +async fn failed_prepare_preserves_exposed_blob_files_after_releasing_contexts() { + let table = table(false, true).await; + let exposed = Arc::new(Mutex::new(Vec::new())); + let captured = exposed.clone(); + let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); + writer + .with_blob_consumer(Some(Arc::new( + move |_: &str, descriptor: Option<&BlobDescriptor>| { + if let Some(descriptor) = descriptor { + captured.lock().unwrap().push(descriptor.uri().to_string()); + } + Ok(false) + }, + ))) + .unwrap(); + writer.write_arrow_batch(&batch(1, true)).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + // Model the successful partition outputs collected before another + // partition's prepare failed. These files have not been handed off. + writer.fail_prepare(&messages).await; + let exposed = exposed.lock().unwrap().clone(); + assert!(!exposed.is_empty()); + for path in exposed { + assert!(table.file_io().exists(&path).await.unwrap()); + } + assert!(writer.append_write_configs.is_empty()); + assert!(writer.blob_consumer.is_none()); + assert!(!table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap() + .iter() + .any(|file| file.path.ends_with(".parquet"))); +} diff --git a/crates/paimon/tests/table_partial_write_test.rs b/crates/paimon/tests/table_partial_write_test.rs index 9c735c62c..5030d90a7 100644 --- a/crates/paimon/tests/table_partial_write_test.rs +++ b/crates/paimon/tests/table_partial_write_test.rs @@ -563,6 +563,7 @@ async fn append_blob_consumer_is_captured_per_partition_across_prepares() { ); } +#[cfg(feature = "internal-api")] #[tokio::test] async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() { for fixed in [false, true] { @@ -595,7 +596,7 @@ async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() .unwrap() .len(); assert_eq!( - names(writer.write_fields_for_row(&routing(1)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), vec!["pt", "id"] ); assert!(writer.prepare_commit().await.unwrap().is_empty()); @@ -604,7 +605,7 @@ async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() .with_write_type(vec!["pt".into(), "omitted".into()]) .unwrap(); assert_eq!( - names(writer.write_fields_for_row(&routing(1)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), vec!["pt", "omitted"] ); assert_eq!( @@ -627,31 +628,33 @@ async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() .with_write_type(vec!["pt".into(), "omitted".into()]) .unwrap(); assert_eq!( - names(writer.write_fields_for_row(&routing(1)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), vec!["pt", "id"] ); assert_eq!( - names(writer.write_fields_for_row(&routing(2)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(2)).unwrap()), vec!["pt", "omitted"] ); let messages = writer.prepare_commit().await.unwrap(); assert_eq!( - names(writer.write_fields_for_row(&routing(1)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), vec!["pt", "id"] ); for invalid in [ routing(1).slice(0, 0), partial(vec![1, 2], vec![Some(1), Some(2)]), ] { - assert!(writer.write_fields_for_row(&invalid).is_err()); + assert!(paimon::table::internal::resolve_row_write_fields(&writer, &invalid).is_err()); } let missing_partition = RecordBatch::try_from_iter([("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef)]) .unwrap(); - assert!(writer.write_fields_for_row(&missing_partition).is_err()); + assert!( + paimon::table::internal::resolve_row_write_fields(&writer, &missing_partition).is_err() + ); // Failed queries leave prepared outputs and future factory settings intact. assert_eq!( - names(writer.write_fields_for_row(&routing(2)).unwrap()), + names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(2)).unwrap()), vec!["pt", "omitted"] ); writer.close().await; @@ -660,6 +663,7 @@ async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() } } +#[cfg(feature = "internal-api")] #[tokio::test] async fn unpartitioned_row_field_resolution_accepts_zero_column_routing_row() { let schema = Schema::builder() @@ -675,10 +679,15 @@ async fn unpartitioned_row_field_resolution_accepts_zero_column_routing_row() { &arrow_array::RecordBatchOptions::new().with_row_count(Some(1)), ) .unwrap(); - assert_eq!(writer.write_fields_for_row(&routing).unwrap().len(), 2); + assert_eq!( + paimon::table::internal::resolve_row_write_fields(&writer, &routing) + .unwrap() + .len(), + 2 + ); writer.with_write_type(vec!["id".into()]).unwrap(); assert_eq!( - writer.write_fields_for_row(&routing).unwrap()[0].name(), + paimon::table::internal::resolve_row_write_fields(&writer, &routing).unwrap()[0].name(), "id" ); assert!(writer.prepare_commit().await.unwrap().is_empty()); From cc9f9727f2d1449aff7a078d77d17286d2a94cf8 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 10:02:15 +0800 Subject: [PATCH 4/4] fix: keep one write type per table writer and remove row schema queries --- .../python/pypaimon_rust/datafusion.pyi | 2 - bindings/python/src/write.rs | 35 - bindings/python/tests/test_write.py | 44 +- .../src/table/bucket_assigner_constant.rs | 16 +- .../paimon/src/table/bucket_assigner_fixed.rs | 16 +- crates/paimon/src/table/internal.rs | 11 - crates/paimon/src/table/table_write.rs | 71 +- crates/paimon/src/table/table_write_append.rs | 339 -------- crates/paimon/src/table/table_write_input.rs | 98 ++- .../src/table/table_write_lifecycle_tests.rs | 33 +- .../src/table/table_write_precomputed.rs | 3 - .../tests/filtered_write_validation_test.rs | 69 +- .../paimon/tests/table_partial_write_test.rs | 800 +++++------------- 13 files changed, 369 insertions(+), 1168 deletions(-) delete mode 100644 crates/paimon/src/table/table_write_append.rs diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index 0fef05bb9..4d4dcefb6 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -446,7 +446,6 @@ class CommitMessage: ... class BatchTableWrite: - def _row_write_columns(self, routing_batch: pyarrow.RecordBatch) -> List[str]: ... def with_blob_uri_reader_factory( self, factory: Optional[_UriReaderFactory] ) -> "BatchTableWrite": @@ -576,7 +575,6 @@ class StreamTableUpdate: ... class StreamTableWrite: - def _row_write_columns(self, routing_batch: pyarrow.RecordBatch) -> List[str]: ... def with_blob_uri_reader_factory( self, factory: Optional[_UriReaderFactory] ) -> "StreamTableWrite": ... diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index dfc33002c..1ac21b058 100644 --- a/bindings/python/src/write.rs +++ b/bindings/python/src/write.rs @@ -347,20 +347,6 @@ impl WriteTarget { } } - fn row_write_columns(&self, routing: &RecordBatch) -> paimon::Result> { - match self { - Self::Table(writer) => Ok(paimon::table::internal::resolve_row_write_fields( - writer, routing, - )? - .iter() - .map(|field| field.name().to_owned()) - .collect()), - Self::PostponeFixed(_) => Err(paimon::Error::Unsupported { - message: "_row_write_columns requires a Paimon append table".into(), - }), - } - } - async fn write_arrow( &mut self, batch: &RecordBatch, @@ -937,15 +923,6 @@ impl WriteState { drop(references); } - fn row_write_columns(&self, routing: &Bound<'_, PyAny>) -> PyResult> { - let routing = crate::arrow_input::write_batch_from_pyarrow(routing)?; - let inner = self - .inner - .as_ref() - .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?; - inner.row_write_columns(&routing).map_err(to_py_err) - } - fn write_result(&self, result: paimon::Result<()>) -> PyResult<()> { result.map_err(to_callback_error) } @@ -1290,10 +1267,6 @@ impl PyBatchTableUpdate { #[pymethods] impl PyBatchTableWrite { - fn _row_write_columns(&self, routing_batch: &Bound<'_, PyAny>) -> PyResult> { - self.state.row_write_columns(routing_batch) - } - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.state.traverse_callbacks(visit) } @@ -1369,10 +1342,6 @@ pub struct PyStreamTableWrite { #[pymethods] impl PyStreamTableWrite { - fn _row_write_columns(&self, routing_batch: &Bound<'_, PyAny>) -> PyResult> { - self.state.row_write_columns(routing_batch) - } - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.state.traverse_callbacks(visit) } @@ -1859,8 +1828,6 @@ writer.with_blob_consumer(Callback(writer, 'first')) if batch { let mut writer = object.extract::>().unwrap(); let inner = writer.state.inner.as_mut().unwrap(); - assert_eq!(inner.row_write_columns(&input).unwrap(), vec!["id"]); - assert!(inner.row_write_columns(&input.slice(0, 0)).is_err()); py.detach(|| runtime().block_on(inner.write_arrow(&input, None))) .unwrap(); } else { @@ -1868,8 +1835,6 @@ writer.with_blob_consumer(Callback(writer, 'first')) .extract::>() .unwrap(); let inner = writer.state.inner.as_mut().unwrap(); - assert_eq!(inner.row_write_columns(&input).unwrap(), vec!["id"]); - assert!(inner.row_write_columns(&input.slice(0, 0)).is_err()); py.detach(|| runtime().block_on(inner.write_arrow(&input, None))) .unwrap(); } diff --git a/bindings/python/tests/test_write.py b/bindings/python/tests/test_write.py index 289ed4b6c..14284b3f6 100644 --- a/bindings/python/tests/test_write.py +++ b/bindings/python/tests/test_write.py @@ -77,22 +77,60 @@ def test_partial_write_type_roundtrip_and_lifecycle(tmp_path, streaming): empty = pa.record_batch( [pa.array([], type=pa.int32()), pa.array([], type=pa.string())], names=["id", "pt"]) writer.write_arrow(empty) + assert writer.with_write_type(["name", "pt"]) is writer assert writer.with_write_type(["id", "pt"]) is writer for columns in [[], ["missing"], ["id", "id"]]: with pytest.raises(ValueError): writer.with_write_type(columns) schema = pa.schema([("id", pa.int32()), ("name", pa.string()), ("pt", pa.string())]) writer.write_arrow(pa.record_batch([[1, 2], ["ignored", "ignored"], ["old", "old"]], schema=schema)) - assert writer.with_write_type(["name", "pt"]) is writer - writer.write_arrow(pa.record_batch([[3, 4], ["ignored", "new"], ["old", "new"]], schema=schema)) + assert writer.with_write_type(["id", "pt"]) is writer + for columns in [["name", "pt"], ["pt", "id"]]: + with pytest.raises(ValueError, match="before any write"): + writer.with_write_type(columns) + writer.write_arrow(pa.record_batch( + [pa.array([3, 4], type=pa.int32()), pa.array(["old", "new"])], names=["id", "pt"])) messages = writer.prepare_commit(True, 7) if streaming else writer.prepare_commit() commit = builder.new_commit() commit.commit(7, messages) if streaming else commit.commit(messages) + if streaming: + assert writer.prepare_commit(True, 8) == [] + assert writer.with_write_type(["id", "pt"]) is writer + with pytest.raises(ValueError, match="before any write"): + writer.with_write_type(["name", "pt"]) writer.close() with pytest.raises(RuntimeError, match="closed"): writer.with_write_type(["id"]) assert pa.Table.from_batches(ctx.sql("SELECT id, name, pt FROM paimon.wdb.t ORDER BY pt, id")).to_pydict() == { - "id": [None, 1, 2, 3], "name": ["new", None, None, None], "pt": ["new", "old", "old", "old"]} + "id": [4, 1, 2, 3], "name": [None] * 4, "pt": ["new", "old", "old", "old"]} + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_ignored_rows_do_not_lock_write_type(tmp_path, streaming): + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": str(tmp_path)}) + ctx.sql("CREATE SCHEMA paimon.wdb") + ctx.sql("CREATE TABLE paimon.wdb.t (id INT, name STRING, op STRING) " + "WITH ('rowkind.field' = 'op', 'ignore-delete' = 'true')") + table = _get_table(str(tmp_path)) + builder = table.new_stream_write_builder() if streaming else table.new_batch_write_builder() + writer = builder.new_write() + schema = pa.schema([("id", pa.int32()), ("name", pa.string()), ("op", pa.string())]) + try: + writer.with_write_type(["id", "op"]) + writer.write_arrow(pa.record_batch([[1], ["ignored"], ["-D"]], schema=schema)) + assert writer.with_write_type(["name", "op"]) is writer + writer.write_arrow(pa.record_batch([[2], ["accepted"], ["+I"]], schema=schema)) + assert writer.with_write_type(["name", "op"]) is writer + with pytest.raises(ValueError, match="before any write"): + writer.with_write_type(["id", "op"]) + messages = writer.prepare_commit(True, 1) if streaming else writer.prepare_commit() + commit = builder.new_commit() + commit.commit(1, messages) if streaming else commit.commit(messages) + finally: + writer.close() + assert pa.Table.from_batches(ctx.sql("SELECT id, name FROM paimon.wdb.t")).to_pydict() == { + "id": [None], "name": ["accepted"]} @pytest.mark.parametrize("streaming", [False, True]) diff --git a/crates/paimon/src/table/bucket_assigner_constant.rs b/crates/paimon/src/table/bucket_assigner_constant.rs index b8d6e7eb3..843d2464a 100644 --- a/crates/paimon/src/table/bucket_assigner_constant.rs +++ b/crates/paimon/src/table/bucket_assigner_constant.rs @@ -45,9 +45,11 @@ impl ConstantBucketAssigner { pub fn bucket(&self) -> i32 { self.bucket } +} - pub(super) fn assign_batch_sync( - &self, +impl BucketAssigner for ConstantBucketAssigner { + async fn assign_batch( + &mut self, batch: &RecordBatch, fields: &[DataField], ) -> Result { @@ -65,16 +67,6 @@ impl ConstantBucketAssigner { skips: Vec::new(), }) } -} - -impl BucketAssigner for ConstantBucketAssigner { - async fn assign_batch( - &mut self, - batch: &RecordBatch, - fields: &[DataField], - ) -> Result { - self.assign_batch_sync(batch, fields) - } async fn prepare_commit_index( &mut self, diff --git a/crates/paimon/src/table/bucket_assigner_fixed.rs b/crates/paimon/src/table/bucket_assigner_fixed.rs index 042ec0ebd..684a98a84 100644 --- a/crates/paimon/src/table/bucket_assigner_fixed.rs +++ b/crates/paimon/src/table/bucket_assigner_fixed.rs @@ -52,9 +52,11 @@ impl FixedBucketAssigner { total_buckets, } } +} - pub(super) fn assign_batch_sync( - &self, +impl BucketAssigner for FixedBucketAssigner { + async fn assign_batch( + &mut self, batch: &RecordBatch, fields: &[DataField], ) -> Result { @@ -80,16 +82,6 @@ impl FixedBucketAssigner { skips: Vec::new(), }) } -} - -impl BucketAssigner for FixedBucketAssigner { - async fn assign_batch( - &mut self, - batch: &RecordBatch, - fields: &[DataField], - ) -> Result { - self.assign_batch_sync(batch, fields) - } async fn prepare_commit_index( &mut self, diff --git a/crates/paimon/src/table/internal.rs b/crates/paimon/src/table/internal.rs index 4d17261b7..ff7d4eaf1 100644 --- a/crates/paimon/src/table/internal.rs +++ b/crates/paimon/src/table/internal.rs @@ -19,17 +19,6 @@ //! These operations are available only with `internal-api`. use super::{CommitMessage, TableWrite}; -use crate::spec::DataField; -use arrow_array::RecordBatch; - -/// Inspect only routing metadata before a binding encodes a named row. -/// This does not create a writer or capture a new bucket configuration. -pub fn resolve_row_write_fields( - writer: &TableWrite, - routing_batch: &RecordBatch, -) -> crate::Result> { - writer.resolve_row_write_fields(routing_batch) -} /// Forward an identifier already supplied by the batch or stream facade. pub async fn prepare_commit_with_identifier( diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 891c7623a..1e8c2217d 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -52,8 +52,6 @@ use arrow_array::{Array, ArrayRef, Int8Array, RecordBatch}; use std::collections::{HashMap, HashSet}; use std::sync::Arc; -#[path = "table_write_append.rs"] -mod append; #[path = "table_write_input.rs"] mod input; #[cfg(test)] @@ -72,8 +70,6 @@ enum AppendPrepareIdentifier { /// Keep only these settings across prepares, not completed file resources. #[derive(Clone)] struct AppendWriteConfig { - schema: Arc, - fields: Vec, blob_consumer: Option>, blob_uri_reader_factory: Option>, last_modified_prepare: Option, @@ -166,7 +162,6 @@ pub struct TableWrite { written: bool, partition_writers: HashMap, append_write_configs: HashMap, - append_write_type_changed: bool, partition_computer: PartitionComputer, partition_keys: Vec, schema_id: i64, @@ -238,7 +233,6 @@ impl TableWrite { written: false, partition_writers: HashMap::new(), append_write_configs: HashMap::new(), - append_write_type_changed: false, partition_computer: PartitionComputer::new( schema.partition_keys(), schema.fields(), @@ -540,7 +534,6 @@ impl TableWrite { written: false, partition_writers: HashMap::new(), append_write_configs: HashMap::new(), - append_write_type_changed: false, partition_computer, partition_keys, schema_id: schema.id(), @@ -669,12 +662,9 @@ impl TableWrite { /// Full-table inputs may provide routing columns which are not written. /// Selected inputs must include partition columns and, unless a bucket is /// supplied through [`Self::write_arrow`], the columns used for hashing. - /// Changes apply to partition/bucket writers created afterwards. Existing - /// writers retain their selected columns across active prepares, until an - /// idle prepare can safely retire them. A full - /// table batch is projected to each writer's captured columns; a selected - /// batch must match the destination writer's names and order, optionally - /// accompanied by routing and configured row-kind fields. + /// Configure this before the first effective write. Afterwards the same + /// ordered selection is a no-op and a different selection is rejected, + /// including after prepares. Empty or ignored input does not lock it. pub fn with_write_type(&mut self, write_columns: Vec) -> Result<&mut Self> { self.ensure_active()?; if self.format_writer.is_some() || !self.table.schema().primary_keys().is_empty() { @@ -710,6 +700,15 @@ impl TableWrite { }) }) .collect::>>()?; + if fields == self.write_fields { + return Ok(self); + } + if self.written { + return Err(crate::Error::DataInvalid { + message: "with_write_type must be called before any write operation".into(), + source: None, + }); + } let indices = |names: &[String]| { names .iter() @@ -749,7 +748,6 @@ impl TableWrite { )) }; let write_schema = build_target_arrow_schema(&fields)?; - self.append_write_type_changed |= self.written && self.write_schema != write_schema; self.bucket_assigner = assigner; self.write_schema = write_schema; self.write_fields = fields; @@ -804,9 +802,6 @@ impl TableWrite { return writer.write(batch).await; } self.ensure_active()?; - if self.has_reconfigured_append_type() { - return self.write_reconfigured_append(batch, None).await; - } let filtered = self.filter_write_input(batch)?; let input = filtered.as_ref().unwrap_or(batch); let Some(batch) = self.normalize_filtered_write_batch(input)? else { @@ -906,23 +901,6 @@ impl TableWrite { // Reject malformed schemas even when every event would be ignored. // Empty arrays check equivalent layouts without examining row values. self.validate_write_batch_schema(&RecordBatch::new_empty(batch.schema()))?; - let projected = self.project_full_write_input(batch)?; - self.filter_write_input_with_schema( - batch, - projected.as_ref().unwrap_or(batch), - &self.write_schema, - ) - } - - fn filter_write_input_with_schema( - &self, - batch: &RecordBatch, - physical: &RecordBatch, - expected_schema: &arrow_schema::Schema, - ) -> Result> { - if self.primary_key_indices.is_empty() && self.row_kind_generator.is_none() { - return Ok(None); - } let kinds = if let Some(generator) = &self.row_kind_generator { generator.generate_batch(batch)? } else if let Some(column) = batch.column_by_name(VALUE_KIND_FIELD_NAME) { @@ -941,7 +919,9 @@ impl TableWrite { // still identify the event; other engines and partition migrations // require every declared top-level NOT NULL field. let delete_required_fields = self.delete_not_null_fields(); - for (index, field) in expected_schema.fields().iter().enumerate() { + let projected = self.project_full_write_input(batch)?; + let physical = projected.as_ref().unwrap_or(batch); + for (index, field) in self.write_schema.fields().iter().enumerate() { if field.is_nullable() || physical.column(index).null_count() == 0 { continue; } @@ -1722,7 +1702,11 @@ impl TableWrite { .append_write_configs .get(&key) .cloned() - .unwrap_or_else(|| self.append_write_config()); + .unwrap_or_else(|| AppendWriteConfig { + blob_consumer: self.blob_consumer.clone(), + blob_uri_reader_factory: self.blob_uri_reader_factory.clone(), + last_modified_prepare: None, + }); let writer = self.create_append_writer(partition_path, bucket, &config)?; self.append_write_configs.insert(key, config); writer @@ -1775,8 +1759,8 @@ impl TableWrite { self.file_format.clone(), self.vector_target_file_size, self.vector_file_format.as_deref(), - &config.schema, - &config.fields, + &self.write_schema, + &self.write_fields, fields, self.table.schema().options(), &self.blob_inline_fields, @@ -1800,22 +1784,21 @@ impl TableWrite { self.file_compression_zstd_level, self.write_buffer_size, self.file_format.clone(), - config.fields.clone(), + self.write_fields.clone(), self.table.schema().options().clone(), Some(0), None, (!super::data_evolution_fields::can_omit_normal_write_cols( self.table.schema().fields(), - &config - .fields + &self + .write_fields .iter() .map(|field| field.name().to_string()) .collect::>(), &CoreOptions::new(self.table.schema().options()), )) .then(|| { - config - .fields + self.write_fields .iter() .map(|field| field.name().to_string()) .collect() @@ -1824,7 +1807,7 @@ impl TableWrite { .with_file_index( self.file_index_options .as_ref() - .and_then(|options| options.project_to_fields(&config.fields)) + .and_then(|options| options.project_to_fields(&self.write_fields)) .map(Arc::new), ) .with_target_file_row_num(self.target_file_row_num) diff --git a/crates/paimon/src/table/table_write_append.rs b/crates/paimon/src/table/table_write_append.rs deleted file mode 100644 index 9163c7f35..000000000 --- a/crates/paimon/src/table/table_write_append.rs +++ /dev/null @@ -1,339 +0,0 @@ -// 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. - -//! Append writer factory configuration survives physical-file prepares. - -use super::*; - -impl TableWrite { - /// Resolve the columns captured by the destination append writer without - /// creating files or changing bucket assignment state. Row adapters can - /// encode only these columns, leaving unused payloads untouched. The input - /// must contain exactly one row and the table's routing columns. - #[cfg(any(test, feature = "internal-api"))] - pub(crate) fn resolve_row_write_fields( - &self, - routing_batch: &RecordBatch, - ) -> Result> { - self.ensure_active()?; - if self.format_writer.is_some() || !self.primary_key_indices.is_empty() { - return Err(crate::Error::Unsupported { - message: "Row write field resolution requires a Paimon append table".into(), - }); - } - if routing_batch.num_rows() != 1 { - return Err(crate::Error::DataInvalid { - message: "Row write field resolution requires exactly one routing row".into(), - source: None, - }); - } - let names = self.automatic_routing_names(); - let key = if names.is_empty() { - // A zero-column routing batch still represents one row. Avoid a - // schema normalizer which would infer zero rows from no arrays. - let bucket = CoreOptions::new(self.table.schema().options()).bucket(); - ( - EMPTY_SERIALIZED_ROW.clone(), - if bucket == POSTPONE_BUCKET { - POSTPONE_BUCKET - } else { - 0 - }, - ) - } else { - let routing = self.normalize_routing_columns(routing_batch, &names)?; - let mut output = self.assign_append_routing(routing)?; - (output.partition_bytes.remove(0), output.buckets[0]) - }; - Ok(self - .append_write_configs - .get(&key) - .map(|config| &config.fields) - .unwrap_or(&self.write_fields) - .clone()) - } - - pub(super) fn append_write_config(&self) -> AppendWriteConfig { - AppendWriteConfig { - schema: self.write_schema.clone(), - fields: self.write_fields.clone(), - blob_consumer: self.blob_consumer.clone(), - blob_uri_reader_factory: self.blob_uri_reader_factory.clone(), - last_modified_prepare: None, - } - } - - pub(super) fn has_reconfigured_append_type(&self) -> bool { - self.append_write_type_changed - } - - /// Java captures the physical write type per partition/bucket. Arrow's - /// named columns must agree with that captured type, rather than being - /// reinterpreted positionally as Java's InternalRow can be. - pub(super) async fn write_reconfigured_append( - &mut self, - input: &RecordBatch, - bucket: Option, - ) -> Result<()> { - let full_schema = build_target_arrow_schema(self.table.schema().fields())?; - let actual_schema = input.schema(); - let names_match = |expected: &arrow_schema::Schema| { - actual_schema - .fields() - .iter() - .map(|f| f.name()) - .eq(expected.fields().iter().map(|f| f.name())) - }; - let mut routing_names = if bucket.is_some() { - self.partition_keys.clone() - } else { - self.automatic_routing_names() - }; - if let Some(field) = CoreOptions::new(self.table.schema().options()).rowkind_field() { - routing_names.push(field.to_owned()); - } - let routing_transport = |expected: &arrow_schema::Schema| { - is_routing_transport(&actual_schema, expected, &routing_names) - }; - let full_input = names_match(&full_schema); - let candidates = std::iter::once(self.write_schema.as_ref()) - .chain( - self.append_write_configs - .values() - .map(|config| config.schema.as_ref()), - ) - .filter(|schema| names_match(schema) || routing_transport(schema)) - .collect::>(); - let logical_schema = if full_input - || names_match(&self.write_schema) - || routing_transport(&self.write_schema) - { - self.write_schema.clone() - } else if let Some(exact) = candidates - .iter() - .copied() - .find(|schema| names_match(schema)) - { - // A selected Arrow layout explicitly declares its logical fields. - Arc::new(exact.clone()) - } else if candidates.is_empty() { - return Err(crate::Error::DataInvalid { - message: "Input columns must match the table or a selected append write type, with optional routing columns".into(), - source: None, - }); - } else { - // Historical types can differ only in routing/event fields while - // admitting the same transport. Do not choose their logical NULL - // checks by HashMap order: retain the common selected fields here, - // then validate each destination's complete physical type below. - let fields = self - .table - .schema() - .fields() - .iter() - .filter(|field| { - candidates - .iter() - .all(|candidate| candidate.field_with_name(field.name()).is_ok()) - }) - .cloned() - .collect::>(); - build_target_arrow_schema(&fields)? - }; - // Validate every provided transport field, without inspecting values of - // omitted payloads. Row marshalling may add routing/event fields to a - // captured subset, but arbitrary payload extras are not accepted. - let transport_fields = actual_schema - .fields() - .iter() - .map(|field| { - self.table - .schema() - .fields() - .iter() - .find(|candidate| candidate.name() == field.name()) - .cloned() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("Column '{}' is not in table schema", field.name()), - source: None, - }) - }) - .collect::>>()?; - let transport_schema = build_target_arrow_schema(&transport_fields)?; - self.normalize_batch_schema(&input.slice(0, 0), &transport_schema)?; - if input.num_rows() == 0 { - return Ok(()); - } - let logical = project(input, &logical_schema)?; - // The facade resolves events and checks its selected NOT NULL fields - // before routing, as Java TableWriteImpl does. Ignored events must not - // cast routing values or open Blob payloads. - let filtered = self.filter_write_input_with_schema(input, &logical, &logical_schema)?; - let input = filtered.as_ref().unwrap_or(input); - if input.num_rows() == 0 { - return Ok(()); - } - let groups = if let Some(bucket) = bucket { - let partitions = self.precomputed_partition_bytes(input, input)?; - let mut groups: HashMap> = HashMap::new(); - for (row, partition) in partitions.into_iter().enumerate() { - groups.entry((partition, bucket)).or_default().push(row); - } - groups - .into_iter() - .map(|(key, rows)| Ok((key, take_rows(input, &rows)?))) - .collect::>>()? - } else { - let routing = self.normalize_routing_columns(input, &self.automatic_routing_names())?; - self.divide_write_input(input, Some(routing)).await? - }; - let mut prepared = Vec::with_capacity(groups.len()); - for (key, batch) in groups { - let expected = self - .append_write_configs - .get(&key) - .map(|config| &config.schema) - .unwrap_or(&self.write_schema); - let physical = if full_input || routing_transport(expected) { - project(&batch, expected)? - } else { - // Selected inputs retain their named layout. Never reinterpret - // equal-width input positionally for a previously opened writer. - batch - }; - if physical.num_rows() > 0 { - prepared.push((key, self.normalize_batch_schema(&physical, expected)?)); - } - } - if bucket.is_some() && prepared.len() > 1 { - return Err(crate::Error::DataInvalid { - message: "A precomputed bucket group contained multiple partitions".into(), - source: None, - }); - } - // Preflight every destination before creating files or modifying any - // existing writer. A bad subset cannot partially stage a mixed batch. - for ((partition, bucket), batch) in prepared { - self.written = true; - self.write_bucket(partition, bucket, batch).await?; - } - Ok(()) - } -} - -fn project(batch: &RecordBatch, expected: &arrow_schema::Schema) -> Result { - let indices = expected - .fields() - .iter() - .map(|field| { - batch - .schema() - .index_of(field.name()) - .map_err(|error| crate::Error::DataInvalid { - message: format!("Cannot project captured append write type: {error}"), - source: None, - }) - }) - .collect::>>()?; - batch - .project(&indices) - .map_err(|error| crate::Error::DataInvalid { - message: format!("Cannot project captured append write type: {error}"), - source: None, - }) -} - -/// Routing/event fields may accompany a named subset. Equal-width reordered -/// subsets remain invalid Arrow layouts; only additional routing fields permit -/// projection, and every captured payload field must still be present. -fn is_routing_transport( - actual: &arrow_schema::Schema, - expected: &arrow_schema::Schema, - routing_names: &[String], -) -> bool { - if actual.fields().len() <= expected.fields().len() { - return false; - } - let names = actual - .fields() - .iter() - .map(|field| field.name()) - .collect::>(); - names.len() == actual.fields().len() - && expected - .fields() - .iter() - .all(|field| names.contains(field.name())) - && actual.fields().iter().all(|field| { - expected - .fields() - .iter() - .any(|candidate| candidate.name() == field.name()) - || routing_names.contains(field.name()) - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::catalog::Identifier; - use crate::io::FileIOBuilder; - use crate::spec::{DataType, IntType, Schema, TableSchema}; - - #[test] - fn row_write_field_resolution_does_not_capture_a_writer() { - let schema = Schema::builder() - .column("first", DataType::Int(IntType::new())) - .column("second", DataType::Int(IntType::new())) - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "row_fields"), - "memory:/row_fields".into(), - TableSchema::new(0, &schema), - None, - ); - let mut writer = TableWrite::new(&table, "row-fields".into()).unwrap(); - let routing = |rows| { - RecordBatch::try_new_with_options( - Arc::new(arrow_schema::Schema::empty()), - vec![], - &arrow_array::RecordBatchOptions::new().with_row_count(Some(rows)), - ) - .unwrap() - }; - assert_eq!( - writer.resolve_row_write_fields(&routing(1)).unwrap().len(), - 2 - ); - writer.with_write_type(vec!["first".into()]).unwrap(); - assert_eq!( - writer.resolve_row_write_fields(&routing(1)).unwrap()[0].name(), - "first" - ); - writer.with_write_type(vec!["second".into()]).unwrap(); - assert_eq!( - writer.resolve_row_write_fields(&routing(1)).unwrap()[0].name(), - "second" - ); - assert!(writer.resolve_row_write_fields(&routing(0)).is_err()); - assert!(writer.resolve_row_write_fields(&routing(2)).is_err()); - assert!(writer.partition_writers.is_empty()); - assert!(writer.append_write_configs.is_empty()); - } -} diff --git a/crates/paimon/src/table/table_write_input.rs b/crates/paimon/src/table/table_write_input.rs index 5c549a78b..a5987da44 100644 --- a/crates/paimon/src/table/table_write_input.rs +++ b/crates/paimon/src/table/table_write_input.rs @@ -26,7 +26,8 @@ pub(super) struct RoutingInput { } impl TableWrite { - /// Accept the declared table layout as well as the selected write layout. + /// Accept the declared table layout or the selected write layout with + /// additional routing/event columns. The selected columns keep their order. /// The empty schema check intentionally does not inspect unwritten values: /// an omitted Blob must never open a URI or invoke a Blob consumer. pub(super) fn project_full_write_input( @@ -42,14 +43,59 @@ impl TableWrite { .map(|field| field.name()) .eq(expected.iter().map(|field| field.name())) }; - if !self.table.schema().primary_keys().is_empty() - || names_match(&self.write_fields) - || !names_match(fields) - { + if !self.table.schema().primary_keys().is_empty() || names_match(&self.write_fields) { return Ok(None); } - let full_schema = build_target_arrow_schema(fields)?; - self.normalize_batch_schema(&input.slice(0, 0), &full_schema)?; + let transport_fields = if names_match(fields) { + fields.to_vec() + } else { + let selected = self + .write_fields + .iter() + .map(|field| field.name()) + .collect::>(); + let mut routing = self + .partition_keys + .iter() + .map(String::as_str) + .collect::>(); + let bucket_keys = self.table.schema().bucket_keys(); + let options = CoreOptions::new(self.table.schema().options()); + if options.bucket() > 1 { + routing.extend(bucket_keys.iter().map(String::as_str)); + } + if let Some(field) = options.rowkind_field() { + routing.insert(field); + } + let schema = input.schema(); + let mut names = HashSet::new(); + if !schema.fields().iter().all(|field| { + names.insert(field.name().as_str()) + && (selected.contains(field.name().as_str()) + || routing.contains(field.name().as_str())) + }) || !schema + .fields() + .iter() + .filter(|field| selected.contains(field.name().as_str())) + .map(|field| field.name().as_str()) + .eq(self.write_fields.iter().map(|field| field.name())) + { + return Ok(None); + } + schema + .fields() + .iter() + .map(|field| { + fields + .iter() + .find(|candidate| candidate.name() == field.name()) + .unwrap() + .clone() + }) + .collect() + }; + let transport_schema = build_target_arrow_schema(&transport_fields)?; + self.normalize_batch_schema(&input.slice(0, 0), &transport_schema)?; let indices = self .write_fields .iter() @@ -64,7 +110,7 @@ impl TableWrite { }) } - pub(super) fn normalize_routing_columns( + fn normalize_routing_columns( &self, input: &RecordBatch, names: &[String], @@ -113,19 +159,14 @@ impl TableWrite { Ok(RoutingInput { batch, fields }) } - pub(super) fn automatic_routing_names(&self) -> Vec { - let mut names = self.partition_keys.clone(); - if CoreOptions::new(self.table.schema().options()).bucket() > 1 { - names.extend(self.table.schema().bucket_keys()); - } - names - } - pub(super) fn normalize_automatic_routing_input( &self, input: &RecordBatch, ) -> Result> { - let names = self.automatic_routing_names(); + let mut names = self.partition_keys.clone(); + if CoreOptions::new(self.table.schema().options()).bucket() > 1 { + names.extend(self.table.schema().bucket_keys()); + } if names .iter() .all(|name| self.write_fields.iter().any(|field| field.name() == name)) @@ -146,11 +187,6 @@ impl TableWrite { let Some(routing) = routing else { return self.divide_by_partition_bucket(physical).await; }; - let output = self.assign_append_routing(routing)?; - Self::group_assigned_rows(physical, output) - } - - pub(super) fn assign_append_routing(&self, routing: RoutingInput) -> Result { let indices = |names: &[String]| { names .iter() @@ -165,25 +201,27 @@ impl TableWrite { }; let partitions = indices(&self.partition_keys); let options = CoreOptions::new(self.table.schema().options()); - if options.bucket() > 1 { - FixedBucketAssigner::new( + let mut assigner = if options.bucket() > 1 { + BucketAssignerEnum::Fixed(FixedBucketAssigner::new( partitions, indices(&self.table.schema().bucket_keys()), options.bucket_function_type()?, options.bucket(), - ) - .assign_batch_sync(&routing.batch, &routing.fields) + )) } else { - ConstantBucketAssigner::new( + BucketAssignerEnum::Constant(ConstantBucketAssigner::new( partitions, if options.bucket() == POSTPONE_BUCKET { POSTPONE_BUCKET } else { 0 }, - ) - .assign_batch_sync(&routing.batch, &routing.fields) - } + )) + }; + let output = assigner + .assign_batch(&routing.batch, &routing.fields) + .await?; + Self::group_assigned_rows(physical, output) } pub(super) fn precomputed_partition_bytes( diff --git a/crates/paimon/src/table/table_write_lifecycle_tests.rs b/crates/paimon/src/table/table_write_lifecycle_tests.rs index b0ce2da4c..0389cd447 100644 --- a/crates/paimon/src/table/table_write_lifecycle_tests.rs +++ b/crates/paimon/src/table/table_write_lifecycle_tests.rs @@ -118,7 +118,7 @@ async fn snapshot(table: &Table, id: i64, user: &str, identifier: i64) { } #[tokio::test] -async fn unaware_append_recaptures_type_only_after_an_idle_prepare() { +async fn unaware_append_retirement_does_not_unlock_write_type() { for idle_prepares in 0..=2 { let table = table(false, false).await; let mut writer = TableWrite::new(&table, "owner".into()).unwrap(); @@ -127,17 +127,17 @@ async fn unaware_append_recaptures_type_only_after_an_idle_prepare() { .unwrap(); write(&mut writer, false, 1).await; assert_columns(&writer.prepare_commit().await.unwrap(), "id"); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); for _ in 0..idle_prepares { assert!(writer.prepare_commit().await.unwrap().is_empty()); } + assert!(writer + .with_write_type(vec!["pt".into(), "value".into()]) + .is_err()); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); write(&mut writer, false, 1).await; - assert_columns( - &writer.prepare_commit().await.unwrap(), - if idle_prepares == 0 { "id" } else { "value" }, - ); + assert_columns(&writer.prepare_commit().await.unwrap(), "id"); } } @@ -213,9 +213,6 @@ async fn fixed_append_waits_for_a_strictly_newer_committed_identifier_of_its_use &writer.prepare_commit_with_identifier(1).await.unwrap(), "id", ); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); for (id, user, identifier) in [(1, "owner", 1), (2, "owner", 1), (3, "other", 99)] { snapshot(&table, id, user, identifier).await; assert!(writer @@ -231,10 +228,11 @@ async fn fixed_append_waits_for_a_strictly_newer_committed_identifier_of_its_use .await .unwrap() .is_empty()); + assert!(writer.append_write_configs.is_empty()); write(&mut writer, true, 1).await; assert_columns( &writer.prepare_commit_with_identifier(3).await.unwrap(), - "value", + "id", ); } @@ -251,9 +249,6 @@ async fn fixed_append_keeps_unknown_prepare_identifiers() { .prepare_commit_internal(initial_identifier) .await .unwrap(); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); snapshot(&table, 1, "owner", 2).await; assert!(writer.prepare_commit().await.unwrap().is_empty()); if initial_identifier.is_none() { @@ -277,9 +272,6 @@ async fn fixed_append_stops_at_an_expired_snapshot_gap() { .unwrap(); write(&mut writer, true, 1).await; writer.prepare_commit_with_identifier(1).await.unwrap(); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); snapshot(&table, 1, "owner", 2).await; snapshot(&table, 3, "other", 99).await; assert!(writer @@ -303,9 +295,6 @@ async fn fixed_append_metadata_errors_leave_pending_output_available() { .unwrap(); write(&mut writer, true, 1).await; writer.prepare_commit_with_identifier(1).await.unwrap(); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); write(&mut writer, true, 2).await; let manager = table.snapshot_manager(); table @@ -323,7 +312,7 @@ async fn fixed_append_metadata_errors_leave_pending_output_available() { snapshot(&table, 1, "owner", 1).await; assert_columns( &writer.prepare_commit_with_identifier(2).await.unwrap(), - "value", + "id", ); } diff --git a/crates/paimon/src/table/table_write_precomputed.rs b/crates/paimon/src/table/table_write_precomputed.rs index 3656777f8..3730d3d0f 100644 --- a/crates/paimon/src/table/table_write_precomputed.rs +++ b/crates/paimon/src/table/table_write_precomputed.rs @@ -63,9 +63,6 @@ impl TableWrite { upper_bound - 1, ))); } - if self.has_reconfigured_append_type() { - return self.write_reconfigured_append(batch, Some(bucket)).await; - } let filtered = self.filter_write_input(batch)?; let input = filtered.as_ref().unwrap_or(batch); let Some(batch) = self.normalize_filtered_write_batch(input)? else { diff --git a/crates/paimon/tests/filtered_write_validation_test.rs b/crates/paimon/tests/filtered_write_validation_test.rs index efced043c..97d0deee0 100644 --- a/crates/paimon/tests/filtered_write_validation_test.rs +++ b/crates/paimon/tests/filtered_write_validation_test.rs @@ -365,6 +365,9 @@ async fn completely_filtered_batches_do_not_lock_writer_configuration() { writer.with_blob_uri_reader_factory(None).unwrap(); writer.write_arrow_batch(&batch("+I")).await.unwrap(); writer.with_write_type(vec!["id".into()]).unwrap(); + assert!(writer + .with_write_type(vec!["id".into(), "op".into()]) + .is_err()); writer.with_blob_consumer(None).unwrap(); assert!(writer.with_blob_uri_reader_factory(None).is_err()); let messages = writer.prepare_commit().await.unwrap(); @@ -397,69 +400,3 @@ async fn completely_filtered_batches_do_not_lock_writer_configuration() { .iter() .all(|batch| batch.column_by_name("op").unwrap().is_null(0))); } - -#[tokio::test] -async fn reconfigured_append_filters_events_before_nested_value_validation() { - let (value_type, invalid_values) = payload("array", false); - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column("value", value_type) - .column("pt", DataType::Int(IntType::new())) - .column("op", DataType::VarChar(VarCharType::string_type())) - .partition_keys(["pt"]) - .option("rowkind.field", "op") - .option("ignore-delete", "true") - .option("data-evolution.enabled", "true") - .option("row-tracking.enabled", "true") - .build() - .unwrap(); - let (io, table) = memory_table( - "memory:/reconfigured-filtered", - TableSchema::new(0, &schema), - ); - setup_dirs(&io, table.location()).await; - persist_table_schema(&io, table.location(), table.schema()).await; - let mut writer = table.new_write_builder().new_write().unwrap(); - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - let input = |ids, values: ArrayRef, partitions, kinds| { - RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), - ("value", values), - ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), - ("op", Arc::new(StringArray::from(kinds)) as ArrayRef), - ]) - .unwrap() - }; - writer - .write_arrow_batch(&input( - vec![11], - invalid_values.slice(0, 1), - vec![1], - vec!["+I"], - )) - .await - .unwrap(); - writer - .with_write_type(vec!["pt".into(), "value".into()]) - .unwrap(); - writer - .write_arrow_batch(&input( - vec![99, 88], - invalid_values, - vec![1, 2], - vec!["-D", "+I"], - )) - .await - .unwrap(); - let messages = writer.prepare_commit().await.unwrap(); - assert_eq!( - messages - .iter() - .flat_map(|m| &m.new_files) - .map(|f| f.row_count) - .sum::(), - 2 - ); -} diff --git a/crates/paimon/tests/table_partial_write_test.rs b/crates/paimon/tests/table_partial_write_test.rs index 5030d90a7..2e305d27f 100644 --- a/crates/paimon/tests/table_partial_write_test.rs +++ b/crates/paimon/tests/table_partial_write_test.rs @@ -158,8 +158,8 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { .with_write_type(columns.into_iter().map(str::to_string).collect()) .is_err()); } - // Invalid requests preserve the factory configuration. Valid requests - // remain allowed after writing and preparing existing bucket writers. + // Empty batches do not lock the configuration; nonempty ones do, even + // after their files have been prepared. writer .write_arrow_batch(&partial(vec![], vec![])) .await @@ -173,12 +173,220 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { .unwrap(); let messages = writer.prepare_commit().await.unwrap(); writer - .with_write_type(vec!["id".into(), "omitted".into(), "pt".into()]) + .with_write_type(vec!["pt".into(), "id".into()]) .unwrap(); + assert!(writer + .with_write_type(vec!["id".into(), "omitted".into(), "pt".into()]) + .is_err()); builder.new_commit().commit(messages).await.unwrap(); assert_eq!(rows(&table).await, vec![vec![Some(1), None, Some(2)]]); } +#[tokio::test] +async fn append_write_type_stays_locked_across_partitions_and_prepares() { + for fixed in [false, true] { + let table = table(!fixed, fixed, false).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + writer + .write_arrow(&[partial(vec![11], vec![Some(1)])], fixed.then_some(0)) + .await + .unwrap(); + for columns in [vec!["pt", "omitted"], vec!["id", "pt"]] { + let error = writer + .with_write_type(columns.into_iter().map(str::to_string).collect()) + .err() + .unwrap(); + assert!(error.to_string().contains("before any write operation")); + } + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + let full = RecordBatch::try_from_iter([ + ("id", Arc::new(Int32Array::from(vec![22, 33])) as ArrayRef), + ( + "omitted", + Arc::new(Int32Array::from(vec![99, 99])) as ArrayRef, + ), + ("pt", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef), + ]) + .unwrap(); + if fixed { + for row in 0..2 { + writer + .write_arrow(&[full.slice(row, 1)], Some(0)) + .await + .unwrap(); + } + } else { + writer.write_arrow_batch(&full).await.unwrap(); + } + let mut messages = writer.prepare_commit().await.unwrap(); + assert!(writer.prepare_commit().await.unwrap().is_empty()); + assert!(writer.with_write_type(vec!["omitted".into()]).is_err()); + writer + .write_arrow(&[partial(vec![44], vec![Some(3)])], fixed.then_some(0)) + .await + .unwrap(); + messages.extend(writer.prepare_commit().await.unwrap()); + builder.new_commit().commit(messages).await.unwrap(); + assert_eq!( + rows(&table).await, + vec![ + vec![Some(11), None, Some(1)], + vec![Some(22), None, Some(1)], + vec![Some(33), None, Some(2)], + vec![Some(44), None, Some(3)], + ] + ); + } +} + +#[tokio::test] +async fn selected_input_accepts_only_current_columns_and_routing_fields() { + use paimon::spec::{BlobType, VarCharType}; + for fixed in [false, true] { + let mut schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::with_nullable(false))) + .column( + "unused", + if fixed { + DataType::Int(IntType::new()) + } else { + DataType::Blob(BlobType::new()) + }, + ) + .column("pt", DataType::Int(IntType::with_nullable(false))) + .column("op", DataType::VarChar(VarCharType::string_type())) + .partition_keys(["pt"]) + .option("rowkind.field", "op") + .option("ignore-delete", "true") + .option("data-evolution.enabled", (!fixed).to_string()) + .option("row-tracking.enabled", (!fixed).to_string()); + if fixed { + schema = schema.option("bucket", "4").option("bucket-key", "id"); + } + let (io, table) = memory_table( + "memory:/selected_transport", + TableSchema::new(0, &schema.build().unwrap()), + ); + setup_dirs(&io, table.location()).await; + persist_table_schema(&io, table.location(), table.schema()).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.with_write_type(vec!["value".into()]).unwrap(); + let transport = |values: Vec>, kinds: Vec<&str>, extra_payload| { + let count = values.len(); + let mut columns = vec![ + ("value", Arc::new(Int32Array::from(values)) as ArrayRef), + ( + "pt", + Arc::new(Int32Array::from(vec![Some(1); count])) as ArrayRef, + ), + ( + "op", + Arc::new(arrow_array::StringArray::from(kinds)) as ArrayRef, + ), + ]; + if fixed { + columns.push(( + "id", + Arc::new(Int32Array::from(vec![10; count])) as ArrayRef, + )); + } + if extra_payload { + let unused: ArrayRef = if fixed { + Arc::new(Int32Array::from(vec![99; count])) + } else { + Arc::new(arrow_array::LargeBinaryArray::from(vec![ + b"bad".as_slice(); + count + ])) + }; + columns.push(("unused", unused)); + } + RecordBatch::try_from_iter(columns).unwrap() + }; + // Ignored events do not validate nested or omitted payload values, but + // the current selected NOT NULL fields still have their logical rule. + assert!(writer + .write_arrow_batch(&transport(vec![None], vec!["-D"], false)) + .await + .is_err()); + writer + .write_arrow_batch(&transport(vec![Some(1)], vec!["-D"], false)) + .await + .unwrap(); + writer + .with_write_type(vec!["value".into(), "pt".into()]) + .unwrap(); + writer.with_write_type(vec!["value".into()]).unwrap(); + writer + .write_arrow_batch(&transport( + vec![Some(22), Some(33)], + vec!["+I", "-D"], + false, + )) + .await + .unwrap(); + assert!(writer + .write_arrow_batch(&transport(vec![Some(99)], vec!["+I"], true)) + .await + .is_err()); + assert!(writer.with_write_type(vec!["id".into()]).is_err()); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + messages + .iter() + .flat_map(|m| &m.new_files) + .map(|file| file.row_count) + .sum::(), + 1 + ); + assert!(messages + .iter() + .flat_map(|m| &m.new_files) + .all(|file| !file.file_name.ends_with(".blob"))); + builder.new_commit().commit(messages).await.unwrap(); + let mut read = table.new_read_builder(); + read.with_projection(&["value", "pt"]).unwrap(); + let plan = read.new_scan().plan().await.unwrap(); + let batches: Vec = read + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!( + batches[0] + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 22 + ); + assert_eq!( + batches[0] + .column_by_name("pt") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + } +} + #[tokio::test] async fn omitted_managed_blob_does_not_create_dedicated_files() { let table = table(true, false, true).await; @@ -337,176 +545,6 @@ async fn selected_columns_validate_input_before_staging() { ); } -fn selected_value(values: Vec, partitions: Vec) -> RecordBatch { - RecordBatch::try_from_iter([ - ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), - ("omitted", Arc::new(Int32Array::from(values)) as ArrayRef), - ]) - .unwrap() -} - -fn full_values(ids: Vec, values: Vec, partitions: Vec) -> RecordBatch { - RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), - ("omitted", Arc::new(Int32Array::from(values)) as ArrayRef), - ("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef), - ]) - .unwrap() -} - -#[tokio::test] -async fn append_write_type_is_captured_per_bucket_across_prepare_rounds() { - for explicit_bucket in [false, true] { - let table = table(!explicit_bucket, explicit_bucket, false).await; - let builder = table.new_write_builder(); - let mut writer = builder.new_write().unwrap(); - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - let bucket = explicit_bucket.then_some(0); - writer - .write_arrow(&[partial(vec![11], vec![Some(1)])], bucket) - .await - .unwrap(); - // Repeating the same type does not lock an already active writer. - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - writer - .with_write_type(vec!["pt".into(), "omitted".into()]) - .unwrap(); - writer - .write_arrow(&[selected_value(vec![22], vec![2])], bucket) - .await - .unwrap(); - // A selected input cannot be reinterpreted positionally in an old - // bucket. Validation must preserve both its old data and the factory. - assert!(writer - .write_arrow(&[selected_value(vec![999], vec![1])], bucket) - .await - .is_err()); - writer - .write_arrow(&[partial(vec![33], vec![Some(1)])], bucket) - .await - .unwrap(); - let first = writer.prepare_commit().await.unwrap(); - assert!(first - .iter() - .flat_map(|m| &m.new_files) - .any(|f| f.write_cols == Some(vec!["pt".into(), "id".into()]))); - assert!(first - .iter() - .flat_map(|m| &m.new_files) - .any(|f| f.write_cols == Some(vec!["pt".into(), "omitted".into()]))); - builder.new_commit().commit(first).await.unwrap(); - // Finished physical resources are gone, but captured settings remain. - writer - .with_write_type(vec!["pt".into(), "omitted".into()]) - .unwrap(); - writer - .write_arrow(&[partial(vec![44], vec![Some(1)])], bucket) - .await - .unwrap(); - writer - .write_arrow(&[selected_value(vec![55], vec![2])], bucket) - .await - .unwrap(); - writer - .write_arrow(&[selected_value(vec![66], vec![3])], bucket) - .await - .unwrap(); - builder - .new_commit() - .commit(writer.prepare_commit().await.unwrap()) - .await - .unwrap(); - assert_eq!( - rows(&table).await, - vec![ - vec![None, Some(22), Some(2)], - vec![None, Some(55), Some(2)], - vec![None, Some(66), Some(3)], - vec![Some(11), None, Some(1)], - vec![Some(33), None, Some(1)], - vec![Some(44), None, Some(1)], - ] - ); - } -} - -#[tokio::test] -async fn reconfigured_mixed_input_is_preflighted_and_full_input_projects_each_bucket() { - let table = table(true, false, false).await; - let builder = table.new_write_builder(); - let mut writer = builder.new_write().unwrap(); - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - writer - .write_arrow_batch(&partial(vec![11], vec![Some(1)])) - .await - .unwrap(); - writer - .with_write_type(vec!["pt".into(), "omitted".into()]) - .unwrap(); - let before = table - .file_io() - .list_status_recursive(table.location()) - .await - .unwrap() - .len(); - assert!(writer - .write_arrow_batch(&selected_value(vec![999, 888], vec![1, 2])) - .await - .is_err()); - assert_eq!( - table - .file_io() - .list_status_recursive(table.location()) - .await - .unwrap() - .len(), - before - ); - writer - .write_arrow_batch(&selected_value(vec![22], vec![2])) - .await - .unwrap(); - writer - .write_arrow_batch(&full_values(vec![33, 44], vec![333, 444], vec![1, 3])) - .await - .unwrap(); - builder - .new_commit() - .commit(writer.prepare_commit().await.unwrap()) - .await - .unwrap(); - // Subsequent valid changes do not alter existing buckets' captured order. - writer - .with_write_type(vec!["id".into(), "pt".into()]) - .unwrap(); - writer - .write_arrow_batch(&full_values(vec![55, 66], vec![555, 666], vec![1, 2])) - .await - .unwrap(); - builder - .new_commit() - .commit(writer.prepare_commit().await.unwrap()) - .await - .unwrap(); - assert_eq!( - rows(&table).await, - vec![ - vec![None, Some(22), Some(2)], - vec![None, Some(444), Some(3)], - vec![None, Some(666), Some(2)], - vec![Some(11), None, Some(1)], - vec![Some(33), None, Some(1)], - vec![Some(55), None, Some(1)], - ] - ); -} - #[tokio::test] async fn append_blob_consumer_is_captured_per_partition_across_prepares() { use paimon::spec::BlobDescriptor; @@ -562,419 +600,3 @@ async fn append_blob_consumer_is_captured_per_partition_across_prepares() { vec!["first", "first", "first", "second", "second"] ); } - -#[cfg(feature = "internal-api")] -#[tokio::test] -async fn append_row_field_resolution_is_read_only_and_retains_captured_layouts() { - for fixed in [false, true] { - let table = table(false, fixed, false).await; - let builder = table.new_write_builder(); - let mut writer = builder.new_write().unwrap(); - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - let routing = |partition| { - RecordBatch::try_from_iter([ - ( - "pt", - Arc::new(Int32Array::from(vec![partition])) as ArrayRef, - ), - ("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef), - ]) - .unwrap() - }; - let names = |fields: Vec| { - fields - .into_iter() - .map(|f| f.name().to_owned()) - .collect::>() - }; - let before = table - .file_io() - .list_status_recursive(table.location()) - .await - .unwrap() - .len(); - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), - vec!["pt", "id"] - ); - assert!(writer.prepare_commit().await.unwrap().is_empty()); - // A query does not create or freeze the partition writer. - writer - .with_write_type(vec!["pt".into(), "omitted".into()]) - .unwrap(); - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), - vec!["pt", "omitted"] - ); - assert_eq!( - table - .file_io() - .list_status_recursive(table.location()) - .await - .unwrap() - .len(), - before - ); - writer - .with_write_type(vec!["pt".into(), "id".into()]) - .unwrap(); - writer - .write_arrow_batch(&partial(vec![11], vec![Some(1)])) - .await - .unwrap(); - writer - .with_write_type(vec!["pt".into(), "omitted".into()]) - .unwrap(); - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), - vec!["pt", "id"] - ); - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(2)).unwrap()), - vec!["pt", "omitted"] - ); - let messages = writer.prepare_commit().await.unwrap(); - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(1)).unwrap()), - vec!["pt", "id"] - ); - for invalid in [ - routing(1).slice(0, 0), - partial(vec![1, 2], vec![Some(1), Some(2)]), - ] { - assert!(paimon::table::internal::resolve_row_write_fields(&writer, &invalid).is_err()); - } - let missing_partition = - RecordBatch::try_from_iter([("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef)]) - .unwrap(); - assert!( - paimon::table::internal::resolve_row_write_fields(&writer, &missing_partition).is_err() - ); - // Failed queries leave prepared outputs and future factory settings intact. - assert_eq!( - names(paimon::table::internal::resolve_row_write_fields(&writer, &routing(2)).unwrap()), - vec!["pt", "omitted"] - ); - writer.close().await; - builder.new_commit().commit(messages).await.unwrap(); - assert_eq!(rows(&table).await, vec![vec![Some(11), None, Some(1)]]); - } -} - -#[cfg(feature = "internal-api")] -#[tokio::test] -async fn unpartitioned_row_field_resolution_accepts_zero_column_routing_row() { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column("omitted", DataType::Int(IntType::new())) - .build() - .unwrap(); - let (_, table) = memory_table("memory:/row_fields", TableSchema::new(0, &schema)); - let mut writer = table.new_write_builder().new_write().unwrap(); - let routing = RecordBatch::try_new_with_options( - Arc::new(arrow_schema::Schema::empty()), - vec![], - &arrow_array::RecordBatchOptions::new().with_row_count(Some(1)), - ) - .unwrap(); - assert_eq!( - paimon::table::internal::resolve_row_write_fields(&writer, &routing) - .unwrap() - .len(), - 2 - ); - writer.with_write_type(vec!["id".into()]).unwrap(); - assert_eq!( - paimon::table::internal::resolve_row_write_fields(&writer, &routing).unwrap()[0].name(), - "id" - ); - assert!(writer.prepare_commit().await.unwrap().is_empty()); -} - -#[tokio::test] -async fn captured_row_transport_only_allows_additional_routing_and_event_fields() { - for fixed in [false, true] { - let mut schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column("new", DataType::Int(IntType::with_nullable(false))) - .column("unused", DataType::Int(IntType::new())) - .column("pt", DataType::Int(IntType::new())) - .column( - "op", - DataType::VarChar(paimon::spec::VarCharType::string_type()), - ) - .partition_keys(["pt"]) - .option("rowkind.field", "op") - .option("ignore-delete", "true") - .option("data-evolution.enabled", (!fixed).to_string()) - .option("row-tracking.enabled", (!fixed).to_string()); - if fixed { - schema = schema.option("bucket", "4").option("bucket-key", "id"); - } - let (io, table) = memory_table( - "memory:/captured_row_transport", - TableSchema::new(0, &schema.build().unwrap()), - ); - setup_dirs(&io, table.location()).await; - persist_table_schema(&io, table.location(), table.schema()).await; - let builder = table.new_write_builder(); - let mut writer = builder.new_write().unwrap(); - writer.with_write_type(vec!["id".into()]).unwrap(); - let original = RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from(vec![11])) as ArrayRef), - ("new", Arc::new(Int32Array::from(vec![22])) as ArrayRef), - ("unused", Arc::new(Int32Array::from(vec![33])) as ArrayRef), - ("pt", Arc::new(Int32Array::from(vec![1])) as ArrayRef), - ( - "op", - Arc::new(arrow_array::StringArray::from(vec!["+I"])) as ArrayRef, - ), - ]) - .unwrap(); - writer.write_arrow_batch(&original).await.unwrap(); - writer.with_write_type(vec!["new".into()]).unwrap(); - let old = |partitions: Vec, kinds: Vec<&str>, extra_payload| { - let mut columns = vec![( - "id", - Arc::new(Int32Array::from(vec![11; partitions.len()])) as ArrayRef, - )]; - if extra_payload { - columns.push(( - "unused", - Arc::new(Int32Array::from(vec![99; partitions.len()])) as ArrayRef, - )); - } - columns.push(("pt", Arc::new(Int32Array::from(partitions)) as ArrayRef)); - columns.push(( - "op", - Arc::new(arrow_array::StringArray::from(kinds)) as ArrayRef, - )); - RecordBatch::try_from_iter(columns).unwrap() - }; - let before = io - .list_status_recursive(table.location()) - .await - .unwrap() - .len(); - // The old subset is not valid for the new partition; no group may be staged. - assert!(writer - .write_arrow_batch(&old(vec![1, 2], vec!["+I", "+I"], false)) - .await - .is_err()); - // A non-routing payload cannot be smuggled into a selected layout. - assert!(writer - .write_arrow_batch(&old(vec![1], vec!["+I"], true)) - .await - .is_err()); - assert_eq!( - io.list_status_recursive(table.location()) - .await - .unwrap() - .len(), - before - ); - // Missing the new factory's NOT NULL column is valid for the captured - // old subset. Ignored rows neither create a bucket nor inspect payloads. - writer - .write_arrow_batch(&old(vec![1, 2], vec!["+I", "-D"], false)) - .await - .unwrap(); - let mut columns = vec![]; - if fixed { - columns.push(("id", original.column(0).clone())); - } - columns.extend([ - ("new", original.column(1).clone()), - ("pt", Arc::new(Int32Array::from(vec![2])) as ArrayRef), - ("op", original.column(4).clone()), - ]); - writer - .write_arrow_batch(&RecordBatch::try_from_iter(columns).unwrap()) - .await - .unwrap(); - let messages = writer.prepare_commit().await.unwrap(); - assert_eq!( - messages - .iter() - .flat_map(|m| &m.new_files) - .map(|f| f.row_count) - .sum::(), - 3 - ); - builder.new_commit().commit(messages).await.unwrap(); - writer - .write_arrow_batch(&old(vec![1], vec!["+I"], false)) - .await - .unwrap(); - builder - .new_commit() - .commit(writer.prepare_commit().await.unwrap()) - .await - .unwrap(); - let mut reader = table.new_read_builder(); - reader - .with_projection(&["id", "new", "unused", "pt"]) - .unwrap(); - let plan = reader.new_scan().plan().await.unwrap(); - let batches: Vec = reader - .new_read() - .unwrap() - .to_arrow(plan.splits()) - .unwrap() - .try_collect() - .await - .unwrap(); - let mut actual = Vec::new(); - for batch in batches { - for row in 0..batch.num_rows() { - actual.push( - batch - .columns() - .iter() - .map(|column| { - let array = column.as_any().downcast_ref::().unwrap(); - array.is_valid(row).then(|| array.value(row)) - }) - .collect::>(), - ); - } - } - actual.sort(); - assert_eq!( - actual, - vec![ - vec![None, Some(22), None, Some(2)], - vec![Some(11), None, None, Some(1)], - vec![Some(11), None, None, Some(1)], - vec![Some(11), None, None, Some(1)], - ] - ); - } -} - -#[tokio::test] -async fn historical_row_layouts_do_not_randomly_validate_unwritten_routing_fields() { - for _ in 0..4 { - let schema = Schema::builder() - .column("id", DataType::Int(IntType::with_nullable(false))) - .column("new", DataType::Int(IntType::with_nullable(false))) - .column("pt", DataType::Int(IntType::with_nullable(false))) - .column( - "op", - DataType::VarChar(paimon::spec::VarCharType::string_type()), - ) - .partition_keys(["pt"]) - .option("rowkind.field", "op") - .option("ignore-delete", "true") - .option("row-tracking.enabled", "true") - .option("data-evolution.enabled", "true") - .build() - .unwrap(); - let (io, table) = memory_table("memory:/historical_routing", TableSchema::new(0, &schema)); - setup_dirs(&io, table.location()).await; - persist_table_schema(&io, table.location(), table.schema()).await; - let builder = table.new_write_builder(); - let mut writer = builder.new_write().unwrap(); - let input = |id, partition| { - RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from(vec![id])) as ArrayRef), - ("new", Arc::new(Int32Array::from(vec![99])) as ArrayRef), - ( - "pt", - Arc::new(Int32Array::from(vec![partition])) as ArrayRef, - ), - ( - "op", - Arc::new(arrow_array::StringArray::from(vec!["+I"])) as ArrayRef, - ), - ]) - .unwrap() - }; - // Omitted routing columns do not impose their physical NOT NULL rule. - writer.with_write_type(vec!["id".into()]).unwrap(); - writer.write_arrow_batch(&input(11, None)).await.unwrap(); - writer - .with_write_type(vec!["id".into(), "pt".into()]) - .unwrap(); - writer.write_arrow_batch(&input(22, Some(1))).await.unwrap(); - writer.with_write_type(vec!["new".into()]).unwrap(); - // id,pt,op can represent either historical type. Actual destinations - // decide whether pt is selected, rather than an arbitrary map entry. - let rows = RecordBatch::try_from_iter([ - ( - "id", - Arc::new(Int32Array::from(vec![33, 44, 55, 66])) as ArrayRef, - ), - ( - "pt", - Arc::new(Int32Array::from(vec![None, Some(1), None, Some(1)])) as ArrayRef, - ), - ( - "op", - Arc::new(arrow_array::StringArray::from(vec!["+I", "+I", "-D", "-D"])) as ArrayRef, - ), - ]) - .unwrap(); - let ignored_null_payload = RecordBatch::try_from_iter([ - ("id", Arc::new(Int32Array::from(vec![None])) as ArrayRef), - ("pt", Arc::new(Int32Array::from(vec![None])) as ArrayRef), - ( - "op", - Arc::new(arrow_array::StringArray::from(vec!["-D"])) as ArrayRef, - ), - ]) - .unwrap(); - // Logical NOT NULL payload checks still precede ignore-delete. - assert!(writer - .write_arrow_batch(&ignored_null_payload) - .await - .is_err()); - fn require_send(future: T) -> T { - future - } - require_send(writer.write_arrow_batch(&rows)).await.unwrap(); - builder - .new_commit() - .commit(writer.prepare_commit().await.unwrap()) - .await - .unwrap(); - let mut reader = table.new_read_builder(); - reader.with_projection(&["id", "pt"]).unwrap(); - let plan = reader.new_scan().plan().await.unwrap(); - let batches: Vec = reader - .new_read() - .unwrap() - .to_arrow(plan.splits()) - .unwrap() - .try_collect() - .await - .unwrap(); - let mut actual = Vec::new(); - for batch in batches { - for row in 0..batch.num_rows() { - actual.push( - batch - .columns() - .iter() - .map(|column| { - let array = column.as_any().downcast_ref::().unwrap(); - array.is_valid(row).then(|| array.value(row)) - }) - .collect::>(), - ); - } - } - actual.sort(); - assert_eq!( - actual, - vec![ - vec![Some(11), None], - vec![Some(22), Some(1)], - vec![Some(33), None], - vec![Some(44), Some(1)] - ] - ); - } -}