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 29727bfe Support BlobConsumer for native append writes (#1063)
29727bfe is described below
commit 29727bfe9bffd4a8d775999bc35015547772280c
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 21:18:10 2026 +0800
Support BlobConsumer for native append writes (#1063)
---
.../python/python/pypaimon_rust/datafusion.pyi | 9 +
bindings/python/src/write.rs | 110 +++++++-
bindings/python/tests/test_blob_consumer.py | 124 +++++++++
crates/paimon/src/arrow/format/blob.rs | 2 +-
crates/paimon/src/arrow/format/blob/writer.rs | 176 +++++++++++--
.../src/arrow/format/blob/writer_consumer_tests.rs | 286 +++++++++++++++++++++
crates/paimon/src/arrow/format/mod.rs | 21 +-
crates/paimon/src/io/file_io.rs | 198 +++++++++++++-
crates/paimon/src/spec/blob_consumer.rs | 41 +++
crates/paimon/src/spec/core_options.rs | 16 ++
crates/paimon/src/spec/mod.rs | 3 +
crates/paimon/src/table/blob_consumer_tests.rs | 247 ++++++++++++++++++
crates/paimon/src/table/data_file_writer.rs | 33 ++-
.../src/table/dedicated_format_file_writer.rs | 102 ++++++--
crates/paimon/src/table/managed_blob_writer.rs | 14 +-
crates/paimon/src/table/mod.rs | 2 +
crates/paimon/src/table/table_write.rs | 134 +++++++++-
17 files changed, 1452 insertions(+), 66 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index d76b942f..ebb823c8 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -237,6 +237,12 @@ class CommitMessage:
class BatchTableWrite:
def with_write_type(self, columns: List[str]) -> "BatchTableWrite": ...
+ def with_blob_consumer(
+ self, callback: Optional[Callable[[str, Optional[bytes]], bool]]
+ ) -> "BatchTableWrite":
+ """Receive serialized payload descriptors (None for NULL fields).
+ Returning true requests a flush after the complete record."""
+ ...
def close(self) -> None: ...
def write_arrow(self, batch: pyarrow.RecordBatch, bucket: Optional[int] =
None) -> None: ...
def prepare_commit(self) -> List[CommitMessage]:
@@ -292,6 +298,9 @@ class StreamTableUpdate:
class StreamTableWrite:
def with_write_type(self, columns: List[str]) -> "StreamTableWrite": ...
+ def with_blob_consumer(
+ self, callback: Optional[Callable[[str, Optional[bytes]], bool]]
+ ) -> "StreamTableWrite": ...
def close(self) -> None: ...
def write_arrow(self, batch: pyarrow.RecordBatch, bucket: Optional[int] =
None) -> None: ...
def prepare_commit(self, wait_compaction: bool, commit_identifier: int) ->
List[CommitMessage]:
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 20abb7df..bfd23051 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -21,7 +21,7 @@ use std::sync::{Arc, Mutex};
use arrow::datatypes::Schema as ArrowSchema;
use arrow::pyarrow::FromPyArrow;
use arrow::record_batch::RecordBatch;
-use paimon::spec::{CoreOptions, DataType, Datum};
+use paimon::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataType, Datum};
use paimon::table::{
CommitMessage, PostponeBucketPlan, PostponeFixedBucketTableWrite, Table,
TableCommit,
TableUpdate, TableUpdateByRowId, TableWrite,
COMMIT_MESSAGE_SERIALIZER_VERSION,
@@ -84,6 +84,7 @@ impl WriteContext {
))),
table_location: self.table.location().to_string(),
commit_user: self.commit_user.clone(),
+ blob_consumer: None,
})
}
}
@@ -314,6 +315,20 @@ enum WriteTarget {
}
impl WriteTarget {
+ fn with_blob_consumer(
+ &mut self,
+ consumer: Option<Arc<dyn BlobConsumer>>,
+ ) -> paimon::Result<()> {
+ match self {
+ Self::Table(writer) => {
+ writer.with_blob_consumer(consumer)?;
+ Ok(())
+ }
+ // Like Java AbstractFileStoreWrite, PK-only writers ignore it.
+ Self::PostponeFixed(_) => Ok(()),
+ }
+ }
+
fn with_write_type(&mut self, columns: Vec<String>) -> paimon::Result<()> {
match self {
Self::Table(writer) => {
@@ -444,6 +459,7 @@ impl PyPostponeFixedBucketWriteBuilder {
))),
table_location: context.table.location().to_string(),
commit_user: context.commit_user.clone(),
+ blob_consumer: None,
},
prepared: false,
})
@@ -458,6 +474,40 @@ struct WriteState {
inner: Option<WriteTarget>,
table_location: String,
commit_user: String,
+ blob_consumer: Option<Arc<PythonBlobConsumer>>,
+}
+
+/// Descriptor serialization and exception transport only; record boundaries,
+/// callback timing and requested flushes belong to Rust core.
+struct PythonBlobConsumer {
+ callback: Py<PyAny>,
+ error: Mutex<Option<PyErr>>,
+}
+
+impl BlobConsumer for PythonBlobConsumer {
+ fn accept(
+ &self,
+ field_name: &str,
+ descriptor: Option<&BlobDescriptor>,
+ ) -> paimon::Result<bool> {
+ let result = tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ let encoded =
+ descriptor.map(|descriptor| PyBytes::new(py,
&descriptor.serialize()));
+ self.callback
+ .bind(py)
+ .call1((field_name, encoded))?
+ .is_truthy()
+ })
+ });
+ result.map_err(|error| {
+ *self.error.lock().unwrap() = Some(error);
+ paimon::Error::DataInvalid {
+ message: "BLOB consumer failed".into(),
+ source: None,
+ }
+ })
+ }
}
fn wrap_messages(
@@ -661,6 +711,40 @@ impl UpdateContext {
}
impl WriteState {
+ fn with_blob_consumer(&mut self, callback: Option<Py<PyAny>>, py:
Python<'_>) -> PyResult<()> {
+ let consumer = callback
+ .map(|callback| {
+ if !callback.bind(py).is_callable() {
+ return Err(PyTypeError::new_err("blob_consumer must be
callable"));
+ }
+ Ok(Arc::new(PythonBlobConsumer {
+ callback,
+ error: Mutex::new(None),
+ }))
+ })
+ .transpose()?;
+ self.inner
+ .as_mut()
+ .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?
+ .with_blob_consumer(
+ consumer
+ .clone()
+ .map(|consumer| consumer as Arc<dyn BlobConsumer>),
+ )
+ .map_err(to_py_err)?;
+ self.blob_consumer = consumer;
+ Ok(())
+ }
+
+ fn write_result(&self, result: paimon::Result<()>) -> PyResult<()> {
+ if let Some(consumer) = &self.blob_consumer {
+ if let Some(error) = consumer.error.lock().unwrap().take() {
+ return Err(error);
+ }
+ }
+ result.map_err(to_py_err)
+ }
+
fn with_write_type(&mut self, columns: Vec<String>) -> PyResult<()> {
self.inner
.as_mut()
@@ -680,8 +764,8 @@ impl WriteState {
.inner
.as_mut()
.ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?;
- py.detach(|| runtime().block_on(inner.write_arrow(&batch, bucket)))
- .map_err(to_py_err)
+ let result = py.detach(|| runtime().block_on(inner.write_arrow(&batch,
bucket)));
+ self.write_result(result)
}
fn prepare_commit(&mut self, py: Python<'_>) ->
PyResult<Vec<PyCommitMessage>> {
@@ -898,6 +982,16 @@ impl PyBatchTableUpdate {
#[pymethods]
impl PyBatchTableWrite {
+ #[pyo3(signature = (callback))]
+ fn with_blob_consumer(
+ mut slf: PyRefMut<'_, Self>,
+ callback: Option<Py<PyAny>>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.state.with_blob_consumer(callback, py)?;
+ Ok(slf)
+ }
+
fn with_write_type(
mut slf: PyRefMut<'_, Self>,
columns: Vec<String>,
@@ -944,6 +1038,16 @@ pub struct PyStreamTableWrite {
#[pymethods]
impl PyStreamTableWrite {
+ #[pyo3(signature = (callback))]
+ fn with_blob_consumer(
+ mut slf: PyRefMut<'_, Self>,
+ callback: Option<Py<PyAny>>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.state.with_blob_consumer(callback, py)?;
+ Ok(slf)
+ }
+
fn with_write_type(
mut slf: PyRefMut<'_, Self>,
columns: Vec<String>,
diff --git a/bindings/python/tests/test_blob_consumer.py
b/bindings/python/tests/test_blob_consumer.py
new file mode 100644
index 00000000..d2e8f068
--- /dev/null
+++ b/bindings/python/tests/test_blob_consumer.py
@@ -0,0 +1,124 @@
+# 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.
+
+import struct
+from urllib.parse import urlparse
+
+import pyarrow as pa
+import pytest
+
+from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
+
+
+def _table(tmp_path):
+ 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 (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true',
+ 'blob.copy-buffer-size' = '2 B')""")
+ 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 _descriptor_payload(encoded):
+ version, magic, length = struct.unpack_from('<BQi', encoded)
+ assert (version, magic) == (2, 0x424C4F4244455343)
+ uri = encoded[13:13 + length].decode('utf-8')
+ offset, size = struct.unpack_from('<qq', encoded, 13 + length)
+ assert uri.endswith('.blob') and size >= 0
+ path = urlparse(uri).path if uri.startswith('file:') else uri
+ with open(path, 'rb') as source:
+ source.seek(offset)
+ return source.read(size)
+
+
[email protected]('stream', [False, True])
+def test_blob_consumer_descriptor_bytes_and_nulls(tmp_path, stream):
+ context, table = _table(tmp_path)
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ received = []
+
+ def callback(name, encoded):
+ received.append((name, encoded))
+ return True
+
+ 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)
+ 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 [_descriptor_payload(encoded) if encoded is not None else None
+ for _, encoded in received] == [b'payload', None, b'']
+ 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'']}
+ finally:
+ writer.close()
+ with pytest.raises(RuntimeError, match='closed'):
+ writer.with_blob_consumer(None)
+
+
[email protected]('stream', [False, True])
+def test_blob_consumer_exception_is_not_wrapped_or_retried(tmp_path, stream):
+ _, table = _table(tmp_path)
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ calls = []
+ failure = KeyError('callback failure')
+
+ def callback(name, encoded):
+ calls.append((name, encoded))
+ raise failure
+
+ try:
+ writer.with_blob_consumer(callback)
+ with pytest.raises(KeyError) as caught:
+ writer.write_arrow(_batch())
+ assert caught.value is failure
+ assert [name for name, _ in calls] == ['payload']
+ assert list(tmp_path.rglob('*.blob'))
+ assert _descriptor_payload(calls[0][1]) == b'payload'
+ assert not list(tmp_path.rglob('*.parquet'))
+ with pytest.raises(ValueError, match='write failure'):
+ writer.write_arrow(_batch())
+ assert [name for name, _ in calls] == ['payload']
+ finally:
+ writer.close()
+
+
+def test_blob_consumer_replacement_and_clear(tmp_path):
+ _, table = _table(tmp_path)
+ writer = table.new_batch_write_builder().new_write()
+ try:
+ writer.with_blob_consumer(lambda name, encoded: pytest.fail('cleared
callback invoked'))
+ with pytest.raises(TypeError, match='callable'):
+ writer.with_blob_consumer(123)
+ writer.with_blob_consumer(None)
+ writer.write_arrow(_batch())
+ finally:
+ writer.close()
diff --git a/crates/paimon/src/arrow/format/blob.rs
b/crates/paimon/src/arrow/format/blob.rs
index 36d538bb..608938db 100644
--- a/crates/paimon/src/arrow/format/blob.rs
+++ b/crates/paimon/src/arrow/format/blob.rs
@@ -43,7 +43,7 @@ use std::ops::Range;
use std::sync::Arc;
mod writer;
-pub(crate) use writer::BlobFormatWriter;
+pub(crate) use writer::{BlobFormatWriter, BlobWriterFactory};
pub(crate) struct BlobFormatReader {
descriptor_mode: bool,
diff --git a/crates/paimon/src/arrow/format/blob/writer.rs
b/crates/paimon/src/arrow/format/blob/writer.rs
index 7aedc97b..6b69c8fb 100644
--- a/crates/paimon/src/arrow/format/blob/writer.rs
+++ b/crates/paimon/src/arrow/format/blob/writer.rs
@@ -22,9 +22,9 @@ use super::{
BLOB_ARRAY_MAGIC_NUMBER, BLOB_ARRAY_VERSION, BLOB_FORMAT_VERSION,
BLOB_MAGIC_NUMBER_BYTES,
BLOB_MAP_MAGIC_NUMBER, BLOB_MAP_VERSION,
};
-use crate::arrow::format::{FormatFileWriter, FormatWriteResult};
-use crate::io::{FileIO, FileWrite, OutputFile};
-use crate::spec::{BlobDescriptor, DataField, DataType};
+use crate::arrow::format::{FormatFileWriter, FormatWriteResult,
FormatWriterFactory};
+use crate::io::{FileIO, FileRead, FileWrite, OutputFile};
+use crate::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataField,
DataType};
use crate::{Error, Result};
use arrow_array::{
Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array,
Int16Array, Int32Array,
@@ -34,14 +34,79 @@ use arrow_array::{
use async_trait::async_trait;
use bytes::Bytes;
use crc32fast::Hasher;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
-const COPY_BUFFER_SIZE: u64 = 8 * 1024 * 1024;
+// Remote FileRead ranges open new requests. Keep read-ahead independent of
+// the Java-sized copy buffer so a 4 KiB buffer does not mean 4 KiB GETs.
+const SOURCE_READ_SIZE: u64 = 8 * 1024 * 1024;
+
+#[cfg(test)]
+#[path = "writer_consumer_tests.rs"]
+mod consumer_tests;
+
+pub(crate) struct BlobWriterFactory {
+ file_io: Option<FileIO>,
+ field: Option<DataField>,
+ consumer: Option<Arc<dyn BlobConsumer>>,
+ copy_buffer_size: usize,
+}
+
+impl BlobWriterFactory {
+ pub(crate) fn new(
+ file_io: Option<FileIO>,
+ field: Option<&DataField>,
+ options: Option<&HashMap<String, String>>,
+ ) -> Result<Self> {
+ let defaults = HashMap::new();
+ Ok(Self {
+ file_io,
+ field: field.cloned(),
+ consumer: None,
+ copy_buffer_size: CoreOptions::new(options.unwrap_or(&defaults))
+ .blob_copy_buffer_size()?,
+ })
+ }
+
+ pub(crate) fn with_consumer(mut self, consumer: Arc<dyn BlobConsumer>) ->
Self {
+ self.consumer = Some(consumer);
+ self
+ }
+}
+
+#[async_trait]
+impl FormatWriterFactory for BlobWriterFactory {
+ async fn create_writer(
+ &self,
+ output: &OutputFile,
+ _compression: &str,
+ ) -> Result<Box<dyn FormatFileWriter>> {
+ let stream = if self.consumer.is_some() {
+ output.flushable_writer().await?
+ } else {
+ output.writer().await?
+ };
+ Ok(Box::new(
+ BlobFormatWriter::from_stream(
+ stream,
+ output,
+ self.file_io.clone(),
+ self.field.as_ref(),
+ )?
+ .with_copy_buffer_size(self.copy_buffer_size)
+ .with_consumer(self.consumer.clone()),
+ ))
+ }
+}
pub(crate) struct BlobFormatWriter {
writer: Box<dyn FileWrite>,
file_io: Option<FileIO>,
kind: BlobFieldKind,
+ path: String,
+ field_name: String,
+ consumer: Option<Arc<dyn BlobConsumer>>,
+ copy_buffer_size: usize,
bytes_written: u64,
lengths: Vec<i64>,
}
@@ -58,6 +123,15 @@ impl BlobFormatWriter {
output: &OutputFile,
file_io: Option<FileIO>,
field: Option<&DataField>,
+ ) -> Result<Self> {
+ Self::from_stream(output.writer().await?, output, file_io, field)
+ }
+
+ fn from_stream(
+ writer: Box<dyn FileWrite>,
+ output: &OutputFile,
+ file_io: Option<FileIO>,
+ field: Option<&DataField>,
) -> Result<Self> {
let kind = match field {
Some(field) =>
validate_read_fields(std::slice::from_ref(field))?.unwrap(),
@@ -67,14 +141,43 @@ impl BlobFormatWriter {
validate_key_type(key_type)?;
}
Ok(Self {
- writer: output.writer().await?,
+ writer,
file_io,
kind,
+ path: output.location().to_string(),
+ field_name: field.map_or_else(String::new, |field|
field.name().to_string()),
+ consumer: None,
+ copy_buffer_size: 4 * 1024,
bytes_written: 0,
lengths: Vec::new(),
})
}
+ pub(crate) fn with_copy_buffer_size(mut self, size: usize) -> Self {
+ debug_assert!(size > 0);
+ self.copy_buffer_size = size;
+ self
+ }
+
+ fn with_consumer(mut self, consumer: Option<Arc<dyn BlobConsumer>>) ->
Self {
+ self.consumer = consumer;
+ self
+ }
+
+ fn accept(&self, offset: u64, length: i64) -> Result<bool> {
+ match &self.consumer {
+ Some(consumer) => consumer.accept(
+ &self.field_name,
+ Some(&BlobDescriptor::new(
+ self.path.clone(),
+ i64::try_from(offset).map_err(|_| invalid("BLOB offset
exceeds i64"))?,
+ length,
+ )),
+ ),
+ None => Ok(false),
+ }
+ }
+
/// The callback range excludes the common entry magic and trailer.
pub(crate) async fn write_managed_value(&mut self, value: &[u8]) ->
Result<(i64, i64)> {
if !matches!(self.kind, BlobFieldKind::Scalar) {
@@ -129,14 +232,19 @@ impl BlobFormatWriter {
let (start, mut hasher) = self.start_record().await?;
let length = self.write_payload(value, &mut hasher).await?;
self.finish_record(start, hasher).await?;
+ if self.accept(start + 4, length)? {
+ self.writer.flush().await?;
+ }
Ok(length)
}
/// Resolve descriptors without materializing a complete referenced object.
async fn write_payload(&mut self, value: &[u8], hasher: &mut Hasher) ->
Result<i64> {
if !BlobDescriptor::is_blob_descriptor(value) {
- self.write_bytes(Bytes::copy_from_slice(value), hasher)
- .await?;
+ for chunk in value.chunks(self.copy_buffer_size) {
+ self.write_bytes(Bytes::copy_from_slice(chunk), hasher)
+ .await?;
+ }
return i64::try_from(value.len()).map_err(|_| invalid("BLOB
payload exceeds i64"));
}
let descriptor = BlobDescriptor::deserialize(value)?;
@@ -168,9 +276,22 @@ impl BlobFormatWriter {
return Ok(0);
}
let reader = input.reader_for_range(offset..end).await?;
+ self.copy_payload(reader.as_ref(), offset..end, hasher,
descriptor.uri())
+ .await?;
+ Ok(length_i64)
+ }
+
+ async fn copy_payload(
+ &mut self,
+ reader: &dyn FileRead,
+ range: std::ops::Range<u64>,
+ hasher: &mut Hasher,
+ uri: &str,
+ ) -> Result<()> {
+ let (offset, end) = (range.start, range.end);
let mut position = offset;
while position < end {
- let chunk_end = position.saturating_add(COPY_BUFFER_SIZE).min(end);
+ let chunk_end = position.saturating_add(SOURCE_READ_SIZE).min(end);
let chunk =
reader
.read(position..chunk_end)
@@ -178,33 +299,42 @@ impl BlobFormatWriter {
.map_err(|error| Error::UnexpectedError {
message: format!(
"Failed to read BlobDescriptor '{}' range
{position}..{chunk_end}: {error}",
-
crate::io::uri_reader::sanitize_blob_uri(descriptor.uri())
+ crate::io::uri_reader::sanitize_blob_uri(uri)
),
source: Some(Box::new(error)),
})?;
if chunk.len() as u64 != chunk_end - position {
return Err(invalid(format!("Failed to read BlobDescriptor
'{}': short read for range {position}..{chunk_end}, expected={} bytes,
actual={} bytes",
-
crate::io::uri_reader::sanitize_blob_uri(descriptor.uri()), chunk_end -
position, chunk.len())));
+ crate::io::uri_reader::sanitize_blob_uri(uri), chunk_end -
position, chunk.len())));
+ }
+ for start in (0..chunk.len()).step_by(self.copy_buffer_size) {
+ let end =
start.saturating_add(self.copy_buffer_size).min(chunk.len());
+ self.write_bytes(chunk.slice(start..end), hasher).await?;
}
- self.write_bytes(chunk, hasher).await?;
position = chunk_end;
}
- Ok(length_i64)
+ Ok(())
}
async fn write_blob_values(
&mut self,
values: &LargeBinaryArray,
hasher: &mut Hasher,
- ) -> Result<Vec<i64>> {
+ ) -> Result<(Vec<i64>, bool)> {
let mut lengths = Vec::with_capacity(values.len());
+ let mut flush = false;
for value in values {
lengths.push(match value {
None => -1,
- Some(value) => self.write_payload(value, hasher).await?,
+ Some(value) => {
+ let position = self.bytes_written;
+ let length = self.write_payload(value, hasher).await?;
+ flush |= self.accept(position, length)?;
+ length
+ }
});
}
- Ok(lengths)
+ Ok((lengths, flush))
}
async fn write_index(&mut self, values: &[i64], hasher: &mut Hasher) ->
Result<i32> {
@@ -245,7 +375,7 @@ impl BlobFormatWriter {
} else {
None
};
- let value_lengths = self.write_blob_values(values, &mut hasher).await?;
+ let (value_lengths, flush) = self.write_blob_values(values, &mut
hasher).await?;
let key_index_length = if let Some(lengths) = key_lengths {
Some(self.write_index(&lengths, &mut hasher).await?)
} else {
@@ -261,7 +391,11 @@ impl BlobFormatWriter {
&mut hasher,
)
.await?;
- self.finish_record(start, hasher).await
+ self.finish_record(start, hasher).await?;
+ if flush {
+ self.writer.flush().await?;
+ }
+ Ok(())
}
}
@@ -276,6 +410,10 @@ impl FormatFileWriter for BlobFormatWriter {
for row in 0..batch.num_rows() {
if column.is_null(row) {
self.lengths.push(-1);
+ if let Some(consumer) = &self.consumer {
+ // Java ignores the return value for NULL fields.
+ consumer.accept(&self.field_name, None)?;
+ }
continue;
}
match &kind {
@@ -314,7 +452,7 @@ impl FormatFileWriter for BlobFormatWriter {
0
}
async fn flush(&mut self) -> Result<()> {
- Ok(())
+ self.writer.flush().await
}
async fn close(mut self: Box<Self>) -> Result<FormatWriteResult> {
let index = encode_delta_varints_write(&self.lengths);
diff --git a/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
b/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
new file mode 100644
index 00000000..cd0df2f6
--- /dev/null
+++ b/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
@@ -0,0 +1,286 @@
+// 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::io::FileIOBuilder;
+use arrow_array::ArrayRef;
+use std::sync::Mutex;
+
+#[derive(Default)]
+struct OutputState {
+ bytes: Vec<u8>,
+ writes: Vec<usize>,
+ flushes: Vec<usize>,
+ fail_flush: bool,
+}
+
+struct RecordingOutput(Arc<Mutex<OutputState>>);
+
+#[async_trait]
+impl FileWrite for RecordingOutput {
+ async fn write(&mut self, bytes: Bytes) -> Result<()> {
+ let mut output = self.0.lock().unwrap();
+ output.writes.push(bytes.len());
+ output.bytes.extend_from_slice(&bytes);
+ Ok(())
+ }
+
+ async fn flush(&mut self) -> Result<()> {
+ let mut output = self.0.lock().unwrap();
+ let position = output.bytes.len();
+ output.flushes.push(position);
+ if output.fail_flush {
+ Err(invalid("flush failed"))
+ } else {
+ Ok(())
+ }
+ }
+
+ async fn close(&mut self) -> Result<()> {
+ Ok(())
+ }
+}
+
+fn writer(kind: BlobFieldKind, output: &Arc<Mutex<OutputState>>) ->
BlobFormatWriter {
+ BlobFormatWriter {
+ writer: Box::new(RecordingOutput(output.clone())),
+ file_io: None,
+ kind,
+ path: "memory:/payload.blob".into(),
+ field_name: "payload".into(),
+ consumer: None,
+ copy_buffer_size: 4 * 1024,
+ bytes_written: 0,
+ lengths: Vec::new(),
+ }
+}
+
+type Event = (Option<BlobDescriptor>, usize);
+
+struct RecordingSource {
+ bytes: Bytes,
+ ranges: Mutex<Vec<std::ops::Range<u64>>>,
+}
+
+#[async_trait]
+impl FileRead for RecordingSource {
+ async fn read(&self, range: std::ops::Range<u64>) -> Result<Bytes> {
+ self.ranges.lock().unwrap().push(range.clone());
+ Ok(self.bytes.slice(range.start as usize..range.end as usize))
+ }
+}
+
+fn consumer(
+ events: &Arc<Mutex<Vec<Event>>>,
+ output: &Arc<Mutex<OutputState>>,
+) -> Arc<dyn BlobConsumer> {
+ let events = events.clone();
+ let output = output.clone();
+ Arc::new(move |field: &str, descriptor: Option<&BlobDescriptor>| {
+ assert_eq!(field, "payload");
+ let position = output.lock().unwrap().bytes.len();
+ events.lock().unwrap().push((descriptor.cloned(), position));
+ // Only the first callback requests a flush. Later callbacks must
+ // still run, and a collection must remember the earlier request.
+ Ok(events.lock().unwrap().len() == 1 || descriptor.is_none())
+ })
+}
+
+fn check_payloads(output: &OutputState, events: &[Event], expected:
&[Option<&[u8]>]) {
+ assert_eq!(events.len(), expected.len());
+ for ((descriptor, _), expected) in events.iter().zip(expected) {
+ match (descriptor, expected) {
+ (Some(descriptor), Some(payload)) => {
+ assert_eq!(descriptor.uri(), "memory:/payload.blob");
+ assert_eq!(descriptor.length(), payload.len() as i64);
+ let start = descriptor.offset() as usize;
+ assert_eq!(&output.bytes[start..start + payload.len()],
*payload);
+ }
+ (None, None) => {}
+ other => panic!("unexpected callback: {other:?}"),
+ }
+ }
+}
+
+#[tokio::test]
+async fn scalar_consumer_runs_after_record_and_null_does_not_flush() {
+ let output = Arc::new(Mutex::new(OutputState::default()));
+ let events = Arc::new(Mutex::new(Vec::new()));
+ let mut writer =
+ writer(BlobFieldKind::Scalar,
&output).with_consumer(Some(consumer(&events, &output)));
+ let batch = RecordBatch::try_from_iter([(
+ "payload",
+ Arc::new(LargeBinaryArray::from(vec![
+ Some(b"abc".as_slice()),
+ Some(b""),
+ None,
+ ])) as ArrayRef,
+ )])
+ .unwrap();
+ writer.write(&batch).await.unwrap();
+ let output = output.lock().unwrap();
+ let events = events.lock().unwrap();
+ check_payloads(&output, &events, &[Some(b"abc"), Some(b""), None]);
+ assert_eq!(
+ events
+ .iter()
+ .map(|(_, position)| *position)
+ .collect::<Vec<_>>(),
+ [19, 35, 35]
+ );
+ assert_eq!(output.flushes, [19]);
+ assert_eq!(writer.lengths, [19, 16, -1]);
+}
+
+#[tokio::test]
+async fn collection_consumer_skips_null_elements_and_flushes_after_trailer() {
+ for keys in [
+ None,
+ Some(vec![
+ b"key-a".to_vec(),
+ b"key-null".to_vec(),
+ b"key-b".to_vec(),
+ ]),
+ ] {
+ let output = Arc::new(Mutex::new(OutputState::default()));
+ let events = Arc::new(Mutex::new(Vec::new()));
+ let mut writer =
+ writer(BlobFieldKind::Array,
&output).with_consumer(Some(consumer(&events, &output)));
+ let values = LargeBinaryArray::from(vec![Some(b"aaa".as_slice()),
None, Some(b"bb")]);
+ writer.write_collection(&values, keys).await.unwrap();
+ let bytes_after_record = writer.bytes_written as usize;
+ writer
+
.write_collection(&LargeBinaryArray::from(Vec::<Option<&[u8]>>::new()), None)
+ .await
+ .unwrap();
+ let output = output.lock().unwrap();
+ let events = events.lock().unwrap();
+ check_payloads(&output, &events, &[Some(b"aaa"), Some(b"bb")]);
+ assert!(events
+ .iter()
+ .all(|(_, position)| *position < bytes_after_record - 12));
+ assert_eq!(output.flushes, [bytes_after_record]);
+ }
+}
+
+#[tokio::test]
+async fn callback_and_flush_errors_stop_writing() {
+ for fail_flush in [false, true] {
+ let output = Arc::new(Mutex::new(OutputState {
+ fail_flush,
+ ..Default::default()
+ }));
+ let calls = Arc::new(Mutex::new(0));
+ let callback_calls = calls.clone();
+ let consumer = Arc::new(move |_: &str, _: Option<&BlobDescriptor>| {
+ *callback_calls.lock().unwrap() += 1;
+ if fail_flush {
+ Ok(true)
+ } else {
+ Err(invalid("callback failed"))
+ }
+ });
+ let mut writer = writer(BlobFieldKind::Scalar,
&output).with_consumer(Some(consumer));
+ let batch = RecordBatch::try_from_iter([(
+ "payload",
+ Arc::new(LargeBinaryArray::from_iter_values([b"a", b"b"])) as
ArrayRef,
+ )])
+ .unwrap();
+ let error = writer.write(&batch).await.unwrap_err();
+ assert!(error.to_string().contains(if fail_flush {
+ "flush failed"
+ } else {
+ "callback failed"
+ }));
+ assert_eq!(*calls.lock().unwrap(), 1);
+ assert_eq!(writer.lengths, [17]);
+ }
+}
+
+#[tokio::test]
+async fn copy_buffer_bounds_inline_and_descriptor_payload_chunks() {
+ let io = FileIOBuilder::new("memory").build().unwrap();
+ let source = "memory:/blob-consumer-copy-source";
+ let payload = b"0123456789abcdef";
+ io.new_output(source)
+ .unwrap()
+ .write(Bytes::copy_from_slice(payload))
+ .await
+ .unwrap();
+ for descriptor in [false, true] {
+ let output = Arc::new(Mutex::new(OutputState::default()));
+ let mut writer = writer(BlobFieldKind::Scalar,
&output).with_copy_buffer_size(5);
+ writer.file_io = Some(io.clone());
+ let encoded = BlobDescriptor::new(source.into(), 2, 12).serialize();
+ let value = if descriptor {
+ encoded.as_slice()
+ } else {
+ &payload[2..14]
+ };
+ writer.write_scalar(value).await.unwrap();
+ let output = output.lock().unwrap();
+ assert_eq!(&output.bytes[4..16], &payload[2..14]);
+ assert_eq!(output.writes, [4, 5, 5, 2, 8, 4]);
+ }
+}
+
+#[test]
+fn copy_buffer_size_matches_java_default_and_positive_int_bounds() {
+ let defaults = HashMap::new();
+ assert_eq!(
+ CoreOptions::new(&defaults).blob_copy_buffer_size().unwrap(),
+ 4096
+ );
+ for (text, bytes) in [
+ ("1 B", 1),
+ ("8 kb", 8192),
+ ("2147483647 B", i32::MAX as usize),
+ ] {
+ let options = HashMap::from([("blob.copy-buffer-size".into(),
text.into())]);
+ assert_eq!(
+ CoreOptions::new(&options).blob_copy_buffer_size().unwrap(),
+ bytes
+ );
+ }
+ for text in ["0 B", "-1 B", "2147483648 B", "no-size"] {
+ let options = HashMap::from([("blob.copy-buffer-size".into(),
text.into())]);
+ assert!(BlobWriterFactory::new(None, None, Some(&options)).is_err());
+ }
+}
+
+#[tokio::test]
+async fn small_copy_buffer_does_not_reopen_a_remote_source_per_chunk() {
+ let size = SOURCE_READ_SIZE as usize + 17;
+ let source = RecordingSource {
+ bytes: Bytes::from(vec![b'x'; size]),
+ ranges: Mutex::new(Vec::new()),
+ };
+ let output = Arc::new(Mutex::new(OutputState::default()));
+ let mut writer = writer(BlobFieldKind::Scalar,
&output).with_copy_buffer_size(4096);
+ writer
+ .copy_payload(&source, 0..size as u64, &mut Hasher::new(), "source")
+ .await
+ .unwrap();
+ assert_eq!(
+ *source.ranges.lock().unwrap(),
+ [0..SOURCE_READ_SIZE, SOURCE_READ_SIZE..size as u64]
+ );
+ let output = output.lock().unwrap();
+ assert_eq!(output.bytes.len(), size);
+ assert!(output.writes.iter().all(|size| *size <= 4096));
+ assert_eq!(output.bytes.as_slice(), source.bytes.as_ref());
+}
diff --git a/crates/paimon/src/arrow/format/mod.rs
b/crates/paimon/src/arrow/format/mod.rs
index 8e141314..dc42feaf 100644
--- a/crates/paimon/src/arrow/format/mod.rs
+++ b/crates/paimon/src/arrow/format/mod.rs
@@ -521,10 +521,16 @@ pub(crate) fn create_format_writer_factory(
crate::spec::map_shredding::validate_format("file.format",
file_format)?;
}
}
+ if file_format.eq_ignore_ascii_case("blob") {
+ return Ok(Arc::new(blob::BlobWriterFactory::new(
+ file_io,
+ write_fields.and_then(|fields| fields.first()),
+ format_options,
+ )?));
+ }
Ok(Arc::new(PlainFormatWriterFactory {
schema,
zstd_level,
- file_io,
write_fields: write_fields.map(<[_]>::to_vec),
format_options: format_options.cloned(),
}))
@@ -557,7 +563,6 @@ pub(crate) async fn create_format_writer(
struct PlainFormatWriterFactory {
schema: SchemaRef,
zstd_level: i32,
- file_io: Option<FileIO>,
write_fields: Option<Vec<DataField>>,
format_options: Option<HashMap<String, String>>,
}
@@ -571,21 +576,11 @@ impl FormatWriterFactory for PlainFormatWriterFactory {
) -> crate::Result<Box<dyn FormatFileWriter>> {
let schema = self.schema.clone();
let zstd_level = self.zstd_level;
- let file_io = self.file_io.clone();
let write_fields = self.write_fields.as_deref();
let format_options = self.format_options.as_ref();
let path = output.location();
let lower = path.to_ascii_lowercase();
- if lower.ends_with(".blob") {
- Ok(Box::new(
- blob::BlobFormatWriter::new(
- output,
- file_io,
- write_fields.and_then(|fields| fields.first()),
- )
- .await?,
- ))
- } else if lower.ends_with(".orc") {
+ if lower.ends_with(".orc") {
if !matches!(
compression.to_ascii_lowercase().as_str(),
"" | "none" | "uncompressed"
diff --git a/crates/paimon/src/io/file_io.rs b/crates/paimon/src/io/file_io.rs
index dbcebd60..d7e6ce19 100644
--- a/crates/paimon/src/io/file_io.rs
+++ b/crates/paimon/src/io/file_io.rs
@@ -30,9 +30,10 @@ use bytes::Bytes;
use chrono::{DateTime, Utc};
use futures::stream::BoxStream;
use futures::{StreamExt, TryStreamExt};
-use opendal::raw::{normalize_path, normalize_root};
+use opendal::raw::{new_std_io_error, normalize_path, normalize_root};
use opendal::Operator;
use sha2::{Digest, Sha256};
+use tokio::io::AsyncWriteExt;
use tokio_util::compat::FuturesAsyncWriteCompatExt;
use url::Url;
@@ -346,14 +347,22 @@ impl FileIO {
fn file_source(&self, path: &str) -> crate::Result<FileSource> {
match &self.backend {
FileIOBackend::Provider(provider) =>
Ok(FileSource::Provider(provider.clone())),
- FileIOBackend::Storage(_) => {
+ FileIOBackend::Storage(storage) => {
let (op, relative_path) = self.create_static(path)?;
let namespace = self.cache_namespace_for_path(path)?;
let cache_path = cache_object_path(&namespace, &op,
&relative_path);
+ #[cfg(feature = "storage-fs")]
+ let local_fs = matches!(storage.as_ref(), Storage::LocalFs {
.. });
+ #[cfg(not(feature = "storage-fs"))]
+ let local_fs = {
+ let _ = storage;
+ false
+ };
Ok(FileSource::Static {
op,
relative_path,
cache_path,
+ local_fs,
})
}
}
@@ -1084,6 +1093,12 @@ impl FileRead for InputFileReader {
pub trait FileWrite: Send + Unpin + 'static {
async fn write(&mut self, bs: Bytes) -> crate::Result<()>;
+ /// Drain stream-level buffering without closing or publishing the file.
+ /// Writers whose `write` already drains their stream need no extra work.
+ async fn flush(&mut self) -> crate::Result<()> {
+ Ok(())
+ }
+
async fn close(&mut self) -> crate::Result<()>;
}
@@ -1105,12 +1120,121 @@ struct CacheInvalidatingWriter {
path: String,
}
+/// OpenDAL's local position writer retains its last write until close. Use a
+/// flushable stream for the built-in local filesystem only; custom operators
+/// and providers must continue to execute their own storage implementations.
+struct LocalFileWriter(tokio::io::BufWriter<tokio::fs::File>);
+
+impl LocalFileWriter {
+ async fn new(op: &Operator, relative_path: &str) -> crate::Result<Self> {
+ let relative_path = normalize_path(relative_path);
+ let relative =
std::path::Path::new(relative_path.trim_end_matches('/'));
+ if relative
+ .components()
+ .any(|component| matches!(component,
std::path::Component::ParentDir))
+ {
+ return Err(opendal::Error::new(
+ opendal::ErrorKind::NotFound,
+ "path escapes the configured root via `..`",
+ )
+ .into());
+ }
+ let path = std::path::Path::new(&op.info().root()).join(relative);
+ if let Some(parent) = path.parent() {
+ tokio::fs::create_dir_all(parent)
+ .await
+ .map_err(new_std_io_error)?;
+ }
+ let file = tokio::fs::OpenOptions::new()
+ .create(true)
+ .write(true)
+ .truncate(true)
+ .open(path)
+ .await
+ .map_err(new_std_io_error)?;
+ Ok(Self(tokio::io::BufWriter::with_capacity(
+ 8 * 1024 * 1024,
+ file,
+ )))
+ }
+}
+
+#[async_trait::async_trait]
+impl FileWrite for LocalFileWriter {
+ async fn write(&mut self, bytes: Bytes) -> crate::Result<()> {
+ self.0.write_all(&bytes).await.map_err(new_std_io_error)?;
+ Ok(())
+ }
+
+ async fn flush(&mut self) -> crate::Result<()> {
+ self.0.flush().await.map_err(new_std_io_error)?;
+ Ok(())
+ }
+
+ async fn close(&mut self) -> crate::Result<()> {
+ self.flush().await?;
+ self.0
+ .get_ref()
+ .sync_all()
+ .await
+ .map_err(new_std_io_error)?;
+ Ok(())
+ }
+}
+
+/// Buffer outside OpenDAL's write generator so flush can drain the pending
+/// stream without closing it. Backend multipart uploads still publish on
close.
+struct FlushableFileWriter {
+ delegate: opendal::Writer,
+ pending: Vec<Bytes>,
+ pending_bytes: usize,
+}
+
+#[async_trait::async_trait]
+impl FileWrite for FlushableFileWriter {
+ async fn write(&mut self, bytes: Bytes) -> crate::Result<()> {
+ // OpenDAL's non-contiguous buffer does not skip empty segments. Keep
+ // empty collection indexes out so the backend can always advance.
+ if bytes.is_empty() {
+ return Ok(());
+ }
+ self.pending_bytes += bytes.len();
+ self.pending.push(bytes);
+ if self.pending_bytes >= 8 * 1024 * 1024 {
+ self.flush().await?;
+ }
+ Ok(())
+ }
+
+ async fn flush(&mut self) -> crate::Result<()> {
+ if !self.pending.is_empty() {
+ self.pending_bytes = 0;
+ self.delegate
+ .write(std::mem::take(&mut self.pending))
+ .await?;
+ }
+ Ok(())
+ }
+
+ async fn close(&mut self) -> crate::Result<()> {
+ self.flush().await?;
+ self.delegate.close().await?;
+ Ok(())
+ }
+}
+
#[async_trait::async_trait]
impl FileWrite for CacheInvalidatingWriter {
async fn write(&mut self, bs: Bytes) -> crate::Result<()> {
self.delegate.write(bs).await
}
+ async fn flush(&mut self) -> crate::Result<()> {
+ self.delegate.flush().await?;
+ self.cache.invalidate_path(&self.path).await;
+ Ok(())
+ }
+
async fn close(&mut self) -> crate::Result<()> {
self.delegate.close().await?;
self.cache.invalidate_path(&self.path).await;
@@ -1189,6 +1313,7 @@ enum FileSource {
op: Operator,
relative_path: String,
cache_path: String,
+ local_fs: bool,
},
Provider(Arc<dyn FileIOProvider>),
}
@@ -1207,6 +1332,7 @@ impl FileSource {
op,
relative_path,
cache_path,
+ ..
} => Ok((op.clone(), relative_path.clone(),
Some(cache_path.clone()))),
}
}
@@ -1341,14 +1467,40 @@ impl OutputFile {
.chunk(8 * 1024 * 1024)
.await?,
);
+ Ok(self.with_cache_invalidation(writer, cache_path))
+ }
+
+ /// A stream-level flush is needed by BlobConsumer. Keep the ordinary
+ /// bulk output path unchanged; do not install OpenDAL's fixed-size chunk
+ /// queue here because it cannot drain a partial chunk before close.
+ pub(crate) async fn flushable_writer(&self) -> crate::Result<Box<dyn
FileWrite>> {
+ let (op, relative_path, cache_path) =
self.source.resolve(&self.path).await?;
+ let writer: Box<dyn FileWrite> =
+ if matches!(self.source, FileSource::Static { local_fs: true, ..
}) {
+ Box::new(LocalFileWriter::new(&op, &relative_path).await?)
+ } else {
+ Box::new(FlushableFileWriter {
+ delegate: op.writer(&relative_path).await?,
+ pending: Vec::new(),
+ pending_bytes: 0,
+ })
+ };
+ Ok(self.with_cache_invalidation(writer, cache_path))
+ }
+
+ fn with_cache_invalidation(
+ &self,
+ writer: Box<dyn FileWrite>,
+ cache_path: Option<String>,
+ ) -> Box<dyn FileWrite> {
let (Some(cache), Some(cache_path)) = (&self.cache, cache_path) else {
- return Ok(writer);
+ return writer;
};
- Ok(Box::new(CacheInvalidatingWriter {
+ Box::new(CacheInvalidatingWriter {
delegate: writer,
cache: cache.clone(),
path: cache_path,
- }))
+ })
}
/// Get an async streaming writer for format-level writes (e.g. parquet).
@@ -2055,6 +2207,42 @@ mod input_output_test {
FileIOBuilder::new("file").build().unwrap()
}
+ #[tokio::test]
+ async fn
test_flushable_local_output_drains_before_close_and_invalidates_cache() {
+ for cached in [false, true] {
+ let directory = tempfile::tempdir().unwrap();
+ let io = if cached {
+ setup_cached_fs_file_io(&directory.path().join("cache"))
+ } else {
+ setup_fs_file_io()
+ };
+ let path = directory
+ .path()
+ .join("output.blob")
+ .to_string_lossy()
+ .into_owned();
+ std::fs::write(&path, b"stale payload").unwrap();
+ let input = io.new_input(&path).unwrap();
+ assert_eq!(input.read().await.unwrap().as_ref(), b"stale payload");
+ let mut writer = io
+ .new_output(&path)
+ .unwrap()
+ .flushable_writer()
+ .await
+ .unwrap();
+ writer.write(Bytes::new()).await.unwrap();
+ writer.write(Bytes::from_static(b"first")).await.unwrap();
+ writer.flush().await.unwrap();
+ assert_eq!(input.read().await.unwrap().as_ref(), b"first");
+ writer.write(Bytes::new()).await.unwrap();
+ writer.write(Bytes::from_static(b"second")).await.unwrap();
+ writer.flush().await.unwrap();
+ assert_eq!(input.read().await.unwrap().as_ref(), b"firstsecond");
+ writer.close().await.unwrap();
+ assert_eq!(input.read().await.unwrap().as_ref(), b"firstsecond");
+ }
+ }
+
fn setup_cached_fs_file_io(cache_directory: &std::path::Path) -> FileIO {
let mut options = Options::new();
options.set(CatalogOptions::LOCAL_CACHE_ENABLED, "true");
diff --git a/crates/paimon/src/spec/blob_consumer.rs
b/crates/paimon/src/spec/blob_consumer.rs
new file mode 100644
index 00000000..67b44a72
--- /dev/null
+++ b/crates/paimon/src/spec/blob_consumer.rs
@@ -0,0 +1,41 @@
+// 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::BlobDescriptor;
+use crate::Result;
+
+/// Receive the physical payload range of each written BLOB, as Java's
+/// `BlobConsumer` does. A NULL field reports `None`; NULL elements inside an
+/// ARRAY or MAP do not invoke the consumer.
+///
+/// Returning `true` requests an output flush after the complete record is
+/// written. Descriptors refer to uncommitted files and do not publish a table
+/// snapshot. A failed callback fails the write and is never retried.
+/// Writers retain their BLOB files on automatic cleanup whenever a consumer
+/// is installed, since an exposed descriptor may already be used elsewhere.
+pub trait BlobConsumer: Send + Sync {
+ fn accept(&self, field_name: &str, descriptor: Option<&BlobDescriptor>) ->
Result<bool>;
+}
+
+impl<F> BlobConsumer for F
+where
+ F: Fn(&str, Option<&BlobDescriptor>) -> Result<bool> + Send + Sync,
+{
+ fn accept(&self, field_name: &str, descriptor: Option<&BlobDescriptor>) ->
Result<bool> {
+ self(field_name, descriptor)
+ }
+}
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index beee183a..3bacbaa3 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -1536,6 +1536,22 @@ impl<'a> CoreOptions<'a> {
.unwrap_or_else(|| self.target_file_size())
}
+ /// Java's BlobFormatWriter copies through a positive, int-sized buffer.
+ pub fn blob_copy_buffer_size(&self) -> crate::Result<usize> {
+ let raw = self.options.get("blob.copy-buffer-size");
+ let bytes = raw
+ .map(|value| parse_memory_size(value))
+ .unwrap_or(Some(4 * 1024))
+ .filter(|size| (1..=i64::from(i32::MAX)).contains(size))
+ .ok_or_else(|| crate::Error::ConfigInvalid {
+ message: format!(
+ "blob.copy-buffer-size must be between 1 and {} bytes, got
{raw:?}",
+ i32::MAX
+ ),
+ })?;
+ Ok(bytes as usize)
+ }
+
/// Dedicated vector-store file format, if configured.
///
/// Java leaves this unset by default. When present, vector columns are
diff --git a/crates/paimon/src/spec/mod.rs b/crates/paimon/src/spec/mod.rs
index 472fa7c2..2a0df88b 100644
--- a/crates/paimon/src/spec/mod.rs
+++ b/crates/paimon/src/spec/mod.rs
@@ -25,6 +25,9 @@ pub use binary_row::*;
mod blob_descriptor;
pub use blob_descriptor::BlobDescriptor;
+mod blob_consumer;
+pub use blob_consumer::BlobConsumer;
+
mod video_frame_descriptor;
pub use video_frame_descriptor::VideoFrameDescriptor;
diff --git a/crates/paimon/src/table/blob_consumer_tests.rs
b/crates/paimon/src/table/blob_consumer_tests.rs
new file mode 100644
index 00000000..9c3ef0a5
--- /dev/null
+++ b/crates/paimon/src/table/blob_consumer_tests.rs
@@ -0,0 +1,247 @@
+// 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::{Table, TableWrite};
+use crate::catalog::Identifier;
+use crate::io::{FileIOBuilder, FileRead};
+use crate::spec::{BlobDescriptor, BlobType, DataType, IntType, Schema,
TableSchema};
+use arrow_array::{ArrayRef, Int32Array, LargeBinaryArray, RecordBatch};
+use std::sync::{Arc, Mutex};
+
+fn table(path: &str, primary_key: bool, options: &[(&str, &str)]) -> Table {
+ let io = FileIOBuilder::new("memory").build().unwrap();
+ let mut schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("payload", DataType::Blob(BlobType::new()));
+ if primary_key {
+ schema = schema.primary_key(["id"]).option("bucket", "1");
+ } else {
+ schema = schema
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true");
+ }
+ for (name, value) in options {
+ schema = schema.option(*name, *value);
+ }
+ Table::new(
+ io,
+ Identifier::new("db", "t"),
+ path.into(),
+ TableSchema::new(0, &schema.build().unwrap()),
+ None,
+ )
+}
+
+fn batch() -> RecordBatch {
+ RecordBatch::try_from_iter([
+ ("id", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef),
+ (
+ "payload",
+ Arc::new(LargeBinaryArray::from(vec![
+ Some(b"payload".as_slice()),
+ None,
+ ])) as ArrayRef,
+ ),
+ ])
+ .unwrap()
+}
+
+#[tokio::test]
+async fn append_consumer_survives_file_rolls_and_checkpoint_recreation() {
+ let table = table(
+ "memory:/consumer-checkpoints",
+ false,
+ &[
+ ("target-file-row-num", "1"),
+ ("blob.copy-buffer-size", "2 B"),
+ ],
+ );
+ let events = Arc::new(Mutex::new(Vec::new()));
+ let received = events.clone();
+ let mut writer = TableWrite::new(&table, "consumer".into()).unwrap();
+ writer
+ .with_blob_consumer(Some(Arc::new(
+ move |name: &str, descriptor: Option<&BlobDescriptor>| {
+ assert_eq!(name, "payload");
+ received.lock().unwrap().push(descriptor.cloned());
+ Ok(true)
+ },
+ )))
+ .unwrap();
+ for _ in 0..2 {
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(
+ messages[0]
+ .new_files
+ .iter()
+ .filter(|file| file.file_name.ends_with(".blob"))
+ .map(|file| file.row_count)
+ .sum::<i64>(),
+ 2
+ );
+ assert!(writer.with_blob_consumer(None).is_err());
+ }
+ writer.close().await;
+ let events = events.lock().unwrap().clone();
+ assert_eq!(events.len(), 4);
+ for (index, descriptor) in events.into_iter().enumerate() {
+ if index % 2 == 1 {
+ assert!(descriptor.is_none());
+ } else {
+ let descriptor = descriptor.unwrap();
+ assert_eq!(descriptor.length(), 7);
+ let reader = table
+ .file_io()
+ .new_input(descriptor.uri())
+ .unwrap()
+ .reader()
+ .await
+ .unwrap();
+ let start = descriptor.offset() as u64;
+ assert_eq!(
+ reader.read(start..start + 7).await.unwrap().as_ref(),
+ b"payload"
+ );
+ }
+ }
+}
+
+#[tokio::test]
+async fn callback_failure_preserves_files_from_previous_prepare() {
+ use std::sync::atomic::{AtomicBool, Ordering};
+ let table = table("memory:/consumer-failed-checkpoint", false, &[]);
+ let fail = Arc::new(AtomicBool::new(false));
+ let callback_fail = fail.clone();
+ let mut writer = TableWrite::new(&table, "consumer".into()).unwrap();
+ writer
+ .with_blob_consumer(Some(Arc::new(
+ move |_: &str, _: Option<&BlobDescriptor>| {
+ if callback_fail.load(Ordering::Relaxed) {
+ Err(crate::Error::DataInvalid {
+ message: "consumer failed".into(),
+ source: None,
+ })
+ } else {
+ Ok(false)
+ }
+ },
+ )))
+ .unwrap();
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ let mut paths = Vec::new();
+ for file in &messages[0].new_files {
+ let path = format!("{}/bucket-0/{}", table.location(), file.file_name);
+ paths.push((
+ path.clone(),
+ table
+ .file_io()
+ .new_input(&path)
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ ));
+ }
+ fail.store(true, Ordering::Relaxed);
+ assert!(writer
+ .write_arrow_batch(&batch())
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("consumer failed"));
+ writer.close().await;
+ for (path, expected) in paths {
+ assert_eq!(
+ table
+ .file_io()
+ .new_input(&path)
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ expected
+ );
+ }
+}
+
+#[tokio::test]
+async fn explicit_commit_abort_can_discard_consumer_files() {
+ let table = table("memory:/consumer-explicit-abort", false, &[]);
+ let mut writer = TableWrite::new(&table, "consumer".into()).unwrap();
+ writer
+ .with_blob_consumer(Some(Arc::new(|_: &str, _:
Option<&BlobDescriptor>| {
+ Ok(false)
+ })))
+ .unwrap();
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ writer.close().await;
+ assert!(!table
+ .file_io()
+ .list_status_recursive(table.location())
+ .await
+ .unwrap()
+ .is_empty());
+ let commit = super::TableCommit::new(table.clone(), "consumer".into());
+ commit.abort(&messages).await.unwrap();
+ assert!(table
+ .file_io()
+ .list_status_recursive(table.location())
+ .await
+ .unwrap()
+ .is_empty());
+}
+
+#[tokio::test]
+async fn primary_key_blob_consumer_is_ignored_like_java() {
+ let table = table(
+ "memory:/consumer-primary-key",
+ true,
+ &[("blob.copy-buffer-size", "2 B")],
+ );
+ let mut writer = TableWrite::new(&table, "consumer".into()).unwrap();
+ writer
+ .with_blob_consumer(Some(Arc::new(|_: &str, _:
Option<&BlobDescriptor>| {
+ panic!("Java PK managed packs do not use the append BlobConsumer")
+ })))
+ .unwrap();
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ assert!(!writer.prepare_commit().await.unwrap()[0]
+ .new_files
+ .is_empty());
+ writer.close().await;
+}
+
+#[test]
+fn invalid_copy_buffers_are_rejected_before_data_is_written() {
+ for primary_key in [false, true] {
+ for value in ["0 B", "-1 B", "2147483648 B", "invalid"] {
+ let table = table(
+ "memory:/consumer-invalid-buffer",
+ primary_key,
+ &[("blob.copy-buffer-size", value)],
+ );
+ assert!(TableWrite::new(&table, "consumer".into())
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("blob.copy-buffer-size"));
+ }
+ }
+}
diff --git a/crates/paimon/src/table/data_file_writer.rs
b/crates/paimon/src/table/data_file_writer.rs
index 17946951..6743bdd3 100644
--- a/crates/paimon/src/table/data_file_writer.rs
+++ b/crates/paimon/src/table/data_file_writer.rs
@@ -79,8 +79,10 @@ pub(crate) struct DataFileWriter {
row_sidecar_enabled: bool,
current_row_sidecar: Option<RowSidecarWriter>,
resources: Option<ResourceContext>,
- /// Paths owned by this write until prepare_commit hands them to the
caller.
+ /// Created paths; cleanup respects the format's deletion policy until
+ /// prepare_commit hands the files to the caller.
created_paths: Vec<String>,
+ delete_file_upon_abort: bool,
}
impl DataFileWriter {
@@ -146,6 +148,7 @@ impl DataFileWriter {
current_index: None,
resources: None,
created_paths: Vec::new(),
+ delete_file_upon_abort: true,
})
}
@@ -159,6 +162,27 @@ impl DataFileWriter {
self
}
+ /// Configure only the dedicated BLOB factory; ordinary format factories
+ /// and their write paths do not carry callback state.
+ pub(super) fn set_blob_consumer(
+ &mut self,
+ consumer: Arc<dyn crate::spec::BlobConsumer>,
+ ) -> Result<()> {
+ debug_assert!(self.current_writer.is_none());
+ self.format_writer_factory = Some(Arc::new(
+ crate::arrow::format::blob::BlobWriterFactory::new(
+ Some(self.file_io.clone()),
+ self.write_fields.first(),
+ Some(&self.format_options),
+ )?
+ .with_consumer(consumer),
+ ));
+ // Java BlobFormatWriter relinquishes deletion rights whenever a
+ // consumer can expose a descriptor, even if a later callback fails.
+ self.delete_file_upon_abort = false;
+ Ok(())
+ }
+
/// Java's dedicated-format writer does not create auxiliary ROW files,
/// including for its normal columns.
pub(super) fn without_row_sidecar(mut self) -> Self {
@@ -503,12 +527,17 @@ impl DataFileWriter {
self.current_file_path = None;
while self.in_flight_closes.join_next().await.is_some() {}
for path in self.created_paths.drain(..) {
- let _ = self.file_io.delete_file(&path).await;
+ if self.delete_file_upon_abort {
+ let _ = self.file_io.delete_file(&path).await;
+ }
}
self.written_files.clear();
}
pub(super) async fn delete_files(&mut self, files: &[DataFileMeta]) {
+ if !self.delete_file_upon_abort {
+ return;
+ }
for file in files {
for path in file.collect_files(self.bucket_dir()) {
let _ = self.file_io.delete_file(&path).await;
diff --git a/crates/paimon/src/table/dedicated_format_file_writer.rs
b/crates/paimon/src/table/dedicated_format_file_writer.rs
index 975aee76..b2874be8 100644
--- a/crates/paimon/src/table/dedicated_format_file_writer.rs
+++ b/crates/paimon/src/table/dedicated_format_file_writer.rs
@@ -58,13 +58,18 @@ pub(crate) struct AppendDedicatedFormatFileWriter {
vector_writer: Option<VectorFieldWriter>,
normal_column_indices: Vec<usize>,
normal_schema: Arc<arrow_schema::Schema>,
- // Completed normal + dedicated file groups, still owned until
prepare_commit.
+ // Completed groups awaiting prepare_commit; each format owns its abort
policy.
written_files: Vec<DataFileMeta>,
target_file_row_num: i64,
current_group_row_count: i64,
}
impl AppendDedicatedFormatFileWriter {
+ #[cfg(test)]
+ pub(super) fn inject_blob_close_failure(&mut self) {
+ self.blob_writers[0].writer.inject_close_failure();
+ }
+
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
file_io: FileIO,
@@ -253,6 +258,18 @@ impl AppendDedicatedFormatFileWriter {
self
}
+ pub(crate) fn with_blob_consumer(
+ mut self,
+ consumer: Option<Arc<dyn crate::spec::BlobConsumer>>,
+ ) -> Result<Self> {
+ if let Some(consumer) = consumer {
+ for blob in &mut self.blob_writers {
+ blob.writer.set_blob_consumer(consumer.clone())?;
+ }
+ }
+ Ok(self)
+ }
+
pub(crate) async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
if batch.num_rows() == 0 {
return Ok(());
@@ -332,15 +349,7 @@ impl AppendDedicatedFormatFileWriter {
}
pub(crate) async fn abort(&mut self) {
- let cleanup = self
- .normal_writer
- .as_mut()
- .or_else(|| self.blob_writers.first_mut().map(|blob| &mut
blob.writer))
- .or_else(|| self.vector_writer.as_mut().map(|vector| &mut
vector.writer));
- if let Some(writer) = cleanup {
- writer.delete_files(&self.written_files).await;
- }
- self.written_files.clear();
+ self.delete_completed_files().await;
self.current_group_row_count = 0;
if let Some(normal) = &mut self.normal_writer {
normal.abort().await;
@@ -353,6 +362,23 @@ impl AppendDedicatedFormatFileWriter {
}
}
+ /// Each physical writer owns its abort policy. A normal writer must not
+ /// delete Blob files whose descriptors may have escaped through a
consumer.
+ async fn delete_completed_files(&mut self) {
+ for file in std::mem::take(&mut self.written_files) {
+ let blob = self.blob_writers.iter_mut().find(|blob| {
+ file.write_cols.as_deref() ==
Some(std::slice::from_ref(&blob.field_name))
+ });
+ let writer = blob
+ .map(|blob| &mut blob.writer)
+ .or(self.normal_writer.as_mut())
+ .or_else(|| self.vector_writer.as_mut().map(|vector| &mut
vector.writer));
+ if let Some(writer) = writer {
+ writer.delete_files(std::slice::from_ref(&file)).await;
+ }
+ }
+ }
+
pub(crate) async fn prepare_commit(&mut self) -> Result<Vec<DataFileMeta>>
{
self.close_group().await?;
Ok(std::mem::take(&mut self.written_files))
@@ -387,13 +413,13 @@ impl AppendDedicatedFormatFileWriter {
#[cfg(test)]
mod tests {
use super::*;
- use crate::io::FileIOBuilder;
+ use crate::io::{FileIOBuilder, FileRead};
use crate::spec::{BlobType, IntType};
use arrow_array::{ArrayRef, Int32Array, LargeBinaryArray};
#[tokio::test]
- async fn failed_blob_close_cleans_every_physical_column() {
- for (external, rolled, normal) in [
+ async fn failed_blob_close_respects_each_physical_writers_abort_policy() {
+ for (external, rolled, normal, has_consumer) in [
(false, false, true),
(true, false, true),
(false, true, true),
@@ -402,7 +428,11 @@ mod tests {
(true, false, false),
(false, true, false),
(true, true, false),
- ] {
+ ]
+ .into_iter()
+ .flat_map(|(external, rolled, normal)| {
+ [false, true].map(|consumer| (external, rolled, normal, consumer))
+ }) {
let io = FileIOBuilder::new("memory").build().unwrap();
let batch = RecordBatch::try_from_iter([
("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef),
@@ -472,6 +502,18 @@ mod tests {
&HashSet::new(),
)
.unwrap();
+ let descriptors = Arc::new(std::sync::Mutex::new(Vec::new()));
+ if has_consumer {
+ let received = descriptors.clone();
+ writer = writer
+ .with_blob_consumer(Some(Arc::new(
+ move |_: &str, descriptor:
Option<&crate::spec::BlobDescriptor>| {
+
received.lock().unwrap().push(descriptor.unwrap().clone());
+ Ok(false)
+ },
+ )))
+ .unwrap();
+ }
writer.write(&batch).await.unwrap();
writer.blob_writers[1].writer.inject_close_failure();
let error = if rolled {
@@ -483,11 +525,33 @@ mod tests {
writer.prepare_commit().await.unwrap_err()
};
assert!(error.to_string().contains("injected close failure"));
- assert!(io
- .list_status_recursive("memory:/")
- .await
- .unwrap()
- .is_empty());
+ let remaining =
io.list_status_recursive("memory:/").await.unwrap();
+ if has_consumer {
+ assert!(!remaining.is_empty());
+ assert!(remaining
+ .iter()
+ .all(|status| status.path.ends_with(".blob")));
+ let descriptors = descriptors.lock().unwrap().clone();
+ for descriptor in descriptors {
+ let reader = io
+ .new_input(descriptor.uri())
+ .unwrap()
+ .reader()
+ .await
+ .unwrap();
+ let start = descriptor.offset() as u64;
+ assert_eq!(
+ reader
+ .read(start..start + descriptor.length() as u64)
+ .await
+ .unwrap()
+ .len(),
+ 1
+ );
+ }
+ } else {
+ assert!(remaining.is_empty());
+ }
}
}
}
diff --git a/crates/paimon/src/table/managed_blob_writer.rs
b/crates/paimon/src/table/managed_blob_writer.rs
index ed951934..402ddc0b 100644
--- a/crates/paimon/src/table/managed_blob_writer.rs
+++ b/crates/paimon/src/table/managed_blob_writer.rs
@@ -87,6 +87,7 @@ pub(crate) struct ManagedBlobWriter {
bucket_dir: String,
file_prefix: String,
target_file_size: u64,
+ copy_buffer_size: usize,
fields: Vec<ManagedBlobField>,
uncommitted_paths: Vec<String>,
}
@@ -104,13 +105,16 @@ impl ManagedBlobWriteState {
value_fields: &[DataField],
options: &CoreOptions<'_>,
) -> Result<Self> {
- let writer = ManagedBlobWriter::new(
+ let mut writer = ManagedBlobWriter::new(
file_io.clone(),
bucket_dir.to_string(),
file_prefix,
options.blob_target_file_size(),
managed_blob_fields(value_fields, options),
)?;
+ if let Some(writer) = &mut writer {
+ writer.copy_buffer_size = options.blob_copy_buffer_size()?;
+ }
Ok(Self {
writer: writer.map(|writer|
tokio::sync::Mutex::new(Box::new(writer))),
})
@@ -164,6 +168,7 @@ impl ManagedBlobWriter {
bucket_dir,
file_prefix: file_prefix.to_string(),
target_file_size,
+ copy_buffer_size: 4 * 1024,
fields: fields
.into_iter()
.map(|(index, kind)| ManagedBlobField {
@@ -356,8 +361,11 @@ impl ManagedBlobWriter {
uuid::Uuid::new_v4()
);
let output = self.file_io.new_output(&path)?;
- let writer =
- Box::new(BlobFormatWriter::new(&output,
Some(self.file_io.clone()), None).await?);
+ let writer = Box::new(
+ BlobFormatWriter::new(&output, Some(self.file_io.clone()),
None)
+ .await?
+ .with_copy_buffer_size(self.copy_buffer_size),
+ );
self.uncommitted_paths.push(path.clone());
self.fields[field_index].current = Some(ManagedBlobPack { path,
writer });
}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 93243947..ed7e6a06 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -24,6 +24,8 @@ pub(crate) mod bin_pack;
mod bitmap_global_index_format;
mod bitmap_global_index_reader;
mod bitmap_global_index_writer;
+#[cfg(test)]
+mod blob_consumer_tests;
mod blob_resolver;
mod branch_manager;
#[cfg(test)]
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index f36323d1..9d7d2f9e 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -177,6 +177,7 @@ pub struct TableWrite {
blob_inline_fields: HashSet<String>,
/// Whether the table has non-descriptor blob fields requiring a
dedicated-format writer.
has_blob_fields: bool,
+ blob_consumer: Option<Arc<dyn crate::spec::BlobConsumer>>,
/// Dedicated vector-store file format, when configured.
vector_file_format: Option<String>,
/// Whether the table has VECTOR fields requiring dedicated vector files.
@@ -247,6 +248,7 @@ impl TableWrite {
blob_view_fields: HashSet::new(),
blob_inline_fields: HashSet::new(),
has_blob_fields: false,
+ blob_consumer: None,
vector_file_format: None,
has_dedicated_vector_fields: false,
row_kind_generator: None,
@@ -474,6 +476,9 @@ impl TableWrite {
.fields()
.iter()
.any(|f| f.data_type().is_blob_file_field() &&
!blob_inline_fields.contains(f.name()));
+ if has_blob_fields {
+ core_options.blob_copy_buffer_size()?;
+ }
let has_dedicated_vector_fields = vector_file_format.is_some()
&& schema
.fields()
@@ -526,6 +531,7 @@ impl TableWrite {
blob_view_fields,
blob_inline_fields,
has_blob_fields,
+ blob_consumer: None,
vector_file_format,
has_dedicated_vector_fields,
row_kind_generator,
@@ -707,6 +713,30 @@ 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;
+ /// 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(),
+ });
+ }
+ self.blob_consumer = consumer;
+ Ok(self)
+ }
+
/// Write an Arrow RecordBatch. Rows are routed to the correct partition
and bucket.
pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) ->
Result<()> {
if let Some(writer) = self.format_writer.as_mut() {
@@ -1280,8 +1310,20 @@ impl TableWrite {
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()
{
+ // 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
+ // available when the caller knows the files are no longer needed.
+ for message in &mut cleanup {
+ message
+ .new_files
+ .retain(|file| !file.file_name.ends_with(".blob"));
+ }
+ }
let commit = super::TableCommit::new(self.table.clone(),
self.commit_user.clone());
- if let Err(error) = commit.abort(messages).await {
+ if let Err(error) = commit.abort(&cleanup).await {
log::warn!("Failed to discard unreturned prepared files: {error}");
}
}
@@ -1341,6 +1383,7 @@ impl TableWrite {
self.table.schema().options(),
&self.blob_inline_fields,
)?
+ .with_blob_consumer(self.blob_consumer.clone())?
.with_resources(self.resources.clone()),
)))
} else {
@@ -1518,6 +1561,95 @@ pub(in crate::table) mod tests {
FileIOBuilder::new("memory").build().unwrap()
}
+ #[tokio::test]
+ async fn failed_partition_prepare_preserves_consumer_blob_files() {
+ use crate::io::FileRead;
+ use crate::spec::BlobDescriptor;
+ use arrow_array::{ArrayRef, LargeBinaryArray};
+ use std::sync::Mutex;
+
+ for has_consumer in [false, true] {
+ let io = test_file_io();
+ let schema = Schema::builder()
+ .column("pt", DataType::Int(IntType::new()))
+ .column("payload", DataType::Blob(BlobType::new()))
+ .partition_keys(["pt"])
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ io.clone(),
+ Identifier::new("db", "partitioned_blob"),
+ "memory:/partitioned-blob".into(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+ let mut writer = TableWrite::new(&table,
"consumer".into()).unwrap();
+ let descriptors = Arc::new(Mutex::new(Vec::new()));
+ if has_consumer {
+ let received = descriptors.clone();
+ writer
+ .with_blob_consumer(Some(Arc::new(
+ move |_: &str, descriptor: Option<&BlobDescriptor>| {
+
received.lock().unwrap().push(descriptor.unwrap().clone());
+ Ok(false)
+ },
+ )))
+ .unwrap();
+ }
+ let batch = RecordBatch::try_from_iter([
+ ("pt", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef),
+ (
+ "payload",
+ Arc::new(LargeBinaryArray::from(vec![
+ Some(b"payload".as_slice()),
+ Some(b"payload".as_slice()),
+ ])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ writer.write_arrow_batch(&batch).await.unwrap();
+ assert_eq!(writer.partition_writers.len(), 2);
+ let FileWriter::AppendDedicated(partition) =
+ writer.partition_writers.values_mut().next().unwrap()
+ else {
+ panic!("expected a dedicated Blob writer");
+ };
+ partition.inject_blob_close_failure();
+ assert!(writer
+ .prepare_commit()
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("injected close failure"));
+ assert!(writer.prepare_commit().await.is_err());
+ writer.close().await;
+ let files =
io.list_status_recursive(table.location()).await.unwrap();
+ if has_consumer {
+ assert_eq!(files.len(), 2);
+ assert!(files.iter().all(|file| file.path.ends_with(".blob")));
+ let descriptors = descriptors.lock().unwrap().clone();
+ assert_eq!(descriptors.len(), 2);
+ for descriptor in descriptors {
+ let reader = io
+ .new_input(descriptor.uri())
+ .unwrap()
+ .reader()
+ .await
+ .unwrap();
+ let start = descriptor.offset() as u64;
+ assert_eq!(
+ reader.read(start..start + 7).await.unwrap().as_ref(),
+ b"payload"
+ );
+ }
+ } else {
+ assert!(files.is_empty());
+ }
+ }
+ }
+
fn test_schema() -> TableSchema {
let schema = Schema::builder()
.column("id", DataType::Int(IntType::new()))
