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 77728893 feat: support native shard column updates in core (#1099)
77728893 is described below
commit 77728893e9dd1a6dd553253a01316ba0e6b3808d
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Oct 11 13:33:35 2026 +0800
feat: support native shard column updates in core (#1099)
---
bindings/python/README.md | 19 +
.../python/python/pypaimon_rust/datafusion.pyi | 28 +
bindings/python/src/context.rs | 1 +
bindings/python/src/lib.rs | 1 +
bindings/python/src/read.rs | 8 +
bindings/python/src/shard_update.rs | 80 +++
bindings/python/src/write.rs | 52 +-
crates/paimon/src/table/data_evolution_writer.rs | 10 +
crates/paimon/src/table/mod.rs | 2 +
crates/paimon/src/table/shard_table_updater.rs | 437 +++++++++++++
.../paimon/src/table/shard_table_updater/tests.rs | 714 +++++++++++++++++++++
crates/paimon/src/table/table_update.rs | 28 +
12 files changed, 1378 insertions(+), 2 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index b4710b85..76ea23d8 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -132,6 +132,25 @@ print(f"\nRead: {batches_tt[0].num_rows} rows")
print(batches_tt[0])
```
+### Shard column updates
+
+The batch and stream update objects expose
+`new_shard_updator(shard_num, total_shard_count, read_type=...)`. The optional
+`read_type` is a Paimon ROW type encoded as JSON, preserving field IDs and
nested
+read metadata. It belongs to this shard operation and is independent of columns
+configured by `with_update_type`. Omitting it reads all table fields; an empty
+ROW reads zero columns with row counts preserved. Core captures one snapshot
and
+balances complete Data Evolution file groups by row ID. The returned
+`arrow_reader()` yields projected batches; `update_by_arrow_batch()` receives
+replacement columns in that reader's visible row order, including when deletion
+vectors remove physical rows. Input batch boundaries may differ from read
batches.
+
+Finish the reader and supply exactly the same number of rows before calling
+`prepare_commit()`. Core buffers one file group, reuses the normal
partial-column
+writer and returns messages with the captured snapshot's conflict check.
Repeated
+preparation reuses those messages. `close()` preserves prepared files; commit
the
+messages with the matching builder's committer.
+
### Prepared MERGE operations
Batch and stream update objects expose `merge_into` to prepare data-evolution
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index a0c81bfa..585f0f71 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -334,6 +334,15 @@ class BatchTableUpdate:
def update_by_arrow_batches_with_row_id(
self, tables: Iterable[pyarrow.Table]
) -> List[CommitMessage]: ...
+ def new_shard_updator(
+ self, shard_num: int, total_shard_count: int, *, read_type:
Optional[str] = None
+ ) -> ShardTableUpdator:
+ """Create a shard operation with a Paimon ROW read type encoded as
JSON.
+
+ None reads all table fields; an empty ROW reads zero columns while
+ preserving row counts. Read fields are independent of with_update_type.
+ """
+ ...
def with_update_type(self, update_cols: List[str]) -> "BatchTableUpdate":
...
def delete_by_row_id(self, row_ids: Sequence[int]) -> List[CommitMessage]:
...
def upsert_by_arrow_with_key(
@@ -377,6 +386,15 @@ class StreamTableUpdate:
def update_by_arrow_with_row_id(
self, table: pyarrow.Table, commit_identifier: int
) -> List[CommitMessage]: ...
+ def new_shard_updator(
+ self, shard_num: int, total_shard_count: int, *, read_type:
Optional[str] = None
+ ) -> ShardTableUpdator:
+ """Create a shard operation with a Paimon ROW read type encoded as
JSON.
+
+ None reads all table fields; an empty ROW reads zero columns while
+ preserving row counts. Read fields are independent of with_update_type.
+ """
+ ...
def with_update_type(self, update_cols: List[str]) -> "StreamTableUpdate":
...
def delete_by_row_id(
self, row_ids: Sequence[int], commit_identifier: int
@@ -660,3 +678,13 @@ class FullTextRead:
class FullTextSearchResult:
def __len__(self) -> int: ...
def row_ids(self) -> Dict[int, float]: ...
+
+class ShardTableUpdator:
+ def arrow_reader(self) -> RecordBatchReader:
+ """Read one captured shard once; values supplied later follow this
reader's order."""
+ ...
+ def update_by_arrow_batch(self, data: pyarrow.RecordBatch) -> None: ...
+ def prepare_commit(self) -> List[CommitMessage]:
+ """Require a finished read and exactly as many update rows; repeated
calls reuse messages."""
+ ...
+ def close(self) -> None: ...
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index ba8c3d87..117dc05e 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -760,6 +760,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::write::PyBatchTableWrite>()?;
this.add_class::<crate::write::PyBatchTableUpdate>()?;
this.add_class::<crate::write::PyTableUpdateByRowId>()?;
+ this.add_class::<crate::shard_update::PyShardTableUpdator>()?;
this.add_class::<crate::write::PyStreamTableUpdate>()?;
this.add_class::<crate::write::PyStreamTableWrite>()?;
this.add_class::<crate::write::PyBatchTableCommit>()?;
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index 87a3362c..86e7341b 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -29,6 +29,7 @@ mod oss_cpp;
mod predicate;
mod read;
mod schema;
+mod shard_update;
mod table;
mod udf;
mod update_assignment;
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index deb87eca..e101e93b 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -776,6 +776,14 @@ pub struct PyRecordBatchReader {
}
impl PyRecordBatchReader {
+ pub(crate) fn from_stream(stream: ArrowRecordBatchStream) -> Self {
+ Self {
+ stream: Mutex::new(Some(stream)),
+ closed: AtomicBool::new(false),
+ close_notify: Notify::new(),
+ }
+ }
+
fn next_record_batch(&self) -> PyResult<Option<RecordBatch>> {
let mut stream = self
.stream
diff --git a/bindings/python/src/shard_update.rs
b/bindings/python/src/shard_update.rs
new file mode 100644
index 00000000..dade9828
--- /dev/null
+++ b/bindings/python/src/shard_update.rs
@@ -0,0 +1,80 @@
+// 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.
+
+//! Arrow transport for core positional shard updates.
+
+use arrow::pyarrow::FromPyArrow;
+use arrow::record_batch::RecordBatch;
+use paimon::table::ShardTableUpdater;
+use paimon_datafusion::runtime::runtime;
+use pyo3::prelude::*;
+
+use crate::error::to_py_err;
+use crate::read::PyRecordBatchReader;
+use crate::write::{wrap_messages, PyCommitMessage};
+
+#[pyclass(name = "ShardTableUpdator", module = "pypaimon_rust.datafusion")]
+pub struct PyShardTableUpdator {
+ inner: ShardTableUpdater,
+ table_location: String,
+ commit_user: String,
+}
+
+impl PyShardTableUpdator {
+ pub(crate) fn new(
+ inner: ShardTableUpdater,
+ table_location: String,
+ commit_user: String,
+ ) -> Self {
+ Self {
+ inner,
+ table_location,
+ commit_user,
+ }
+ }
+}
+
+#[pymethods]
+impl PyShardTableUpdator {
+ fn arrow_reader(&mut self) -> PyResult<PyRecordBatchReader> {
+ self.inner
+ .arrow_reader()
+ .map(PyRecordBatchReader::from_stream)
+ .map_err(to_py_err)
+ }
+
+ fn update_by_arrow_batch(&mut self, py: Python<'_>, data: &Bound<'_,
PyAny>) -> PyResult<()> {
+ let batch = RecordBatch::from_pyarrow_bound(data)?;
+ py.detach(||
runtime().block_on(self.inner.update_by_arrow_batch(batch)))
+ .map_err(to_py_err)
+ }
+
+ fn prepare_commit(&mut self, py: Python<'_>) ->
PyResult<Vec<PyCommitMessage>> {
+ let messages = py
+ .detach(|| runtime().block_on(self.inner.prepare_commit()))
+ .map_err(to_py_err)?;
+ Ok(wrap_messages(
+ messages,
+ &self.table_location,
+ &self.commit_user,
+ ))
+ }
+
+ fn close(&mut self) {
+ self.inner.close();
+ }
+}
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index 6f6901d4..7f1d0fd8 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -21,7 +21,7 @@ use std::sync::{Arc, Mutex};
use arrow::datatypes::Schema as ArrowSchema;
use arrow::pyarrow::FromPyArrow;
use arrow::record_batch::RecordBatch;
-use paimon::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataType, Datum};
+use paimon::spec::{BlobConsumer, BlobDescriptor, CoreOptions, DataType, Datum,
RowType};
use paimon::table::{
CommitMessage, PostponeBucketPlan, PostponeFixedBucketTableWrite, Table,
TableCommit,
TableUpdate, TableUpdateByRowId, TableWrite,
COMMIT_MESSAGE_SERIALIZER_VERSION,
@@ -516,7 +516,7 @@ impl BlobConsumer for PythonBlobConsumer {
}
}
-fn wrap_messages(
+pub(crate) fn wrap_messages(
messages: Vec<CommitMessage>,
table_location: &str,
commit_user: &str,
@@ -736,6 +736,30 @@ impl UpdateContext {
))
}
+ fn new_shard_updator(
+ &self,
+ py: Python<'_>,
+ shard: usize,
+ count: usize,
+ read_type: Option<&str>,
+ ) -> PyResult<crate::shard_update::PyShardTableUpdator> {
+ let read_type = read_type
+ .map(|json| {
+ serde_json::from_str::<RowType>(json)
+ .map(|row| row.fields().to_vec())
+ .map_err(|e| PyValueError::new_err(format!("invalid Paimon
read type: {e}")))
+ })
+ .transpose()?;
+ let inner = py
+ .detach(|| runtime().block_on(self.inner.new_shard_updater(shard,
count, read_type)))
+ .map_err(to_py_err)?;
+ Ok(crate::shard_update::PyShardTableUpdator::new(
+ inner,
+ self.table_location.clone(),
+ self.commit_user.clone(),
+ ))
+ }
+
fn new_update_by_row_id(&self, py: Python<'_>) ->
PyResult<PyTableUpdateByRowId> {
let inner = py
.detach(|| runtime().block_on(self.inner.new_update_by_row_id()))
@@ -988,6 +1012,18 @@ impl PyStreamTableUpdate {
self.context.new_update_by_row_id(py)
}
+ #[pyo3(signature = (shard_num, total_shard_count, *, read_type=None))]
+ fn new_shard_updator(
+ &self,
+ py: Python<'_>,
+ shard_num: usize,
+ total_shard_count: usize,
+ read_type: Option<&str>,
+ ) -> PyResult<crate::shard_update::PyShardTableUpdator> {
+ self.context
+ .new_shard_updator(py, shard_num, total_shard_count, read_type)
+ }
+
fn with_update_type<'py>(
mut slf: PyRefMut<'py, Self>,
update_cols: Vec<String>,
@@ -1079,6 +1115,18 @@ impl PyBatchTableUpdate {
self.context.new_update_by_row_id(py)
}
+ #[pyo3(signature = (shard_num, total_shard_count, *, read_type=None))]
+ fn new_shard_updator(
+ &self,
+ py: Python<'_>,
+ shard_num: usize,
+ total_shard_count: usize,
+ read_type: Option<&str>,
+ ) -> PyResult<crate::shard_update::PyShardTableUpdator> {
+ self.context
+ .new_shard_updator(py, shard_num, total_shard_count, read_type)
+ }
+
fn with_update_type<'py>(
mut slf: PyRefMut<'py, Self>,
update_cols: Vec<String>,
diff --git a/crates/paimon/src/table/data_evolution_writer.rs
b/crates/paimon/src/table/data_evolution_writer.rs
index 05bf3ee5..a2069eb8 100644
--- a/crates/paimon/src/table/data_evolution_writer.rs
+++ b/crates/paimon/src/table/data_evolution_writer.rs
@@ -1316,6 +1316,16 @@ impl RowIdFileIndex {
})
}
+ /// Anchor of the normal-file group owning a visible row.
+ pub(super) fn first_row_id_for(&self, row_id: i64) -> Result<i64> {
+ find_owning_file(&self.files, row_id)
+ .map(|(_, file)| file.first_row_id)
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: format!("No file found for _ROW_ID {row_id}"),
+ source: None,
+ })
+ }
+
pub(super) fn matched_first_row_ids(&self, batches: &[RecordBatch]) ->
Result<HashSet<i64>> {
if batches.iter().all(|batch| batch.num_rows() == 0) {
return Ok(HashSet::new());
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index ce1b01f9..cf1f164a 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -137,6 +137,7 @@ mod row_position_selection;
mod row_sidecar;
mod scan_trace;
pub(crate) mod schema_manager;
+mod shard_table_updater;
pub(crate) mod snapshot_commit;
mod snapshot_manager;
mod sort_merge;
@@ -216,6 +217,7 @@ pub use read_builder::ReadBuilder;
pub use rest_env::RESTEnv;
pub use scan_trace::ScanTrace;
pub use schema_manager::SchemaManager;
+pub use shard_table_updater::ShardTableUpdater;
pub use snapshot_commit::{RESTSnapshotCommit, RenamingSnapshotCommit,
SnapshotCommit};
pub use snapshot_manager::SnapshotManager;
pub use source::{
diff --git a/crates/paimon/src/table/shard_table_updater.rs
b/crates/paimon/src/table/shard_table_updater.rs
new file mode 100644
index 00000000..0c4853be
--- /dev/null
+++ b/crates/paimon/src/table/shard_table_updater.rs
@@ -0,0 +1,437 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Positional column updates over a shard of complete Data Evolution file
groups.
+
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::sync::{Arc, Mutex};
+
+use arrow_array::{Array, ArrayRef, Int64Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use async_stream::try_stream;
+use futures::TryStreamExt;
+
+use crate::spec::{row_id_data_field, DataField};
+
+use super::data_evolution_writer::RowIdFileIndex;
+use super::source::is_data_evolution_normal_file;
+use super::update_input::{cast_update_value, unique_column_index, CastMode};
+use super::{ArrowRecordBatchStream, CommitMessage, DataEvolutionWriter,
DataSplit, Table};
+
+const ROW_ID: &str = "_ROW_ID";
+
+fn invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+#[derive(Default)]
+struct ReadState {
+ pending: VecDeque<i64>,
+ complete: bool,
+ failed: bool,
+ closed: bool,
+}
+
+/// Read and replace columns of one balanced shard without exposing row-ID
mapping
+/// to the caller. The read uses one captured snapshot. Updates follow reader
order
+/// and may use different batch boundaries; deleted rows never consume input
values.
+///
+/// Finish the reader and supply exactly one output row per visible input row
before
+/// preparing. Core reuses the row-ID update writer, preserving Java's
column-file
+/// metadata, original physical ranges and conflict checks. Buffer one file
group
+/// at a time; returned messages belong to the caller and are never removed on
close.
+pub struct ShardTableUpdater {
+ table: Table,
+ splits: Vec<DataSplit>,
+ read_type: Vec<DataField>,
+ write_schema: SchemaRef,
+ index: RowIdFileIndex,
+ writer: Option<DataEvolutionWriter>,
+ write_columns: Vec<String>,
+ current_group: Option<i64>,
+ finished_groups: HashSet<i64>,
+ staged: Vec<CommitMessage>,
+ blob_uri_reader_factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ state: Arc<Mutex<ReadState>>,
+ reader_started: bool,
+ failed: bool,
+ messages: Option<Vec<CommitMessage>>,
+}
+
+impl ShardTableUpdater {
+ pub(super) async fn new(
+ table: &Table,
+ read_type: Option<Vec<DataField>>,
+ write_columns: Vec<String>,
+ shard: usize,
+ shard_count: usize,
+ ) -> crate::Result<Self> {
+ if shard_count == 0 || shard >= shard_count {
+ return Err(invalid(
+ "Shard index must be smaller than a positive shard count",
+ ));
+ }
+ if write_columns.is_empty() {
+ return Err(invalid("Shard update columns must not be empty"));
+ }
+ let writer = DataEvolutionWriter::new(table, write_columns.clone())?;
+ table.schema().core_options().ensure_read_authorized()?;
+ let snapshot = super::time_travel::resolve_snapshot(table).await?;
+ let table = table
+ .copy_with_options(HashMap::from([(
+ "scalar-index.search-mode".into(),
+ "FULL".into(),
+ )]))
+ .copy_with_pinned_snapshot(snapshot.as_ref());
+ let plan = table
+ .new_read_builder()
+ .new_scan()
+ .with_scan_all_files()
+ .plan()
+ .await?;
+ let splits = select_shard(plan.splits(), shard, shard_count)?;
+ let index = RowIdFileIndex::from_splits(table.clone(), &splits)?;
+ let read_type = read_type.unwrap_or_else(||
table.schema().fields().to_vec());
+ // Validate the same canonical request used by ordinary table readers.
+ table
+ .new_read_builder()
+ .with_read_type(with_row_id(&read_type)?)
+ .new_read()?;
+ let fields = super::data_evolution_fields::project_by_paths(
+ table.schema().fields(),
+ &write_columns,
+ )?;
+ let write_schema = crate::arrow::build_target_arrow_schema(&fields)?;
+ let state = ReadState {
+ complete: splits.is_empty(),
+ ..Default::default()
+ };
+ Ok(Self {
+ table,
+ splits,
+ read_type,
+ write_schema,
+ index,
+ writer: Some(writer),
+ write_columns,
+ current_group: None,
+ finished_groups: HashSet::new(),
+ staged: Vec::new(),
+ blob_uri_reader_factory: None,
+ state: Arc::new(Mutex::new(state)),
+ reader_started: false,
+ failed: false,
+ messages: None,
+ })
+ }
+
+ fn ensure_active(&self) -> crate::Result<()> {
+ let state = self
+ .state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?;
+ if state.closed || state.failed || self.failed ||
self.messages.is_some() {
+ return Err(invalid(
+ "Shard updater is closed, failed or already prepared",
+ ));
+ }
+ Ok(())
+ }
+
+ pub(super) fn with_blob_uri_reader_factory(
+ &mut self,
+ factory: Option<Arc<dyn crate::io::UriReaderFactory>>,
+ ) {
+ self.writer
+ .as_mut()
+ .unwrap()
+ .with_blob_uri_reader_factory(factory.clone());
+ self.blob_uri_reader_factory = factory;
+ }
+
+ /// Open the shard's single reader. Only requested columns are returned;
+ /// `_ROW_ID` is retained internally even when absent from the read type.
+ pub fn arrow_reader(&mut self) -> crate::Result<ArrowRecordBatchStream> {
+ self.ensure_active()?;
+ if self.reader_started {
+ return Err(invalid("Shard reader can only be opened once"));
+ }
+ let mut builder = self.table.new_read_builder();
+ builder.with_read_type(with_row_id(&self.read_type)?);
+ let mut input = builder.new_read()?.to_arrow(&self.splits)?;
+ let state = self.state.clone();
+ let output_columns = self.read_type.len();
+ self.reader_started = true;
+ Ok(Box::pin(try_stream! {
+ loop {
+ let next = input.try_next().await;
+ if next.is_err() { finish_read(&state, true)?; }
+ let Some(batch) = next? else { finish_read(&state, false)?;
break; };
+ let output = record_rows(&state, batch, output_columns);
+ if output.is_err() { finish_read(&state, true)?; }
+ yield output?;
+ }
+ }))
+ }
+
+ /// Supply column values in visible reader order. Validate the complete
batch
+ /// before consuming any row IDs, so a malformed batch can be corrected.
+ pub async fn update_by_arrow_batch(&mut self, data: RecordBatch) ->
crate::Result<()> {
+ self.ensure_active()?;
+ if data.num_columns() != self.write_schema.fields().len() {
+ return Err(invalid(
+ "Shard input must contain exactly the configured update
columns",
+ ));
+ }
+ let mut columns = Vec::<ArrayRef>::new();
+ for field in self.write_schema.fields() {
+ let index = unique_column_index(data.schema().as_ref(),
field.name())?;
+ columns.push(cast_update_value(
+ data.column(index),
+ field.data_type(),
+ CastMode::Assignment,
+ )?);
+ }
+ let ids = {
+ let state = self
+ .state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?;
+ if data.num_rows() > state.pending.len() {
+ return Err(invalid(
+ "Shard update contains more rows than have been read",
+ ));
+ }
+ state
+ .pending
+ .iter()
+ .take(data.num_rows())
+ .copied()
+ .collect::<Vec<_>>()
+ };
+ // Resolve all group boundaries before staging or consuming this input.
+ let groups = ids
+ .iter()
+ .map(|id| self.index.first_row_id_for(*id))
+ .collect::<crate::Result<Vec<_>>>()?;
+ let mut arrays = vec![Arc::new(Int64Array::from(ids)) as ArrayRef];
+ arrays.extend(columns);
+ let mut fields = vec![Arc::new(Field::new(ROW_ID, DataType::Int64,
false))];
+ fields.extend(self.write_schema.fields().iter().cloned());
+ let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays)
+ .map_err(|error| invalid(error.to_string()))?;
+ let result = self.stage_groups(batch, &groups).await;
+ if result.is_err() {
+ self.failed = true;
+ }
+ result?;
+ self.state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?
+ .pending
+ .drain(..data.num_rows());
+ Ok(())
+ }
+
+ /// Prepare after complete consumption. Repeated calls return the same
messages.
+ pub async fn prepare_commit(&mut self) ->
crate::Result<Vec<CommitMessage>> {
+ if let Some(messages) = &self.messages {
+ return Ok(messages.clone());
+ }
+ self.ensure_active()?;
+ {
+ let state = self
+ .state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?;
+ if !state.complete || !state.pending.is_empty() {
+ return Err(invalid("Read and update exactly the same number of
rows from this shard before preparing"));
+ }
+ }
+ let result = self.flush_group().await;
+ if result.is_err() {
+ self.failed = true;
+ }
+ result?;
+ let messages = std::mem::take(&mut self.staged);
+ self.messages = Some(messages.clone());
+ Ok(messages)
+ }
+
+ async fn stage_groups(&mut self, batch: RecordBatch, groups: &[i64]) ->
crate::Result<()> {
+ let mut offset = 0;
+ while offset < groups.len() {
+ let group = groups[offset];
+ if self.current_group.is_some_and(|current| current != group) {
+ self.flush_group().await?;
+ }
+ if self.finished_groups.contains(&group) {
+ return Err(invalid(
+ "Shard reader revisited an already updated file group",
+ ));
+ }
+ if self.writer.is_none() {
+ let mut writer = DataEvolutionWriter::new(&self.table,
self.write_columns.clone())?;
+
writer.with_blob_uri_reader_factory(self.blob_uri_reader_factory.clone());
+ self.writer = Some(writer);
+ }
+ self.current_group = Some(group);
+ let count = groups[offset..]
+ .iter()
+ .take_while(|id| **id == group)
+ .count();
+ self.writer
+ .as_mut()
+ .unwrap()
+ .add_matched_batch(batch.slice(offset, count))?;
+ offset += count;
+ }
+ Ok(())
+ }
+
+ async fn flush_group(&mut self) -> crate::Result<()> {
+ if let Some(writer) = self.writer.take() {
+ // Cancellation also drops this consuming writer. Poison first so a
+ // cancelled prepare/update cannot later return a partial success.
+ self.failed = true;
+ self.staged
+ .extend(writer.prepare_commit_with_index(&self.index).await?);
+ self.failed = false;
+ }
+ if let Some(group) = self.current_group.take() {
+ self.finished_groups.insert(group);
+ }
+ Ok(())
+ }
+
+ /// Stop this operation. Prepared CommitMessage files may already be
published.
+ pub fn close(&mut self) {
+ self.writer.take();
+ if let Ok(mut state) = self.state.lock() {
+ state.closed = true;
+ state.pending.clear();
+ }
+ }
+}
+
+/// The caller's read fields must not change the system column used for
mapping.
+fn with_row_id(read_type: &[DataField]) -> crate::Result<Vec<DataField>> {
+ let row_id = row_id_data_field();
+ let mut ids = read_type.iter().filter(|field| field.name() == ROW_ID);
+ if let Some(field) = ids.next() {
+ if ids.next().is_some()
+ || field.id() != row_id.id()
+ || !matches!(field.data_type(), crate::spec::DataType::BigInt(_))
+ {
+ return Err(invalid(
+ "Read type must use the canonical BIGINT _ROW_ID field",
+ ));
+ }
+ return Ok(read_type.to_vec());
+ }
+ let mut fields = read_type.to_vec();
+ fields.push(row_id);
+ Ok(fields)
+}
+
+fn select_shard(splits: &[DataSplit], shard: usize, count: usize) ->
crate::Result<Vec<DataSplit>> {
+ let mut ranges = splits
+ .iter()
+ .flat_map(|split| split.data_files())
+ .filter(|file| is_data_evolution_normal_file(file))
+ .map(|file| {
+ file.row_id_range()
+ .ok_or_else(|| invalid("Data Evolution file is missing its
row-ID range"))
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+ ranges.sort_unstable();
+ let mut groups = Vec::<(i64, i64)>::new();
+ for (start, end) in ranges {
+ if let Some(last) = groups.last_mut().filter(|last| start <= last.1) {
+ last.1 = last.1.max(end);
+ } else {
+ groups.push((start, end));
+ }
+ }
+ let size = groups.len() / count;
+ let remainder = groups.len() % count;
+ let begin = shard * size + shard.min(remainder);
+ let end = begin + size + usize::from(shard < remainder);
+ let Some(&(first, _)) = groups.get(begin).filter(|_| begin < end) else {
+ return Ok(Vec::new());
+ };
+ let last = groups[end - 1].1;
+ // The selected groups are a contiguous slice in global row-ID order.
+ // One interval avoids a scan of every selected group for each split.
+ Ok(splits
+ .iter()
+ .filter(|split| {
+ split
+ .data_files()
+ .iter()
+ .filter(|file| is_data_evolution_normal_file(file))
+ .filter_map(|file| file.row_id_range())
+ .any(|(start, end)| first <= end && last >= start)
+ })
+ .map(|split| {
+ split
+ .clone()
+ .with_selected_row_ranges(vec![super::RowRange::new(first,
last)])
+ })
+ .collect())
+}
+
+fn record_rows(
+ state: &Mutex<ReadState>,
+ batch: RecordBatch,
+ output_columns: usize,
+) -> crate::Result<RecordBatch> {
+ let ids = batch
+ .column_by_name(ROW_ID)
+ .and_then(|array| array.as_any().downcast_ref::<Int64Array>())
+ .ok_or_else(|| invalid("Shard reader must produce Int64 row IDs"))?;
+ if ids.null_count() != 0 {
+ return Err(invalid("Shard reader produced NULL row IDs"));
+ }
+ let output = batch
+ .project(&(0..output_columns).collect::<Vec<_>>())
+ .map_err(|error| invalid(error.to_string()))?;
+ let mut state = state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?;
+ if state.closed {
+ return Err(invalid("Shard updater was closed"));
+ }
+ state.pending.extend(ids.values().iter().copied());
+ Ok(output)
+}
+
+fn finish_read(state: &Mutex<ReadState>, failed: bool) -> crate::Result<()> {
+ let mut state = state
+ .lock()
+ .map_err(|_| invalid("Shard reader state poisoned"))?;
+ state.failed |= failed;
+ state.complete = !state.failed;
+ Ok(())
+}
+
+#[cfg(test)]
+#[path = "shard_table_updater/tests.rs"]
+mod tests;
diff --git a/crates/paimon/src/table/shard_table_updater/tests.rs
b/crates/paimon/src/table/shard_table_updater/tests.rs
new file mode 100644
index 00000000..4f2c2ffe
--- /dev/null
+++ b/crates/paimon/src/table/shard_table_updater/tests.rs
@@ -0,0 +1,714 @@
+// 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.
+
+use super::*;
+use crate::catalog::Identifier;
+use crate::io::FileIOBuilder;
+use crate::spec::{
+ DataType as PaimonType, IntType, RowType, Schema as TableDefinition,
TableSchema, VarCharType,
+};
+use crate::table::TableCommit;
+use arrow_array::{Int32Array, StringArray, StructArray};
+
+fn table(partitioned: bool) -> Table {
+ table_with_io(partitioned, FileIOBuilder::new("memory").build().unwrap())
+}
+
+fn table_with_io(partitioned: bool, file_io: crate::io::FileIO) -> Table {
+ let mut schema = TableDefinition::builder()
+ .column("p", PaimonType::VarChar(VarCharType::string_type()))
+ .column("id", PaimonType::Int(IntType::new()))
+ .column("value", PaimonType::Int(IntType::new()))
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true")
+ .option("deletion-vectors.enabled", "true");
+ if partitioned {
+ schema = schema.partition_keys(["p"]);
+ }
+ Table::new(
+ file_io,
+ Identifier::new("db", "shard"),
+ "memory:/shard".into(),
+ TableSchema::new(0, &schema.build().unwrap()),
+ None,
+ )
+}
+
+fn batch(partitions: Vec<&str>, ids: Vec<i32>, values: Vec<i32>) ->
RecordBatch {
+ RecordBatch::try_from_iter([
+ ("p", Arc::new(StringArray::from(partitions)) as ArrayRef),
+ ("id", Arc::new(Int32Array::from(ids)) as ArrayRef),
+ ("value", Arc::new(Int32Array::from(values)) as ArrayRef),
+ ])
+ .unwrap()
+}
+
+fn values(values: Vec<i32>) -> RecordBatch {
+ RecordBatch::try_from_iter([("value", Arc::new(Int32Array::from(values))
as ArrayRef)]).unwrap()
+}
+
+async fn commit(table: &Table, messages: Vec<CommitMessage>) {
+ TableCommit::new(table.clone(), "shard".into())
+ .commit(messages)
+ .await
+ .unwrap();
+}
+
+async fn seed(table: &Table, data: RecordBatch) {
+ let mut writer = table
+ .new_write_builder()
+ .with_commit_user("shard")
+ .unwrap()
+ .new_write()
+ .unwrap();
+ writer.write_arrow_batch(&data).await.unwrap();
+ commit(table, writer.prepare_commit().await.unwrap()).await;
+ writer.close().await;
+}
+
+async fn updater(table: &Table, index: usize, count: usize) ->
ShardTableUpdater {
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ update
+ .new_shard_updater(index, count,
Some(vec![table.schema().fields()[1].clone()]))
+ .await
+ .unwrap()
+}
+
+async fn rows(table: &Table) -> Vec<(i32, i32)> {
+ let builder = table.new_read_builder();
+ let plan = builder.new_scan().plan().await.unwrap();
+ let mut read =
builder.new_read().unwrap().to_arrow(plan.splits()).unwrap();
+ let mut result = Vec::new();
+ while let Some(batch) = read.try_next().await.unwrap() {
+ 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();
+ result.extend(
+ ids.values()
+ .iter()
+ .zip(values.values())
+ .map(|(a, b)| (*a, *b)),
+ );
+ }
+ result.sort();
+ result
+}
+
+#[tokio::test]
+async fn balanced_shards_keep_overlapping_column_files_together() {
+ let table = table(false);
+ for id in 0..5 {
+ seed(&table, batch(vec!["p"], vec![id], vec![id * 10])).await;
+ }
+ // Construct every shard against one snapshot, before any replacement
files exist.
+ let mut shards = Vec::new();
+ for shard in 0..3 {
+ shards.push(updater(&table, shard, 3).await);
+ }
+ let mut counts = Vec::new();
+ for mut shard in shards {
+ let mut read = shard.arrow_reader().unwrap();
+ let mut count = 0;
+ while let Some(batch) = read.try_next().await.unwrap() {
+ assert_eq!(batch.schema().fields().len(), 1);
+ let ids = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ count += ids.len();
+ shard
+ .update_by_arrow_batch(values(ids.values().iter().map(|id| id
+ 100).collect()))
+ .await
+ .unwrap();
+ }
+ counts.push(count);
+ let messages = shard.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.check_from_snapshot == Some(5)));
+ for message in &messages {
+ assert_eq!(message.bucket, 0);
+ assert!(message
+ .new_files
+ .iter()
+ .all(|file| file.write_cols == Some(vec!["value".into()])));
+ }
+ commit(&table, messages).await;
+ }
+ assert_eq!(counts, vec![2, 2, 1]);
+ assert_eq!(
+ rows(&table).await,
+ (0..5).map(|id| (id, id + 100)).collect::<Vec<_>>()
+ );
+ // A subsequent scan has two column files per range, which still form five
shards.
+ let mut sizes = Vec::new();
+ for shard in 0..7 {
+ let mut update = updater(&table, shard, 7).await;
+ sizes.push(
+ update
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap()
+ .iter()
+ .map(RecordBatch::num_rows)
+ .sum::<usize>(),
+ );
+ }
+ assert_eq!(sizes, vec![1, 1, 1, 1, 1, 0, 0]);
+}
+
+#[tokio::test]
+async fn partitioned_shard_accepts_rebatched_values_and_pins_snapshot() {
+ let table = table(true);
+ seed(
+ &table,
+ batch(vec!["a", "b", "a"], vec![1, 2, 3], vec![10, 20, 30]),
+ )
+ .await;
+ let mut update = updater(&table, 0, 1).await;
+ seed(&table, batch(vec!["a"], vec![4], vec![40])).await;
+ let data = update
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let ids = data
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(ids.len(), 3);
+ update
+ .update_by_arrow_batch(values(ids.iter().map(|id| id * 100).collect()))
+ .await
+ .unwrap();
+ let messages = update.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 2);
+ assert!(messages
+ .iter()
+ .all(|message| message.check_from_snapshot == Some(1)));
+ assert!(messages.iter().all(|message| message
+ .new_files
+ .iter()
+ .all(|file| file.first_row_id.is_some())));
+ commit(&table, messages).await;
+ assert_eq!(
+ rows(&table).await,
+ vec![(1, 100), (2, 200), (3, 300), (4, 40)]
+ );
+}
+
+#[tokio::test]
+async fn deleted_rows_do_not_shift_positional_assignments() {
+ let table = table(false);
+ seed(
+ &table,
+ batch(vec!["p"; 4], vec![1, 2, 3, 4], vec![10, 20, 30, 40]),
+ )
+ .await;
+ let messages = table
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .delete_by_row_id(vec![1, 3])
+ .await
+ .unwrap();
+ commit(&table, messages).await;
+ let mut update = updater(&table, 0, 1).await;
+ let data = update
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(data.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+ update
+ .update_by_arrow_batch(values(vec![111, 333]))
+ .await
+ .unwrap();
+ let messages = update.prepare_commit().await.unwrap();
+ assert!(messages
+ .iter()
+ .flat_map(|m| &m.new_files)
+ .all(|file| file.row_count == 4 && file.first_row_id == Some(0)));
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 111), (3, 333)]);
+}
+
+#[tokio::test]
+async fn count_and_schema_failures_preserve_pending_rows_for_correction() {
+ let table = table(false);
+ seed(&table, batch(vec!["p"; 3], vec![1, 2, 3], vec![10, 20, 30])).await;
+ let mut update = updater(&table, 0, 1).await;
+ assert!(update.update_by_arrow_batch(values(vec![9])).await.is_err());
+ update
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let bad =
+ RecordBatch::try_from_iter([("wrong",
Arc::new(Int32Array::from(vec![1])) as ArrayRef)])
+ .unwrap();
+ assert!(update.update_by_arrow_batch(bad).await.is_err());
+ assert!(update
+ .update_by_arrow_batch(values(vec![1, 2, 3, 4]))
+ .await
+ .is_err());
+ update
+ .update_by_arrow_batch(values(vec![100]))
+ .await
+ .unwrap();
+ assert!(update
+ .prepare_commit()
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("same number"));
+ update
+ .update_by_arrow_batch(values(vec![200, 300]))
+ .await
+ .unwrap();
+ let messages = update.prepare_commit().await.unwrap();
+ let repeated = update.prepare_commit().await.unwrap();
+ assert_eq!(
+ messages[0].new_files[0].file_name,
+ repeated[0].new_files[0].file_name
+ );
+ assert!(update.update_by_arrow_batch(values(vec![])).await.is_err());
+ update.close();
+ // Prepared files survive explicit close and can still be committed.
+ commit(&table, messages).await;
+ assert_eq!(rows(&table).await, vec![(1, 100), (2, 200), (3, 300)]);
+}
+
+#[tokio::test]
+async fn incomplete_reader_cannot_publish_partial_shard() {
+ let table = table(false);
+ seed(&table, batch(vec!["p"], vec![1], vec![10])).await;
+ seed(&table, batch(vec!["p"], vec![2], vec![20])).await;
+ let mut update = updater(&table, 0, 1).await;
+ let mut read = update.arrow_reader().unwrap();
+ let batch = read.try_next().await.unwrap().unwrap();
+ update
+ .update_by_arrow_batch(values(vec![99; batch.num_rows()]))
+ .await
+ .unwrap();
+ drop(read);
+ assert!(update.arrow_reader().is_err());
+ assert!(update.prepare_commit().await.is_err());
+ update.close();
+ assert!(update.prepare_commit().await.is_err());
+ assert_eq!(rows(&table).await, vec![(1, 10), (2, 20)]);
+}
+
+#[tokio::test]
+async fn empty_shards_and_invalid_configuration() {
+ let table = table(false);
+ let mut empty = updater(&table, 2, 3).await;
+ assert!(empty.prepare_commit().await.unwrap().is_empty());
+ let update = table.new_write_builder().new_update().unwrap();
+ assert!(update.new_shard_updater(0, 0, None).await.is_err());
+ assert!(update.new_shard_updater(2, 2, None).await.is_err());
+ let mut bad_columns = update;
+ bad_columns.with_update_type(vec![]).unwrap();
+ assert!(bad_columns.new_shard_updater(0, 1, None).await.is_err());
+}
+
+#[tokio::test]
+async fn returned_reader_survives_updater_scope_and_close_stops_it() {
+ let table = table(false);
+ seed(&table, batch(vec!["p"], vec![1], vec![10])).await;
+ let mut read = updater(&table, 0, 1).await.arrow_reader().unwrap();
+ assert_eq!(read.try_next().await.unwrap().unwrap().num_rows(), 1);
+ assert!(read.try_next().await.unwrap().is_none());
+ let mut update = updater(&table, 0, 1).await;
+ let mut read = update.arrow_reader().unwrap();
+ update.close();
+ assert!(read.try_next().await.is_err());
+ assert!(update.prepare_commit().await.is_err());
+}
+
+#[test]
+fn reader_eof_cannot_clear_failure() {
+ let state = Mutex::new(ReadState::default());
+ finish_read(&state, true).unwrap();
+ finish_read(&state, false).unwrap();
+ let state = state.lock().unwrap();
+ assert!(state.failed);
+ assert!(!state.complete);
+}
+
+#[derive(Debug)]
+struct ControlledIo {
+ operator: opendal::Operator,
+ mode: std::sync::atomic::AtomicU8,
+ entered: tokio::sync::Notify,
+}
+
+impl ControlledIo {
+ fn table() -> (Table, Arc<Self>) {
+ let provider = Arc::new(Self {
+ operator:
opendal::Operator::from_config(opendal::services::MemoryConfig::default())
+ .unwrap(),
+ mode: std::sync::atomic::AtomicU8::new(0),
+ entered: tokio::sync::Notify::new(),
+ });
+ let io = FileIOBuilder::new("unused")
+ .with_provider(provider.clone())
+ .build()
+ .unwrap();
+ (table_with_io(false, io), provider)
+ }
+}
+
+#[async_trait::async_trait]
+impl crate::io::FileIOProvider for ControlledIo {
+ async fn create(&self, path: &str) -> crate::Result<(opendal::Operator,
String)> {
+ use std::sync::atomic::Ordering;
+ if path.ends_with(".parquet") {
+ match self.mode.load(Ordering::SeqCst) {
+ 1 => {
+ self.entered.notify_one();
+ std::future::pending::<()>().await;
+ }
+ 2 => {
+ self.mode.store(0, Ordering::SeqCst);
+ return Err(invalid("Injected shard I/O failure"));
+ }
+ _ => {}
+ }
+ }
+ Ok((
+ self.operator.clone(),
+ path.trim_start_matches("memory:")
+ .trim_start_matches('/')
+ .into(),
+ ))
+ }
+}
+
+#[tokio::test]
+async fn cancelled_prepare_cannot_retry_as_an_empty_success() {
+ use std::sync::atomic::Ordering;
+ let (table, provider) = ControlledIo::table();
+ seed(&table, batch(vec!["p"], vec![1], vec![10])).await;
+ let mut update = updater(&table, 0, 1).await;
+ update
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ update
+ .update_by_arrow_batch(values(vec![99]))
+ .await
+ .unwrap();
+ provider.mode.store(1, Ordering::SeqCst);
+ let mut prepare = Box::pin(update.prepare_commit());
+ tokio::select! {
+ result = &mut prepare => panic!("Expected stalled I/O, got
{result:?}"),
+ result = tokio::time::timeout(std::time::Duration::from_secs(5),
provider.entered.notified()) => {
+ result.expect("prepare never reached storage I/O");
+ }
+ }
+ drop(prepare);
+ provider.mode.store(0, Ordering::SeqCst);
+ assert!(update
+ .prepare_commit()
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("failed"));
+ assert!(update.update_by_arrow_batch(values(vec![])).await.is_err());
+ assert_eq!(rows(&table).await, vec![(1, 10)]);
+ assert_eq!(
+ table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .unwrap()
+ .id(),
+ 1
+ );
+}
+
+#[tokio::test]
+async fn reader_eof_does_not_resurrect_a_failed_group_write() {
+ use std::sync::atomic::Ordering;
+ let (table, provider) = ControlledIo::table();
+ seed(&table, batch(vec!["p"], vec![1], vec![10])).await;
+ seed(&table, batch(vec!["p"], vec![2], vec![20])).await;
+ let mut update = updater(&table, 0, 1).await;
+ let mut read = update.arrow_reader().unwrap();
+ assert_eq!(read.try_next().await.unwrap().unwrap().num_rows(), 1);
+ update
+ .update_by_arrow_batch(values(vec![100]))
+ .await
+ .unwrap();
+ assert_eq!(read.try_next().await.unwrap().unwrap().num_rows(), 1);
+ provider.mode.store(2, Ordering::SeqCst);
+ let error = update
+ .update_by_arrow_batch(values(vec![200]))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("Injected shard I/O failure"));
+ assert!(read.try_next().await.unwrap().is_none());
+ assert!(update
+ .update_by_arrow_batch(values(vec![200]))
+ .await
+ .is_err());
+ assert!(update.prepare_commit().await.is_err());
+ assert_eq!(rows(&table).await, vec![(1, 10), (2, 20)]);
+}
+
+#[tokio::test]
+async fn read_types_belong_to_each_operation_and_are_independent_of_updates() {
+ let table = table(false);
+ seed(&table, batch(vec!["p", "p"], vec![1, 2], vec![10, 20])).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let id_field = table.schema().fields()[1].clone();
+ let mut ids = update
+ .new_shard_updater(0, 1, Some(vec![id_field]))
+ .await
+ .unwrap();
+ let mut originals = update
+ .new_shard_updater(
+ 0,
+ 1,
+ Some(vec![
+ table.schema().fields()[2].clone(),
+ row_id_data_field(),
+ ]),
+ )
+ .await
+ .unwrap();
+ let read = originals
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(read[0].schema().field(0).name(), "value");
+ assert_eq!(
+ read[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values(),
+ &[10, 20]
+ );
+ assert_eq!(
+ read[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values(),
+ &[0, 1]
+ );
+ originals.close();
+ let mut read = ids.arrow_reader().unwrap();
+ while let Some(batch) = read.try_next().await.unwrap() {
+ assert_eq!(batch.schema().fields().len(), 1);
+ assert_eq!(batch.schema().field(0).name(), "id");
+ let source = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ ids.update_by_arrow_batch(values(source.values().iter().map(|id| id *
100).collect()))
+ .await
+ .unwrap();
+ }
+ commit(&table, ids.prepare_commit().await.unwrap()).await;
+ assert_eq!(rows(&table).await, vec![(1, 100), (2, 200)]);
+}
+
+#[tokio::test]
+async fn default_and_empty_read_types_preserve_row_counts() {
+ let table = table(false);
+ seed(&table, batch(vec!["p", "p"], vec![1, 2], vec![10, 20])).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let mut full = update.new_shard_updater(0, 1, None).await.unwrap();
+ let batches = full
+ .arrow_reader()
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(
+ batches[0]
+ .schema()
+ .fields()
+ .iter()
+ .map(|f| f.name().as_str())
+ .collect::<Vec<_>>(),
+ vec!["p", "id", "value"]
+ );
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+ full.close();
+ let mut empty = update.new_shard_updater(0, 1,
Some(vec![])).await.unwrap();
+ let mut read = empty.arrow_reader().unwrap();
+ while let Some(batch) = read.try_next().await.unwrap() {
+ assert_eq!(batch.num_columns(), 0);
+ assert_eq!(batch.num_rows(), 2);
+ empty
+ .update_by_arrow_batch(values(vec![99; batch.num_rows()]))
+ .await
+ .unwrap();
+ }
+ commit(&table, empty.prepare_commit().await.unwrap()).await;
+ assert_eq!(rows(&table).await, vec![(1, 99), (2, 99)]);
+}
+
+#[tokio::test]
+async fn nested_read_type_prunes_children_and_retains_their_ids() {
+ let definition = TableDefinition::builder()
+ .column("p", PaimonType::VarChar(VarCharType::string_type()))
+ .column("id", PaimonType::Int(IntType::new()))
+ .column("value", PaimonType::Int(IntType::new()))
+ .column(
+ "payload",
+ PaimonType::Row(RowType::new(vec![
+ DataField::new(0, "a".into(), PaimonType::Int(IntType::new())),
+ DataField::new(1, "b".into(), PaimonType::Int(IntType::new())),
+ ])),
+ )
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("db", "nested_shard"),
+ "memory:/nested_shard".into(),
+ TableSchema::new(0, &definition),
+ None,
+ );
+ let data = batch(vec!["p", "p"], vec![1, 2], vec![10, 20]);
+ let payload = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Int32, true)),
+ Arc::new(Int32Array::from(vec![100, 200])) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Int32, true)),
+ Arc::new(Int32Array::from(vec![3, 4])) as ArrayRef,
+ ),
+ ]);
+ let mut fields = data.schema().fields().to_vec();
+ fields.push(Arc::new(Field::new(
+ "payload",
+ payload.data_type().clone(),
+ true,
+ )));
+ let mut columns = data.columns().to_vec();
+ columns.push(Arc::new(payload));
+ seed(
+ &table,
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(),
+ )
+ .await;
+ let source = &table.schema().fields()[3];
+ let PaimonType::Row(row) = source.data_type() else {
+ panic!("expected ROW")
+ };
+ let read_type = vec![DataField::new(
+ source.id(),
+ source.name().into(),
+ PaimonType::Row(RowType::new(vec![row.fields()[1].clone()])),
+ )];
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let mut shard = update
+ .new_shard_updater(0, 1, Some(read_type))
+ .await
+ .unwrap();
+ let mut read = shard.arrow_reader().unwrap();
+ while let Some(batch) = read.try_next().await.unwrap() {
+ let payload = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(payload.num_columns(), 1);
+ assert_eq!(payload.fields()[0].name(), "b");
+ let source = payload
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(source.values(), &[3, 4]);
+ shard
+ .update_by_arrow_batch(values(source.values().iter().map(|id| id *
100).collect()))
+ .await
+ .unwrap();
+ }
+ commit(&table, shard.prepare_commit().await.unwrap()).await;
+ assert_eq!(rows(&table).await, vec![(1, 300), (2, 400)]);
+}
+
+#[tokio::test]
+async fn read_type_cannot_override_the_internal_row_id() {
+ let table = table(false);
+ seed(&table, batch(vec!["p"], vec![1], vec![10])).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ for fields in [
+ vec![DataField::new(
+ crate::spec::ROW_ID_FIELD_ID,
+ ROW_ID.into(),
+ PaimonType::Int(IntType::new()),
+ )],
+ vec![row_id_data_field().with_id(0)],
+ vec![row_id_data_field(), row_id_data_field()],
+ ] {
+ let error = update
+ .new_shard_updater(0, 1, Some(fields))
+ .await
+ .err()
+ .unwrap();
+ assert!(error.to_string().contains("canonical BIGINT _ROW_ID"));
+ }
+ assert_eq!(rows(&table).await, vec![(1, 10)]);
+}
diff --git a/crates/paimon/src/table/table_update.rs
b/crates/paimon/src/table/table_update.rs
index a516fb70..af9b283a 100644
--- a/crates/paimon/src/table/table_update.rs
+++ b/crates/paimon/src/table/table_update.rs
@@ -25,6 +25,7 @@ use std::collections::HashSet;
use arrow_array::RecordBatch;
+use crate::spec::DataField;
use crate::table::{CommitMessage, Table, TableUpdateByRowId};
const ROW_ID: &str = "_ROW_ID";
@@ -101,6 +102,33 @@ impl TableUpdate {
.collect()
}
+ /// Read and replace columns over one balanced shard of whole file groups.
+ ///
+ /// `read_type` configures this operation's reader, independently of the
+ /// update columns. It preserves field IDs and nested read metadata, just
+ /// like
[`ReadBuilder::with_read_type`](super::ReadBuilder::with_read_type).
+ /// `None` reads all table fields; `Some(vec![])` returns zero-column
batches
+ /// with their row counts preserved. Row IDs are tracked internally.
+ pub async fn new_shard_updater(
+ &self,
+ shard: usize,
+ shard_count: usize,
+ read_type: Option<Vec<DataField>>,
+ ) -> crate::Result<super::ShardTableUpdater> {
+ let mut updater = super::ShardTableUpdater::new(
+ &self.table,
+ read_type,
+ self.update_cols
+ .clone()
+ .unwrap_or_else(|| self.all_fields()),
+ shard,
+ shard_count,
+ )
+ .await?;
+
updater.with_blob_uri_reader_factory(self.blob_uri_reader_factory.clone());
+ Ok(updater)
+ }
+
fn columns_for_batch(&self, batch: &RecordBatch) -> Vec<String> {
self.update_cols.clone().unwrap_or_else(|| {
batch
