vortex_array/arrow/
record_batch.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::RecordBatch;
5use arrow_array::cast::AsArray;
6use arrow_schema::DataType;
7use arrow_schema::Schema;
8use vortex_error::VortexResult;
9
10use crate::LEGACY_SESSION;
11use crate::VortexSessionExecute;
12use crate::array::IntoArray;
13use crate::arrays::StructArray;
14use crate::arrow::ArrowArrayExecutor;
15
16impl StructArray {
17    pub fn into_record_batch_with_schema(
18        self,
19        schema: impl AsRef<Schema>,
20    ) -> VortexResult<RecordBatch> {
21        let data_type = DataType::Struct(schema.as_ref().fields.clone());
22        let array_ref = self
23            .into_array()
24            .execute_arrow(Some(&data_type), &mut LEGACY_SESSION.create_execution_ctx())?;
25        Ok(RecordBatch::from(array_ref.as_struct()))
26    }
27}
28
29#[cfg(test)]
30mod tests {
31    use std::sync::Arc;
32
33    use arrow_schema::DataType;
34    use arrow_schema::Field;
35    use arrow_schema::FieldRef;
36    use arrow_schema::Schema;
37
38    use crate::arrow::record_batch::StructArray;
39    use crate::builders::ArrayBuilder;
40    use crate::builders::ListBuilder;
41    use crate::dtype::DType;
42    use crate::dtype::Nullability;
43    use crate::dtype::PType;
44    use crate::scalar::Scalar;
45
46    #[test]
47    fn test_into_rb_with_schema() {
48        let mut xs = ListBuilder::<u32>::new(
49            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
50            Nullability::Nullable,
51        );
52
53        xs.append_scalar(&Scalar::list(
54            xs.element_dtype().clone(),
55            vec![1i32.into(), 2i32.into(), 3i32.into()],
56            Nullability::Nullable,
57        ))
58        .unwrap();
59        xs.append_null();
60        xs.append_zero();
61
62        let xs = xs.finish();
63
64        let array = StructArray::from_fields(&[("xs", xs)]).unwrap();
65
66        // Explicitly request a conversion to LargeListView type instead of the preferred type.
67        let arrow_schema = Arc::new(Schema::new(vec![Field::new(
68            "xs",
69            DataType::LargeListView(FieldRef::new(Field::new_list_field(DataType::Int32, false))),
70            true,
71        )]));
72        let rb = array.into_record_batch_with_schema(arrow_schema).unwrap();
73
74        let xs = rb.column(0);
75        assert_eq!(
76            xs.data_type(),
77            &DataType::LargeListView(FieldRef::new(Field::new_list_field(DataType::Int32, false)))
78        );
79    }
80}
vortex_array/arrow/record_batch.rs

vortex_array/arrow/
record_batch.rs
