diff --git a/bindings/python/src/predicate.rs b/bindings/python/src/predicate.rs index df15470bd..85d066e5f 100644 --- a/bindings/python/src/predicate.rs +++ b/bindings/python/src/predicate.rs @@ -302,6 +302,64 @@ fn decimal_datum(value: &Bound<'_, PyAny>, dec: &DecimalType) -> PyResult }) } +/// Predicate literals are values to compare, not values to store in the column. +/// Java Decimal.compareTo permits a different precision and scale. Keep the +/// strict column conversion above for partition writes and other stored values. +fn predicate_datum(value: &Bound<'_, PyAny>, data_type: &DataType) -> PyResult { + let DataType::Decimal(column) = data_type else { + return py_to_datum(value, data_type); + }; + if let Ok(value) = decimal_datum(value, column) { + return Ok(value); + } + let decimal_class = value.py().import("decimal")?.getattr("Decimal")?; + let value = if value.is_instance_of::() && !value.is_instance_of::() { + decimal_class.call1((value,))? + } else if value.is_instance(&decimal_class)? { + value.clone() + } else { + return decimal_datum(value, column); + }; + let tuple = value.call_method0("as_tuple")?; + let sign: u8 = tuple.getattr("sign")?.extract()?; + let mut exponent: i64 = tuple + .getattr("exponent")? + .extract() + .map_err(|_| PyValueError::new_err("expected a finite decimal predicate literal"))?; + let digits: Vec = tuple.getattr("digits")?.extract()?; + let start = digits.iter().position(|d| *d != 0).unwrap_or(digits.len()); + let mut digits = &digits[start..]; + if digits.is_empty() { + return Ok(Datum::Decimal { + unscaled: 0, + precision: column.precision(), + scale: column.scale(), + }); + } + while digits.last() == Some(&0) { + digits = &digits[..digits.len() - 1]; + exponent = exponent + .checked_add(1) + .ok_or_else(|| PyValueError::new_err("decimal exponent overflow"))?; + } + let scale = exponent.checked_neg().unwrap_or(i64::MAX).max(0); + let precision = (digits.len() as i64) + .saturating_add(exponent.max(0)) + .max(scale); + if precision > 38 || scale > 38 { + return Err(PyValueError::new_err( + "decimal predicate literal exceeds Java DECIMAL precision", + )); + } + let magnitude = digits.iter().fold(0_i128, |v, d| v * 10 + i128::from(*d)) + * 10_i128.pow(exponent.max(0) as u32); + Ok(Datum::Decimal { + unscaled: if sign == 1 { -magnitude } else { magnitude }, + precision: precision as u32, + scale: scale as u32, + }) +} + /// Rescale a finite `decimal.Decimal` to `scale` exactly via `as_tuple()` /// (sign, digits, exponent), avoiding Decimal arithmetic whose context /// precision could silently round. @@ -494,7 +552,7 @@ fn leaf_to_predicate( continue; } out.push(Some( - py_to_datum(&item, &data_type) + predicate_datum(&item, &data_type) .map_err(|e| with_field_context(e, &field, &data_type))?, )); } @@ -1540,6 +1598,79 @@ value = datetime.datetime(2024, 1, 1, tzinfo=FloatingTz())", DataType::Decimal(paimon::spec::DecimalType::new(precision, scale).unwrap()) } + #[test] + fn decimal_predicate_literals_preserve_comparison_value() { + Python::attach(|py| { + for (text, unscaled, precision, scale) in [ + ("1.235", 1235, 4, 3), + ("-1.23500000000000000000000000000000000000000", -1235, 4, 3), + ("123456789.12", 12345678912, 11, 2), + ("1E-38", 1, 38, 38), + ("1E+37", 10_i128.pow(37), 38, 0), + ("0E+100", 0, 10, 2), + ] { + let value = py_decimal(py, text); + assert_eq!( + predicate_datum(&value, &decimal_type(10, 2)).unwrap(), + Datum::Decimal { + unscaled, + precision, + scale + }, + "{text}" + ); + // A comparison literal must not relax stored partition values. + assert!(py_to_datum(&value, &decimal_type(10, 2)).is_err(), "{text}"); + } + let value = 123456789_i64.into_pyobject(py).unwrap(); + assert_eq!( + predicate_datum(&value, &decimal_type(10, 2)).unwrap(), + Datum::Decimal { + unscaled: 123456789, + precision: 9, + scale: 0 + } + ); + for text in ["1E-39", "1E+38"] { + assert!(predicate_datum(&py_decimal(py, text), &decimal_type(10, 2)).is_err()); + } + }); + } + + #[test] + fn decimal_comparison_and_membership_use_literal_scale() { + Python::attach(|py| { + let fields = vec![DataField::new(0, "amount".into(), decimal_type(10, 2))]; + for method in [ + "equal", + "notEqual", + "lessThan", + "lessOrEqual", + "greaterThan", + "greaterOrEqual", + "in", + "notIn", + ] { + let node = leaf_dict(py, method, "amount", &[]); + let literals = PyList::empty(py); + literals.append(py_decimal(py, "1.235")).unwrap(); + node.set_item("literals", literals).unwrap(); + let predicate = dict_to_predicate(&node, &fields, true).unwrap(); + let Predicate::Leaf { literals, .. } = predicate else { + panic!("expected leaf") + }; + assert_eq!( + literals, + vec![Datum::Decimal { + unscaled: 1235, + precision: 4, + scale: 3 + }] + ); + } + }); + } + #[test] fn decimal_field_accepts_exact_scale() { Python::attach(|py| { diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index 2183942af..4e7d199d2 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -1032,17 +1032,38 @@ def test_filter_decimal_literal_prunes_and_reads(): assert t.column("id").to_pylist() == [1] -def test_filter_decimal_rejects_float_and_rounding(): - from decimal import Decimal - +def test_filter_decimal_rejects_float(): with tempfile.TemporaryDirectory() as warehouse: table = _make_temporal_table(warehouse) b = table.new_read_builder() with pytest.raises(ValueError): b.with_filter({"method": "equal", "field": "dec", "literals": [12.34]}) - with pytest.raises(ValueError): - b.with_filter( - {"method": "equal", "field": "dec", "literals": [Decimal("0.005")]}) + + +@pytest.mark.parametrize("method,literal,expected_ids", [ + ("equal", "0.005", []), + ("equal", "12.345", []), + ("equal", "12.340", [1]), + ("lessThan", "12.345", [1]), + ("greaterThan", "56.775", [2]), + ("lessThan", "100000000.00", [1, 2]), + ("greaterThan", "-100000000.00", [1, 2]), +]) +def test_filter_decimal_preserves_literal_precision_and_scale( + method, literal, expected_ids): + from decimal import Decimal + + with tempfile.TemporaryDirectory() as warehouse: + table = _make_temporal_table(warehouse) + builder = table.new_read_builder().with_filter( + {"method": method, "field": "dec", "literals": [Decimal(literal)]}) + filtered_splits = builder.new_scan().plan().splits() + all_splits = table.new_read_builder().new_scan().plan().splits() + # Check both file-stat pruning and row filtering without file pruning. + for splits in (filtered_splits, all_splits): + ids = [row["id"] for batch in builder.new_read().read(splits) + for row in batch.to_pylist()] + assert sorted(ids) == expected_ids def test_filter_unknown_field_raises(): diff --git a/crates/paimon/src/table/key_normalization.rs b/crates/paimon/src/table/key_normalization.rs index 50b78dc84..16fcef2f7 100644 --- a/crates/paimon/src/table/key_normalization.rs +++ b/crates/paimon/src/table/key_normalization.rs @@ -23,10 +23,13 @@ use arrow_schema::DataType; use super::update_input::{cast_update_value, decode_dictionary, CastMode}; use super::write_batch_normalize::normalize_write_array; +mod numeric; + /// 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. +/// indices are batch-local; numeric widths, decimal scales and temporal units +/// may differ only when conversion is exact. Binary keys compare their bytes. +/// Callers validate logical type families and decide whether NULL keys match. +/// Payload write rules are separate. pub(super) fn normalize_match_key( array: &ArrayRef, expected: &DataType, @@ -36,19 +39,54 @@ pub(super) fn normalize_match_key( return Ok(new_null_array(expected, decoded.len())); } let source = decoded.data_type(); - 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(_) - ) + if source == expected { + return Ok(decoded); + } + if source.is_numeric() && expected.is_numeric() { + return numeric::normalize(&decoded, expected); + } + let binary = |t: &DataType| { + matches!( + t, + DataType::Binary | DataType::LargeBinary | DataType::FixedSizeBinary(_) ) - || matches!((source, expected), (DataType::Timestamp(_, from), DataType::Timestamp(_, to)) + }; + if binary(source) && binary(expected) { + if source == &DataType::LargeBinary && expected == &DataType::Binary { + return super::write_batch_normalize::normalize_large_binary( + decoded.as_any().downcast_ref().unwrap(), + ); + } + // Key identity is its bytes, independently of Arrow offset width. + // Ordinary append validation still sees the original source array. + let values = if matches!(source, DataType::FixedSizeBinary(_)) { + normalize_write_array(&decoded, &DataType::Binary)? + } else { + decoded.clone() + }; + return arrow_cast::cast_with_options( + values.as_ref(), + expected, + &arrow_cast::CastOptions { + safe: false, + ..Default::default() + }, + ) + .map_err(|error| crate::Error::DataInvalid { + message: error.to_string(), + source: None, + }); + } + let exact_family = 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) @@ -62,7 +100,8 @@ mod tests { use super::*; use arrow_array::types::Int8Type; use arrow_array::{ - Array, Decimal128Array, DictionaryArray, Int32Array, Int64Array, Int8Array, StringArray, + Array, BinaryArray, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, Float32Array, + Float64Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray, }; use arrow_schema::TimeUnit; @@ -102,6 +141,92 @@ mod tests { assert!(normalize_match_key(&text, &DataType::Int32).is_err()); } + #[test] + fn floating_keys_change_width_only_when_identity_is_preserved() { + let input: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(99.0), + Some(1.5), + None, + Some(-2.25), + ])); + let input = input.slice(1, 3); + let narrow = normalize_match_key(&input, &DataType::Float32).unwrap(); + assert_eq!( + narrow.as_any().downcast_ref::().unwrap(), + &Float32Array::from(vec![Some(1.5), None, Some(-2.25)]) + ); + assert_eq!( + normalize_match_key(&narrow, &DataType::Float64) + .unwrap() + .to_data(), + input.to_data() + ); + for value in [1.00000001, f64::MAX, f64::MIN_POSITIVE] { + let lossy: ArrayRef = Arc::new(Float64Array::from(vec![value])); + assert!(normalize_match_key(&lossy, &DataType::Float32).is_err()); + } + let dictionary: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), None, Some(1), Some(0)]), + Arc::new(Float32Array::from(vec![Some(1.5), None])), + ) + .unwrap(), + ); + let decoded = normalize_match_key(&dictionary, &DataType::Float64).unwrap(); + assert_eq!( + decoded.as_any().downcast_ref::().unwrap(), + &Float64Array::from(vec![Some(1.5), None, None, Some(1.5)]) + ); + } + + #[test] + fn binary_keys_compare_bytes_across_offsets_and_dictionaries() { + let expected = + BinaryArray::from(vec![Some(b"\xff\0".as_slice()), None, Some(b"".as_slice())]); + let large: ArrayRef = Arc::new(LargeBinaryArray::from(vec![ + Some(b"padding".as_slice()), + Some(b"\xff\0".as_slice()), + None, + Some(b"".as_slice()), + ])); + let output = normalize_match_key(&large.slice(1, 3), &DataType::Binary).unwrap(); + assert_eq!( + output.as_any().downcast_ref::().unwrap(), + &expected + ); + let restored = normalize_match_key(&output, &DataType::LargeBinary).unwrap(); + assert_eq!(restored.to_data(), large.slice(1, 3).to_data()); + + let fixed: ArrayRef = Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(b"ab".as_slice()), None, Some(b"cd".as_slice())].into_iter(), + 2, + ) + .unwrap(), + ); + for target in [DataType::Binary, DataType::LargeBinary] { + let normalized = normalize_match_key(&fixed.slice(1, 2), &target).unwrap(); + assert_eq!( + normalized.to_data(), + arrow_cast::cast(&fixed.slice(1, 2), &target) + .unwrap() + .to_data() + ); + } + let dictionary: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(0), None, Some(1)]), + Arc::new(LargeBinaryArray::from(vec![Some(b"ab".as_slice()), None])), + ) + .unwrap(), + ); + let output = normalize_match_key(&dictionary, &DataType::Binary).unwrap(); + assert_eq!( + output.as_any().downcast_ref::().unwrap(), + &BinaryArray::from(vec![Some(b"ab".as_slice()), None, None]) + ); + } + #[test] fn decimal_keys_cannot_round_into_existing_keys() { let array = |value| { diff --git a/crates/paimon/src/table/key_normalization/numeric.rs b/crates/paimon/src/table/key_normalization/numeric.rs new file mode 100644 index 000000000..5b49bd2fb --- /dev/null +++ b/crates/paimon/src/table/key_normalization/numeric.rs @@ -0,0 +1,329 @@ +// 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. + +//! Exact numeric equality for external Arrow match keys. Assignment casts may +//! round; a matching key must retain its mathematical value across families. + +use std::sync::Arc; + +use arrow_array::{cast::AsArray, Array, ArrayRef, Decimal128Array, Float32Array, Float64Array}; +use arrow_buffer::i256; +use arrow_schema::DataType; + +use super::super::update_input::{cast_update_value, CastMode}; + +fn invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +pub(super) fn normalize(array: &ArrayRef, target: &DataType) -> crate::Result { + let source = array.data_type(); + if source.is_floating() { + if let DataType::Decimal128(precision, scale) = target { + let values = arrow_cast::cast(array, &DataType::Float64) + .map_err(|error| invalid(error.to_string()))?; + let values = values + .as_primitive::() + .iter() + .map(|value| { + value + .map(|value| { + let coefficient = float_coefficient(value, *scale) + .and_then(i256::to_i128) + .filter(|value| { + value.unsigned_abs() < 10_u128.pow(u32::from(*precision)) + }) + .ok_or_else(|| { + invalid("Floating match key is not an exact DECIMAL value") + })?; + Ok(coefficient) + }) + .transpose() + }) + .collect::>>()?; + return Decimal128Array::from(values) + .with_precision_and_scale(*precision, *scale) + .map(|array| Arc::new(array) as ArrayRef) + .map_err(|error| invalid(error.to_string())); + } + } + if source.is_decimal() { + if target.is_floating() { + return decimal_to_float(array, target); + } + if target.is_integer() { + // Canonicalize first to avoid scale arithmetic overflow in generic + // Arrow integer casts, including external Decimal256 zero values. + let integral = + cast_update_value(array, &DataType::Decimal128(38, 0), CastMode::Assignment)?; + return cast_update_value(&integral, target, CastMode::Assignment); + } + } + if source.is_floating() || target.is_floating() { + // Unlike assignment's consecutive-integer range, identity only needs + // an exact round trip (e.g. 2^53 + 2 is a valid DOUBLE key). + let options = arrow_cast::CastOptions { + safe: false, + ..Default::default() + }; + let converted = arrow_cast::cast_with_options(array.as_ref(), target, &options) + .map_err(|error| invalid(error.to_string()))?; + let restored = arrow_cast::cast_with_options(converted.as_ref(), source, &options) + .map_err(|error| invalid(error.to_string()))?; + if restored.to_data() != array.to_data() { + return Err(invalid( + "Numeric match key conversion would change key identity", + )); + } + return Ok(converted); + } + cast_update_value(array, target, CastMode::Assignment) +} + +fn decimal_to_float(array: &ArrayRef, target: &DataType) -> crate::Result { + let scale = match array.data_type() { + DataType::Decimal32(_, s) + | DataType::Decimal64(_, s) + | DataType::Decimal128(_, s) + | DataType::Decimal256(_, s) => *s, + _ => unreachable!("validated DECIMAL key"), + }; + let values = (0..array.len()) + .map(|row| { + if array.is_null(row) { + return Ok(None); + } + let coefficient = match array.data_type() { + DataType::Decimal32(..) => i256::from_i128(i128::from( + array + .as_primitive::() + .value(row), + )), + DataType::Decimal64(..) => i256::from_i128(i128::from( + array + .as_primitive::() + .value(row), + )), + DataType::Decimal128(..) => i256::from_i128( + array + .as_primitive::() + .value(row), + ), + DataType::Decimal256(..) => array + .as_primitive::() + .value(row), + _ => unreachable!("validated DECIMAL key"), + }; + let text = format!("{coefficient}e{}", -i16::from(scale)); + let value = if target == &DataType::Float32 { + f64::from( + text.parse::() + .map_err(|error| invalid(error.to_string()))?, + ) + } else { + text.parse::() + .map_err(|error| invalid(error.to_string()))? + }; + if float_coefficient(value, scale) != Some(coefficient) { + return Err(invalid("DECIMAL match key is not an exact floating value")); + } + Ok(Some(value)) + }) + .collect::>>()?; + Ok(if target == &DataType::Float32 { + Arc::new(Float32Array::from_iter( + values.into_iter().map(|value| value.map(|v| v as f32)), + )) as ArrayRef + } else { + Arc::new(Float64Array::from(values)) as ArrayRef + }) +} + +/// Return the exact decimal coefficient of a binary float at the given scale. +/// Check divisibility before multiplying, so a fractional key cannot round +/// into an existing row. All work is bounded by the Arrow coefficient width. +fn float_coefficient(value: f64, scale: i8) -> Option { + if !value.is_finite() { + return None; + } + if value == 0.0 { + return Some(i256::ZERO); + } + let bits = value.abs().to_bits(); + let raw_exponent = ((bits >> 52) & 0x7ff) as i32; + let mut mantissa = bits & ((1_u64 << 52) - 1); + let mut exponent = if raw_exponent == 0 { + -1074 + } else { + mantissa |= 1_u64 << 52; + raw_exponent - 1023 - 52 + }; + let zeros = mantissa.trailing_zeros(); + mantissa >>= zeros; + exponent += zeros as i32 + i32::from(scale); + if exponent < 0 { + return None; + } + let five = if scale < 0 { + for _ in 0..-i16::from(scale) { + if !mantissa.is_multiple_of(5) { + return None; + } + mantissa /= 5; + } + i256::ONE + } else { + i256::from_i128(5).checked_pow(scale as u32)? + }; + let coefficient = i256::from_i128(i128::from(mantissa)) + .checked_mul(five)? + .checked_mul(i256::from_i128(2).checked_pow(exponent as u32)?)?; + if value < 0.0 { + coefficient.checked_neg() + } else { + Some(coefficient) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ + Decimal256Array, Decimal32Array, Decimal64Array, Int32Array, Int64Array, UInt64Array, + }; + + #[test] + fn integer_float_keys_require_identity_without_assignment_range_limits() { + let integer: ArrayRef = Arc::new(Int64Array::from(vec![Some((1_i64 << 53) + 2), None])); + let converted = normalize(&integer, &DataType::Float64).unwrap(); + assert_eq!( + normalize(&converted, &DataType::Int64).unwrap().to_data(), + integer.to_data() + ); + let fractional: ArrayRef = Arc::new(Float64Array::from(vec![1.00000001])); + assert!(normalize(&fractional, &DataType::Int32).is_err()); + let rounded: ArrayRef = Arc::new(Int64Array::from(vec![(1_i64 << 53) + 1])); + assert!(normalize(&rounded, &DataType::Float64).is_err()); + let overflow: ArrayRef = Arc::new(UInt64Array::from(vec![u64::MAX])); + assert!(normalize(&overflow, &DataType::Float64).is_err()); + let overflow: ArrayRef = Arc::new(Float64Array::from(vec![2_f64.powi(63)])); + assert!(normalize(&overflow, &DataType::Int64).is_err()); + } + + #[test] + fn decimal_integer_keys_share_the_exact_rescaler() { + let decimal: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(12500), None]) + .with_precision_and_scale(15, 4) + .unwrap(), + ); + assert!(normalize(&decimal, &DataType::Int32).is_err()); + let integer: ArrayRef = Arc::new(Int32Array::from(vec![Some(-1), None, Some(2)])); + let converted = normalize(&integer, &DataType::Decimal128(10, 2)).unwrap(); + assert_eq!( + normalize(&converted, &DataType::Int32).unwrap().to_data(), + integer.to_data() + ); + let zero: ArrayRef = Arc::new( + Decimal256Array::from(vec![Some(i256::ZERO), None]) + .with_precision_and_scale(76, -128) + .unwrap(), + ); + let actual = normalize(&zero, &DataType::Int64).unwrap(); + assert_eq!( + actual.to_data(), + Int64Array::from(vec![Some(0), None]).to_data() + ); + } + + #[test] + fn decimal_float_keys_do_not_use_rounded_decimal_equality() { + let floating: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.5), None, Some(-2.25)])); + let decimal = normalize(&floating, &DataType::Decimal128(10, 2)).unwrap(); + assert_eq!( + decimal.to_data(), + Decimal128Array::from(vec![Some(150), None, Some(-225)]) + .with_precision_and_scale(10, 2) + .unwrap() + .to_data() + ); + assert_eq!( + normalize(&decimal, &DataType::Float64).unwrap().to_data(), + floating.to_data() + ); + for value in [0.1, 1.235, 1.0 + f64::EPSILON, f64::MAX, f64::MIN_POSITIVE] { + let input: ArrayRef = Arc::new(Float64Array::from(vec![value])); + assert!(normalize(&input, &DataType::Decimal128(10, 2)).is_err()); + } + let rounded: ArrayRef = Arc::new( + Decimal128Array::from(vec![10]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + assert!(normalize(&rounded, &DataType::Float64).is_err()); + assert!(normalize(&rounded, &DataType::Float32).is_err()); + let small: ArrayRef = Arc::new( + Decimal32Array::from(vec![Some(150), None]) + .with_precision_and_scale(9, 2) + .unwrap(), + ); + let medium: ArrayRef = Arc::new( + Decimal64Array::from(vec![Some(150), None]) + .with_precision_and_scale(18, 2) + .unwrap(), + ); + for input in [small, medium] { + assert_eq!( + normalize(&input, &DataType::Float64).unwrap().to_data(), + Float64Array::from(vec![Some(1.5), None]).to_data() + ); + } + } + + #[test] + fn exact_binary_values_support_wide_and_negative_decimal_scales() { + assert_eq!(float_coefficient(1e20, -20), Some(i256::ONE)); + assert_eq!(float_coefficient(-1e20, -20), Some(-i256::ONE)); + assert_eq!(float_coefficient(1e20, -21), None); + assert_eq!(float_coefficient(0.0, -128), Some(i256::ZERO)); + assert_eq!(float_coefficient(0.0, 127), Some(i256::ZERO)); + // 2^-70 needs 70 decimal places, beyond Decimal128's coefficient width. + let coefficient = i256::from_i128(5).checked_pow(70).unwrap(); + let input: ArrayRef = Arc::new( + Decimal256Array::from(vec![Some(coefficient), None, Some(-coefficient)]) + .with_precision_and_scale(76, 70) + .unwrap(), + ); + let actual = normalize(&input, &DataType::Float64).unwrap(); + assert_eq!( + actual.to_data(), + Float64Array::from(vec![Some(2_f64.powi(-70)), None, Some(-2_f64.powi(-70))]).to_data() + ); + let large: ArrayRef = Arc::new( + Decimal256Array::from(vec![Some(i256::ONE), None]) + .with_precision_and_scale(76, -20) + .unwrap(), + ); + assert_eq!( + normalize(&large, &DataType::Float64).unwrap().to_data(), + Float64Array::from(vec![Some(1e20), None]).to_data() + ); + } +} diff --git a/crates/paimon/src/table/table_update.rs b/crates/paimon/src/table/table_update.rs index fdaa256ef..490010808 100644 --- a/crates/paimon/src/table/table_update.rs +++ b/crates/paimon/src/table/table_update.rs @@ -245,6 +245,9 @@ impl TableUpdate { /// 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. + /// Numeric keys may use different Arrow types only when conversion keeps + /// their exact mathematical value. This matching conversion does not alter + /// the original key columns used as update or append payloads. /// 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/write_batch_normalize.rs b/crates/paimon/src/table/write_batch_normalize.rs index 44b689b0b..d732215f2 100644 --- a/crates/paimon/src/table/write_batch_normalize.rs +++ b/crates/paimon/src/table/write_batch_normalize.rs @@ -26,9 +26,9 @@ use crate::{Error, Result}; use arrow_array::{ Array, ArrayRef, BinaryArray, Decimal128Array, FixedSizeBinaryArray, FixedSizeListArray, - LargeStringArray, ListArray, MapArray, StringArray, StructArray, + LargeBinaryArray, LargeStringArray, ListArray, MapArray, StringArray, StructArray, }; -use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow_buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field}; use std::sync::Arc; @@ -435,28 +435,52 @@ fn normalize_visible_array( /// in its backing buffer, so casting the original absolute offsets is wrong. /// The value buffer remains shared and only the offset buffer is rebuilt. fn normalize_large_string(array: &LargeStringArray) -> Result { + let (offsets, values) = narrow_byte_range(array.value_offsets(), array.values())?; + StringArray::try_new(offsets, values, array.nulls().cloned()) + .map(|array| Arc::new(array) as ArrayRef) + .map_err(normalization_error) +} + +/// Offset adaptation for matching keys; ordinary append type validation remains +/// separate so this does not broaden accepted write schemas. +pub(super) fn normalize_large_binary(array: &LargeBinaryArray) -> Result { let source = array.value_offsets(); + if array.null_count() != 0 && source[source.len() - 1] - source[0] > i64::from(i32::MAX) { + // NULL slots can retain arbitrarily large, unobservable byte ranges. + // Compact only when they prevent narrowing; ordinary arrays continue + // to share the visible value buffer below. + let visible_len = array + .iter() + .flatten() + .try_fold(0_usize, |len, value| len.checked_add(value.len())); + if visible_len.is_none_or(|len| len > i32::MAX as usize) { + return Err(Error::DataInvalid { + message: "Arrow byte range exceeds 32-bit offsets".into(), + source: None, + }); + } + return Ok(Arc::new(BinaryArray::from_iter(array.iter()))); + } + let (offsets, values) = narrow_byte_range(array.value_offsets(), array.values())?; + BinaryArray::try_new(offsets, values, array.nulls().cloned()) + .map(|array| Arc::new(array) as ArrayRef) + .map_err(normalization_error) +} + +fn narrow_byte_range(source: &[i64], values: &Buffer) -> Result<(OffsetBuffer, Buffer)> { let start = source[0]; let end = source[source.len() - 1]; let offsets = source .iter() .map(|offset| { i32::try_from(offset - start).map_err(|_| Error::DataInvalid { - message: "Arrow STRING byte range exceeds 32-bit offsets".into(), + message: "Arrow byte range exceeds 32-bit offsets".into(), source: None, }) }) .collect::>>()?; - let values = array - .values() - .slice_with_length(start as usize, (end - start) as usize); - StringArray::try_new( - OffsetBuffer::new(offsets.into()), - values, - array.nulls().cloned(), - ) - .map(|array| Arc::new(array) as ArrayRef) - .map_err(normalization_error) + let values = values.slice_with_length(start as usize, (end - start) as usize); + Ok((OffsetBuffer::new(offsets.into()), values)) } /// Keep only children referenced by visible parents. Ordinary inputs and @@ -604,6 +628,56 @@ mod tests { }); let empty = normalize_write_array(&input.slice(3, 0), &DataType::Utf8).unwrap(); assert!(empty.is_empty()); + let text = input.as_any().downcast_ref::().unwrap(); + let binary: ArrayRef = Arc::new(LargeBinaryArray::new( + text.offsets().clone(), + text.values().clone(), + text.nulls().cloned(), + )); + let normalized = + crate::table::key_normalization::normalize_match_key(&binary, &DataType::Binary) + .unwrap(); + assert_eq!( + normalized.to_data(), + BinaryArray::from(vec![Some("abc".as_bytes()), None, Some("中文".as_bytes())]) + .to_data() + ); + assert!(normalize_write_array(&binary, &DataType::Binary).is_err()); + for mixed in [false, true] { + let offsets = if mixed { + vec![0, start, start + 3] + } else { + vec![0, start] + }; + let input: ArrayRef = Arc::new(LargeBinaryArray::new( + OffsetBuffer::new(offsets.into()), + text.values().clone(), + Some(NullBuffer::from(if mixed { + vec![false, true] + } else { + vec![false] + })), + )); + let output = + crate::table::key_normalization::normalize_match_key(&input, &DataType::Binary) + .unwrap(); + let expected = if mixed { + vec![None, Some(b"abc".as_slice())] + } else { + vec![None] + }; + assert_eq!(output.to_data(), BinaryArray::from(expected).to_data()); + } + let visible_overflow: ArrayRef = Arc::new(LargeBinaryArray::new( + OffsetBuffer::new(vec![0, start + 3, start + 3].into()), + text.values().clone(), + Some(NullBuffer::from(vec![true, false])), + )); + assert!(crate::table::key_normalization::normalize_match_key( + &visible_overflow, + &DataType::Binary, + ) + .is_err()); // A visible range larger than i32::MAX still cannot be represented. // SAFETY: the mapped bytes are zeroes followed by a valid UTF-8 tail; // offsets and NULL bitmap are valid. Skip scanning the 2 GiB prefix.