[core] Fix VARIANT timestamp projection and scalar conversions

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::<Vec<_>>();
@@ -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<ArrayRef> {
+        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::<StructArray>()
+            .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::<Float64Array>().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::<StructArray>().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);
+        }
+    }
 }

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<bool> {
 
 fn cast_variant_to_i64(variant: VariantRef<'_>) -> Option<i64> {
     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<i64> {
 
 fn cast_variant_to_f64(variant: VariantRef<'_>) -> Option<f64> {
     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<String> {
 
 fn cast_variant_to_decimal(variant: VariantRef<'_>, precision: u8, scale: i8) -> Option<i128> {
     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()?;