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

vortex_array/arrow/compute/to_arrow/
list.rs
