From a9feac586d6e77fe510835715ae7d67568a97f72 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 17 Oct 2026 09:19:34 -0700 Subject: [PATCH] [core] Fix VARIANT timestamp projection and scalar conversions --- crates/paimon/src/arrow/shredding/variant.rs | 171 ++++++++++++++++++- crates/paimon/src/variant.rs | 30 +++- 2 files changed, 193 insertions(+), 8 deletions(-) diff --git a/crates/paimon/src/arrow/shredding/variant.rs b/crates/paimon/src/arrow/shredding/variant.rs index 519510f47..c0e5201df 100644 --- a/crates/paimon/src/arrow/shredding/variant.rs +++ b/crates/paimon/src/arrow/shredding/variant.rs @@ -1546,7 +1546,12 @@ fn timestamp_array( let values = values .iter() .map(|value| match value { - Some(ShreddedValue::Timestamp(v)) => Some(*v), + Some(ShreddedValue::Timestamp(v)) => match precision { + 0 => Some(v.div_euclid(1_000_000)), + 1..=3 => Some(v.div_euclid(1_000)), + 4..=6 => Some(*v), + _ => v.checked_mul(1_000), + }, _ => None, }) .collect::>(); @@ -2757,4 +2762,168 @@ mod tests { .data_type() ); } + fn parquet_roundtrip(array: ArrayRef) -> ArrayRef { + use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter}; + let schema = Arc::new(arrow_schema::Schema::new(vec![ArrowField::new( + "v", + array.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + let mut output = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut output, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(output)) + .unwrap() + .build() + .unwrap() + .next() + .unwrap() + .unwrap() + .column(0) + .clone() + } + + fn extract_root(array: &dyn Array, data_type: DataType, strict: bool) -> Result { + let field = DataField::new(0, "result".into(), data_type).with_description(Some( + crate::spec::build_variant_metadata("$", strict, "UTC").unwrap(), + )); + let result = assemble_variant_extraction_array(array, &RowType::new(vec![field]))?; + Ok(result + .as_any() + .downcast_ref::() + .unwrap() + .column(0) + .clone()) + } + + #[test] + fn variant_timestamp_projection_units_and_overflow() { + use crate::spec::{LocalZonedTimestampType, TimestampType}; + let micros = [1_234_567_890_123_i64, -1, -1_234_567, 0, i64::MAX, i64::MIN]; + for zoned in [false, true] { + let dtype = |precision| { + if zoned { + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(precision).unwrap()) + } else { + DataType::Timestamp(TimestampType::new(precision).unwrap()) + } + }; + let values = micros + .iter() + .map(|value| { + let mut raw = vec![if zoned { 12 << 2 } else { 13 << 2 }]; + raw.extend_from_slice(&value.to_le_bytes()); + Some(GenericVariant::from_parts(raw, vec![1, 0, 0]).unwrap()) + }) + .chain([Some(GenericVariant::parse_json("null").unwrap()), None]) + .collect(); + let plain = variant_array(values).unwrap(); + let physical = variant_shredding_type(&dtype(3)).unwrap(); + let shredded = shred_variant_array(plain.as_ref(), &physical).unwrap(); + for input in [plain, shredded] { + let input = parquet_roundtrip(input); + for precision in [0, 3, 6, 9] { + let result = extract_root(input.as_ref(), dtype(precision), false).unwrap(); + assert_eq!( + result.data_type(), + &paimon_type_to_arrow(&dtype(precision)).unwrap() + ); + for (row, value) in micros.iter().enumerate() { + let expected = match precision { + 0 => Some(value.div_euclid(1_000_000)), + 3 => Some(value.div_euclid(1_000)), + 6 => Some(*value), + _ => value.checked_mul(1_000), + }; + match expected { + Some(value) => { + assert!(!result.is_null(row)); + assert_eq!( + timestamp_value_at(result.as_ref(), row).unwrap(), + value + ); + } + None => assert!(result.is_null(row)), + } + } + assert!(result.is_null(micros.len())); + assert!(result.is_null(micros.len() + 1)); + let strict = extract_root(input.as_ref(), dtype(precision), true); + assert_eq!(strict.is_err(), precision == 9); + let valid = input.slice(0, 4); + assert!(extract_root(valid.as_ref(), dtype(precision), true).is_ok()); + } + } + } + } + + #[test] + fn variant_boolean_numeric_projection() { + use crate::spec::{ + BigIntType, BooleanType, DecimalType, DoubleType, SmallIntType, TinyIntType, + }; + let values = ["true", "false", "null"] + .iter() + .map(|value| Some(GenericVariant::parse_json(value).unwrap())) + .chain([None]) + .collect(); + let plain = variant_array(values).unwrap(); + let physical = variant_shredding_type(&DataType::Boolean(BooleanType::new())).unwrap(); + let shredded = shred_variant_array(plain.as_ref(), &physical).unwrap(); + for input in [plain, shredded] { + let input = parquet_roundtrip(input); + for dtype in [ + DataType::TinyInt(TinyIntType::new()), + DataType::SmallInt(SmallIntType::new()), + DataType::Int(IntType::new()), + DataType::BigInt(BigIntType::new()), + DataType::Float(FloatType::new()), + DataType::Double(DoubleType::new()), + DataType::Decimal(DecimalType::new(10, 2).unwrap()), + ] { + for strict in [false, true] { + let output = extract_root(input.as_ref(), dtype.clone(), strict).unwrap(); + assert_eq!(output.null_count(), 2); + assert!(output.is_null(2)); + assert!(output.is_null(3)); + let numeric = + arrow_cast::cast(output.as_ref(), &ArrowDataType::Float64).unwrap(); + let numeric = numeric.as_any().downcast_ref::().unwrap(); + assert_eq!(numeric.value(0), 1.0); + assert_eq!(numeric.value(1), 0.0); + } + } + } + } + + #[test] + fn variant_decimal_overflow_preserves_residual() { + use crate::spec::DecimalType; + let values = ["10", "-10", "9223372036854775807", "-9223372036854775808"]; + let input = variant_array( + values + .iter() + .map(|v| Some(GenericVariant::parse_json(v).unwrap())) + .collect(), + ) + .unwrap(); + let physical = + variant_shredding_type(&DataType::Decimal(DecimalType::new(38, 38).unwrap())).unwrap(); + let output = shred_variant_array(input.as_ref(), &physical).unwrap(); + let output = parquet_roundtrip(output); + let storage = output.as_any().downcast_ref::().unwrap(); + assert_eq!( + storage.column_by_name("typed_value").unwrap().null_count(), + values.len() + ); + assert_eq!(storage.column_by_name("value").unwrap().null_count(), 0); + let restored = assemble_shredded_variant_array(output.as_ref()).unwrap(); + for row in 0..values.len() { + let expected = variant_from_array(input.as_ref(), row).unwrap().unwrap(); + let actual = variant_from_array(restored.as_ref(), row).unwrap().unwrap(); + assert_eq!(actual, expected); + } + } } diff --git a/crates/paimon/src/variant.rs b/crates/paimon/src/variant.rs index d9616315b..20aa42b21 100644 --- a/crates/paimon/src/variant.rs +++ b/crates/paimon/src/variant.rs @@ -2774,7 +2774,20 @@ pub(crate) fn cast_variant_to_shredded_value( message: format!("Unsupported Variant extraction target type: {data_type:?}"), }); }; - match cast_variant_to_extraction_value(variant, &scalar) { + let value = cast_variant_to_extraction_value(variant, &scalar).filter(|value| { + let precision = match data_type { + DataType::Timestamp(ts) => ts.precision(), + DataType::LocalZonedTimestamp(ts) => ts.precision(), + _ => return true, + }; + match value { + ShreddedValue::Timestamp(micros) if precision > 6 => { + micros.checked_mul(1_000).is_some() + } + _ => true, + } + }); + match value { Some(value) => Ok(Some(value)), None if fail_on_error => Err(Error::DataInvalid { message: format!("Cannot cast Variant value to {data_type:?}"), @@ -2827,6 +2840,7 @@ fn cast_variant_to_bool(variant: VariantRef<'_>) -> Option { fn cast_variant_to_i64(variant: VariantRef<'_>) -> Option { match variant.kind().ok()? { + VariantKind::Boolean => Some(i64::from(variant.get_boolean().ok()?)), VariantKind::Long | VariantKind::Date | VariantKind::Timestamp @@ -2873,6 +2887,7 @@ fn f64_to_i64_truncating(value: f64) -> Option { fn cast_variant_to_f64(variant: VariantRef<'_>) -> Option { match variant.kind().ok()? { + VariantKind::Boolean => Some(u8::from(variant.get_boolean().ok()?) as f64), VariantKind::Long | VariantKind::Date | VariantKind::Timestamp @@ -2931,6 +2946,9 @@ fn cast_variant_to_string(variant: VariantRef<'_>) -> Option { fn cast_variant_to_decimal(variant: VariantRef<'_>, precision: u8, scale: i8) -> Option { let unscaled = match variant.kind().ok()? { + VariantKind::Boolean => { + rescale_decimal_exact(i128::from(variant.get_boolean().ok()?), 0, scale)? + } VariantKind::Long | VariantKind::Date | VariantKind::Timestamp @@ -3153,12 +3171,10 @@ fn try_typed_shred( Some(ShreddedValue::Int64(variant.get_long()?)) } (VariantKind::Long, VariantScalarSchema::Decimal { precision, scale }) => { - let unscaled = rescale_decimal(variant.get_long()? as i128, 0, *scale)?; - if decimal_precision(unscaled) <= *precision { - Some(ShreddedValue::Decimal128(unscaled)) - } else { - None - } + rescale_decimal_exact(variant.get_long()? as i128, 0, *scale).and_then(|unscaled| { + (decimal_precision(unscaled) <= *precision) + .then_some(ShreddedValue::Decimal128(unscaled)) + }) } (VariantKind::Decimal, VariantScalarSchema::Decimal { precision, scale }) => { let decimal = variant.get_decimal()?;