This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 77b1d49a [core] Preserve append writer configuration per bucket across
prepares (#1188)
77b1d49a is described below
commit 77b1d49a7c414dac0c79af5eb24da3732777d122
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Oct 22 13:16:14 2026 +0800
[core] Preserve append writer configuration per bucket across prepares
(#1188)
---
bindings/python/src/write.rs | 339 +++++++++++++++++++-
bindings/python/tests/test_blob_consumer.py | 39 ++-
bindings/python/tests/test_write.py | 62 +++-
crates/paimon/src/table/blob_consumer_tests.rs | 4 +-
crates/paimon/src/table/internal.rs | 18 +-
crates/paimon/src/table/table_write.rs | 202 ++++++++++--
crates/paimon/src/table/table_write_input.rs | 60 +++-
.../src/table/table_write_lifecycle_tests.rs | 354 +++++++++++++++++++++
.../paimon/tests/filtered_write_validation_test.rs | 7 +-
crates/paimon/tests/table_partial_write_test.rs | 264 +++++++++++++++
10 files changed, 1273 insertions(+), 76 deletions(-)
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index e0dce175..1ac21b05 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<Vec<CommitMessage>> {
+ async fn prepare_commit(
+ &mut self,
+ commit_identifier: i64,
+ ) -> paimon::Result<Vec<CommitMessage>> {
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<WriteTarget>,
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<Arc<PythonBlobConsumer>>,
}
/// Descriptor serialization and exception transport only; record boundaries,
/// callback timing and requested flushes belong to Rust core.
struct PythonBlobConsumer {
- callback: Py<PyAny>,
+ callback: Mutex<Option<Py<PyAny>>>,
}
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<dyn BlobConsumer>),
)
.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::<Vec<_>>();
+ // 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<Vec<PyCommitMessage>> {
+ fn prepare_commit(
+ &mut self,
+ py: Python<'_>,
+ commit_identifier: i64,
+ ) -> PyResult<Vec<PyCommitMessage>> {
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<Vec<PyCommitMessage>> {
- 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::<PyRefMut<'_,
PyBatchTableWrite>>().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::<PyRefMut<'_, PyStreamTableWrite>>()
+ .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 d2e8f068..d86a2ee8 100644
--- a/bindings/python/tests/test_blob_consumer.py
+++ b/bindings/python/tests/test_blob_consumer.py
@@ -24,20 +24,25 @@ import pytest
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 test_blob_consumer_descriptor_bytes_and_nulls(tmp_path,
stream):
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 a915f314..14284b3f 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"]}
+
+
[email protected]("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 9c3ef0a5..d11b28ee 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::<i64>(),
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 06b6a80c..800f5b83 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<GlobalIndexResult> {
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<Vec<CommitMessage>> {
+ 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 39842497..1e8c2217 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<Arc<dyn crate::spec::BlobConsumer>>,
+ blob_uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ last_modified_prepare: Option<AppendPrepareIdentifier>,
+}
+
/// 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<DataField>,
written: bool,
partition_writers: HashMap<PartitionBucketKey, FileWriter>,
+ append_write_configs: HashMap<PartitionBucketKey, AppendWriteConfig>,
partition_computer: PartitionComputer,
partition_keys: Vec<String>,
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<String>) ->
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::<Result<Vec<_>>>()?;
+ 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<Arc<dyn crate::spec::BlobConsumer>>,
) -> 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<Vec<CommitMessage>> {
+ 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<Vec<CommitMessage>> {
+ self.prepare_commit_internal(Some(commit_identifier)).await
+ }
+
+ async fn prepare_commit_internal(
+ &mut self,
+ commit_identifier: Option<i64>,
+ ) -> Result<Vec<CommitMessage>> {
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::<HashSet<_>>();
+ 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<Option<i64>> {
+ // 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<FileWriter> {
+ fn create_append_writer(
+ &self,
+ partition_path: String,
+ bucket: i32,
+ config: &AppendWriteConfig,
+ ) -> Result<FileWriter> {
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 b2e99b4c..a5987da4 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::<HashSet<_>>();
+ let mut routing = self
+ .partition_keys
+ .iter()
+ .map(String::as_str)
+ .collect::<HashSet<_>>();
+ 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 00000000..0389cd44
--- /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<AtomicUsize>| {
+ Arc::new(move |_: &str, descriptor: Option<&BlobDescriptor>| {
+ if descriptor.is_some() {
+ calls.fetch_add(1, Ordering::Relaxed);
+ }
+ Ok(false)
+ }) as Arc<dyn crate::spec::BlobConsumer>
+ };
+ 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 4d51a78e..97d0deee 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 fd0c87b3..2e305d27 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<Option<i32>>, 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::<i64>(),
+ 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<RecordBatch> = read
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
1);
+ assert_eq!(
+ batches[0]
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 22
+ );
+ assert_eq!(
+ batches[0]
+ .column_by_name("pt")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .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<dyn paimon::spec::BlobConsumer>
+ };
+ 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"]
+ );
+}
