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 b71af884 Complete native FUSE writes, postpone planning and MERGE
callbacks (#1165)
b71af884 is described below
commit b71af88414ba0020a9d0d0557c51eb2e43da8cae
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon Oct 19 08:56:04 2026 +0800
Complete native FUSE writes, postpone planning and MERGE callbacks (#1165)
---
bindings/c/src/tests.rs | 6 +-
bindings/c/src/write.rs | 6 +-
bindings/go/postpone_fixed_bucket_write.go | 6 +-
bindings/go/tests/paimon_test.go | 6 +-
bindings/python/README.md | 15 +-
.../python/python/pypaimon_rust/datafusion.pyi | 26 +-
bindings/python/src/merge_into.rs | 7 +-
bindings/python/src/write.rs | 25 +-
bindings/python/tests/test_table_commit.py | 23 ++
bindings/python/tests/test_update.py | 47 +++
crates/paimon/src/io/file_io/fuse.rs | 99 +++++
crates/paimon/src/spec/core_options.rs | 43 ++
crates/paimon/src/table/branch_write_tests.rs | 117 ++++++
crates/paimon/src/table/mod.rs | 38 +-
crates/paimon/src/table/partition_spec.rs | 284 +++++++++++++
crates/paimon/src/table/postpone_bucket_plan.rs | 6 +
crates/paimon/src/table/postpone_bucket_planner.rs | 437 +++++++++++++++++++++
.../src/table/postpone_fixed_bucket_router.rs | 4 +
.../src/table/postpone_fixed_bucket_write.rs | 57 ++-
.../table/postpone_fixed_bucket_write_builder.rs | 47 ++-
20 files changed, 1239 insertions(+), 60 deletions(-)
diff --git a/bindings/c/src/tests.rs b/bindings/c/src/tests.rs
index 8a8d535f..f445de6c 100644
--- a/bindings/c/src/tests.rs
+++ b/bindings/c/src/tests.rs
@@ -2528,9 +2528,9 @@ fn test_distributed_postpone_writers_share_bucket_plan() {
assert!(fixed.error.is_null());
assert!(!fixed.write_builder.is_null());
let write =
paimon_postpone_fixed_bucket_write_builder_new_write(fixed.write_builder);
- assert!(write.write.is_null());
- assert!(error_message(write.error).contains("bucket plan is
required"));
- paimon_error_free(write.error);
+ assert!(write.error.is_null());
+ assert!(!write.write.is_null());
+ paimon_postpone_fixed_bucket_table_write_free(write.write);
paimon_postpone_fixed_bucket_write_builder_free(fixed.write_builder);
let wb1 =
paimon_table_new_postpone_fixed_bucket_write_builder_with_commit_user(
diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs
index 86f86e18..3282e0d5 100644
--- a/bindings/c/src/write.rs
+++ b/bindings/c/src/write.rs
@@ -167,7 +167,8 @@ pub unsafe extern "C" fn paimon_table_new_write_builder(
}
/// Create a one-shot fixed-bucket WriteBuilder for a postpone table.
-/// A bucket plan must be set before creating a writer.
+/// An explicit bucket plan is optional. Without one, core reuses existing
+/// layouts and uses the configured default or infers new partition counts.
///
/// # Safety
/// `table` must be a valid table pointer, or null (returns error).
@@ -195,7 +196,8 @@ pub unsafe extern "C" fn
paimon_table_new_write_builder_with_commit_user(
}
/// Create a fixed-bucket WriteBuilder with a stable commit identity.
-/// A bucket plan must be set before creating a writer.
+/// Share an explicit bucket plan across distributed writers. Without one, core
+/// reuses existing layouts and infers new counts independently per writer.
///
/// # Safety
/// `table` must be a valid table pointer. `commit_user` must be a valid UTF-8
diff --git a/bindings/go/postpone_fixed_bucket_write.go
b/bindings/go/postpone_fixed_bucket_write.go
index d3445c16..dc279286 100644
--- a/bindings/go/postpone_fixed_bucket_write.go
+++ b/bindings/go/postpone_fixed_bucket_write.go
@@ -29,7 +29,9 @@ import (
)
// PostponeFixedBucketWriteBuilder creates fixed-bucket writers for bucket=-2
-// tables. A resolved bucket plan must be supplied before NewWrite.
+// tables. Without an explicit plan, core reuses known layouts and applies the
+// configured default or infers counts for new partitions across this writer's
+// input batches. Distributed writers should share one resolved bucket plan.
type PostponeFixedBucketWriteBuilder struct {
ctx context.Context
lib *libRef
@@ -151,7 +153,7 @@ func (wb *PostponeFixedBucketWriteBuilder) NewCommit()
(*PostponeFixedBucketTabl
return &PostponeFixedBucketTableCommit{ctx: wb.ctx, lib: wb.lib, inner:
inner}, nil
}
-// PostponeFixedBucketTableWrite writes rows according to a resolved bucket
plan.
+// PostponeFixedBucketTableWrite writes rows using explicit or inferred bucket
counts.
type PostponeFixedBucketTableWrite struct {
ctx context.Context
lib *libRef
diff --git a/bindings/go/tests/paimon_test.go b/bindings/go/tests/paimon_test.go
index e7b0c943..8a359b0c 100644
--- a/bindings/go/tests/paimon_test.go
+++ b/bindings/go/tests/paimon_test.go
@@ -688,9 +688,11 @@ func TestMultiplePostponeFixedBucketWritersSharePlan(t
*testing.T) {
defer builder.Close()
}
- if _, err := builders[0].NewWrite(); err == nil ||
!strings.Contains(err.Error(), "bucket plan is required") {
- t.Fatalf("Expected missing bucket plan error, got: %v", err)
+ automatic, err := builders[0].NewWrite()
+ if err != nil {
+ t.Fatalf("Failed to create automatically planned writer: %v",
err)
}
+ automatic.Close()
plan := makePartitionedBucketPlan(t, []string{"2026-08-14",
"2026-08-15"}, 1)
for index, builder := range builders {
if err := builder.WithBucketPlan(plan); err != nil {
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 5ce313ee..adebdc1e 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -294,11 +294,16 @@ for a partition and bucket must have a single writer
owner. Both duplicate
owners in a commit and stale overlapping writers are rejected. Pending files
(`bucket=-2`) can coexist with real buckets; they remain pending until
assigned.
-Without an explicit plan, set `postpone.default-bucket-num` to a positive
integer.
-The count is used exactly for new partitions, and append reuses existing
layouts
-from the same snapshot used to initialize sequence numbers. Overwrite uses the
-configured default for replaced partitions. No staging, automatic rescaling or
-compaction is performed. Deletion-vector tables are rejected by this writer.
+Without an explicit plan, `postpone.default-bucket-num` can provide an exact
count
+for new partitions. Otherwise core buffers unknown partitions until preparation
+and infers their counts from this writer's input: the target row count takes
+precedence over estimated BinaryRow bytes. Counts round up to a power of two,
+then cap at the configured maximum parallelism. Append reuses existing layouts
+from the same snapshot used to initialize sequence numbers and writes known
+partitions directly. Overwrite plans replaced partitions from the new input.
+Distributed writers should share one resolved plan. This does not implement
+Spark's temporary-file measurement, automatic rescaling or compaction.
+Deletion-vector tables are rejected by this writer.
A failed write cannot be prepared, and close deletes only unprepared output.
Configure overwrite on the batch builder:
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 5842f1fb..badf0a44 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -462,8 +462,10 @@ class BatchTableUpdate:
A Table source reads its selected snapshot or numeric timestamp window
in core, independently of the target. Arrow input preserves batches
and source row order.
- Clauses carry assignments as (column, source/target/literal, value)
- tuples, a delete flag, and optional SQL condition metadata.
+ Clauses carry assignments as (column, source/target/literal/callable,
value)
+ tuples, a delete flag, and optional SQL condition metadata. Callables
+ consume the complete clause's qualified t./s. Arrow columns and return
+ an Array or ChunkedArray with the same row count; empty clauses are
lazy.
Commit returned messages through the matching builder's committer.
"""
...
@@ -520,8 +522,10 @@ class StreamTableUpdate:
A Table source reads its selected snapshot or numeric timestamp window
in core, independently of the target. Arrow input preserves batches
and source row order.
- Clauses carry assignments as (column, source/target/literal, value)
- tuples, a delete flag, and optional SQL condition metadata.
+ Clauses carry assignments as (column, source/target/literal/callable,
value)
+ tuples, a delete flag, and optional SQL condition metadata. Callables
+ consume the complete clause's qualified t./s. Arrow columns and return
+ an Array or ChunkedArray with the same row count; empty clauses are
lazy.
Commit returned messages through the matching builder's committer.
"""
...
@@ -717,7 +721,8 @@ class BatchWriteBuilder:
replace touched partitions, irrespective of the static spec; empty
input
removes nothing. With the option false, match the spec ({} means all).
An unpartitioned empty overwrite truncates the table. Partition values
- use the schema's Python types; None or the default partition name is
null.
+ use the schema's Python types or Java textual specs; None or the
default
+ partition name is null.
"""
...
def with_restore_snapshot(self, snapshot_id: int) -> "BatchWriteBuilder":
...
@@ -728,10 +733,15 @@ class BatchWriteBuilder:
class PostponeFixedBucketWriteBuilder:
"""Real-bucket batch writes for a bucket=-2 primary-key table.
- Supply a shared Arrow plan or configure postpone.default-bucket-num.
+ Supply a shared Arrow plan, configure postpone.default-bucket-num, or let
+ core infer new partition counts across all input batches at prepare_commit.
+ Inference uses target rows, otherwise estimated BinaryRow bytes, rounds up
+ to a power of two and caps at the configured maximum parallelism.
The default applies exactly to new partitions; append reuses existing
- layouts from the writer's baseline snapshot. Overwrite uses the new
default.
- Automatic staging, rescaling and compaction are not performed.
+ layouts from the writer's baseline snapshot. Overwrite plans replaced
+ partitions from the new input or configured default.
+ 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_bucket_plan(self, plan: pyarrow.RecordBatch) ->
"PostponeFixedBucketWriteBuilder":
diff --git a/bindings/python/src/merge_into.rs
b/bindings/python/src/merge_into.rs
index 17e8683b..4d237d2f 100644
--- a/bindings/python/src/merge_into.rs
+++ b/bindings/python/src/merge_into.rs
@@ -67,7 +67,12 @@ fn assignments(
let assignment = match kind.as_str() {
"source" =>
MergeAssignment::SourceColumn(value.bind(py).extract()?),
"target" =>
MergeAssignment::TargetColumn(value.bind(py).extract()?),
- "literal" => {
+ "literal" | "callable" => {
+ if kind == "callable" && !value.bind(py).is_callable() {
+ return Err(PyValueError::new_err(
+ "MERGE callable assignment must be callable",
+ ));
+ }
let values = PyDict::new(py);
values.set_item(&name, value)?;
let assignment =
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index f9a1716a..0ad668c5 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -103,29 +103,14 @@ fn boolean_option(table: &Table, key: &str, default:
bool) -> PyResult<bool> {
}
}
-/// Java partition specs may encode numeric values as strings. Other Python
+/// Java partition specs encode schema-typed values as strings. Other Python
/// literals use the same schema-driven conversion as the predicate API.
-fn partition_value(value: &Bound<'_, PyAny>, data_type: &DataType) ->
PyResult<Datum> {
+fn partition_value(value: &Bound<'_, PyAny>, data_type: &DataType) ->
PyResult<Option<Datum>> {
if let Ok(text) = value.cast::<PyString>() {
let text = text.to_str()?;
- let invalid = || {
- PyValueError::new_err(format!(
- "Invalid partition value '{text}' for {data_type:?}"
- ))
- };
- match data_type {
- DataType::TinyInt(_) => return
text.parse().map(Datum::TinyInt).map_err(|_| invalid()),
- DataType::SmallInt(_) => {
- return text.parse().map(Datum::SmallInt).map_err(|_| invalid())
- }
- DataType::Int(_) => return
text.parse().map(Datum::Int).map_err(|_| invalid()),
- DataType::BigInt(_) => return
text.parse().map(Datum::Long).map_err(|_| invalid()),
- DataType::Float(_) => return
text.parse().map(Datum::Float).map_err(|_| invalid()),
- DataType::Double(_) => return
text.parse().map(Datum::Double).map_err(|_| invalid()),
- _ => {}
- }
+ return paimon::table::parse_partition_spec_value(text,
data_type).map_err(to_py_err);
}
- py_to_datum(value, data_type)
+ py_to_datum(value, data_type).map(Some)
}
fn partition_spec(
@@ -154,7 +139,7 @@ fn partition_spec(
let datum = if value.is_none() || is_default {
None
} else {
- Some(partition_value(value, field.data_type())?)
+ partition_value(value, field.data_type())?
};
Ok((key, datum))
})
diff --git a/bindings/python/tests/test_table_commit.py
b/bindings/python/tests/test_table_commit.py
index 2de95d8b..e6ee4bdd 100644
--- a/bindings/python/tests/test_table_commit.py
+++ b/bindings/python/tests/test_table_commit.py
@@ -404,6 +404,29 @@ def
test_partition_truncation_matches_java_lifecycle(tmp_path):
assert _rows(table) == []
[email protected]("overwrite", [False, True])
+def test_decimal_partition_precision_overflow_matches_java_null_spec(tmp_path,
overwrite):
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": str(tmp_path)})
+ ctx.sql("CREATE SCHEMA paimon.db")
+ ctx.sql("""CREATE TABLE paimon.db.t (id INT, pt DECIMAL(3, 0)) PARTITIONED
BY (pt)
+ WITH ('dynamic-partition-overwrite' = 'false')""")
+ ctx.sql("INSERT INTO paimon.db.t (id, pt) VALUES (1, NULL), (2, 12)")
+ table = PaimonCatalog({"warehouse": str(tmp_path)}).get_table("db.t")
+ if overwrite:
+ table.new_batch_write_builder().with_overwrite({"pt":
"1000"}).new_commit().commit([])
+ else:
+
table.new_batch_write_builder().new_commit().truncate_partitions([{"pt":
"1000"}])
+ assert _rows(table) == [2]
+
+
+def test_integer_partition_spec_accepts_java_unicode_digits(tmp_path):
+ table = _table(tmp_path, partitioned=True)
+ _append(table, [1, 2], [12, 13])
+ table.new_batch_write_builder().new_commit().truncate_partitions([{"pt":
"１２"}])
+ assert _rows(table) == [2]
+
+
@pytest.mark.parametrize("spec", [{"id": 1}, {"pt": "bad"}, {"pt": True}])
def test_partition_spec_validation_precedes_mutation(tmp_path, spec):
table = _table(tmp_path, partitioned=True)
diff --git a/bindings/python/tests/test_update.py
b/bindings/python/tests/test_update.py
index 68157654..069a3a34 100644
--- a/bindings/python/tests/test_update.py
+++ b/bindings/python/tests/test_update.py
@@ -508,3 +508,50 @@ def
test_prepared_merge_accepts_native_table_sources(tmp_path, stream, historica
actual = pa.Table.from_batches(context.sql(
'SELECT id, value FROM paimon.table_merge.target ORDER BY
id')).to_pylist()
assert actual == [dict(id=1, value=11), dict(id=2, value=20 if historical
else 22), dict(id=3, value=33)]
+
+
[email protected]('stream', [False, True])
+def
test_prepared_merge_callable_uses_complete_clause_and_preserves_exception(tmp_path,
stream):
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.callable_merge')
+ context.sql("""CREATE TABLE paimon.callable_merge.t (id INT, value INT)
WITH (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true')""")
+ for values in ['(1, 10)', '(2, 20)']:
+ context.sql('INSERT INTO paimon.callable_merge.t (id, value) VALUES '
+ values)
+ table = PaimonCatalog({'warehouse':
str(tmp_path)}).get_table('callable_merge.t')
+ builder =
(table.new_stream_write_builder().with_commit_user('callable-merge')
+ if stream else table.new_batch_write_builder())
+ kwargs = dict(commit_identifier=42) if stream else {}
+ update = builder.new_update()
+ source = pa.table({'id': pa.array([1, 2, 3], pa.int32()), 'delta':
pa.array([4, 5, 6], pa.int32())})
+ calls = []
+
+ def values(batch):
+ calls.append(batch.num_rows)
+ assert {'t.id', 't.value', 's.id', 's.delta'} <=
set(batch.column_names)
+ return pa.chunked_array([[100], [200]], type=pa.int32())
+
+ messages = update.merge_into(source, on=[('id', 'id')], when_matched=[
+ dict(delete=False, assignments=[('value', 'callable', values)])],
when_not_matched=[
+ dict(assignments=[('id', 'source', 'id'), ('value', 'callable', lambda
batch: batch['s.delta'])])],
+ **kwargs)
+ assert calls == [2]
+ commit = builder.new_commit()
+ commit.commit(42, messages) if stream else commit.commit(messages)
+ actual = pa.Table.from_batches(context.sql(
+ 'SELECT id, value FROM paimon.callable_merge.t ORDER BY
id')).to_pylist()
+ assert actual == [dict(id=1, value=100), dict(id=2, value=200), dict(id=3,
value=6)]
+ error = RuntimeError('callback error')
+
+ def fail(batch):
+ raise error
+
+ clause = [dict(delete=False, assignments=[('value', 'callable', fail)])]
+ update = builder.new_update()
+ assert update.merge_into(pa.table({'id': pa.array([99], pa.int32())}),
on=[('id', 'id')],
+ when_matched=clause, when_not_matched=[],
**kwargs) == []
+ with pytest.raises(RuntimeError) as caught:
+ update.merge_into(source, on=[('id', 'id')], when_matched=clause,
+ when_not_matched=[], **kwargs)
+ assert caught.value is error
diff --git a/crates/paimon/src/io/file_io/fuse.rs
b/crates/paimon/src/io/file_io/fuse.rs
index 660b0784..d0a287db 100644
--- a/crates/paimon/src/io/file_io/fuse.rs
+++ b/crates/paimon/src/io/file_io/fuse.rs
@@ -96,6 +96,105 @@ mod tests {
base
}
+ #[tokio::test]
+ async fn bound_mount_supports_table_write_commit_and_truncate() {
+ use crate::catalog::Identifier;
+ use crate::spec::{DataType, IntType, Schema, TableSchema};
+ use crate::table::{Table, TableCommit};
+ use arrow_array::{Int32Array, RecordBatch};
+ use futures::TryStreamExt;
+
+ let mount = tempfile::tempdir().unwrap();
+ let root = "oss://bucket/warehouse/table";
+ let io = FileIOBuilder::new("memory")
+ .build()
+ .unwrap()
+ .with_fuse(root, mount.path().to_str().unwrap())
+ .unwrap();
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .build()
+ .unwrap();
+ let table = Table::new(
+ io.clone(),
+ Identifier::new("db", "table"),
+ root.into(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+ // Construction binds schema and data I/O to the same provider.
+ let schema_path = format!("{root}/schema/schema-0");
+ io.new_output(&schema_path)
+ .unwrap()
+ .write(bytes::Bytes::from(
+ serde_json::to_vec(table.schema()).unwrap(),
+ ))
+ .await
+ .unwrap();
+ assert_eq!(table.schema_manager().schema(0).await.unwrap().id(), 0);
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ let batch = RecordBatch::try_from_iter(vec![(
+ "id",
+ Arc::new(Int32Array::from(vec![1, 2])) as arrow_array::ArrayRef,
+ )])
+ .unwrap();
+ writer.write_arrow_batch(&batch).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ let files = messages
+ .iter()
+ .flat_map(|m| {
+ m.new_files
+ .iter()
+ .map(move |f| format!("{root}/bucket-{}/{}", m.bucket,
f.file_name))
+ })
+ .collect::<Vec<_>>();
+ let committer = TableCommit::new(table.clone(),
builder.commit_user().into());
+ committer.commit_with_identifier(messages, 1).await.unwrap();
+ writer.close().await;
+ assert!(mount.path().join("snapshot/snapshot-1").exists());
+ for file in &files {
+ assert!(io.exists(file).await.unwrap());
+ }
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let mut batches =
reader.new_read().unwrap().to_arrow(plan.splits()).unwrap();
+ let mut count = 0;
+ while let Some(batch) = batches.try_next().await.unwrap() {
+ count += batch.num_rows();
+ }
+ assert_eq!(count, 2);
+ committer.truncate_table_with_identifier(17).await.unwrap();
+ assert!(table
+ .new_read_builder()
+ .new_scan()
+ .plan()
+ .await
+ .unwrap()
+ .splits()
+ .is_empty());
+ // Logical truncation must retain files referenced by older snapshots.
+ for file in &files {
+ assert!(io.exists(file).await.unwrap());
+ }
+ let mut writer = builder.new_write().unwrap();
+ writer.write_arrow_batch(&batch).await.unwrap();
+ committer
+ .commit_with_identifier(writer.prepare_commit().await.unwrap(), 18)
+ .await
+ .unwrap();
+ // Retrying an uncertain truncate cannot remove the later append.
+ committer.truncate_table_with_identifier(17).await.unwrap();
+ assert!(!table
+ .new_read_builder()
+ .new_scan()
+ .plan()
+ .await
+ .unwrap()
+ .splits()
+ .is_empty());
+ }
+
#[tokio::test]
async fn
mounted_cache_reuses_files_isolates_mounts_and_invalidates_writes() {
let first = tempfile::tempdir().unwrap();
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index cba2dab6..c9850975 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -1618,6 +1618,49 @@ impl<'a> CoreOptions<'a> {
.unwrap_or(DEFAULT_TARGET_FILE_SIZE)
}
+ /// Maximum automatically inferred postpone bucket count (Java default:
2048).
+ pub fn postpone_batch_write_fixed_bucket_max_parallelism(&self) ->
crate::Result<i32> {
+ let key = "postpone.batch-write-fixed-bucket.max-parallelism";
+ let Some(raw) = self.options.get(key) else {
+ return Ok(2048);
+ };
+ raw.parse::<i32>()
+ .ok()
+ .filter(|n| *n > 0)
+ .ok_or_else(|| crate::Error::ConfigInvalid {
+ message: format!("{key} must be a positive 32-bit integer, got
'{raw}'"),
+ })
+ }
+
+ pub fn postpone_target_row_num_per_bucket(&self) ->
crate::Result<Option<i64>> {
+ self.options
+ .get("postpone.target-row-num-per-bucket")
+ .map(|raw| {
+ raw.parse::<i64>().ok().filter(|n| *n > 0).ok_or_else(|| {
+ crate::Error::ConfigInvalid {
+ message: format!(
+ "postpone.target-row-num-per-bucket must be
positive, got '{raw}'"
+ ),
+ }
+ })
+ })
+ .transpose()
+ }
+
+ pub fn postpone_target_size_per_bucket(&self) -> crate::Result<i64> {
+ self.options
+ .get("postpone.target-size-per-bucket")
+ .map_or(Ok(1024 * 1024 * 1024), |raw| {
+ parse_memory_size(raw).filter(|n| *n > 0).ok_or_else(|| {
+ crate::Error::ConfigInvalid {
+ message: format!(
+ "postpone.target-size-per-bucket must be positive,
got '{raw}'"
+ ),
+ }
+ })
+ })
+ }
+
/// Maximum rows in a newly written data file. Java defaults to
`Long.MAX_VALUE`.
pub fn target_file_row_num(&self) -> crate::Result<i64> {
let Some(raw) = self.options.get("target-file-row-num") else {
diff --git a/crates/paimon/src/table/branch_write_tests.rs
b/crates/paimon/src/table/branch_write_tests.rs
index 612df694..c1538848 100644
--- a/crates/paimon/src/table/branch_write_tests.rs
+++ b/crates/paimon/src/table/branch_write_tests.rs
@@ -563,6 +563,123 @@ async fn
branch_write_schema_id_is_resolved_in_its_own_namespace() {
);
}
+#[tokio::test]
+async fn rest_branch_copy_retains_bound_file_io_and_loads_branch_identity() {
+ use axum::{extract::Path, routing::get, Json, Router};
+
+ let mount = tempfile::tempdir().unwrap();
+ let location = "oss://bucket/warehouse/db/t";
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("query-auth.enabled", "true")
+ .build()
+ .unwrap();
+ let response = move |name: &str| {
+ serde_json::json!({
+ "database": "db", "name": name, "id": format!("uuid-{name}"),
+ "path": location, "isExternal": false,
+ "schemaId": if name == "t" { 0 } else { 1 }, "schema": schema
+ })
+ };
+ let original = response("t");
+ let app = Router::new().route(
+ "/v1/databases/:db/tables/:table",
+ get(move |Path((_db, name)): Path<(String, String)>| {
+ let body = response(&name);
+ async move { Json(body) }
+ }),
+ );
+ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let address = listener.local_addr().unwrap();
+ let server = tokio::spawn(async move { axum::serve(listener,
app).await.unwrap() });
+ let mut options = crate::Options::new();
+ for (key, value) in [
+ ("uri", format!("http://{address}";)),
+ ("warehouse", "test".into()),
+ ("token.provider", "bear".into()),
+ ("token", "test-token".into()),
+ ("fuse.enabled", "true".into()),
+ ("fuse.mode", "pvfs".into()),
+ ("fuse.root", mount.path().to_string_lossy().into_owned()),
+ ("fuse.validation-mode", "none".into()),
+ ] {
+ options.set(key, value);
+ }
+ let main = Table::from_rest_response(
+ Identifier::new("db", "t"),
+ serde_json::from_value(original).unwrap(),
+ options,
+ )
+ .await
+ .unwrap();
+ let branch_schema = TableSchema::new(
+ 1,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .option("query-auth.enabled", "true")
+ .build()
+ .unwrap(),
+ );
+ let schema_path = format!("{location}/branch/branch-dev/schema/schema-1");
+ main.file_io()
+ .new_output(&schema_path)
+ .unwrap()
+ .write(bytes::Bytes::from(
+ serde_json::to_vec(&branch_schema).unwrap(),
+ ))
+ .await
+ .unwrap();
+ // The Python binding's existing Table.copy delegates to this core API.
+ let copied = main
+ .copy_with_time_travel_strict(HashMap::from([
+ ("branch".into(), "dev".into()),
+ ("read.batch-size".into(), "17".into()),
+ ]))
+ .await
+ .unwrap();
+ assert_eq!(copied.branch(), "dev");
+ assert_eq!(copied.schema().id(), 1);
+ assert_eq!(
+ copied.rest_env().unwrap().identifier().object(),
+ "t$branch_dev"
+ );
+ assert_eq!(copied.rest_env().unwrap().uuid(), "uuid-t$branch_dev");
+ assert!(copied.query_auth_session.is_some());
+ assert_ne!(copied.query_auth_session, main.query_auth_session);
+ assert_eq!(copied.schema().options()["read.batch-size"], "17");
+ assert_eq!(
+ copied.schema_manager().schema(1).await.unwrap().as_ref(),
+ &branch_schema
+ );
+ assert!(!mount
+ .path()
+ .join("db/t$branch_dev/branch/branch-dev/schema/schema-1")
+ .exists());
+ let main_schema_path = format!("{location}/schema/schema-0");
+ main.file_io()
+ .new_output(&main_schema_path)
+ .unwrap()
+ .write(bytes::Bytes::from(
+ serde_json::to_vec(main.schema()).unwrap(),
+ ))
+ .await
+ .unwrap();
+ let restored = copied
+ .copy_with_time_travel_strict(HashMap::from([("branch".into(),
"main".into())]))
+ .await
+ .unwrap();
+ assert_eq!(restored.branch(), "main");
+ assert_eq!(restored.schema().id(), 0);
+ assert_eq!(restored.schema().options()["read.batch-size"], "17");
+ assert_eq!(restored.rest_env().unwrap().uuid(), "uuid-t");
+ assert_ne!(restored.query_auth_session, copied.query_auth_session);
+ assert_eq!(
+ restored.schema_manager().schema(0).await.unwrap().as_ref(),
+ main.schema()
+ );
+ server.abort();
+}
+
#[tokio::test]
async fn
branch_write_rest_publisher_routes_by_argument_and_preserves_payload() {
use crate::api::rest_api::RESTApi;
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 19b9315d..42d1a16b 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -103,6 +103,7 @@ mod mosaic_table_write_tests;
mod object_table;
mod partition_filter;
mod partition_row_count;
+mod partition_spec;
mod partition_stat;
#[cfg(feature = "fulltext")]
mod pk_full_text_bucket_search;
@@ -123,6 +124,7 @@ mod pk_vector_read;
mod pk_vector_scan;
mod pk_vector_search_params;
mod postpone_bucket_plan;
+mod postpone_bucket_planner;
mod postpone_file_writer;
mod postpone_fixed_bucket_router;
mod postpone_fixed_bucket_write;
@@ -221,6 +223,7 @@ pub use incremental_scan::{
pub use instant::Instant;
pub use object_table::{ObjectEntry, ObjectTable};
pub use partition_row_count::PartitionRowCount;
+pub use partition_spec::parse_partition_spec_value;
pub use partition_stat::PartitionStat;
pub use pk_vector_bucket_split::{BucketVectorPayload, BucketVectorSearchSplit};
pub use postpone_bucket_plan::{PostponeBucketPlan,
POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD};
@@ -867,6 +870,8 @@ impl Table {
/// table's fields and keys come from that snapshot's schema while the
/// runtime options stay merged. Column declarations follow the historical
/// schema unless explicitly overridden, as in Java.
+ /// A different `branch` option first loads that branch's schema and
catalog
+ /// identity, retaining the bound FileIO and preceding dynamic overrides.
/// Like Java, resolution failures fall back silently to the current
/// schema (the `if let Ok` below swallows them); an invalid selector
/// still fails later at scan planning.
@@ -947,7 +952,28 @@ impl Table {
CoreOptions::new(self.schema.options())
.ensure_type_paimon_served(&self.identifier.full_name())?;
self.ensure_format_table_partition_options_unchanged(&extra)?;
- let mut table = self.copy_with_options(extra);
+ // Resolve a different branch before applying dynamic read options so
+ // its schema, REST identity and metadata namespace stay consistent.
+ let mut table = match extra.get("branch") {
+ Some(branch) if branch != self.branch() => {
+ let target = self.copy_with_branch(branch).await?;
+ // Earlier copy() overrides remain in effect on the new branch;
+ // other defaults come from its own schema.
+ let mut dynamic: HashMap<_, _> = self
+ .applied_dynamic_option_keys
+ .iter()
+ .filter_map(|key| {
+ self.schema
+ .options()
+ .get(key)
+ .map(|value| (key.clone(), value.clone()))
+ })
+ .collect();
+ dynamic.extend(extra);
+ target.copy_with_options(dynamic)
+ }
+ _ => self.copy_with_options(extra),
+ };
// Reject unimplemented scan options on the merged view before any IO,
so
// both table-level and per-read options are covered.
CoreOptions::new(table.schema().options()).validate_scan_options()?;
@@ -1028,7 +1054,8 @@ impl Table {
};
if let Some(env) = &self.rest_env {
// Catalog branches may have no schema directory in the data store.
- // Load their schema, identity and FileIO through the same catalog.
+ // Load their schema and identity through the same catalog. A table
+ // copy retains the original FileIO, including its bound FUSE
mount.
let identifier = env.identifier().with_branch(&branch)?;
let (current, mut table) = tokio::try_join!(
env.current_table_identity_checked(),
@@ -1045,6 +1072,13 @@ impl Table {
source: None,
});
}
+ table.file_io = self.file_io.clone();
+ let schema_manager = SchemaManager::new(table.file_io.clone(),
table.location.clone());
+ table.schema_manager = if branch == DEFAULT_MAIN_BRANCH {
+ schema_manager
+ } else {
+ schema_manager.with_branch(&branch)
+ };
let mut options = table.schema.options().clone();
options.insert("branch".to_string(), branch);
table.schema = table.schema.copy_with_replaced_options(options);
diff --git a/crates/paimon/src/table/partition_spec.rs
b/crates/paimon/src/table/partition_spec.rs
new file mode 100644
index 00000000..1d1c8615
--- /dev/null
+++ b/crates/paimon/src/table/partition_spec.rs
@@ -0,0 +1,284 @@
+// 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.
+
+//! Java TypeUtils conversion for textual overwrite and truncate partition
specs.
+
+use crate::spec::{DataType, Datum};
+use arrow_array::{ArrayRef, RecordBatch, StringArray};
+use arrow_schema::{Field, Schema};
+use std::sync::Arc;
+
+/// Convert a Java partition spec, retaining DECIMAL precision overflow as
NULL.
+pub fn parse_partition_spec_value(
+ text: &str,
+ data_type: &DataType,
+) -> crate::Result<Option<Datum>> {
+ let invalid = || crate::Error::DataInvalid {
+ message: format!("Invalid partition value '{text}' for {data_type:?}"),
+ source: None,
+ };
+ match data_type {
+ // TypeUtils uses valueOf, not the decimal-truncating SQL cast rules.
+ DataType::TinyInt(_) | DataType::SmallInt(_) | DataType::Int(_) |
DataType::BigInt(_) => {
+ let text = if text.is_ascii() {
+ std::borrow::Cow::Borrowed(text)
+ } else {
+ // Integer.valueOf scans Character.digit, unlike SQL's
+ // BinaryStringUtils integer cast with its fractional grammar.
+ std::borrow::Cow::Owned(
+ text.chars()
+ .map(|ch| {
+ crate::common::number::java_decimal_digit(ch)
+ .map_or(ch, |digit| char::from(b'0' + digit))
+ })
+ .collect::<String>(),
+ )
+ };
+ return match data_type {
+ DataType::TinyInt(_) => text.parse().map(Datum::TinyInt),
+ DataType::SmallInt(_) => text.parse().map(Datum::SmallInt),
+ DataType::Int(_) => text.parse().map(Datum::Int),
+ DataType::BigInt(_) => text.parse().map(Datum::Long),
+ _ => unreachable!(),
+ }
+ .map(Some)
+ .map_err(|_| invalid());
+ }
+ DataType::Float(_) => {
+ use crate::common::number::{format_java_float, parse_java_float};
+ let double = parse_java_float(text, false)?;
+ let single = parse_java_float(text, true)? as f32;
+ return if double == f64::from(single)
+ || format_java_float(single) == format_java_float(double)
+ {
+ Ok(Some(Datum::Float(single)))
+ } else {
+ Err(invalid())
+ };
+ }
+ DataType::Double(_) => {
+ return crate::common::number::parse_java_float(text, false)
+ .map(Datum::Double)
+ .map(Some)
+ }
+ DataType::Time(_) => return
parse_time_spec(text).map(Datum::Time).map(Some),
+ DataType::Char(ty) if text.chars().count() > ty.length() => return
Err(invalid()),
+ DataType::VarChar(ty) if text.chars().count() > ty.length() as usize
=> {
+ return Err(invalid())
+ }
+ DataType::Char(_) | DataType::VarChar(_) => return
Ok(Some(Datum::String(text.into()))),
+ DataType::Binary(_) => return
Ok(Some(Datum::Bytes(text.as_bytes().to_vec()))),
+ DataType::VarBinary(ty) => {
+ return if text.len() <= ty.length() as usize {
+ Ok(Some(Datum::Bytes(text.as_bytes().to_vec())))
+ } else {
+ Err(invalid())
+ };
+ }
+ DataType::Boolean(_)
+ | DataType::Date(_)
+ | DataType::Timestamp(_)
+ | DataType::LocalZonedTimestamp(_)
+ | DataType::Decimal(_) => {}
+ _ => return Err(invalid()),
+ }
+ let input: ArrayRef = Arc::new(StringArray::from(vec![text]));
+ let array = crate::arrow::cast_array_for_schema_evolution(
+ &input,
+ &DataType::VarChar(crate::spec::VarCharType::string_type()),
+ data_type,
+ )?;
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "partition",
+ array.data_type().clone(),
+ true,
+ )])),
+ vec![array],
+ )
+ .map_err(|_| invalid())?;
+ crate::spec::extract_datum_from_arrow(&batch, 0, 0, data_type)
+}
+
+fn parse_time_spec(text: &str) -> crate::Result<i32> {
+ let invalid = || crate::Error::DataInvalid {
+ message: format!("Invalid TIME partition value '{text}'"),
+ source: None,
+ };
+ // BinaryStringUtils.toTime accepts the legacy millis representation.
+ let digits = text.strip_prefix('-').unwrap_or(text);
+ if !digits.is_empty() && digits.bytes().all(|c| c.is_ascii_digit()) {
+ return text.parse().map_err(|_| invalid());
+ }
+ // DateTimeUtils.parseTime accepts partial clocks and adds the signed zone
+ // offset, without restricting clock components to a single day's range.
+ let component = |part: &str| -> crate::Result<i32> {
+ let part = part.trim_matches(|c| c <= ' ');
+ if part.is_empty() || !part.bytes().all(|c| c.is_ascii_digit()) {
+ return Err(invalid());
+ }
+ part.parse().map_err(|_| invalid())
+ };
+ let zone_at = text.find('-').or_else(|| text.find('+'));
+ let (clock, offset) = match zone_at {
+ Some(at) => {
+ let sign: i32 = if text.as_bytes()[at] == b'-' { -1 } else { 1 };
+ let (hours, minutes) = text[at + 1..]
+ .split_once(':')
+ .unwrap_or((&text[at + 1..], "0"));
+ let millis = component(hours)?
+ .wrapping_mul(3_600_000)
+ .wrapping_add(component(minutes)?.wrapping_mul(60_000));
+ (&text[..at], sign.wrapping_mul(millis))
+ }
+ None => (text, 0),
+ };
+ let parts = clock.split(':').collect::<Vec<_>>();
+ if parts.len() > 3 || parts.len() < 3 && zone_at.is_some_and(|at| text[at
+ 1..].contains(':'))
+ {
+ // Java searches for the clock's colons in the complete string. With
+ // missing seconds a colon in the offset is not a valid second colon.
+ return Err(invalid());
+ }
+ let hours = component(parts[0])?;
+ let minutes = parts.get(1).map(|v| component(v)).transpose()?.unwrap_or(0);
+ let (seconds, fraction) = match parts.get(2) {
+ Some(value) => {
+ let (seconds, fraction) = value.split_once('.').unwrap_or((value,
""));
+ (component(seconds)?, fraction.trim_matches(|c| c <= ' '))
+ }
+ None => (0, ""),
+ };
+ // Java's fraction parser rounds using the fourth character.
+ let mut millis = 0;
+ let mut chars = fraction.chars();
+ for multiplier in [100, 10, 1] {
+ let Some(ch) = chars.next() else { break };
+ millis += ch.to_digit(10).filter(|_| ch.is_ascii()).unwrap_or(0) as
i32 * multiplier;
+ }
+ if chars.next().is_some_and(|ch| ch >= '5') {
+ millis += 1;
+ }
+ Ok(hours
+ .wrapping_mul(3_600_000)
+ .wrapping_add(minutes.wrapping_mul(60_000))
+ .wrapping_add(seconds.wrapping_mul(1000))
+ .wrapping_add(millis)
+ .wrapping_add(offset))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::{
+ BooleanType, DateType, DecimalType, FloatType, IntType, TimeType,
TimestampType,
+ };
+
+ #[test]
+ fn partition_specs_accept_java_unicode_and_decimal_null() {
+ let integer = DataType::Int(IntType::new());
+ for (value, expected) in [("１２", 12), ("١٢", 12), ("-１２", -12),
("+١٢", 12)] {
+ assert_eq!(
+ parse_partition_spec_value(value, &integer).unwrap(),
+ Some(Datum::Int(expected))
+ );
+ }
+ let decimal = DataType::Decimal(DecimalType::new(3, 0).unwrap());
+ assert_eq!(parse_partition_spec_value("1000", &decimal).unwrap(),
None);
+ assert_eq!(parse_partition_spec_value("999.5", &decimal).unwrap(),
None);
+ assert!(parse_partition_spec_value("not a decimal",
&decimal).is_err());
+ let time = DataType::Time(TimeType::new(3).unwrap());
+ for value in ["12:34+02:30", "12+02:30"] {
+ assert!(parse_partition_spec_value(value, &time).is_err());
+ }
+ }
+
+ #[test]
+ fn partition_specs_use_java_typed_text_conversion() {
+ assert_eq!(
+ parse_partition_spec_value("TRUE",
&DataType::Boolean(BooleanType::new()))
+ .unwrap()
+ .unwrap(),
+ Datum::Bool(true)
+ );
+ assert_eq!(
+ parse_partition_spec_value("1970-01-02",
&DataType::Date(DateType::new()))
+ .unwrap()
+ .unwrap(),
+ Datum::Date(1)
+ );
+ assert_eq!(
+ parse_partition_spec_value("12:34:56.789",
&DataType::Time(TimeType::new(3).unwrap()))
+ .unwrap()
+ .unwrap(),
+ Datum::Time(45_296_789)
+ );
+ assert_eq!(
+ parse_partition_spec_value(
+ "1970-01-01 00:00:00.123456",
+ &DataType::Timestamp(TimestampType::new(1).unwrap())
+ )
+ .unwrap()
+ .unwrap(),
+ Datum::Timestamp {
+ millis: 100,
+ nanos: 0
+ }
+ );
+ assert!(matches!(
+ parse_partition_spec_value(
+ "12.345",
+ &DataType::Decimal(DecimalType::new(6, 2).unwrap())
+ )
+ .unwrap()
+ .unwrap(),
+ Datum::Decimal { unscaled: 1235, .. }
+ ));
+ for text in ["42.9", "2147483648", " 42"] {
+ assert!(parse_partition_spec_value(text,
&DataType::Int(IntType::new())).is_err());
+ }
+ let float = DataType::Float(FloatType::new());
+ assert_eq!(
+ parse_partition_spec_value("1.2", &float).unwrap().unwrap(),
+ Datum::Float(1.2)
+ );
+ assert_eq!(
+ parse_partition_spec_value("0x1p2", &float)
+ .unwrap()
+ .unwrap(),
+ Datum::Float(4.0)
+ );
+ assert!(parse_partition_spec_value("1.23456789", &float).is_err());
+ let time = DataType::Time(TimeType::new(0).unwrap());
+ for (text, millis) in [
+ ("45296789", 45_296_789),
+ ("-1", -1),
+ ("12:34", 45_240_000),
+ ("12:34:56.7896+02:30", 54_296_790),
+ ("12:34:56.7896-02:30", 36_296_790),
+ ("24:00:00", 86_400_000),
+ ] {
+ assert_eq!(
+ parse_partition_spec_value(text, &time).unwrap().unwrap(),
+ Datum::Time(millis)
+ );
+ }
+ for text in ["+1", "12:bad", "12:34:56:78", "2147483648:00"] {
+ assert!(parse_partition_spec_value(text, &time).is_err());
+ }
+ }
+}
diff --git a/crates/paimon/src/table/postpone_bucket_plan.rs
b/crates/paimon/src/table/postpone_bucket_plan.rs
index f17c6db6..83bba51f 100644
--- a/crates/paimon/src/table/postpone_bucket_plan.rs
+++ b/crates/paimon/src/table/postpone_bucket_plan.rs
@@ -38,6 +38,12 @@ pub struct PostponeBucketPlan {
}
impl PostponeBucketPlan {
+ pub(super) fn from_counts(bucket_counts: HashMap<Vec<u8>, i32>) -> Self {
+ Self {
+ bucket_counts: Arc::new(bucket_counts),
+ default_bucket_num: None,
+ }
+ }
pub fn from_arrow(table: &Table, batch: &RecordBatch) -> Result<Self> {
let partition_fields = table.schema().partition_fields();
let partition_count = partition_fields.len();
diff --git a/crates/paimon/src/table/postpone_bucket_planner.rs
b/crates/paimon/src/table/postpone_bucket_planner.rs
new file mode 100644
index 00000000..5efac025
--- /dev/null
+++ b/crates/paimon/src/table/postpone_bucket_planner.rs
@@ -0,0 +1,437 @@
+// 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.
+
+//! Snapshot-bound sizing for fixed-bucket batch writes.
+
+use super::postpone_bucket_plan::data_invalid;
+use crate::spec::{batch_to_serialized_bytes, BinaryRow, CoreOptions,
DataField, DataType};
+use crate::table::{PostponeBucketPlan, Table, TableScan};
+use crate::Result;
+use arrow_array::{
+ Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, MapArray,
RecordBatch,
+ StructArray,
+};
+use arrow_schema::{DataType as ArrowType, Schema};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+pub(super) struct PostponeBucketPlanner {
+ fields: Vec<DataField>,
+ partitions: Vec<usize>,
+ known: HashMap<Vec<u8>, i32>,
+ input: HashMap<Vec<u8>, (u128, u128)>,
+ target_rows: Option<u128>,
+ target_size: u128,
+ maximum: u128,
+ batches: Vec<RecordBatch>,
+}
+
+impl PostponeBucketPlanner {
+ pub(super) fn new(table: &Table) -> Result<Self> {
+ let options = CoreOptions::new(table.schema().options());
+ let target_rows = options
+ .postpone_target_row_num_per_bucket()?
+ .map(|n| n as u128);
+ Ok(Self {
+ fields: table.schema().fields().to_vec(),
+ partitions: table
+ .schema()
+ .partition_keys()
+ .iter()
+ .map(|name| {
+ table
+ .schema()
+ .fields()
+ .iter()
+ .position(|f| f.name() == name)
+ .unwrap()
+ })
+ .collect(),
+ known: HashMap::new(),
+ input: HashMap::new(),
+ target_rows,
+ // Java ignores byte sizing when a row target is supplied.
+ target_size: if target_rows.is_none() {
+ options.postpone_target_size_per_bucket()? as u128
+ } else {
+ 1
+ },
+ maximum:
options.postpone_batch_write_fixed_bucket_max_parallelism()? as u128,
+ batches: Vec::new(),
+ })
+ }
+
+ pub(super) async fn load_metadata(&mut self, table: &Table, snapshot_id:
i64) -> Result<()> {
+ if snapshot_id == 0 {
+ return Ok(());
+ }
+ let snapshot =
table.snapshot_manager().get_snapshot(snapshot_id).await?;
+ let entries = TableScan::new(table, None, vec![], None, None, None)
+ .with_scan_all_files()
+ .plan_manifest_entries(&snapshot)
+ .await?;
+ for entry in entries {
+ let partition = entry.partition().to_vec();
+ if entry.bucket() >= 0 {
+ let count = entry.total_buckets();
+ if count <= 0 || entry.bucket() >= count {
+ return Err(data_invalid(
+ "Invalid total bucket count in existing postpone
files",
+ ));
+ }
+ if self
+ .known
+ .insert(partition, count)
+ .is_some_and(|previous| previous != count)
+ {
+ return Err(data_invalid("Existing postpone files contain
conflicting total bucket counts for one partition"));
+ }
+ }
+ }
+ Ok(())
+ }
+
+ pub(super) fn known_plan(&self) -> PostponeBucketPlan {
+ PostponeBucketPlan::from_counts(self.known.clone())
+ }
+
+ pub(super) fn add(&mut self, batch: RecordBatch) ->
Result<Option<RecordBatch>> {
+ let partitions = batch_to_serialized_bytes(&batch, &self.partitions,
&self.fields)?;
+ let mut known_rows = Vec::new();
+ let mut unknown_rows = Vec::new();
+ for (row, partition) in partitions.into_iter().enumerate() {
+ if self.known.contains_key(&partition) {
+ known_rows.push(row);
+ continue;
+ }
+ unknown_rows.push(row);
+ let size = if self.target_rows.is_none() {
+ row_size(&batch, row, &self.fields)?
+ } else {
+ 0
+ };
+ let stats = self.input.entry(partition).or_default();
+ stats.0 += 1;
+ stats.1 += size as u128;
+ }
+ if unknown_rows.is_empty() {
+ return Ok(Some(batch));
+ }
+ if known_rows.is_empty() {
+ self.batches.push(batch);
+ return Ok(None);
+ }
+ let known = super::table_write::take_rows(&batch, &known_rows)?;
+ self.batches
+ .push(super::table_write::take_rows(&batch, &unknown_rows)?);
+ Ok(Some(known))
+ }
+
+ pub(super) fn finish(mut self) -> (PostponeBucketPlan, Vec<RecordBatch>) {
+ for (partition, (rows, size)) in self.input {
+ let required = match self.target_rows {
+ Some(target) => rows.div_ceil(target),
+ None => size.div_ceil(self.target_size),
+ };
+ // PostponeUtils rounds inferred counts up before applying the cap.
+ let count = required.max(1).next_power_of_two().min(self.maximum)
as i32;
+ self.known.insert(partition, count);
+ }
+ (PostponeBucketPlan::from_counts(self.known), self.batches)
+ }
+}
+
+fn round_word(size: usize) -> usize {
+ size.div_ceil(8) * 8
+}
+
+fn row_size(batch: &RecordBatch, row: usize, fields: &[DataField]) ->
Result<usize> {
+ let mut size = BinaryRow::cal_fix_part_size_in_bytes(fields.len() as i32)
as usize;
+ for (column, field) in batch.columns().iter().zip(fields) {
+ size += variable_size(column, row, field.data_type())?;
+ }
+ Ok(size)
+}
+
+fn binary_length(array: &ArrayRef, row: usize) -> Result<usize> {
+ macro_rules! length {
+ ($ty:ty) => {
+ if let Some(values) = array.as_any().downcast_ref::<$ty>() {
+ return Ok(values.value(row).len());
+ }
+ };
+ }
+ length!(arrow_array::BinaryArray);
+ length!(arrow_array::LargeBinaryArray);
+ length!(arrow_array::FixedSizeBinaryArray);
+ length!(arrow_array::StringArray);
+ length!(arrow_array::LargeStringArray);
+ length!(arrow_array::StringViewArray);
+ Err(data_invalid(
+ "Expected binary or string input while sizing postpone data",
+ ))
+}
+
+fn array_width(data_type: &DataType) -> usize {
+ match data_type {
+ DataType::Boolean(_) | DataType::TinyInt(_) => 1,
+ DataType::SmallInt(_) => 2,
+ DataType::Int(_) | DataType::Float(_) | DataType::Date(_) |
DataType::Time(_) => 4,
+ _ => 8,
+ }
+}
+
+fn array_size(values: &ArrayRef, data_type: &DataType) -> Result<usize> {
+ let count = values.len();
+ let mut size = round_word(4 + count.div_ceil(32) * 4 + count *
array_width(data_type));
+ for row in 0..count {
+ if !values.is_null(row) {
+ size += variable_size(values, row, data_type)?;
+ }
+ }
+ Ok(size)
+}
+
+fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) ->
Result<usize> {
+ // Decimal and non-compact timestamp writers reserve their variable slot
+ // even for NULL, as Java BinaryRowWriter does.
+ if let DataType::Decimal(decimal) = data_type {
+ return Ok(if decimal.precision() > 18 { 16 } else { 0 });
+ }
+ if let DataType::Timestamp(timestamp) = data_type {
+ return Ok(if timestamp.precision() > 3 { 8 } else { 0 });
+ }
+ if let DataType::LocalZonedTimestamp(timestamp) = data_type {
+ return Ok(if timestamp.precision() > 3 { 8 } else { 0 });
+ }
+ if array.is_null(row) {
+ return Ok(0);
+ }
+ match data_type {
+ DataType::Blob(_)
+ | DataType::Binary(_)
+ | DataType::VarBinary(_)
+ | DataType::Char(_)
+ | DataType::VarChar(_) => {
+ let length = binary_length(array, row)?;
+ Ok(if length <= 7 { 0 } else { round_word(length) })
+ }
+ DataType::Row(row_type) => {
+ let values = array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| data_invalid("Expected ROW input while sizing
postpone data"))?;
+ let mut size =
+ BinaryRow::cal_fix_part_size_in_bytes(row_type.fields().len()
as i32) as usize;
+ for (value, field) in
values.columns().iter().zip(row_type.fields()) {
+ size += variable_size(value, row, field.data_type())?;
+ }
+ Ok(round_word(size))
+ }
+ DataType::Array(element) => {
+ let values = if let Some(values) =
array.as_any().downcast_ref::<ListArray>() {
+ values.value(row)
+ } else if let Some(values) =
array.as_any().downcast_ref::<LargeListArray>() {
+ values.value(row)
+ } else if let Some(values) =
array.as_any().downcast_ref::<FixedSizeListArray>() {
+ values.value(row)
+ } else {
+ return Err(data_invalid(
+ "Expected ARRAY input while sizing postpone data",
+ ));
+ };
+ Ok(round_word(array_size(&values, element.element_type())?))
+ }
+ DataType::Map(map) => map_size(array, row, map.key_type(),
map.value_type()),
+ DataType::Multiset(multiset) => map_size(
+ array,
+ row,
+ multiset.element_type(),
+ &DataType::Int(crate::spec::IntType::new()),
+ ),
+ DataType::Vector(vector) => {
+ let values = array
+ .as_any()
+ .downcast_ref::<FixedSizeListArray>()
+ .ok_or_else(|| data_invalid("Expected VECTOR input while
sizing postpone data"))?;
+ Ok(round_word(
+ 4 + values.value(row).len() *
array_width(vector.element_type()),
+ ))
+ }
+ _ => {
+ // Reuse the existing Java BinaryRow encoder for scalar layouts,
+ // including inline strings/BINARY and Variant serialization.
+ if matches!(
+ array.data_type(),
+ ArrowType::Boolean
+ | ArrowType::Int8
+ | ArrowType::Int16
+ | ArrowType::Int32
+ | ArrowType::Int64
+ | ArrowType::Float32
+ | ArrowType::Float64
+ | ArrowType::Date32
+ | ArrowType::Time32(_)
+ ) {
+ return Ok(0);
+ }
+ let field = DataField::new(0, "value".into(), data_type.clone());
+ let batch = RecordBatch::try_new(
+ Arc::new(Schema::new(vec![arrow_schema::Field::new(
+ "value",
+ array.data_type().clone(),
+ true,
+ )])),
+ vec![array.slice(row, 1)],
+ )
+ .map_err(|e| data_invalid(e.to_string()))?;
+ let encoded = batch_to_serialized_bytes(&batch, &[0], &[field])?;
+ Ok(encoded[0].len() - 4 - BinaryRow::cal_fix_part_size_in_bytes(1)
as usize)
+ }
+ }
+}
+
+fn map_size(array: &ArrayRef, row: usize, key: &DataType, value: &DataType) ->
Result<usize> {
+ let entries = array
+ .as_any()
+ .downcast_ref::<MapArray>()
+ .ok_or_else(|| data_invalid("Expected MAP input while sizing postpone
data"))?
+ .value(row);
+ Ok(round_word(
+ 4 + array_size(entries.column(0), key)? +
array_size(entries.column(1), value)?,
+ ))
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::{BlobType, DecimalType};
+ use crate::table::table_write::tests::{
+ make_batch, setup_dirs, test_file_io, test_postpone_pk_table,
+ };
+ use arrow_array::{Decimal128Array, LargeBinaryArray};
+
+ #[test]
+ fn binary_row_sizing_handles_blob_and_null_array_values() {
+ let blob: ArrayRef = Arc::new(LargeBinaryArray::from(vec![
+ Some(b"123456789".as_slice()),
+ None,
+ ]));
+ let blob_type = DataType::Blob(BlobType::new());
+ assert_eq!(variable_size(&blob, 0, &blob_type).unwrap(), 16);
+ assert_eq!(variable_size(&blob, 1, &blob_type).unwrap(), 0);
+ let decimal = DataType::Decimal(DecimalType::new(38, 0).unwrap());
+ let values: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![None])
+ .with_precision_and_scale(38, 0)
+ .unwrap(),
+ );
+ assert_eq!(variable_size(&values, 0, &decimal).unwrap(), 16);
+ assert_eq!(array_size(&values, &decimal).unwrap(), 16);
+ let values: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![Some(123i128)])
+ .with_precision_and_scale(38, 0)
+ .unwrap(),
+ );
+ assert_eq!(array_size(&values, &decimal).unwrap(), 32);
+ }
+
+ #[tokio::test]
+ async fn
inference_excludes_old_pending_files_and_streams_known_partitions() {
+ let io = test_file_io();
+ let path = "memory:/auto_postpone_pending";
+ setup_dirs(&io, path).await;
+ let table = test_postpone_pk_table(&io,
path).copy_with_options(HashMap::from([
+ ("postpone.target-row-num-per-bucket".into(), "1".into()),
+ (
+ "postpone.batch-write-fixed-bucket.max-parallelism".into(),
+ "8".into(),
+ ),
+ ]));
+ let pending = table.new_write_builder();
+ let mut writer = pending.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1, 2, 3, 4], vec![10, 20, 30,
40]))
+ .await
+ .unwrap();
+ pending
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let mut planner = PostponeBucketPlanner::new(&table).unwrap();
+ planner.load_metadata(&table, 1).await.unwrap();
+ planner.add(make_batch(vec![5], vec![50])).unwrap();
+ let (plan, batches) = planner.finish();
+ assert_eq!(
+ plan.total_buckets(&crate::spec::EMPTY_SERIALIZED_ROW),
+ Some(1)
+ );
+ assert_eq!(batches.len(), 1);
+ let mut planner = PostponeBucketPlanner::new(&table).unwrap();
+ planner
+ .known
+ .insert(crate::spec::EMPTY_SERIALIZED_ROW.clone(), 3);
+ let known = planner
+ .add(make_batch(vec![6, 7], vec![60, 70]))
+ .unwrap()
+ .unwrap();
+ assert_eq!(known.num_rows(), 2);
+ assert!(
+ planner.batches.is_empty(),
+ "known partitions must not retain the whole input"
+ );
+ }
+
+ #[test]
+ fn size_inference_rounds_up_then_caps_and_row_target_takes_precedence() {
+ let io = test_file_io();
+ let table = test_postpone_pk_table(&io,
"memory:/auto_postpone_size").copy_with_options(
+ HashMap::from([
+ ("postpone.target-size-per-bucket".into(), "20 b".into()),
+ (
+ "postpone.batch-write-fixed-bucket.max-parallelism".into(),
+ "3".into(),
+ ),
+ ]),
+ );
+ let mut planner = PostponeBucketPlanner::new(&table).unwrap();
+ planner
+ .add(make_batch(vec![1, 2, 3], vec![10, 20, 30]))
+ .unwrap();
+ let (plan, _) = planner.finish();
+ assert_eq!(
+ plan.total_buckets(&crate::spec::EMPTY_SERIALIZED_ROW),
+ Some(3)
+ );
+ let table = table.copy_with_options(HashMap::from([
+ ("postpone.target-row-num-per-bucket".into(), "2".into()),
+ ("postpone.target-size-per-bucket".into(), "invalid".into()),
+ ]));
+ let mut planner = PostponeBucketPlanner::new(&table).unwrap();
+ planner
+ .add(make_batch(vec![1, 2, 3], vec![10, 20, 30]))
+ .unwrap();
+ assert_eq!(
+ planner
+ .finish()
+ .0
+ .total_buckets(&crate::spec::EMPTY_SERIALIZED_ROW),
+ Some(2)
+ );
+ }
+}
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_router.rs
b/crates/paimon/src/table/postpone_fixed_bucket_router.rs
index 56b2c2e2..fe116c7f 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_router.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_router.rs
@@ -43,6 +43,10 @@ pub(crate) struct PostponeFixedBucketRouter {
}
impl PostponeFixedBucketRouter {
+ pub(super) fn set_plan(&mut self, plan: PostponeBucketPlan) {
+ self.plan = plan;
+ }
+
pub(crate) fn new(table: &Table, plan: PostponeBucketPlan) -> Result<Self>
{
validate_postpone_fixed_bucket_table(table)?;
let schema = table.schema();
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
index 44565fc6..4f55a352 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
@@ -16,6 +16,7 @@
// under the License.
use super::postpone_bucket_plan::data_invalid;
+use super::postpone_bucket_planner::PostponeBucketPlanner;
use super::postpone_fixed_bucket_router::{
validate_postpone_fixed_bucket_table, PostponeFixedBucketRouter,
};
@@ -34,6 +35,7 @@ pub struct PostponeFixedBucketTableWrite {
overwrite: bool,
check_from_snapshot: Option<i64>,
prepare_started: bool,
+ planner: Option<PostponeBucketPlanner>,
}
impl PostponeFixedBucketTableWrite {
@@ -42,6 +44,11 @@ impl PostponeFixedBucketTableWrite {
factory: Option<std::sync::Arc<dyn crate::io::UriReaderFactory>>,
) -> Result<&mut Self> {
self.ensure_writable()?;
+ if self.check_from_snapshot.is_some() {
+ return Err(data_invalid(
+ "with_blob_uri_reader_factory must be called before any write
operation",
+ ));
+ }
self.inner.with_blob_uri_reader_factory(factory)?;
Ok(self)
}
@@ -49,12 +56,17 @@ impl PostponeFixedBucketTableWrite {
pub(crate) fn new(
table: &Table,
commit_user: String,
- plan: PostponeBucketPlan,
+ plan: Option<PostponeBucketPlan>,
load_existing_buckets: bool,
overwrite: bool,
resources: Option<ResourceContext>,
) -> Result<Self> {
validate_postpone_fixed_bucket_table(table)?;
+ let planner = if plan.is_none() {
+ Some(PostponeBucketPlanner::new(table)?)
+ } else {
+ None
+ };
let mut inner = TableWrite::new(table, commit_user)?;
if let Some(resources) = resources {
inner = inner.with_resources(resources);
@@ -69,10 +81,11 @@ impl PostponeFixedBucketTableWrite {
} else {
inner
},
- router: PostponeFixedBucketRouter::new(table, plan)?,
+ router: PostponeFixedBucketRouter::new(table,
plan.unwrap_or_default())?,
overwrite,
check_from_snapshot: None,
prepare_started: false,
+ planner,
})
}
@@ -94,13 +107,28 @@ impl PostponeFixedBucketTableWrite {
if self.check_from_snapshot.is_none() {
let snapshot_id = self.inner.pin_sequence_snapshot().await?;
if self.load_existing_buckets {
- self.router
- .load_existing_buckets(&self.table, snapshot_id)
- .await?;
+ if let Some(planner) = &mut self.planner {
+ planner.load_metadata(&self.table, snapshot_id).await?;
+ self.router.set_plan(planner.known_plan());
+ } else {
+ self.router
+ .load_existing_buckets(&self.table, snapshot_id)
+ .await?;
+ }
}
self.check_from_snapshot = Some(snapshot_id);
}
- for routed in self.router.route(&batch)? {
+ if let Some(planner) = &mut self.planner {
+ if let Some(known) = planner.add(batch)? {
+ return self.write_routed_batch(&known).await;
+ }
+ return Ok(());
+ }
+ self.write_routed_batch(&batch).await
+ }
+
+ async fn write_routed_batch(&mut self, batch: &RecordBatch) -> Result<()> {
+ for routed in self.router.route(batch)? {
self.inner
.write_partition_bucket_batch(routed.partition, routed.bucket,
routed.batch)
.await?;
@@ -118,6 +146,22 @@ impl PostponeFixedBucketTableWrite {
pub async fn prepare_commit(&mut self) -> Result<Vec<CommitMessage>> {
self.ensure_writable()?;
self.prepare_started = true;
+ let result = self.prepare_commit_inner().await;
+ if result.is_err() {
+ self.failed = true;
+ self.inner.close().await;
+ }
+ result
+ }
+
+ async fn prepare_commit_inner(&mut self) -> Result<Vec<CommitMessage>> {
+ if let Some(planner) = self.planner.take() {
+ let (plan, batches) = planner.finish();
+ self.router.set_plan(plan);
+ for batch in batches {
+ self.write_routed_batch(&batch).await?;
+ }
+ }
let mut messages = self.inner.prepare_commit().await?;
for message in &mut messages {
message.total_buckets =
self.router.total_buckets(&message.partition);
@@ -132,6 +176,7 @@ impl PostponeFixedBucketTableWrite {
/// Discard unprepared files. Prepared messages and their files belong to
the caller.
pub async fn close(&mut self) {
self.closed = true;
+ self.planner = None;
self.inner.close().await;
}
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 530fa496..661f3207 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
-use super::postpone_bucket_plan::data_invalid;
use super::postpone_fixed_bucket_router::validate_postpone_fixed_bucket_table;
use crate::resource::ResourceContext;
use crate::spec::CoreOptions;
@@ -85,11 +84,10 @@ impl<'a> PostponeFixedBucketWriteBuilder<'a> {
ensure_table_write_allowed(self.table)?;
let load_existing_buckets = self.bucket_plan.is_none() &&
!self.overwrite;
let plan = match &self.bucket_plan {
- Some(plan) => plan.clone(),
- None => PostponeBucketPlan::with_default_bucket_num(
-
CoreOptions::new(self.table.schema().options()).postpone_default_bucket_num()?
- .ok_or_else(|| data_invalid("A resolved postpone bucket
plan is required unless postpone.default-bucket-num is configured"))?,
- ),
+ Some(plan) => Some(plan.clone()),
+ None => CoreOptions::new(self.table.schema().options())
+ .postpone_default_bucket_num()?
+ .map(PostponeBucketPlan::with_default_bucket_num),
};
PostponeFixedBucketTableWrite::new(
self.table,
@@ -132,13 +130,11 @@ mod tests {
setup_dirs(&file_io, table_path).await;
let table = test_postpone_pk_table(&file_io, table_path);
- let error = table
+ table
.new_postpone_fixed_bucket_write_builder()
.unwrap()
.new_write()
- .err()
.unwrap();
- assert!(error.to_string().contains("bucket plan is required"));
let builder = table
.new_postpone_fixed_bucket_write_builder()
@@ -200,4 +196,37 @@ mod tests {
.all(|message| message.total_buckets == Some(3) &&
(0..3).contains(&message.bucket)));
builder.new_commit().commit(messages).await.unwrap();
}
+
+ #[tokio::test]
+ async fn automatic_bucket_count_uses_the_whole_input() {
+ let io = test_file_io();
+ let path = "memory:/automatic_postpone_count";
+ setup_dirs(&io, path).await;
+ let table =
+ test_postpone_pk_table(&io,
path).copy_with_options(std::collections::HashMap::from([
+ ("postpone.target-row-num-per-bucket".into(), "2".into()),
+ (
+ "postpone.batch-write-fixed-bucket.max-parallelism".into(),
+ "8".into(),
+ ),
+ ]));
+ let builder = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1, 2], vec![10, 20]))
+ .await
+ .unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![3, 4, 5], vec![30, 40, 50]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(!messages.is_empty());
+ assert!(messages.iter().all(|m| m.total_buckets == Some(4)));
+ builder.new_commit().commit(messages).await.unwrap();
+ assert_eq!(
+ super::super::table_write::tests::read_id_value_rows(&table).await,
+ vec![(1, 10), (2, 20), (3, 30), (4, 40), (5, 50)]
+ );
+ }
}
