17100981141

Committed 20 Aug 2025 02:12PM UTC coverage: 88.059% (+0.3%) from 87.714%

Build # 17100981141

Build Type

Pull #4285

github

Committed by

web-flow

Commit Message

Merge ab585a447 into a014ab253

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

Coverage Stats

3 of 3 new or added lines in 1 file covered. (100.0%)

11 existing lines in 1 file now uncovered.

56845 of 64553 relevant lines covered (88.06%)

623884.71 hits per line

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

71.11

/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::{
    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, None),
            ) => to_arrow_temporal::<TimestampSecondType>(&array),
            (
                TemporalMetadata::Timestamp(TimeUnit::Ms, _),
                DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
            ) => to_arrow_temporal::<TimestampMillisecondType>(&array),
            (
                TemporalMetadata::Timestamp(TimeUnit::Us, _),
                DataType::Timestamp(ArrowTimeUnit::Microsecond, None),
            ) => to_arrow_temporal::<TimestampMicrosecondType>(&array),
            (
                TemporalMetadata::Timestamp(TimeUnit::Ns, _),
                DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
            ) => to_arrow_temporal::<TimestampNanosecondType>(&array),
            _ => 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<T: ArrowPrimitiveType>(array: &TemporalArray) -> VortexResult<ArrowArrayRef>
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(Arc::new(ArrowPrimitiveArray::<T>::new(values, nulls)))
}

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	Date32Type, Date64Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
8	Time64NanosecondType, TimestampMicrosecondType, TimestampMillisecondType,
9	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,059✔
31	let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;	228,059✔
32
33	if !array	228,059✔
34	.as_opt::<ExtensionVTable>()	228,059✔
35	.is_some_and(\|ext\| is_temporal_ext_type(ext.ext_dtype().id()))	228,059✔
36	{
37	// This kernel only handles temporal arrays.
38	return Ok(None);	227,930✔
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) => {
UNCOV 54	to_arrow_temporal::<Date64Type>(&array)	×
55	}
56	(TemporalMetadata::Time(TimeUnit::S), DataType::Time32(ArrowTimeUnit::Second)) => {
UNCOV 57	to_arrow_temporal::<Time32SecondType>(&array)	×
58	}
59	(
60	TemporalMetadata::Time(TimeUnit::Ms),
61	DataType::Time32(ArrowTimeUnit::Millisecond),
UNCOV 62	) => to_arrow_temporal::<Time32MillisecondType>(&array),	×
63	(
64	TemporalMetadata::Time(TimeUnit::Us),
65	DataType::Time64(ArrowTimeUnit::Microsecond),
UNCOV 66	) => to_arrow_temporal::<Time64MicrosecondType>(&array),	×
67
68	(TemporalMetadata::Time(TimeUnit::Ns), DataType::Time64(ArrowTimeUnit::Nanosecond)) => {
UNCOV 69	to_arrow_temporal::<Time64NanosecondType>(&array)	×
70	}
71	(
72	TemporalMetadata::Timestamp(TimeUnit::S, _),
73	DataType::Timestamp(ArrowTimeUnit::Second, None),
UNCOV 74	) => to_arrow_temporal::<TimestampSecondType>(&array),	×
75	(
76	TemporalMetadata::Timestamp(TimeUnit::Ms, _),
77	DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
UNCOV 78	) => to_arrow_temporal::<TimestampMillisecondType>(&array),	×
79	(
80	TemporalMetadata::Timestamp(TimeUnit::Us, _),
81	DataType::Timestamp(ArrowTimeUnit::Microsecond, None),
UNCOV 82	) => to_arrow_temporal::<TimestampMicrosecondType>(&array),	×
83	(
84	TemporalMetadata::Timestamp(TimeUnit::Ns, _),
85	DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
86	) => to_arrow_temporal::<TimestampNanosecondType>(&array),	×
87	_ => vortex_bail!(	×
UNCOV 88	"Cannot convert {} array to Arrow type {}",	×
UNCOV 89	array.dtype(),	×
90	arrow_type,
91	),
UNCOV 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,059✔
100	}
101
102	fn to_arrow_temporal<T: ArrowPrimitiveType>(array: &TemporalArray) -> VortexResult<ArrowArrayRef>	129✔
103	where	129✔
104	T::Native: NativePType,	129✔
105	{
106	let values_dtype = DType::Primitive(T::Native::PTYPE, array.dtype().nullability());	129✔
107	let values = cast(array.temporal_values(), &values_dtype)?	129✔
108	.to_primitive()?	129✔
109	.into_buffer()	129✔
110	.into_arrow_scalar_buffer();	129✔
111	let nulls = array.temporal_values().validity_mask()?.to_null_buffer();	129✔
112
113	Ok(Arc::new(ArrowPrimitiveArray::<T>::new(values, nulls)))	129✔
114	}	129✔

vortex-data / vortex / 17100981141

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