vortex_compute/arrow/
struct_.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use arrow_array::Array;
7use arrow_array::ArrayRef;
8use arrow_array::StructArray;
9use arrow_schema::Field;
10use arrow_schema::Fields;
11use vortex_error::VortexResult;
12use vortex_vector::Vector;
13use vortex_vector::VectorOps;
14use vortex_vector::struct_::StructVector;
15
16use crate::arrow::IntoArrow;
17use crate::arrow::IntoVector;
18use crate::arrow::nulls_to_mask;
19
20impl IntoArrow for StructVector {
21    type Output = StructArray;
22
23    fn into_arrow(self) -> VortexResult<Self::Output> {
24        let len = self.len();
25        let (fields, validity) = self.into_parts();
26        let arrow_fields = fields
27            .iter()
28            .map(|field| field.clone().into_arrow())
29            .collect::<Result<Vec<ArrayRef>, _>>()?;
30
31        // We need to make up the field names since vectors are unnamed, so we just use the field
32        // indices.
33        let fields = Fields::from(
34            (0..arrow_fields.len())
35                .map(|i| {
36                    Field::new(
37                        i.to_string(),
38                        arrow_fields[i].data_type().clone(),
39                        true, // Vectors are always nullable.
40                    )
41                })
42                .collect::<Vec<Field>>(),
43        );
44
45        // SAFETY: Since all of these components came from a valid `StructVector`, we know that all
46        // of the lengths of the vectors are correct. Additionally, all extra metadata is directly
47        // derived from the existing components so all invariants are upheld.
48        Ok(unsafe {
49            StructArray::new_unchecked_with_length(fields, arrow_fields, validity.into(), len)
50        })
51    }
52}
53
54impl IntoVector for &StructArray {
55    type Output = StructVector;
56
57    fn into_vector(self) -> VortexResult<Self::Output> {
58        let fields: Box<[Vector]> = self
59            .columns()
60            .iter()
61            .map(|col| col.as_ref().into_vector())
62            .collect::<Result<_, _>>()?;
63
64        let validity = nulls_to_mask(self.nulls(), self.len());
65
66        Ok(StructVector::new(Arc::new(fields), validity))
67    }
68}