[core] Support logical Arrow keys and selected upsert values

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()
+        );
+    }
+}

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;

crates/paimon/src/table/table_merge_into.rs

@@ -427,23 +427,37 @@ 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 =
             RowConverter::new(fields).map_err(|error| invalid(error.to_string()))?;
         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()))?;