Support converting TimestampTZ to and from duckdb (#4226)

fix #4208 

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

Cargo.lock

@@ -6,6 +6,7 @@ use std::time::{Duration, Instant};
 use anyhow::Result;
 use log::trace;
 use url::Url;
+use vortex::error::VortexExpect;
 use vortex_duckdb::duckdb::{Connection, Database};
 
 use crate::{BenchmarkDataset, Format, IdempotentPath};
@@ -71,7 +72,10 @@ impl DuckDBCtx {
         let query_time = time_instant.elapsed();
         trace!("query completed in {:.3}s", query_time.as_secs_f64());
 
-        Ok((query_time, result.row_count()?))
+        Ok((
+            query_time,
+            usize::try_from(result.row_count()).vortex_expect("row count overflow"),
+        ))
     }
 
     /// Register tables for benchmarks using the internal connection

bench-vortex/src/engines/ddb/mod.rs

@@ -3,7 +3,6 @@
 
 use std::mem::MaybeUninit;
 
-use arrow_buffer::ArrowNativeType;
 use fastlanes::BitPacking;
 use itertools::Itertools;
 use num_traits::{AsPrimitive, PrimInt};
@@ -133,10 +132,7 @@ pub unsafe fn bitpack_unchecked(
 /// Bitpack a slice of primitives down to the given width.
 ///
 /// See `bitpack` for more caller information.
-pub fn bitpack_primitive<T: NativePType + BitPacking + ArrowNativeType>(
-    array: &[T],
-    bit_width: u8,
-) -> Buffer<T> {
+pub fn bitpack_primitive<T: NativePType + BitPacking>(array: &[T], bit_width: u8) -> Buffer<T> {
     if bit_width == 0 {
         return Buffer::<T>::empty();
     }

encodings/fastlanes/src/bitpacking/compress.rs

@@ -4,7 +4,6 @@
 use std::mem;
 use std::mem::MaybeUninit;
 
-use arrow_buffer::ArrowNativeType;
 use fastlanes::BitPacking;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::compute::{FilterKernel, FilterKernelAdapter, filter};
@@ -43,12 +42,12 @@ register_kernel!(FilterKernelAdapter(BitPackedVTable).lift());
 /// FastLanes decompression is so fast and the bookkeepping necessary to decompress individual
 /// elements is relatively slow. If you prefer to never fully decompress, use
 /// [filter_primitive_no_decompression].
-fn filter_primitive<T: NativePType + BitPacking + ArrowNativeType>(
+fn filter_primitive<T: NativePType + BitPacking>(
     array: &BitPackedArray,
     mask: &Mask,
 ) -> VortexResult<PrimitiveArray> {
     // Short-circuit if the selectivity is high enough.
-    let full_decompression_threshold = match T::get_byte_width() {
+    let full_decompression_threshold = match size_of::<T>() {
         1 => 0.03,
         2 => 0.03,
         4 => 0.075,
@@ -67,7 +66,7 @@ fn filter_primitive<T: NativePType + BitPacking + ArrowNativeType>(
 /// Filter a bit-packed array, without using full decompression.
 ///
 /// You should probably use [filter_primitive].
-fn filter_primitive_no_decompression<T: NativePType + BitPacking + ArrowNativeType>(
+fn filter_primitive_no_decompression<T: NativePType + BitPacking>(
     array: &BitPackedArray,
     mask: &Mask,
 ) -> VortexResult<PrimitiveArray> {
@@ -96,7 +95,7 @@ fn filter_primitive_no_decompression<T: NativePType + BitPacking + ArrowNativeTy
     Ok(values)
 }
 
-fn filter_indices<T: NativePType + BitPacking + ArrowNativeType>(
+fn filter_indices<T: NativePType + BitPacking>(
     array: &BitPackedArray,
     indices_len: usize,
     indices: impl Iterator<Item = usize>,

encodings/fastlanes/src/bitpacking/compute/filter.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use arrow_buffer::ArrowNativeType;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::{BoolArray, DecimalArray, PrimitiveArray, StructArray, VarBinViewArray};
 use vortex_array::builders::{ArrayBuilderExt, builder_with_capacity};
@@ -131,7 +130,7 @@ pub fn take_canonical_array(
     }
 }
 
-fn take_primitive<T: NativePType + ArrowNativeType>(
+fn take_primitive<T: NativePType>(
     primitive_array: PrimitiveArray,
     validity: Validity,
     indices: &[usize],

fuzz/src/array/take.rs

@@ -1,8 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use arrow_buffer::ArrowNativeType;
 use itertools::Itertools as _;
+use num_traits::AsPrimitive;
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::{DecimalDType, NativePType, match_each_integer_ptype};
 use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
@@ -58,7 +58,7 @@ fn patch_typed<I, ValuesDVT, PatchDVT>(
     patched_validity: Validity,
 ) -> VortexResult<DecimalArray>
 where
-    I: NativePType + ArrowNativeType,
+    I: NativePType + AsPrimitive<usize>,
     PatchDVT: NativeDecimalType,
     ValuesDVT: NativeDecimalType,
 {
@@ -71,7 +71,7 @@ where
     }
 
     for (idx, value) in patch_indices.iter().zip_eq(patch_values.into_iter()) {
-        buffer[idx.as_usize() - patch_indices_offset] = <ValuesDVT as BigCast>::from(value).vortex_expect(
+        buffer[idx.as_() - patch_indices_offset] = <ValuesDVT as BigCast>::from(value).vortex_expect(
             "values of a given DecimalDType are representable in all compatible NativeDecimalType",
         );
     }

vortex-array/src/arrays/decimal/patch.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use arrow_buffer::ArrowNativeType;
+use num_traits::AsPrimitive;
 use vortex_dtype::{NativePType, match_each_integer_ptype, match_each_native_ptype};
 use vortex_error::VortexResult;
 
@@ -43,15 +43,15 @@ impl PrimitiveArray {
         patched_validity: Validity,
     ) -> VortexResult<Self>
     where
-        T: NativePType + ArrowNativeType,
-        I: NativePType + ArrowNativeType,
+        T: NativePType,
+        I: NativePType + AsPrimitive<usize>,
     {
         let mut own_values = self.into_buffer_mut::<T>();
 
         let patch_indices = patch_indices.as_slice::<I>();
         let patch_values = patch_values.as_slice::<T>();
         for (idx, value) in itertools::zip_eq(patch_indices, patch_values) {
-            own_values[idx.as_usize() - patch_indices_offset] = *value;
+            own_values[idx.as_() - patch_indices_offset] = *value;
         }
         Ok(Self::new(own_values, patched_validity))
     }

vortex-array/src/arrays/primitive/patch.rs

@@ -40,6 +40,7 @@ vortex = { workspace = true, features = ["files"] }
 vortex-file = { workspace = true, features = ["tokio", "object_store"] }
 
 [dev-dependencies]
+jiff = { workspace = true }
 rstest = { workspace = true }
 
 [lints]

vortex-duckdb/Cargo.toml

@@ -93,7 +93,7 @@ impl LogicalType {
         unsafe {
             let ptr = cpp::duckdb_create_list_type(element_type.as_ptr());
             if ptr.is_null() {
-                return Err(vortex_err!("Failed to create list type"));
+                vortex_bail!("Failed to create list type");
             }
             Ok(Self::own(ptr))
         }
@@ -119,24 +119,28 @@ impl LogicalType {
         let duckdb_type = match temporal_metadata {
             TemporalMetadata::Date(TimeUnit::D) => DUCKDB_TYPE::DUCKDB_TYPE_DATE,
             TemporalMetadata::Date(time_unit) => {
-                return Err(vortex_err!("Invalid TimeUnit {} for date", time_unit));
+                vortex_bail!("Invalid TimeUnit {} for date", time_unit);
             }
             TemporalMetadata::Time(TimeUnit::Us) => DUCKDB_TYPE::DUCKDB_TYPE_TIME,
             TemporalMetadata::Time(time_unit) => {
-                return Err(vortex_err!("Invalid TimeUnit {} for time", time_unit));
+                vortex_bail!("Invalid TimeUnit {} for time", time_unit);
             }
-            TemporalMetadata::Timestamp(time_unit, tz) => {
-                if tz.is_some() {
-                    return Err(vortex_err!("Timestamp with timezone is not yet supported"));
+            TemporalMetadata::Timestamp(time_unit, tz) => match time_unit {
+                TimeUnit::Ns => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS,
+                TimeUnit::Us => {
+                    if let Some(tz) = tz {
+                        if tz != "UTC" {
+                            vortex_bail!("Invalid timezone for timestamp: {tz}");
+                        }
+                        DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ
+                    } else {
+                        DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP
+                    }
                 }
-                match time_unit {
-                    TimeUnit::Ns => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS,
-                    TimeUnit::Us => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP,
-                    TimeUnit::Ms => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_MS,
-                    TimeUnit::S => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_S,
-                    _ => return Err(vortex_err!("Invalid TimeUnit {} for timestamp", time_unit)),
-                }
-            }
+                TimeUnit::Ms => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_MS,
+                TimeUnit::S => DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_S,
+                _ => vortex_bail!("Invalid TimeUnit {} for timestamp", time_unit),
+            },
         };
 
         Ok(Self::new(duckdb_type))
@@ -171,7 +175,6 @@ impl FromLogicalType for DType {
             DUCKDB_TYPE::DUCKDB_TYPE_UHUGEINT => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_FLOAT => DType::Primitive(F32, nullability),
             DUCKDB_TYPE::DUCKDB_TYPE_DOUBLE => DType::Primitive(F64, nullability),
-
             DUCKDB_TYPE::DUCKDB_TYPE_VARCHAR => DType::Utf8(nullability),
             DUCKDB_TYPE::DUCKDB_TYPE_BLOB => DType::Binary(nullability),
             DUCKDB_TYPE::DUCKDB_TYPE_DECIMAL => {
@@ -206,11 +209,17 @@ impl FromLogicalType for DType {
                 Arc::new(DType::Primitive(I64, nullability)),
                 Some(TemporalMetadata::Timestamp(TimeUnit::Us, None).into()),
             ))),
+            DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ => DType::Extension(Arc::new(ExtDType::new(
+                TIMESTAMP_ID.clone(),
+                Arc::new(DType::Primitive(I64, nullability)),
+                Some(TemporalMetadata::Timestamp(TimeUnit::Us, Some("UTC".to_string())).into()),
+            ))),
             DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_NS => DType::Extension(Arc::new(ExtDType::new(
                 TIMESTAMP_ID.clone(),
                 Arc::new(DType::Primitive(I64, nullability)),
                 Some(TemporalMetadata::Timestamp(TimeUnit::Ns, None).into()),
             ))),
+            DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_INTERVAL => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_ENUM => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_LIST => todo!(),
@@ -220,8 +229,6 @@ impl FromLogicalType for DType {
             DUCKDB_TYPE::DUCKDB_TYPE_UUID => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_UNION => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_BIT => todo!(),
-            DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ => todo!(),
-            DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_ANY => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_VARINT => todo!(),
             DUCKDB_TYPE::DUCKDB_TYPE_STRING_LITERAL => todo!(),
@@ -526,6 +533,26 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_timestamp_with_timezone() {
+        use std::sync::Arc;
+
+        use vortex::dtype::datetime::{TIMESTAMP_ID, TemporalMetadata, TimeUnit};
+        use vortex::dtype::{ExtDType, PType};
+
+        let ext_dtype = ExtDType::new(
+            TIMESTAMP_ID.clone(),
+            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
+            Some(TemporalMetadata::Timestamp(TimeUnit::Us, Some("UTC".to_string())).into()),
+        );
+        let dtype = DType::Extension(Arc::new(ext_dtype));
+
+        assert_eq!(
+            LogicalType::try_from(&dtype).unwrap().as_type_id(),
+            cpp::DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ
+        );
+    }
+
     #[test]
     fn test_temporal_extension_invalid_time_units() {
         use std::sync::Arc;
@@ -552,23 +579,6 @@ mod tests {
         assert!(LogicalType::try_from(&dtype).is_err());
     }
 
-    #[test]
-    fn test_timestamp_with_timezone_unsupported() {
-        use std::sync::Arc;
-
-        use vortex::dtype::datetime::{TIMESTAMP_ID, TemporalMetadata, TimeUnit};
-        use vortex::dtype::{ExtDType, PType};
-
-        let ext_dtype = ExtDType::new(
-            TIMESTAMP_ID.clone(),
-            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
-            Some(TemporalMetadata::Timestamp(TimeUnit::Us, Some("UTC".to_string())).into()),
-        );
-        let dtype = DType::Extension(Arc::new(ext_dtype));
-
-        assert!(LogicalType::try_from(&dtype).is_err());
-    }
-
     #[test]
     fn test_unsupported_extension_type() {
         use vortex::dtype::{ExtDType, ExtID, PType};

vortex-duckdb/src/convert/dtype.rs

@@ -6,16 +6,16 @@ use std::sync::Arc;
 
 use arrow_array::builder::GenericBinaryBuilder;
 use arrow_array::types::{
-    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType,
-    UInt8Type, UInt16Type, UInt32Type, UInt64Type,
+    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type,
+    UInt32Type, UInt64Type,
 };
 use arrow_array::{
     Array, BooleanArray, Date32Array, Decimal128Array, PrimitiveArray, StringArray,
-    TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-    TimestampSecondArray,
+    Time64MicrosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
+    TimestampNanosecondArray, TimestampSecondArray,
 };
 use arrow_buffer::buffer::BooleanBuffer;
-use bitvec::macros::internal::funty::Fundamental;
+use num_traits::AsPrimitive;
 use vortex::ArrayRef;
 use vortex::arrays::StructArray;
 use vortex::arrow::FromArrowArray;
@@ -116,13 +116,13 @@ pub fn flat_vector_to_arrow_array(
 
             Ok(Arc::new(structs))
         }
-        DUCKDB_TYPE::DUCKDB_TYPE_TIME_TZ => {
+        DUCKDB_TYPE::DUCKDB_TYPE_TIMESTAMP_TZ => {
             let data = vector.as_slice_with_len::<duckdb_timestamp>(len);
-            let micros = data.iter().map(|duckdb_timestamp { micros }| *micros);
             let structs = TimestampMicrosecondArray::from_iter_values_with_nulls(
-                micros,
+                data.iter().map(|duckdb_timestamp { micros }| *micros),
                 vector.validity_ref(data.len()).to_null_buffer(),
-            );
+            )
+            .with_timezone("UTC");
 
             Ok(Arc::new(structs))
         }
@@ -181,7 +181,7 @@ pub fn flat_vector_to_arrow_array(
             let data = vector.as_slice_with_len::<duckdb_time>(len);
 
             Ok(Arc::new(
-                PrimitiveArray::<Time64MicrosecondType>::from_iter_values_with_nulls(
+                Time64MicrosecondArray::from_iter_values_with_nulls(
                     data.iter().map(|duckdb_time { micros }| *micros),
                     vector.validity_ref(data.len()).to_null_buffer(),
                 ),
@@ -311,7 +311,7 @@ pub fn flat_vector_to_arrow_array(
 
             Ok(Arc::new(decimal_array))
         }
-        _ => todo!("missing impl for {:?}", type_id),
+        _ => todo!("missing impl for {type_id:?}"),
     }
 }
 
@@ -323,12 +323,13 @@ pub fn data_chunk_to_arrow(field_names: &FieldNames, chunk: &DataChunk) -> Vorte
         .map(|(i, name)| {
             let mut vector = chunk.get_vector(i);
             vector.flatten(len);
-            flat_vector_to_arrow_array(&mut vector, len.as_usize())
+            flat_vector_to_arrow_array(&mut vector, len.as_())
                 .map(|array_data| {
-                    assert_eq!(array_data.len(), chunk.len().as_usize());
+                    let chunk_len: usize = chunk.len().as_();
+                    assert_eq!(array_data.len(), chunk_len);
                     (name, ArrayRef::from_arrow(array_data.as_ref(), true))
                 })
-                .map_err(|e| vortex_err!("duckdb to arrow conversion failure {}", e.to_string()))
+                .map_err(|e| vortex_err!("duckdb to arrow conversion failure {e}"))
         })
         .collect::<VortexResult<Vec<_>>>()?;
     StructArray::try_from_iter(columns).map(|a| a.to_array())