16935267080

Committed 13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%

Build # 16935267080

Build Type

Pull #4226

github

Committed by

web-flow

Commit Message

Merge 81b48c7fb into baa6ea202

Pull Request Pull Request #4226: Support converting TimestampTZ to and from duckdb

Run Details

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/vortex-array/src/arrays/datetime/mod.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

#[cfg(test)]
mod test;

use std::sync::Arc;

use vortex_dtype::datetime::{DATE_ID, TIME_ID, TIMESTAMP_ID, TemporalMetadata, TimeUnit};
use vortex_dtype::{DType, ExtDType};
use vortex_error::{VortexError, vortex_err, vortex_panic};

use crate::arrays::{ExtensionArray, ExtensionVTable};
use crate::{Array, ArrayRef, IntoArray};

/// An array wrapper for primitive values that have an associated temporal meaning.
///
/// This is a wrapper around ExtensionArrays containing numeric types, each of which corresponds to
/// either a timestamp or julian date (both referenced to UNIX epoch), OR a time since midnight.
///
/// ## Arrow compatibility
///
/// TemporalArray can be created from Arrow arrays containing the following datatypes:
/// * `Time32`
/// * `Time64`
/// * `Timestamp`
/// * `Date32`
/// * `Date64`
///
/// Anything that can be constructed and held in a `TemporalArray` can also be zero-copy converted
/// back to the relevant Arrow datatype.
#[derive(Clone, Debug)]
pub struct TemporalArray {
    /// The underlying Vortex extension array holding all the numeric values.
    ext: ExtensionArray,

    /// In-memory representation of the ExtMetadata that is held by the underlying extension array.
    ///
    /// We hold this directly to avoid needing to deserialize the metadata to access things like
    /// timezone and TimeUnit of the underlying array.
    temporal_metadata: TemporalMetadata,
}

macro_rules! assert_width {
    ($width:ty, $array:expr) => {{
        let DType::Primitive(ptype, _) = $array.dtype() else {
            panic!("array must have primitive type");
        };

        assert_eq!(
            <$width as vortex_dtype::NativePType>::PTYPE,
            *ptype,
            "invalid ptype {} for array, expected {}",
            <$width as vortex_dtype::NativePType>::PTYPE,
            *ptype
        );
    }};
}

impl TemporalArray {
    /// Create a new `TemporalArray` holding either i32 day offsets, or i64 millisecond offsets
    /// that are evenly divisible by the number of 86,400,000.
    ///
    /// This is equivalent to the data described by either of the `Date32` or `Date64` data types
    /// from Arrow.
    ///
    /// # Panics
    ///
    /// If the time unit is milliseconds, and the array is not of primitive I64 type, it panics.
    ///
    /// If the time unit is days, and the array is not of primitive I32 type, it panics.
    ///
    /// If any other time unit is provided, it panics.
    pub fn new_date(array: ArrayRef, time_unit: TimeUnit) -> Self {
        match time_unit {
            TimeUnit::D => {
                assert_width!(i32, array);
            }
            TimeUnit::Ms => {
                assert_width!(i64, array);
            }
            _ => vortex_panic!("invalid TimeUnit {time_unit} for vortex.date"),
        };

        let ext_dtype = ExtDType::new(
            DATE_ID.clone(),
            Arc::new(array.dtype().clone()),
            Some(TemporalMetadata::Date(time_unit).into()),
        );

        Self {
            ext: ExtensionArray::new(Arc::new(ext_dtype), array),
            temporal_metadata: TemporalMetadata::Date(time_unit),
        }
    }

    /// Create a new `TemporalArray` holding one of the following values:
    ///
    /// * `i32` values representing seconds since midnight
    /// * `i32` values representing milliseconds since midnight
    /// * `i64` values representing microseconds since midnight
    /// * `i64` values representing nanoseconds since midnight
    ///
    /// Note, this is equivalent to the set of values represented by the Time32 or Time64 types
    /// from Arrow.
    ///
    /// # Panics
    ///
    /// If the time unit is seconds, and the array is not of primitive I32 type, it panics.
    ///
    /// If the time unit is milliseconds, and the array is not of primitive I32 type, it panics.
    ///
    /// If the time unit is microseconds, and the array is not of primitive I64 type, it panics.
    ///
    /// If the time unit is nanoseconds, and the array is not of primitive I64 type, it panics.
    pub fn new_time(array: ArrayRef, time_unit: TimeUnit) -> Self {
        match time_unit {
            TimeUnit::S | TimeUnit::Ms => assert_width!(i32, array),
            TimeUnit::Us | TimeUnit::Ns => assert_width!(i64, array),
            TimeUnit::D => vortex_panic!("invalid unit D for vortex.time data"),
        }

        let temporal_metadata = TemporalMetadata::Time(time_unit);
        Self {
            ext: ExtensionArray::new(
                Arc::new(ExtDType::new(
                    TIME_ID.clone(),
                    Arc::new(array.dtype().clone()),
                    Some(temporal_metadata.clone().into()),
                )),
                array,
            ),
            temporal_metadata,
        }
    }

    /// Create a new `TemporalArray` holding Arrow spec compliant Timestamp data, with an
    /// optional timezone.
    ///
    /// # Panics
    ///
    /// If `array` does not hold Primitive i64 data, the function will panic.
    ///
    /// If the time_unit is days, the function will panic.
    pub fn new_timestamp(array: ArrayRef, time_unit: TimeUnit, time_zone: Option<String>) -> Self {
        assert_width!(i64, array);

        let temporal_metadata = TemporalMetadata::Timestamp(time_unit, time_zone);

        Self {
            ext: ExtensionArray::new(
                Arc::new(ExtDType::new(
                    TIMESTAMP_ID.clone(),
                    Arc::new(array.dtype().clone()),
                    Some(temporal_metadata.clone().into()),
                )),
                array,
            ),
            temporal_metadata,
        }
    }
}

impl TemporalArray {
    /// Access the underlying temporal values in the underlying ExtensionArray storage.
    ///
    /// These values are to be interpreted based on the time unit and optional time-zone stored
    /// in the TemporalMetadata.
    pub fn temporal_values(&self) -> &ArrayRef {
        self.ext.storage()
    }

    /// Retrieve the temporal metadata.
    ///
    /// The metadata is used to provide semantic meaning to the temporal values Array, for example
    /// to understand the granularity of the samples and if they have an associated timezone.
    pub fn temporal_metadata(&self) -> &TemporalMetadata {
        &self.temporal_metadata
    }

    /// Retrieve the extension DType associated with the underlying array.
    pub fn ext_dtype(&self) -> Arc<ExtDType> {
        self.ext.ext_dtype().clone()
    }

    /// Retrieve the DType of the array. This will be a `DType::Extension` variant.
    pub fn dtype(&self) -> &DType {
        self.ext.dtype()
    }
}

impl From<TemporalArray> for ArrayRef {
    fn from(value: TemporalArray) -> Self {
        value.ext.into_array()
    }
}

impl IntoArray for TemporalArray {
    fn into_array(self) -> ArrayRef {
        self.into()
    }
}

impl TryFrom<ArrayRef> for TemporalArray {
    type Error = VortexError;

    /// Try to specialize a generic Vortex array as a TemporalArray.
    ///
    /// # Errors
    ///
    /// If the provided Array does not have `vortex.ext` encoding, an error will be returned.
    ///
    /// If the provided Array does not have recognized ExtMetadata corresponding to one of the known
    /// `TemporalMetadata` variants, an error is returned.
    fn try_from(value: ArrayRef) -> Result<Self, Self::Error> {
        let ext = value
            .as_opt::<ExtensionVTable>()
            .ok_or_else(|| vortex_err!("array must be an ExtensionArray"))?;
        let temporal_metadata = TemporalMetadata::try_from(ext.ext_dtype())?;
        Ok(Self {
            ext: ext.clone(),
            temporal_metadata,
        })
    }
}

// Conversions to/from ExtensionArray
impl From<&TemporalArray> for ExtensionArray {
    fn from(value: &TemporalArray) -> Self {
        value.ext.clone()
    }
}

impl From<TemporalArray> for ExtensionArray {
    fn from(value: TemporalArray) -> Self {
        value.ext
    }
}

impl TryFrom<ExtensionArray> for TemporalArray {
    type Error = VortexError;

    fn try_from(ext: ExtensionArray) -> Result<Self, Self::Error> {
        let temporal_metadata = TemporalMetadata::try_from(ext.ext_dtype().as_ref())?;
        Ok(Self {
            ext,
            temporal_metadata,
        })
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	#[cfg(test)]
5	mod test;
6
7	use std::sync::Arc;
8
9	use vortex_dtype::datetime::{DATE_ID, TIME_ID, TIMESTAMP_ID, TemporalMetadata, TimeUnit};
10	use vortex_dtype::{DType, ExtDType};
11	use vortex_error::{VortexError, vortex_err, vortex_panic};
12
13	use crate::arrays::{ExtensionArray, ExtensionVTable};
14	use crate::{Array, ArrayRef, IntoArray};
15
16	/// An array wrapper for primitive values that have an associated temporal meaning.
17	///
18	/// This is a wrapper around ExtensionArrays containing numeric types, each of which corresponds to
19	/// either a timestamp or julian date (both referenced to UNIX epoch), OR a time since midnight.
20	///
21	/// ## Arrow compatibility
22	///
23	/// TemporalArray can be created from Arrow arrays containing the following datatypes:
24	/// * `Time32`
25	/// * `Time64`
26	/// * `Timestamp`
27	/// * `Date32`
28	/// * `Date64`
29	///
30	/// Anything that can be constructed and held in a `TemporalArray` can also be zero-copy converted
31	/// back to the relevant Arrow datatype.
32	#[derive(Clone, Debug)]
33	pub struct TemporalArray {
34	/// The underlying Vortex extension array holding all the numeric values.
35	ext: ExtensionArray,
36
37	/// In-memory representation of the ExtMetadata that is held by the underlying extension array.
38	///
39	/// We hold this directly to avoid needing to deserialize the metadata to access things like
40	/// timezone and TimeUnit of the underlying array.
41	temporal_metadata: TemporalMetadata,
42	}
43
44	macro_rules! assert_width {
45	($width:ty, $array:expr) => {{
46	let DType::Primitive(ptype, _) = $array.dtype() else {
47	panic!("array must have primitive type");
48	};
49
50	assert_eq!(
51	<$width as vortex_dtype::NativePType>::PTYPE,
52	*ptype,
53	"invalid ptype {} for array, expected {}",
54	<$width as vortex_dtype::NativePType>::PTYPE,
55	*ptype
56	);
57	}};
58	}
59
60	impl TemporalArray {
61	/// Create a new `TemporalArray` holding either i32 day offsets, or i64 millisecond offsets
62	/// that are evenly divisible by the number of 86,400,000.
63	///
64	/// This is equivalent to the data described by either of the `Date32` or `Date64` data types
65	/// from Arrow.
66	///
67	/// # Panics
68	///
69	/// If the time unit is milliseconds, and the array is not of primitive I64 type, it panics.
70	///
71	/// If the time unit is days, and the array is not of primitive I32 type, it panics.
72	///
73	/// If any other time unit is provided, it panics.
UNCOV 74	pub fn new_date(array: ArrayRef, time_unit: TimeUnit) -> Self {	×
UNCOV 75	match time_unit {	×
76	TimeUnit::D => {
UNCOV 77	assert_width!(i32, array);	×
78	}
79	TimeUnit::Ms => {
UNCOV 80	assert_width!(i64, array);	×
81	}
82	_ => vortex_panic!("invalid TimeUnit {time_unit} for vortex.date"),	×
83	};
84
UNCOV 85	let ext_dtype = ExtDType::new(	×
UNCOV 86	DATE_ID.clone(),	×
UNCOV 87	Arc::new(array.dtype().clone()),	×
UNCOV 88	Some(TemporalMetadata::Date(time_unit).into()),	×
89	);
90
UNCOV 91	Self {	×
UNCOV 92	ext: ExtensionArray::new(Arc::new(ext_dtype), array),	×
UNCOV 93	temporal_metadata: TemporalMetadata::Date(time_unit),	×
UNCOV 94	}	×
UNCOV 95	}	×
96
97	/// Create a new `TemporalArray` holding one of the following values:
98	///
99	/// * `i32` values representing seconds since midnight
100	/// * `i32` values representing milliseconds since midnight
101	/// * `i64` values representing microseconds since midnight
102	/// * `i64` values representing nanoseconds since midnight
103	///
104	/// Note, this is equivalent to the set of values represented by the Time32 or Time64 types
105	/// from Arrow.
106	///
107	/// # Panics
108	///
109	/// If the time unit is seconds, and the array is not of primitive I32 type, it panics.
110	///
111	/// If the time unit is milliseconds, and the array is not of primitive I32 type, it panics.
112	///
113	/// If the time unit is microseconds, and the array is not of primitive I64 type, it panics.
114	///
115	/// If the time unit is nanoseconds, and the array is not of primitive I64 type, it panics.
UNCOV 116	pub fn new_time(array: ArrayRef, time_unit: TimeUnit) -> Self {	×
UNCOV 117	match time_unit {	×
UNCOV 118	TimeUnit::S \| TimeUnit::Ms => assert_width!(i32, array),	×
UNCOV 119	TimeUnit::Us \| TimeUnit::Ns => assert_width!(i64, array),	×
120	TimeUnit::D => vortex_panic!("invalid unit D for vortex.time data"),	×
121	}
122
UNCOV 123	let temporal_metadata = TemporalMetadata::Time(time_unit);	×
UNCOV 124	Self {	×
UNCOV 125	ext: ExtensionArray::new(	×
UNCOV 126	Arc::new(ExtDType::new(	×
UNCOV 127	TIME_ID.clone(),	×
UNCOV 128	Arc::new(array.dtype().clone()),	×
UNCOV 129	Some(temporal_metadata.clone().into()),	×
UNCOV 130	)),	×
UNCOV 131	array,	×
UNCOV 132	),	×
UNCOV 133	temporal_metadata,	×
UNCOV 134	}	×
UNCOV 135	}	×
136
137	/// Create a new `TemporalArray` holding Arrow spec compliant Timestamp data, with an
138	/// optional timezone.
139	///
140	/// # Panics
141	///
142	/// If `array` does not hold Primitive i64 data, the function will panic.
143	///
144	/// If the time_unit is days, the function will panic.
UNCOV 145	pub fn new_timestamp(array: ArrayRef, time_unit: TimeUnit, time_zone: Option<String>) -> Self {	×
UNCOV 146	assert_width!(i64, array);	×
147
UNCOV 148	let temporal_metadata = TemporalMetadata::Timestamp(time_unit, time_zone);	×
149
UNCOV 150	Self {	×
UNCOV 151	ext: ExtensionArray::new(	×
UNCOV 152	Arc::new(ExtDType::new(	×
UNCOV 153	TIMESTAMP_ID.clone(),	×
UNCOV 154	Arc::new(array.dtype().clone()),	×
UNCOV 155	Some(temporal_metadata.clone().into()),	×
UNCOV 156	)),	×
UNCOV 157	array,	×
UNCOV 158	),	×
UNCOV 159	temporal_metadata,	×
UNCOV 160	}	×
UNCOV 161	}	×
162	}
163
164	impl TemporalArray {
165	/// Access the underlying temporal values in the underlying ExtensionArray storage.
166	///
167	/// These values are to be interpreted based on the time unit and optional time-zone stored
168	/// in the TemporalMetadata.
UNCOV 169	pub fn temporal_values(&self) -> &ArrayRef {	×
UNCOV 170	self.ext.storage()	×
UNCOV 171	}	×
172
173	/// Retrieve the temporal metadata.
174	///
175	/// The metadata is used to provide semantic meaning to the temporal values Array, for example
176	/// to understand the granularity of the samples and if they have an associated timezone.
UNCOV 177	pub fn temporal_metadata(&self) -> &TemporalMetadata {	×
UNCOV 178	&self.temporal_metadata	×
UNCOV 179	}	×
180
181	/// Retrieve the extension DType associated with the underlying array.
UNCOV 182	pub fn ext_dtype(&self) -> Arc<ExtDType> {	×
UNCOV 183	self.ext.ext_dtype().clone()	×
UNCOV 184	}	×
185
186	/// Retrieve the DType of the array. This will be a `DType::Extension` variant.
UNCOV 187	pub fn dtype(&self) -> &DType {	×
UNCOV 188	self.ext.dtype()	×
UNCOV 189	}	×
190	}
191
192	impl From<TemporalArray> for ArrayRef {
UNCOV 193	fn from(value: TemporalArray) -> Self {	×
UNCOV 194	value.ext.into_array()	×
UNCOV 195	}	×
196	}
197
198	impl IntoArray for TemporalArray {
199	fn into_array(self) -> ArrayRef {	×
200	self.into()	×
201	}	×
202	}
203
204	impl TryFrom<ArrayRef> for TemporalArray {
205	type Error = VortexError;
206
207	/// Try to specialize a generic Vortex array as a TemporalArray.
208	///
209	/// # Errors
210	///
211	/// If the provided Array does not have `vortex.ext` encoding, an error will be returned.
212	///
213	/// If the provided Array does not have recognized ExtMetadata corresponding to one of the known
214	/// `TemporalMetadata` variants, an error is returned.
UNCOV 215	fn try_from(value: ArrayRef) -> Result<Self, Self::Error> {	×
UNCOV 216	let ext = value	×
UNCOV 217	.as_opt::<ExtensionVTable>()	×
UNCOV 218	.ok_or_else(\|\| vortex_err!("array must be an ExtensionArray"))?;	×
UNCOV 219	let temporal_metadata = TemporalMetadata::try_from(ext.ext_dtype())?;	×
UNCOV 220	Ok(Self {	×
UNCOV 221	ext: ext.clone(),	×
UNCOV 222	temporal_metadata,	×
UNCOV 223	})	×
UNCOV 224	}	×
225	}
226
227	// Conversions to/from ExtensionArray
228	impl From<&TemporalArray> for ExtensionArray {
229	fn from(value: &TemporalArray) -> Self {	×
230	value.ext.clone()	×
231	}	×
232	}
233
234	impl From<TemporalArray> for ExtensionArray {
UNCOV 235	fn from(value: TemporalArray) -> Self {	×
UNCOV 236	value.ext	×
UNCOV 237	}	×
238	}
239
240	impl TryFrom<ExtensionArray> for TemporalArray {
241	type Error = VortexError;
242
UNCOV 243	fn try_from(ext: ExtensionArray) -> Result<Self, Self::Error> {	×
UNCOV 244	let temporal_metadata = TemporalMetadata::try_from(ext.ext_dtype().as_ref())?;	×
UNCOV 245	Ok(Self {	×
UNCOV 246	ext,	×
UNCOV 247	temporal_metadata,	×
UNCOV 248	})	×
UNCOV 249	}	×
250	}

vortex-data / vortex / 16935267080

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous