vortex_array/arrow/
record_batch.rs

1use arrow_array::RecordBatch;
2use arrow_array::cast::AsArray;
3use arrow_schema::{DataType, Schema};
4use vortex_error::{VortexError, VortexResult, vortex_err};
5
6use crate::arrays::StructArray;
7use crate::arrow::{FromArrowArray, IntoArrowArray};
8use crate::validity::Validity;
9use crate::{Array, ArrayRef, ToCanonical, TryIntoArray};
10
11impl TryIntoArray for RecordBatch {
12    fn try_into_array(self) -> VortexResult<ArrayRef> {
13        Ok(StructArray::try_new(
14            self.schema()
15                .fields()
16                .iter()
17                .map(|f| f.name().as_str().into())
18                .collect(),
19            self.columns()
20                .iter()
21                .zip(self.schema().fields())
22                .map(|(array, field)| ArrayRef::from_arrow(array.clone(), field.is_nullable()))
23                .collect(),
24            self.num_rows(),
25            Validity::NonNullable, // Must match FromArrowType<SchemaRef> for DType
26        )?
27        .into_array())
28    }
29}
30
31impl TryFrom<&dyn Array> for RecordBatch {
32    type Error = VortexError;
33
34    fn try_from(value: &dyn Array) -> VortexResult<Self> {
35        let struct_arr = value.to_struct().map_err(|err| {
36            vortex_err!("RecordBatch can only be constructed from a Vortex StructArray: {err}")
37        })?;
38
39        struct_arr.into_record_batch()
40    }
41}
42
43impl StructArray {
44    pub fn into_record_batch(self) -> VortexResult<RecordBatch> {
45        let array_ref = self.into_array().into_arrow_preferred()?;
46        Ok(RecordBatch::from(array_ref.as_struct()))
47    }
48
49    pub fn into_record_batch_with_schema(self, schema: &Schema) -> VortexResult<RecordBatch> {
50        let data_type = DataType::Struct(schema.fields.clone());
51        let array_ref = self.into_array().into_arrow(&data_type)?;
52        Ok(RecordBatch::from(array_ref.as_struct()))
53    }
54}
vortex_array/arrow/record_batch.rs

vortex_array/arrow/
record_batch.rs
