[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,22 @@ impl<'a> GetOptions<'a> {
 
 #[cfg(test)]
 mod test {
-    use std::sync::Arc;
-
+    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, BooleanArray, Date32Array, FixedSizeBinaryArray,
-        Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
-        StringArray, StructArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+        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
@@ -454,9 +452,8 @@ mod test {
                     .with_field("value", Arc::new(values), true)
                     .with_nulls(nulls)
                     .build();
-                Arc::new(
-                    VariantArray::try_new(Arc::new(struct_array))
-                        .expect("should create variant array"),
+                ArrayRef::from(
+                    VariantArray::try_new(&struct_array).expect("should create variant array"),
                 )
             }
         };
@@ -924,6 +921,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
     ///
@@ -955,10 +1102,11 @@ mod test {
                 None, // row 2 is a string, so no typed value
                 Some(<$primitive_type>::try_from(100u8).unwrap()), // row 3 is shredded, so it has a value
             ]));
+        };
     }
-    
-    macro_rules partially_shredded_variant_array_gen {
-         ($func:ident) => {
+
+    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
@@ -987,12 +1135,7 @@ mod test {
                     None,                // row 3 is shredded, so no value
                 ]);
 
-                let typed_value = $array_type::from(vec![
-                    Some(<$primitive_type>::try_from(34u8).unwrap()), // row 0 is shredded, so it has a value
-                    None,                                             // row 1 is null, so no value
-                    None, // row 2 is a string, so no typed value
-                    Some(<$primitive_type>::try_from(100u8).unwrap()), // row 3 is shredded, so it has a value
-                ]);
+                let typed_value = $typed_array_gen();
 
                 let struct_array = StructArrayBuilder::new()
                     .with_field("metadata", Arc::new(metadata), false)
@@ -1001,9 +1144,8 @@ mod test {
                     .with_nulls(nulls)
                     .build();
 
-                Arc::new(
-                    VariantArray::try_new(Arc::new(struct_array))
-                        .expect("should create variant array"),
+                ArrayRef::from(
+                    VariantArray::try_new(&struct_array).expect("should create variant array"),
                 )
             }
         };
@@ -1090,24 +1232,23 @@ mod test {
                 false, // row 2 is string
                 true,  // row 3 has value
             ]);
-            let typed_value = arrow::array::FixedSizeBinaryArray::try_new(
+            arrow::array::FixedSizeBinaryArray::try_new(
                 3, // byte width
                 arrow::buffer::Buffer::from(data),
                 Some(typed_value_nulls),
             )
-            .expect("should create fixed size binary array");
-
-    }
+            .expect("should create fixed size binary array")
+        }
+    );
 
-    /// Return a VariantArray that represents a partially "shredded" variant for UTF8
-    partially_shredded_variant_array_gen!(partially_shredded_utf8_variant_array, ||
+    partially_shredded_variant_array_gen!(partially_shredded_utf8_variant_array, || {
         StringArray::from(vec![
             Some("hello"), // row 0 is shredded
             None,          // row 1 is null
             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 {