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 0452c6ee [core] Support precomputed bucket writes and snapshot
restoration (#1056)
0452c6ee is described below
commit 0452c6eef7caeaf5fc0f7b28ae0d4f098b3acc54
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 16:38:41 2026 +0800
[core] Support precomputed bucket writes and snapshot restoration (#1056)
---
.../python/python/pypaimon_rust/datafusion.pyi | 6 +-
bindings/python/src/write.rs | 86 ++-
bindings/python/tests/test_write.py | 33 +
crates/paimon/src/table/bucket_assigner_dynamic.rs | 158 ++++-
crates/paimon/src/table/table_write.rs | 75 ++-
crates/paimon/src/table/table_write_precomputed.rs | 102 +++
crates/paimon/src/table/write_builder.rs | 47 +-
.../paimon/tests/precomputed_bucket_write_test.rs | 731 +++++++++++++++++++++
8 files changed, 1190 insertions(+), 48 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 517222e3..efac2bc3 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -237,7 +237,7 @@ class CommitMessage:
class BatchTableWrite:
def with_write_type(self, columns: List[str]) -> "BatchTableWrite": ...
def close(self) -> None: ...
- def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ...
+ def write_arrow(self, batch: pyarrow.RecordBatch, bucket: Optional[int] =
None) -> None: ...
def prepare_commit(self) -> List[CommitMessage]:
"""Prepare once per instance, including empty or failed attempts."""
...
@@ -290,7 +290,7 @@ class StreamTableUpdate:
class StreamTableWrite:
def with_write_type(self, columns: List[str]) -> "StreamTableWrite": ...
def close(self) -> None: ...
- def write_arrow(self, batch: pyarrow.RecordBatch) -> None: ...
+ def write_arrow(self, batch: pyarrow.RecordBatch, bucket: Optional[int] =
None) -> None: ...
def prepare_commit(self, wait_compaction: bool, commit_identifier: int) ->
List[CommitMessage]:
"""Flush a checkpoint; submit its messages with the same identifier.
@@ -341,6 +341,7 @@ class BatchWriteBuilder:
use the schema's Python types; None or the default partition name is
null.
"""
...
+ def with_restore_snapshot(self, snapshot_id: int) -> "BatchWriteBuilder":
...
def new_write(self) -> BatchTableWrite: ...
def new_update(self) -> BatchTableUpdate: ...
def new_commit(self) -> BatchTableCommit: ...
@@ -368,6 +369,7 @@ class PostponeFixedBucketWriteBuilder:
class StreamWriteBuilder:
def commit_user(self) -> str: ...
def with_commit_user(self, commit_user: str) -> "StreamWriteBuilder": ...
+ def with_restore_snapshot(self, snapshot_id: int) -> "StreamWriteBuilder":
...
def new_write(self) -> StreamTableWrite: ...
def new_update(self) -> StreamTableUpdate: ...
def new_commit(self) -> StreamTableCommit: ...
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 3057a440..eea7870a 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -41,12 +41,26 @@ type PythonPartitionSpec = HashMap<String, Py<PyAny>>;
struct WriteContext {
table: Arc<Table>,
commit_user: String,
+ restore_snapshot_id: Option<i64>,
}
impl WriteContext {
fn new(table: Arc<Table>) -> Self {
let commit_user = table.new_write_builder().commit_user().to_string();
- Self { table, commit_user }
+ Self {
+ table,
+ commit_user,
+ restore_snapshot_id: None,
+ }
+ }
+
+ fn with_restore_snapshot(&mut self, snapshot_id: i64) -> PyResult<()> {
+ self.table
+ .new_write_builder()
+ .with_restore_snapshot(snapshot_id)
+ .map_err(to_py_err)?;
+ self.restore_snapshot_id = Some(snapshot_id);
+ Ok(())
}
fn new_write(&self, overwrite: bool) -> PyResult<WriteState> {
@@ -60,6 +74,10 @@ impl WriteContext {
} else {
builder
};
+ let builder = match self.restore_snapshot_id {
+ Some(id) => builder.with_restore_snapshot(id).map_err(to_py_err)?,
+ None => builder,
+ };
Ok(WriteState {
inner: Some(WriteTarget::Table(Box::new(
builder.new_write().map_err(to_py_err)?,
@@ -184,6 +202,14 @@ impl PyBatchWriteBuilder {
Ok(slf)
}
+ fn with_restore_snapshot(
+ mut slf: PyRefMut<'_, Self>,
+ snapshot_id: i64,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ slf.context.with_restore_snapshot(snapshot_id)?;
+ Ok(slf)
+ }
+
fn new_write(&self) -> PyResult<PyBatchTableWrite> {
Ok(PyBatchTableWrite {
state: self.context.new_write(self.static_partition.is_some())?,
@@ -253,6 +279,14 @@ impl PyStreamWriteBuilder {
Ok(slf)
}
+ fn with_restore_snapshot(
+ mut slf: PyRefMut<'_, Self>,
+ snapshot_id: i64,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ slf.context.with_restore_snapshot(snapshot_id)?;
+ Ok(slf)
+ }
+
fn new_write(&self) -> PyResult<PyStreamTableWrite> {
Ok(PyStreamTableWrite {
state: self.context.new_write(false)?,
@@ -292,10 +326,23 @@ 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 write_arrow(
+ &mut self,
+ batch: &RecordBatch,
+ bucket: Option<i32>,
+ ) -> paimon::Result<()> {
+ match (self, bucket) {
+ (Self::Table(writer), bucket) => {
+ writer
+ .write_arrow(std::slice::from_ref(batch), bucket)
+ .await
+ }
+ (Self::PostponeFixed(writer), None) =>
writer.write_arrow_batch(batch).await,
+ (Self::PostponeFixed(_), Some(_)) =>
Err(paimon::Error::DataInvalid {
+ message: "Precomputed bucket writes require HASH_FIXED or
HASH_DYNAMIC tables"
+ .into(),
+ source: None,
+ }),
}
}
@@ -610,13 +657,18 @@ impl WriteState {
.map_err(to_py_err)
}
- fn write_arrow(&mut self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<()> {
+ fn write_arrow(
+ &mut self,
+ py: Python<'_>,
+ batch: &Bound<'_, PyAny>,
+ bucket: Option<i32>,
+ ) -> PyResult<()> {
let batch = RecordBatch::from_pyarrow_bound(batch)?;
let inner = self
.inner
.as_mut()
.ok_or_else(|| PyRuntimeError::new_err("TableWrite is closed"))?;
- py.detach(|| runtime().block_on(inner.write_arrow_batch(&batch)))
+ py.detach(|| runtime().block_on(inner.write_arrow(&batch, bucket)))
.map_err(to_py_err)
}
@@ -848,8 +900,14 @@ impl PyBatchTableWrite {
}
}
- fn write_arrow(&mut self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<()> {
- self.state.write_arrow(py, batch)
+ #[pyo3(signature = (batch, bucket=None))]
+ fn write_arrow(
+ &mut self,
+ py: Python<'_>,
+ batch: &Bound<'_, PyAny>,
+ bucket: Option<i32>,
+ ) -> PyResult<()> {
+ self.state.write_arrow(py, batch, bucket)
}
fn prepare_commit(&mut self, py: Python<'_>) ->
PyResult<Vec<PyCommitMessage>> {
@@ -888,8 +946,14 @@ impl PyStreamTableWrite {
}
}
- fn write_arrow(&mut self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<()> {
- self.state.write_arrow(py, batch)
+ #[pyo3(signature = (batch, bucket=None))]
+ fn write_arrow(
+ &mut self,
+ py: Python<'_>,
+ batch: &Bound<'_, PyAny>,
+ bucket: Option<i32>,
+ ) -> PyResult<()> {
+ self.state.write_arrow(py, batch, bucket)
}
/// Rust currently flushes synchronously and has no background compaction.
diff --git a/bindings/python/tests/test_write.py
b/bindings/python/tests/test_write.py
index 900a1774..a915f314 100644
--- a/bindings/python/tests/test_write.py
+++ b/bindings/python/tests/test_write.py
@@ -89,6 +89,39 @@ def
test_partial_write_type_roundtrip_and_lifecycle(tmp_path, streaming):
"id": [1, 2], "name": [None, None]}
[email protected]("streaming", [False, True])
[email protected]("dynamic", [False, True])
+def test_precomputed_bucket_api_roundtrip_and_closed_writer(tmp_path,
streaming, dynamic):
+ ctx = SQLContext()
+ ctx.register_catalog("paimon", {"warehouse": str(tmp_path)})
+ ctx.sql("CREATE SCHEMA paimon.wdb")
+ bucket = -1 if dynamic else 4
+ ctx.sql("CREATE TABLE paimon.wdb.t (id INT, name STRING, PRIMARY KEY (id))
"
+ "WITH ('bucket' = '{}')".format(bucket))
+ table = _get_table(str(tmp_path))
+ builder = table.new_stream_write_builder() if streaming else
table.new_batch_write_builder()
+ if dynamic:
+ assert builder.with_restore_snapshot(0) is builder
+ else:
+ with pytest.raises(ValueError, match="HASH_DYNAMIC"):
+ builder.with_restore_snapshot(0)
+ writer = builder.new_write()
+ try:
+ writer.write_arrow(_batch([1, 1, 2], ["old", "new", "other"]),
bucket=3)
+ with pytest.raises(ValueError, match="Bucket id"):
+ writer.write_arrow(_batch([3], ["wrong"]), bucket=-1)
+ messages = writer.prepare_commit(True, 1) if streaming else
writer.prepare_commit()
+ commit = builder.new_commit()
+ commit.commit(1, messages) if streaming else commit.commit(messages)
+ assert len(messages) == 1
+ finally:
+ writer.close()
+ with pytest.raises(RuntimeError, match="closed"):
+ writer.write_arrow(_batch([4], ["closed"]), bucket=3)
+ assert pa.Table.from_batches(ctx.sql("SELECT id, name FROM paimon.wdb.t
ORDER BY id")).to_pydict() == {
+ "id": [1, 2], "name": ["new", "other"]}
+
+
@pytest.mark.parametrize("primary_key", [False, True])
def test_custom_data_file_prefix_matches_table_option(tmp_path, primary_key):
diff --git a/crates/paimon/src/table/bucket_assigner_dynamic.rs
b/crates/paimon/src/table/bucket_assigner_dynamic.rs
index ee30e5ef..d6422b60 100644
--- a/crates/paimon/src/table/bucket_assigner_dynamic.rs
+++ b/crates/paimon/src/table/bucket_assigner_dynamic.rs
@@ -276,6 +276,67 @@ impl PartitionIndex {
}
}
+ /// Merge buckets restored or modified by direct writers. The base snapshot
+ /// is shared, so an already-notified maintainer retains its newer state.
+ fn merge(&mut self, other: Self) -> Result<()> {
+ for (hash, bucket) in other.hash_to_bucket {
+ self.record_mapping(hash, bucket)?;
+ }
+ self.bucket_maintainers.extend(other.bucket_maintainers);
+ Ok(())
+ }
+
+ fn record_mapping(&mut self, hash: i32, bucket: i32) -> Result<()> {
+ if let Some(&previous) = self.hash_to_bucket.get(&hash) {
+ if previous != bucket {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Precomputed hash {hash} belongs to bucket {previous},
not {bucket}"
+ ),
+ source: None,
+ });
+ }
+ return Ok(());
+ }
+ self.hash_to_bucket.insert(hash, bucket);
+ if self.all_buckets.insert(bucket) {
+ self.bucket_ids.push(bucket);
+ self.non_full_buckets.insert(bucket, 0);
+ }
+ if let Some(count) = self.non_full_buckets.get_mut(&bucket) {
+ *count += 1;
+ if *count >= self.target_bucket_row_number {
+ self.non_full_buckets.remove(&bucket);
+ }
+ }
+ Ok(())
+ }
+
+ fn notify_precomputed(&mut self, bucket: i32, hashes: &[i32]) ->
Result<()> {
+ let notified: HashSet<i32> = hashes.iter().copied().collect();
+ // Validate the entire group before changing the assignment or
maintainer.
+ for &hash in hashes {
+ if let Some(&previous) = self.hash_to_bucket.get(&hash) {
+ if previous != bucket {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Precomputed hash {hash} belongs to bucket
{previous}, not {bucket}"
+ ),
+ source: None,
+ });
+ }
+ }
+ }
+ for hash in notified {
+ self.record_mapping(hash, bucket)?;
+ self.bucket_maintainers
+ .entry(bucket)
+ .or_insert_with(|| DynamicBucketIndexMaintainer::new(vec![]))
+ .notify_new_record(hash);
+ }
+ Ok(())
+ }
+
/// Load partition index from existing hash index files.
///
/// Reads all HASH-type index entries for this partition and reconstructs
@@ -480,10 +541,14 @@ pub(crate) struct DynamicBucketAssigner {
/// Schema fields for BinaryRow extraction.
fields: Vec<DataField>,
partition_indexes: HashMap<Vec<u8>, PartitionIndex>,
+ /// Direct writers restore only the supplied buckets. Assignment upgrades
+ /// these partial indexes to a full partition index on first use.
+ precomputed_indexes: HashMap<Vec<u8>, PartitionIndex>,
target_bucket_row_number: i64,
table: Table,
max_buckets: i32,
snapshot: Option<Snapshot>,
+ restore_snapshot_id: Option<i64>,
/// Cached index manifest entries from the latest snapshot (loaded once).
cached_index_entries: Option<Vec<IndexManifestEntry>>,
/// Overwrite mode: skip loading existing index entries.
@@ -515,9 +580,11 @@ impl DynamicBucketAssigner {
primary_key_indices,
fields: table.schema().fields().to_vec(),
partition_indexes: HashMap::new(),
+ precomputed_indexes: HashMap::new(),
target_bucket_row_number: options.dynamic_bucket_target_row_num(),
max_buckets: options.dynamic_bucket_max_buckets()?,
snapshot: None,
+ restore_snapshot_id: None,
cached_index_entries: None,
is_overwrite,
partition_computer,
@@ -530,6 +597,59 @@ impl DynamicBucketAssigner {
self.is_overwrite = is_overwrite;
}
+ /// Set by the builder before any index state is loaded.
+ pub fn set_restore_snapshot(&mut self, snapshot_id: i64) {
+ self.restore_snapshot_id = Some(snapshot_id);
+ }
+
+ /// Java's DynamicBucketIndexMaintainer sees keys after bucket assignment.
+ /// Keep that same full-file replacement lifecycle for precomputed groups.
+ pub async fn notify_precomputed_batch(
+ &mut self,
+ batch: &RecordBatch,
+ partition: &[u8],
+ bucket: i32,
+ ) -> Result<()> {
+ let hashes = batch_hash_codes(batch, &self.primary_key_indices,
&self.fields)?;
+ if let Some(index) = self.partition_indexes.get_mut(partition) {
+ return index.notify_precomputed(bucket, &hashes);
+ }
+ let loaded = self
+ .precomputed_indexes
+ .get(partition)
+ .is_some_and(|index|
index.bucket_maintainers.contains_key(&bucket));
+ if !loaded {
+ self.ensure_index_entries_loaded().await?;
+ let entries: Vec<_> = self
+ .cached_index_entries
+ .as_deref()
+ .unwrap_or(&[])
+ .iter()
+ .filter(|entry| {
+ entry.partition == partition
+ && entry.bucket == bucket
+ && entry.index_file.index_type == HASH_INDEX
+ })
+ .cloned()
+ .collect();
+ let mut restored = self.load_hash_indexes(partition,
&entries).await?;
+ // Remember an empty bucket too, so subsequent batches do not
restore
+ // it repeatedly. It remains unmodified until a key is notified.
+ restored
+ .bucket_maintainers
+ .entry(bucket)
+ .or_insert_with(|| DynamicBucketIndexMaintainer::new(vec![]));
+ self.precomputed_indexes
+ .entry(partition.to_vec())
+ .or_insert_with(||
PartitionIndex::empty(self.target_bucket_row_number))
+ .merge(restored)?;
+ }
+ self.precomputed_indexes
+ .get_mut(partition)
+ .unwrap()
+ .notify_precomputed(bucket, &hashes)
+ }
+
/// Load all index manifest entries from the latest snapshot (cached).
/// Overwrite mode skips loading — old index is irrelevant.
async fn ensure_index_entries_loaded(&mut self) -> Result<()> {
@@ -541,7 +661,11 @@ impl DynamicBucketAssigner {
return Ok(());
}
let snapshot_manager = self.table.snapshot_manager();
- let latest_snapshot = snapshot_manager.get_latest_snapshot().await?;
+ let latest_snapshot = match self.restore_snapshot_id {
+ Some(0) => None,
+ Some(id) => Some(snapshot_manager.get_snapshot(id).await?),
+ None => snapshot_manager.get_latest_snapshot().await?,
+ };
let entries = if let Some(snapshot) = &latest_snapshot {
if let Some(index_manifest_name) = snapshot.index_manifest() {
@@ -578,6 +702,15 @@ impl DynamicBucketAssigner {
self.validate_data_buckets(partition_bytes, &partition_entries)
.await?;
+ self.load_hash_indexes(partition_bytes, &partition_entries)
+ .await
+ }
+
+ async fn load_hash_indexes(
+ &self,
+ partition_bytes: &[u8],
+ partition_entries: &[IndexManifestEntry],
+ ) -> Result<PartitionIndex> {
if !partition_entries.is_empty() {
let partition_path = self.partition_path(partition_bytes)?;
let options = CoreOptions::new(self.table.schema().options());
@@ -590,7 +723,7 @@ impl DynamicBucketAssigner {
return PartitionIndex::load(
self.table.file_io(),
&layout,
- &partition_entries,
+ partition_entries,
self.target_bucket_row_number,
)
.await;
@@ -606,6 +739,9 @@ impl DynamicBucketAssigner {
partition: &[u8],
indexes: &[IndexManifestEntry],
) -> Result<()> {
+ if self.is_overwrite {
+ return Ok(());
+ }
let Some(snapshot) = &self.snapshot else {
return Ok(());
};
@@ -658,7 +794,10 @@ impl BucketAssigner for DynamicBucketAssigner {
self.ensure_index_entries_loaded().await?;
}
for partition_bytes in unseen {
- let index = self.load_partition_index(&partition_bytes).await?;
+ let mut index = self.load_partition_index(&partition_bytes).await?;
+ if let Some(notified) =
self.precomputed_indexes.remove(&partition_bytes) {
+ index.merge(notified)?;
+ }
self.partition_indexes.insert(partition_bytes, index);
}
@@ -687,7 +826,12 @@ impl BucketAssigner for DynamicBucketAssigner {
let mut result = HashMap::new();
let table_path =
self.table.location().trim_end_matches('/').to_string();
let index_file_in_data_file_dir = self.index_file_in_data_file_dir;
- let partition_keys: Vec<Vec<u8>> =
self.partition_indexes.keys().cloned().collect();
+ let partition_keys: Vec<Vec<u8>> = self
+ .partition_indexes
+ .keys()
+ .chain(self.precomputed_indexes.keys())
+ .cloned()
+ .collect();
let mut partition_paths = Vec::with_capacity(partition_keys.len());
for partition_bytes in &partition_keys {
partition_paths.push(self.partition_path(partition_bytes)?);
@@ -701,7 +845,11 @@ impl BucketAssigner for DynamicBucketAssigner {
partition_path: &partition_path,
index_file_in_data_file_dir,
};
- if let Some(partition_index) =
self.partition_indexes.get_mut(&partition_bytes) {
+ if let Some(partition_index) = self
+ .partition_indexes
+ .get_mut(&partition_bytes)
+ .or_else(||
self.precomputed_indexes.get_mut(&partition_bytes))
+ {
let bucket_files = partition_index
.prepare_commit(
file_io,
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index d5d1a7e2..f36323d1 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -52,6 +52,9 @@ use arrow_array::{ArrayRef, RecordBatch};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
+#[path = "table_write_precomputed.rs"]
+mod precomputed;
+
/// Enum to hold either an append-only writer, a key-value writer, or a
postpone writer.
enum FileWriter {
Append(DataFileWriter),
@@ -101,9 +104,11 @@ impl FileWriter {
if result.is_err() {
self.abort().await;
}
- // A pending-bucket writer keeps its first-retract validation and
- // arrival sequence across stream checkpoints, as Java does.
- let retain = result.is_ok() && matches!(&self,
FileWriter::Postpone(_));
+ // Java MergeTreeWriter keeps its sequence counter across checkpoints,
+ // including prepares whose messages have not been committed yet.
+ // Pending-bucket writers also retain first-retract validation.
+ let retain =
+ result.is_ok() && matches!(&self, FileWriter::KeyValue(_) |
FileWriter::Postpone(_));
(result, retain.then_some(self))
}
@@ -158,6 +163,7 @@ pub struct TableWrite {
changelog_file_compression: String,
partition_seq_cache: HashMap<Vec<u8>, HashMap<i32, i64>>,
sequence_snapshot: Option<Option<Snapshot>>,
+ restore_snapshot_id: Option<i64>,
commit_user: String,
/// Shared by this writer's postpone files, as in Java's per-writer
writeId.
postpone_write_id: i32,
@@ -230,6 +236,7 @@ impl TableWrite {
changelog_file_compression: String::new(),
partition_seq_cache: HashMap::new(),
sequence_snapshot: None,
+ restore_snapshot_id: None,
commit_user,
postpone_write_id: 0,
bucket_assigner:
BucketAssignerEnum::Constant(ConstantBucketAssigner::new(
@@ -511,6 +518,7 @@ impl TableWrite {
changelog_file_compression,
partition_seq_cache: HashMap::new(),
sequence_snapshot: None,
+ restore_snapshot_id: None,
commit_user,
postpone_write_id: (uuid::Uuid::new_v4().as_u128() % i32::MAX as
u128) as i32,
bucket_assigner,
@@ -728,7 +736,12 @@ impl TableWrite {
let batch = self.validate_write_batch_schema(batch)?;
// Java filters row kinds before extracting Blob values. Ignored rows
// need neither valid descriptor bytes nor physical files.
- let batch = self.enrich_rowkind_batch(&batch)?;
+ let batch = self.generate_rowkind_batch(&batch)?;
+ let selection = self.rowkind_selection(&batch)?;
+ let batch = match &selection {
+ Some(rows) => take_rows(&batch, rows)?,
+ None => batch,
+ };
if batch.num_rows() == 0 {
return Ok(None);
}
@@ -1001,9 +1014,9 @@ impl TableWrite {
})
}
- fn enrich_rowkind_batch(&self, batch: &RecordBatch) -> Result<RecordBatch>
{
+ fn generate_rowkind_batch(&self, batch: &RecordBatch) ->
Result<RecordBatch> {
let Some(generator) = &self.row_kind_generator else {
- return self.filter_rowkind_batch(batch);
+ return Ok(batch.clone());
};
if batch
.schema()
@@ -1020,18 +1033,17 @@ impl TableWrite {
let kinds = (0..batch.num_rows())
.map(|row| generator.generate(batch, row).map(|kind|
kind.to_value()))
.collect::<Result<Vec<_>>>()?;
- let batch = Self::add_per_row_value_kind_column(batch, kinds)?;
- self.filter_rowkind_batch(&batch)
+ Self::add_per_row_value_kind_column(batch, kinds)
}
/// Both generated row kinds and explicit caller-provided kinds use the
/// same Java RowKindFilter before bucket assignment and changelog writing.
- fn filter_rowkind_batch(&self, batch: &RecordBatch) -> Result<RecordBatch>
{
+ fn rowkind_selection(&self, batch: &RecordBatch) ->
Result<Option<Vec<usize>>> {
let Some(filter) = &self.row_kind_filter else {
- return Ok(batch.clone());
+ return Ok(None);
};
let Some(column) = batch.column_by_name(VALUE_KIND_FIELD_NAME) else {
- return Ok(batch.clone());
+ return Ok(None);
};
let kinds = column
.as_any()
@@ -1047,7 +1059,7 @@ impl TableWrite {
keep_rows.push(row);
}
}
- take_rows(batch, &keep_rows)
+ Ok((keep_rows.len() != batch.num_rows()).then_some(keep_rows))
}
fn add_per_row_value_kind_column(
@@ -1159,10 +1171,19 @@ impl TableWrite {
Ok(())
}
- /// Write multiple Arrow RecordBatches.
- pub async fn write_arrow(&mut self, batches: &[RecordBatch]) -> Result<()>
{
+ /// Write Arrow batches, optionally using a bucket assigned upstream, as
+ /// in Java TableWrite.write(row, bucket). Dynamic HASH indexes observe
every
+ /// surviving key after RowKind filtering.
+ pub async fn write_arrow(
+ &mut self,
+ batches: &[RecordBatch],
+ bucket: Option<i32>,
+ ) -> Result<()> {
for batch in batches {
- self.write_arrow_batch(batch).await?;
+ match bucket {
+ Some(bucket) => self.write_precomputed_bucket(batch,
bucket).await?,
+ None => self.write_arrow_batch(batch).await?,
+ }
}
Ok(())
}
@@ -1179,17 +1200,12 @@ impl TableWrite {
}
}
- /// Close all writers and collect CommitMessages for use with TableCommit.
- /// Writers are cleared after this call, allowing the TableWrite to be
reused.
+ /// Flush all writers and collect CommitMessages for use with TableCommit.
+ /// Append writers are closed; primary-key writers retain their sequence
+ /// counters after handing off the prepared files, as Java does.
///
- /// The per-bucket primary-key sequence cache is invalidated here too: it
- /// memoizes `max_sequence_number + 1` scanned from the latest snapshot,
which
- /// this commit advances. Keeping it would make the next reuse cycle
re-seed
- /// from the pre-commit value and assign sequence numbers that overlap the
- /// just-written files, so the highest-sequence-wins merge would silently
drop
- /// the newer rows -- Java's `MergeTreeWriter` advances its counter across
- /// commits. (`sequence_snapshot` is pinned only by the postpone path,
which
- /// forbids reuse, so it is left untouched.)
+ /// Invalidate the snapshot sequence cache used to seed previously unseen
+ /// buckets. Already-open key-value writers keep their own advancing
counter.
pub async fn prepare_commit(&mut self) -> Result<Vec<CommitMessage>> {
if let Some(writer) = self.format_writer.as_mut() {
return writer.prepare_commit().await;
@@ -1412,6 +1428,15 @@ impl TableWrite {
bucket: i32,
partition_bytes: &[u8],
) -> Result<FileWriter> {
+ if !self.is_overwrite && self.sequence_snapshot.is_none() {
+ if let Some(id) = self.restore_snapshot_id {
+ self.sequence_snapshot = Some(if id == 0 {
+ None
+ } else {
+ Some(self.table.snapshot_manager().get_snapshot(id).await?)
+ });
+ }
+ }
// Lazily scan partition sequence numbers on first writer creation per
partition.
// Overwrite mode skips this — old data will be replaced, so seq
starts at 0.
if !self.is_overwrite &&
!self.partition_seq_cache.contains_key(partition_bytes) {
diff --git a/crates/paimon/src/table/table_write_precomputed.rs
b/crates/paimon/src/table/table_write_precomputed.rs
new file mode 100644
index 00000000..ca627c3e
--- /dev/null
+++ b/crates/paimon/src/table/table_write_precomputed.rs
@@ -0,0 +1,102 @@
+// 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 TableWrite.write(row, bucket) for groups assigned upstream.
+
+use super::{BucketAssignerEnum, CoreOptions, RecordBatch, Result, TableWrite};
+use crate::spec::{batch_to_serialized_bytes, EMPTY_SERIALIZED_ROW,
MAX_DYNAMIC_BUCKETS};
+
+fn invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+impl TableWrite {
+ pub(in crate::table) fn set_restore_snapshot(&mut self, snapshot_id: i64) {
+ self.restore_snapshot_id = Some(snapshot_id);
+ if let BucketAssignerEnum::Dynamic(assigner) = &mut
self.bucket_assigner {
+ assigner.set_restore_snapshot(snapshot_id);
+ }
+ // Java WriteRestore restores data files and indexes from the same
snapshot.
+ }
+
+ pub(super) async fn write_precomputed_bucket(
+ &mut self,
+ batch: &RecordBatch,
+ bucket: i32,
+ ) -> Result<()> {
+ self.ensure_active()?;
+ let dynamic = matches!(self.bucket_assigner,
BucketAssignerEnum::Dynamic(_));
+ let total_buckets =
CoreOptions::new(self.table.schema().options()).bucket();
+ if self.format_writer.is_some()
+ || matches!(self.bucket_assigner,
BucketAssignerEnum::CrossPartition(_))
+ || (!dynamic && total_buckets <= 0)
+ {
+ return Err(invalid(
+ "Precomputed bucket writes are only valid for HASH_FIXED or
HASH_DYNAMIC tables",
+ ));
+ }
+ let upper_bound = if dynamic {
+ MAX_DYNAMIC_BUCKETS
+ } else {
+ total_buckets
+ };
+ if !(0..upper_bound).contains(&bucket) {
+ return Err(invalid(format!(
+ "Bucket id must be between 0 and {}, but was {bucket}",
+ upper_bound - 1,
+ )));
+ }
+ let Some(batch) = self.normalize_write_batch(batch)? else {
+ return Ok(());
+ };
+ let partition_indices = self
+ .partition_keys
+ .iter()
+ .map(|name| {
+ self.write_fields
+ .iter()
+ .position(|field| field.name() == name)
+ .unwrap()
+ })
+ .collect::<Vec<_>>();
+ let partitions = if partition_indices.is_empty() {
+ vec![EMPTY_SERIALIZED_ROW.clone(); batch.num_rows()]
+ } else {
+ batch_to_serialized_bytes(&batch, &partition_indices,
&self.write_fields)?
+ };
+ let partition = &partitions[0];
+ if partitions.iter().any(|actual| actual != partition) {
+ return Err(invalid(
+ "A precomputed bucket group contained multiple partitions",
+ ));
+ }
+ self.written = true;
+ if let BucketAssignerEnum::Dynamic(assigner) = &mut
self.bucket_assigner {
+ if let Err(error) = assigner
+ .notify_precomputed_batch(&batch, partition, bucket)
+ .await
+ {
+ self.fail_write().await;
+ return Err(error);
+ }
+ }
+ self.write_bucket(partition.clone(), bucket, batch).await
+ }
+}
diff --git a/crates/paimon/src/table/write_builder.rs
b/crates/paimon/src/table/write_builder.rs
index 41712038..bbde338f 100644
--- a/crates/paimon/src/table/write_builder.rs
+++ b/crates/paimon/src/table/write_builder.rs
@@ -78,6 +78,39 @@ impl<'a> WriteBuilder<'a> {
}
}
+ /// Restore dynamic-bucket write state from the coordinator's snapshot.
+ /// Zero denotes an empty base; without this option the writer restores
+ /// latest. Data-file sequence numbers and HASH indexes share this base,
+ /// as in Java FileSystemWriteRestore. Overwrite writers start empty.
+ pub fn with_restore_snapshot(self, snapshot_id: i64) ->
crate::Result<Self> {
+ if snapshot_id < 0 {
+ return Err(crate::Error::DataInvalid {
+ message: "Restore snapshot id must not be negative".into(),
+ source: None,
+ });
+ }
+ match self.0 {
+ WriteBuilderKind::Paimon(mut builder)
+ if
crate::spec::CoreOptions::new(builder.table.schema().options()).bucket()
+ == -1
+ && !builder.table.schema().primary_keys().is_empty()
+ && builder
+ .table
+ .schema()
+ .partition_keys()
+ .iter()
+ .all(|key|
builder.table.schema().primary_keys().contains(key)) =>
+ {
+ builder.restore_snapshot_id = Some(snapshot_id);
+ Ok(Self(WriteBuilderKind::Paimon(builder)))
+ }
+ _ => Err(crate::Error::DataInvalid {
+ message: "Restore snapshots are only valid for HASH_DYNAMIC
tables".into(),
+ source: None,
+ }),
+ }
+ }
+
/// Share a memory budget across table writers created by `new_write`.
pub fn with_resources(self, resources: ResourceContext) -> Self {
match self.0 {
@@ -147,6 +180,7 @@ struct PaimonWriteBuilder<'a> {
commit_user: String,
overwrite: bool,
resources: Option<ResourceContext>,
+ restore_snapshot_id: Option<i64>,
}
impl<'a> PaimonWriteBuilder<'a> {
@@ -156,6 +190,7 @@ impl<'a> PaimonWriteBuilder<'a> {
commit_user: Uuid::new_v4().to_string(),
overwrite: false,
resources: None,
+ restore_snapshot_id: None,
}
}
@@ -216,11 +251,13 @@ impl<'a> PaimonWriteBuilder<'a> {
if let Some(resources) = &self.resources {
write = write.with_resources(resources.clone());
}
- Ok(if self.overwrite {
- write.with_overwrite()
- } else {
- write
- })
+ if self.overwrite {
+ write = write.with_overwrite();
+ }
+ if let Some(snapshot_id) = self.restore_snapshot_id {
+ write.set_restore_snapshot(snapshot_id);
+ }
+ Ok(write)
}
/// Create a configured table update object.
diff --git a/crates/paimon/tests/precomputed_bucket_write_test.rs
b/crates/paimon/tests/precomputed_bucket_write_test.rs
new file mode 100644
index 00000000..5a2f58fe
--- /dev/null
+++ b/crates/paimon/tests/precomputed_bucket_write_test.rs
@@ -0,0 +1,731 @@
+// 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 std::sync::Arc;
+
+use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
+use common::incremental_helpers::{
+ make_batch, make_batch_with_kinds, memory_table, persist_table_schema,
pk_schema, setup_dirs,
+};
+use futures::TryStreamExt;
+use paimon::spec::{
+ batch_hash_codes, DataType, IndexManifest, IntType, Schema, TableSchema,
VarCharType,
+};
+use paimon::table::{CommitMessage, Table};
+
+async fn table(options: &[(&str, &str)]) -> Table {
+ let (io, table) = memory_table("memory:/precomputed", pk_schema(options));
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ table
+}
+
+fn hashes(table: &Table, batch: &RecordBatch) -> Vec<i32> {
+ batch_hash_codes(batch, &[0], table.schema().fields()).unwrap()
+}
+
+async fn rows(table: &Table) -> Vec<(i32, i32)> {
+ let builder = table.new_read_builder();
+ let plan = builder.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = builder
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ let ids = batch
+ .column_by_name("id")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let values = batch
+ .column_by_name("value")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ rows.extend((0..batch.num_rows()).map(|row| (ids.value(row),
values.value(row))));
+ }
+ rows.sort();
+ rows
+}
+
+async fn assert_index_hashes(table: &Table, bucket: i32, expected: &[i32]) {
+ let manager = table.snapshot_manager();
+ let snapshot = manager.get_latest_snapshot().await.unwrap().unwrap();
+ let path = format!(
+ "{}/{}",
+ manager.manifest_dir(),
+ snapshot.index_manifest().unwrap()
+ );
+ let entries = IndexManifest::read(table.file_io(), &path).await.unwrap();
+ let entries: Vec<_> = entries
+ .iter()
+ .filter(|e| e.index_file.index_type == "HASH" && e.bucket == bucket)
+ .collect();
+ assert_eq!(entries.len(), 1);
+ assert_eq!(entries[0].index_file.row_count, expected.len() as i64);
+ let path = entries[0]
+ .index_file
+ .external_path
+ .clone()
+ .unwrap_or_else(|| {
+ format!(
+ "{}/index/{}",
+ table.location(),
+ entries[0].index_file.file_name
+ )
+ });
+ let bytes = table
+ .file_io()
+ .new_input(&path)
+ .unwrap()
+ .read()
+ .await
+ .unwrap();
+ let mut actual = bytes
+ .as_chunks::<4>()
+ .0
+ .iter()
+ .map(|bytes| i32::from_be_bytes(*bytes))
+ .collect::<Vec<_>>();
+ let mut expected = expected.to_vec();
+ actual.sort();
+ expected.sort();
+ assert_eq!(actual, expected);
+}
+
+async fn commit(table: &Table, messages: Vec<CommitMessage>) {
+ table
+ .new_write_builder()
+ .new_commit()
+ .commit(messages)
+ .await
+ .unwrap();
+}
+
+#[tokio::test]
+async fn fixed_bucket_write_uses_upstream_bucket_and_shared_sequence() {
+ let table = table(&[("bucket", "4"), ("changelog-producer",
"input")]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![10])),
+ Some(3),
+ )
+ .await
+ .unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1, 2], vec![20, 30])),
+ Some(3),
+ )
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, 3);
+ // Deduplication discards id=1's first version, while its changelog
remains.
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 1);
+ assert_eq!(messages[0].new_files[0].max_sequence_number, 2);
+ assert!(!messages[0].new_changelog_files.is_empty());
+ assert!(messages[0].new_index_files.is_empty());
+ commit(&table, messages).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![40])),
+ Some(3),
+ )
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 3);
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 40), (2, 30)]);
+}
+
+#[tokio::test]
+async fn dynamic_upstream_bucket_survives_mixed_writes_and_checkpoints() {
+ let table = table(&[("bucket", "-1"), ("dynamic-bucket.target-row-num",
"1")]).await;
+ let mut writer = table
+ .new_write_builder()
+ .with_restore_snapshot(0)
+ .unwrap()
+ .new_write()
+ .unwrap();
+ let first = make_batch(vec![1], vec![10]);
+ let first_hashes = hashes(&table, &first);
+ writer
+ .write_arrow(std::slice::from_ref(&first), Some(17))
+ .await
+ .unwrap();
+ // Ordinary assignment must recognize a mapping notified by an upstream
group.
+ writer
+ .write_arrow_batch(&make_batch(vec![1], vec![20]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, 17);
+ commit(&table, messages).await;
+ assert_index_hashes(&table, 17, &first_hashes).await;
+ let second = make_batch(vec![1, 2], vec![30, 40]);
+ let second_hashes = hashes(&table, &second);
+ writer
+ .write_arrow(std::slice::from_ref(&second), Some(17))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_index_files[0].row_count, 2);
+ commit(&table, messages).await;
+ assert_index_hashes(&table, 17, &second_hashes).await;
+ assert_eq!(rows(&table).await, vec![(1, 30), (2, 40)]);
+ let mut restarted = table.new_write_builder().new_write().unwrap();
+ restarted
+ .write_arrow_batch(&make_batch(vec![1, 2], vec![50, 60]))
+ .await
+ .unwrap();
+ let messages = restarted.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, 17);
+ assert!(messages[0].new_index_files.is_empty());
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 50), (2, 60)]);
+}
+
+#[tokio::test]
+async fn repeated_hashes_do_not_rewrite_unchanged_index() {
+ let table = table(&[("bucket", "-1")]).await;
+ let batch = make_batch(vec![1], vec![10]);
+ let key_hashes = hashes(&table, &batch);
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(std::slice::from_ref(&batch), Some(7))
+ .await
+ .unwrap();
+ commit(&table, writer.prepare_commit().await.unwrap()).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(std::slice::from_ref(&batch), Some(7))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert!(messages[0].new_index_files.is_empty());
+ commit(&table, messages).await;
+ assert_index_hashes(&table, 7, &key_hashes).await;
+}
+
+#[tokio::test]
+async fn row_kind_filter_precedes_index_hash_computation() {
+ let table = table(&[
+ ("bucket", "-1"),
+ ("ignore-delete", "true"),
+ ("changelog-producer", "input"),
+ ])
+ .await;
+ // Ignored deletes must not suppress index notification for surviving
inserts.
+ let batch = make_batch_with_kinds(vec![1, 1, 2, 2], vec![0, 10, 0, 20],
vec![3, 0, 3, 0]);
+ let key_hashes = hashes(&table, &batch);
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(std::slice::from_ref(&batch), Some(9))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_index_files[0].row_count, 2);
+ commit(&table, messages).await;
+ assert_index_hashes(&table, 9, &[key_hashes[1], key_hashes[3]]).await;
+ assert_eq!(rows(&table).await, vec![(1, 10), (2, 20)]);
+}
+
+#[tokio::test]
+async fn
partition_group_and_projected_write_type_are_validated_before_staging() {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .column("pt", DataType::VarChar(VarCharType::string_type()))
+ .partition_keys(["pt"])
+ .option("bucket", "4")
+ .option("bucket-key", "id")
+ .build()
+ .unwrap();
+ let (io, table) = memory_table("memory:/precomputed_partial",
TableSchema::new(0, &schema));
+ setup_dirs(&io, table.location()).await;
+ persist_table_schema(&io, table.location(), table.schema()).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ let batch = RecordBatch::try_from_iter([
+ (
+ "pt",
+ Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
+ ),
+ ("id", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef),
+ ])
+ .unwrap();
+ let error = writer
+ .write_arrow(std::slice::from_ref(&batch), Some(2))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("multiple partitions"));
+ // Rejected input did not freeze the write type or stage half a group.
+ writer
+ .with_write_type(vec!["pt".into(), "id".into()])
+ .unwrap();
+ writer
+ .write_arrow(std::slice::from_ref(&batch.slice(0, 1)), Some(2))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, 2);
+ assert_eq!(messages[0].new_files[0].row_count, 1);
+ assert_eq!(
+ messages[0].new_files[0].write_cols,
+ Some(vec!["pt".into(), "id".into()])
+ );
+}
+
+#[tokio::test]
+async fn configuration_and_bucket_errors_do_not_stage_data() {
+ for bucket in ["1", "4", "-1"] {
+ let table = table(&[("bucket", bucket)]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ let batch = make_batch(vec![1, 2], vec![10, 20]);
+ for invalid in [
+ -1,
+ if bucket == "-1" {
+ 32768
+ } else {
+ bucket.parse().unwrap()
+ },
+ ] {
+ assert!(writer
+ .write_arrow(std::slice::from_ref(&batch), Some(invalid))
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("Bucket id"));
+ }
+ if bucket != "-1" {
+ assert!(table
+ .new_write_builder()
+ .with_restore_snapshot(0)
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("HASH_DYNAMIC"));
+ } else {
+
assert!(table.new_write_builder().with_restore_snapshot(-1).is_err());
+ let mut missing = table
+ .new_write_builder()
+ .with_restore_snapshot(99)
+ .unwrap()
+ .new_write()
+ .unwrap();
+ assert!(missing
+ .write_arrow(std::slice::from_ref(&batch), Some(0))
+ .await
+ .is_err());
+ missing.close().await;
+ }
+ assert!(writer.prepare_commit().await.unwrap().is_empty());
+ }
+}
+
+#[tokio::test]
+async fn builder_configuration_is_copied_to_created_writers() {
+ let table = table(&[("bucket", "-1")]).await;
+ let builder = table.new_write_builder().with_restore_snapshot(0).unwrap();
+ let mut writer = builder.new_write().unwrap();
+ let builder = builder.with_restore_snapshot(99).unwrap();
+ writer
+ .write_arrow(std::slice::from_ref(&make_batch(vec![], vec![])),
Some(0))
+ .await
+ .unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![10])),
+ Some(0),
+ )
+ .await
+ .unwrap();
+ assert_eq!(writer.prepare_commit().await.unwrap().len(), 1);
+ let mut missing = builder.new_write().unwrap();
+ assert!(missing
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![2], vec![20])),
+ Some(0)
+ )
+ .await
+ .is_err());
+}
+
+#[tokio::test]
+async fn pinned_base_does_not_silently_follow_later_snapshots() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![10])),
+ Some(0),
+ )
+ .await
+ .unwrap();
+ commit(&table, writer.prepare_commit().await.unwrap()).await;
+ let mut pinned = table
+ .new_write_builder()
+ .with_restore_snapshot(1)
+ .unwrap()
+ .new_write()
+ .unwrap();
+ let second = make_batch(vec![2], vec![20]);
+ writer
+ .write_arrow(std::slice::from_ref(&second), Some(0))
+ .await
+ .unwrap();
+ commit(&table, writer.prepare_commit().await.unwrap()).await;
+ pinned
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![3], vec![30])),
+ Some(0),
+ )
+ .await
+ .unwrap();
+ let messages = pinned.prepare_commit().await.unwrap();
+ // Java restores both data-file sequences and HASH indexes from snapshot 1.
+ assert_eq!(messages[0].new_index_files[0].row_count, 2);
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 1);
+ table
+ .new_write_builder()
+ .new_commit()
+ .abort(&messages)
+ .await
+ .unwrap();
+ assert_eq!(rows(&table).await, vec![(1, 10), (2, 20)]);
+}
+
+#[tokio::test]
+async fn
restore_snapshot_pins_data_file_sequences_and_default_restores_latest() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut first = table.new_write_builder().new_write().unwrap();
+ first
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![10])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ commit(&table, first.prepare_commit().await.unwrap()).await;
+ let mut pinned = table
+ .new_write_builder()
+ .with_restore_snapshot(1)
+ .unwrap()
+ .new_write()
+ .unwrap();
+ first
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1, 1, 1], vec![20, 30, 40])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ commit(&table, first.prepare_commit().await.unwrap()).await;
+ pinned
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![50])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ let messages = pinned.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 1);
+ // The supplied historical base is a coordinated restore, not permission to
+ // overwrite another owner's newer writes. Inspect and abort these
messages.
+ table
+ .new_write_builder()
+ .new_commit()
+ .abort(&messages)
+ .await
+ .unwrap();
+ pinned.close().await;
+ assert_eq!(rows(&table).await, vec![(1, 40)]);
+ let mut current = table.new_write_builder().new_write().unwrap();
+ current
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![50])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ let messages = current.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 4);
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 50)]);
+}
+
+#[tokio::test]
+async fn conflicting_upstream_mapping_never_publishes_a_second_bucket() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ let batch = make_batch(vec![1], vec![10]);
+ writer
+ .write_arrow(std::slice::from_ref(&batch), Some(7))
+ .await
+ .unwrap();
+ let error = writer
+ .write_arrow(std::slice::from_ref(&batch), Some(8))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("belongs to bucket 7"));
+ assert!(writer.prepare_commit().await.is_err());
+ assert!(table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .is_none());
+}
+
+#[tokio::test]
+async fn overwrite_ignores_old_hashes_and_sequence_numbers() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut first = table.new_write_builder().new_write().unwrap();
+ first
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1, 2], vec![10, 20])),
+ Some(5),
+ )
+ .await
+ .unwrap();
+ commit(&table, first.prepare_commit().await.unwrap()).await;
+ let builder = table
+ .new_write_builder()
+ .with_overwrite()
+ .with_restore_snapshot(1)
+ .unwrap();
+ let mut writer = builder.new_write().unwrap();
+ let batch = make_batch(vec![3], vec![30]);
+ let key_hashes = hashes(&table, &batch);
+ writer
+ .write_arrow(std::slice::from_ref(&batch), Some(5))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_files[0].min_sequence_number, 0);
+ assert_eq!(messages[0].new_index_files[0].row_count, 1);
+ builder
+ .new_commit()
+ .overwrite(messages, None)
+ .await
+ .unwrap();
+ assert_index_hashes(&table, 5, &key_hashes).await;
+ assert_eq!(rows(&table).await, vec![(3, 30)]);
+}
+
+#[tokio::test]
+async fn prepared_checkpoints_keep_sequence_progress_before_publication() {
+ let table = table(&[("bucket", "-1")]).await;
+ let builder = table.new_write_builder().with_restore_snapshot(0).unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![2, 1], vec![20, 10])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ let first = writer.prepare_commit().await.unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![1], vec![30])),
+ Some(7),
+ )
+ .await
+ .unwrap();
+ let second = writer.prepare_commit().await.unwrap();
+ assert_eq!(first[0].new_files[0].max_sequence_number, 1);
+ assert_eq!(second[0].new_files[0].min_sequence_number, 2);
+ builder
+ .new_commit()
+ .commit_with_identifier(first, 1)
+ .await
+ .unwrap();
+ builder
+ .new_commit()
+ .commit_with_identifier(second, 2)
+ .await
+ .unwrap();
+ assert_eq!(rows(&table).await, vec![(1, 30), (2, 20)]);
+}
+
+#[tokio::test]
+async fn configured_overwrite_can_use_ordinary_and_direct_assignment() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut first = table.new_write_builder().new_write().unwrap();
+ first
+ .write_arrow_batch(&make_batch(vec![1, 2], vec![10, 20]))
+ .await
+ .unwrap();
+ commit(&table, first.prepare_commit().await.unwrap()).await;
+ let builder = table
+ .new_write_builder()
+ .with_overwrite()
+ .with_restore_snapshot(1)
+ .unwrap();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![3], vec![30]))
+ .await
+ .unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![3, 4], vec![40, 50])),
+ Some(0),
+ )
+ .await
+ .unwrap();
+ builder
+ .new_commit()
+ .overwrite(writer.prepare_commit().await.unwrap(), None)
+ .await
+ .unwrap();
+ assert_eq!(rows(&table).await, vec![(3, 40), (4, 50)]);
+}
+
+#[tokio::test]
+async fn precomputed_writer_restores_only_its_bucket() {
+ let table = table(&[("bucket", "-1")]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ for (id, bucket) in [(1, 7), (2, 8)] {
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(vec![id], vec![id * 10])),
+ Some(bucket),
+ )
+ .await
+ .unwrap();
+ }
+ let messages = writer.prepare_commit().await.unwrap();
+ let other_index = messages
+ .iter()
+ .find(|message| message.bucket == 8)
+ .unwrap()
+ .new_index_files[0]
+ .file_name
+ .clone();
+ commit(&table, messages).await;
+ // An unrelated bucket's payload must never be opened by a direct writer.
+ // Ordinary assignment still needs all hashes and detects the corruption.
+ table
+ .file_io()
+ .new_output(&format!("{}/index/{other_index}", table.location()))
+ .unwrap()
+ .write(bytes::Bytes::from_static(b"bad"))
+ .await
+ .unwrap();
+ let mut direct = table.new_write_builder().new_write().unwrap();
+ let batch = make_batch(vec![1, 3], vec![100, 30]);
+ direct
+ .write_arrow(std::slice::from_ref(&batch), Some(7))
+ .await
+ .unwrap();
+ let messages = direct.prepare_commit().await.unwrap();
+ assert_eq!(messages[0].new_index_files[0].row_count, 2);
+ let mut assign = table.new_write_builder().new_write().unwrap();
+ assert!(assign
+ .write_arrow_batch(&batch)
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("Corrupt HASH index"));
+ table
+ .new_write_builder()
+ .new_commit()
+ .abort(&messages)
+ .await
+ .unwrap();
+}
+
+#[tokio::test]
+async fn bucket_local_index_paths_survive_precomputed_writer_restart() {
+ let table = table(&[
+ ("bucket", "-1"),
+ ("index-file-in-data-file-dir", "true"),
+ ("data-file.path-directory", "data/nested"),
+ ])
+ .await;
+ for ids in [vec![1, 2], vec![2, 3]] {
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow(
+ std::slice::from_ref(&make_batch(ids.clone(), ids)),
+ Some(19),
+ )
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ let meta = &messages[0].new_index_files[0];
+ let path = format!(
+ "{}/data/nested/bucket-19/{}",
+ table.location(),
+ meta.file_name
+ );
+ assert!(table.file_io().exists(&path).await.unwrap());
+ commit(&table, messages).await;
+ }
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer
+ .write_arrow_batch(&make_batch(vec![1, 2, 3], vec![10, 20, 30]))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, 19);
+ assert!(messages[0].new_index_files.is_empty());
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 10), (2, 20), (3, 30)]);
+}
+
+#[tokio::test]
+async fn unified_arrow_writes_multiple_batches_with_or_without_bucket() {
+ for bucket in [None, Some(0)] {
+ let table = table(&[("bucket", "4")]).await;
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ let batches = vec![
+ make_batch(vec![1, 2], vec![10, 20]),
+ make_batch(vec![1], vec![30]),
+ ];
+ writer.write_arrow(&batches, bucket).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ if let Some(bucket) = bucket {
+ assert_eq!(messages.len(), 1);
+ assert_eq!(messages[0].bucket, bucket);
+ }
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 30), (2, 20)]);
+ }
+}
