vortex_array/arrays/list/compute/
to_arrow.rs

1use std::sync::Arc;
2
3use arrow_array::ArrayRef;
4use arrow_schema::{DataType, Field, FieldRef};
5use vortex_dtype::PType;
6use vortex_error::{VortexResult, vortex_bail};
7
8use crate::arrays::{ListArray, ListEncoding};
9use crate::arrow::IntoArrowArray;
10use crate::compute::{ToArrowFn, cast};
11use crate::variants::PrimitiveArrayTrait;
12use crate::{Array, ToCanonical};
13
14impl ToArrowFn<&ListArray> for ListEncoding {
15    fn to_arrow(&self, array: &ListArray, data_type: &DataType) -> VortexResult<Option<ArrayRef>> {
16        let (cast_ptype, element_dtype) = match data_type {
17            DataType::List(field) => (PType::I32, field.data_type()),
18            DataType::LargeList(field) => (PType::I64, field.data_type()),
19            _ => {
20                vortex_bail!("Unsupported data type: {data_type}");
21            }
22        };
23
24        let offsets = array
25            .offsets()
26            .to_primitive()
27            .map_err(|err| err.with_context("Failed to canonicalize offsets"))?;
28
29        let arrow_offsets = cast(&offsets, cast_ptype.into())
30            .map_err(|err| err.with_context("Failed to cast offsets to PrimitiveArray"))?
31            .to_primitive()?;
32
33        let values = array.elements().clone().into_arrow(element_dtype)?;
34
35        let field_ref = FieldRef::new(Field::new_list_field(
36            values.data_type().clone(),
37            array.elements().dtype().nullability().into(),
38        ));
39
40        let nulls = array.validity_mask()?.to_null_buffer();
41
42        Ok(Some(match arrow_offsets.ptype() {
43            PType::I32 => Arc::new(arrow_array::ListArray::try_new(
44                field_ref,
45                arrow_offsets.buffer::<i32>().into_arrow_offset_buffer(),
46                values,
47                nulls,
48            )?),
49            PType::I64 => Arc::new(arrow_array::LargeListArray::try_new(
50                field_ref,
51                arrow_offsets.buffer::<i64>().into_arrow_offset_buffer(),
52                values,
53                nulls,
54            )?),
55            _ => vortex_bail!("Invalid offsets type {}", arrow_offsets.ptype()),
56        }))
57    }
58}
vortex_array/arrays/list/compute/to_arrow.rs

vortex_array/arrays/list/compute/
to_arrow.rs
