diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md index 44633f866..86c408608 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 b36002b71..9838f8d29 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 ad573e03c..427e98104 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/uri_reader.py b/bindings/python/python/pypaimon_rust/uri_reader.py new file mode 100644 index 000000000..3beb5d7e3 --- /dev/null +++ b/bindings/python/python/pypaimon_rust/uri_reader.py @@ -0,0 +1,34 @@ +# 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. + +"""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 1ac21b058..299730b0f 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> { @@ -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> { + 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> { @@ -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>, - ) -> PyResult<()> { - let factory = optional_uri_reader_factory(factory, py)?; - self.inner.with_blob_uri_reader_factory(factory); - Ok(()) + ) -> PyResult> { + 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>, - ) -> PyResult<()> { - self.context.with_blob_uri_reader_factory(factory, py) + ) -> PyResult> { + 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>, - ) -> PyResult<()> { - self.context.with_blob_uri_reader_factory(factory, py) + ) -> PyResult> { + 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 6c0fbaacd..54bbc54a2 100644 --- a/bindings/python/tests/test_blob_uri_reader.py +++ b/bindings/python/tests/test_blob_uri_reader.py @@ -210,7 +210,7 @@ def create(self, uri): 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 e6ee4bdda..1fa5ff4a8 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 000000000..d1039e749 --- /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 59d5a0df1..c78a7b15d 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::() + .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 83bba51fc..ff0896507 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 { + 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::>(); + 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::>>()?; + let mut columns = Vec::::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::>>()?; + 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::>(), + ) + .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 661f32075..21ebc7141 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 { + 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::() + .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::() + .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 e18521ea2..c78abe560 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;