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

vortex_array/arrow/
record_batch.rs
