diff --git a/crates/paimon/src/table/key_normalization.rs b/crates/paimon/src/table/key_normalization.rs new file mode 100644 index 000000000..3ea04dc59 --- /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 { + 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::::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::().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::().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::() + .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::().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::() + .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::() + .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 33382b119..4b0b155df 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 0004e3b7a..00652277c 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::>()?; 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::>(); + .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::>>()?; 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 b20c47530..2c520d217 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) -> WhenMatched { } } +fn dictionary_batch(ids: Vec>, values: Vec>) -> 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 546dc8a22..fdaa256ef 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 8869e527e..4b51ff66b 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 { + 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 { + 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::>(); + 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::>(); - 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::>>()?; + RecordBatch::try_new(key_schema.clone(), columns) .map_err(|error| invalid(error.to_string())) }) .collect::>>()?; @@ -277,6 +345,7 @@ impl TableUpsert { fn append_groups(&self, indices: &[usize]) -> crate::Result>> { let mut groups: IndexMap, Vec> = IndexMap::new(); for (batch, _) in self.selected_batches(indices)? { + let batch = self.normalize_selected_batch(batch, true)?; let mut by_partition: IndexMap, Vec> = 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 = 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, 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, Option)> { + 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::() + .unwrap(); + let keys = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let keep = batch + .column(3) + .as_any() + .downcast_ref::() + .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::(), 1); + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .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 d68de1185..f6a809db9 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 crate::Result { +pub(super) fn decode_dictionary(array: &ArrayRef) -> crate::Result { 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 e7de7ec7e..44b689b0b 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 { 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::>>()?; - 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 { @@ -165,6 +214,40 @@ fn normalize_blob_map_key(array: &ArrayRef, expected: &DataType) -> Result Result { + 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::>(); + 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::() + .unwrap(); + let entries = list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + entries + .column(0) + .as_any() + .downcast_ref::() + .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(