vortex_array/arrays/primitive/vtable/
mod.rs

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

vortex_array/arrays/primitive/vtable/
mod.rs
