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, Schema};
7use vortex_error::{VortexError, VortexResult, vortex_err};
8
9use crate::arrays::StructArray;
10use crate::arrow::compute::{to_arrow, to_arrow_preferred};
11use crate::{Array, ToCanonical};
12
13impl TryFrom<&dyn Array> for RecordBatch {
14    type Error = VortexError;
15
16    fn try_from(value: &dyn Array) -> VortexResult<Self> {
17        let struct_arr = value.to_struct().map_err(|err| {
18            vortex_err!("RecordBatch can only be constructed from a Vortex StructArray: {err}")
19        })?;
20
21        struct_arr.into_record_batch()
22    }
23}
24
25impl StructArray {
26    pub fn into_record_batch(self) -> VortexResult<RecordBatch> {
27        let array_ref = to_arrow_preferred(self.as_ref())?;
28        Ok(RecordBatch::from(array_ref.as_struct()))
29    }
30
31    pub fn into_record_batch_with_schema(self, schema: &Schema) -> VortexResult<RecordBatch> {
32        let data_type = DataType::Struct(schema.fields.clone());
33        let array_ref = to_arrow(self.as_ref(), &data_type)?;
34        Ok(RecordBatch::from(array_ref.as_struct()))
35    }
36}