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 c60d70ee [core] Expose postpone bucket plans through the write builder
(#1200)
c60d70ee is described below
commit c60d70eedeeab7a62780c3f0216d12edece5c4ef
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Oct 22 22:31:25 2026 +0800
[core] Expose postpone bucket plans through the write builder (#1200)
---
bindings/python/project-description.md | 32 +++
bindings/python/python/pypaimon_rust/__init__.py | 7 +
.../python/python/pypaimon_rust/datafusion.pyi | 41 ++--
.../pypaimon_rust/{__init__.py => uri_reader.py} | 18 +-
bindings/python/src/write.rs | 213 ++++++++++++++++++--
bindings/python/tests/test_blob_uri_reader.py | 4 +-
bindings/python/tests/test_table_commit.py | 12 +-
bindings/python/tests/test_uri_reader_protocols.py | 61 ++++++
crates/paimon/src/table/authorized_write_tests.rs | 53 +++++
crates/paimon/src/table/postpone_bucket_plan.rs | 136 ++++++++++++-
.../table/postpone_fixed_bucket_write_builder.rs | 217 ++++++++++++++++++++-
crates/paimon/src/table/table_upsert.rs | 7 +-
12 files changed, 748 insertions(+), 53 deletions(-)
diff --git a/bindings/python/project-description.md
b/bindings/python/project-description.md
index 44633f86..86c40860 100644
--- a/bindings/python/project-description.md
+++ b/bindings/python/project-description.md
@@ -53,6 +53,38 @@ ctx.sql("INSERT INTO paimon.my_db.users VALUES (1, 'alice'),
(2, 'bob')")
batches = ctx.sql("SELECT * FROM paimon.my_db.users")
```
+### Blob Source Streams and Commit Identity
+
+The write and update APIs accept a `UriReaderFactory` for custom Blob reference
+sources. Both stream protocols are public runtime imports and support
structural
+typing; implementations do not need to inherit from them.
+
+```python
+from io import BytesIO
+from pypaimon_rust.datafusion import UriReader, UriReaderFactory
+
+payloads = {"custom://example": b"content"}
+
+class Reader:
+ def new_input_stream(self, uri: str):
+ return BytesIO(payloads[uri])
+
+class Factory:
+ def create(self, uri: str) -> UriReader:
+ return Reader()
+
+factory: UriReaderFactory = Factory()
+builder = table.new_batch_write_builder().with_commit_user("my-job")
+writer = builder.new_write().with_blob_uri_reader_factory(factory)
+update = builder.new_update().with_blob_uri_reader_factory(factory)
+```
+
+Rust copies the requested Blob ranges and closes the source streams. Pass
`None`
+to clear the factory. Configure a writer before its first write. Batch and
stream
+updates, including incremental row-ID updaters, expose the same chainable
setter.
+`with_commit_user` is available on batch, stream, and postpone fixed bucket
write
+builders, and the identity applies to both writers and committers from that
builder.
+
### Multimodal SQL Helpers
`SQLContext` also registers Python scalar UDFs for common BLOB media and vector
diff --git a/bindings/python/python/pypaimon_rust/__init__.py
b/bindings/python/python/pypaimon_rust/__init__.py
index b36002b7..9838f8d2 100644
--- a/bindings/python/python/pypaimon_rust/__init__.py
+++ b/bindings/python/python/pypaimon_rust/__init__.py
@@ -16,3 +16,10 @@
# under the License.
from .pypaimon_rust import *
+from .pypaimon_rust import datafusion as datafusion
+from .uri_reader import UriReader as UriReader, UriReaderFactory as
UriReaderFactory
+
+# The native submodule is registered during extension initialization. Export
the
+# Python stream protocols beside its writers and update APIs.
+datafusion.UriReader = UriReader
+datafusion.UriReaderFactory = UriReaderFactory
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index ad573e03..427e9810 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -18,22 +18,18 @@
from datetime import datetime
from os import PathLike
from typing import (
- Any, BinaryIO, Callable, Dict, Iterable, List, Literal, Optional,
Protocol, Sequence, Tuple,
+ Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple,
TypeAlias, Union,
)
import pyarrow
+from .uri_reader import UriReader as UriReader, UriReaderFactory as
UriReaderFactory
+
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: ...
@@ -462,7 +458,7 @@ class CommitMessage:
class BatchTableWrite:
def with_blob_uri_reader_factory(
- self, factory: Optional[_UriReaderFactory]
+ self, factory: Optional[UriReaderFactory]
) -> "BatchTableWrite":
"""Use custom source streams for Blob references; Rust copies bounded
ranges, reuses consecutive sources and closes them on completion."""
@@ -481,7 +477,9 @@ class BatchTableWrite:
...
class BatchTableUpdate:
- def _with_blob_uri_reader_factory(self, factory: Optional[Any]) -> None:
...
+ def with_blob_uri_reader_factory(
+ self, factory: Optional[UriReaderFactory]
+ ) -> "BatchTableUpdate": ...
def merge_into(
self, source: Optional[Union[pyarrow.Table, Table]], *, on:
Sequence[Tuple[str, str]],
when_matched: Sequence[Dict[str, Any]], when_not_matched:
Sequence[Dict[str, Any]]
@@ -531,7 +529,9 @@ class BatchTableUpdate:
...
class TableUpdateByRowId:
- def _with_blob_uri_reader_factory(self, factory: Optional[Any]) -> None:
...
+ def with_blob_uri_reader_factory(
+ self, factory: Optional[UriReaderFactory]
+ ) -> "TableUpdateByRowId": ...
def update_columns(self, data: pyarrow.Table, column_names: List[str]) ->
List[CommitMessage]:
"""Stage one input table; return the cumulative messages to commit
once."""
...
@@ -540,7 +540,9 @@ class TableUpdateByRowId:
def _abort(self) -> None: ...
class StreamTableUpdate:
- def _with_blob_uri_reader_factory(self, factory: Optional[Any]) -> None:
...
+ def with_blob_uri_reader_factory(
+ self, factory: Optional[UriReaderFactory]
+ ) -> "StreamTableUpdate": ...
def merge_into(
self, source: Optional[Union[pyarrow.Table, Table]], *, on:
Sequence[Tuple[str, str]],
when_matched: Sequence[Dict[str, Any]], when_not_matched:
Sequence[Dict[str, Any]],
@@ -591,7 +593,7 @@ class StreamTableUpdate:
class StreamTableWrite:
def with_blob_uri_reader_factory(
- self, factory: Optional[_UriReaderFactory]
+ self, factory: Optional[UriReaderFactory]
) -> "StreamTableWrite": ...
def with_write_type(self, columns: List[str]) -> "StreamTableWrite": ...
def with_blob_consumer(
@@ -740,8 +742,8 @@ class StreamTableCommit:
...
class BatchWriteBuilder:
- def _with_commit_user(self, commit_user: str) -> "BatchWriteBuilder":
- """Internal PyPaimon bridge for commits produced by an external Python
writer."""
+ def with_commit_user(self, commit_user: str) -> "BatchWriteBuilder":
+ """Set the commit identity shared by this builder's writers and
committers."""
...
def with_overwrite(self, static_partition: Optional[Dict[str, Any]] = {})
-> "BatchWriteBuilder":
"""Configure both writer and committer. Explicit None restores append.
@@ -772,7 +774,16 @@ class PostponeFixedBucketWriteBuilder:
Unknown partitions buffer until preparation; known layouts write directly.
Rescaling and compaction are not performed.
"""
- def _with_commit_user(self, commit_user: str) ->
"PostponeFixedBucketWriteBuilder": ...
+ def with_commit_user(self, commit_user: str) ->
"PostponeFixedBucketWriteBuilder": ...
+ def resolve_bucket_plan(self) -> pyarrow.RecordBatch:
+ """Known partition layouts from the latest writer snapshot, without
SELECT.
+
+ Uses the with_bucket_plan schema. Explicit plans take precedence;
+ overwrite otherwise returns no existing layouts. Configured defaults
+ do not create partition rows in this metadata plan.
+ This method does not infer bucket counts for new input partitions.
+ """
+ ...
def with_bucket_plan(self, plan: pyarrow.RecordBatch) ->
"PostponeFixedBucketWriteBuilder":
"""Partition columns in schema order followed by total_buckets: int32.
diff --git a/bindings/python/python/pypaimon_rust/__init__.py
b/bindings/python/python/pypaimon_rust/uri_reader.py
similarity index 62%
copy from bindings/python/python/pypaimon_rust/__init__.py
copy to bindings/python/python/pypaimon_rust/uri_reader.py
index b36002b7..3beb5d7e 100644
--- a/bindings/python/python/pypaimon_rust/__init__.py
+++ b/bindings/python/python/pypaimon_rust/uri_reader.py
@@ -15,4 +15,20 @@
# specific language governing permissions and limitations
# under the License.
-from .pypaimon_rust import *
+"""Source stream contracts for managed Blob URI reads."""
+
+from typing import BinaryIO, Protocol
+
+__all__ = ["UriReader", "UriReaderFactory"]
+
+
+class UriReader(Protocol):
+ """Open readable binary streams; Rust closes each stream after use."""
+
+ def new_input_stream(self, uri: str) -> BinaryIO: ...
+
+
+class UriReaderFactory(Protocol):
+ """Select a URI reader while leaving payload copying to Rust Core."""
+
+ def create(self, uri: str) -> UriReader: ...
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 1ac21b05..299730b0 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -19,7 +19,7 @@ use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use arrow::datatypes::Schema as ArrowSchema;
-use arrow::pyarrow::FromPyArrow;
+use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use arrow::record_batch::RecordBatch;
use paimon::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataType, Datum,
RowType};
use paimon::table::{
@@ -165,8 +165,8 @@ impl PyBatchWriteBuilder {
#[pymethods]
impl PyBatchWriteBuilder {
- /// Internal PyPaimon bridge: retain the identity of the external Python
writer.
- fn _with_commit_user(
+ /// Set the commit identity shared by this builder's writers and
committers.
+ fn with_commit_user(
mut slf: PyRefMut<'_, Self>,
commit_user: String,
) -> PyResult<PyRefMut<'_, Self>> {
@@ -412,7 +412,28 @@ impl PyPostponeFixedBucketWriteBuilder {
#[pymethods]
impl PyPostponeFixedBucketWriteBuilder {
- fn _with_commit_user(
+ /// Existing partitions and their bucket counts, using writer metadata
access.
+ fn resolve_bucket_plan(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
+ let table = &self.batch.context.table;
+ let mut builder = table
+ .new_postpone_fixed_bucket_write_builder()
+ .map_err(to_py_err)?;
+ if let Some(plan) = &self.plan {
+ builder = builder.with_bucket_plan(plan.clone());
+ }
+ if self.batch.static_partition.is_some() {
+ builder = builder.with_overwrite();
+ }
+ let plan = py
+ .detach(|| runtime().block_on(builder.resolve_bucket_plan()))
+ .map_err(to_py_err)?;
+ plan.to_arrow(table)
+ .map_err(to_py_err)?
+ .to_pyarrow(py)
+ .map(Bound::unbind)
+ }
+
+ fn with_commit_user(
mut slf: PyRefMut<'_, Self>,
commit_user: String,
) -> PyResult<PyRefMut<'_, Self>> {
@@ -994,14 +1015,14 @@ pub struct PyTableUpdateByRowId {
#[pymethods]
impl PyTableUpdateByRowId {
- fn _with_blob_uri_reader_factory(
- &mut self,
- py: Python<'_>,
+ #[pyo3(signature = (factory))]
+ fn with_blob_uri_reader_factory(
+ mut slf: PyRefMut<'_, Self>,
factory: Option<Py<PyAny>>,
- ) -> PyResult<()> {
- let factory = optional_uri_reader_factory(factory, py)?;
- self.inner.with_blob_uri_reader_factory(factory);
- Ok(())
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let factory = optional_uri_reader_factory(factory, slf.py())?;
+ slf.inner.with_blob_uri_reader_factory(factory);
+ Ok(slf)
}
fn update_columns(
@@ -1043,12 +1064,14 @@ pub struct PyStreamTableUpdate {
#[pymethods]
impl PyStreamTableUpdate {
- fn _with_blob_uri_reader_factory(
- &mut self,
- py: Python<'_>,
+ #[pyo3(signature = (factory))]
+ fn with_blob_uri_reader_factory(
+ mut slf: PyRefMut<'_, Self>,
factory: Option<Py<PyAny>>,
- ) -> PyResult<()> {
- self.context.with_blob_uri_reader_factory(factory, py)
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.context.with_blob_uri_reader_factory(factory, py)?;
+ Ok(slf)
}
#[pyo3(signature = (source, *, on, when_matched, when_not_matched,
commit_identifier))]
@@ -1170,12 +1193,14 @@ impl PyBatchTableUpdate {
self.context.delete_by_predicate(py, predicate)
}
- fn _with_blob_uri_reader_factory(
- &mut self,
- py: Python<'_>,
+ #[pyo3(signature = (factory))]
+ fn with_blob_uri_reader_factory(
+ mut slf: PyRefMut<'_, Self>,
factory: Option<Py<PyAny>>,
- ) -> PyResult<()> {
- self.context.with_blob_uri_reader_factory(factory, py)
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let py = slf.py();
+ slf.context.with_blob_uri_reader_factory(factory, py)?;
+ Ok(slf)
}
#[pyo3(signature = (source, *, on, when_matched, when_not_matched))]
@@ -1759,6 +1784,152 @@ mod tests {
use paimon::spec::{BigIntType, IntType, Schema, TableSchema, VarCharType};
use pyo3::ffi::c_str;
+ #[test]
+ fn public_uri_factory_and_commit_user_setters_are_chainable() {
+ Python::attach(|py| {
+ let io = FileIOBuilder::new("memory").build().unwrap();
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ io.clone(),
+ Identifier::new("default", "public_uri_setters"),
+ "memory:/public_uri_setters".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let fixed_schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .primary_key(["id"])
+ .option("bucket", "-2")
+ .build()
+ .unwrap();
+ let fixed = Arc::new(Table::new(
+ io,
+ Identifier::new("default", "public_commit_user"),
+ "memory:/public_commit_user".into(),
+ TableSchema::new(0, &fixed_schema),
+ None,
+ ));
+ let batch = Py::new(py,
PyBatchWriteBuilder::new(table.clone())).unwrap();
+ let stream = Py::new(py,
PyStreamWriteBuilder::new(table)).unwrap();
+ let fixed =
+ Py::new(py,
PyPostponeFixedBucketWriteBuilder::new(fixed).unwrap()).unwrap();
+ let locals = PyDict::new(py);
+ locals.set_item("batch", batch).unwrap();
+ locals.set_item("stream", stream).unwrap();
+ locals.set_item("fixed", fixed).unwrap();
+ py.run(
+ c_str!(
+ r#"
+import io
+
+class Reader:
+ def new_input_stream(self, uri):
+ return io.BytesIO(uri.encode())
+
+class Factory:
+ def create(self, uri):
+ return Reader()
+
+factory = Factory()
+for builder in [batch, stream, fixed]:
+ assert builder.with_commit_user('public-job') is builder
+ assert not hasattr(builder, '_with_commit_user')
+ for invalid in ['', '..', 'name/part']:
+ try:
+ builder.with_commit_user(invalid)
+ except ValueError:
+ pass
+ else:
+ raise AssertionError('invalid commit user accepted')
+
+for builder in [batch, stream]:
+ writer = builder.new_write()
+ assert writer.with_blob_uri_reader_factory(factory) is writer
+ assert writer.with_blob_uri_reader_factory(None) is writer
+ writer.close()
+ update = builder.new_update()
+ row_id = (update.new_update_by_row_id() if builder is batch else
+ update.new_update_by_row_id(7))
+ for operation in [update, row_id]:
+ assert operation.with_blob_uri_reader_factory(factory) is operation
+ try:
+ operation.with_blob_uri_reader_factory(type('Invalid', (),
{'create': None})())
+ except TypeError:
+ pass
+ else:
+ raise AssertionError('factory without callable create accepted')
+ assert operation.with_blob_uri_reader_factory(None) is operation
+ assert not hasattr(operation, '_with_blob_uri_reader_factory')
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
+ #[test]
+ fn
postpone_resolve_bucket_plan_binding_round_trips_arrow_and_overwrite_precedence()
{
+ Python::attach(|py| {
+ let schema = Schema::builder()
+ .column("pt", DataType::VarChar(VarCharType::string_type()))
+ .column("id", DataType::Int(IntType::new()))
+ .partition_keys(["pt"])
+ .primary_key(["pt", "id"])
+ .option("bucket", "-2")
+ .option("postpone.default-bucket-num", "7")
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "bucket_plan_binding"),
+ "memory:/bucket_plan_binding".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let builder = Py::new(
+ py,
+ PyPostponeFixedBucketWriteBuilder::new(table.clone()).unwrap(),
+ )
+ .unwrap();
+ let overwrite =
+ Py::new(py,
PyPostponeFixedBucketWriteBuilder::new(table).unwrap()).unwrap();
+ let locals = PyDict::new(py);
+ locals.set_item("builder", builder).unwrap();
+ locals.set_item("overwrite", overwrite).unwrap();
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+empty = builder.resolve_bucket_plan()
+assert not hasattr(builder, 'bucket_plan')
+assert isinstance(empty, pa.RecordBatch)
+assert empty.schema.names == ['pt', 'total_buckets']
+assert empty.num_rows == 0
+assert empty.schema.field('total_buckets') == pa.field('total_buckets',
pa.int32(), nullable=False)
+plan = pa.RecordBatch.from_arrays([pa.array(['b', 'a']), pa.array([3, 2],
type=pa.int32())],
+ names=['pt', 'total_buckets'])
+builder.with_bucket_plan(plan)
+assert sorted(zip(*builder.resolve_bucket_plan().to_pydict().values())) ==
[('a', 2), ('b', 3)]
+builder.with_overwrite()
+assert builder.resolve_bucket_plan().num_rows == 2
+overwrite.with_overwrite()
+assert overwrite.resolve_bucket_plan().num_rows == 0
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
#[test]
fn replaced_blob_consumers_are_visible_to_writer_cycle_collection() {
Python::attach(|py| {
diff --git a/bindings/python/tests/test_blob_uri_reader.py
b/bindings/python/tests/test_blob_uri_reader.py
index 6c0fbaac..54bbc54a 100644
--- a/bindings/python/tests/test_blob_uri_reader.py
+++ b/bindings/python/tests/test_blob_uri_reader.py
@@ -210,7 +210,7 @@ def
test_update_factory_errors_keep_identity_and_existing_data(tmp_path, stream,
builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
update = builder.new_update()
- update._with_blob_uri_reader_factory(Factory())
+ assert update.with_blob_uri_reader_factory(Factory()) is update
reference = _reference(source_file.as_uri())
source = pa.Table.from_batches([_batch([0], [reference])])
data = pa.table({'_ROW_ID': [row_id], 'payload': pa.array([reference],
type=pa.large_binary())})
@@ -252,7 +252,7 @@ def
test_predicate_updates_share_blob_factory_across_matched_file_groups(tmp_pat
factory = ReaderFactory()
builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
update = builder.new_update()
- update._with_blob_uri_reader_factory(factory)
+ assert update.with_blob_uri_reader_factory(factory) is update
reference = descriptor(2, -1)
read_columns = None
if assignment == 'scalar':
diff --git a/bindings/python/tests/test_table_commit.py
b/bindings/python/tests/test_table_commit.py
index e6ee4bdd..1fa5ff4a 100644
--- a/bindings/python/tests/test_table_commit.py
+++ b/bindings/python/tests/test_table_commit.py
@@ -85,7 +85,7 @@ def test_public_api_separates_batch_and_stream(tmp_path):
factory(**kwargs)
batch = table.new_batch_write_builder()
stream = table.new_stream_write_builder()
- assert not hasattr(batch, "with_commit_user")
+ assert batch.with_commit_user("batch-job") is batch
assert not hasattr(stream, "with_overwrite")
for obj in (batch, stream, batch.new_commit(), stream.new_commit()):
assert not hasattr(obj, "deserialize_commit_message")
@@ -208,11 +208,11 @@ def test_empty_stream_snapshot_properties(tmp_path):
@pytest.mark.parametrize("overwrite", [False, True])
@pytest.mark.parametrize("serialized", [False, True])
-def test_batch_bridge_preserves_external_commit_user(tmp_path, overwrite,
serialized):
+def test_public_batch_commit_user_survives_serialization(tmp_path, overwrite,
serialized):
table = _table(tmp_path, primary_key=True, options={"bucket": "1"})
_append(table, [1], [10])
builder = table.new_batch_write_builder()
- assert builder._with_commit_user("python-batch-job") is builder
+ assert builder.with_commit_user("python-batch-job") is builder
if overwrite:
builder.with_overwrite()
messages = _prepare(builder, [2], [20])
@@ -226,11 +226,11 @@ def
test_batch_bridge_preserves_external_commit_user(tmp_path, overwrite, serial
@pytest.mark.parametrize("user", ["", "../job", "a/b"])
-def test_invalid_batch_bridge_commit_user_preserves_identity(tmp_path, user):
+def test_invalid_public_batch_commit_user_preserves_identity(tmp_path, user):
table = _table(tmp_path)
- builder =
table.new_batch_write_builder()._with_commit_user("python-batch-job")
+ builder =
table.new_batch_write_builder().with_commit_user("python-batch-job")
with pytest.raises(ValueError):
- builder._with_commit_user(user)
+ builder.with_commit_user(user)
builder.new_commit().commit(_prepare(builder, [1], [10]))
assert _snapshot(table)["commitUser"] == "python-batch-job"
assert _rows(table) == [1]
diff --git a/bindings/python/tests/test_uri_reader_protocols.py
b/bindings/python/tests/test_uri_reader_protocols.py
new file mode 100644
index 00000000..d1039e74
--- /dev/null
+++ b/bindings/python/tests/test_uri_reader_protocols.py
@@ -0,0 +1,61 @@
+# 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.
+
+"""The package bootstrap exports stream protocols without an SDK rebuild."""
+
+from pathlib import Path
+import subprocess
+import sys
+
+
+def test_public_uri_protocols_are_runtime_imports():
+ source = Path(__file__).resolve().parents[1] / 'python'
+ code = r"""
+import io
+import sys
+import types
+
+# Model the dynamic submodule created by the extension. Load the real Python
+# package bootstrap and protocol source without loading an installed extension.
+engine = types.ModuleType('pypaimon_rust.pypaimon_rust')
+engine.datafusion = types.ModuleType('pypaimon_rust.datafusion')
+engine.__all__ = ['datafusion']
+sys.modules[engine.__name__] = engine
+sys.modules[engine.datafusion.__name__] = engine.datafusion
+sys.path.insert(0, sys.argv[1])
+
+from pypaimon_rust import UriReader as TopReader, UriReaderFactory as
TopFactory
+from pypaimon_rust.datafusion import UriReader, UriReaderFactory
+from pypaimon_rust.uri_reader import UriReader as SourceReader
+assert UriReader is TopReader is SourceReader
+assert UriReaderFactory is TopFactory
+
+class Reader(UriReader):
+ def new_input_stream(self, uri):
+ return io.BytesIO(uri.encode())
+
+class Factory(UriReaderFactory):
+ def create(self, uri) -> UriReader:
+ return Reader()
+
+with Factory().create('custom://source').new_input_stream('custom://source')
as stream:
+ assert stream.read() == b'custom://source'
+
+assert not hasattr(engine.datafusion, '_UriReader')
+assert not hasattr(engine.datafusion, '_UriReaderFactory')
+"""
+ subprocess.run([sys.executable, '-c', code, str(source)], check=True)
diff --git a/crates/paimon/src/table/authorized_write_tests.rs
b/crates/paimon/src/table/authorized_write_tests.rs
index 59d5a0df..c78a7b15 100644
--- a/crates/paimon/src/table/authorized_write_tests.rs
+++ b/crates/paimon/src/table/authorized_write_tests.rs
@@ -296,6 +296,59 @@ impl CommitCallback for Callback {
}
}
+#[tokio::test]
+async fn
rest_postpone_resolve_bucket_plan_requires_write_identity_without_select_permission()
{
+ let fixture = RestTable::new(
+ true,
+ &[("bucket", "-2"), ("postpone.default-bucket-num", "2")],
+ )
+ .await;
+ fixture
+ .append("fixed", &[(1, "a", 10), (2, "b", 20)], true)
+ .await;
+ let table = fixture.table.copy_with_pinned_snapshot(None);
+ // The REST response already carries the selected snapshot. The filesystem
+ // copy is not required to inspect its manifests through a writer builder.
+ table
+ .file_io()
+ .delete_file(&table.snapshot_manager().snapshot_path(1))
+ .await
+ .unwrap();
+ let exported = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&table)
+ .unwrap();
+ assert_eq!(exported.num_rows(), 2);
+ assert_eq!(exported.schema().field(0).name(), "p");
+ assert_eq!(
+ exported
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values(),
+ &[2, 2]
+ );
+ assert_eq!(fixture.state.query_requests.load(Ordering::SeqCst), 0);
+ let mut schema = serde_json::to_value(table.schema()).unwrap();
+ schema["fields"][2]["name"] = serde_json::json!("another_value");
+ let unbound = table
+ .copy_with_resolved_schema(serde_json::from_value(schema).unwrap(),
"main")
+ .unwrap();
+ assert!(unbound.query_auth_session().is_none());
+ assert!(unbound
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .resolve_bucket_plan()
+ .await
+ .is_err());
+ assert_eq!(fixture.state.query_requests.load(Ordering::SeqCst), 0);
+}
+
#[tokio::test]
async fn rest_stream_recovery_and_callbacks_do_not_require_select_permission()
{
let fixture = RestTable::new(false, &[]).await;
diff --git a/crates/paimon/src/table/postpone_bucket_plan.rs
b/crates/paimon/src/table/postpone_bucket_plan.rs
index 83bba51f..ff089650 100644
--- a/crates/paimon/src/table/postpone_bucket_plan.rs
+++ b/crates/paimon/src/table/postpone_bucket_plan.rs
@@ -15,10 +15,11 @@
// specific language governing permissions and limitations
// under the License.
-use crate::spec::batch_to_serialized_bytes;
+use crate::spec::{batch_to_serialized_bytes, BinaryRow, Snapshot};
use crate::table::{Table, TableScan};
use crate::Result;
-use arrow_array::{Array, Int32Array, RecordBatch};
+use arrow_array::{new_empty_array, Array, ArrayRef, Int32Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema};
use std::collections::HashMap;
use std::sync::Arc;
@@ -124,6 +125,60 @@ impl PostponeBucketPlan {
})
}
+ /// Export known partitions using the same schema accepted by `from_arrow`.
+ /// A configured default is not an existing partition and is not exported.
+ pub fn to_arrow(&self, table: &Table) -> Result<RecordBatch> {
+ let partition_fields = table.schema().partition_fields();
+ let mut fields =
crate::arrow::build_target_arrow_schema(&partition_fields)?
+ .fields()
+ .to_vec();
+ fields.push(Arc::new(Field::new(
+ POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD,
+ DataType::Int32,
+ false,
+ )));
+ let mut entries = self.bucket_counts.iter().collect::<Vec<_>>();
+ entries.sort_unstable_by_key(|(partition, _)| *partition);
+ let partitions = entries
+ .iter()
+ .map(|(bytes, _)| {
+ let row = BinaryRow::from_serialized_bytes(bytes)?;
+ if row.arity() as usize != partition_fields.len() {
+ return Err(data_invalid(
+ "Postpone bucket plan partition arity does not match
the table schema",
+ ));
+ }
+ Ok(row)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let mut columns = Vec::<ArrayRef>::with_capacity(fields.len());
+ for (index, field) in partition_fields.iter().enumerate() {
+ let arrays = partitions
+ .iter()
+ .map(|partition| {
+ crate::arrow::partition::partition_array(partition, index,
field.data_type(), 1)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let column = if arrays.is_empty() {
+ new_empty_array(fields[index].data_type())
+ } else {
+ arrow_select::concat::concat(
+ &arrays
+ .iter()
+ .map(|array| array.as_ref())
+ .collect::<Vec<_>>(),
+ )
+ .map_err(|error| data_invalid(error.to_string()))?
+ };
+ columns.push(column);
+ }
+ columns.push(Arc::new(Int32Array::from_iter_values(
+ entries.iter().map(|(_, count)| **count),
+ )));
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+ .map_err(|error| data_invalid(error.to_string()))
+ }
+
pub(crate) fn with_default_bucket_num(default_bucket_num: i32) -> Self {
Self {
default_bucket_num: Some(default_bucket_num),
@@ -142,9 +197,18 @@ impl PostponeBucketPlan {
return Ok(());
}
let snapshot =
table.snapshot_manager().get_snapshot(snapshot_id).await?;
+ self.load_existing_buckets_from_snapshot(table, &snapshot)
+ .await
+ }
+
+ pub(super) async fn load_existing_buckets_from_snapshot(
+ &mut self,
+ table: &Table,
+ snapshot: &Snapshot,
+ ) -> Result<()> {
let entries = TableScan::new(table, None, vec![], None, None, None)
.with_scan_all_files()
- .plan_manifest_entries(&snapshot)
+ .plan_manifest_entries(snapshot)
.await?;
let mut counts = HashMap::new();
for entry in entries {
@@ -174,3 +238,69 @@ impl PostponeBucketPlan {
.or(self.default_bucket_num)
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::table::table_write::tests::test_file_io;
+ use arrow_array::StringArray;
+
+ #[test]
+ fn arrow_plan_round_trip_keeps_null_partition_and_excludes_default_rows() {
+ let schema = crate::spec::Schema::builder()
+ .column(
+ "pt",
+
crate::spec::DataType::VarChar(crate::spec::VarCharType::string_type()),
+ )
+ .column(
+ "id",
+ crate::spec::DataType::Int(crate::spec::IntType::new()),
+ )
+ .primary_key(["pt", "id"])
+ .partition_keys(["pt"])
+ .option("bucket", "-2")
+ .option("primary-key.nullable", "true")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ test_file_io(),
+ crate::catalog::Identifier::new("default", "bucket_plan_arrow"),
+ "memory:/bucket_plan_arrow".into(),
+ crate::spec::TableSchema::new(0, &schema),
+ None,
+ );
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("pt", DataType::Utf8, true),
+ Field::new(
+ POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD,
+ DataType::Int32,
+ false,
+ ),
+ ])),
+ vec![
+ Arc::new(StringArray::from(vec![Some("b"), None, Some("a")])),
+ Arc::new(Int32Array::from(vec![2, 3, 4])),
+ ],
+ )
+ .unwrap();
+ let plan = PostponeBucketPlan::from_arrow(&table, &batch).unwrap();
+ let exported = plan.to_arrow(&table).unwrap();
+ assert_eq!(exported.schema().field(0).name(), "pt");
+ assert_eq!(exported.schema().field(0).data_type(), &DataType::Utf8);
+ assert_eq!(exported.schema().field(1), batch.schema().field(1));
+ assert_eq!(exported.num_rows(), 3);
+ assert_eq!(
+ PostponeBucketPlan::from_arrow(&table, &exported)
+ .unwrap()
+ .bucket_counts,
+ plan.bucket_counts
+ );
+ let expected_schema = exported.schema();
+ let default = PostponeBucketPlan::with_default_bucket_num(7);
+ let exported = default.to_arrow(&table).unwrap();
+ assert_eq!(exported.schema(), expected_schema);
+ assert_eq!(exported.num_rows(), 0);
+ assert_eq!(default.total_buckets(&[]), Some(7));
+ }
+}
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
index 661f3207..21ebc714 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
@@ -18,7 +18,9 @@
use super::postpone_fixed_bucket_router::validate_postpone_fixed_bucket_table;
use crate::resource::ResourceContext;
use crate::spec::CoreOptions;
-use crate::table::write_builder::{ensure_table_write_allowed,
validate_commit_user};
+use crate::table::write_builder::{
+ ensure_table_write_allowed, ensure_table_write_authorized,
validate_commit_user,
+};
use crate::table::{
PostponeBucketPlan, PostponeFixedBucketTableCommit,
PostponeFixedBucketTableWrite, Table,
};
@@ -71,6 +73,29 @@ impl<'a> PostponeFixedBucketWriteBuilder<'a> {
self
}
+ /// Resolve routing metadata before distributed input is materialized.
+ ///
+ /// Existing partitions are read from the latest write baseline,
independently
+ /// of historical read selectors. This uses writer metadata access and does
+ /// not require a SELECT grant. Overwrite starts with no existing layout;
+ /// an explicitly supplied plan takes precedence over that default.
+ /// Bucket counts for new input partitions are not inferred by this method.
+ pub async fn resolve_bucket_plan(&self) -> Result<PostponeBucketPlan> {
+ ensure_table_write_authorized(self.table)?;
+ if let Some(plan) = &self.bucket_plan {
+ return Ok(plan.clone());
+ }
+ if self.overwrite {
+ return Ok(PostponeBucketPlan::default());
+ }
+ let mut plan = PostponeBucketPlan::default();
+ if let Some(snapshot) =
self.table.snapshot_manager().get_latest_snapshot().await? {
+ plan.load_existing_buckets_from_snapshot(self.table, &snapshot)
+ .await?;
+ }
+ Ok(plan)
+ }
+
pub fn new_commit(&self) -> PostponeFixedBucketTableCommit {
PostponeFixedBucketTableCommit::new(self.table,
self.commit_user.clone(), self.overwrite)
}
@@ -123,6 +148,196 @@ mod tests {
PostponeBucketPlan::from_arrow(table, &batch).unwrap()
}
+ #[tokio::test]
+ async fn
resolve_bucket_plan_obeys_explicit_and_overwrite_precedence_without_default_rows()
{
+ let io = test_file_io();
+ let path = "memory:/builder_bucket_plan_modes";
+ setup_dirs(&io, path).await;
+ let table =
+ test_postpone_pk_table(&io,
path).copy_with_options(std::collections::HashMap::from([
+ ("postpone.default-bucket-num".into(), "7".into()),
+ ]));
+ let builder = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ assert_eq!(
+ builder
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&table)
+ .unwrap()
+ .num_rows(),
+ 0
+ );
+ let explicit = bucket_plan(&table, 3);
+ let builder = builder.with_overwrite().with_bucket_plan(explicit);
+ let exported = builder
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&table)
+ .unwrap();
+ assert_eq!(exported.num_rows(), 1);
+ assert_eq!(
+ exported
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 3
+ );
+ assert_eq!(
+ table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .with_overwrite()
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&table)
+ .unwrap()
+ .num_rows(),
+ 0
+ );
+ }
+
+ #[tokio::test]
+ async fn
resolve_bucket_plan_uses_latest_write_baseline_and_ignores_pending_files() {
+ let io = test_file_io();
+ let path = "memory:/builder_bucket_plan_latest";
+ setup_dirs(&io, path).await;
+ let table = test_postpone_pk_table(&io, path);
+ let pending = table.new_write_builder();
+ let mut writer = pending.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1], vec![10]))
+ .await
+ .unwrap();
+ pending
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ assert_eq!(
+ table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&table)
+ .unwrap()
+ .num_rows(),
+ 0
+ );
+ let fixed = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .with_bucket_plan(bucket_plan(&table, 3));
+ let mut writer = fixed.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![2], vec![20]))
+ .await
+ .unwrap();
+ fixed
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let historical =
table.copy_with_options(std::collections::HashMap::from([(
+ "scan.snapshot-id".into(),
+ "1".into(),
+ )]));
+ let exported = historical
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&historical)
+ .unwrap();
+ assert_eq!(exported.num_rows(), 1);
+ assert_eq!(
+ exported
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 3
+ );
+ assert_eq!(
+ historical
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .with_overwrite()
+ .resolve_bucket_plan()
+ .await
+ .unwrap()
+ .to_arrow(&historical)
+ .unwrap()
+ .num_rows(),
+ 0
+ );
+ }
+
+ #[tokio::test]
+ async fn resolve_bucket_plan_rejects_conflicting_entries_in_one_bucket() {
+ use crate::spec::{Manifest, ManifestEntry, ManifestList};
+ let io = test_file_io();
+ let path = "memory:/builder_bucket_plan_conflict";
+ setup_dirs(&io, path).await;
+ let table = test_postpone_pk_table(&io, path);
+ let fixed = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .with_bucket_plan(bucket_plan(&table, 1));
+ let mut writer = fixed.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1], vec![10]))
+ .await
+ .unwrap();
+ fixed
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let manager = table.snapshot_manager();
+ let snapshot = manager.get_snapshot(1).await.unwrap();
+ let list_path = manager.manifest_path(snapshot.delta_manifest_list());
+ let mut metas = ManifestList::read(&io, &list_path).await.unwrap();
+ let manifest_path = manager.manifest_path(metas[0].file_name());
+ let mut entries = Manifest::read(&io, &manifest_path).await.unwrap();
+ let original = &entries[0];
+ let mut file = original.file().clone();
+ file.file_name = "different-bucket-count.parquet".into();
+ entries.push(ManifestEntry::new(
+ *original.kind(),
+ original.partition().to_vec(),
+ original.bucket(),
+ 2,
+ file,
+ 2,
+ ));
+ Manifest::write(&io, &manifest_path, &entries)
+ .await
+ .unwrap();
+ let mut meta = serde_json::to_value(&metas[0]).unwrap();
+ meta["_NUM_ADDED_FILES"] = serde_json::json!(entries.len());
+ meta["_TOTAL_BUCKETS"] = serde_json::Value::Null;
+ metas[0] = serde_json::from_value(meta).unwrap();
+ ManifestList::write(&io, &list_path, &metas).await.unwrap();
+ let error = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .resolve_bucket_plan()
+ .await
+ .unwrap_err();
+ assert!(error
+ .to_string()
+ .contains("conflicting total bucket counts"));
+ }
+
#[tokio::test]
async fn test_postpone_fixed_bucket_builder_modes() {
let file_io = test_file_io();
diff --git a/crates/paimon/src/table/table_upsert.rs
b/crates/paimon/src/table/table_upsert.rs
index e18521ea..c78abe56 100644
--- a/crates/paimon/src/table/table_upsert.rs
+++ b/crates/paimon/src/table/table_upsert.rs
@@ -34,8 +34,7 @@ use super::update_input::unique_column_index;
use super::upsert_key_matcher::UpsertKeyMatcher;
use super::write_batch_normalize::{normalize_blob_write_field,
normalize_write_array};
use crate::spec::{
- batch_to_serialized_bytes, extract_datum_from_arrow, BinaryRow,
CoreOptions, Predicate,
- PredicateBuilder,
+ batch_to_serialized_bytes, extract_datum_from_arrow, CoreOptions,
Predicate, PredicateBuilder,
};
use crate::table::source::{DataSplit, Plan};
use crate::table::{CommitMessage, Table};
@@ -562,8 +561,8 @@ mod tests {
use crate::catalog::Identifier;
use crate::io::FileIOBuilder;
use crate::spec::{
- DataEvolutionIndexSourceMeta, DataType as PaimonType, IndexManifest,
TimestampType,
- VarCharType,
+ BinaryRow, DataEvolutionIndexSourceMeta, DataType as PaimonType,
IndexManifest,
+ TimestampType, VarCharType,
};
use crate::spec::{IntType, Schema as PaimonSchema, TableSchema};
use crate::table::TableCommit;
