16810706267

Committed 07 Aug 2025 04:52PM UTC coverage: 84.907% (+0.05%) from 84.856%

Build # 16810706267

Build Type

Pull #4159

github

Committed by

web-flow

Commit Message

Merge 80c87ff6a into cc5d64403

Pull Request Pull Request #4159: chore: Bump msrv to 1.89

Run Details

107 of 157 new or added lines in 34 files covered. (68.15%)

2 existing lines in 2 files now uncovered.

50642 of 59644 relevant lines covered (84.91%)

567883.04 hits per line

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

88.74

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

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

mod canonical;
mod temporal;
mod varbin;

use std::any::Any;
use std::sync::LazyLock;

use arcref::ArcRef;
use arrow_array::ArrayRef as ArrowArrayRef;
use arrow_schema::DataType;
use vortex_dtype::DType;
use vortex_dtype::arrow::FromArrowType;
use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};

use crate::Array;
use crate::arrow::array::{ArrowArray, ArrowVTable};
use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
use crate::vtable::VTable;

static TO_ARROW_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("to_arrow".into(), ArcRef::new_ref(&ToArrow));

    // Register the kernels we ship ourselves
    compute.register_kernel(ArcRef::new_ref(&canonical::ToArrowCanonical));
    compute.register_kernel(ArcRef::new_ref(&temporal::ToArrowTemporal));

    for kernel in inventory::iter::<ToArrowKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// Convert a Vortex array to an Arrow array with the encoding's preferred `DataType`.
///
/// For example, a `VarBinArray` will be converted to an Arrow `VarBin` array, instead of the
/// canonical `VarBinViewArray`.
///
/// Warning: do not use this to convert a Vortex [`crate::stream::ArrayStream`] since each array
/// may have a different preferred Arrow type. Use [`to_arrow`] instead.
pub fn to_arrow_preferred(array: &dyn Array) -> VortexResult<ArrowArrayRef> {
    to_arrow_opts(array, &ToArrowOptions { arrow_type: None })
}

/// Convert a Vortex array to an Arrow array of the given type.
pub fn to_arrow(array: &dyn Array, arrow_type: &DataType) -> VortexResult<ArrowArrayRef> {
    to_arrow_opts(
        array,
        &ToArrowOptions {
            arrow_type: Some(arrow_type.clone()),
        },
    )
}

pub fn to_arrow_opts(array: &dyn Array, options: &ToArrowOptions) -> VortexResult<ArrowArrayRef> {
    let arrow = TO_ARROW_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into()],
            options,
        })?
        .unwrap_array()?
        .as_opt::<ArrowVTable>()
        .ok_or_else(|| vortex_err!("ToArrow compute kernels must return a Vortex ArrowArray"))?
        .inner()
        .clone();

    if let Some(arrow_type) = &options.arrow_type
        && arrow.data_type() != arrow_type
    {
        vortex_bail!(
            "Arrow array type mismatch: expected {:?}, got {:?}",
            &options.arrow_type,
            arrow.data_type()
        );
    }

    Ok(arrow)
}

pub struct ToArrowOptions {
    /// The Arrow data type to convert to, if specified.
    pub arrow_type: Option<DataType>,
}

impl Options for ToArrowOptions {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

struct ToArrow;

impl ComputeFnVTable for ToArrow {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;

        for kernel in kernels {
            if let Some(output) = kernel.invoke(args)? {
                return Ok(output);
            }
        }
        if let Some(output) = array.invoke(&TO_ARROW_FN, args)? {
            return Ok(output);
        }

        // Fall back to canonicalizing and then converting.
        if !array.is_canonical() {
            let canonical_array = array.to_canonical()?;
            let arrow_array = to_arrow_opts(
                canonical_array.as_ref(),
                &ToArrowOptions {
                    arrow_type: arrow_type.cloned(),
                },
            )?;
            return Ok(ArrowArray::new(arrow_array, array.dtype().nullability())
                .to_array()
                .into());
        }

        vortex_bail!(
            "Failed to convert array {} to Arrow {:?}",
            array.encoding_id(),
            arrow_type
        );
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;
        Ok(arrow_type
            .map(|arrow_type| DType::from_arrow((arrow_type, array.dtype().nullability())))
            .unwrap_or_else(|| array.dtype().clone()))
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        let ToArrowArgs { array, .. } = ToArrowArgs::try_from(args)?;
        Ok(array.len())
    }

    fn is_elementwise(&self) -> bool {
        false
    }
}

pub struct ToArrowArgs<'a> {
    array: &'a dyn Array,
    arrow_type: Option<&'a DataType>,
}

impl<'a> TryFrom<&InvocationArgs<'a>> for ToArrowArgs<'a> {
    type Error = VortexError;

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 1 {
            vortex_bail!("Expected 1 input, found {}", value.inputs.len());
        }
        let array = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
        let options = value
            .options
            .as_any()
            .downcast_ref::<ToArrowOptions>()
            .vortex_expect("Expected options to be ToArrowOptions");

        Ok(ToArrowArgs {
            array,
            arrow_type: options.arrow_type.as_ref(),
        })
    }
}

pub struct ToArrowKernelRef(pub ArcRef<dyn Kernel>);
inventory::collect!(ToArrowKernelRef);

pub trait ToArrowKernel: VTable {
    fn to_arrow(
        &self,
        arr: &Self::Array,
        arrow_type: Option<&DataType>,
    ) -> VortexResult<Option<ArrowArrayRef>>;
}

#[derive(Debug)]
pub struct ToArrowKernelAdapter<V: VTable>(pub V);

impl<V: VTable + ToArrowKernel> ToArrowKernelAdapter<V> {
    pub const fn lift(&'static self) -> ToArrowKernelRef {
        ToArrowKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + ToArrowKernel> Kernel for ToArrowKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let inputs = ToArrowArgs::try_from(args)?;
        let Some(array) = inputs.array.as_opt::<V>() else {
            return Ok(None);
        };

        let Some(arrow_array) = V::to_arrow(&self.0, array, inputs.arrow_type)? else {
            return Ok(None);
        };

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

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_array::types::Int32Type;
    use arrow_array::{ArrayRef, PrimitiveArray, StringViewArray, StructArray};
    use arrow_buffer::NullBuffer;

    use super::to_arrow;
    use crate::{IntoArray, arrays};

    #[test]
    fn test_to_arrow() {
        let array = arrays::StructArray::from_fields(
            vec![
                (
                    "a",
                    arrays::PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2)])
                        .into_array(),
                ),
                (
                    "b",
                    arrays::VarBinViewArray::from_iter_str(vec!["a", "b", "c"]).into_array(),
                ),
            ]
            .as_slice(),
        )
        .unwrap();

        let arrow_array: ArrayRef = Arc::new(
            StructArray::try_from(vec![
                (
                    "a",
                    Arc::new(PrimitiveArray::<Int32Type>::from_iter_values_with_nulls(
                        vec![1, 0, 2],
                        Some(NullBuffer::from(vec![true, false, true])),
                    )) as ArrayRef,
                ),
                (
                    "b",
                    Arc::new(StringViewArray::from(vec![Some("a"), Some("b"), Some("c")])),
                ),
            ])
            .unwrap(),
        );

        assert_eq!(
            &to_arrow(array.as_ref(), &array.dtype().to_arrow_dtype().unwrap()).unwrap(),
            &arrow_array
        );
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	mod canonical;
5	mod temporal;
6	mod varbin;
7
8	use std::any::Any;
9	use std::sync::LazyLock;
10
11	use arcref::ArcRef;
12	use arrow_array::ArrayRef as ArrowArrayRef;
13	use arrow_schema::DataType;
14	use vortex_dtype::DType;
15	use vortex_dtype::arrow::FromArrowType;
16	use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
17
18	use crate::Array;
19	use crate::arrow::array::{ArrowArray, ArrowVTable};
20	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
21	use crate::vtable::VTable;
22
23	static TO_ARROW_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	10,290✔
24	let compute = ComputeFn::new("to_arrow".into(), ArcRef::new_ref(&ToArrow));	10,290✔
25
26	// Register the kernels we ship ourselves
27	compute.register_kernel(ArcRef::new_ref(&canonical::ToArrowCanonical));	10,290✔
28	compute.register_kernel(ArcRef::new_ref(&temporal::ToArrowTemporal));	10,290✔
29
30	for kernel in inventory::iter::<ToArrowKernelRef> {	20,580✔
31	compute.register_kernel(kernel.0.clone());	10,290✔
32	}	10,290✔
33	compute	10,290✔
34	});	10,290✔
35
36	/// Convert a Vortex array to an Arrow array with the encoding's preferred `DataType`.
37	///
38	/// For example, a `VarBinArray` will be converted to an Arrow `VarBin` array, instead of the
39	/// canonical `VarBinViewArray`.
40	///
41	/// Warning: do not use this to convert a Vortex [`crate::stream::ArrayStream`] since each array
42	/// may have a different preferred Arrow type. Use [`to_arrow`] instead.
43	pub fn to_arrow_preferred(array: &dyn Array) -> VortexResult<ArrowArrayRef> {	×
44	to_arrow_opts(array, &ToArrowOptions { arrow_type: None })	×
45	}	×
46
47	/// Convert a Vortex array to an Arrow array of the given type.
48	pub fn to_arrow(array: &dyn Array, arrow_type: &DataType) -> VortexResult<ArrowArrayRef> {	1,436✔
49	to_arrow_opts(	1,436✔
50	array,	1,436✔
51	&ToArrowOptions {	1,436✔
52	arrow_type: Some(arrow_type.clone()),	1,436✔
53	},	1,436✔
54	)
55	}	1,436✔
56
57	pub fn to_arrow_opts(array: &dyn Array, options: &ToArrowOptions) -> VortexResult<ArrowArrayRef> {	517,591✔
58	let arrow = TO_ARROW_FN	517,591✔
59	.invoke(&InvocationArgs {	517,591✔
60	inputs: &[array.into()],	517,591✔
61	options,	517,591✔
62	})?	517,591✔
63	.unwrap_array()?	517,589✔
64	.as_opt::<ArrowVTable>()	517,589✔
65	.ok_or_else(\|\| vortex_err!("ToArrow compute kernels must return a Vortex ArrowArray"))?	517,589✔
66	.inner()	517,589✔
67	.clone();	517,589✔
68
69	if let Some(arrow_type) = &options.arrow_type	517,589✔
70	&& arrow.data_type() != arrow_type	19,314✔
71	{
NEW 72	vortex_bail!(	×
NEW 73	"Arrow array type mismatch: expected {:?}, got {:?}",	×
NEW 74	&options.arrow_type,	×
NEW 75	arrow.data_type()	×
76	);
77	}	517,589✔
78
79	Ok(arrow)	517,589✔
80	}	517,591✔
81
82	pub struct ToArrowOptions {
83	/// The Arrow data type to convert to, if specified.
84	pub arrow_type: Option<DataType>,
85	}
86
87	impl Options for ToArrowOptions {
88	fn as_any(&self) -> &dyn Any {	2,524,094✔
89	self	2,524,094✔
90	}	2,524,094✔
91	}
92
93	struct ToArrow;
94
95	impl ComputeFnVTable for ToArrow {
96	fn invoke(	517,591✔
97	&self,	517,591✔
98	args: &InvocationArgs,	517,591✔
99	kernels: &[ArcRef<dyn Kernel>],	517,591✔
100	) -> VortexResult<Output> {	517,591✔
101	let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;	517,591✔
102
103	for kernel in kernels {	1,195,413✔
104	if let Some(output) = kernel.invoke(args)? {	971,321✔
105	return Ok(output);	293,497✔
106	}	677,822✔
107	}
108	if let Some(output) = array.invoke(&TO_ARROW_FN, args)? {	224,092✔
109	return Ok(output);	×
110	}	224,092✔
111
112	// Fall back to canonicalizing and then converting.
113	if !array.is_canonical() {	224,092✔
114	let canonical_array = array.to_canonical()?;	224,092✔
115	let arrow_array = to_arrow_opts(	224,092✔
116	canonical_array.as_ref(),	224,092✔
117	&ToArrowOptions {	224,092✔
118	arrow_type: arrow_type.cloned(),	224,092✔
119	},	224,092✔
120	)?;	×
121	return Ok(ArrowArray::new(arrow_array, array.dtype().nullability())	224,092✔
122	.to_array()	224,092✔
123	.into());	224,092✔
124	}	×
125
126	vortex_bail!(	×
127	"Failed to convert array {} to Arrow {:?}",	×
128	array.encoding_id(),	×
129	arrow_type
130	);
131	}	517,591✔
132
133	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	517,591✔
134	let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;	517,591✔
135	Ok(arrow_type	517,591✔
136	.map(\|arrow_type\| DType::from_arrow((arrow_type, array.dtype().nullability())))	517,591✔
137	.unwrap_or_else(\|\| array.dtype().clone()))	517,591✔
138	}	517,591✔
139
140	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	517,591✔
141	let ToArrowArgs { array, .. } = ToArrowArgs::try_from(args)?;	517,591✔
142	Ok(array.len())	517,591✔
143	}	517,591✔
144
145	fn is_elementwise(&self) -> bool {	517,591✔
146	false	517,591✔
147	}	517,591✔
148	}
149
150	pub struct ToArrowArgs<'a> {
151	array: &'a dyn Array,
152	arrow_type: Option<&'a DataType>,
153	}
154
155	impl<'a> TryFrom<&InvocationArgs<'a>> for ToArrowArgs<'a> {
156	type Error = VortexError;
157
158	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	2,524,094✔
159	if value.inputs.len() != 1 {	2,524,094✔
160	vortex_bail!("Expected 1 input, found {}", value.inputs.len());	×
161	}	2,524,094✔
162	let array = value.inputs[0]	2,524,094✔
163	.array()	2,524,094✔
164	.ok_or_else(\|\| vortex_err!("Expected input 0 to be an array"))?;	2,524,094✔
165	let options = value	2,524,094✔
166	.options	2,524,094✔
167	.as_any()	2,524,094✔
168	.downcast_ref::<ToArrowOptions>()	2,524,094✔
169	.vortex_expect("Expected options to be ToArrowOptions");	2,524,094✔
170
171	Ok(ToArrowArgs {	2,524,094✔
172	array,	2,524,094✔
173	arrow_type: options.arrow_type.as_ref(),	2,524,094✔
174	})	2,524,094✔
175	}	2,524,094✔
176	}
177
178	pub struct ToArrowKernelRef(pub ArcRef<dyn Kernel>);
179	inventory::collect!(ToArrowKernelRef);
180
181	pub trait ToArrowKernel: VTable {
182	fn to_arrow(
183	&self,
184	arr: &Self::Array,
185	arrow_type: Option<&DataType>,
186	) -> VortexResult<Option<ArrowArrayRef>>;
187	}
188
189	#[derive(Debug)]
190	pub struct ToArrowKernelAdapter<V: VTable>(pub V);
191
192	impl<V: VTable + ToArrowKernel> ToArrowKernelAdapter<V> {
193	pub const fn lift(&'static self) -> ToArrowKernelRef {	×
194	ToArrowKernelRef(ArcRef::new_ref(self))	×
195	}	×
196	}
197
198	impl<V: VTable + ToArrowKernel> Kernel for ToArrowKernelAdapter<V> {
199	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	226,775✔
200	let inputs = ToArrowArgs::try_from(args)?;	226,775✔
201	let Some(array) = inputs.array.as_opt::<V>() else {	226,775✔
202	return Ok(None);	224,014✔
203	};
204
205	let Some(arrow_array) = V::to_arrow(&self.0, array, inputs.arrow_type)? else {	2,761✔
206	return Ok(None);	78✔
207	};
208
209	Ok(Some(	2,683✔
210	ArrowArray::new(arrow_array, array.dtype().nullability())	2,683✔
211	.to_array()	2,683✔
212	.into(),	2,683✔
213	))	2,683✔
214	}	226,775✔
215	}
216
217	#[cfg(test)]
218	mod tests {
219	use std::sync::Arc;
220
221	use arrow_array::types::Int32Type;
222	use arrow_array::{ArrayRef, PrimitiveArray, StringViewArray, StructArray};
223	use arrow_buffer::NullBuffer;
224
225	use super::to_arrow;
226	use crate::{IntoArray, arrays};
227
228	#[test]
229	fn test_to_arrow() {	1✔
230	let array = arrays::StructArray::from_fields(	1✔
231	vec![	1✔
232	(	1✔
233	"a",	1✔
234	arrays::PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2)])	1✔
235	.into_array(),	1✔
236	),	1✔
237	(	1✔
238	"b",	1✔
239	arrays::VarBinViewArray::from_iter_str(vec!["a", "b", "c"]).into_array(),	1✔
240	),	1✔
241	]	1✔
242	.as_slice(),	1✔
243	)
244	.unwrap();	1✔
245
246	let arrow_array: ArrayRef = Arc::new(	1✔
247	StructArray::try_from(vec![	1✔
248	(	1✔
249	"a",	1✔
250	Arc::new(PrimitiveArray::<Int32Type>::from_iter_values_with_nulls(	1✔
251	vec![1, 0, 2],	1✔
252	Some(NullBuffer::from(vec![true, false, true])),	1✔
253	)) as ArrayRef,	1✔
254	),	1✔
255	(	1✔
256	"b",	1✔
257	Arc::new(StringViewArray::from(vec![Some("a"), Some("b"), Some("c")])),	1✔
258	),	1✔
259	])	1✔
260	.unwrap(),	1✔
261	);	1✔
262
263	assert_eq!(	1✔
264	&to_arrow(array.as_ref(), &array.dtype().to_arrow_dtype().unwrap()).unwrap(),	1✔
265	&arrow_array	1✔
266	);
267	}	1✔
268	}

vortex-data / vortex / 16810706267

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