vortex-array 0.68.0

Vortex in memory columnar data format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use arrow_array::RecordBatch;
use arrow_array::cast::AsArray;
use arrow_schema::DataType;
use arrow_schema::Schema;
use vortex_error::VortexError;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;

use crate::ArrayRef;
use crate::Canonical;
use crate::LEGACY_SESSION;
use crate::VortexSessionExecute;
use crate::array::IntoArray;
use crate::arrays::StructArray;
use crate::arrow::ArrowArrayExecutor;

impl TryFrom<&ArrayRef> for RecordBatch {
    type Error = VortexError;

    fn try_from(value: &ArrayRef) -> VortexResult<Self> {
        let Canonical::Struct(struct_array) = value.to_canonical()? else {
            vortex_bail!("RecordBatch can only be constructed from ")
        };

        vortex_ensure!(
            struct_array.all_valid()?,
            "RecordBatch can only be constructed from StructArray with no nulls"
        );

        let data_type = struct_array.dtype().to_arrow_dtype()?;
        let array_ref = struct_array
            .into_array()
            .execute_arrow(Some(&data_type), &mut LEGACY_SESSION.create_execution_ctx())?;
        Ok(RecordBatch::from(array_ref.as_struct()))
    }
}

impl StructArray {
    pub fn into_record_batch_with_schema(
        self,
        schema: impl AsRef<Schema>,
    ) -> VortexResult<RecordBatch> {
        let data_type = DataType::Struct(schema.as_ref().fields.clone());
        let array_ref = self
            .into_array()
            .execute_arrow(Some(&data_type), &mut LEGACY_SESSION.create_execution_ctx())?;
        Ok(RecordBatch::from(array_ref.as_struct()))
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_schema::DataType;
    use arrow_schema::Field;
    use arrow_schema::FieldRef;
    use arrow_schema::Schema;

    use crate::arrow::record_batch::StructArray;
    use crate::builders::ArrayBuilder;
    use crate::builders::ListBuilder;
    use crate::dtype::DType;
    use crate::dtype::Nullability;
    use crate::dtype::PType;
    use crate::scalar::Scalar;

    #[test]
    fn test_into_rb_with_schema() {
        let mut xs = ListBuilder::<u32>::new(
            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
            Nullability::Nullable,
        );

        xs.append_scalar(&Scalar::list(
            xs.element_dtype().clone(),
            vec![1i32.into(), 2i32.into(), 3i32.into()],
            Nullability::Nullable,
        ))
        .unwrap();
        xs.append_null();
        xs.append_zero();

        let xs = xs.finish();

        let array = StructArray::from_fields(&[("xs", xs)]).unwrap();

        // Explicitly request a conversion to LargeListView type instead of the preferred type.
        let arrow_schema = Arc::new(Schema::new(vec![Field::new(
            "xs",
            DataType::LargeListView(FieldRef::new(Field::new_list_field(DataType::Int32, false))),
            true,
        )]));
        let rb = array.into_record_batch_with_schema(arrow_schema).unwrap();

        let xs = rb.column(0);
        assert_eq!(
            xs.data_type(),
            &DataType::LargeListView(FieldRef::new(Field::new_list_field(DataType::Int32, false)))
        );
    }
}