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 b72adaee [core] Support logical Arrow keys and selected upsert values
(#1147)
b72adaee is described below
commit b72adaeeecea9ae946082c8a218afcdf3fb9bb51
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Oct 17 21:19:10 2026 +0800
[core] Support logical Arrow keys and selected upsert values (#1147)
---
crates/paimon/src/table/key_normalization.rs | 222 ++++++++++++
crates/paimon/src/table/mod.rs | 1 +
crates/paimon/src/table/table_merge_into.rs | 22 +-
crates/paimon/src/table/table_merge_into_tests.rs | 145 ++++++++
crates/paimon/src/table/table_update.rs | 2 +
crates/paimon/src/table/table_upsert.rs | 401 ++++++++++++++++++++--
crates/paimon/src/table/update_input.rs | 23 +-
crates/paimon/src/table/write_batch_normalize.rs | 173 +++++++++-
8 files changed, 957 insertions(+), 32 deletions(-)
diff --git a/crates/paimon/src/table/key_normalization.rs
b/crates/paimon/src/table/key_normalization.rs
new file mode 100644
index 00000000..3ea04dc5
--- /dev/null
+++ b/crates/paimon/src/table/key_normalization.rs
@@ -0,0 +1,222 @@
+// 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.
+
+//! Logical keys used by row-level matching, independent of Arrow encoding.
+
+use arrow_array::{new_null_array, ArrayRef};
+use arrow_schema::DataType;
+
+use super::update_input::{cast_update_value, decode_dictionary, CastMode};
+use super::write_batch_normalize::normalize_write_array;
+
+/// Unify Arrow representations without changing key identity. Dictionary
+/// indices are batch-local; integer widths, decimal scales and temporal units
+/// may differ only when conversion is exact. Callers validate the logical type
+/// families and decide whether NULL keys match. Payload write rules are
separate.
+pub(super) fn normalize_match_key(
+ array: &ArrayRef,
+ expected: &DataType,
+) -> crate::Result<ArrayRef> {
+ let decoded = decode_dictionary(array)?;
+ if decoded.data_type() == &DataType::Null {
+ return Ok(new_null_array(expected, decoded.len()));
+ }
+ let source = decoded.data_type();
+ if source.is_decimal() && matches!(expected, DataType::Decimal128(..)) {
+ let source_scale = match source {
+ DataType::Decimal32(_, scale)
+ | DataType::Decimal64(_, scale)
+ | DataType::Decimal128(_, scale)
+ | DataType::Decimal256(_, scale) => *scale,
+ _ => unreachable!(),
+ };
+ let DataType::Decimal128(_, target_scale) = expected else {
+ unreachable!()
+ };
+ let difference = i16::from(*target_scale) - i16::from(source_scale);
+ let text_fallback = || {
+ // Arrow's decimal rescaler can reject a large scale difference
+ // before inspecting values, even for zero and NULL. Reuse the
+ // exact text parser as a fallback; it still rejects nonzero values
+ // that exceed precision or lose fractional digits.
+ let text = cast_update_value(&decoded, &DataType::Utf8,
CastMode::Assignment)?;
+ cast_update_value(&text, expected, CastMode::Assignment)
+ };
+ // Arrow computes scale differences in i8 and uses a Decimal128 power
+ // table. Skip that kernel before either can overflow.
+ return if !(-38..=38).contains(&difference) {
+ text_fallback()
+ } else {
+ cast_update_value(&decoded, expected,
CastMode::Assignment).or_else(|_| text_fallback())
+ };
+ }
+ let exact_family = source.is_integer() && expected.is_integer()
+ || source.is_decimal() && expected.is_decimal()
+ || matches!(
+ (source, expected),
+ (
+ DataType::Date32 | DataType::Date64,
+ DataType::Date32 | DataType::Date64
+ ) | (
+ DataType::Time32(_) | DataType::Time64(_),
+ DataType::Time32(_) | DataType::Time64(_)
+ )
+ )
+ || matches!((source, expected), (DataType::Timestamp(_, from),
DataType::Timestamp(_, to))
+ if from.is_some() == to.is_some());
+ if exact_family {
+ cast_update_value(&decoded, expected, CastMode::Assignment)
+ } else {
+ normalize_write_array(&decoded, expected)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::types::Int8Type;
+ use arrow_array::{
+ Array, Decimal128Array, DictionaryArray, Int32Array, Int64Array,
Int8Array, StringArray,
+ TimestampMicrosecondArray, TimestampMillisecondArray,
+ };
+ use arrow_schema::TimeUnit;
+ use std::sync::Arc;
+
+ #[test]
+ fn dictionary_keys_use_value_validity_as_well_as_index_validity() {
+ let input: ArrayRef = Arc::new(
+ DictionaryArray::<Int8Type>::try_new(
+ Int8Array::from(vec![Some(0), None, Some(1), Some(0)]),
+ Arc::new(StringArray::from(vec![None, Some("key")])),
+ )
+ .unwrap(),
+ );
+ let output = normalize_match_key(&input, &DataType::Utf8).unwrap();
+ assert_eq!(
+ output.as_any().downcast_ref::<StringArray>().unwrap(),
+ &StringArray::from(vec![None, None, Some("key"), None])
+ );
+ }
+
+ #[test]
+ fn integer_keys_are_exact_and_keep_nulls() {
+ let input: ArrayRef = Arc::new(Int64Array::from(vec![
+ Some(-1),
+ None,
+ Some(i32::MAX as i64),
+ ]));
+ let output = normalize_match_key(&input, &DataType::Int32).unwrap();
+ assert_eq!(
+ output.as_any().downcast_ref::<Int32Array>().unwrap(),
+ &Int32Array::from(vec![Some(-1), None, Some(i32::MAX)])
+ );
+ let overflow: ArrayRef = Arc::new(Int64Array::from(vec![i32::MAX as
i64 + 1]));
+ assert!(normalize_match_key(&overflow, &DataType::Int32).is_err());
+ let text: ArrayRef = Arc::new(StringArray::from(vec!["42"]));
+ assert!(normalize_match_key(&text, &DataType::Int32).is_err());
+ }
+
+ #[test]
+ fn decimal_keys_cannot_round_into_existing_keys() {
+ let array = |value| {
+ Arc::new(
+ Decimal128Array::from(vec![value])
+ .with_precision_and_scale(15, 4)
+ .unwrap(),
+ ) as ArrayRef
+ };
+ let expected = DataType::Decimal128(10, 2);
+ let output = normalize_match_key(&array(12500), &expected).unwrap();
+ assert_eq!(
+ output
+ .as_any()
+ .downcast_ref::<Decimal128Array>()
+ .unwrap()
+ .value(0),
+ 125
+ );
+ assert!(normalize_match_key(&array(12501), &expected).is_err());
+ }
+
+ #[test]
+ fn decimal_zero_and_null_keys_rescale_across_large_scale_differences() {
+ for (from, to) in [(-2, 38), (38, -2), (-90, 38), (-128, 38), (38,
-128)] {
+ let input: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![Some(0), None])
+ .with_precision_and_scale(38, from)
+ .unwrap(),
+ );
+ let expected = DataType::Decimal128(38, to);
+ let output = normalize_match_key(&input, &expected).unwrap();
+ assert_eq!(output.data_type(), &expected);
+ assert_eq!(
+ output.as_any().downcast_ref::<Decimal128Array>().unwrap(),
+ &Decimal128Array::from(vec![Some(0), None])
+ .with_precision_and_scale(38, to)
+ .unwrap()
+ );
+ let nonzero: ArrayRef = Arc::new(
+ Decimal128Array::from(vec![1])
+ .with_precision_and_scale(38, from)
+ .unwrap(),
+ );
+ assert!(normalize_match_key(&nonzero, &expected).is_err());
+ }
+ }
+
+ #[test]
+ fn temporal_keys_require_lossless_units_and_preserve_instant() {
+ let input: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![
+ Some(-1000),
+ None,
+ Some(2000),
+ ]));
+ let output =
+ normalize_match_key(&input,
&DataType::Timestamp(TimeUnit::Millisecond, None)).unwrap();
+ assert_eq!(
+ output
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap(),
+ &TimestampMillisecondArray::from(vec![Some(-1), None, Some(2)])
+ );
+ let fractional: ArrayRef =
Arc::new(TimestampMicrosecondArray::from(vec![-1001]));
+ assert!(normalize_match_key(
+ &fractional,
+ &DataType::Timestamp(TimeUnit::Millisecond, None)
+ )
+ .is_err());
+ let local: ArrayRef =
+
Arc::new(TimestampMicrosecondArray::from(vec![0]).with_timezone("America/New_York"));
+ let output = normalize_match_key(
+ &local,
+ &DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
+ )
+ .unwrap();
+ assert_eq!(
+ output
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ 0
+ );
+ assert!(
+ normalize_match_key(&local,
&DataType::Timestamp(TimeUnit::Microsecond, None)).is_err()
+ );
+ }
+}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 33382b11..4b0b155d 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -133,6 +133,7 @@ mod query_auth;
mod system_table;
pub use system_table::SystemRead;
+mod key_normalization;
mod read_builder;
mod read_limit;
pub mod referenced_files;
diff --git a/crates/paimon/src/table/table_merge_into.rs
b/crates/paimon/src/table/table_merge_into.rs
index 0004e3b7..00652277 100644
--- a/crates/paimon/src/table/table_merge_into.rs
+++ b/crates/paimon/src/table/table_merge_into.rs
@@ -427,14 +427,18 @@ pub(super) async fn merge_into(
.map_err(|error| invalid(error.to_string()))?
.data_type()
};
- if source_type != target_type
+ let mut logical_source_type = source_type;
+ while let DataType::Dictionary(_, value) = logical_source_type
{
+ logical_source_type = value;
+ }
+ if logical_source_type != target_type
||
!super::upsert_key_matcher::supported_key_type(target_type)
{
return Err(invalid(
"MERGE ON key types must match and support row
encoding",
));
}
- Ok(SortField::new(source_type.clone()))
+ Ok(SortField::new(target_type.clone()))
})
.collect::<crate::Result<_>>()?;
let key_converter =
@@ -442,8 +446,18 @@ pub(super) async fn merge_into(
for (chunk, batch) in source.batches.iter().enumerate() {
let keys = on
.iter()
- .map(|(_, name)| batch.column_by_name(name).unwrap().clone())
- .collect::<Vec<_>>();
+ .map(|(target, name)| {
+ let expected = if target == ROW_ID {
+ &DataType::Int64
+ } else {
+ schema.field_with_name(target).unwrap().data_type()
+ };
+ super::key_normalization::normalize_match_key(
+ batch.column_by_name(name).unwrap(),
+ expected,
+ )
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
let rows = key_converter
.convert_columns(&keys)
.map_err(|error| invalid(error.to_string()))?;
diff --git a/crates/paimon/src/table/table_merge_into_tests.rs
b/crates/paimon/src/table/table_merge_into_tests.rs
index b20c4753..2c520d21 100644
--- a/crates/paimon/src/table/table_merge_into_tests.rs
+++ b/crates/paimon/src/table/table_merge_into_tests.rs
@@ -107,6 +107,151 @@ fn update(condition: Option<MergeCondition>) ->
WhenMatched {
}
}
+fn dictionary_batch(ids: Vec<Option<i32>>, values: Vec<Option<i32>>) ->
RecordBatch {
+ let keys = arrow_cast::cast(
+ &Int32Array::from(ids),
+ &arrow_schema::DataType::Dictionary(
+ Box::new(arrow_schema::DataType::Int8),
+ Box::new(arrow_schema::DataType::Int32),
+ ),
+ )
+ .unwrap();
+ let schema = arrow_schema::Schema::new(vec![
+ arrow_schema::Field::new("id", keys.data_type().clone(), true),
+ arrow_schema::Field::new("value", arrow_schema::DataType::Int32, true),
+ ]);
+ RecordBatch::try_new(
+ Arc::new(schema),
+ vec![keys, Arc::new(Int32Array::from(values))],
+ )
+ .unwrap()
+}
+
+#[tokio::test]
+async fn dictionary_merge_keys_use_values_and_sql_null_semantics() {
+ let target = table_at("dictionary_merge");
+ seed(
+ &target,
+ &batch(
+ vec![Some(1), Some(1), Some(2), None],
+ vec![Some(10), Some(11), Some(20), Some(90)],
+ ),
+ )
+ .await;
+ // Every non-null source key has dictionary index zero in its own chunk.
+ let source = vec![
+ dictionary_batch(vec![Some(1)], vec![Some(42)]),
+ dictionary_batch(vec![Some(3)], vec![Some(30)]),
+ dictionary_batch(vec![None], vec![Some(91)]),
+ ];
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Batches(source),
+ vec![("id".into(), "id".into())],
+ vec![update(None)],
+ vec![WhenNotMatched {
+ condition: None,
+ assignments: vec![
+ ("id".into(), MergeAssignment::SourceColumn("id".into())),
+ (
+ "value".into(),
+ MergeAssignment::SourceColumn("value".into()),
+ ),
+ ],
+ }],
+ )
+ .await
+ .unwrap();
+ TableCommit::new(target.clone(), "merge".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(
+ rows(&target).await,
+ vec![
+ (None, Some(90)),
+ (None, Some(91)),
+ (Some(1), Some(42)),
+ (Some(1), Some(42)),
+ (Some(2), Some(20)),
+ (Some(3), Some(30))
+ ]
+ );
+}
+
+#[tokio::test]
+async fn dictionary_merge_cardinality_is_checked_before_conditions() {
+ let target = table_at("dictionary_cardinality");
+ seed(&target, &batch(vec![Some(1)], vec![Some(10)])).await;
+ let source = vec![
+ dictionary_batch(vec![Some(1)], vec![Some(11)]),
+ dictionary_batch(vec![Some(3), Some(1)], vec![Some(30), Some(12)]),
+ ];
+ let calls = Arc::new(AtomicUsize::new(0));
+ let observed = calls.clone();
+ let condition = MergeCondition {
+ target_columns: Vec::new(),
+ source_columns: Vec::new(),
+ evaluate: Arc::new(move |batch| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ Box::pin(async move { Ok(BooleanArray::from(vec![false;
batch.num_rows()])) })
+ }),
+ };
+ let error = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Batches(source),
+ vec![("id".into(), "id".into())],
+ vec![update(Some(condition))],
+ Vec::new(),
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("multiple source rows"));
+ assert_eq!(calls.load(Ordering::Relaxed), 0);
+ assert_eq!(rows(&target).await, vec![(Some(1), Some(10))]);
+}
+
+#[tokio::test]
+async fn dictionary_merge_duplicate_unconditional_delete_is_allowed() {
+ let target = table_at("dictionary_delete");
+ seed(
+ &target,
+ &batch(vec![Some(1), Some(2)], vec![Some(10), Some(20)]),
+ )
+ .await;
+ let source = vec![
+ dictionary_batch(vec![Some(1)], vec![Some(11)]),
+ dictionary_batch(vec![Some(3), Some(1)], vec![Some(30), Some(12)]),
+ ];
+ let messages = target
+ .new_write_builder()
+ .new_update()
+ .unwrap()
+ .merge_into(
+ MergeSource::Batches(source),
+ vec![("id".into(), "id".into())],
+ vec![WhenMatched {
+ condition: None,
+ delete: true,
+ assignments: Vec::new(),
+ }],
+ Vec::new(),
+ )
+ .await
+ .unwrap();
+ TableCommit::new(target.clone(), "merge".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(rows(&target).await, vec![(Some(2), Some(20))]);
+}
+
#[tokio::test]
async fn core_merge_preserves_row_ids_and_handles_sql_null_keys() {
let table = table();
diff --git a/crates/paimon/src/table/table_update.rs
b/crates/paimon/src/table/table_update.rs
index 546dc8a2..fdaa256e 100644
--- a/crates/paimon/src/table/table_update.rs
+++ b/crates/paimon/src/table/table_update.rs
@@ -243,6 +243,8 @@ 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.
+ /// After last-write-wins, matched values follow row-ID update conversion
+ /// rules; new rows must satisfy the ordinary Arrow write type.
/// 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.
diff --git a/crates/paimon/src/table/table_upsert.rs
b/crates/paimon/src/table/table_upsert.rs
index 8869e527..4b51ff66 100644
--- a/crates/paimon/src/table/table_upsert.rs
+++ b/crates/paimon/src/table/table_upsert.rs
@@ -150,9 +150,9 @@ 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,
- /// but may carry different non-key columns, preserving absent fields.
+ /// differ from the table schema. Keep original payloads, including keys,
+ /// until last-write-wins and matching have selected their action.
+ /// Batches 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();
@@ -170,48 +170,116 @@ impl TableUpsert {
}
let mut columns = Vec::with_capacity(target.fields().len());
let mut fields = Vec::new();
+ for field in target.fields() {
+ let Some(column) = batch.column_by_name(field.name()) else {
+ continue;
+ };
+ columns.push(column.clone());
+
fields.push(input_schema.field_with_name(field.name()).unwrap().clone());
+ }
+ let ordered = RecordBatch::try_new(Arc::new(Schema::new(fields)),
columns)
+ .map_err(|error| invalid(format!("cannot order upsert columns:
{error}")))?;
+ self.source.push(ordered);
+ Ok(())
+ }
+
+ /// Appends use the ordinary write contract; matched values use the row-ID
+ /// updater's conversion rules. Neither action touches discarded source
rows.
+ fn normalize_selected_batch(
+ &self,
+ batch: RecordBatch,
+ append: bool,
+ ) -> crate::Result<RecordBatch> {
+ let target =
crate::arrow::build_target_arrow_schema(self.table.schema().fields())?;
+ let update_fields = super::data_evolution_fields::project_by_paths(
+ self.table.schema().fields(),
+ &self.update_columns,
+ )?;
let inline =
CoreOptions::new(self.table.schema().options()).blob_inline_fields();
+ let mut fields = Vec::new();
+ let mut columns = Vec::new();
for (index, field) in target.fields().iter().enumerate() {
let Some(column) = batch.column_by_name(field.name()) else {
+ // The row-ID updater diagnoses missing update columns. Absent
+ // append columns stay absent in the file's write_cols
metadata.
continue;
};
let logical = &self.table.schema().fields()[index];
- let normalized =
- if logical.data_type().is_blob_file_field() &&
!inline.contains(logical.name()) {
+ let update_field = update_fields
+ .iter()
+ .find(|updated| updated.name() == field.name());
+ let normalized = if !append && update_field.is_none() {
+ Ok((
+ batch
+ .schema()
+ .field_with_name(field.name())
+ .unwrap()
+ .clone(),
+ column.clone(),
+ ))
+ } else if logical.data_type().is_blob_file_field()
+ && !inline.contains(logical.name())
+ && (append || matches!(logical.data_type(),
crate::spec::DataType::Map(_)))
+ {
+ // Named-row Blob maps carry nullable keys and decimal text in
+ // their dedicated physical layout. Keep that transport intact.
+ if append {
normalize_blob_write_field(column, field)
} else {
- normalize_write_array(column, field.data_type())
- .map(|column| (field.as_ref().clone(), column))
- };
+
super::write_batch_normalize::normalize_blob_update_field(column, field)
+ }
+ } else if append {
+ normalize_write_array(column, field.data_type())
+ .map(|column| (field.as_ref().clone(), column))
+ } else {
+ Ok((
+ batch
+ .schema()
+ .field_with_name(field.name())
+ .unwrap()
+ .clone(),
+ column.clone(),
+ ))
+ };
let (field, column) = normalized.map_err(|error| {
invalid(format!("Invalid upsert column '{}': {error}",
field.name()))
})?;
+ fields.push(field);
columns.push(column);
- fields.push(Arc::new(field));
}
- let ordered = RecordBatch::try_new(Arc::new(Schema::new(fields)),
columns)
- .map_err(|error| invalid(format!("cannot order upsert columns:
{error}")))?;
- self.source.push(ordered);
- Ok(())
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
+ .map_err(|error| invalid(format!("cannot normalize selected upsert
rows: {error}")))
}
/// 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 target =
crate::arrow::build_target_arrow_schema(self.table.schema().fields())?;
+ let fields = target
+ .fields()
+ .iter()
+ .filter(|field| self.keys.contains(field.name()))
+ .cloned()
+ .collect::<Vec<_>>();
+ let key_schema = Arc::new(Schema::new(fields));
let batches = self
.source
.iter()
.map(|batch| {
- let columns = batch
- .schema()
+ let columns = key_schema
.fields()
.iter()
- .enumerate()
- .filter(|(_, field)| self.keys.contains(field.name()))
- .map(|(index, _)| index)
- .collect::<Vec<_>>();
- batch
- .project(&columns)
+ .map(|field| {
+ super::key_normalization::normalize_match_key(
+ batch.column_by_name(field.name()).unwrap(),
+ field.data_type(),
+ )
+ .map_err(|error| {
+ invalid(format!("Invalid upsert key '{}':
{error}", field.name()))
+ })
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+ RecordBatch::try_new(key_schema.clone(), columns)
.map_err(|error| invalid(error.to_string()))
})
.collect::<crate::Result<Vec<_>>>()?;
@@ -277,6 +345,7 @@ impl TableUpsert {
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 batch = self.normalize_selected_batch(batch, true)?;
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);
@@ -357,6 +426,7 @@ impl TableUpsert {
update.pin_read_snapshot(snapshot_id);
}
for (selected, range) in matched_batches {
+ let selected = self.normalize_selected_batch(selected,
false)?;
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();
@@ -404,7 +474,296 @@ impl TableUpsert {
#[cfg(test)]
mod tests {
use super::*;
+ use crate::catalog::Identifier;
+ use crate::io::FileIOBuilder;
use crate::spec::{DataType as PaimonType, DecimalType, TimestampType,
VarCharType};
+ use crate::spec::{IntType, Schema as PaimonSchema, TableSchema};
+ use crate::table::TableCommit;
+ use arrow_array::{Array, Int32Array, StringArray};
+
+ async fn numeric_table(path: &str) -> Table {
+ let schema = PaimonSchema::builder()
+ .column("key", PaimonType::Int(IntType::new()))
+ .column("value", PaimonType::Int(IntType::new()))
+ .column("keep", 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", "upsert"),
+ path.into(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+ let initial = RecordBatch::try_from_iter([
+ ("key", Arc::new(Int32Array::from(vec![1, 1, 2])) as ArrayRef),
+ ("value", Arc::new(Int32Array::from(vec![0; 3])) as ArrayRef),
+ (
+ "keep",
+ Arc::new(Int32Array::from(vec![10, 11, 12])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ let mut write = table.new_write_builder().new_write().unwrap();
+ write.write_arrow_batch(&initial).await.unwrap();
+ TableCommit::new(table.clone(), "seed".into())
+ .commit(write.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ write.close().await;
+ table
+ }
+
+ fn numeric_source(keys: Vec<i32>, value: ArrayRef) -> RecordBatch {
+ RecordBatch::try_from_iter([
+ ("key", Arc::new(Int32Array::from(keys)) as ArrayRef),
+ ("value", value),
+ ])
+ .unwrap()
+ }
+
+ async fn numeric_rows(table: &Table) -> Vec<(i64, i32, Option<i32>,
Option<i32>)> {
+ let mut builder = table.new_read_builder();
+ builder
+ .with_projection(&["_ROW_ID", "key", "value", "keep"])
+ .unwrap();
+ let plan = builder.new_scan().plan().await.unwrap();
+ let batches: Vec<_> = 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(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let keys = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let values = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ let keep = batch
+ .column(3)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ rows.extend((0..batch.num_rows()).map(|row| {
+ (
+ ids.value(row),
+ keys.value(row),
+ (!values.is_null(row)).then(|| values.value(row)),
+ (!keep.is_null(row)).then(|| keep.value(row)),
+ )
+ }));
+ }
+ rows.sort();
+ rows
+ }
+
+ #[tokio::test]
+ async fn
upsert_casts_surviving_matched_values_and_ignores_other_payloads() {
+ let table = numeric_table("memory:/upsert_selected").await;
+ let first = numeric_source(vec![1],
Arc::new(StringArray::from(vec!["not-an-int"])));
+ let mut last = numeric_source(vec![1],
Arc::new(StringArray::from(vec!["42"])));
+ last = RecordBatch::try_from_iter([
+ ("key", last.column(0).clone()),
+ ("value", last.column(1).clone()),
+ (
+ "keep",
+ Arc::new(StringArray::from(vec!["unwritten"])) as ArrayRef,
+ ),
+ ])
+ .unwrap();
+ 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![first, last], vec!["key".into()])
+ .await
+ .unwrap();
+ assert!(messages
+ .iter()
+ .all(|message| message.check_from_snapshot == Some(1)));
+ TableCommit::new(table.clone(), "update".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(
+ numeric_rows(&table).await,
+ vec![
+ (0, 1, Some(42), Some(10)),
+ (1, 1, Some(42), Some(11)),
+ (2, 2, Some(0), Some(12)),
+ ]
+ );
+ }
+
+ #[tokio::test]
+ async fn upsert_keeps_update_casts_separate_from_append_validation() {
+ let table = numeric_table("memory:/upsert_mixed_types").await;
+ let matched = numeric_source(vec![1],
Arc::new(StringArray::from(vec!["42"])));
+ let append = numeric_source(vec![3],
Arc::new(Int32Array::from(vec![7])));
+ 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![matched, append],
vec!["key".into()])
+ .await
+ .unwrap();
+ TableCommit::new(table.clone(), "update".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(
+ numeric_rows(&table).await,
+ vec![
+ (0, 1, Some(42), Some(10)),
+ (1, 1, Some(42), Some(11)),
+ (2, 2, Some(0), Some(12)),
+ (3, 3, Some(7), None),
+ ]
+ );
+ }
+
+ #[tokio::test]
+ async fn upsert_rejects_invalid_appends_before_staging_matched_updates() {
+ let table = numeric_table("memory:/upsert_invalid_append").await;
+ let batch = numeric_source(vec![1, 3],
Arc::new(StringArray::from(vec!["42", "7"])));
+ let mut update = table.new_write_builder().new_update().unwrap();
+ update.with_update_type(vec!["value".into()]).unwrap();
+ let error = update
+ .upsert_by_arrow_with_key(vec![batch], vec!["key".into()])
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("Invalid upsert column 'value'"));
+ assert_eq!(
+ table
+ .snapshot_manager()
+ .get_latest_snapshot_id()
+ .await
+ .unwrap(),
+ Some(1)
+ );
+ assert_eq!(
+ numeric_rows(&table).await,
+ vec![
+ (0, 1, Some(0), Some(10)),
+ (1, 1, Some(0), Some(11)),
+ (2, 2, Some(0), Some(12)),
+ ]
+ );
+ }
+
+ #[tokio::test]
+ async fn upsert_matched_binary_values_use_existing_row_id_casts() {
+ use arrow_array::{BinaryArray, LargeBinaryArray};
+ let schema = PaimonSchema::builder()
+ .column("key", PaimonType::Int(IntType::new()))
+ .column(
+ "value",
+ PaimonType::VarBinary(crate::spec::VarBinaryType::default()),
+ )
+ .option("data-evolution.enabled", "true")
+ .option("row-tracking.enabled", "true")
+ .build()
+ .unwrap();
+ let table = Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("db", "upsert"),
+ "memory:/upsert_binary".into(),
+ TableSchema::new(0, &schema),
+ None,
+ );
+ let initial = numeric_source(
+ vec![1],
+ Arc::new(BinaryArray::from(vec![b"old".as_slice()])),
+ );
+ let mut writer = table.new_write_builder().new_write().unwrap();
+ writer.write_arrow_batch(&initial).await.unwrap();
+ TableCommit::new(table.clone(), "seed".into())
+ .commit(writer.prepare_commit().await.unwrap())
+ .await
+ .unwrap();
+ writer.close().await;
+ let update = table.new_write_builder().new_update().unwrap();
+ let input = numeric_source(
+ vec![1],
+ Arc::new(LargeBinaryArray::from(vec![b"\xff\0new".as_slice()])),
+ );
+ let messages = update
+ .upsert_by_arrow_with_key(vec![input], vec!["key".into()])
+ .await
+ .unwrap();
+ TableCommit::new(table.clone(), "update".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ let mut builder = table.new_read_builder();
+ builder.with_projection(&["value"]).unwrap();
+ let plan = builder.new_scan().plan().await.unwrap();
+ let batches: Vec<_> = builder
+ .new_read()
+ .unwrap()
+ .to_arrow(plan.splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
1);
+ let values = batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap();
+ assert_eq!(values.value(0), b"\xff\0new");
+ }
+
+ #[tokio::test]
+ async fn upsert_matches_integer_key_widths_but_keeps_append_input_types() {
+ let table = numeric_table("memory:/upsert_integer_key_widths").await;
+ let source = |key, value| {
+ RecordBatch::try_from_iter([
+ ("key", Arc::new(Int64Array::from(vec![key])) as ArrayRef),
+ ("value", Arc::new(Int32Array::from(vec![value])) as ArrayRef),
+ ])
+ .unwrap()
+ };
+ 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![source(1, 42)], vec!["key".into()])
+ .await
+ .unwrap();
+ TableCommit::new(table.clone(), "update".into())
+ .commit(messages)
+ .await
+ .unwrap();
+ assert_eq!(
+ numeric_rows(&table).await,
+ vec![
+ (0, 1, Some(42), Some(10)),
+ (1, 1, Some(42), Some(11)),
+ (2, 2, Some(0), Some(12))
+ ]
+ );
+ let error = update
+ .upsert_by_arrow_with_key(vec![source(3, 7)], vec!["key".into()])
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("Invalid upsert column 'key'"));
+ assert_eq!(numeric_rows(&table).await.len(), 3);
+ }
#[tokio::test]
async fn upsert_normalizes_string_offsets_before_matching_and_writing() {
diff --git a/crates/paimon/src/table/update_input.rs
b/crates/paimon/src/table/update_input.rs
index d68de118..f6a809db 100644
--- a/crates/paimon/src/table/update_input.rs
+++ b/crates/paimon/src/table/update_input.rs
@@ -98,7 +98,7 @@ pub(super) fn normalize_row_ids(batch: RecordBatch) ->
crate::Result<RecordBatch
RecordBatch::try_new(schema, columns).map_err(|error|
invalid(error.to_string()))
}
-fn decode_dictionary(array: &ArrayRef) -> crate::Result<ArrayRef> {
+pub(super) fn decode_dictionary(array: &ArrayRef) -> crate::Result<ArrayRef> {
let mut array = array.clone();
while let DataType::Dictionary(_, values) = array.data_type() {
array = arrow_cast::cast_with_options(
@@ -2014,6 +2014,27 @@ mod tests {
assert!(cast_assignment(&input, &target).is_err());
}
+ #[test]
+ fn row_update_casts_allow_text_numeric_and_null_binary_values() {
+ let text: ArrayRef = Arc::new(StringArray::from(vec![Some("42"),
None]));
+ let binary = cast_update_value(&text, &DataType::Binary,
CastMode::RowUpdate).unwrap();
+ let integers = cast_update_value(&binary, &DataType::Int32,
CastMode::RowUpdate).unwrap();
+ assert_eq!(
+ integers.to_data(),
+ Int32Array::from(vec![Some(42), None]).to_data()
+ );
+ assert_eq!(
+ cast_update_value(
+ &new_null_array(&DataType::Null, 2),
+ &DataType::Binary,
+ CastMode::RowUpdate
+ )
+ .unwrap()
+ .null_count(),
+ 2
+ );
+ }
+
#[test]
fn binary_assignments_parse_text_and_never_reinterpret_numeric_bytes() {
let binary: ArrayRef = Arc::new(arrow_array::BinaryArray::from(vec![
diff --git a/crates/paimon/src/table/write_batch_normalize.rs
b/crates/paimon/src/table/write_batch_normalize.rs
index e7de7ec7..44b689b0 100644
--- a/crates/paimon/src/table/write_batch_normalize.rs
+++ b/crates/paimon/src/table/write_batch_normalize.rs
@@ -65,6 +65,33 @@ pub(super) fn blob_map_row_type(expected: &DataType) ->
Option<DataType> {
pub(super) fn normalize_blob_write_field(
array: &ArrayRef,
expected: &Field,
+) -> Result<(Field, ArrayRef)> {
+ normalize_blob_field(array, expected, true)
+}
+
+/// Normalize Blob map keys and their physical transport, retaining the value
+/// input type for the row-ID updater's conversion at the existing read layout.
+pub(super) fn normalize_blob_update_field(
+ array: &ArrayRef,
+ expected: &Field,
+) -> Result<(Field, ArrayRef)> {
+ let array = super::update_input::decode_dictionary(array)?;
+ if matches!(array.data_type(), DataType::Map(..) | DataType::List(_)) {
+ normalize_blob_field(&array, expected, false)
+ } else {
+ // NULLs and other row-ID constructor inputs retain their conversion
+ // at the existing DE writer position, rather than the declared schema.
+ Ok((
+ expected.clone().with_data_type(array.data_type().clone()),
+ array,
+ ))
+ }
+}
+
+fn normalize_blob_field(
+ array: &ArrayRef,
+ expected: &Field,
+ normalize_values: bool,
) -> Result<(Field, ArrayRef)> {
let Some(target) = blob_map_row_type(expected.data_type()) else {
return Ok((
@@ -121,17 +148,39 @@ pub(super) fn normalize_blob_write_field(
.enumerate()
.map(|(index, (column, field))| {
if index == 0 {
- normalize_blob_map_key(column, field.data_type())
- } else {
+ if normalize_values {
+ normalize_blob_map_key(column, field.data_type())
+ } else {
+ normalize_blob_update_key(column, field.data_type())
+ }
+ } else if normalize_values {
normalize_write_array(column, field.data_type())
+ } else {
+ Ok(column.clone())
}
})
.collect::<Result<Vec<_>>>()?;
- let entries =
- StructArray::try_new(fields.clone(), columns,
None).map_err(normalization_error)?;
- let list = ListArray::try_new(child.clone(), offsets, Arc::new(entries),
nulls.cloned())
+ let mut fields = fields.to_vec();
+ fields[1] = Arc::new(
+ fields[1]
+ .as_ref()
+ .clone()
+ .with_data_type(columns[1].data_type().clone()),
+ );
+ let fields = fields.into();
+ let entries = StructArray::try_new(fields, columns,
None).map_err(normalization_error)?;
+ let child = Arc::new(
+ child
+ .as_ref()
+ .clone()
+ .with_data_type(entries.data_type().clone()),
+ );
+ let list = ListArray::try_new(child, offsets, Arc::new(entries),
nulls.cloned())
.map_err(normalization_error)?;
- Ok((expected.clone().with_data_type(target), Arc::new(list)))
+ Ok((
+ expected.clone().with_data_type(list.data_type().clone()),
+ Arc::new(list),
+ ))
}
fn normalize_blob_map_key(array: &ArrayRef, expected: &DataType) ->
Result<ArrayRef> {
@@ -165,6 +214,40 @@ fn normalize_blob_map_key(array: &ArrayRef, expected:
&DataType) -> Result<Array
normalize_write_array(array, expected)
}
+fn normalize_blob_update_key(array: &ArrayRef, expected: &DataType) ->
Result<ArrayRef> {
+ let decoded = super::update_input::decode_dictionary(array)?;
+ if matches!(expected, DataType::Decimal128(..)) && decoded.data_type() !=
&DataType::Null {
+ // Blob map keys follow Java Decimal.fromBigDecimal (HALF_UP), rather
+ // than the exact decimal conversion used to match table upsert keys.
+ let text = if decoded.data_type().is_decimal() {
+ super::update_input::cast_update_value(
+ &decoded,
+ &DataType::Utf8,
+ super::update_input::CastMode::Assignment,
+ )?
+ } else {
+ decoded
+ };
+ return normalize_blob_map_key(&text, expected);
+ }
+ if matches!(
+ (decoded.data_type(), expected),
+ (
+ DataType::Time32(_) | DataType::Time64(_),
+ DataType::Time32(_) | DataType::Time64(_)
+ )
+ ) {
+ // Java/Python serialize TIME Blob keys in milliseconds. This is a
+ // value normalization step, unlike lossless table-key matching.
+ return super::update_input::cast_update_value(
+ &decoded,
+ expected,
+ super::update_input::CastMode::RowUpdate,
+ );
+ }
+ super::key_normalization::normalize_match_key(&decoded, expected)
+}
+
/// Decimal.fromBigDecimal: round magnitude HALF_UP, then check precision.
/// Discarded fractional digits are never accumulated into i128, so source
/// decimals can have more digits than the target's maximum precision of 38.
@@ -972,6 +1055,84 @@ mod tests {
assert!(normalize_blob_write_field(&visible_overflow,
&expected).is_err());
}
+ #[test]
+ fn blob_map_updates_normalize_keys_and_transport_before_value_casts() {
+ let entries = StructArray::new(
+ vec![
+ Field::new("k", DataType::Utf8, false),
+ Field::new("v", DataType::Utf8, true),
+ ]
+ .into(),
+ vec![
+ Arc::new(StringArray::from(vec!["1.235"])),
+ Arc::new(StringArray::from(vec!["new"])),
+ ],
+ None,
+ );
+ let entry_field = Arc::new(Field::new("input",
entries.data_type().clone(), false));
+ let offsets = OffsetBuffer::new(vec![0_i32, 1].into());
+ let map: ArrayRef = Arc::new(MapArray::new(
+ entry_field.clone(),
+ offsets.clone(),
+ entries.clone(),
+ None,
+ false,
+ ));
+ let list: ArrayRef = Arc::new(ListArray::new(
+ entry_field,
+ offsets,
+ Arc::new(entries),
+ None,
+ ));
+ let expected = Field::new(
+ "payload",
+ DataType::Map(
+ Arc::new(Field::new(
+ "entries",
+ DataType::Struct(
+ vec![
+ Field::new("key", DataType::Decimal128(10, 2),
false),
+ Field::new("value", DataType::LargeBinary, true),
+ ]
+ .into(),
+ ),
+ false,
+ )),
+ false,
+ ),
+ true,
+ );
+ let normalized = [map, list]
+ .iter()
+ .map(|array| normalize_blob_update_field(array,
&expected).unwrap())
+ .collect::<Vec<_>>();
+ assert_eq!(normalized[0].0, normalized[1].0);
+ assert_eq!(normalized[0].1.to_data(), normalized[1].1.to_data());
+ let list = normalized[0]
+ .1
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let entries = list
+ .values()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(
+ entries
+ .column(0)
+ .as_any()
+ .downcast_ref::<Decimal128Array>()
+ .unwrap()
+ .value(0),
+ 124
+ );
+ // Preserve the original input value type until the row-ID updater
+ // casts against the actual read layout, including its nullability.
+ assert_eq!(entries.column(1).data_type(), &DataType::Utf8);
+ assert!(normalize_blob_write_field(&normalized[0].1,
&expected).is_err());
+ }
+
#[test]
fn sliced_blob_map_inputs_drop_unreferenced_entries() {
let input = blob_map_rows(
