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 86f71e52 feat(blob): support custom URI readers in native writes
(#1065)
86f71e52 is described below
commit 86f71e52bc28f1dbfaaf296fba0356916697e4c3
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 23:19:32 2026 +0800
feat(blob): support custom URI readers in native writes (#1065)
---
Cargo.lock | 2 +
bindings/python/Cargo.toml | 2 +
bindings/python/README.md | 15 +
.../python/python/pypaimon_rust/datafusion.pyi | 17 +-
bindings/python/src/blob_uri_reader.rs | 209 ++++++++++++
bindings/python/src/lib.rs | 1 +
bindings/python/src/write.rs | 81 +++--
bindings/python/tests/test_blob_uri_reader.py | 107 +++++++
crates/paimon/src/arrow/format/blob/writer.rs | 152 ++++++++-
.../src/arrow/format/blob/writer_consumer_tests.rs | 2 +
.../arrow/format/blob/writer_uri_reader_tests.rs | 350 +++++++++++++++++++++
crates/paimon/src/io/mod.rs | 1 +
crates/paimon/src/io/uri_reader.rs | 4 +
crates/paimon/src/io/uri_reader/custom.rs | 122 +++++++
crates/paimon/src/table/blob_uri_reader_tests.rs | 149 +++++++++
crates/paimon/src/table/data_file_writer.rs | 29 +-
.../src/table/dedicated_format_file_writer.rs | 23 +-
crates/paimon/src/table/kv_file_writer.rs | 8 +
crates/paimon/src/table/managed_blob_writer.rs | 20 +-
crates/paimon/src/table/mod.rs | 2 +
crates/paimon/src/table/postpone_file_writer.rs | 8 +
.../src/table/postpone_fixed_bucket_write.rs | 9 +
crates/paimon/src/table/table_write.rs | 39 ++-
23 files changed, 1292 insertions(+), 60 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 3347b5bb..ba6864b1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5383,6 +5383,8 @@ version = "0.4.0"
dependencies = [
"arrow",
"arrow-row",
+ "async-trait",
+ "bytes",
"chrono",
"datafusion",
"datafusion-ffi",
diff --git a/bindings/python/Cargo.toml b/bindings/python/Cargo.toml
index 7f985951..2739e62d 100644
--- a/bindings/python/Cargo.toml
+++ b/bindings/python/Cargo.toml
@@ -30,6 +30,8 @@ doc = false
[dependencies]
arrow = { workspace = true, features = ["pyarrow"] }
arrow-row = { workspace = true }
+async-trait = "0.1.81"
+bytes = "1.7.1"
chrono = "0.4.38"
datafusion = { workspace = true }
datafusion-ffi = { workspace = true }
diff --git a/bindings/python/README.md b/bindings/python/README.md
index ad3230c4..f436eeef 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -108,6 +108,21 @@ print(f"\nRead: {batches_tt[0].num_rows} rows")
print(batches_tt[0])
```
+### Custom Blob URI sources
+
+Before the first Arrow write, call
+`writer.with_blob_uri_reader_factory(factory)` to copy referenced Blob payloads
+from custom sources. The factory provides `create(uri)`, returning a reader
with
+`new_input_stream(uri)`. Streams provide bounded `read(size)`, `close()`, and
+optionally `seek(offset)`; an empty read means EOF. Return the same reader
object
+when sources share reader identity so consecutive reference windows can reuse
an
+open stream. Pass `None` to restore the default URI reader before writing.
+
+Rust handles descriptor ranges, successful short reads, stream reuse, reopening
+and closure for scalar, array and map Blob writes, including managed
primary-key
+Blob packs. Exceptions raised by Python source callbacks are propagated to the
+caller. The Python layer does not copy payloads or cache source streams.
+
### Native commit from serialized messages
The Python binding follows Java's batch/stream builder structure. Use
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index ebb823c8..ed8a1b4a 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -17,7 +17,7 @@
from os import PathLike
from typing import (
- Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple,
+ Any, BinaryIO, Callable, Dict, Iterable, List, Literal, Optional,
Protocol, Sequence, Tuple,
TypeAlias, Union,
)
@@ -27,6 +27,12 @@ ArrowTypeLike: TypeAlias = Union[pyarrow.DataType,
pyarrow.Field, str]
InputFieldsLike: TypeAlias = Union[ArrowTypeLike, Sequence[ArrowTypeLike]]
VolatilityLike: TypeAlias = Union[str, Any]
+class _UriReader(Protocol):
+ def new_input_stream(self, uri: str) -> BinaryIO: ...
+
+class _UriReaderFactory(Protocol):
+ def create(self, uri: str) -> _UriReader: ...
+
class DataField:
def id(self) -> int: ...
def name(self) -> str: ...
@@ -236,6 +242,12 @@ class CommitMessage:
...
class BatchTableWrite:
+ def with_blob_uri_reader_factory(
+ self, factory: Optional[_UriReaderFactory]
+ ) -> "BatchTableWrite":
+ """Use custom source streams for Blob references; Rust copies bounded
+ ranges, reuses consecutive sources and closes them on completion."""
+ ...
def with_write_type(self, columns: List[str]) -> "BatchTableWrite": ...
def with_blob_consumer(
self, callback: Optional[Callable[[str, Optional[bytes]], bool]]
@@ -297,6 +309,9 @@ class StreamTableUpdate:
...
class StreamTableWrite:
+ def with_blob_uri_reader_factory(
+ self, factory: Optional[_UriReaderFactory]
+ ) -> "StreamTableWrite": ...
def with_write_type(self, columns: List[str]) -> "StreamTableWrite": ...
def with_blob_consumer(
self, callback: Optional[Callable[[str, Optional[bytes]], bool]]
diff --git a/bindings/python/src/blob_uri_reader.rs
b/bindings/python/src/blob_uri_reader.rs
new file mode 100644
index 00000000..46608317
--- /dev/null
+++ b/bindings/python/src/blob_uri_reader.rs
@@ -0,0 +1,209 @@
+// 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.
+
+//! Python URI objects and exception transport; source reuse and copy semantics
+//! are implemented by Rust core.
+
+use bytes::Bytes;
+use paimon::io::{UriInputStream, UriReader, UriReaderFactory};
+use pyo3::exceptions::PyTypeError;
+use pyo3::prelude::*;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, Weak};
+
+struct PythonCallbackError(PyErr);
+
+impl std::fmt::Debug for PythonCallbackError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.write_str("PythonCallbackError")
+ }
+}
+
+impl std::fmt::Display for PythonCallbackError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.write_str("Python Blob callback failed")
+ }
+}
+impl std::error::Error for PythonCallbackError {}
+
+pub(crate) fn callback_error(error: PyErr) -> paimon::Error {
+ paimon::Error::UnexpectedError {
+ message: "Python Blob callback failed".into(),
+ source: Some(Box::new(PythonCallbackError(error))),
+ }
+}
+
+pub(crate) fn to_callback_error(error: paimon::Error) -> PyErr {
+ let mut cause: &(dyn std::error::Error + 'static) = &error;
+ loop {
+ if let Some(original) = cause.downcast_ref::<PythonCallbackError>() {
+ return Python::attach(|py| original.0.clone_ref(py));
+ }
+ // Paimon UnexpectedError deliberately omits its source from the
+ // standard chain. Inspect that explicit source without stringifying
+ // Python errors (which can contain URI credentials).
+ let next = match cause.downcast_ref::<paimon::Error>() {
+ Some(
+ paimon::Error::UnexpectedError { source, .. }
+ | paimon::Error::DataInvalid { source, .. },
+ ) => source
+ .as_deref()
+ .map(|source| source as &(dyn std::error::Error + 'static)),
+ _ => cause.source(),
+ };
+ match next {
+ Some(source) => cause = source,
+ None => break,
+ }
+ }
+ crate::error::to_py_err(error)
+}
+
+pub(crate) struct PythonUriReaderFactory {
+ factory: Py<PyAny>,
+ // Preserve object identity across interleaved physical Blob writers
without
+ // retaining readers after core releases them. URI caching belongs to the
factory.
+ readers: Mutex<HashMap<usize, Weak<PythonUriReader>>>,
+}
+
+impl PythonUriReaderFactory {
+ pub(crate) fn new(factory: Py<PyAny>, py: Python<'_>) -> PyResult<Self> {
+ if !factory.bind(py).getattr("create")?.is_callable() {
+ return Err(PyTypeError::new_err(
+ "URI reader factory must provide callable create",
+ ));
+ }
+ Ok(Self {
+ factory,
+ readers: Mutex::new(HashMap::new()),
+ })
+ }
+}
+
+impl UriReaderFactory for PythonUriReaderFactory {
+ fn create(&self, uri: &str) -> paimon::Result<Arc<dyn UriReader>> {
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ let reader = self
+ .factory
+ .bind(py)
+ .call_method1("create", (uri,))?
+ .unbind();
+ if !reader.bind(py).getattr("new_input_stream")?.is_callable()
{
+ return Err(PyTypeError::new_err(
+ "URI reader must provide callable new_input_stream",
+ ));
+ }
+ let identity = reader.as_ptr() as usize;
+ let mut cached = self.readers.lock().unwrap();
+ if let Some(previous) =
cached.get(&identity).and_then(Weak::upgrade) {
+ return Ok(previous as Arc<dyn UriReader>);
+ }
+ cached.retain(|_, reader| reader.strong_count() != 0);
+ let reader = Arc::new(PythonUriReader { reader });
+ cached.insert(identity, Arc::downgrade(&reader));
+ Ok(reader as Arc<dyn UriReader>)
+ })
+ })
+ .map_err(callback_error)
+ }
+}
+
+struct PythonUriReader {
+ reader: Py<PyAny>,
+}
+
+#[async_trait::async_trait]
+impl UriReader for PythonUriReader {
+ async fn new_input_stream(&self, uri: &str) -> paimon::Result<Box<dyn
UriInputStream>> {
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ let stream = self
+ .reader
+ .bind(py)
+ .call_method1("new_input_stream", (uri,))?
+ .unbind();
+ Ok(Box::new(PythonUriInputStream {
+ stream: Some(stream),
+ }) as Box<dyn UriInputStream>)
+ })
+ })
+ .map_err(callback_error)
+ }
+}
+
+struct PythonUriInputStream {
+ stream: Option<Py<PyAny>>,
+}
+
+#[async_trait::async_trait]
+impl UriInputStream for PythonUriInputStream {
+ async fn read(&mut self, length: usize) -> paimon::Result<Bytes> {
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ let bytes = self
+ .stream
+ .as_ref()
+ .expect("open stream")
+ .bind(py)
+ .call_method1("read", (length,))?
+ .extract::<Vec<u8>>()?;
+ Ok(Bytes::from(bytes))
+ })
+ })
+ .map_err(callback_error)
+ }
+
+ async fn seek(&mut self, offset: u64) -> paimon::Result<()> {
+ tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ self.stream
+ .as_ref()
+ .expect("open stream")
+ .bind(py)
+ .call_method1("seek", (offset,))?;
+ Ok(())
+ })
+ })
+ .map_err(callback_error)
+ }
+
+ async fn close(&mut self) -> paimon::Result<()> {
+ match self.stream.take() {
+ Some(stream) => tokio::task::block_in_place(|| {
+ Python::attach(|py| {
+ stream.bind(py).call_method0("close")?;
+ Ok(())
+ })
+ })
+ .map_err(callback_error),
+ None => Ok(()),
+ }
+ }
+}
+
+impl Drop for PythonUriInputStream {
+ fn drop(&mut self) {
+ // Async close is explicit in core. On cancellation, still release a
+ // Python stream without replacing the triggering error.
+ if let Some(stream) = self.stream.take() {
+ Python::attach(|py| {
+ let _ = stream.bind(py).call_method0("close");
+ });
+ }
+ }
+}
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index a4c8eaa1..30446a3f 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -18,6 +18,7 @@
use pyo3::prelude::*;
mod blob;
+mod blob_uri_reader;
mod context;
mod error;
mod oss_cpp;
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index bfd23051..d1de7f8a 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -31,6 +31,7 @@ use pyo3::exceptions::{PyRuntimeError, PyTypeError,
PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyString};
+use crate::blob_uri_reader::{callback_error, to_callback_error,
PythonUriReaderFactory};
use crate::error::to_py_err;
use crate::predicate::{dict_to_table_predicate, py_to_datum};
@@ -84,7 +85,6 @@ impl WriteContext {
))),
table_location: self.table.location().to_string(),
commit_user: self.commit_user.clone(),
- blob_consumer: None,
})
}
}
@@ -315,6 +315,21 @@ enum WriteTarget {
}
impl WriteTarget {
+ fn with_blob_uri_reader_factory(
+ &mut self,
+ factory: Option<Arc<dyn paimon::io::UriReaderFactory>>,
+ ) -> paimon::Result<()> {
+ match self {
+ Self::Table(writer) => {
+ writer.with_blob_uri_reader_factory(factory)?;
+ }
+ Self::PostponeFixed(writer) => {
+ writer.with_blob_uri_reader_factory(factory)?;
+ }
+ }
+ Ok(())
+ }
+
fn with_blob_consumer(
&mut self,
consumer: Option<Arc<dyn BlobConsumer>>,
@@ -459,7 +474,6 @@ impl PyPostponeFixedBucketWriteBuilder {
))),
table_location: context.table.location().to_string(),
commit_user: context.commit_user.clone(),
- blob_consumer: None,
},
prepared: false,
})
@@ -474,14 +488,12 @@ 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 {
@@ -500,13 +512,7 @@ impl BlobConsumer for PythonBlobConsumer {
.is_truthy()
})
});
- result.map_err(|error| {
- *self.error.lock().unwrap() = Some(error);
- paimon::Error::DataInvalid {
- message: "BLOB consumer failed".into(),
- source: None,
- }
- })
+ result.map_err(callback_error)
}
}
@@ -711,16 +717,31 @@ impl UpdateContext {
}
impl WriteState {
+ fn with_blob_uri_reader_factory(
+ &mut self,
+ factory: Option<Py<PyAny>>,
+ py: Python<'_>,
+ ) -> PyResult<()> {
+ let factory = factory
+ .map(|factory| {
+ PythonUriReaderFactory::new(factory, py)
+ .map(|factory| Arc::new(factory) as Arc<dyn
paimon::io::UriReaderFactory>)
+ })
+ .transpose()?;
+ self.inner
+ .as_mut()
+ .ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?
+ .with_blob_uri_reader_factory(factory)
+ .map_err(to_py_err)
+ }
+
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),
- }))
+ Ok(Arc::new(PythonBlobConsumer { callback }))
})
.transpose()?;
self.inner
@@ -732,17 +753,11 @@ impl WriteState {
.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)
+ result.map_err(to_callback_error)
}
fn with_write_type(&mut self, columns: Vec<String>) -> PyResult<()> {
@@ -775,7 +790,7 @@ impl WriteState {
.ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?;
let messages = py
.detach(|| runtime().block_on(inner.prepare_commit()))
- .map_err(to_py_err)?;
+ .map_err(to_callback_error)?;
Ok(wrap_messages(
messages,
&self.table_location,
@@ -982,6 +997,16 @@ impl PyBatchTableUpdate {
#[pymethods]
impl PyBatchTableWrite {
+ #[pyo3(signature = (factory))]
+ fn with_blob_uri_reader_factory(
+ mut slf: PyRefMut<'_, Self>,
+ factory: Option<Py<PyAny>>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.state.with_blob_uri_reader_factory(factory, py)?;
+ Ok(slf)
+ }
+
#[pyo3(signature = (callback))]
fn with_blob_consumer(
mut slf: PyRefMut<'_, Self>,
@@ -1038,6 +1063,16 @@ pub struct PyStreamTableWrite {
#[pymethods]
impl PyStreamTableWrite {
+ #[pyo3(signature = (factory))]
+ fn with_blob_uri_reader_factory(
+ mut slf: PyRefMut<'_, Self>,
+ factory: Option<Py<PyAny>>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.state.with_blob_uri_reader_factory(factory, py)?;
+ Ok(slf)
+ }
+
#[pyo3(signature = (callback))]
fn with_blob_consumer(
mut slf: PyRefMut<'_, Self>,
diff --git a/bindings/python/tests/test_blob_uri_reader.py
b/bindings/python/tests/test_blob_uri_reader.py
new file mode 100644
index 00000000..0e20ed09
--- /dev/null
+++ b/bindings/python/tests/test_blob_uri_reader.py
@@ -0,0 +1,107 @@
+# 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 io
+import struct
+
+import pyarrow as pa
+import pytest
+
+from test_blob_consumer import _table
+
+
+def descriptor(offset, length):
+ uri = b'custom://source'
+ return struct.pack('<BQi', 2, 0x424C4F4244455343, len(uri)) + uri +
struct.pack('<qq', offset, length)
+
+
+class ReaderFactory:
+ def __init__(self, no_seek=False):
+ self.no_seek = no_seek
+ self.opens = self.closes = 0
+ self.failure = None
+
+ def create(self, uri):
+ assert uri == 'custom://source'
+ return self
+
+ def new_input_stream(self, uri):
+ self.opens += 1
+ owner = self
+
+ class Stream(io.BytesIO):
+ def read(self, length=-1):
+ if owner.failure:
+ raise owner.failure
+ return super().read(min(length, 2))
+
+ def seek(self, offset, whence=0):
+ if owner.no_seek:
+ raise io.UnsupportedOperation('cannot seek')
+ return super().seek(offset, whence)
+
+ def close(self):
+ if not self.closed:
+ owner.closes += 1
+ super().close()
+
+ return Stream(b'0123456789')
+
+
[email protected]('stream', [False, True])
[email protected]('no_seek', [False, True])
+def
test_custom_uri_streams_are_core_bounded_and_rewind_failures_recover(tmp_path,
stream, no_seek):
+ context, table = _table(tmp_path)
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ factory = ReaderFactory(no_seek)
+ try:
+ assert writer.with_blob_uri_reader_factory(factory) is writer
+ data = pa.record_batch([pa.array([1, 2, 3], type=pa.int32()),
+ pa.array([descriptor(0, 3), descriptor(3, 2),
descriptor(0, 3)],
+ type=pa.large_binary())],
names=['id', 'payload'])
+ writer.write_arrow(data)
+ with pytest.raises(ValueError, match='before any write'):
+ writer.with_blob_uri_reader_factory(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 pa.Table.from_batches(context.sql(
+ 'SELECT id, payload FROM paimon.consumer.t ORDER BY
id')).to_pydict() == {
+ 'id': [1, 2, 3], 'payload': [b'012', b'34', b'012']}
+ assert factory.opens == factory.closes == (2 if no_seek else 1)
+ finally:
+ writer.close()
+
+
+def
test_custom_uri_read_error_preserves_original_exception_and_closes_once(tmp_path):
+ _, table = _table(tmp_path)
+ writer = table.new_batch_write_builder().new_write()
+ factory = ReaderFactory()
+ factory.failure = LookupError('opaque source failed')
+ try:
+ writer.with_blob_uri_reader_factory(factory)
+ data = pa.record_batch([pa.array([1], type=pa.int32()),
+ pa.array([descriptor(0, 3)],
type=pa.large_binary())],
+ names=['id', 'payload'])
+ with pytest.raises(LookupError) as caught:
+ writer.write_arrow(data)
+ assert caught.value is factory.failure
+ assert factory.opens == factory.closes == 1
+ finally:
+ writer.close()
+ assert not list(tmp_path.rglob('*.blob'))
diff --git a/crates/paimon/src/arrow/format/blob/writer.rs
b/crates/paimon/src/arrow/format/blob/writer.rs
index 6b69c8fb..5219c5a2 100644
--- a/crates/paimon/src/arrow/format/blob/writer.rs
+++ b/crates/paimon/src/arrow/format/blob/writer.rs
@@ -23,7 +23,8 @@ use super::{
BLOB_MAP_MAGIC_NUMBER, BLOB_MAP_VERSION,
};
use crate::arrow::format::{FormatFileWriter, FormatWriteResult,
FormatWriterFactory};
-use crate::io::{FileIO, FileRead, FileWrite, OutputFile};
+use crate::io::uri_reader::ReusingBlobRefStreamProvider;
+use crate::io::{FileIO, FileRead, FileWrite, OutputFile, UriInputStream,
UriReaderFactory};
use crate::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataField,
DataType};
use crate::{Error, Result};
use arrow_array::{
@@ -45,10 +46,15 @@ const SOURCE_READ_SIZE: u64 = 8 * 1024 * 1024;
#[path = "writer_consumer_tests.rs"]
mod consumer_tests;
+#[cfg(test)]
+#[path = "writer_uri_reader_tests.rs"]
+mod uri_reader_tests;
+
pub(crate) struct BlobWriterFactory {
file_io: Option<FileIO>,
field: Option<DataField>,
consumer: Option<Arc<dyn BlobConsumer>>,
+ uri_reader_factory: Option<Arc<dyn UriReaderFactory>>,
copy_buffer_size: usize,
}
@@ -63,6 +69,7 @@ impl BlobWriterFactory {
file_io,
field: field.cloned(),
consumer: None,
+ uri_reader_factory: None,
copy_buffer_size: CoreOptions::new(options.unwrap_or(&defaults))
.blob_copy_buffer_size()?,
})
@@ -72,6 +79,11 @@ impl BlobWriterFactory {
self.consumer = Some(consumer);
self
}
+
+ pub(crate) fn with_uri_reader_factory(mut self, factory: Arc<dyn
UriReaderFactory>) -> Self {
+ self.uri_reader_factory = Some(factory);
+ self
+ }
}
#[async_trait]
@@ -94,6 +106,7 @@ impl FormatWriterFactory for BlobWriterFactory {
self.field.as_ref(),
)?
.with_copy_buffer_size(self.copy_buffer_size)
+ .with_uri_reader_factory(self.uri_reader_factory.clone())
.with_consumer(self.consumer.clone()),
))
}
@@ -106,6 +119,8 @@ pub(crate) struct BlobFormatWriter {
path: String,
field_name: String,
consumer: Option<Arc<dyn BlobConsumer>>,
+ uri_reader_factory: Option<Arc<dyn UriReaderFactory>>,
+ reference_streams: ReusingBlobRefStreamProvider,
copy_buffer_size: usize,
bytes_written: u64,
lengths: Vec<i64>,
@@ -147,6 +162,8 @@ impl BlobFormatWriter {
path: output.location().to_string(),
field_name: field.map_or_else(String::new, |field|
field.name().to_string()),
consumer: None,
+ uri_reader_factory: None,
+ reference_streams: ReusingBlobRefStreamProvider::default(),
copy_buffer_size: 4 * 1024,
bytes_written: 0,
lengths: Vec::new(),
@@ -159,6 +176,14 @@ impl BlobFormatWriter {
self
}
+ pub(crate) fn with_uri_reader_factory(
+ mut self,
+ factory: Option<Arc<dyn UriReaderFactory>>,
+ ) -> Self {
+ self.uri_reader_factory = factory;
+ self
+ }
+
fn with_consumer(mut self, consumer: Option<Arc<dyn BlobConsumer>>) ->
Self {
self.consumer = consumer;
self
@@ -249,6 +274,11 @@ impl BlobFormatWriter {
}
let descriptor = BlobDescriptor::deserialize(value)?;
let range = descriptor.range_spec()?;
+ if let Some(factory) = self.uri_reader_factory.clone() {
+ return self
+ .copy_custom_reference(factory.as_ref(), &descriptor, hasher)
+ .await;
+ }
let file_io = self.file_io.as_ref().ok_or_else(|| {
invalid("BlobFormatWriter received a BlobDescriptor but has no
FileIO to resolve it")
})?;
@@ -281,6 +311,94 @@ impl BlobFormatWriter {
Ok(length_i64)
}
+ async fn copy_custom_reference(
+ &mut self,
+ factory: &dyn UriReaderFactory,
+ descriptor: &BlobDescriptor,
+ hasher: &mut Hasher,
+ ) -> Result<i64> {
+ let range = descriptor.range_spec()?;
+ let reader = factory.create(descriptor.uri())?;
+ let length = match range.length() {
+ Some(length) => {
+ let mut sources = std::mem::take(&mut self.reference_streams);
+ let result = async {
+ sources
+ .prepare(reader, descriptor.uri(), range.offset())
+ .await?;
+ let result = self
+ .copy_uri_stream(sources.stream(), Some(length),
hasher)
+ .await;
+ match result {
+ Ok(length) => {
+ sources.advance(length);
+ Ok(length)
+ }
+ Err(error) => {
+ let _ = sources.close().await;
+ Err(error)
+ }
+ }
+ }
+ .await;
+ self.reference_streams = sources;
+ result?
+ }
+ None => {
+ // Unknown-length references read to EOF on their own stream;
+ // do not consume or replace the cached bounded source.
+ let mut stream =
reader.new_input_stream(descriptor.uri()).await?;
+ let result = async {
+ if range.offset() != 0 {
+ stream.seek(range.offset()).await?;
+ }
+ self.copy_uri_stream(stream.as_mut(), None, hasher).await
+ }
+ .await;
+ let close = stream.close().await;
+ let length = result?;
+ close?;
+ length
+ }
+ };
+ i64::try_from(length).map_err(|_| invalid("BLOB payload exceeds i64"))
+ }
+
+ async fn copy_uri_stream(
+ &mut self,
+ stream: &mut dyn UriInputStream,
+ length: Option<u64>,
+ hasher: &mut Hasher,
+ ) -> Result<u64> {
+ let mut copied = 0_u64;
+ loop {
+ let requested = length.map_or(self.copy_buffer_size, |length| {
+ (length - copied).min(self.copy_buffer_size as u64) as usize
+ });
+ if requested == 0 {
+ return Ok(copied);
+ }
+ let bytes = stream.read(requested).await?;
+ if bytes.len() > requested {
+ return Err(invalid("URI stream returned more bytes than
requested"));
+ }
+ if bytes.is_empty() {
+ if length.is_some_and(|length| length != copied) {
+ return Err(invalid(format!(
+ "Unexpected EOF copying BLOB payload: expected {}
bytes, received {copied}",
+ length.unwrap()
+ )));
+ }
+ return Ok(copied);
+ }
+ copied = copied
+ .checked_add(bytes.len() as u64)
+ .filter(|length| *length <= i64::MAX as u64)
+ .ok_or_else(|| invalid("BLOB payload exceeds i64"))?;
+ self.write_bytes(bytes, hasher).await?;
+ }
+ }
+
async fn copy_payload(
&mut self,
reader: &dyn FileRead,
@@ -455,17 +573,27 @@ impl FormatFileWriter for BlobFormatWriter {
self.writer.flush().await
}
async fn close(mut self: Box<Self>) -> Result<FormatWriteResult> {
- let index = encode_delta_varints_write(&self.lengths);
- let length =
- i32::try_from(index.len()).map_err(|_| invalid("BLOB file index
exceeds i32"))?;
- self.writer.write(index.into()).await?;
- self.writer
- .write(Bytes::copy_from_slice(&length.to_le_bytes()))
- .await?;
- self.writer
- .write(Bytes::from_static(&[BLOB_FORMAT_VERSION]))
- .await?;
- self.writer.close().await?;
+ let footer = async {
+ let index = encode_delta_varints_write(&self.lengths);
+ let length =
+ i32::try_from(index.len()).map_err(|_| invalid("BLOB file
index exceeds i32"))?;
+ self.writer.write(index.into()).await?;
+ self.writer
+ .write(Bytes::copy_from_slice(&length.to_le_bytes()))
+ .await?;
+ self.writer
+ .write(Bytes::from_static(&[BLOB_FORMAT_VERSION]))
+ .await?;
+ Ok::<_, Error>(length)
+ }
+ .await;
+ // Close both source and destination even if the footer or one close
+ // fails, keeping the first failure primary, like Java's finally
blocks.
+ let source_close = self.reference_streams.close().await;
+ let output_close = self.writer.close().await;
+ let length = footer?;
+ source_close?;
+ output_close?;
Ok(FormatWriteResult::new(
self.bytes_written + length as u64 + 5,
))
diff --git a/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
b/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
index cd0df2f6..010bc5c6 100644
--- a/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
+++ b/crates/paimon/src/arrow/format/blob/writer_consumer_tests.rs
@@ -63,6 +63,8 @@ fn writer(kind: BlobFieldKind, output:
&Arc<Mutex<OutputState>>) -> BlobFormatWr
path: "memory:/payload.blob".into(),
field_name: "payload".into(),
consumer: None,
+ uri_reader_factory: None,
+ reference_streams: ReusingBlobRefStreamProvider::default(),
copy_buffer_size: 4 * 1024,
bytes_written: 0,
lengths: Vec::new(),
diff --git a/crates/paimon/src/arrow/format/blob/writer_uri_reader_tests.rs
b/crates/paimon/src/arrow/format/blob/writer_uri_reader_tests.rs
new file mode 100644
index 00000000..027d24b4
--- /dev/null
+++ b/crates/paimon/src/arrow/format/blob/writer_uri_reader_tests.rs
@@ -0,0 +1,350 @@
+// 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, UriReader};
+use arrow_array::ArrayRef;
+use std::sync::Mutex;
+
+#[derive(Default)]
+struct SourceState {
+ events: Vec<String>,
+ no_seek: bool,
+ fail_read: bool,
+ fail_close: bool,
+ oversized_read: bool,
+}
+
+struct Reader(Arc<Mutex<SourceState>>);
+struct Factory(Arc<Reader>);
+struct Stream {
+ state: Arc<Mutex<SourceState>>,
+ bytes: Bytes,
+ position: usize,
+ uri: String,
+}
+
+impl UriReaderFactory for Factory {
+ fn create(&self, _: &str) -> Result<Arc<dyn UriReader>> {
+ Ok(self.0.clone())
+ }
+}
+
+#[async_trait]
+impl UriReader for Reader {
+ async fn new_input_stream(&self, uri: &str) -> Result<Box<dyn
UriInputStream>> {
+ self.0.lock().unwrap().events.push(format!("open:{uri}"));
+ Ok(Box::new(Stream {
+ state: self.0.clone(),
+ bytes: Bytes::from_static(b"0123456789"),
+ position: 0,
+ uri: uri.into(),
+ }))
+ }
+}
+
+#[async_trait]
+impl UriInputStream for Stream {
+ async fn read(&mut self, length: usize) -> Result<Bytes> {
+ let mut state = self.state.lock().unwrap();
+ state
+ .events
+ .push(format!("read:{}:{length}", self.position));
+ if state.fail_read {
+ return Err(invalid("source read failed"));
+ }
+ if state.oversized_read {
+ return Ok(Bytes::from(vec![b'x'; length + 1]));
+ }
+ // Short, successful reads are valid InputStream behavior.
+ let end = (self.position + length.min(2)).min(self.bytes.len());
+ let bytes = self.bytes.slice(self.position..end);
+ self.position = end;
+ Ok(bytes)
+ }
+ async fn seek(&mut self, offset: u64) -> Result<()> {
+ let mut state = self.state.lock().unwrap();
+ state.events.push(format!("seek:{offset}"));
+ if state.no_seek {
+ return Err(invalid("seek unsupported"));
+ }
+ self.position = offset as usize;
+ Ok(())
+ }
+ async fn close(&mut self) -> Result<()> {
+ let mut state = self.state.lock().unwrap();
+ state.events.push(format!("close:{}", self.uri));
+ if state.fail_close {
+ return Err(invalid("source close failed"));
+ }
+ Ok(())
+ }
+}
+
+fn reference(uri: &str, offset: i64, length: i64) -> Vec<u8> {
+ BlobDescriptor::new(uri.into(), offset, length).serialize()
+}
+
+fn batch(values: &[Vec<u8>]) -> RecordBatch {
+ RecordBatch::try_from_iter([(
+ "payload",
+ Arc::new(LargeBinaryArray::from(
+ values
+ .iter()
+ .map(|value| Some(value.as_slice()))
+ .collect::<Vec<_>>(),
+ )) as ArrayRef,
+ )])
+ .unwrap()
+}
+
+async fn writer(state: &Arc<Mutex<SourceState>>) -> (FileIO, Box<dyn
FormatFileWriter>) {
+ let io = FileIOBuilder::new("memory").build().unwrap();
+ let output = io.new_output("memory:/output.blob").unwrap();
+ let factory = BlobWriterFactory::new(None, None, None)
+ .unwrap()
+
.with_uri_reader_factory(Arc::new(Factory(Arc::new(Reader(state.clone())))));
+ (io, factory.create_writer(&output, "").await.unwrap())
+}
+
+fn payloads(bytes: &Bytes, values: &[&[u8]]) {
+ let mut position = 0;
+ for value in values {
+ assert_eq!(&bytes[position + 4..position + 4 + value.len()], *value);
+ position += value.len() + 16;
+ }
+}
+
+#[tokio::test]
+async fn references_reuse_seekable_sources_without_reading_outside_windows() {
+ let state = Arc::new(Mutex::new(SourceState::default()));
+ let (io, mut writer) = writer(&state).await;
+ writer
+ .write(&batch(&[
+ reference("custom:/a", 2, 3),
+ reference("custom:/a", 5, 2),
+ reference("custom:/a", 1, 2),
+ b"inline".to_vec(),
+ reference("custom:/a", 3, 0),
+ ]))
+ .await
+ .unwrap();
+ writer.close().await.unwrap();
+ payloads(
+ &io.new_input("memory:/output.blob")
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ &[b"234", b"56", b"12", b"inline", b""],
+ );
+ assert_eq!(
+ state.lock().unwrap().events,
+ [
+ "open:custom:/a",
+ "seek:2",
+ "read:2:3",
+ "read:4:1",
+ "read:5:2",
+ "seek:1",
+ "read:1:2",
+ "close:custom:/a",
+ ]
+ );
+}
+
+#[tokio::test]
+async fn non_seekable_sources_reuse_adjacent_windows_and_reopen_to_rewind() {
+ let state = Arc::new(Mutex::new(SourceState {
+ no_seek: true,
+ ..Default::default()
+ }));
+ let (io, mut writer) = writer(&state).await;
+ writer
+ .write(&batch(&[
+ reference("custom:/a", 0, 2),
+ reference("custom:/a", 2, 2),
+ reference("custom:/a", 0, 2),
+ ]))
+ .await
+ .unwrap();
+ writer.close().await.unwrap();
+ payloads(
+ &io.new_input("memory:/output.blob")
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ &[b"01", b"23", b"01"],
+ );
+ assert_eq!(
+ state.lock().unwrap().events,
+ [
+ "open:custom:/a",
+ "read:0:2",
+ "read:2:2",
+ "seek:0",
+ "close:custom:/a",
+ "open:custom:/a",
+ "read:0:2",
+ "close:custom:/a",
+ ]
+ );
+}
+
+#[tokio::test]
+async fn unknown_length_reads_to_eof_and_preserves_the_cached_bounded_source()
{
+ let state = Arc::new(Mutex::new(SourceState::default()));
+ let (io, mut writer) = writer(&state).await;
+ writer
+ .write(&batch(&[
+ reference("custom:/a", 0, 2),
+ reference("custom:/a", 3, -1),
+ reference("custom:/a", 2, 2),
+ ]))
+ .await
+ .unwrap();
+ writer.close().await.unwrap();
+ payloads(
+ &io.new_input("memory:/output.blob")
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ &[b"01", b"3456789", b"23"],
+ );
+ let events = &state.lock().unwrap().events;
+ assert_eq!(
+ events
+ .iter()
+ .filter(|event| event.starts_with("open:"))
+ .count(),
+ 2
+ );
+ assert_eq!(
+ events
+ .iter()
+ .filter(|event| event.starts_with("close:"))
+ .count(),
+ 2
+ );
+ assert_eq!(&events[events.len() - 2..], ["read:2:2", "close:custom:/a"]);
+}
+
+#[tokio::test]
+async fn copy_failures_discard_sources_and_keep_the_copy_error_primary() {
+ for (fail_read, oversized_read, length, message) in [
+ (true, false, 3, "source read failed"),
+ (false, false, 12, "Unexpected EOF"),
+ (false, true, 3, "more bytes than requested"),
+ ] {
+ let state = Arc::new(Mutex::new(SourceState {
+ fail_read,
+ oversized_read,
+ fail_close: true,
+ ..Default::default()
+ }));
+ let (_, mut writer) = writer(&state).await;
+ let error = writer
+ .write(&batch(&[reference("custom:/a", 0, length)]))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains(message), "{error}");
+ writer.close().await.unwrap();
+ assert_eq!(
+ state
+ .lock()
+ .unwrap()
+ .events
+ .iter()
+ .filter(|event| event.starts_with("close:"))
+ .count(),
+ 1
+ );
+ }
+}
+
+#[tokio::test]
+async fn a_previous_source_close_failure_prevents_opening_a_different_uri() {
+ let state = Arc::new(Mutex::new(SourceState {
+ fail_close: true,
+ ..Default::default()
+ }));
+ let (_, mut writer) = writer(&state).await;
+ let error = writer
+ .write(&batch(&[
+ reference("custom:/a", 0, 2),
+ reference("custom:/b", 0, 2),
+ ]))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("source close failed"));
+ writer.close().await.unwrap();
+ assert_eq!(
+ state.lock().unwrap().events,
+ ["open:custom:/a", "read:0:2", "close:custom:/a"]
+ );
+}
+
+#[tokio::test]
+async fn failed_initial_seek_closes_the_opened_source() {
+ let state = Arc::new(Mutex::new(SourceState {
+ no_seek: true,
+ fail_close: true,
+ ..Default::default()
+ }));
+ let (_, mut writer) = writer(&state).await;
+ let error = writer
+ .write(&batch(&[reference("custom:/a", 1, 2)]))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("seek unsupported"));
+ writer.close().await.unwrap();
+ assert_eq!(
+ state.lock().unwrap().events,
+ ["open:custom:/a", "seek:1", "close:custom:/a"]
+ );
+}
+
+#[tokio::test]
+async fn source_close_failure_surfaces_at_format_prepare() {
+ let state = Arc::new(Mutex::new(SourceState {
+ fail_close: true,
+ ..Default::default()
+ }));
+ let (io, mut writer) = writer(&state).await;
+ writer
+ .write(&batch(&[reference("custom:/a", 0, 2)]))
+ .await
+ .unwrap();
+ assert!(writer
+ .close()
+ .await
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("source close failed"));
+ // The destination still closes and drains its buffered payload.
+ payloads(
+ &io.new_input("memory:/output.blob")
+ .unwrap()
+ .read()
+ .await
+ .unwrap(),
+ &[b"01"],
+ );
+}
diff --git a/crates/paimon/src/io/mod.rs b/crates/paimon/src/io/mod.rs
index c8dd33de..7b3023d2 100644
--- a/crates/paimon/src/io/mod.rs
+++ b/crates/paimon/src/io/mod.rs
@@ -19,6 +19,7 @@ pub(crate) mod cache;
mod file_io;
pub(crate) mod uri_reader;
pub use file_io::*;
+pub use uri_reader::{UriInputStream, UriReader, UriReaderFactory};
mod storage;
pub use storage::*;
diff --git a/crates/paimon/src/io/uri_reader.rs
b/crates/paimon/src/io/uri_reader.rs
index 6324671d..d6e25cfe 100644
--- a/crates/paimon/src/io/uri_reader.rs
+++ b/crates/paimon/src/io/uri_reader.rs
@@ -25,6 +25,10 @@ use reqwest::header::CONTENT_ENCODING;
use std::ops::Range;
use std::sync::{Arc, LazyLock};
+mod custom;
+pub(crate) use custom::ReusingBlobRefStreamProvider;
+pub use custom::{UriInputStream, UriReader, UriReaderFactory};
+
pub(crate) enum UriInput {
File(InputFile),
Http(HttpReader),
diff --git a/crates/paimon/src/io/uri_reader/custom.rs
b/crates/paimon/src/io/uri_reader/custom.rs
new file mode 100644
index 00000000..264b25fc
--- /dev/null
+++ b/crates/paimon/src/io/uri_reader/custom.rs
@@ -0,0 +1,122 @@
+// 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.
+
+//! Custom URI readers and Java's consecutive BlobRef source reuse.
+
+use crate::{Error, Result};
+use bytes::Bytes;
+use std::sync::Arc;
+
+/// Select a URI reader, corresponding to Java's UriReaderFactory.create.
+/// Return the same Arc when references share the same reader identity.
+pub trait UriReaderFactory: Send + Sync {
+ fn create(&self, uri: &str) -> Result<Arc<dyn UriReader>>;
+}
+
+/// Open a source stream, corresponding to Java's UriReader.newInputStream.
+#[async_trait::async_trait]
+pub trait UriReader: Send + Sync {
+ async fn new_input_stream(&self, uri: &str) -> Result<Box<dyn
UriInputStream>>;
+}
+
+/// A Blob source stream. Empty reads mean EOF, and reads must not exceed the
+/// requested length. Seek is optional for offset-zero or consecutive windows.
+#[async_trait::async_trait]
+pub trait UriInputStream: Send {
+ async fn read(&mut self, length: usize) -> Result<Bytes>;
+
+ async fn seek(&mut self, _offset: u64) -> Result<()> {
+ Err(Error::Unsupported {
+ message: "URI stream does not support seek".into(),
+ })
+ }
+
+ async fn close(&mut self) -> Result<()> {
+ Ok(())
+ }
+}
+
+struct Source {
+ reader: Arc<dyn UriReader>,
+ uri: String,
+ stream: Box<dyn UriInputStream>,
+ position: u64,
+}
+
+/// Keep one bounded-reference source per physical Blob writer. Unknown-length
+/// references open their own streams, as Java ReusingBlobRefStreamProvider
does.
+#[derive(Default)]
+pub(crate) struct ReusingBlobRefStreamProvider {
+ source: Option<Source>,
+}
+
+impl ReusingBlobRefStreamProvider {
+ pub(crate) async fn prepare(
+ &mut self,
+ reader: Arc<dyn UriReader>,
+ uri: &str,
+ offset: u64,
+ ) -> Result<()> {
+ if let Some(source) = &mut self.source {
+ if Arc::ptr_eq(&source.reader, &reader) && source.uri == uri {
+ if source.position == offset {
+ return Ok(());
+ }
+ if source.stream.seek(offset).await.is_ok() {
+ source.position = offset;
+ return Ok(());
+ }
+ }
+ // A previous source's close failure must be surfaced before
opening
+ // this reference; do not mistake cleanup failure for a fetch
failure.
+ self.close().await?;
+ }
+ let mut stream = reader.new_input_stream(uri).await?;
+ if offset != 0 {
+ if let Err(error) = stream.seek(offset).await {
+ let _ = stream.close().await;
+ return Err(error);
+ }
+ }
+ self.source = Some(Source {
+ reader,
+ uri: uri.into(),
+ stream,
+ position: offset,
+ });
+ Ok(())
+ }
+
+ pub(crate) fn stream(&mut self) -> &mut dyn UriInputStream {
+ self.source
+ .as_mut()
+ .expect("source prepared")
+ .stream
+ .as_mut()
+ }
+
+ pub(crate) fn advance(&mut self, length: u64) {
+ self.source.as_mut().expect("source prepared").position += length;
+ }
+
+ pub(crate) async fn close(&mut self) -> Result<()> {
+ match self.source.take() {
+ Some(mut source) => source.stream.close().await,
+ None => Ok(()),
+ }
+ }
+}
diff --git a/crates/paimon/src/table/blob_uri_reader_tests.rs
b/crates/paimon/src/table/blob_uri_reader_tests.rs
new file mode 100644
index 00000000..e9b4f868
--- /dev/null
+++ b/crates/paimon/src/table/blob_uri_reader_tests.rs
@@ -0,0 +1,149 @@
+// 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, UriInputStream, UriReader, UriReaderFactory};
+use crate::spec::{BlobDescriptor, BlobType, DataType, IntType, Schema,
TableSchema};
+use crate::Result;
+use arrow_array::{ArrayRef, Int32Array, LargeBinaryArray, RecordBatch};
+use bytes::Bytes;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+#[derive(Default)]
+struct Reader {
+ opened: AtomicUsize,
+ closed: Arc<AtomicUsize>,
+}
+
+struct Factory(Arc<Reader>);
+
+impl UriReaderFactory for Factory {
+ fn create(&self, uri: &str) -> Result<Arc<dyn UriReader>> {
+ assert_eq!(uri, "custom://payload");
+ Ok(self.0.clone())
+ }
+}
+
+#[async_trait::async_trait]
+impl UriReader for Reader {
+ async fn new_input_stream(&self, _uri: &str) -> Result<Box<dyn
UriInputStream>> {
+ self.opened.fetch_add(1, Ordering::Relaxed);
+ Ok(Box::new(Stream {
+ position: 0,
+ closed: self.closed.clone(),
+ }))
+ }
+}
+
+struct Stream {
+ position: usize,
+ closed: Arc<AtomicUsize>,
+}
+
+#[async_trait::async_trait]
+impl UriInputStream for Stream {
+ async fn read(&mut self, length: usize) -> Result<Bytes> {
+ let bytes = Bytes::from_static(b"payload");
+ let end = (self.position + length).min(bytes.len());
+ let result = bytes.slice(self.position..end);
+ self.position = end;
+ Ok(result)
+ }
+
+ async fn close(&mut self) -> Result<()> {
+ self.closed.fetch_add(1, Ordering::Relaxed);
+ Ok(())
+ }
+}
+
+fn table(bucket: &str) -> Table {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("payload", DataType::Blob(BlobType::new()))
+ .primary_key(["id"])
+ .option("bucket", bucket)
+ .option("postpone.default-bucket-num", "2")
+ .option("blob.target-file-size", "1 MB")
+ .build()
+ .unwrap();
+ Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("db", "t"),
+ format!("memory:/uri-reader-{bucket}"),
+ TableSchema::new(0, &schema),
+ None,
+ )
+}
+
+fn batch() -> RecordBatch {
+ let descriptor = BlobDescriptor::new("custom://payload".into(), 0,
7).serialize();
+ RecordBatch::try_from_iter([
+ ("id", Arc::new(Int32Array::from(vec![1])) as ArrayRef),
+ (
+ "payload",
+ Arc::new(LargeBinaryArray::from(vec![descriptor.as_slice()])) as
ArrayRef,
+ ),
+ ])
+ .unwrap()
+}
+
+#[tokio::test]
+async fn
managed_blob_abort_closes_custom_source_before_removing_unprepared_files() {
+ for bucket in ["1", "-1", "-2"] {
+ let table = table(bucket);
+ let reader = Arc::new(Reader::default());
+ let mut writer = TableWrite::new(&table, "uri-reader".into()).unwrap();
+ writer
+
.with_blob_uri_reader_factory(Some(Arc::new(Factory(reader.clone()))))
+ .unwrap();
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ assert_eq!(reader.opened.load(Ordering::Relaxed), 1);
+ assert_eq!(reader.closed.load(Ordering::Relaxed), 0);
+ writer.close().await;
+ assert_eq!(reader.closed.load(Ordering::Relaxed), 1, "bucket
{bucket}");
+ writer.close().await;
+ assert_eq!(reader.closed.load(Ordering::Relaxed), 1);
+ assert!(table
+ .file_io()
+ .list_status_recursive(table.location())
+ .await
+ .unwrap()
+ .is_empty());
+ }
+}
+
+#[tokio::test]
+async fn postpone_fixed_blob_reader_factory_cannot_change_after_writing() {
+ let table = table("-2");
+ let reader = Arc::new(Reader::default());
+ let mut writer = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .new_write()
+ .unwrap();
+ writer
+ .with_blob_uri_reader_factory(Some(Arc::new(Factory(reader.clone()))))
+ .unwrap();
+ writer.write_arrow_batch(&batch()).await.unwrap();
+ let error = writer.with_blob_uri_reader_factory(None).err().unwrap();
+ assert!(error.to_string().contains("before any write"));
+ writer.close().await;
+ assert_eq!(reader.opened.load(Ordering::Relaxed), 1);
+ assert_eq!(reader.closed.load(Ordering::Relaxed), 1);
+}
diff --git a/crates/paimon/src/table/data_file_writer.rs
b/crates/paimon/src/table/data_file_writer.rs
index 6743bdd3..14efe59d 100644
--- a/crates/paimon/src/table/data_file_writer.rs
+++ b/crates/paimon/src/table/data_file_writer.rs
@@ -163,23 +163,28 @@ impl DataFileWriter {
}
/// Configure only the dedicated BLOB factory; ordinary format factories
- /// and their write paths do not carry callback state.
- pub(super) fn set_blob_consumer(
+ /// and their write paths do not carry Blob-specific state.
+ pub(super) fn set_blob_writer_options(
&mut self,
- consumer: Arc<dyn crate::spec::BlobConsumer>,
+ consumer: Option<Arc<dyn crate::spec::BlobConsumer>>,
+ uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
) -> 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),
- ));
+ let mut factory = crate::arrow::format::blob::BlobWriterFactory::new(
+ Some(self.file_io.clone()),
+ self.write_fields.first(),
+ Some(&self.format_options),
+ )?;
// Java BlobFormatWriter relinquishes deletion rights whenever a
// consumer can expose a descriptor, even if a later callback fails.
- self.delete_file_upon_abort = false;
+ self.delete_file_upon_abort = consumer.is_none();
+ if let Some(consumer) = consumer {
+ factory = factory.with_consumer(consumer);
+ }
+ if let Some(reader) = uri_reader_factory {
+ factory = factory.with_uri_reader_factory(reader);
+ }
+ self.format_writer_factory = Some(Arc::new(factory));
Ok(())
}
diff --git a/crates/paimon/src/table/dedicated_format_file_writer.rs
b/crates/paimon/src/table/dedicated_format_file_writer.rs
index b2874be8..a07d9829 100644
--- a/crates/paimon/src/table/dedicated_format_file_writer.rs
+++ b/crates/paimon/src/table/dedicated_format_file_writer.rs
@@ -258,13 +258,15 @@ impl AppendDedicatedFormatFileWriter {
self
}
- pub(crate) fn with_blob_consumer(
+ pub(crate) fn with_blob_writer_options(
mut self,
consumer: Option<Arc<dyn crate::spec::BlobConsumer>>,
+ uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
) -> Result<Self> {
- if let Some(consumer) = consumer {
+ if consumer.is_some() || uri_reader_factory.is_some() {
for blob in &mut self.blob_writers {
- blob.writer.set_blob_consumer(consumer.clone())?;
+ blob.writer
+ .set_blob_writer_options(consumer.clone(),
uri_reader_factory.clone())?;
}
}
Ok(self)
@@ -506,12 +508,15 @@ mod tests {
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)
- },
- )))
+ .with_blob_writer_options(
+ Some(Arc::new(
+ move |_: &str, descriptor:
Option<&crate::spec::BlobDescriptor>| {
+
received.lock().unwrap().push(descriptor.unwrap().clone());
+ Ok(false)
+ },
+ )),
+ None,
+ )
.unwrap();
}
writer.write(&batch).await.unwrap();
diff --git a/crates/paimon/src/table/kv_file_writer.rs
b/crates/paimon/src/table/kv_file_writer.rs
index eee20858..d12c2b26 100644
--- a/crates/paimon/src/table/kv_file_writer.rs
+++ b/crates/paimon/src/table/kv_file_writer.rs
@@ -201,6 +201,14 @@ impl KeyValueFileWriter {
})
}
+ pub(super) fn with_blob_uri_reader_factory(
+ mut self,
+ factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ ) -> Self {
+ self.managed_blob_writer.set_uri_reader_factory(factory);
+ self
+ }
+
/// Construct one reusable factory for this flush's data or changelog
output.
/// Statistics retain logical value fields independently of the physical
KV schema.
fn new_format_writer_factory(
diff --git a/crates/paimon/src/table/managed_blob_writer.rs
b/crates/paimon/src/table/managed_blob_writer.rs
index 402ddc0b..dc7a8d29 100644
--- a/crates/paimon/src/table/managed_blob_writer.rs
+++ b/crates/paimon/src/table/managed_blob_writer.rs
@@ -88,6 +88,7 @@ pub(crate) struct ManagedBlobWriter {
file_prefix: String,
target_file_size: u64,
copy_buffer_size: usize,
+ uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
fields: Vec<ManagedBlobField>,
uncommitted_paths: Vec<String>,
}
@@ -124,6 +125,15 @@ impl ManagedBlobWriteState {
self.writer.is_some()
}
+ pub(crate) fn set_uri_reader_factory(
+ &mut self,
+ factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ ) {
+ if let Some(writer) = &mut self.writer {
+ writer.get_mut().uri_reader_factory = factory;
+ }
+ }
+
pub(crate) async fn externalize(&mut self, batch: RecordBatch) ->
Result<RecordBatch> {
match &mut self.writer {
Some(writer) => writer.get_mut().externalize(&batch).await,
@@ -169,6 +179,7 @@ impl ManagedBlobWriter {
file_prefix: file_prefix.to_string(),
target_file_size,
copy_buffer_size: 4 * 1024,
+ uri_reader_factory: None,
fields: fields
.into_iter()
.map(|(index, kind)| ManagedBlobField {
@@ -364,7 +375,8 @@ impl ManagedBlobWriter {
let writer = Box::new(
BlobFormatWriter::new(&output, Some(self.file_io.clone()),
None)
.await?
- .with_copy_buffer_size(self.copy_buffer_size),
+ .with_copy_buffer_size(self.copy_buffer_size)
+ .with_uri_reader_factory(self.uri_reader_factory.clone()),
);
self.uncommitted_paths.push(path.clone());
self.fields[field_index].current = Some(ManagedBlobPack { path,
writer });
@@ -395,7 +407,11 @@ impl ManagedBlobWriter {
pub(crate) async fn abort(&mut self) {
for field in &mut self.fields {
- field.current.take();
+ if let Some(pack) = field.current.take() {
+ // Finish source and destination cleanup before removing an
+ // unprepared pack, like Java abortCurrent's closeQuietly.
+ let _ = pack.writer.close().await;
+ }
}
for path in self.uncommitted_paths.drain(..) {
let _ = self.file_io.delete_file(&path).await;
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index ed7e6a06..168f3d91 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -27,6 +27,8 @@ mod bitmap_global_index_writer;
#[cfg(test)]
mod blob_consumer_tests;
mod blob_resolver;
+#[cfg(test)]
+mod blob_uri_reader_tests;
mod branch_manager;
#[cfg(test)]
mod branch_write_tests;
diff --git a/crates/paimon/src/table/postpone_file_writer.rs
b/crates/paimon/src/table/postpone_file_writer.rs
index 89f3257b..449d4471 100644
--- a/crates/paimon/src/table/postpone_file_writer.rs
+++ b/crates/paimon/src/table/postpone_file_writer.rs
@@ -94,6 +94,14 @@ pub(crate) struct PostponeFileWriter {
}
impl PostponeFileWriter {
+ pub(super) fn with_blob_uri_reader_factory(
+ mut self,
+ factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ ) -> Self {
+ self.managed_blob_writer.set_uri_reader_factory(factory);
+ self
+ }
+
pub(crate) fn new(file_io: FileIO, config: PostponeWriteConfig) ->
Result<Self> {
let paths = DataFilePathFactory::new(
&config.table_location,
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
index 32ba67fe..6f2f59a9 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
@@ -38,6 +38,15 @@ pub struct PostponeFixedBucketTableWrite {
}
impl PostponeFixedBucketTableWrite {
+ pub fn with_blob_uri_reader_factory(
+ &mut self,
+ factory: Option<std::sync::Arc<dyn crate::io::UriReaderFactory>>,
+ ) -> Result<&mut Self> {
+ self.ensure_writable()?;
+ self.inner.with_blob_uri_reader_factory(factory)?;
+ Ok(self)
+ }
+
pub(crate) fn new(
table: &Table,
commit_user: String,
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index 9d7d2f9e..4e8032be 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -178,6 +178,7 @@ pub struct TableWrite {
/// 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>>,
+ blob_uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
/// Dedicated vector-store file format, when configured.
vector_file_format: Option<String>,
/// Whether the table has VECTOR fields requiring dedicated vector files.
@@ -249,6 +250,7 @@ impl TableWrite {
blob_inline_fields: HashSet::new(),
has_blob_fields: false,
blob_consumer: None,
+ blob_uri_reader_factory: None,
vector_file_format: None,
has_dedicated_vector_fields: false,
row_kind_generator: None,
@@ -532,6 +534,7 @@ impl TableWrite {
blob_inline_fields,
has_blob_fields,
blob_consumer: None,
+ blob_uri_reader_factory: None,
vector_file_format,
has_dedicated_vector_fields,
row_kind_generator,
@@ -737,6 +740,30 @@ impl TableWrite {
Ok(self)
}
+ /// Select URI readers for referenced Blob payloads before writing any
data.
+ /// Source selection, bounded copies, stream reuse and cleanup are handled
+ /// by core, as with Java UriReaderFactory and BlobRef.
+ pub fn with_blob_uri_reader_factory(
+ &mut self,
+ factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ ) -> Result<&mut Self> {
+ self.ensure_active()?;
+ if self.written {
+ return Err(crate::Error::DataInvalid {
+ message: "with_blob_uri_reader_factory must be called before
any write operation"
+ .into(),
+ source: None,
+ });
+ }
+ if self.format_writer.is_some() {
+ return Err(crate::Error::Unsupported {
+ message: "with_blob_uri_reader_factory requires a Paimon
table".into(),
+ });
+ }
+ self.blob_uri_reader_factory = factory;
+ 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() {
@@ -790,6 +817,11 @@ impl TableWrite {
bucket: i32,
batch: RecordBatch,
) -> Result<()> {
+ self.ensure_active()?;
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+ self.written = true;
self.write_bucket(partition, bucket, batch).await
}
@@ -1383,7 +1415,10 @@ impl TableWrite {
self.table.schema().options(),
&self.blob_inline_fields,
)?
- .with_blob_consumer(self.blob_consumer.clone())?
+ .with_blob_writer_options(
+ self.blob_consumer.clone(),
+ self.blob_uri_reader_factory.clone(),
+ )?
.with_resources(self.resources.clone()),
)))
} else {
@@ -1460,6 +1495,7 @@ impl TableWrite {
file_index_options: self.file_index_options.clone(),
},
)?
+ .with_blob_uri_reader_factory(self.blob_uri_reader_factory.clone())
.with_resources(self.resources.clone()),
))
}
@@ -1530,6 +1566,7 @@ impl TableWrite {
},
next_seq,
)?
+ .with_blob_uri_reader_factory(self.blob_uri_reader_factory.clone())
.with_resources(self.resources.clone()),
))
}
