Extension DType vtables (#6081)

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

Cargo.lock

@@ -11,8 +11,8 @@ use vortex_array::validity::Validity;
 use vortex_buffer::BufferMut;
 use vortex_dtype::DType;
 use vortex_dtype::PType;
-use vortex_dtype::datetime::TemporalMetadata;
 use vortex_dtype::datetime::TimeUnit;
+use vortex_dtype::datetime::Timestamp;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexExpect as _;
 use vortex_error::VortexResult;
@@ -31,11 +31,11 @@ pub fn decode_to_temporal(
         vortex_panic!(ComputeError: "expected dtype to be DType::Extension variant")
     };
 
-    let Ok(temporal_metadata) = TemporalMetadata::try_from(ext.as_ref()) else {
+    let Some(options) = ext.metadata_opt::<Timestamp>() else {
         vortex_panic!(ComputeError: "must decode TemporalMetadata from extension metadata");
     };
 
-    let divisor = match temporal_metadata.time_unit() {
+    let divisor = match options.unit {
         TimeUnit::Nanoseconds => 1_000_000_000,
         TimeUnit::Microseconds => 1_000_000,
         TimeUnit::Milliseconds => 1_000,
@@ -98,8 +98,8 @@ pub fn decode_to_temporal(
     Ok(TemporalArray::new_timestamp(
         PrimitiveArray::new(values.freeze(), Validity::copy_from_array(array.as_ref())?)
             .into_array(),
-        temporal_metadata.time_unit(),
-        temporal_metadata.time_zone().map(ToString::to_string),
+        options.unit,
+        options.tz.clone(),
     ))
 }
 
@@ -142,7 +142,7 @@ mod test {
         let date_times = DateTimePartsArray::try_from(TemporalArray::new_timestamp(
             milliseconds.clone().into_array(),
             TimeUnit::Milliseconds,
-            Some("UTC".to_string()),
+            Some("UTC".into()),
         ))
         .unwrap();
 

encodings/datetime-parts/src/canonical.rs

@@ -100,11 +100,8 @@ mod tests {
             validity.clone(),
         )
         .into_array();
-        let temporal_array = TemporalArray::new_timestamp(
-            milliseconds,
-            TimeUnit::Milliseconds,
-            Some("UTC".to_string()),
-        );
+        let temporal_array =
+            TemporalArray::new_timestamp(milliseconds, TimeUnit::Milliseconds, Some("UTC".into()));
         let TemporalParts {
             days,
             seconds,

encodings/datetime-parts/src/compress.rs

@@ -66,7 +66,7 @@ mod tests {
             )
             .into_array(),
             TimeUnit::Milliseconds,
-            Some("UTC".to_string()),
+            Some("UTC".into()),
         ))
         .unwrap()
         .into_array()
@@ -99,7 +99,7 @@ mod tests {
         let result = cast(&array, &DType::Bool(Nullability::NonNullable));
         assert!(
             result.as_ref().is_err_and(|err| err.to_string().contains(
-                "No compute kernel to cast array vortex.ext with dtype ext(vortex.timestamp, i64, ExtMetadata([2, 3, 0, 85, 84, 67]))? to bool"
+                "No compute kernel to cast array vortex.ext with dtype vortex.timestamp[unit=ms, tz=UTC](i64?) to bool"
             )),
             "Got error: {result:?}"
         );
@@ -126,7 +126,7 @@ mod tests {
             345_600_000, // 4 days in ms
         ].into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     #[case(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         PrimitiveArray::from_option_iter([
@@ -137,12 +137,12 @@ mod tests {
             None,
         ]).into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     #[case(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![86_400_000_000_000i64].into_array(), // 1 day in ns
         TimeUnit::Nanoseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     fn test_cast_datetime_parts_conformance(#[case] array: DateTimePartsArray) {
         use vortex_array::compute::conformance::cast::test_cast_conformance;

encodings/datetime-parts/src/compute/cast.rs

@@ -15,7 +15,7 @@ use vortex_array::compute::or;
 use vortex_array::register_kernel;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
-use vortex_dtype::datetime::TemporalMetadata;
+use vortex_dtype::datetime::Timestamp;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
 
@@ -50,8 +50,10 @@ impl CompareKernel for DateTimePartsVTable {
 
         let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
 
-        let temporal_metadata = TemporalMetadata::try_from(ext_dtype.as_ref())?;
-        let ts_parts = timestamp::split(timestamp, temporal_metadata.time_unit())?;
+        let Some(options) = ext_dtype.metadata_opt::<Timestamp>() else {
+            return Ok(None);
+        };
+        let ts_parts = timestamp::split(timestamp, options.unit)?;
 
         match operator {
             Operator::Eq => compare_eq(lhs, &ts_parts, nullability),
@@ -218,7 +220,7 @@ mod test {
         DateTimePartsArray::try_from(TemporalArray::new_timestamp(
             PrimitiveArray::new(buffer![value], validity).into_array(),
             TimeUnit::Seconds,
-            Some("UTC".to_string()),
+            Some("UTC".into()),
         ))
         .expect("Failed to construct DateTimePartsArray from TemporalArray")
     }
@@ -293,7 +295,7 @@ mod test {
         let temporal_array = TemporalArray::new_timestamp(
             PrimitiveArray::new(buffer![0i64], lhs_validity.clone()).into_array(),
             TimeUnit::Seconds,
-            Some("UTC".to_string()),
+            Some("UTC".into()),
         );
 
         let lhs = DateTimePartsArray::try_new(

encodings/datetime-parts/src/compute/compare.rs

@@ -48,11 +48,8 @@ mod test {
         ]
         .into_array();
 
-        let temporal = TemporalArray::new_timestamp(
-            timestamps,
-            TimeUnit::Milliseconds,
-            Some("UTC".to_string()),
-        );
+        let temporal =
+            TemporalArray::new_timestamp(timestamps, TimeUnit::Milliseconds, Some("UTC".into()));
 
         let array = DateTimePartsArray::try_from(temporal).unwrap();
         test_filter_conformance(array.as_ref());
@@ -67,11 +64,8 @@ mod test {
         ])
         .into_array();
 
-        let temporal = TemporalArray::new_timestamp(
-            timestamps,
-            TimeUnit::Milliseconds,
-            Some("UTC".to_string()),
-        );
+        let temporal =
+            TemporalArray::new_timestamp(timestamps, TimeUnit::Milliseconds, Some("UTC".into()));
 
         let array = DateTimePartsArray::try_from(temporal).unwrap();
         test_filter_conformance(array.as_ref());

encodings/datetime-parts/src/compute/filter.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::compute::MaskKernel;
 use vortex_array::compute::MaskKernelAdapter;
@@ -23,17 +24,14 @@ impl MaskKernel for DateTimePartsVTable {
         // through the days component.
 
         let masked_days = mask(array.days(), mask_array)?;
+        assert!(masked_days.dtype().is_nullable());
 
         // Keep seconds and subseconds unchanged since they must remain non-nullable
         let seconds = array.seconds().clone();
         let subseconds = array.subseconds().clone();
 
         // Update the dtype to reflect the new nullability of days
-        let new_dtype = if masked_days.dtype().is_nullable() {
-            array.dtype().as_nullable()
-        } else {
-            array.dtype().clone()
-        };
+        let new_dtype = array.dtype().as_nullable();
 
         DateTimePartsArray::try_new(new_dtype, masked_days, seconds, subseconds)
             .map(|a| a.to_array())

encodings/datetime-parts/src/compute/mask.rs

@@ -26,46 +26,46 @@ mod tests {
     #[case::datetime_seconds(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![0i64, 86400, 172800, 259200, 345600].into_array(),
         TimeUnit::Seconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     #[case::datetime_millis(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![0i64, 86400000, 172800000].into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     #[case::datetime_micros(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![0i64, 86400000000, 172800000000].into_array(),
         TimeUnit::Microseconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     #[case::datetime_nanos(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![0i64, 86400000000000].into_array(),
         TimeUnit::Nanoseconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     // Nullable arrays
     #[case::datetime_nullable_seconds(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         PrimitiveArray::from_option_iter([Some(0i64), None, Some(86400), Some(172800), None]).into_array(),
         TimeUnit::Seconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     // Edge cases
     #[case::datetime_single(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![1234567890i64].into_array(),
         TimeUnit::Seconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     // Large arrays (> 1024 elements)
     #[case::datetime_large(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         PrimitiveArray::from_iter((0..1500).map(|i| i as i64 * 86400)).into_array(),
         TimeUnit::Seconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
     // Different time patterns
     #[case::datetime_with_subseconds(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![123456789i64, 234567890, 345678901, 456789012, 567890123].into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string()),
+        Some("UTC".into()),
     )).unwrap())]
 
     fn test_datetime_parts_consistency(#[case] array: DateTimePartsArray) {

encodings/datetime-parts/src/compute/mod.rs

@@ -18,7 +18,7 @@ use vortex_array::optimizer::ArrayOptimizer;
 use vortex_array::optimizer::rules::ArrayParentReduceRule;
 use vortex_array::optimizer::rules::ParentRuleSet;
 use vortex_dtype::DType;
-use vortex_dtype::datetime::TemporalMetadata;
+use vortex_dtype::datetime::Timestamp;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
@@ -161,8 +161,8 @@ fn try_extract_days_constant(array: &ArrayRef) -> Option<i64> {
         return None;
     };
 
-    let temporal_metadata = TemporalMetadata::try_from(ext_dtype.as_ref()).ok()?;
-    let ts_parts = timestamp::split(timestamp, temporal_metadata.time_unit()).ok()?;
+    let options = ext_dtype.metadata::<Timestamp>();
+    let ts_parts = timestamp::split(timestamp, options.unit).ok()?;
 
     // Only allow pushdown if seconds and subseconds are zero
     if ts_parts.seconds != 0 || ts_parts.subseconds != 0 {
@@ -190,8 +190,8 @@ mod tests {
     use vortex_array::optimizer::ArrayOptimizer;
     use vortex_array::validity::Validity;
     use vortex_buffer::Buffer;
-    use vortex_buffer::buffer;
     use vortex_dtype::datetime::TimeUnit;
+    use vortex_dtype::datetime::TimestampOptions;
     use vortex_scalar::Scalar;
 
     use super::*;
@@ -230,12 +230,13 @@ mod tests {
             TimeUnit::Days => panic!("Days not supported"),
         };
         let timestamp = day * SECONDS_PER_DAY * multiplier;
-        let temporal = TemporalArray::new_timestamp(
-            PrimitiveArray::new(buffer![timestamp], Validity::NonNullable).into_array(),
-            time_unit,
-            None,
+        let scalar = Scalar::extension::<Timestamp>(
+            TimestampOptions {
+                unit: time_unit,
+                tz: None,
+            },
+            timestamp.into(),
         );
-        let scalar = Scalar::extension(temporal.ext_dtype(), timestamp.into());
         ConstantArray::new(scalar, len).into_array()
     }
 
@@ -249,12 +250,13 @@ mod tests {
             TimeUnit::Days => panic!("Days not supported"),
         };
         let timestamp = (day * SECONDS_PER_DAY + seconds) * multiplier;
-        let temporal = TemporalArray::new_timestamp(
-            PrimitiveArray::new(buffer![timestamp], Validity::NonNullable).into_array(),
-            time_unit,
-            None,
+        let scalar = Scalar::extension::<Timestamp>(
+            TimestampOptions {
+                unit: time_unit,
+                tz: None,
+            },
+            timestamp.into(),
         );
-        let scalar = Scalar::extension(temporal.ext_dtype(), timestamp.into());
         ConstantArray::new(scalar, len).into_array()
     }
 

encodings/datetime-parts/src/compute/rules.rs

@@ -117,7 +117,7 @@ mod tests {
             345_600_000, // 4 days in ms
         ].into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     #[case(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         PrimitiveArray::from_option_iter([
@@ -128,12 +128,12 @@ mod tests {
             None,
         ]).into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     #[case(DateTimePartsArray::try_from(TemporalArray::new_timestamp(
         buffer![86_400_000i64].into_array(),
         TimeUnit::Milliseconds,
-        Some("UTC".to_string())
+        Some("UTC".into())
     )).unwrap())]
     fn test_take_datetime_parts_conformance(#[case] array: DateTimePartsArray) {
         test_take_conformance(array.as_ref());