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

vortex_array/arrow/compute/to_arrow/
list.rs
