vortex_array/array/list/compute/
to_arrow.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
use std::sync::Arc;

use arrow_array::ArrayRef;
use arrow_schema::{DataType, Field, FieldRef};
use vortex_dtype::PType;
use vortex_error::{vortex_bail, VortexResult};

use crate::array::{ListArray, ListEncoding};
use crate::arrow::IntoArrowArray;
use crate::compute::{try_cast, ToArrowFn};
use crate::variants::PrimitiveArrayTrait;
use crate::IntoArrayVariant;

impl ToArrowFn<ListArray> for ListEncoding {
    fn to_arrow(&self, array: &ListArray, data_type: &DataType) -> VortexResult<Option<ArrayRef>> {
        let (cast_ptype, element_dtype) = match data_type {
            DataType::List(field) => (PType::I32, field.data_type()),
            DataType::LargeList(field) => (PType::I64, field.data_type()),
            _ => {
                vortex_bail!("Unsupported data type: {data_type}");
            }
        };

        let offsets = array
            .offsets()
            .into_primitive()
            .map_err(|err| err.with_context("Failed to canonicalize offsets"))?;

        let arrow_offsets = try_cast(offsets, cast_ptype.into())
            .map_err(|err| err.with_context("Failed to cast offsets to PrimitiveArray"))?
            .into_primitive()?;

        let values = array.elements().into_arrow(element_dtype)?;

        let field_ref = FieldRef::new(Field::new_list_field(
            values.data_type().clone(),
            array.elements().dtype().nullability().into(),
        ));

        let nulls = array.logical_validity()?.to_null_buffer();

        Ok(Some(match arrow_offsets.ptype() {
            PType::I32 => Arc::new(arrow_array::ListArray::try_new(
                field_ref,
                arrow_offsets.buffer::<i32>().into_arrow_offset_buffer(),
                values,
                nulls,
            )?),
            PType::I64 => Arc::new(arrow_array::LargeListArray::try_new(
                field_ref,
                arrow_offsets.buffer::<i64>().into_arrow_offset_buffer(),
                values,
                nulls,
            )?),
            _ => vortex_bail!("Invalid offsets type {}", arrow_offsets.ptype()),
        }))
    }
}