diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index e0dce1757..1ac21b058 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(), }) } } @@ -366,9 +367,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, } } @@ -464,6 +471,7 @@ impl PyPostponeFixedBucketWriteBuilder { ))), table_location: context.table.location().to_string(), commit_user: context.commit_user.clone(), + blob_consumers: Vec::new(), }, prepared: false, }) @@ -478,12 +486,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 +507,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 +873,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 +887,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.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())?; + } + 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 write_result(&self, result: paimon::Result<()>) -> PyResult<()> { result.map_err(to_callback_error) } @@ -900,14 +950,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, @@ -1213,6 +1267,14 @@ impl PyBatchTableUpdate { #[pymethods] impl PyBatchTableWrite { + 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 +1307,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))] @@ -1264,7 +1327,7 @@ impl PyBatchTableWrite { )); } self.prepared = true; - self.state.prepare_commit(py) + self.state.prepare_commit(py, i64::MAX) } } @@ -1279,6 +1342,14 @@ pub struct PyStreamTableWrite { #[pymethods] impl PyStreamTableWrite { + 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 +1382,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))] @@ -1324,15 +1396,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) } } @@ -1687,6 +1759,239 @@ 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(); + 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(); + 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 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| { @@ -2208,7 +2513,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..14284b3f6 100644 --- a/bindings/python/tests/test_write.py +++ b/bindings/python/tests/test_write.py @@ -66,27 +66,71 @@ 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(["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) - 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(["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 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": [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/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 index 06b6a80c1..800f5b83b 100644 --- a/crates/paimon/src/table/internal.rs +++ b/crates/paimon/src/table/internal.rs @@ -15,15 +15,17 @@ // specific language governing permissions and limitations // under the License. -//! Unstable scheduling bridge for the repository's language bindings. +//! Unstable execution bridge for the repository's language bindings. //! //! Source plans are the public input to search readers. Candidate and scoring //! requests below are internal execution stages: Core owns their ordering and -//! global result coordination. This module is available only with `internal-api`. +//! global result coordination. Writer adapters forward their existing commit +//! identifiers through the same bridge. This module is available only with `internal-api`. pub use super::search_task::{SearchTask, SearchTaskExecutor, SearchTaskResult}; use super::{ - BatchVectorRead, GlobalIndexResult, HybridSearchBuilder, Table, VectorRead, VectorScanPlan, + BatchVectorRead, CommitMessage, GlobalIndexResult, HybridSearchBuilder, Table, TableWrite, + VectorRead, VectorScanPlan, }; pub async fn execute_search_task( @@ -65,3 +67,13 @@ pub async fn execute_hybrid_with_executor( ) -> crate::Result { builder.execute_with_executor(executor).await } + +/// 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/table_write.rs b/crates/paimon/src/table/table_write.rs index 39842497a..1e8c2217d 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -54,9 +54,27 @@ use std::sync::Arc; #[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)] +struct AppendWriteConfig { + 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. enum FileWriter { Append(DataFileWriter), @@ -143,6 +161,7 @@ pub struct TableWrite { write_fields: Vec, written: bool, partition_writers: HashMap, + append_write_configs: HashMap, partition_computer: PartitionComputer, partition_keys: Vec, schema_id: i64, @@ -213,6 +232,7 @@ impl TableWrite { write_fields: schema.fields().to_vec(), written: false, partition_writers: HashMap::new(), + append_write_configs: HashMap::new(), partition_computer: PartitionComputer::new( schema.partition_keys(), schema.fields(), @@ -513,6 +533,7 @@ impl TableWrite { write_fields: fields.to_vec(), written: false, partition_writers: HashMap::new(), + append_write_configs: HashMap::new(), partition_computer, partition_keys, schema_id: schema.id(), @@ -641,15 +662,11 @@ 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. + /// 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.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(), @@ -683,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() @@ -728,21 +754,15 @@ impl TableWrite { 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(), @@ -1457,19 +1477,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(); @@ -1531,16 +1602,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.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 @@ -1561,7 +1697,19 @@ 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(|| 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 } else if bucket == POSTPONE_BUCKET { self.create_postpone_writer(partition_path, bucket)? } else { @@ -1584,7 +1732,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() @@ -1613,8 +1766,8 @@ impl TableWrite { &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()), ))) @@ -1800,7 +1953,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 +1995,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_input.rs b/crates/paimon/src/table/table_write_input.rs index b2e99b4ce..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() 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..0389cd447 --- /dev/null +++ b/crates/paimon/src/table/table_write_lifecycle_tests.rs @@ -0,0 +1,354 @@ +// 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_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(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); + write(&mut writer, false, 1).await; + assert_columns(&writer.prepare_commit().await.unwrap(), "id"); + 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(), "id"); + } +} + +#[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", + ); + 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()); + assert!(writer.append_write_configs.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_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(); + 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(); + 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(); + 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(), + "id", + ); +} + +#[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/filtered_write_validation_test.rs b/crates/paimon/tests/filtered_write_validation_test.rs index 4d51a78ed..97d0deee0 100644 --- a/crates/paimon/tests/filtered_write_validation_test.rs +++ b/crates/paimon/tests/filtered_write_validation_test.rs @@ -364,8 +364,11 @@ 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(); + 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(); assert_eq!( diff --git a/crates/paimon/tests/table_partial_write_test.rs b/crates/paimon/tests/table_partial_write_test.rs index fd0c87b32..2e305d27f 100644 --- a/crates/paimon/tests/table_partial_write_test.rs +++ b/crates/paimon/tests/table_partial_write_test.rs @@ -172,6 +172,9 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { .await .unwrap(); let messages = writer.prepare_commit().await.unwrap(); + writer + .with_write_type(vec!["pt".into(), "id".into()]) + .unwrap(); assert!(writer .with_write_type(vec!["id".into(), "omitted".into(), "pt".into()]) .is_err()); @@ -179,6 +182,211 @@ async fn invalid_write_type_does_not_replace_previous_configuration() { 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; @@ -336,3 +544,59 @@ async fn selected_columns_validate_input_before_staging() { 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"] + ); +}