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 66b14a0d Extend native cross-partition and postpone writes (#1017)
66b14a0d is described below
commit 66b14a0d74c12f98c4ee5ccb4a2239797be8c456
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Oct 4 08:39:55 2026 +0800
Extend native cross-partition and postpone writes (#1017)
---
bindings/python/README.md | 34 ++
.../python/python/pypaimon_rust/datafusion.pyi | 21 ++
bindings/python/src/context.rs | 1 +
bindings/python/src/table.rs | 8 +-
bindings/python/src/write.rs | 133 +++++++-
crates/paimon/src/spec/core_options.rs | 12 +
crates/paimon/src/table/bucket_assigner_cross.rs | 246 ++++++++-------
crates/paimon/src/table/postpone_bucket_plan.rs | 54 +++-
crates/paimon/src/table/postpone_file_writer.rs | 120 ++++---
.../src/table/postpone_fixed_bucket_router.rs | 8 +
.../src/table/postpone_fixed_bucket_write.rs | 346 ++++++++++++++++++++-
.../table/postpone_fixed_bucket_write_builder.rs | 36 ++-
crates/paimon/src/table/table_commit.rs | 35 ---
crates/paimon/src/table/table_write.rs | 173 ++++++-----
crates/paimon/tests/cross_partition_write_test.rs | 323 +++++++++++++++++++
crates/paimon/tests/postpone_native_write_test.rs | 137 ++++++++
16 files changed, 1397 insertions(+), 290 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 496aecc9..56613434 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -139,6 +139,40 @@ committer. Batch empty commits follow
`snapshot.ignore-empty-commit` (default
true). `truncate_table()` shares the batch commit guard; `truncate_partitions`
accepts a nonempty list of partition specs as in Java.
+For fixed-bucket writes to a primary-key table with `bucket=-2`, use the
+specialized builder. Routing, merge functions, sequence initialization and
+failure cleanup all run in the core writer:
+
+```python
+import pyarrow as pa
+
+plan = pa.record_batch(
+ [pa.array(["2026-10-01", "2026-10-02"]), pa.array([3, 5],
type=pa.int32())],
+ names=["dt", "total_buckets"],
+)
+builder =
table.new_postpone_fixed_bucket_write_builder().with_bucket_plan(plan)
+writer = builder.new_write()
+try:
+ writer.write_arrow(batch)
+ messages = writer.prepare_commit()
+finally:
+ writer.close()
+builder.new_commit().commit(messages)
+```
+
+The plan contains partition columns in table order and an `int32` count.
+Every input partition must be present. Workers may share a plan, but all rows
+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.
+A failed write cannot be prepared, and close deletes only unprepared output.
+
Configure overwrite on the batch builder:
```python
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index cb707319..5b94ab35 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -186,6 +186,7 @@ class Table:
def schema(self) -> TableSchema: ...
def new_read_builder(self, options: Optional[Dict[str, str]] = None) ->
ReadBuilder: ...
def new_batch_write_builder(self) -> "BatchWriteBuilder": ...
+ def new_postpone_fixed_bucket_write_builder(self) ->
"PostponeFixedBucketWriteBuilder": ...
def new_stream_write_builder(self) -> "StreamWriteBuilder": ...
def latest_snapshot(self) -> Optional[Snapshot]:
"""
@@ -336,6 +337,26 @@ class BatchWriteBuilder:
def new_update(self) -> BatchTableUpdate: ...
def new_commit(self) -> BatchTableCommit: ...
+class PostponeFixedBucketWriteBuilder:
+ """Real-bucket batch writes for a bucket=-2 primary-key table.
+
+ Supply a shared Arrow plan or configure postpone.default-bucket-num.
+ 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.
+ """
+ def _with_commit_user(self, commit_user: str) ->
"PostponeFixedBucketWriteBuilder": ...
+ def with_bucket_plan(self, plan: pyarrow.RecordBatch) ->
"PostponeFixedBucketWriteBuilder":
+ """Partition columns in schema order followed by total_buckets: int32.
+
+ Counts must be positive and non-null. An explicit plan must contain
+ every input partition; the table default does not fill missing entries.
+ """
+ ...
+ def with_overwrite(self, static_partition: Optional[Dict[str, Any]] = {})
-> "PostponeFixedBucketWriteBuilder": ...
+ def new_write(self) -> BatchTableWrite: ...
+ def new_commit(self) -> BatchTableCommit: ...
+
class StreamWriteBuilder:
def commit_user(self) -> str: ...
def with_commit_user(self, commit_user: str) -> "StreamWriteBuilder": ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index c0628460..b661f14b 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -675,6 +675,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<PyPythonScalarUDFObject>()?;
this.add_class::<PySQLContext>()?;
this.add_class::<crate::write::PyBatchWriteBuilder>()?;
+ this.add_class::<crate::write::PyPostponeFixedBucketWriteBuilder>()?;
this.add_class::<crate::write::PyStreamWriteBuilder>()?;
this.add_class::<crate::write::PyBatchTableWrite>()?;
this.add_class::<crate::write::PyBatchTableUpdate>()?;
diff --git a/bindings/python/src/table.rs b/bindings/python/src/table.rs
index aa474a4d..95e1d82f 100644
--- a/bindings/python/src/table.rs
+++ b/bindings/python/src/table.rs
@@ -34,7 +34,7 @@ use crate::read::PyReadBuilder;
use crate::schema::PyTableSchema;
use crate::snapshot::PySnapshot;
use crate::tag::PyTag;
-use crate::write::{PyBatchWriteBuilder, PyStreamWriteBuilder};
+use crate::write::{PyBatchWriteBuilder, PyPostponeFixedBucketWriteBuilder,
PyStreamWriteBuilder};
#[pyclass(name = "Table", module = "pypaimon_rust.datafusion")]
pub struct PyTable {
@@ -175,6 +175,12 @@ impl PyTable {
PyBatchWriteBuilder::new(Arc::clone(&self.inner))
}
+ fn new_postpone_fixed_bucket_write_builder(
+ &self,
+ ) -> PyResult<PyPostponeFixedBucketWriteBuilder> {
+ PyPostponeFixedBucketWriteBuilder::new(Arc::clone(&self.inner))
+ }
+
fn new_stream_write_builder(&self) -> PyStreamWriteBuilder {
PyStreamWriteBuilder::new(Arc::clone(&self.inner))
}
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index bb02e891..6c619e65 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -23,8 +23,8 @@ use arrow::pyarrow::FromPyArrow;
use arrow::record_batch::RecordBatch;
use paimon::spec::{CoreOptions, DataType, Datum};
use paimon::table::{
- CommitMessage, Table, TableCommit, TableUpdate, TableUpdateByRowId,
TableWrite,
- COMMIT_MESSAGE_SERIALIZER_VERSION,
+ CommitMessage, PostponeBucketPlan, PostponeFixedBucketTableWrite, Table,
TableCommit,
+ TableUpdate, TableUpdateByRowId, TableWrite,
COMMIT_MESSAGE_SERIALIZER_VERSION,
};
use paimon_datafusion::runtime::runtime;
use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
@@ -61,7 +61,9 @@ impl WriteContext {
builder
};
Ok(WriteState {
- inner: Some(builder.new_write().map_err(to_py_err)?),
+ inner: Some(WriteTarget::Table(Box::new(
+ builder.new_write().map_err(to_py_err)?,
+ ))),
table_location: self.table.location().to_string(),
commit_user: self.commit_user.clone(),
})
@@ -270,8 +272,131 @@ impl PyStreamWriteBuilder {
}
}
+/// Keep conversion, ownership and commit-message wrapping shared by both
writers.
+/// All routing and validation remain in the Rust core implementations.
+enum WriteTarget {
+ Table(Box<TableWrite>),
+ PostponeFixed(Box<PostponeFixedBucketTableWrite>),
+}
+
+impl WriteTarget {
+ async fn write_arrow_batch(&mut self, batch: &RecordBatch) ->
paimon::Result<()> {
+ match self {
+ Self::Table(writer) => writer.write_arrow_batch(batch).await,
+ Self::PostponeFixed(writer) =>
writer.write_arrow_batch(batch).await,
+ }
+ }
+
+ async fn prepare_commit(&mut self) -> paimon::Result<Vec<CommitMessage>> {
+ match self {
+ Self::Table(writer) => writer.prepare_commit().await,
+ Self::PostponeFixed(writer) => writer.prepare_commit().await,
+ }
+ }
+
+ async fn close(&mut self) {
+ match self {
+ Self::Table(writer) => writer.close().await,
+ Self::PostponeFixed(writer) => writer.close().await,
+ }
+ }
+}
+
+/// The public batch contract is shared with PyPaimon's specialized builder.
+#[pyclass(
+ name = "PostponeFixedBucketWriteBuilder",
+ module = "pypaimon_rust.datafusion"
+)]
+pub struct PyPostponeFixedBucketWriteBuilder {
+ batch: PyBatchWriteBuilder,
+ plan: Option<PostponeBucketPlan>,
+}
+
+impl PyPostponeFixedBucketWriteBuilder {
+ pub fn new(table: Arc<Table>) -> PyResult<Self> {
+ table
+ .new_postpone_fixed_bucket_write_builder()
+ .map_err(to_py_err)?;
+ Ok(Self {
+ batch: PyBatchWriteBuilder::new(table),
+ plan: None,
+ })
+ }
+}
+
+#[pymethods]
+impl PyPostponeFixedBucketWriteBuilder {
+ fn _with_commit_user(
+ mut slf: PyRefMut<'_, Self>,
+ commit_user: String,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ slf.batch
+ .context
+ .table
+ .new_write_builder()
+ .with_commit_user(commit_user.clone())
+ .map_err(to_py_err)?;
+ slf.batch.context.commit_user = commit_user;
+ Ok(slf)
+ }
+
+ /// Partition fields followed by a non-null Int32 total_buckets column.
+ fn with_bucket_plan<'py>(
+ mut slf: PyRefMut<'py, Self>,
+ plan: &Bound<'py, PyAny>,
+ ) -> PyResult<PyRefMut<'py, Self>> {
+ let plan = RecordBatch::from_pyarrow_bound(plan)?;
+ slf.plan = Some(
+ PostponeBucketPlan::from_arrow(&slf.batch.context.table,
&plan).map_err(to_py_err)?,
+ );
+ Ok(slf)
+ }
+
+ #[pyo3(signature = (static_partition=Some(HashMap::new())))]
+ fn with_overwrite<'py>(
+ mut slf: PyRefMut<'py, Self>,
+ py: Python<'py>,
+ static_partition: Option<PythonPartitionSpec>,
+ ) -> PyResult<PyRefMut<'py, Self>> {
+ slf.batch.static_partition = static_partition
+ .map(|spec| partition_spec(py, &slf.batch.context.table, spec))
+ .transpose()?;
+ Ok(slf)
+ }
+
+ fn new_write(&self) -> PyResult<PyBatchTableWrite> {
+ let context = &self.batch.context;
+ let mut builder = context
+ .table
+ .new_postpone_fixed_bucket_write_builder()
+ .map_err(to_py_err)?
+ .with_commit_user(context.commit_user.clone())
+ .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();
+ }
+ Ok(PyBatchTableWrite {
+ state: WriteState {
+ inner: Some(WriteTarget::PostponeFixed(Box::new(
+ builder.new_write().map_err(to_py_err)?,
+ ))),
+ table_location: context.table.location().to_string(),
+ commit_user: context.commit_user.clone(),
+ },
+ prepared: false,
+ })
+ }
+
+ fn new_commit(&self) -> PyResult<PyBatchTableCommit> {
+ self.batch.new_commit()
+ }
+}
+
struct WriteState {
- inner: Option<TableWrite>,
+ inner: Option<WriteTarget>,
table_location: String,
commit_user: String,
}
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index f842ed2e..8aabb1d3 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -1417,6 +1417,18 @@ impl<'a> CoreOptions<'a> {
/// Java leaves `target-file-size` without a default and documents 128 MB
for
/// primary-key tables and 256 MB for append tables; this returns the
append
/// value for both.
+ /// Java's exact default for postpone partitions without an existing
layout.
+ /// Unlike inferred counts this is not capped or rounded to a power of two.
+ pub fn postpone_default_bucket_num(&self) -> crate::Result<Option<i32>> {
+ self.options.get("postpone.default-bucket-num").map(|raw| {
+ raw.parse::<i32>().ok().filter(|count| *count > 0).ok_or_else(|| {
+ crate::Error::ConfigInvalid {
+ message: format!("postpone.default-bucket-num must be a
positive 32-bit integer, got '{raw}'"),
+ }
+ })
+ }).transpose()
+ }
+
pub fn target_file_size(&self) -> i64 {
self.options
.get("target-file-size")
diff --git a/crates/paimon/src/table/bucket_assigner_cross.rs
b/crates/paimon/src/table/bucket_assigner_cross.rs
index 4be9800e..5abb9c33 100644
--- a/crates/paimon/src/table/bucket_assigner_cross.rs
+++ b/crates/paimon/src/table/bucket_assigner_cross.rs
@@ -26,7 +26,7 @@ use crate::table::Table;
use crate::Result;
use arrow_array::RecordBatch;
use futures::TryStreamExt;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap};
/// Result of assigning a bucket for a key in cross-partition mode.
enum AssignResult {
@@ -51,19 +51,23 @@ enum AssignResult {
struct GlobalPartitionIndex {
/// pk_bytes -> (partition_bytes, bucket)
key_to_location: HashMap<Vec<u8>, (Vec<u8>, i32)>,
- /// partition -> { bucket -> row_count } (only non-full buckets)
- partition_non_full_buckets: HashMap<Vec<u8>, HashMap<i32, i64>>,
- /// partition -> all known bucket ids
- partition_all_buckets: HashMap<Vec<u8>, HashSet<i32>>,
- /// partition -> next bucket id to allocate (avoids linear scan)
- partition_next_bucket_id: HashMap<Vec<u8>, i32>,
- /// (partition, bucket) -> total row count
- bucket_row_counts: HashMap<(Vec<u8>, i32), i64>,
+ /// Java BucketAssigner uses a TreeMap: reuse the first non-full bucket,
+ /// then allocate the smallest unused non-negative ID.
+ partition_bucket_counts: HashMap<Vec<u8>, BTreeMap<i32, i64>>,
target_bucket_row_number: i64,
merge_engine: MergeEngine,
}
impl GlobalPartitionIndex {
+ fn empty(target_bucket_row_number: i64, merge_engine: MergeEngine) -> Self
{
+ Self {
+ key_to_location: HashMap::new(),
+ partition_bucket_counts: HashMap::new(),
+ target_bucket_row_number,
+ merge_engine,
+ }
+ }
+
/// Build the global partition index by scanning all data files from the
latest snapshot.
///
/// Uses TableRead to get deduplicated PK rows per split, so DELETE records
@@ -78,7 +82,7 @@ impl GlobalPartitionIndex {
crate::spec::CoreOptions::new(table.schema().options()).ensure_read_authorized()?;
let mut key_to_location: HashMap<Vec<u8>, (Vec<u8>, i32)> =
HashMap::new();
- let mut bucket_row_counts: HashMap<(Vec<u8>, i32), i64> =
HashMap::new();
+ let mut partition_bucket_counts: HashMap<Vec<u8>, BTreeMap<i32, i64>>
= HashMap::new();
let fields = table.schema().fields();
let pk_field_names: Vec<&str> = primary_key_indices
@@ -101,55 +105,33 @@ impl GlobalPartitionIndex {
let partition_bytes = split.partition().to_serialized_bytes();
let bucket = split.bucket();
- let batches: Vec<RecordBatch> = read
- .to_arrow(std::slice::from_ref(split))?
- .try_collect()
- .await?;
-
- let pb_key = (partition_bytes.clone(), bucket);
- for batch in &batches {
+ let mut batches = read.to_arrow(std::slice::from_ref(split))?;
+ while let Some(batch) = batches.try_next().await? {
let pk_bytes_vec =
- batch_to_serialized_bytes(batch, &projected_pk_indices,
&pk_fields)?;
- let count = pk_bytes_vec.len() as i64;
+ batch_to_serialized_bytes(&batch, &projected_pk_indices,
&pk_fields)?;
for pk_bytes in pk_bytes_vec {
- key_to_location.insert(pk_bytes, (partition_bytes.clone(),
bucket));
+ if key_to_location
+ .insert(pk_bytes, (partition_bytes.clone(), bucket))
+ .is_some()
+ {
+ return Err(crate::Error::DataInvalid {
+ message: "Duplicate primary key found while
bootstrapping the cross-partition index".into(),
+ source: None,
+ });
+ }
+ *partition_bucket_counts
+ .entry(partition_bytes.clone())
+ .or_default()
+ .entry(bucket)
+ .or_default() += 1;
}
- *bucket_row_counts.entry(pb_key.clone()).or_insert(0) += count;
}
}
- let mut partition_all_buckets: HashMap<Vec<u8>, HashSet<i32>> =
HashMap::new();
- let mut partition_non_full_buckets: HashMap<Vec<u8>, HashMap<i32,
i64>> = HashMap::new();
- for ((partition, bucket), count) in &bucket_row_counts {
- partition_all_buckets
- .entry(partition.clone())
- .or_default()
- .insert(*bucket);
- if *count < target_bucket_row_number {
- partition_non_full_buckets
- .entry(partition.clone())
- .or_default()
- .insert(*bucket, *count);
- }
- }
-
- let partition_next_bucket_id: HashMap<Vec<u8>, i32> =
partition_all_buckets
- .iter()
- .map(|(p, buckets)| {
- let next = buckets.iter().copied().max().map_or(0, |m| m + 1);
- (p.clone(), next)
- })
- .collect();
-
- Ok(Self {
- key_to_location,
- partition_non_full_buckets,
- partition_all_buckets,
- partition_next_bucket_id,
- bucket_row_counts,
- target_bucket_row_number,
- merge_engine,
- })
+ let mut index = Self::empty(target_bucket_row_number, merge_engine);
+ index.key_to_location = key_to_location;
+ index.partition_bucket_counts = partition_bucket_counts;
+ Ok(index)
}
/// Assign a bucket for the given primary key targeting `new_partition`.
@@ -171,19 +153,14 @@ impl GlobalPartitionIndex {
let old_partition = existing_partition.clone();
let old_bucket = *existing_bucket;
- // Decrement row count for old bucket
- let old_key = (old_partition.clone(), old_bucket);
- let old_count =
self.bucket_row_counts.get(&old_key).copied().unwrap_or(0);
- let new_count = (old_count - 1).max(0);
- self.bucket_row_counts.insert(old_key, new_count);
- if new_count < self.target_bucket_row_number {
- self.partition_non_full_buckets
- .entry(old_partition.clone())
- .or_default()
- .insert(old_bucket, new_count);
- }
-
- let new_bucket =
self.assign_bucket_in_partition(new_partition);
+ let old_count = self
+ .partition_bucket_counts
+ .entry(old_partition.clone())
+ .or_default()
+ .entry(old_bucket)
+ .or_default();
+ *old_count -= 1;
+ let new_bucket =
self.assign_bucket_in_partition(new_partition)?;
self.key_to_location
.insert(pk_bytes.to_vec(), (new_partition.to_vec(),
new_bucket));
@@ -202,60 +179,34 @@ impl GlobalPartitionIndex {
}
}
- let bucket = self.assign_bucket_in_partition(new_partition);
+ let bucket = self.assign_bucket_in_partition(new_partition)?;
self.key_to_location
.insert(pk_bytes.to_vec(), (new_partition.to_vec(), bucket));
Ok(AssignResult::SamePartition { bucket })
}
- fn assign_bucket_in_partition(&mut self, partition: &[u8]) -> i32 {
- let non_full = self
- .partition_non_full_buckets
+ fn assign_bucket_in_partition(&mut self, partition: &[u8]) -> Result<i32> {
+ let buckets = self
+ .partition_bucket_counts
.entry(partition.to_vec())
.or_default();
-
- let mut full_buckets = Vec::new();
- let mut assigned_bucket = None;
- for (&bucket, count) in non_full.iter_mut() {
+ for (&bucket, count) in buckets.iter_mut() {
if *count < self.target_bucket_row_number {
*count += 1;
- assigned_bucket = Some(bucket);
- break;
- } else {
- full_buckets.push(bucket);
+ return Ok(bucket);
}
}
- for b in full_buckets {
- non_full.remove(&b);
+ let mut next = 0;
+ while buckets.contains_key(&next) {
+ next = next
+ .checked_add(1)
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "No bucket IDs available for cross-partition
writes".into(),
+ source: None,
+ })?;
}
-
- let bucket = if let Some(b) = assigned_bucket {
- b
- } else {
- let all = self
- .partition_all_buckets
- .entry(partition.to_vec())
- .or_default();
- let next_id = self
- .partition_next_bucket_id
- .entry(partition.to_vec())
- .or_insert(0);
- let new_bucket = *next_id;
- *next_id += 1;
- all.insert(new_bucket);
- self.partition_non_full_buckets
- .entry(partition.to_vec())
- .or_default()
- .insert(new_bucket, 1);
- new_bucket
- };
-
- *self
- .bucket_row_counts
- .entry((partition.to_vec(), bucket))
- .or_insert(0) += 1;
-
- bucket
+ buckets.insert(next, 1);
+ Ok(next)
}
}
@@ -280,15 +231,41 @@ impl CrossPartitionAssigner {
primary_key_indices: Vec<usize>,
target_bucket_row_number: i64,
merge_engine: MergeEngine,
- ) -> Self {
- Self {
+ ) -> Result<Self> {
+ if !crate::spec::CoreOptions::new(table.schema().options())
+ .sequence_fields()
+ .is_empty()
+ {
+ return Err(crate::Error::DataInvalid {
+ message: "Cannot define 'sequence.field' for cross-partition
update tables".into(),
+ source: None,
+ });
+ }
+ if table.schema().options().contains_key("bucket-key") {
+ return Err(crate::Error::DataInvalid {
+ message: "Cannot define 'bucket-key' for dynamic bucket
tables".into(),
+ source: None,
+ });
+ }
+ if table
+ .schema()
+ .options()
+ .contains_key("cross-partition-upsert.index-ttl")
+ {
+ return Err(crate::Error::Unsupported {
+ message:
+ "Cross-partition writes do not support
'cross-partition-upsert.index-ttl' yet"
+ .into(),
+ });
+ }
+ Ok(Self {
table,
partition_field_indices,
primary_key_indices,
global_partition_index: None,
target_bucket_row_number,
merge_engine,
- }
+ })
}
}
@@ -359,3 +336,50 @@ impl BucketAssigner for CrossPartitionAssigner {
Ok(HashMap::new())
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn bucket_assignment_reuses_lowest_nonfull_bucket_then_first_unused_id() {
+ let mut index = GlobalPartitionIndex::empty(2,
MergeEngine::Deduplicate);
+ index
+ .partition_bucket_counts
+ .insert(vec![1], [(0, 2), (2, 1), (4, 1)].into());
+ assert_eq!(index.assign_bucket_in_partition(&[1]).unwrap(), 2);
+ assert_eq!(index.assign_bucket_in_partition(&[1]).unwrap(), 4);
+ assert_eq!(index.assign_bucket_in_partition(&[1]).unwrap(), 1);
+ assert_eq!(index.assign_bucket_in_partition(&[1]).unwrap(), 1);
+ assert_eq!(index.assign_bucket_in_partition(&[1]).unwrap(), 3);
+ }
+
+ #[test]
+ fn migration_releases_old_bucket_capacity_without_forgetting_other_keys() {
+ let mut index = GlobalPartitionIndex::empty(1,
MergeEngine::Deduplicate);
+ assert!(matches!(
+ index.assign(&[1], &[1]).unwrap(),
+ AssignResult::SamePartition { bucket: 0 }
+ ));
+ assert!(matches!(
+ index.assign(&[2], &[1]).unwrap(),
+ AssignResult::SamePartition { bucket: 1 }
+ ));
+ assert!(matches!(
+ index.assign(&[1], &[2]).unwrap(),
+ AssignResult::CrossPartition {
+ old_bucket: 0,
+ new_bucket: 0,
+ ..
+ }
+ ));
+ assert!(matches!(
+ index.assign(&[3], &[1]).unwrap(),
+ AssignResult::SamePartition { bucket: 0 }
+ ));
+ assert!(matches!(
+ index.assign(&[2], &[1]).unwrap(),
+ AssignResult::SamePartition { bucket: 1 }
+ ));
+ }
+}
diff --git a/crates/paimon/src/table/postpone_bucket_plan.rs
b/crates/paimon/src/table/postpone_bucket_plan.rs
index 6b6d867e..613ba6b4 100644
--- a/crates/paimon/src/table/postpone_bucket_plan.rs
+++ b/crates/paimon/src/table/postpone_bucket_plan.rs
@@ -16,7 +16,7 @@
// under the License.
use crate::spec::batch_to_serialized_bytes;
-use crate::table::Table;
+use crate::table::{SnapshotManager, Table, TableScan};
use crate::Result;
use arrow_array::{Array, Int32Array, RecordBatch};
use std::collections::HashMap;
@@ -34,6 +34,7 @@ pub(crate) fn data_invalid(message: impl Into<String>) ->
crate::Error {
#[derive(Debug, Clone, Default)]
pub struct PostponeBucketPlan {
bucket_counts: Arc<HashMap<Vec<u8>, i32>>,
+ default_bucket_num: Option<i32>,
}
impl PostponeBucketPlan {
@@ -113,10 +114,59 @@ impl PostponeBucketPlan {
}
Ok(Self {
bucket_counts: Arc::new(bucket_counts),
+ default_bucket_num: None,
})
}
+ pub(crate) fn with_default_bucket_num(default_bucket_num: i32) -> Self {
+ Self {
+ default_bucket_num: Some(default_bucket_num),
+ ..Self::default()
+ }
+ }
+
+ /// Read the same baseline used for sequence numbers and commit conflict
checks.
+ /// Pending (-2) files have no real layout and never override the
configured default.
+ pub(crate) async fn load_existing_buckets(
+ &mut self,
+ table: &Table,
+ snapshot_id: i64,
+ ) -> Result<()> {
+ if snapshot_id == 0 {
+ return Ok(());
+ }
+ let snapshot = SnapshotManager::new(table.file_io().clone(),
table.location().to_string())
+ .get_snapshot(snapshot_id)
+ .await?;
+ let entries = TableScan::new(table, None, vec![], None, None, None)
+ .with_scan_all_files()
+ .plan_manifest_entries(&snapshot)
+ .await?;
+ let mut counts = HashMap::new();
+ for entry in entries {
+ if entry.bucket() < 0 {
+ continue;
+ }
+ let count = entry.total_buckets();
+ if count <= 0 || entry.bucket() >= count {
+ return Err(data_invalid(
+ "Invalid total bucket count in existing postpone files",
+ ));
+ }
+ if let Some(previous) = counts.insert(entry.partition().to_vec(),
count) {
+ if previous != count {
+ return Err(data_invalid("Existing postpone files contain
conflicting total bucket counts for one partition"));
+ }
+ }
+ }
+ self.bucket_counts = Arc::new(counts);
+ Ok(())
+ }
+
pub(crate) fn total_buckets(&self, partition: &[u8]) -> Option<i32> {
- self.bucket_counts.get(partition).copied()
+ self.bucket_counts
+ .get(partition)
+ .copied()
+ .or(self.default_bucket_num)
}
}
diff --git a/crates/paimon/src/table/postpone_file_writer.rs
b/crates/paimon/src/table/postpone_file_writer.rs
index ea117501..e89235cb 100644
--- a/crates/paimon/src/table/postpone_file_writer.rs
+++ b/crates/paimon/src/table/postpone_file_writer.rs
@@ -30,7 +30,8 @@ use crate::io::FileIO;
use crate::resource::ResourceContext;
use crate::spec::stats::BinaryTableStats;
use crate::spec::{
- data_file_to_file_index_file_name, DataFileMeta, EMPTY_SERIALIZED_ROW,
VALUE_KIND_FIELD_NAME,
+ data_file_to_file_index_file_name, BinaryRow, DataField, DataFileMeta,
RowKind,
+ EMPTY_SERIALIZED_ROW, VALUE_KIND_FIELD_NAME,
};
use crate::table::data_file_index_writer::{DataFileIndexWriter,
FileIndexOptions};
use crate::table::kv_file_writer::build_physical_schema;
@@ -48,6 +49,9 @@ pub(crate) struct PostponeWriteConfig {
pub bucket: i32,
pub schema_id: i64,
pub target_file_size: i64,
+ pub target_file_row_num: i64,
+ pub primary_key_indices: Vec<usize>,
+ pub value_fields: Vec<DataField>,
pub file_compression: String,
pub file_compression_zstd_level: i32,
pub write_buffer_size: i64,
@@ -61,7 +65,7 @@ pub(crate) struct PostponeWriteConfig {
///
/// Streams data directly to a FormatFileWriter in arrival order (no
sort/dedup),
/// prepending `_SEQUENCE_NUMBER` and `_VALUE_KIND` columns to each batch.
-/// Rolls to a new file when `target_file_size` is reached.
+/// Rolls after a batch reaches the configured size or row limit.
pub(crate) struct PostponeFileWriter {
paths: DataFilePathFactory,
file_io: FileIO,
@@ -71,9 +75,7 @@ pub(crate) struct PostponeFileWriter {
current_index: Option<DataFileIndexWriter>,
current_file_name: Option<String>,
current_file_path: Option<DataFilePath>,
- current_row_count: i64,
- /// Sequence number at which the current file started.
- current_file_start_seq: i64,
+ current_stats: PostponeFileStats,
/// Timestamp captured when the current file was opened (used for
deterministic replay order).
current_file_creation_time: DateTime<Utc>,
written_files: Vec<DataFileMeta>,
@@ -100,8 +102,7 @@ impl PostponeFileWriter {
current_index: None,
current_file_name: None,
current_file_path: None,
- current_row_count: 0,
- current_file_start_seq: 0,
+ current_stats: PostponeFileStats::default(),
current_file_creation_time: Utc::now(),
written_files: Vec::new(),
created_paths: Vec::new(),
@@ -186,10 +187,12 @@ impl PostponeFileWriter {
index.write(&logical_batch)?;
}
self.next_sequence_number = end_seq + 1;
- self.current_row_count += num_rows as i64;
+ self.current_stats
+ .add_batch(batch, &self.config, start_seq, end_seq)?;
// Roll to a new file if target size is reached — close in background
if self.current_writer.as_ref().unwrap().num_bytes() as i64 >=
self.config.target_file_size
+ || self.current_stats.row_count >= self.config.target_file_row_num
{
self.roll_file();
}
@@ -235,6 +238,10 @@ impl PostponeFileWriter {
})??;
self.written_files.push(meta);
}
+ // Java replays postpone files by millisecond creation time, retaining
+ // manifest order for ties. Background closes must not reorder
arrivals.
+ self.written_files
+ .sort_by_key(|file| file.min_sequence_number);
self.created_paths.clear();
Ok(std::mem::take(&mut self.written_files))
}
@@ -255,10 +262,7 @@ impl PostponeFileWriter {
.file_index_options
.as_ref()
.map(|options| options.in_manifest_threshold);
- let row_count = self.current_row_count;
- let min_seq = self.current_file_start_seq;
- let max_seq = self.next_sequence_number - 1;
- self.current_row_count = 0;
+ let stats = std::mem::take(&mut self.current_stats);
let schema_id = self.config.schema_id;
// Capture creation_time from when the file was opened, not when the
async close finishes.
// Java's postpone compaction sorts by creationTime for replay order.
@@ -266,15 +270,7 @@ impl PostponeFileWriter {
self.in_flight_closes.spawn(async move {
let file_size = writer.close().await?.file_size as i64;
- let mut meta = build_meta(
- file_name,
- file_size,
- row_count,
- min_seq,
- max_seq,
- schema_id,
- creation_time,
- );
+ let mut meta = build_meta(file_name, file_size, stats, schema_id,
creation_time);
meta.external_path = location.external_path;
write_index(index, threshold, &file_io, &bucket_dir, &mut
meta).await?;
Ok(meta)
@@ -322,8 +318,7 @@ impl PostponeFileWriter {
self.current_index = index;
self.current_file_name = Some(file_name);
self.current_file_path = Some(location);
- self.current_row_count = 0;
- self.current_file_start_seq = self.next_sequence_number;
+ self.current_stats = PostponeFileStats::default();
self.current_file_creation_time = Utc::now();
Ok(())
}
@@ -335,19 +330,13 @@ impl PostponeFileWriter {
};
let file_name = self.current_file_name.take().unwrap();
let index = self.current_index.take();
- let row_count = self.current_row_count;
- self.current_row_count = 0;
+ let stats = std::mem::take(&mut self.current_stats);
let file_size = writer.close().await?.file_size as i64;
- let min_seq = self.current_file_start_seq;
- let max_seq = self.next_sequence_number - 1;
-
let mut meta = build_meta(
file_name,
file_size,
- row_count,
- min_seq,
- max_seq,
+ stats,
self.config.schema_id,
self.current_file_creation_time,
);
@@ -388,21 +377,72 @@ async fn write_index(
Ok(())
}
+/// Java KeyValueDataFileWriter retains the first/last keys even for unsorted
+/// postpone input. These bounds must still be decodable using the key schema.
+#[derive(Default)]
+struct PostponeFileStats {
+ row_count: i64,
+ delete_row_count: i64,
+ first_key: Vec<u8>,
+ last_key: Vec<u8>,
+ min_sequence_number: i64,
+ max_sequence_number: i64,
+}
+
+impl PostponeFileStats {
+ fn add_batch(
+ &mut self,
+ batch: &RecordBatch,
+ config: &PostponeWriteConfig,
+ min_seq: i64,
+ max_seq: i64,
+ ) -> Result<()> {
+ let key_at = |row| -> Result<Vec<u8>> {
+ Ok(BinaryRow::from_arrow(
+ batch,
+ row,
+ &config.primary_key_indices,
+ &config.value_fields,
+ )?
+ .to_serialized_bytes())
+ };
+ if self.row_count == 0 {
+ self.first_key = key_at(0)?;
+ self.min_sequence_number = min_seq;
+ }
+ self.last_key = key_at(batch.num_rows() - 1)?;
+ self.max_sequence_number = max_seq;
+ self.row_count += batch.num_rows() as i64;
+ if let Some(column) = batch.column_by_name(VALUE_KIND_FIELD_NAME) {
+ let kinds =
column.as_any().downcast_ref::<Int8Array>().ok_or_else(|| {
+ crate::Error::DataInvalid {
+ message: "_VALUE_KIND column must be Int8".into(),
+ source: None,
+ }
+ })?;
+ for kind in kinds.iter().flatten() {
+ if !RowKind::from_value(kind)?.is_add() {
+ self.delete_row_count += 1;
+ }
+ }
+ }
+ Ok(())
+ }
+}
+
fn build_meta(
file_name: String,
file_size: i64,
- row_count: i64,
- min_seq: i64,
- max_seq: i64,
+ stats: PostponeFileStats,
schema_id: i64,
creation_time: DateTime<Utc>,
) -> DataFileMeta {
DataFileMeta {
file_name,
file_size,
- row_count,
- min_key: EMPTY_SERIALIZED_ROW.clone(),
- max_key: EMPTY_SERIALIZED_ROW.clone(),
+ row_count: stats.row_count,
+ min_key: stats.first_key,
+ max_key: stats.last_key,
key_stats: BinaryTableStats::new(
EMPTY_SERIALIZED_ROW.clone(),
EMPTY_SERIALIZED_ROW.clone(),
@@ -413,13 +453,13 @@ fn build_meta(
EMPTY_SERIALIZED_ROW.clone(),
vec![],
),
- min_sequence_number: min_seq,
- max_sequence_number: max_seq,
+ min_sequence_number: stats.min_sequence_number,
+ max_sequence_number: stats.max_sequence_number,
schema_id,
level: 0,
extra_files: vec![],
creation_time: Some(creation_time),
- delete_row_count: Some(0),
+ delete_row_count: Some(stats.delete_row_count),
embedded_index: None,
file_source: Some(0), // FileSource.APPEND
value_stats_cols: Some(vec![]),
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_router.rs
b/crates/paimon/src/table/postpone_fixed_bucket_router.rs
index 27b9cb88..56b2c2e2 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_router.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_router.rs
@@ -65,6 +65,14 @@ impl PostponeFixedBucketRouter {
})
}
+ pub(crate) async fn load_existing_buckets(
+ &mut self,
+ table: &Table,
+ snapshot_id: i64,
+ ) -> Result<()> {
+ self.plan.load_existing_buckets(table, snapshot_id).await
+ }
+
pub(crate) fn route(&self, batch: &RecordBatch) ->
Result<Vec<PostponeBucketBatch>> {
let mut output = Vec::new();
for (partition, batch) in
diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
index beb20bcc..32ba67fe 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs
@@ -27,6 +27,10 @@ use arrow_array::RecordBatch;
pub struct PostponeFixedBucketTableWrite {
inner: TableWrite,
+ table: Table,
+ load_existing_buckets: bool,
+ failed: bool,
+ closed: bool,
router: PostponeFixedBucketRouter,
overwrite: bool,
check_from_snapshot: Option<i64>,
@@ -38,6 +42,7 @@ impl PostponeFixedBucketTableWrite {
table: &Table,
commit_user: String,
plan: PostponeBucketPlan,
+ load_existing_buckets: bool,
overwrite: bool,
resources: Option<ResourceContext>,
) -> Result<Self> {
@@ -47,6 +52,10 @@ impl PostponeFixedBucketTableWrite {
inner = inner.with_resources(resources);
}
Ok(Self {
+ table: table.clone(),
+ load_existing_buckets,
+ failed: false,
+ closed: false,
inner: if overwrite {
inner.with_overwrite()
} else {
@@ -61,11 +70,27 @@ impl PostponeFixedBucketTableWrite {
pub async fn write_arrow_batch(&mut self, batch: &RecordBatch) ->
Result<()> {
self.ensure_writable()?;
+ let result = self.write_batch(batch).await;
+ if result.is_err() {
+ // A rejected batch must never leave earlier batches committable.
+ self.failed = true;
+ self.inner.close().await;
+ }
+ result
+ }
+
+ async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
let Some(batch) = self.inner.normalize_write_batch(batch)? else {
return Ok(());
};
if self.check_from_snapshot.is_none() {
- self.check_from_snapshot =
Some(self.inner.pin_sequence_snapshot().await?);
+ let snapshot_id = self.inner.pin_sequence_snapshot().await?;
+ if self.load_existing_buckets {
+ self.router
+ .load_existing_buckets(&self.table, snapshot_id)
+ .await?;
+ }
+ self.check_from_snapshot = Some(snapshot_id);
}
for routed in self.router.route(&batch)? {
self.inner
@@ -96,7 +121,18 @@ impl PostponeFixedBucketTableWrite {
Ok(messages)
}
+ /// Discard unprepared files. Prepared messages and their files belong to
the caller.
+ pub async fn close(&mut self) {
+ self.closed = true;
+ self.inner.close().await;
+ }
+
fn ensure_writable(&self) -> Result<()> {
+ if self.closed || self.failed {
+ return Err(data_invalid(
+ "Fixed-bucket postpone TableWrite is closed or failed",
+ ));
+ }
if self.prepare_started {
return Err(data_invalid("Fixed-bucket postpone TableWrite only
supports one prepare_commit call; create a new writer for the next batch"));
}
@@ -675,8 +711,38 @@ mod tests {
.unwrap();
}
+ async fn assert_pending_and_fixed_files(table: &Table, pending_rows: i64,
fixed_rows: i64) {
+ let snapshot = table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .unwrap();
+ let entries = TableScan::new(table, None, vec![], None, None, None)
+ .with_scan_all_files()
+ .plan_manifest_entries(&snapshot)
+ .await
+ .unwrap();
+ assert_eq!(
+ entries
+ .iter()
+ .filter(|entry| entry.bucket() == -2)
+ .map(|entry| entry.file().row_count)
+ .sum::<i64>(),
+ pending_rows
+ );
+ assert_eq!(
+ entries
+ .iter()
+ .filter(|entry| entry.bucket() >= 0)
+ .map(|entry| entry.file().row_count)
+ .sum::<i64>(),
+ fixed_rows
+ );
+ }
+
#[tokio::test]
- async fn test_postpone_rejects_negative_bucket_after_fixed_layout() {
+ async fn test_postpone_preserves_negative_bucket_after_fixed_layout() {
let file_io = test_file_io();
let table_path =
"memory:/test_postpone_rejects_negative_bucket_after_fixed";
setup_dirs(&file_io, table_path).await;
@@ -700,17 +766,13 @@ mod tests {
.unwrap();
let messages = normal_write.prepare_commit().await.unwrap();
assert!(messages.iter().all(|message| message.bucket == -2));
- let error = normal_builder
- .new_commit()
- .commit(messages)
- .await
- .unwrap_err();
- assert!(error.to_string().contains("Cannot mix bucket=-2 files"));
+ normal_builder.new_commit().commit(messages).await.unwrap();
+ assert_pending_and_fixed_files(&table, 1, 1).await;
assert_eq!(read_id_value_rows(&table).await, vec![(1, 10)]);
}
#[tokio::test]
- async fn test_postpone_fixed_append_rejects_existing_negative_bucket() {
+ async fn test_postpone_fixed_append_preserves_existing_negative_bucket() {
let file_io = test_file_io();
let table_path =
"memory:/test_postpone_fixed_rejects_existing_negative_bucket";
setup_dirs(&file_io, table_path).await;
@@ -733,11 +795,12 @@ mod tests {
let fixed_messages =
write_fixed_batch(&table, "fixed-append", 1, &make_batch(vec![1],
vec![100])).await;
- let error = TableCommit::new(table.clone(), "fixed-append".to_string())
+ TableCommit::new(table.clone(), "fixed-append".to_string())
.commit(fixed_messages)
.await
- .unwrap_err();
- assert!(error.to_string().contains("Cannot mix bucket=-2 files"));
+ .unwrap();
+ assert_pending_and_fixed_files(&table, 4, 1).await;
+ assert_eq!(read_id_value_rows(&table).await, vec![(1, 100)]);
let overwrite_builder = table
.new_postpone_fixed_bucket_write_builder()
@@ -829,7 +892,7 @@ mod tests {
}
#[tokio::test]
- async fn test_postpone_fixed_commit_rejects_concurrent_negative_bucket() {
+ async fn test_postpone_fixed_commit_preserves_concurrent_negative_bucket()
{
let file_io = test_file_io();
let table_path =
"memory:/test_postpone_fixed_rejects_concurrent_negative_bucket";
setup_dirs(&file_io, table_path).await;
@@ -863,12 +926,13 @@ mod tests {
.await
.unwrap();
- let error = fixed_builder
+ fixed_builder
.new_commit()
.commit(fixed_messages)
.await
- .unwrap_err();
- assert!(error.to_string().contains("Cannot mix bucket=-2 files"));
+ .unwrap();
+ assert_pending_and_fixed_files(&table, 1, 1).await;
+ assert_eq!(read_id_value_rows(&table).await, vec![(1, 10)]);
}
#[tokio::test]
@@ -1008,4 +1072,254 @@ mod tests {
assert_eq!(messages[0].new_files[0].row_count, 1);
assert_eq!(messages[0].new_files[0].delete_row_count, Some(1));
}
+ fn with_default_count(table: &Table, count: &str) -> Table {
+ table.copy_with_options(std::collections::HashMap::from([(
+ "postpone.default-bucket-num".into(),
+ count.into(),
+ )]))
+ }
+
+ #[tokio::test]
+ async fn default_count_reuses_committed_layout_but_overwrite_replaces_it()
{
+ let io = test_file_io();
+ let path = "memory:/fixed_default_reuse";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_pk_table(&io, path),
"3");
+ for (count, overwrite, value, expected_count) in
+ [("3", false, 10, 3), ("5", false, 20, 3), ("5", true, 30, 5)]
+ {
+ let table = with_default_count(&table, count);
+ let mut builder =
table.new_postpone_fixed_bucket_write_builder().unwrap();
+ if overwrite {
+ builder = builder.with_overwrite();
+ }
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1], vec![value]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.total_buckets == Some(expected_count)));
+ builder.new_commit().commit(messages).await.unwrap();
+ writer.close().await;
+ assert_eq!(read_id_value_rows(&table).await, vec![(1, value)]);
+ }
+ }
+
+ #[tokio::test]
+ async fn default_count_uses_one_snapshot_for_all_input_partitions() {
+ let io = test_file_io();
+ let path = "memory:/fixed_default_snapshot";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_partitioned_table(&io,
path), "3");
+ let builder = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_partitioned_batch_3col(vec!["a"],
vec![1], vec![10]))
+ .await
+ .unwrap();
+
+ // A newer snapshot must not change the plan of a writer that already
started.
+ let changed = with_default_count(&table, "5");
+ let concurrent =
changed.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut concurrent_writer = concurrent.new_write().unwrap();
+ concurrent_writer
+ .write_arrow_batch(&make_partitioned_batch_3col(vec!["b"],
vec![2], vec![20]))
+ .await
+ .unwrap();
+ concurrent
+ .new_commit()
+ .commit(concurrent_writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+
+ writer
+ .write_arrow_batch(&make_partitioned_batch_3col(vec!["b"],
vec![3], vec![30]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.total_buckets == Some(3)
+ && message.check_from_snapshot == Some(0)));
+ let error = builder
+ .new_commit()
+ .commit(messages.clone())
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("different total bucket counts"));
+ builder.new_commit().abort(&messages).await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn fixed_plan_failure_poisoning_discards_earlier_batches() {
+ let io = test_file_io();
+ let path = "memory:/fixed_plan_failure";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_partitioned_table(&io,
path), "3");
+ let builder = table
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .with_bucket_plan(make_partition_bucket_plan(&table, vec!["a"],
vec![2]));
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_partitioned_batch_3col(vec!["a"],
vec![1], vec![10]))
+ .await
+ .unwrap();
+ let error = writer
+ .write_arrow_batch(&make_partitioned_batch_3col(
+ vec!["missing"],
+ vec![2],
+ vec![20],
+ ))
+ .await
+ .unwrap_err();
+ assert!(error
+ .to_string()
+ .contains("does not contain an input partition"));
+ assert!(writer
+ .prepare_commit()
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("closed or failed"));
+ assert!(writer
+ .write_arrow_batch(&make_partitioned_batch_3col(vec!["a"],
vec![3], vec![30]))
+ .await
+ .is_err());
+ writer.close().await;
+ assert!(table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .is_none());
+ }
+
+ #[tokio::test]
+ async fn fixed_writer_close_preserves_only_prepared_output() {
+ let io = test_file_io();
+ let path = "memory:/fixed_close_ownership";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_pk_table(&io, path),
"1");
+ let builder = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut discarded = builder.new_write().unwrap();
+ discarded
+ .write_arrow_batch(&make_batch(vec![1], vec![10]))
+ .await
+ .unwrap();
+ discarded.close().await;
+ discarded.close().await;
+ assert!(discarded
+ .prepare_commit()
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("closed or failed"));
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![2], vec![20]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ writer.close().await;
+ builder.new_commit().commit(messages).await.unwrap();
+ assert_eq!(read_id_value_rows(&table).await, vec![(2, 20)]);
+ }
+
+ #[tokio::test]
+ async fn default_count_ignores_pending_files_and_inference_options() {
+ let io = test_file_io();
+ let path = "memory:/fixed_default_pending";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_pk_table(&io, path),
"3").copy_with_options(
+ std::collections::HashMap::from([
+ ("postpone.target-row-num-per-bucket".into(), "0".into()),
+ ("postpone.target-size-per-bucket".into(), "invalid".into()),
+ (
+ "postpone.batch-write-fixed-bucket.max-parallelism".into(),
+ "1".into(),
+ ),
+ ]),
+ );
+ let pending_builder = table.new_write_builder();
+ let mut pending_writer = pending_builder.new_write().unwrap();
+ pending_writer
+ .write_arrow_batch(&make_batch(vec![1, 2, 3], vec![10, 20, 30]))
+ .await
+ .unwrap();
+ pending_builder
+ .new_commit()
+ .commit(pending_writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let builder = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![4], vec![40]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.total_buckets == Some(3)
+ && message.check_from_snapshot == Some(1)));
+ builder.new_commit().commit(messages).await.unwrap();
+ assert_pending_and_fixed_files(&table, 3, 1).await;
+ assert_eq!(read_id_value_rows(&table).await, vec![(4, 40)]);
+ }
+
+ #[tokio::test]
+ async fn default_count_restores_level_zero_first_row_layout() {
+ let io = test_file_io();
+ let path = "memory:/fixed_default_first_row";
+ setup_dirs(&io, path).await;
+ let table = with_default_count(&test_postpone_pk_table(&io, path),
"3").copy_with_options(
+ std::collections::HashMap::from([("merge-engine".into(),
"first-row".into())]),
+ );
+ let first = table.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut writer = first.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1], vec![10]))
+ .await
+ .unwrap();
+ first
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ let changed = with_default_count(&table, "5");
+ let builder =
changed.new_postpone_fixed_bucket_write_builder().unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![2], vec![20]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.total_buckets == Some(3)));
+ builder.new_commit().commit(messages).await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn default_count_rejects_invalid_configuration_before_io() {
+ let io = test_file_io();
+ let table = test_postpone_pk_table(&io,
"memory:/invalid_default_bucket");
+ for raw in ["0", "-1", "2147483648", "3.5", "abc"] {
+ let changed = with_default_count(&table, raw);
+ let error = changed
+ .new_postpone_fixed_bucket_write_builder()
+ .unwrap()
+ .new_write()
+ .err()
+ .unwrap();
+ assert!(
+ error.to_string().contains("postpone.default-bucket-num"),
+ "{error}"
+ );
+ }
+ }
}
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 3439e155..bc1537cf 100644
--- a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
+++ b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs
@@ -18,6 +18,7 @@
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;
use crate::table::write_builder::{ensure_table_write_allowed,
validate_commit_user};
use crate::table::{
PostponeBucketPlan, PostponeFixedBucketTableCommit,
PostponeFixedBucketTableWrite, Table,
@@ -82,14 +83,19 @@ impl<'a> PostponeFixedBucketWriteBuilder<'a> {
pub fn new_write(&self) -> Result<PostponeFixedBucketTableWrite> {
ensure_table_write_allowed(self.table)?;
- let plan = self
- .bucket_plan
- .clone()
- .ok_or_else(|| data_invalid("A resolved postpone bucket plan is
required"))?;
+ 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"))?,
+ ),
+ };
PostponeFixedBucketTableWrite::new(
self.table,
self.commit_user.clone(),
plan,
+ load_existing_buckets,
self.overwrite,
self.resources.clone(),
)
@@ -165,4 +171,26 @@ mod tests {
if message.contains("postpone fixed-bucket writes")
&& message.contains("deletion-vectors.enabled=true")));
}
+
+ #[tokio::test]
+ async fn default_bucket_count_routes_without_an_explicit_plan() {
+ let file_io = test_file_io();
+ let table_path = "memory:/default_postpone_bucket_count";
+ setup_dirs(&file_io, table_path).await;
+ let table = test_postpone_pk_table(&file_io,
table_path).copy_with_options(
+
std::collections::HashMap::from([("postpone.default-bucket-num".into(),
"3".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, 3, 4], vec![10, 20, 30,
40]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(!messages.is_empty());
+ assert!(messages
+ .iter()
+ .all(|message| message.total_buckets == Some(3) &&
(0..3).contains(&message.bucket)));
+ builder.new_commit().commit(messages).await.unwrap();
+ }
}
diff --git a/crates/paimon/src/table/table_commit.rs
b/crates/paimon/src/table/table_commit.rs
index eef8c266..9b2f48e9 100644
--- a/crates/paimon/src/table/table_commit.rs
+++ b/crates/paimon/src/table/table_commit.rs
@@ -2379,7 +2379,6 @@ impl TableCommit {
all_entries.extend(delta_entries.iter().cloned());
let merged_entries = merge_active_entries(all_entries);
self.check_total_bucket_conflicts(&merged_entries)?;
- self.check_postpone_bucket_mixing(&merged_entries)?;
self.check_fixed_bucket_ownership_conflicts(
latest_snapshot,
delta_entries,
@@ -2457,15 +2456,6 @@ impl TableCommit {
if *entry.kind() != FileKind::Add {
continue;
}
- if entry.bucket() == POSTPONE_BUCKET &&
owned_partitions.contains(entry.partition())
- {
- return Err(crate::Error::DataInvalid {
- message: format!(
- "Postpone fixed-bucket writer conflict: another
commit wrote bucket=-2 files for the same partition after snapshot
{check_from_snapshot}"
- ),
- source: None,
- });
- }
if owned_buckets.contains(&(entry.partition(),
entry.bucket())) {
return Err(crate::Error::DataInvalid {
message: format!(
@@ -2480,31 +2470,6 @@ impl TableCommit {
Ok(())
}
- fn check_postpone_bucket_mixing(&self, active_entries: &[ManifestEntry])
-> Result<()> {
- if self.total_buckets != POSTPONE_BUCKET {
- return Ok(());
- }
- let fixed_partitions = active_entries
- .iter()
- .filter(|entry| {
- *entry.kind() == FileKind::Add && entry.bucket() >= 0 &&
entry.total_buckets() > 0
- })
- .map(|entry| entry.partition())
- .collect::<HashSet<_>>();
- if active_entries.iter().any(|entry| {
- *entry.kind() == FileKind::Add
- && entry.bucket() == POSTPONE_BUCKET
- && fixed_partitions.contains(entry.partition())
- }) {
- return Err(crate::Error::DataInvalid {
- message: "Cannot mix bucket=-2 files and fixed buckets in the
same partition; use a complete overwrite to migrate the partition"
- .to_string(),
- source: None,
- });
- }
- Ok(())
- }
-
fn check_total_bucket_conflicts(&self, entries: &[ManifestEntry]) ->
Result<()> {
let mut bucket_counts: HashMap<Vec<u8>, i32> = HashMap::new();
for entry in entries {
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index 3325a9c6..814b5c46 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -29,7 +29,7 @@ use crate::spec::{
DataType, MergeEngine, RowKind, RowKindFilter, EMPTY_SERIALIZED_ROW,
POSTPONE_BUCKET,
VALUE_KIND_FIELD_NAME,
};
-use crate::table::bucket_assigner::{BucketAssignerEnum, PartitionBucketKey};
+use crate::table::bucket_assigner::{BatchAssignOutput, BucketAssignerEnum,
PartitionBucketKey};
use crate::table::bucket_assigner_constant::ConstantBucketAssigner;
use crate::table::bucket_assigner_cross::CrossPartitionAssigner;
use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner;
@@ -406,10 +406,19 @@ impl TableWrite {
BucketAssignerEnum::CrossPartition(Box::new(CrossPartitionAssigner::new(
table.clone(),
partition_field_indices,
- primary_key_indices.clone(),
+ schema
+ .primary_keys()
+ .iter()
+ .map(|name| {
+ fields
+ .iter()
+ .position(|field| field.name() == name)
+ .unwrap()
+ })
+ .collect(),
target_bucket_row_number,
merge_engine,
- )))
+ )?))
} else if is_dynamic_bucket {
BucketAssignerEnum::Dynamic(Box::new(DynamicBucketAssigner::new(
table.clone(),
@@ -773,13 +782,11 @@ impl TableWrite {
let fields = self.table.schema().fields().to_vec();
let output = self.bucket_assigner.assign_batch(batch, &fields).await?;
-
+ if matches!(self.bucket_assigner,
BucketAssignerEnum::CrossPartition(_)) {
+ return self.divide_cross_partition_batch(batch, output);
+ }
let mut groups: HashMap<PartitionBucketKey, Vec<usize>> =
HashMap::new();
- let skip_set: HashSet<usize> = output.skips.into_iter().collect();
for row_idx in 0..batch.num_rows() {
- if skip_set.contains(&row_idx) {
- continue;
- }
groups
.entry((
output.partition_bytes[row_idx].clone(),
@@ -788,52 +795,73 @@ impl TableWrite {
.or_default()
.push(row_idx);
}
+ groups
+ .into_iter()
+ .map(|(key, rows)| Ok((key, take_rows(batch, &rows)?)))
+ .collect()
+ }
- let mut result = Vec::with_capacity(groups.len());
- let batch_has_value_kind = batch
- .schema()
- .column_with_name(VALUE_KIND_FIELD_NAME)
- .is_some();
- // Cross-partition writers must always include _VALUE_KIND to keep the
- // Arrow schema stable across batches (some batches may have deletes,
- // others may not — KeyValueFileWriter's concat_batches requires a
- // consistent schema).
- let needs_value_kind = batch_has_value_kind
- || matches!(self.bucket_assigner,
BucketAssignerEnum::CrossPartition(_))
- || !output.deletes.is_empty();
- for (key, row_indices) in groups {
- let sub_batch = take_rows(batch, &row_indices)?;
- let sub_batch = if matches!(self.bucket_assigner,
BucketAssignerEnum::CrossPartition(_))
- {
- self.with_partition_values(&sub_batch, &key.0)?
- } else {
- sub_batch
- };
- let sub_batch = if needs_value_kind && !batch_has_value_kind {
- Self::add_value_kind_column(&sub_batch, 0)?
- } else {
- sub_batch
- };
- result.push((key, sub_batch));
- }
-
- if !output.deletes.is_empty() {
- let mut delete_groups: HashMap<PartitionBucketKey, Vec<usize>> =
HashMap::new();
- for (row_idx, old_partition, old_bucket) in &output.deletes {
- delete_groups
- .entry((old_partition.clone(), *old_bucket))
- .or_default()
- .push(*row_idx);
+ /// Preserve Java GlobalIndexAssigner's DELETE-before-input event order
+ /// within each bucket, including migrations back to an earlier partition.
+ fn divide_cross_partition_batch(
+ &self,
+ batch: &RecordBatch,
+ output: BatchAssignOutput,
+ ) -> Result<Vec<(PartitionBucketKey, RecordBatch)>> {
+ let mut groups: HashMap<PartitionBucketKey, (Vec<usize>, Vec<i8>)> =
HashMap::new();
+ let skips: HashSet<usize> = output.skips.into_iter().collect();
+ let mut deletes = output.deletes.into_iter().peekable();
+ let kinds = batch.column_by_name(VALUE_KIND_FIELD_NAME).map(|column| {
+ column
+ .as_any()
+ .downcast_ref::<arrow_array::Int8Array>()
+ .unwrap()
+ });
+ for row_idx in 0..batch.num_rows() {
+ if skips.contains(&row_idx) {
+ continue;
}
- for (key, row_indices) in delete_groups {
- let sub_batch = take_rows(batch, &row_indices)?;
- let sub_batch = self.with_partition_values(&sub_batch,
&key.0)?;
- let delete_batch = Self::add_value_kind_column(&sub_batch, 3)?;
- result.push((key, delete_batch));
+ if deletes
+ .peek()
+ .is_some_and(|(index, _, _)| *index == row_idx)
+ {
+ let (_, partition, bucket) = deletes.next().unwrap();
+ let (rows, kinds) = groups.entry((partition,
bucket)).or_default();
+ rows.push(row_idx);
+ kinds.push(RowKind::Delete.to_value());
}
+ let key = (
+ output.partition_bytes[row_idx].clone(),
+ output.buckets[row_idx],
+ );
+ let (rows, row_kinds) = groups.entry(key).or_default();
+ rows.push(row_idx);
+ row_kinds.push(kinds.map_or(RowKind::Insert.to_value(), |kinds|
kinds.value(row_idx)));
}
-
- Ok(result)
+ // Generated row kinds are replaced below; migration DELETEs must not
+ // append a second _VALUE_KIND column to a rowkind.field batch.
+ let data = if let Some((index, _)) =
batch.schema().column_with_name(VALUE_KIND_FIELD_NAME)
+ {
+ batch
+ .project(
+ &(0..batch.num_columns())
+ .filter(|i| *i != index)
+ .collect::<Vec<_>>(),
+ )
+ .map_err(|error| crate::Error::DataInvalid {
+ message: format!("Failed to project cross-partition input:
{error}"),
+ source: Some(Box::new(error)),
+ })?
+ } else {
+ batch.clone()
+ };
+ groups
+ .into_iter()
+ .map(|(key, (rows, kinds))| {
+ let routed = self.with_partition_values(&take_rows(&data,
&rows)?, &key.0)?;
+ Ok((key, Self::add_per_row_value_kind_column(&routed, kinds)?))
+ })
+ .collect()
}
/// Keep physical partition columns consistent with the routed partition.
@@ -867,30 +895,6 @@ impl TableWrite {
})
}
- /// Add a `_VALUE_KIND` column to a batch with the given value for all
rows.
- fn add_value_kind_column(batch: &RecordBatch, value_kind: i8) ->
Result<RecordBatch> {
- use arrow_array::Int8Array;
- use arrow_schema::{DataType as ArrowDataType, Field as ArrowField};
-
- let vk_array = Arc::new(Int8Array::from(vec![value_kind;
batch.num_rows()]));
- let vk_field = Arc::new(ArrowField::new(
- crate::spec::VALUE_KIND_FIELD_NAME,
- ArrowDataType::Int8,
- false,
- ));
-
- let mut fields = batch.schema().fields().to_vec();
- let mut columns: Vec<Arc<dyn arrow_array::Array>> =
batch.columns().to_vec();
- fields.push(vk_field);
- columns.push(vk_array);
-
- let schema = Arc::new(arrow_schema::Schema::new(fields));
- RecordBatch::try_new(schema, columns).map_err(|e|
crate::Error::DataInvalid {
- message: format!("Failed to add _VALUE_KIND column: {e}"),
- source: None,
- })
- }
-
fn enrich_rowkind_batch(&self, batch: &RecordBatch) -> Result<RecordBatch>
{
let Some(generator) = &self.row_kind_generator else {
return self.filter_rowkind_batch(batch);
@@ -1224,6 +1228,9 @@ impl TableWrite {
bucket,
schema_id: self.schema_id,
target_file_size: self.target_file_size,
+ target_file_row_num: self.target_file_row_num,
+ primary_key_indices: self.primary_key_indices.clone(),
+ value_fields: self.table.schema().fields().to_vec(),
file_compression: self.file_compression.clone(),
file_compression_zstd_level:
self.file_compression_zstd_level,
write_buffer_size: self.write_buffer_size,
@@ -4542,9 +4549,21 @@ pub(in crate::table) mod tests {
"Postpone mode should preserve arrival order"
);
- // Empty key stats for postpone mode
- assert_eq!(file.min_key, EMPTY_SERIALIZED_ROW.clone());
- assert_eq!(file.max_key, EMPTY_SERIALIZED_ROW.clone());
+ // Java records the first and last keys without sorting postpone input.
+ assert_eq!(
+ BinaryRow::from_serialized_bytes(&file.min_key)
+ .unwrap()
+ .get_int(0)
+ .unwrap(),
+ 3
+ );
+ assert_eq!(
+ BinaryRow::from_serialized_bytes(&file.max_key)
+ .unwrap()
+ .get_int(0)
+ .unwrap(),
+ 2
+ );
}
// -----------------------------------------------------------------------
diff --git a/crates/paimon/tests/cross_partition_write_test.rs
b/crates/paimon/tests/cross_partition_write_test.rs
new file mode 100644
index 00000000..36b82988
--- /dev/null
+++ b/crates/paimon/tests/cross_partition_write_test.rs
@@ -0,0 +1,323 @@
+// 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.
+
+mod common;
+
+use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
+use common::incremental_helpers::{memory_table, persist_table_schema,
setup_dirs, write_batch};
+use futures::TryStreamExt;
+use paimon::spec::{DataType, IntType, Schema, TableSchema, VarCharType};
+use paimon::table::Table;
+use std::sync::Arc;
+
+async fn table(partition_in_key: bool, options: &[(&str, &str)]) -> Table {
+ let mut schema = Schema::builder()
+ .column("tenant", DataType::VarChar(VarCharType::string_type()))
+ .column("p", DataType::VarChar(VarCharType::string_type()))
+ .column("id", DataType::Int(IntType::new()))
+ .column("v", DataType::Int(IntType::new()))
+ .partition_keys(["tenant", "p"])
+ .primary_key(if partition_in_key {
+ vec!["tenant", "id"]
+ } else {
+ vec!["id"]
+ })
+ .option("bucket", "-1");
+ if options.iter().any(|(key, _)| *key == "rowkind.field") {
+ schema = schema.column("op",
DataType::VarChar(VarCharType::string_type()));
+ }
+ for (key, value) in options {
+ schema = schema.option(*key, *value);
+ }
+ let schema = TableSchema::new(0, &schema.build().unwrap());
+ let (io, table) = memory_table("memory:/cross_partition_write", schema);
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ table
+}
+
+fn batch(rows: &[(&str, &str, i32, i32)]) -> RecordBatch {
+ RecordBatch::try_from_iter([
+ (
+ "tenant",
+ Arc::new(StringArray::from_iter_values(rows.iter().map(|r| r.0)))
as ArrayRef,
+ ),
+ (
+ "p",
+ Arc::new(StringArray::from_iter_values(rows.iter().map(|r| r.1)))
as ArrayRef,
+ ),
+ (
+ "id",
+ Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.2)))
as ArrayRef,
+ ),
+ (
+ "v",
+ Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.3)))
as ArrayRef,
+ ),
+ ])
+ .unwrap()
+}
+
+async fn rows(table: &Table) -> Vec<(String, String, i32, i32)> {
+ let read = table.new_read_builder();
+ let plan = read.new_scan().with_scan_all_files().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = read
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ let tenant = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let p = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let id = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let v = batch
+ .column(3)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ rows.extend((0..batch.num_rows()).map(|i| {
+ (
+ tenant.value(i).into(),
+ p.value(i).into(),
+ id.value(i),
+ v.value(i),
+ )
+ }));
+ }
+ rows.sort();
+ rows
+}
+
+#[tokio::test]
+async fn
repeated_migrations_keep_the_last_input_even_when_returning_to_a_partition() {
+ for chunk_size in [1, 2, 3, 4] {
+ let table = table(false, &[]).await;
+ let input = [
+ ("t", "a", 1, 10),
+ ("t", "b", 1, 20),
+ ("t", "a", 1, 30),
+ ("t", "a", 1, 40),
+ ];
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ for chunk in input.chunks(chunk_size) {
+ writer.write_arrow_batch(&batch(chunk)).await.unwrap();
+ }
+ builder
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ assert_eq!(
+ rows(&table).await,
+ vec![("t".into(), "a".into(), 1, 40)],
+ "chunk_size={chunk_size}"
+ );
+ }
+}
+
+#[tokio::test]
+async fn global_index_keeps_partition_columns_of_composite_primary_keys() {
+ let table = table(true, &[]).await;
+ write_batch(&table, &batch(&[("x", "a", 1, 10), ("y", "a", 1, 20)])).await;
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("x".into(), "a".into(), 1, 10),
+ ("y".into(), "a".into(), 1, 20)
+ ]
+ );
+ // Reconstructing the global index must retain both tenants too.
+ write_batch(&table, &batch(&[("x", "b", 1, 30)])).await;
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("x".into(), "b".into(), 1, 30),
+ ("y".into(), "a".into(), 1, 20)
+ ]
+ );
+}
+
+#[tokio::test]
+async fn dynamic_overwrite_keeps_global_index_for_migration_deletes() {
+ let table = table(false, &[]).await;
+ write_batch(&table, &batch(&[("t", "a", 1, 10), ("t", "c", 2, 20)])).await;
+ let builder = table.new_write_builder().with_overwrite();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&batch(&[("t", "b", 1, 30)]))
+ .await
+ .unwrap();
+ builder
+ .new_commit()
+ .overwrite(writer.prepare_commit().await.unwrap(), None)
+ .await
+ .unwrap();
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("t".into(), "b".into(), 1, 30),
+ ("t".into(), "c".into(), 2, 20)
+ ]
+ );
+ write_batch(&table, &batch(&[("t", "d", 1, 40)])).await;
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("t".into(), "c".into(), 2, 20),
+ ("t".into(), "d".into(), 1, 40)
+ ]
+ );
+}
+
+#[tokio::test]
+async fn merge_engines_route_updates_like_java_after_restart() {
+ for engine in ["deduplicate", "first-row", "partial-update",
"aggregation"] {
+ for restart in [false, true] {
+ let mut options = vec![("merge-engine", engine)];
+ if engine == "aggregation" {
+ options.push(("fields.v.aggregate-function", "sum"));
+ }
+ let table = table(true, &options).await;
+ let first = batch(&[("x", "a", 1, 10), ("y", "a", 1, 100)]);
+ let second = batch(&[("x", "b", 1, 20), ("y", "b", 1, 200)]);
+ if restart {
+ write_batch(&table, &first).await;
+ write_batch(&table, &second).await;
+ } else {
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer.write_arrow_batch(&first).await.unwrap();
+ writer.write_arrow_batch(&second).await.unwrap();
+ builder
+ .new_commit()
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ }
+ let (partition, x, y) = match engine {
+ "deduplicate" => ("b", 20, 200),
+ "first-row" => ("a", 10, 100),
+ "partial-update" => ("a", 20, 200),
+ "aggregation" => ("a", 30, 300),
+ _ => unreachable!(),
+ };
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("x".into(), partition.into(), 1, x),
+ ("y".into(), partition.into(), 1, y)
+ ],
+ "engine={engine}, restart={restart}"
+ );
+ }
+ }
+}
+
+fn with_kinds(rows: &[(&str, &str, i32, i32)], kinds: &[&str]) -> RecordBatch {
+ let data = batch(rows);
+ let mut fields = data.schema().fields().to_vec();
+ fields.push(Arc::new(arrow_schema::Field::new(
+ "op",
+ arrow_schema::DataType::Utf8,
+ true,
+ )));
+ let mut columns = data.columns().to_vec();
+ columns.push(Arc::new(StringArray::from(kinds.to_vec())));
+ RecordBatch::try_new(Arc::new(arrow_schema::Schema::new(fields)),
columns).unwrap()
+}
+
+#[tokio::test]
+async fn generated_row_kinds_keep_delete_and_update_order_during_migrations() {
+ let table = table(false, &[("rowkind.field", "op")]).await;
+ write_batch(&table, &with_kinds(&[("t", "a", 1, 10)], &["+I"])).await;
+ write_batch(
+ &table,
+ &with_kinds(
+ &[
+ ("t", "b", 1, 20),
+ ("t", "a", 1, 30),
+ ("t", "b", 2, 40),
+ ("t", "c", 2, 50),
+ ],
+ &["+U", "+U", "+I", "-D"],
+ ),
+ )
+ .await;
+ assert_eq!(rows(&table).await, vec![("t".into(), "a".into(), 1, 30)]);
+ // The deleted key must not reappear when the global index is rebuilt.
+ write_batch(&table, &with_kinds(&[("t", "c", 2, 60)], &["+I"])).await;
+ assert_eq!(
+ rows(&table).await,
+ vec![
+ ("t".into(), "a".into(), 1, 30),
+ ("t".into(), "c".into(), 2, 60)
+ ]
+ );
+}
+
+#[tokio::test]
+async fn bootstrap_rejects_duplicate_global_keys_before_writing() {
+ let table = table(false, &[]).await;
+ // Simulate existing files with the same full primary key in two
partitions.
+ // Such a table cannot supply an unambiguous global location on restart.
+ let fixed = table.copy_with_options([("bucket".into(),
"1".into())].into());
+ write_batch(&fixed, &batch(&[("t", "a", 1, 10), ("t", "b", 1, 20)])).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ let error = writer
+ .write_arrow_batch(&batch(&[("t", "c", 1, 30)]))
+ .await
+ .unwrap_err();
+ assert!(
+ error.to_string().contains("Duplicate primary key"),
+ "{error}"
+ );
+ assert!(writer.prepare_commit().await.is_err());
+ assert_eq!(rows(&table).await.len(), 2);
+}
+
+#[tokio::test]
+async fn cross_partition_rejects_unsupported_index_configuration() {
+ let table = table(false, &[]).await;
+ for (key, value, message) in [
+ ("bucket-key", "id", "Cannot define 'bucket-key'"),
+ ("sequence.field", "v", "Cannot define 'sequence.field'"),
+ ("cross-partition-upsert.index-ttl", "1 h", "index-ttl"),
+ ] {
+ let configured = table.copy_with_options([(key.into(),
value.into())].into());
+ match configured.new_write_builder().new_write() {
+ Ok(_) => panic!("should reject {key}"),
+ Err(error) => assert!(error.to_string().contains(message),
"{error}"),
+ }
+ }
+}
diff --git a/crates/paimon/tests/postpone_native_write_test.rs
b/crates/paimon/tests/postpone_native_write_test.rs
new file mode 100644
index 00000000..1063e104
--- /dev/null
+++ b/crates/paimon/tests/postpone_native_write_test.rs
@@ -0,0 +1,137 @@
+// 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.
+
+mod common;
+
+use arrow_array::{Array, Int32Array, Int64Array, Int8Array};
+use common::incremental_helpers::{make_batch_with_kinds, memory_table,
pk_schema, setup_dirs};
+use paimon::spec::{BinaryRow, POSTPONE_BUCKET};
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+
+#[tokio::test]
+async fn postpone_rolls_by_batch_row_limit_and_preserves_unsorted_retracts() {
+ let (io, table) = memory_table(
+ "memory:/postpone_row_limit",
+ pk_schema(&[("bucket", "-2"), ("target-file-row-num", "3")]),
+ );
+ setup_dirs(&io, table.location()).await;
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ let ids = [5, 1, 3, 2, 1];
+ let kinds = [0, 1, 2, 3, 0];
+ for range in [0..2, 2..4, 4..5] {
+ writer
+ .write_arrow_batch(&make_batch_with_kinds(
+ ids[range.clone()].to_vec(),
+ vec![10; range.len()],
+ kinds[range].to_vec(),
+ ))
+ .await
+ .unwrap();
+ }
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, POSTPONE_BUCKET);
+ let files = messages[0].new_files.clone();
+ assert_eq!(
+ files.iter().map(|f| f.row_count).collect::<Vec<_>>(),
+ vec![4, 1]
+ );
+ assert_eq!(
+ files.iter().map(|f| f.delete_row_count).collect::<Vec<_>>(),
+ vec![Some(2), Some(0)]
+ );
+ assert_eq!(
+ files
+ .iter()
+ .map(|f| (f.min_sequence_number, f.max_sequence_number))
+ .collect::<Vec<_>>(),
+ vec![(0, 3), (4, 4)]
+ );
+ let bounds: Vec<_> = files
+ .iter()
+ .map(|file| {
+ (
+ BinaryRow::from_serialized_bytes(&file.min_key)
+ .unwrap()
+ .get_int(0)
+ .unwrap(),
+ BinaryRow::from_serialized_bytes(&file.max_key)
+ .unwrap()
+ .get_int(0)
+ .unwrap(),
+ )
+ })
+ .collect();
+ // Java records the first and last arriving keys for pending files.
+ assert_eq!(bounds, vec![(5, 2), (1, 1)]);
+ let mut actual = Vec::new();
+ for file in files {
+ let path = format!("{}/bucket-postpone/{}", table.location(),
file.file_name);
+ let bytes = io.new_input(&path).unwrap().read().await.unwrap();
+ let reader = ParquetRecordBatchReaderBuilder::try_new(bytes)
+ .unwrap()
+ .build()
+ .unwrap();
+ for batch in reader {
+ let batch = batch.unwrap();
+ let seq = batch
+ .column_by_name("_SEQUENCE_NUMBER")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let kind = batch
+ .column_by_name("_VALUE_KIND")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap();
+ let id = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ actual
+ .extend((0..batch.num_rows()).map(|i| (seq.value(i),
kind.value(i), id.value(i))));
+ }
+ }
+ assert_eq!(
+ actual,
+ vec![(0, 0, 5), (1, 1, 1), (2, 2, 3), (3, 3, 2), (4, 0, 1)]
+ );
+}
+
+#[tokio::test]
+async fn postpone_delete_counts_reset_after_prepare_commit() {
+ let (io, table) = memory_table(
+ "memory:/postpone_delete_counts",
+ pk_schema(&[("bucket", "-2")]),
+ );
+ setup_dirs(&io, table.location()).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ for (kinds, expected) in [(vec![0, 3, 1], 2), (vec![0, 2], 0)] {
+ let len = kinds.len();
+ writer
+ .write_arrow_batch(&make_batch_with_kinds(vec![1; len], vec![10;
len], kinds))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_files[0].delete_row_count, Some(expected));
+ }
+}
