[Variant] Support typed access for timestamp(micro&nano)

Author: klion26

parquet-variant-compute/src/variant_array.rs

@@ -17,16 +17,18 @@
 
 //! [`VariantArray`] implementation
 
-use crate::type_conversion::primitive_conversion_single_value;
+use crate::type_conversion::{generic_conversion_single_value, primitive_conversion_single_value};
 use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray, StructArray};
 use arrow::buffer::NullBuffer;
 use arrow::compute::cast;
 use arrow::datatypes::{
     Date32Type, Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
-    UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    TimestampMicrosecondType, TimestampNanosecondType, UInt16Type, UInt32Type, UInt64Type,
+    UInt8Type,
 };
 use arrow_schema::extension::ExtensionType;
-use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields};
+use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, TimeUnit};
+use chrono::DateTime;
 use parquet_variant::Uuid;
 use parquet_variant::Variant;
 use std::sync::Arc;
@@ -864,6 +866,51 @@ fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Variant<'_, '
         DataType::Float64 => {
             primitive_conversion_single_value!(Float64Type, typed_value, index)
         }
+        DataType::Timestamp(timeunit, tz) => {
+            match (timeunit, tz) {
+                (TimeUnit::Microsecond, Some(_)) => {
+                    generic_conversion_single_value!(
+                        TimestampMicrosecondType,
+                        as_primitive,
+                        |v| DateTime::from_timestamp_micros(v).unwrap(),
+                        typed_value,
+                        index
+                    )
+                }
+                (TimeUnit::Microsecond, None) => {
+                    generic_conversion_single_value!(
+                        TimestampMicrosecondType,
+                        as_primitive,
+                        |v| DateTime::from_timestamp_micros(v).unwrap().naive_utc(),
+                        typed_value,
+                        index
+                    )
+                }
+                (TimeUnit::Nanosecond, Some(_)) => {
+                    generic_conversion_single_value!(
+                        TimestampNanosecondType,
+                        as_primitive,
+                        DateTime::from_timestamp_nanos,
+                        typed_value,
+                        index
+                    )
+                }
+                (TimeUnit::Nanosecond, None) => {
+                    generic_conversion_single_value!(
+                        TimestampNanosecondType,
+                        as_primitive,
+                        |v| DateTime::from_timestamp_nanos(v).naive_utc(),
+                        typed_value,
+                        index
+                    )
+                }
+                // Variant timestamp only support time unit with microsecond or nanosecond precision
+                _ => panic!(
+                    "Variant only support timestamp with microsecond or nanosecond precision"
+                ),
+            }
+        }
+
         // todo other types here (note this is very similar to cast_to_variant.rs)
         // so it would be great to figure out how to share this code
         _ => {

parquet-variant-compute/src/variant_get.rs

@@ -34,7 +34,7 @@ pub(crate) enum ShreddedPathStep {
     /// Path step succeeded, return the new shredding state
     Success(ShreddingState),
     /// The path element is not present in the `typed_value` column and there is no `value` column,
-    /// so we we know it does not exist. It, and all paths under it, are all-NULL.
+    /// so we know it does not exist. It, and all paths under it, are all-NULL.
     Missing,
     /// The path element is not present in the `typed_value` column and must be retrieved from the `value`
     /// column instead. The caller should be prepared to handle any value, including the requested
@@ -294,24 +294,18 @@ impl<'a> GetOptions<'a> {
 
 #[cfg(test)]
 mod test {
-    use std::sync::Arc;
-
-    use arrow::array::{
-        Array, ArrayRef, AsArray, BinaryViewArray, Date32Array,
-        Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
-        StringArray, StructArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
-    };
+    use super::{variant_get, GetOptions};
+    use crate::json_to_variant;
+    use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+    use crate::VariantArray;
+    use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray, Date32Array, Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, StringArray, StructArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array};
     use arrow::buffer::NullBuffer;
     use arrow::compute::CastOptions;
     use arrow::datatypes::DataType::{Int16, Int32, Int64, UInt16, UInt32, UInt64, UInt8};
     use arrow_schema::{DataType, Field, FieldRef, Fields};
+    use chrono::DateTime;
     use parquet_variant::{Variant, VariantPath, EMPTY_VARIANT_METADATA_BYTES};
-
-    use crate::json_to_variant;
-    use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
-    use crate::VariantArray;
-
-    use super::{variant_get, GetOptions};
+    use std::sync::Arc;
 
     fn single_variant_get_test(input_json: &str, path: VariantPath, expected_json: &str) {
         // Create input array from JSON string
@@ -923,6 +917,156 @@ mod test {
         f64
     );
 
+    macro_rules! assert_variant_get_as_variant_array_with_default_option {
+        ($variant_array: expr, $array_expected: expr) => {{
+            let options = GetOptions::new();
+            let array = $variant_array;
+            let result = variant_get(&array, options).unwrap();
+
+            // expect the result is a VariantArray
+            let result = VariantArray::try_new(&result).unwrap();
+
+            assert_eq!(result.len(), $array_expected.len());
+
+            for (idx, item) in $array_expected.into_iter().enumerate() {
+                match item {
+                    Some(item) => assert_eq!(result.value(idx), item),
+                    None => assert!(result.is_null(idx)),
+                }
+            }
+        }};
+    }
+
+    partially_shredded_variant_array_gen!(
+        partially_shredded_timestamp_micro_ntz_variant_array,
+        || {
+            arrow::array::TimestampMicrosecondArray::from(vec![
+                Some(-456000),
+                None,
+                None,
+                Some(1758602096000000),
+            ])
+        }
+    );
+
+    #[test]
+    fn get_variant_partial_shredded_timestamp_micro_ntz_as_variant() {
+        let array = partially_shredded_timestamp_micro_ntz_variant_array();
+        assert_variant_get_as_variant_array_with_default_option!(
+            array,
+            vec![
+                Some(Variant::from(
+                    DateTime::from_timestamp_micros(-456000i64)
+                        .unwrap()
+                        .naive_utc(),
+                )),
+                None,
+                Some(Variant::from("n/a")),
+                Some(Variant::from(
+                    DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
+                        .unwrap()
+                        .naive_utc(),
+                )),
+            ]
+        )
+    }
+
+    partially_shredded_variant_array_gen!(partially_shredded_timestamp_micro_variant_array, || {
+        arrow::array::TimestampMicrosecondArray::from(vec![
+            Some(-456000),
+            None,
+            None,
+            Some(1758602096000000),
+        ])
+        .with_timezone("+00:00")
+    });
+
+    #[test]
+    fn get_variant_partial_shredded_timestamp_micro_as_variant() {
+        let array = partially_shredded_timestamp_micro_variant_array();
+        assert_variant_get_as_variant_array_with_default_option!(
+            array,
+            vec![
+                Some(Variant::from(
+                    DateTime::from_timestamp_micros(-456000i64)
+                        .unwrap()
+                        .to_utc(),
+                )),
+                None,
+                Some(Variant::from("n/a")),
+                Some(Variant::from(
+                    DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
+                        .unwrap()
+                        .to_utc(),
+                )),
+            ]
+        )
+    }
+
+    partially_shredded_variant_array_gen!(
+        partially_shredded_timestamp_nano_ntz_variant_array,
+        || {
+            arrow::array::TimestampNanosecondArray::from(vec![
+                Some(-4999999561),
+                None,
+                None,
+                Some(1758602096000000000),
+            ])
+        }
+    );
+
+    #[test]
+    fn get_variant_partial_shredded_timestamp_nano_ntz_as_variant() {
+        let array = partially_shredded_timestamp_nano_ntz_variant_array();
+
+        assert_variant_get_as_variant_array_with_default_option!(
+            array,
+            vec![
+                Some(Variant::from(
+                    DateTime::from_timestamp(-5, 439).unwrap().naive_utc()
+                )),
+                None,
+                Some(Variant::from("n/a")),
+                Some(Variant::from(
+                    DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
+                        .unwrap()
+                        .naive_utc()
+                )),
+            ]
+        )
+    }
+
+    partially_shredded_variant_array_gen!(partially_shredded_timestamp_nano_variant_array, || {
+        arrow::array::TimestampNanosecondArray::from(vec![
+            Some(-4999999561),
+            None,
+            None,
+            Some(1758602096000000000),
+        ])
+        .with_timezone("+00:00")
+    });
+
+    #[test]
+    fn get_variant_partial_shredded_timestamp_nano_as_variant() {
+        let array = partially_shredded_timestamp_nano_variant_array();
+
+        assert_variant_get_as_variant_array_with_default_option!(
+            array,
+            vec![
+                Some(Variant::from(
+                    DateTime::from_timestamp(-5, 439).unwrap().to_utc()
+                )),
+                None,
+                Some(Variant::from("n/a")),
+                Some(Variant::from(
+                    DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
+                        .unwrap()
+                        .to_utc()
+                )),
+            ]
+        )
+    }
+
     /// Return a VariantArray that represents a normal "shredded" variant
     /// for the following example
     ///
@@ -957,6 +1101,52 @@ mod test {
         }
     }
 
+    macro_rules! partially_shredded_variant_array_gen {
+        ($func:ident, $typed_array_gen: expr) => {
+            fn $func() -> ArrayRef {
+                // At the time of writing, the `VariantArrayBuilder` does not support shredding.
+                // so we must construct the array manually.  see https://github.com/apache/arrow-rs/issues/7895
+                let (metadata, string_value) = {
+                    let mut builder = parquet_variant::VariantBuilder::new();
+                    builder.append_value("n/a");
+                    builder.finish()
+                };
+
+                let nulls = NullBuffer::from(vec![
+                    true,  // row 0 non null
+                    false, // row 1 is null
+                    true,  // row 2 non null
+                    true,  // row 3 non null
+                ]);
+
+                // metadata is the same for all rows
+                let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 4));
+
+                // See https://docs.google.com/document/d/1pw0AWoMQY3SjD7R4LgbPvMjG_xSCtXp3rZHkVp9jpZ4/edit?disco=AAABml8WQrY
+                // about why row1 is an empty but non null, value.
+                let values = BinaryViewArray::from(vec![
+                    None,                // row 0 is shredded, so no value
+                    Some(b"" as &[u8]),  // row 1 is null, so empty value (why?)
+                    Some(&string_value), // copy the string value "N/A"
+                    None,                // row 3 is shredded, so no value
+                ]);
+
+                let typed_value = $typed_array_gen();
+
+                let struct_array = StructArrayBuilder::new()
+                    .with_field("metadata", Arc::new(metadata), false)
+                    .with_field("typed_value", Arc::new(typed_value), true)
+                    .with_field("value", Arc::new(values), true)
+                    .with_nulls(nulls)
+                    .build();
+
+                ArrayRef::from(
+                    VariantArray::try_new(&struct_array).expect("should create variant array"),
+                )
+            }
+        };
+    }
+
     numeric_partially_shredded_variant_array_fn!(
         partially_shredded_int8_variant_array,
         Int8Array,
@@ -1054,7 +1244,7 @@ mod test {
             None,          // row 2 is a string
             Some("world"), // row 3 is shredded
         ])
-        );
+    );
 
     /// Return a VariantArray that represents a partially "shredded" variant for Date32
     fn partially_shredded_date32_variant_array() -> ArrayRef {

parquet/tests/variant_integration.rs

@@ -91,23 +91,10 @@ variant_test_case!(16);
 variant_test_case!(17);
 variant_test_case!(18);
 variant_test_case!(19);
-// https://github.com/apache/arrow-rs/issues/8331
-variant_test_case!(
-    20,
-    "Unsupported typed_value type: Timestamp(Microsecond, Some(\"UTC\"))"
-);
-variant_test_case!(
-    21,
-    "Unsupported typed_value type: Timestamp(Microsecond, Some(\"UTC\"))"
-);
-variant_test_case!(
-    22,
-    "Unsupported typed_value type: Timestamp(Microsecond, None)"
-);
-variant_test_case!(
-    23,
-    "Unsupported typed_value type: Timestamp(Microsecond, None)"
-);
+variant_test_case!(20);
+variant_test_case!(21);
+variant_test_case!(22);
+variant_test_case!(23);
 // https://github.com/apache/arrow-rs/issues/8332
 variant_test_case!(24, "Unsupported typed_value type: Decimal128(9, 4)");
 variant_test_case!(25, "Unsupported typed_value type: Decimal128(9, 4)");
@@ -119,23 +106,10 @@ variant_test_case!(30);
 variant_test_case!(31);
 // https://github.com/apache/arrow-rs/issues/8334
 variant_test_case!(32, "Unsupported typed_value type: Time64(Microsecond)");
-// https://github.com/apache/arrow-rs/issues/8331
-variant_test_case!(
-    33,
-    "Unsupported typed_value type: Timestamp(Nanosecond, Some(\"UTC\"))"
-);
-variant_test_case!(
-    34,
-    "Unsupported typed_value type: Timestamp(Nanosecond, Some(\"UTC\"))"
-);
-variant_test_case!(
-    35,
-    "Unsupported typed_value type: Timestamp(Nanosecond, None)"
-);
-variant_test_case!(
-    36,
-    "Unsupported typed_value type: Timestamp(Nanosecond, None)"
-);
+variant_test_case!(33);
+variant_test_case!(34);
+variant_test_case!(35);
+variant_test_case!(36);
 variant_test_case!(37);
 // https://github.com/apache/arrow-rs/issues/8336
 variant_test_case!(38, "Unsupported typed_value type: Struct(");