17093423531

Committed 20 Aug 2025 08:48AM UTC coverage: 87.929% (-0.02%) from 87.949%

Build # 17093423531

Build Type

Pull #4285

github

Committed by

web-flow

Commit Message

Merge 7c36e6b8a into b95a854c4

Pull Request Pull Request #4285: feat[vortex-array]: support converting to timestamp arrow arrays with tz

Coverage Stats

10 of 29 new or added lines in 1 file covered. (34.48%)

1 existing line in 1 file now uncovered.

56752 of 64543 relevant lines covered (87.93%)

623874.13 hits per line

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

57.35

/vortex-array/src/arrow/compute/to_arrow/temporal.rs

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

use std::sync::Arc;

use arrow_array::types::{
    ArrowTimestampType, Date32Type, Date64Type, Time32MillisecondType, Time32SecondType,
    Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType,
    TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
};
use arrow_array::{
    ArrayRef as ArrowArrayRef, ArrowPrimitiveType, PrimitiveArray as ArrowPrimitiveArray,
};
use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
use vortex_dtype::datetime::{TemporalMetadata, TimeUnit, is_temporal_ext_type};
use vortex_dtype::{DType, NativePType};
use vortex_error::{VortexExpect, VortexResult, vortex_bail};

use crate::arrays::{ExtensionVTable, TemporalArray};
use crate::arrow::array::ArrowArray;
use crate::arrow::compute::to_arrow::ToArrowArgs;
use crate::compute::{InvocationArgs, Kernel, Output, cast};
use crate::{Array as _, IntoArray, ToCanonical};

/// Implementation of `ToArrow` kernel for canonical Vortex arrays.
#[derive(Debug)]
pub(super) struct ToArrowTemporal;

impl Kernel for ToArrowTemporal {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;

        if !array
            .as_opt::<ExtensionVTable>()
            .is_some_and(|ext| is_temporal_ext_type(ext.ext_dtype().id()))
        {
            // This kernel only handles temporal arrays.
            return Ok(None);
        }
        let array = TemporalArray::try_from(array.to_array())
            .vortex_expect("Checked above that array is a temporal ExtensionArray");

        // Figure out the target Arrow type, or use the canonical type
        let arrow_type = arrow_type
            .cloned()
            .map(Ok)
            .unwrap_or_else(|| array.dtype().to_arrow_dtype())?;

        let arrow_array: ArrowArrayRef = match (array.temporal_metadata(), &arrow_type) {
            (TemporalMetadata::Date(TimeUnit::D), DataType::Date32) => {
                to_arrow_temporal::<Date32Type>(&array)
            }
            (TemporalMetadata::Date(TimeUnit::Ms), DataType::Date64) => {
                to_arrow_temporal::<Date64Type>(&array)
            }
            (TemporalMetadata::Time(TimeUnit::S), DataType::Time32(ArrowTimeUnit::Second)) => {
                to_arrow_temporal::<Time32SecondType>(&array)
            }
            (
                TemporalMetadata::Time(TimeUnit::Ms),
                DataType::Time32(ArrowTimeUnit::Millisecond),
            ) => to_arrow_temporal::<Time32MillisecondType>(&array),
            (
                TemporalMetadata::Time(TimeUnit::Us),
                DataType::Time64(ArrowTimeUnit::Microsecond),
            ) => to_arrow_temporal::<Time64MicrosecondType>(&array),

            (TemporalMetadata::Time(TimeUnit::Ns), DataType::Time64(ArrowTimeUnit::Nanosecond)) => {
                to_arrow_temporal::<Time64NanosecondType>(&array)
            }
            (
                TemporalMetadata::Timestamp(TimeUnit::S, _),
                DataType::Timestamp(ArrowTimeUnit::Second, arrow_tz),
            ) => to_arrow_timestamp::<TimestampSecondType>(&array, arrow_tz),
            (
                TemporalMetadata::Timestamp(TimeUnit::Ms, _),
                DataType::Timestamp(ArrowTimeUnit::Millisecond, arrow_tz),
            ) => to_arrow_timestamp::<TimestampMillisecondType>(&array, arrow_tz),
            (
                TemporalMetadata::Timestamp(TimeUnit::Us, _),
                DataType::Timestamp(ArrowTimeUnit::Microsecond, arrow_tz),
            ) => to_arrow_timestamp::<TimestampMicrosecondType>(&array, arrow_tz),
            (
                TemporalMetadata::Timestamp(TimeUnit::Ns, _),
                DataType::Timestamp(ArrowTimeUnit::Nanosecond, arrow_tz),
            ) => to_arrow_timestamp::<TimestampNanosecondType>(&array, arrow_tz),
            _ => vortex_bail!(
                "Cannot convert {} array to Arrow type {}",
                array.dtype(),
                arrow_type,
            ),
        }?;

        Ok(Some(
            ArrowArray::new(arrow_array, array.dtype().nullability())
                .into_array()
                .into(),
        ))
    }
}

fn to_arrow_temporal_inner<T: ArrowPrimitiveType>(
    array: &TemporalArray,
) -> VortexResult<ArrowPrimitiveArray<T>>
where
    T::Native: NativePType,
{
    let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability());
    let values = cast(array.temporal_values(), &values_dtype)?
        .to_primitive()?
        .into_buffer()
        .into_arrow_scalar_buffer();
    let nulls = array.temporal_values().validity_mask()?.to_null_buffer();
    Ok(ArrowPrimitiveArray::<T>::new(values, nulls))
}

fn to_arrow_temporal<T: ArrowPrimitiveType>(array: &TemporalArray) -> VortexResult<ArrowArrayRef>
where
    T::Native: NativePType,
{
    Ok(Arc::new(to_arrow_temporal_inner::<T>(array)?))
}

fn to_arrow_timestamp<T: ArrowTimestampType>(
    array: &TemporalArray,
    arrow_tz: &Option<Arc<str>>,
) -> VortexResult<ArrowArrayRef>
where
    T::Native: NativePType,
{
    let mut arrow_array = to_arrow_temporal_inner::<T>(array)?;
    if let Some(tz) = arrow_tz {
        arrow_array = arrow_array.with_timezone(tz.as_ref());
    }
    Ok(Arc::new(arrow_array))
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::sync::Arc;
5
6	use arrow_array::types::{
7	ArrowTimestampType, Date32Type, Date64Type, Time32MillisecondType, Time32SecondType,
8	Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType,
9	TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
10	};
11	use arrow_array::{
12	ArrayRef as ArrowArrayRef, ArrowPrimitiveType, PrimitiveArray as ArrowPrimitiveArray,
13	};
14	use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
15	use vortex_dtype::datetime::{TemporalMetadata, TimeUnit, is_temporal_ext_type};
16	use vortex_dtype::{DType, NativePType};
17	use vortex_error::{VortexExpect, VortexResult, vortex_bail};
18
19	use crate::arrays::{ExtensionVTable, TemporalArray};
20	use crate::arrow::array::ArrowArray;
21	use crate::arrow::compute::to_arrow::ToArrowArgs;
22	use crate::compute::{InvocationArgs, Kernel, Output, cast};
23	use crate::{Array as _, IntoArray, ToCanonical};
24
25	/// Implementation of `ToArrow` kernel for canonical Vortex arrays.
26	#[derive(Debug)]
27	pub(super) struct ToArrowTemporal;
28
29	impl Kernel for ToArrowTemporal {
30	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	228,047✔
31	let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;	228,047✔
32
33	if !array	228,047✔
34	.as_opt::<ExtensionVTable>()	228,047✔
35	.is_some_and(\|ext\| is_temporal_ext_type(ext.ext_dtype().id()))	228,047✔
36	{
37	// This kernel only handles temporal arrays.
38	return Ok(None);	227,918✔
39	}	129✔
40	let array = TemporalArray::try_from(array.to_array())	129✔
41	.vortex_expect("Checked above that array is a temporal ExtensionArray");	129✔
42
43	// Figure out the target Arrow type, or use the canonical type
44	let arrow_type = arrow_type	129✔
45	.cloned()	129✔
46	.map(Ok)	129✔
47	.unwrap_or_else(\|\| array.dtype().to_arrow_dtype())?;	129✔
48
49	let arrow_array: ArrowArrayRef = match (array.temporal_metadata(), &arrow_type) {	129✔
50	(TemporalMetadata::Date(TimeUnit::D), DataType::Date32) => {
51	to_arrow_temporal::<Date32Type>(&array)	129✔
52	}
53	(TemporalMetadata::Date(TimeUnit::Ms), DataType::Date64) => {
54	to_arrow_temporal::<Date64Type>(&array)	×
55	}
56	(TemporalMetadata::Time(TimeUnit::S), DataType::Time32(ArrowTimeUnit::Second)) => {
57	to_arrow_temporal::<Time32SecondType>(&array)	×
58	}
59	(
60	TemporalMetadata::Time(TimeUnit::Ms),
61	DataType::Time32(ArrowTimeUnit::Millisecond),
62	) => to_arrow_temporal::<Time32MillisecondType>(&array),	×
63	(
64	TemporalMetadata::Time(TimeUnit::Us),
65	DataType::Time64(ArrowTimeUnit::Microsecond),
66	) => to_arrow_temporal::<Time64MicrosecondType>(&array),	×
67
68	(TemporalMetadata::Time(TimeUnit::Ns), DataType::Time64(ArrowTimeUnit::Nanosecond)) => {
69	to_arrow_temporal::<Time64NanosecondType>(&array)	×
70	}
71	(
72	TemporalMetadata::Timestamp(TimeUnit::S, _),
NEW 73	DataType::Timestamp(ArrowTimeUnit::Second, arrow_tz),	×
NEW 74	) => to_arrow_timestamp::<TimestampSecondType>(&array, arrow_tz),	×
75	(
76	TemporalMetadata::Timestamp(TimeUnit::Ms, _),
NEW 77	DataType::Timestamp(ArrowTimeUnit::Millisecond, arrow_tz),	×
NEW 78	) => to_arrow_timestamp::<TimestampMillisecondType>(&array, arrow_tz),	×
79	(
80	TemporalMetadata::Timestamp(TimeUnit::Us, _),
NEW 81	DataType::Timestamp(ArrowTimeUnit::Microsecond, arrow_tz),	×
NEW 82	) => to_arrow_timestamp::<TimestampMicrosecondType>(&array, arrow_tz),	×
83	(
84	TemporalMetadata::Timestamp(TimeUnit::Ns, _),
NEW 85	DataType::Timestamp(ArrowTimeUnit::Nanosecond, arrow_tz),	×
NEW 86	) => to_arrow_timestamp::<TimestampNanosecondType>(&array, arrow_tz),	×
87	_ => vortex_bail!(	×
88	"Cannot convert {} array to Arrow type {}",	×
89	array.dtype(),	×
90	arrow_type,
91	),
92	}?;	×
93
94	Ok(Some(	129✔
95	ArrowArray::new(arrow_array, array.dtype().nullability())	129✔
96	.into_array()	129✔
97	.into(),	129✔
98	))	129✔
99	}	228,047✔
100	}
101
102	fn to_arrow_temporal_inner<T: ArrowPrimitiveType>(	129✔
103	array: &TemporalArray,	129✔
104	) -> VortexResult<ArrowPrimitiveArray<T>>	129✔
105	where	129✔
106	T::Native: NativePType,	129✔
107	{
108	let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability());	129✔
109	let values = cast(array.temporal_values(), &values_dtype)?	129✔
110	.to_primitive()?	129✔
111	.into_buffer()	129✔
112	.into_arrow_scalar_buffer();	129✔
113	let nulls = array.temporal_values().validity_mask()?.to_null_buffer();	129✔
114	Ok(ArrowPrimitiveArray::<T>::new(values, nulls))	129✔
115	}	129✔
116
117	fn to_arrow_temporal<T: ArrowPrimitiveType>(array: &TemporalArray) -> VortexResult<ArrowArrayRef>	129✔
118	where	129✔
119	T::Native: NativePType,	129✔
120	{
121	Ok(Arc::new(to_arrow_temporal_inner::<T>(array)?))	129✔
122	}	129✔
123
NEW 124	fn to_arrow_timestamp<T: ArrowTimestampType>(	×
NEW 125	array: &TemporalArray,	×
NEW 126	arrow_tz: &Option<Arc<str>>,	×
NEW 127	) -> VortexResult<ArrowArrayRef>	×
NEW 128	where	×
NEW 129	T::Native: NativePType,	×
130	{
NEW 131	let mut arrow_array = to_arrow_temporal_inner::<T>(array)?;	×
NEW 132	if let Some(tz) = arrow_tz {	×
NEW 133	arrow_array = arrow_array.with_timezone(tz.as_ref());	×
NEW 134	}	×
NEW 135	Ok(Arc::new(arrow_array))	×
UNCOV 136	}	×

vortex-data / vortex / 17093423531

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