• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

vortex-data / vortex / 16331938722

16 Jul 2025 10:49PM UTC coverage: 80.702% (-0.9%) from 81.557%
16331938722

push

github

web-flow
feat: build with stable rust (#3881)

120 of 173 new or added lines in 28 files covered. (69.36%)

174 existing lines in 102 files now uncovered.

41861 of 51871 relevant lines covered (80.7%)

157487.71 hits per line

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

88.82
/vortex-array/src/arrow/compute/to_arrow/mod.rs
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
/// Convert a Vortex array to an Arrow array with the encoding's preferred `DataType`.
24
///
25
/// For example, a `VarBinArray` will be converted to an Arrow `VarBin` array, instead of the
26
/// canonical `VarBinViewArray`.
27
///
28
/// Warning: do not use this to convert a Vortex [`crate::stream::ArrayStream`] since each array
29
/// may have a different preferred Arrow type. Use [`to_arrow`] instead.
30
pub fn to_arrow_preferred(array: &dyn Array) -> VortexResult<ArrowArrayRef> {
×
31
    to_arrow_opts(array, &ToArrowOptions { arrow_type: None })
×
32
}
×
33

34
/// Convert a Vortex array to an Arrow array of the given type.
35
pub fn to_arrow(array: &dyn Array, arrow_type: &DataType) -> VortexResult<ArrowArrayRef> {
974✔
36
    to_arrow_opts(
974✔
37
        array,
974✔
38
        &ToArrowOptions {
974✔
39
            arrow_type: Some(arrow_type.clone()),
974✔
40
        },
974✔
41
    )
42
}
974✔
43

44
pub fn to_arrow_opts(array: &dyn Array, options: &ToArrowOptions) -> VortexResult<ArrowArrayRef> {
35,905✔
45
    let arrow = TO_ARROW_FN
35,905✔
46
        .invoke(&InvocationArgs {
35,905✔
47
            inputs: &[array.into()],
35,905✔
48
            options,
35,905✔
49
        })?
35,905✔
50
        .unwrap_array()?
35,903✔
51
        .as_opt::<ArrowVTable>()
35,903✔
52
        .ok_or_else(|| vortex_err!("ToArrow compute kernels must return a Vortex ArrowArray"))?
35,903✔
53
        .inner()
35,903✔
54
        .clone();
35,903✔
55

56
    if let Some(arrow_type) = &options.arrow_type {
35,903✔
57
        if arrow.data_type() != arrow_type {
8,456✔
58
            vortex_bail!(
×
59
                "Arrow array type mismatch: expected {:?}, got {:?}",
×
60
                &options.arrow_type,
×
61
                arrow.data_type()
×
62
            );
63
        }
8,456✔
64
    }
27,447✔
65

66
    Ok(arrow)
35,903✔
67
}
35,905✔
68

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

74
impl Options for ToArrowOptions {
75
    fn as_any(&self) -> &dyn Any {
173,712✔
76
        self
173,712✔
77
    }
173,712✔
78
}
79

80
struct ToArrow;
81

82
impl ComputeFnVTable for ToArrow {
83
    fn invoke(
35,905✔
84
        &self,
35,905✔
85
        args: &InvocationArgs,
35,905✔
86
        kernels: &[ArcRef<dyn Kernel>],
35,905✔
87
    ) -> VortexResult<Output> {
35,905✔
88
        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;
35,905✔
89

90
        for kernel in kernels {
79,368✔
91
            if let Some(output) = kernel.invoke(args)? {
65,997✔
92
                return Ok(output);
22,532✔
93
            }
43,463✔
94
        }
95
        if let Some(output) = array.invoke(&TO_ARROW_FN, args)? {
13,371✔
96
            return Ok(output);
×
97
        }
13,371✔
98

99
        // Fall back to canonicalizing and then converting.
100
        if !array.is_canonical() {
13,371✔
101
            let canonical_array = array.to_canonical()?;
13,371✔
102
            let arrow_array = to_arrow_opts(
13,371✔
103
                canonical_array.as_ref(),
13,371✔
104
                &ToArrowOptions {
13,371✔
105
                    arrow_type: arrow_type.cloned(),
13,371✔
106
                },
13,371✔
UNCOV
107
            )?;
×
108
            return Ok(ArrowArray::new(arrow_array, array.dtype().nullability())
13,371✔
109
                .to_array()
13,371✔
110
                .into());
13,371✔
111
        }
×
112

113
        vortex_bail!(
×
114
            "Failed to convert array {} to Arrow {:?}",
×
115
            array.encoding_id(),
×
116
            arrow_type
117
        );
118
    }
35,905✔
119

120
    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
35,905✔
121
        let ToArrowArgs { array, arrow_type } = ToArrowArgs::try_from(args)?;
35,905✔
122
        Ok(arrow_type
35,905✔
123
            .map(|arrow_type| DType::from_arrow((arrow_type, array.dtype().nullability())))
35,905✔
124
            .unwrap_or_else(|| array.dtype().clone()))
35,905✔
125
    }
35,905✔
126

127
    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
35,905✔
128
        let ToArrowArgs { array, .. } = ToArrowArgs::try_from(args)?;
35,905✔
129
        Ok(array.len())
35,905✔
130
    }
35,905✔
131

132
    fn is_elementwise(&self) -> bool {
35,905✔
133
        false
35,905✔
134
    }
35,905✔
135
}
136

137
pub static TO_ARROW_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
3,124✔
138
    let compute = ComputeFn::new("to_arrow".into(), ArcRef::new_ref(&ToArrow));
3,124✔
139

140
    // Register the kernels we ship ourselves
141
    compute.register_kernel(ArcRef::new_ref(&canonical::ToArrowCanonical));
3,124✔
142
    compute.register_kernel(ArcRef::new_ref(&temporal::ToArrowTemporal));
3,124✔
143

144
    for kernel in inventory::iter::<ToArrowKernelRef> {
6,248✔
145
        compute.register_kernel(kernel.0.clone());
3,124✔
146
    }
3,124✔
147
    compute
3,124✔
148
});
3,124✔
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> {
173,712✔
159
        if value.inputs.len() != 1 {
173,712✔
160
            vortex_bail!("Expected 1 input, found {}", value.inputs.len());
×
161
        }
173,712✔
162
        let array = value.inputs[0]
173,712✔
163
            .array()
173,712✔
164
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
173,712✔
165
        let options = value
173,712✔
166
            .options
173,712✔
167
            .as_any()
173,712✔
168
            .downcast_ref::<ToArrowOptions>()
173,712✔
169
            .vortex_expect("Expected options to be ToArrowOptions");
173,712✔
170

171
        Ok(ToArrowArgs {
173,712✔
172
            array,
173,712✔
173
            arrow_type: options.arrow_type.as_ref(),
173,712✔
174
        })
173,712✔
175
    }
173,712✔
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>> {
14,986✔
200
        let inputs = ToArrowArgs::try_from(args)?;
14,986✔
201
        let Some(array) = inputs.array.as_opt::<V>() else {
14,986✔
202
            return Ok(None);
13,299✔
203
        };
204

205
        let Some(arrow_array) = V::to_arrow(&self.0, array, inputs.arrow_type)? else {
1,687✔
206
            return Ok(None);
72✔
207
        };
208

209
        Ok(Some(
1,615✔
210
            ArrowArray::new(arrow_array, array.dtype().nullability())
1,615✔
211
                .to_array()
1,615✔
212
                .into(),
1,615✔
213
        ))
1,615✔
214
    }
14,986✔
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
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc