vortex_array/arrays/primitive/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::Alignment;
5use vortex_buffer::Buffer;
6use vortex_dtype::DType;
7use vortex_dtype::PType;
8use vortex_dtype::match_each_native_ptype;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_ensure;
13use vortex_vector::primitive::PVector;
14
15use crate::ArrayRef;
16use crate::EmptyMetadata;
17use crate::arrays::PrimitiveArray;
18use crate::buffer::BufferHandle;
19use crate::executor::ExecutionCtx;
20use crate::serde::ArrayChildren;
21use crate::validity::Validity;
22use crate::vtable;
23use crate::vtable::ArrayVTableExt;
24use crate::vtable::NotSupported;
25use crate::vtable::VTable;
26use crate::vtable::ValidityVTableFromValidityHelper;
27
28mod array;
29mod canonical;
30mod operations;
31pub mod rules;
32mod validity;
33mod visitor;
34
35pub use rules::PrimitiveMaskedValidityRule;
36use vortex_vector::Vector;
37
38use crate::arrays::primitive::vtable::rules::RULES;
39use crate::vtable::ArrayId;
40use crate::vtable::ArrayVTable;
41
42vtable!(Primitive);
43
44impl VTable for PrimitiveVTable {
45    type Array = PrimitiveArray;
46
47    type Metadata = EmptyMetadata;
48
49    type ArrayVTable = Self;
50    type CanonicalVTable = Self;
51    type OperationsVTable = Self;
52    type ValidityVTable = ValidityVTableFromValidityHelper;
53    type VisitorVTable = Self;
54    type ComputeVTable = NotSupported;
55    type EncodeVTable = NotSupported;
56
57    fn id(&self) -> ArrayId {
58        ArrayId::new_ref("vortex.primitive")
59    }
60
61    fn encoding(_array: &Self::Array) -> ArrayVTable {
62        PrimitiveVTable.as_vtable()
63    }
64
65    fn metadata(_array: &PrimitiveArray) -> VortexResult<Self::Metadata> {
66        Ok(EmptyMetadata)
67    }
68
69    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
70        Ok(Some(vec![]))
71    }
72
73    fn deserialize(_buffer: &[u8]) -> VortexResult<Self::Metadata> {
74        Ok(EmptyMetadata)
75    }
76
77    fn build(
78        &self,
79        dtype: &DType,
80        len: usize,
81        _metadata: &Self::Metadata,
82        buffers: &[BufferHandle],
83        children: &dyn ArrayChildren,
84    ) -> VortexResult<PrimitiveArray> {
85        if buffers.len() != 1 {
86            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
87        }
88        let buffer = buffers[0].clone().try_to_bytes()?;
89
90        let validity = if children.is_empty() {
91            Validity::from(dtype.nullability())
92        } else if children.len() == 1 {
93            let validity = children.get(0, &Validity::DTYPE, len)?;
94            Validity::Array(validity)
95        } else {
96            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
97        };
98
99        let ptype = PType::try_from(dtype)?;
100
101        if !buffer.is_aligned(Alignment::new(ptype.byte_width())) {
102            vortex_bail!(
103                "Buffer is not aligned to {}-byte boundary",
104                ptype.byte_width()
105            );
106        }
107        if buffer.len() != ptype.byte_width() * len {
108            vortex_bail!(
109                "Buffer length {} does not match expected length {} for {}, {}",
110                buffer.len(),
111                ptype.byte_width() * len,
112                ptype.byte_width(),
113                len,
114            );
115        }
116
117        match_each_native_ptype!(ptype, |P| {
118            let buffer = Buffer::<P>::from_byte_buffer(buffer);
119            Ok(PrimitiveArray::new(buffer, validity))
120        })
121    }
122
123    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
124        Ok(match_each_native_ptype!(array.ptype(), |T| {
125            PVector::new(array.buffer::<T>(), array.validity_mask()).into()
126        }))
127    }
128
129    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
130        vortex_ensure!(
131            children.len() <= 1,
132            "PrimitiveArray can have at most 1 child (validity), got {}",
133            children.len()
134        );
135
136        array.validity = if children.is_empty() {
137            Validity::from(array.dtype().nullability())
138        } else {
139            Validity::Array(children.into_iter().next().vortex_expect("checked"))
140        };
141
142        Ok(())
143    }
144
145    fn reduce_parent(
146        array: &Self::Array,
147        parent: &ArrayRef,
148        child_idx: usize,
149    ) -> VortexResult<Option<ArrayRef>> {
150        RULES.evaluate(array, parent, child_idx)
151    }
152}
153
154#[derive(Debug)]
155pub struct PrimitiveVTable;