vortex_array/arrow/compute/to_arrow/
list.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::{ArrayRef as ArrowArrayRef, GenericListArray, OffsetSizeTrait};
7use arrow_schema::{DataType, Field, FieldRef};
8use vortex_dtype::{DType, IntegerPType};
9use vortex_error::{VortexResult, vortex_bail};
10
11use crate::arrays::{ListArray, ListVTable, list_view_from_list};
12use crate::arrow::IntoArrowArray;
13use crate::arrow::compute::{ToArrowKernel, ToArrowKernelAdapter};
14use crate::compute::cast;
15use crate::{IntoArray, ToCanonical, register_kernel};
16
17impl ToArrowKernel for ListVTable {
18    fn to_arrow(
19        &self,
20        array: &ListArray,
21        arrow_type: Option<&DataType>,
22    ) -> VortexResult<Option<ArrowArrayRef>> {
23        match arrow_type {
24            None => {
25                // Default to a `ListArray` with `i32` offsets (preferred) when no `arrow_type` is
26                // specified.
27                list_array_to_arrow_list::<i32>(array, None)
28            }
29            Some(DataType::List(field)) => list_array_to_arrow_list::<i32>(array, Some(field)),
30            Some(DataType::LargeList(field)) => list_array_to_arrow_list::<i64>(array, Some(field)),
31            Some(dt @ DataType::ListView(_)) | Some(dt @ DataType::LargeListView(_)) => {
32                // Convert `ListArray` to `ListViewArray`, then use the canonical conversion.
33                let list_view = list_view_from_list(array.clone());
34                Ok(list_view.into_array().into_arrow(dt)?)
35            }
36            _ => vortex_bail!(
37                "Cannot convert `ListArray` to non-list Arrow type: {:?}",
38                arrow_type
39            ),
40        }
41        .map(Some)
42    }
43}
44
45register_kernel!(ToArrowKernelAdapter(ListVTable).lift());
46
47/// Converts a Vortex [`ListArray`] directly into an arrow [`GenericListArray`].
48fn list_array_to_arrow_list<O: IntegerPType + OffsetSizeTrait>(
49    array: &ListArray,
50    element: Option<&FieldRef>,
51) -> VortexResult<ArrowArrayRef> {
52    // First we cast the offsets and sizes into the specified width (determined by `O::PTYPE`).
53    let offsets_dtype = DType::Primitive(O::PTYPE, array.dtype().nullability());
54    let offsets = cast(array.offsets(), &offsets_dtype)
55        .map_err(|err| err.with_context(format!("Failed to cast offsets to {offsets_dtype}")))?
56        .to_primitive();
57
58    // Convert `offsets` and `validity` to Arrow buffers.
59    let arrow_offsets = offsets.buffer::<O>().into_arrow_offset_buffer();
60    let nulls = array.validity_mask().to_null_buffer();
61
62    // Convert the child `elements` array to Arrow.
63    let (elements, element_field) = {
64        if let Some(element) = element {
65            // Convert elements to the specific Arrow type the caller wants.
66            (
67                array.elements().clone().into_arrow(element.data_type())?,
68                element.clone(),
69            )
70        } else {
71            // Otherwise, convert into whatever Arrow prefers.
72            let elements = array.elements().clone().into_arrow_preferred()?;
73            let element_field = Arc::new(Field::new_list_field(
74                elements.data_type().clone(),
75                array.elements().dtype().is_nullable(),
76            ));
77            (elements, element_field)
78        }
79    };
80
81    Ok(Arc::new(GenericListArray::new(
82        element_field,
83        arrow_offsets,
84        elements,
85        nulls,
86    )))
87}
vortex_array/arrow/compute/to_arrow/list.rs

vortex_array/arrow/compute/to_arrow/
list.rs
