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 ca54d4a4 Preserve supplied fields across key upsert batches (#1061)
ca54d4a4 is described below
commit ca54d4a4d4f2baf0a5cb469c2d4dc098ef6e86ac
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 19:32:50 2026 +0800
Preserve supplied fields across key upsert batches (#1061)
---
.../python/python/pypaimon_rust/datafusion.pyi | 8 +-
bindings/python/src/write.rs | 14 +-
bindings/python/tests/test_update.py | 43 ++++
crates/paimon/src/table/table_update.rs | 3 +
crates/paimon/src/table/table_upsert.rs | 177 ++++++++++++-----
crates/paimon/tests/table_update_test.rs | 216 +++++++++++++++++++++
6 files changed, 405 insertions(+), 56 deletions(-)
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 1eb2f4bf..d76b942f 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -256,9 +256,10 @@ class BatchTableUpdate:
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(
- self, table: pyarrow.Table, upsert_keys: List[str]
+ self, table: Union[pyarrow.Table, Sequence[pyarrow.RecordBatch]],
upsert_keys: List[str]
) -> List[CommitMessage]:
- """Prepare one key upsert; commit the returned messages separately."""
+ """Prepare one key upsert. Batches may have different non-key fields;
+ appended rows in each partition must share a field set."""
...
class TableUpdateByRowId:
@@ -283,7 +284,8 @@ class StreamTableUpdate:
self, row_ids: Sequence[int], commit_identifier: int
) -> List[CommitMessage]: ...
def upsert_by_arrow_with_key(
- self, table: pyarrow.Table, upsert_keys: List[str], commit_identifier:
int
+ self, table: Union[pyarrow.Table, Sequence[pyarrow.RecordBatch]],
+ upsert_keys: List[str], commit_identifier: int
) -> List[CommitMessage]:
"""Prepare one stream upsert; commit with the same identifier."""
...
diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs
index eea7870a..20abb7df 100644
--- a/bindings/python/src/write.rs
+++ b/bindings/python/src/write.rs
@@ -492,6 +492,18 @@ fn arrow_table_batches(table: &Bound<'_, PyAny>) ->
PyResult<Vec<RecordBatch>> {
Ok(batches)
}
+/// Preserve each batch's supplied fields; core owns matching, deduplication
+/// and the distinction between missing columns and explicit NULL values.
+fn arrow_upsert_batches(input: &Bound<'_, PyAny>) ->
PyResult<Vec<RecordBatch>> {
+ if input.hasattr("to_batches")? {
+ return arrow_table_batches(input);
+ }
+ input
+ .try_iter()?
+ .map(|batch| RecordBatch::from_pyarrow_bound(&batch?))
+ .collect()
+}
+
struct UpdateContext {
inner: TableUpdate,
table: Arc<Table>,
@@ -621,7 +633,7 @@ impl UpdateContext {
input: &Bound<'_, PyAny>,
keys: Vec<String>,
) -> PyResult<Vec<PyCommitMessage>> {
- let batches = arrow_table_batches(input)?;
+ let batches = arrow_upsert_batches(input)?;
let messages = py
.detach(||
runtime().block_on(self.inner.upsert_by_arrow_with_key(batches, keys)))
.map_err(to_py_err)?;
diff --git a/bindings/python/tests/test_update.py
b/bindings/python/tests/test_update.py
index 33bf8a09..3702f645 100644
--- a/bindings/python/tests/test_update.py
+++ b/bindings/python/tests/test_update.py
@@ -22,6 +22,49 @@ import pypaimon_rust.datafusion as datafusion
from pypaimon_rust.datafusion import PaimonCatalog, SQLContext
+def
test_upsert_batches_preserve_missing_fields_nulls_and_stream_updates(tmp_path):
+ context = SQLContext()
+ context.register_catalog('paimon', {'warehouse': str(tmp_path)})
+ context.sql('CREATE SCHEMA paimon.shapes')
+ context.sql("""CREATE TABLE paimon.shapes.t (id INT, name STRING, score
INT) WITH (
+ 'row-tracking.enabled' = 'true', 'data-evolution.enabled' = 'true')""")
+ context.sql("INSERT INTO paimon.shapes.t (id, name, score) VALUES (1, 'a',
10), (1, 'b', 11), (2, 'c', 20)")
+ table = PaimonCatalog({'warehouse': str(tmp_path)}).get_table('shapes.t')
+
+ def batch(values):
+ schema = pa.schema([(name, pa.string() if name == 'name' else
pa.int32()) for name in values])
+ return pa.RecordBatch.from_pydict(values, schema=schema)
+
+ builder = table.new_batch_write_builder()
+ update = builder.new_update().with_update_type(['name'])
+ messages = update.upsert_by_arrow_with_key([
+ batch({'id': [1], 'score': [999]}),
+ batch({'id': [3], 'score': [30]}),
+ batch({'name': [None], 'id': [1]}),
+ batch({'id': [2], 'name': ['new'], 'score': [999]}),
+ ], ['id'])
+ builder.new_commit().commit(messages)
+ assert pa.Table.from_batches(context.sql(
+ 'SELECT id, name, score FROM paimon.shapes.t ORDER BY id,
score')).to_pylist() == [
+ {'id': 1, 'name': None, 'score': 10},
+ {'id': 1, 'name': None, 'score': 11},
+ {'id': 2, 'name': 'new', 'score': 20},
+ {'id': 3, 'name': None, 'score': 30},
+ ]
+ stream = table.new_stream_write_builder().with_commit_user('shape-stream')
+ messages =
stream.new_update().with_update_type(['score']).upsert_by_arrow_with_key(
+ [batch({'id': [1], 'score': [100]}), batch({'id': [3], 'score': [300],
'name': ['ignored']})],
+ ['id'], 101)
+ stream.new_commit().commit(101, messages)
+ assert pa.Table.from_batches(context.sql(
+ 'SELECT id, name, score FROM paimon.shapes.t ORDER BY id,
score')).to_pylist() == [
+ {'id': 1, 'name': None, 'score': 100},
+ {'id': 1, 'name': None, 'score': 100},
+ {'id': 2, 'name': 'new', 'score': 20},
+ {'id': 3, 'name': None, 'score': 300},
+ ]
+
+
def test_table_upsert_updates_duplicate_targets_and_appends(tmp_path):
assert not hasattr(datafusion, '_match_upsert_keys')
assert not hasattr(datafusion, 'UpsertKeyMatcher')
diff --git a/crates/paimon/src/table/table_update.rs
b/crates/paimon/src/table/table_update.rs
index 49765fed..480b776f 100644
--- a/crates/paimon/src/table/table_update.rs
+++ b/crates/paimon/src/table/table_update.rs
@@ -195,6 +195,9 @@ impl TableUpdate {
/// columns are omitted from `upsert_keys`. Existing keys update every
/// matching row ID; new keys append all input columns. Omitted columns on
/// matched rows retain their values; on new rows they are absent from the
file.
+ /// Input batches may have different non-key columns. After last-write-wins
+ /// deduplication, matched rows must supply their update columns and
appended
+ /// rows in each partition must share a field set.
pub async fn upsert_by_arrow_with_key(
&self,
batches: Vec<RecordBatch>,
diff --git a/crates/paimon/src/table/table_upsert.rs
b/crates/paimon/src/table/table_upsert.rs
index 5ff0534f..0ce5aa75 100644
--- a/crates/paimon/src/table/table_upsert.rs
+++ b/crates/paimon/src/table/table_upsert.rs
@@ -19,12 +19,14 @@
use std::borrow::Cow;
use std::collections::HashSet;
+use std::ops::Range;
use std::sync::Arc;
use arrow_array::{ArrayRef, Int64Array, RecordBatch, UInt32Array};
use arrow_schema::{DataType, Field, Schema};
use arrow_select::{concat::concat_batches, take::take};
use futures::TryStreamExt;
+use indexmap::IndexMap;
use super::update_input::unique_column_index;
use super::upsert_key_matcher::UpsertKeyMatcher;
@@ -140,7 +142,8 @@ impl TableUpsert {
/// Add rows containing the keys and columns to write. Column order may
/// differ from the table schema; names and Arrow layouts follow the same
- /// normalization as ordinary writes. Batches form one logical upsert
input.
+ /// normalization as ordinary writes. Batches form one logical upsert
input,
+ /// but may carry different non-key columns, preserving absent fields.
pub(super) fn add_batch(&mut self, batch: RecordBatch) ->
crate::Result<()> {
let target =
crate::arrow::build_target_arrow_schema(self.table.schema().fields())?;
let input_schema = batch.schema();
@@ -171,41 +174,71 @@ impl TableUpsert {
}
let ordered = RecordBatch::try_new(Arc::new(Schema::new(fields)),
columns)
.map_err(|error| invalid(format!("cannot order upsert columns:
{error}")))?;
- if self
- .source
- .first()
- .is_some_and(|first| first.schema() != ordered.schema())
- {
- return Err(invalid(
- "Arrow batches in one upsert input must have the same columns",
- ));
- }
self.source.push(ordered);
Ok(())
}
- /// Match only the source partitions, as PyPaimon does. Filter planned
splits
- /// after global row IDs have been assigned so pruning cannot renumber
rows.
- fn matching_splits<'a>(
+ /// Concatenate only match keys. Padding payload columns with NULL would
+ /// lose the distinction between an absent field and an explicit NULL.
+ fn key_rows(&self) -> crate::Result<RecordBatch> {
+ let batches = self
+ .source
+ .iter()
+ .map(|batch| {
+ let columns = batch
+ .schema()
+ .fields()
+ .iter()
+ .enumerate()
+ .filter(|(_, field)| self.keys.contains(field.name()))
+ .map(|(index, _)| index)
+ .collect::<Vec<_>>();
+ batch
+ .project(&columns)
+ .map_err(|error| invalid(error.to_string()))
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+ let first = batches
+ .first()
+ .ok_or_else(|| invalid("Input data is empty"))?;
+ concat_batches(&first.schema(), &batches)
+ .map_err(|error| invalid(format!("cannot concatenate upsert keys:
{error}")))
+ }
+
+ /// Matcher indices are sorted by their original input position, with one
+ /// repeated index per matching target row. The range retains alignment
+ /// with the matcher's parallel row-ID vector across source batch
boundaries.
+ fn selected_batches(
&self,
- source: &RecordBatch,
- plan: &'a Plan,
- ) -> crate::Result<Cow<'a, [DataSplit]>> {
- let schema = self.table.schema();
- if schema.partition_keys().is_empty() {
- return Ok(Cow::Borrowed(plan.splits()));
+ indices: &[usize],
+ ) -> crate::Result<Vec<(RecordBatch, Range<usize>)>> {
+ let mut groups = Vec::new();
+ let (mut offset, mut selected) = (0, 0);
+ for source in &self.source {
+ let end = offset + source.num_rows();
+ let count = indices[selected..].partition_point(|index| *index <
end);
+ if count > 0 {
+ let range = selected..selected + count;
+ let rows = indices[range.clone()]
+ .iter()
+ .map(|index| index - offset)
+ .collect::<Vec<_>>();
+ groups.push((selected_rows(source, &rows)?, range));
+ selected += count;
+ }
+ offset = end;
}
+ Ok(groups)
+ }
+
+ fn serialized_partitions(&self, batch: &RecordBatch) ->
crate::Result<Vec<Vec<u8>>> {
+ let schema = self.table.schema();
let indices = schema
.partition_keys()
.iter()
- .map(|name| {
- source
- .schema()
- .index_of(name)
- .map_err(|error| invalid(error.to_string()))
- })
+ .map(|name| unique_column_index(&batch.schema(), name))
.collect::<crate::Result<Vec<_>>>()?;
- let input_fields = source
+ let fields = batch
.schema()
.fields()
.iter()
@@ -218,9 +251,45 @@ impl TableUpsert {
.clone()
})
.collect::<Vec<_>>();
- let partitions: HashSet<_> = batch_to_serialized_bytes(source,
&indices, &input_fields)?
- .into_iter()
- .collect();
+ batch_to_serialized_bytes(batch, &indices, &fields)
+ }
+
+ /// Java append writers use one write type for a partition's input. Like
+ /// PyPaimon named-row upserts, validate only surviving unmatched rows;
+ /// different partitions may append different field sets.
+ fn append_groups(&self, indices: &[usize]) ->
crate::Result<Vec<Vec<RecordBatch>>> {
+ let mut groups: IndexMap<Vec<u8>, Vec<RecordBatch>> = IndexMap::new();
+ for (batch, _) in self.selected_batches(indices)? {
+ let mut by_partition: IndexMap<Vec<u8>, Vec<usize>> =
IndexMap::new();
+ for (row, partition) in
self.serialized_partitions(&batch)?.into_iter().enumerate() {
+ by_partition.entry(partition).or_default().push(row);
+ }
+ for (partition, rows) in by_partition {
+ let group = groups.entry(partition).or_default();
+ if group
+ .first()
+ .is_some_and(|first| first.schema() != batch.schema())
+ {
+ return Err(invalid("upsert_by_key requires appended rows
in the same partition to have the same field set"));
+ }
+ group.push(selected_rows(&batch, &rows)?);
+ }
+ }
+ Ok(groups.into_values().collect())
+ }
+
+ /// Match only the source partitions, as PyPaimon does. Filter planned
splits
+ /// after global row IDs have been assigned so pruning cannot renumber
rows.
+ fn matching_splits<'a>(
+ &self,
+ source: &RecordBatch,
+ plan: &'a Plan,
+ ) -> crate::Result<Cow<'a, [DataSplit]>> {
+ let schema = self.table.schema();
+ if schema.partition_keys().is_empty() {
+ return Ok(Cow::Borrowed(plan.splits()));
+ }
+ let partitions: HashSet<_> =
self.serialized_partitions(source)?.into_iter().collect();
let fields = schema.partition_fields();
let mut splits = Vec::new();
for split in plan.splits() {
@@ -236,11 +305,7 @@ impl TableUpsert {
#[must_use = "commit messages must be passed to TableCommit"]
pub(super) async fn prepare_commit(self) ->
crate::Result<Vec<CommitMessage>> {
CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?;
- let Some(first) = self.source.first() else {
- return Err(invalid("Input data is empty"));
- };
- let source = concat_batches(&first.schema(), &self.source)
- .map_err(|error| invalid(format!("cannot concatenate upsert input:
{error}")))?;
+ let source = self.key_rows()?;
if source.num_rows() == 0 {
return Err(invalid("Input data is empty"));
}
@@ -259,44 +324,52 @@ impl TableUpsert {
}
let (matched_indices, row_ids, new_indices) = matcher.finish();
+ // Validate append field sets before staging any updates or new files.
+ let append_groups = self.append_groups(&new_indices)?;
+ let matched_batches = self.selected_batches(&matched_indices)?;
+
let mut messages = Vec::new();
let result: crate::Result<()> = async {
if !matched_indices.is_empty() {
- let selected = selected_rows(&source, &matched_indices)?;
- let mut fields = selected.schema().fields().to_vec();
- fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64,
false)));
- let mut columns: Vec<ArrayRef> = selected.columns().to_vec();
- columns.push(Arc::new(Int64Array::from(row_ids)));
- let matched =
RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
- .map_err(|error| {
- invalid(format!("cannot build matched upsert rows:
{error}"))
- })?;
- let mut update =
- super::DataEvolutionWriter::for_row_id(&self.table,
self.update_columns)?;
+ let mut update = super::DataEvolutionWriter::for_row_id(
+ &self.table,
+ self.update_columns.clone(),
+ )?;
if let Some(snapshot_id) = plan.snapshot_id() {
update.pin_read_snapshot(snapshot_id);
}
- update.add_matched_batch(matched)?;
+ for (selected, range) in matched_batches {
+ let mut fields = selected.schema().fields().to_vec();
+ fields.push(Arc::new(Field::new(ROW_ID, DataType::Int64,
false)));
+ let mut columns: Vec<ArrayRef> =
selected.columns().to_vec();
+
columns.push(Arc::new(Int64Array::from(row_ids[range].to_vec())));
+ let matched =
RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+ .map_err(|error| {
+ invalid(format!("cannot build matched upsert rows:
{error}"))
+ })?;
+ update.add_matched_batch(matched)?;
+ }
messages.extend(update.prepare_commit().await?);
}
- if !new_indices.is_empty() {
- let new_rows = selected_rows(&source, &new_indices)?;
+ for group in append_groups {
let mut append = self
.table
.new_write_builder()
.with_commit_user(self.commit_user.clone())?
.new_write()?;
append.with_write_type(
- new_rows
+ group[0]
.schema()
.fields()
.iter()
.map(|field| field.name().to_string())
.collect(),
)?;
- if let Err(error) = append.write_arrow_batch(&new_rows).await {
- append.close().await;
- return Err(error);
+ for batch in group {
+ if let Err(error) = append.write_arrow_batch(&batch).await
{
+ append.close().await;
+ return Err(error);
+ }
}
messages.extend(append.prepare_commit().await?);
}
diff --git a/crates/paimon/tests/table_update_test.rs
b/crates/paimon/tests/table_update_test.rs
index 220f7471..8689dd5e 100644
--- a/crates/paimon/tests/table_update_test.rs
+++ b/crates/paimon/tests/table_update_test.rs
@@ -824,6 +824,78 @@ async fn
partial_upsert_updates_matches_and_appends_only_input_columns() {
);
}
+#[tokio::test]
+async fn heterogeneous_upsert_deduplicates_before_matched_column_validation() {
+ let table = evolution_table().await;
+ seed(&table).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![
+ // Missing update columns on a superseded source row are
irrelevant.
+ batch(&[("id", vec![1]), ("score", vec![999])]),
+ // New rows need not contain the matched-row update projection.
+ batch(&[("id", vec![4]), ("score", vec![400])]),
+ batch(&[("id", vec![1]), ("value", vec![111])]),
+ ],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ let appended = messages
+ .iter()
+ .flat_map(|message| &message.new_files)
+ .filter(|file| {
+ file.write_cols
+ .as_ref()
+ .is_some_and(|cols| cols.contains(&"id".into()))
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(appended.len(), 1);
+ assert_eq!(
+ appended[0].write_cols,
+ Some(vec!["id".into(), "score".into()])
+ );
+ commit(&table, messages).await;
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ use arrow_array::Array;
+ let cols = batch
+ .columns()
+ .iter()
+ .map(|array| array.as_any().downcast_ref::<Int32Array>().unwrap())
+ .collect::<Vec<_>>();
+ for row in 0..batch.num_rows() {
+ rows.push(
+ cols.iter()
+ .map(|col| col.is_valid(row).then(|| col.value(row)))
+ .collect::<Vec<_>>(),
+ );
+ }
+ }
+ rows.sort();
+ assert_eq!(
+ rows,
+ vec![
+ vec![Some(1), Some(111), Some(100)],
+ vec![Some(2), Some(20), Some(200)],
+ vec![Some(3), Some(30), Some(300)],
+ vec![Some(4), None, Some(400)],
+ ]
+ );
+}
+
#[tokio::test]
async fn partial_upsert_requires_update_columns_only_for_matches() {
let table = evolution_table().await;
@@ -855,6 +927,150 @@ async fn
partial_upsert_requires_update_columns_only_for_matches() {
commit(&table, messages).await;
}
+#[tokio::test]
+async fn heterogeneous_upsert_keeps_append_shapes_separate_by_partition() {
+ let path = "memory:/upsert_partition_shapes";
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .column("score", DataType::Int(IntType::new()))
+ .column("region", DataType::Int(IntType::new()))
+ .partition_keys(["region"])
+ .option("row-tracking.enabled", "true")
+ .option("data-evolution.enabled", "true")
+ .build()
+ .unwrap();
+ let (io, table) = memory_table(path, TableSchema::new(0, &schema));
+ setup_dirs(&io, path).await;
+ persist_table_schema(&io, path, table.schema()).await;
+ write_batch(
+ &table,
+ &batch(&[
+ ("id", vec![1, 1]),
+ ("value", vec![10, 20]),
+ ("score", vec![100, 200]),
+ ("region", vec![1, 2]),
+ ]),
+ )
+ .await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![
+ batch(&[
+ ("region", vec![1, 1]),
+ ("value", vec![111, 40]),
+ ("id", vec![1, 4]),
+ ]),
+ batch(&[("id", vec![4]), ("score", vec![400]), ("region",
vec![2])]),
+ batch(&[
+ ("id", vec![1]),
+ ("value", vec![222]),
+ ("score", vec![999]),
+ ("region", vec![2]),
+ ]),
+ batch(&[("id", vec![4]), ("value", vec![41]), ("region",
vec![1])]),
+ ],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ let shapes = messages
+ .iter()
+ .flat_map(|message| &message.new_files)
+ .filter_map(|file| file.write_cols.as_ref())
+ .filter(|columns| columns.contains(&"id".into()))
+ .cloned()
+ .collect::<std::collections::HashSet<_>>();
+ assert_eq!(
+ shapes,
+ std::collections::HashSet::from([
+ vec!["id".into(), "value".into(), "region".into()],
+ vec!["id".into(), "score".into(), "region".into()],
+ ])
+ );
+ commit(&table, messages).await;
+ let reader = table.new_read_builder();
+ let plan = reader.new_scan().plan().await.unwrap();
+ let batches: Vec<RecordBatch> = reader
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ let mut rows = Vec::new();
+ for batch in batches {
+ use arrow_array::Array;
+ let cols = batch
+ .columns()
+ .iter()
+ .map(|array| array.as_any().downcast_ref::<Int32Array>().unwrap())
+ .collect::<Vec<_>>();
+ for row in 0..batch.num_rows() {
+ rows.push(
+ cols.iter()
+ .map(|col| col.is_valid(row).then(|| col.value(row)))
+ .collect::<Vec<_>>(),
+ );
+ }
+ }
+ rows.sort();
+ assert_eq!(
+ rows,
+ vec![
+ vec![Some(1), Some(111), Some(100), Some(1)],
+ vec![Some(1), Some(222), Some(200), Some(2)],
+ vec![Some(4), None, Some(400), Some(2)],
+ vec![Some(4), Some(41), None, Some(1)],
+ ]
+ );
+}
+
+#[tokio::test]
+async fn heterogeneous_upsert_validates_only_surviving_shapes_before_staging()
{
+ let table = evolution_table().await;
+ seed(&table).await;
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let before = parquet_files(&table).await;
+ for batches in [
+ // The winning matched row must supply the requested update columns.
+ vec![
+ batch(&[("id", vec![1]), ("value", vec![111])]),
+ batch(&[("id", vec![2]), ("score", vec![222])]),
+ ],
+ // Appended rows in one partition cannot silently acquire NULL padding.
+ vec![
+ batch(&[("id", vec![4]), ("value", vec![40])]),
+ batch(&[("id", vec![5]), ("score", vec![500])]),
+ ],
+ ] {
+ assert!(update
+ .upsert_by_arrow_with_key(batches, vec!["id".into()])
+ .await
+ .is_err());
+ assert_eq!(parquet_files(&table).await, before);
+ }
+ let messages = update
+ .upsert_by_arrow_with_key(
+ vec![
+ batch(&[("id", vec![4]), ("value", vec![40])]),
+ batch(&[("id", vec![4]), ("score", vec![400])]),
+ ],
+ vec!["id".into()],
+ )
+ .await
+ .unwrap();
+ assert_eq!(
+ messages[0].new_files[0].write_cols,
+ Some(vec!["id".into(), "score".into()])
+ );
+ commit(&table, messages).await;
+}
+
#[tokio::test]
async fn partial_upsert_rejects_ambiguous_unknown_and_inconsistent_input() {
let table = evolution_table().await;
