vortex_fastlanes/for/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::hash::Hash;
6use std::hash::Hasher;
7
8use vortex_array::Array;
9use vortex_array::ArrayEq;
10use vortex_array::ArrayHash;
11use vortex_array::ArrayId;
12use vortex_array::ArrayParts;
13use vortex_array::ArrayRef;
14use vortex_array::ArrayView;
15use vortex_array::EqMode;
16use vortex_array::ExecutionCtx;
17use vortex_array::ExecutionResult;
18use vortex_array::IntoArray;
19use vortex_array::arrays::PrimitiveArray;
20use vortex_array::buffer::BufferHandle;
21use vortex_array::dtype::DType;
22use vortex_array::scalar::Scalar;
23use vortex_array::scalar::ScalarValue;
24use vortex_array::serde::ArrayChildren;
25use vortex_array::smallvec::smallvec;
26use vortex_array::vtable::VTable;
27use vortex_array::vtable::ValidityVTableFromChild;
28use vortex_error::VortexResult;
29use vortex_error::vortex_bail;
30use vortex_error::vortex_ensure;
31use vortex_error::vortex_panic;
32use vortex_session::VortexSession;
33use vortex_session::registry::CachedId;
34
35use crate::FoRData;
36use crate::r#for::array::FoRArrayExt;
37use crate::r#for::array::FoRSlots;
38use crate::r#for::array::FoRSlotsView;
39use crate::r#for::array::for_decompress::decompress;
40use crate::r#for::vtable::rules::PARENT_RULES;
41
42mod kernels;
43mod operations;
44mod rules;
45mod slice;
46mod validity;
47
48/// A [`FoR`]-encoded Vortex array.
49pub type FoRArray = Array<FoR>;
50
51pub(crate) fn initialize(session: &VortexSession) {
52    kernels::initialize(session);
53}
54
55impl ArrayHash for FoRData {
56    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
57        self.reference.hash(state);
58    }
59}
60
61impl ArrayEq for FoRData {
62    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
63        self.reference == other.reference
64    }
65}
66
67impl VTable for FoR {
68    type TypedArrayData = FoRData;
69
70    type OperationsVTable = Self;
71    type ValidityVTable = ValidityVTableFromChild;
72
73    fn id(&self) -> ArrayId {
74        static ID: CachedId = CachedId::new("fastlanes.for");
75        *ID
76    }
77
78    fn validate(
79        &self,
80        data: &Self::TypedArrayData,
81        dtype: &DType,
82        len: usize,
83        slots: &[Option<ArrayRef>],
84    ) -> VortexResult<()> {
85        let encoded = FoRSlotsView::from_slots(slots).encoded;
86        validate_parts(encoded.dtype(), encoded.len(), &data.reference, dtype, len)
87    }
88
89    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
90        0
91    }
92
93    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
94        vortex_panic!("FoRArray buffer index {idx} out of bounds")
95    }
96
97    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
98        None
99    }
100
101    fn with_buffers(
102        &self,
103        array: ArrayView<'_, Self>,
104        buffers: &[BufferHandle],
105    ) -> VortexResult<ArrayParts<Self>> {
106        vortex_array::vtable::with_empty_buffers(self, array, buffers)
107    }
108
109    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
110        FoRSlots::NAMES[idx].to_string()
111    }
112
113    fn serialize(
114        array: ArrayView<'_, Self>,
115        _session: &VortexSession,
116    ) -> VortexResult<Option<Vec<u8>>> {
117        // Note that we **only** serialize the optional scalar value (not including the dtype).
118        Ok(Some(ScalarValue::to_proto_bytes(
119            array.reference_scalar().value(),
120        )))
121    }
122
123    fn deserialize(
124        &self,
125        dtype: &DType,
126        len: usize,
127        metadata: &[u8],
128        buffers: &[BufferHandle],
129        children: &dyn ArrayChildren,
130        session: &VortexSession,
131    ) -> VortexResult<ArrayParts<Self>> {
132        vortex_ensure!(
133            buffers.is_empty(),
134            "FoRArray expects 0 buffers, got {}",
135            buffers.len()
136        );
137        if children.len() != 1 {
138            vortex_bail!(
139                "Expected 1 child for FoR encoding, found {}",
140                children.len()
141            )
142        }
143
144        let scalar_value = ScalarValue::from_proto_bytes(metadata, dtype, session)?;
145        let reference = Scalar::try_new(dtype.clone(), scalar_value)?;
146        let encoded = children.get(0, dtype, len)?;
147        let slots = smallvec![Some(encoded)];
148
149        let data = FoRData::try_new(reference)?;
150        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
151    }
152
153    fn reduce_parent(
154        array: ArrayView<'_, Self>,
155        parent: &ArrayRef,
156        child_idx: usize,
157    ) -> VortexResult<Option<ArrayRef>> {
158        PARENT_RULES.evaluate(array, parent, child_idx)
159    }
160
161    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
162        Ok(ExecutionResult::done(decompress(&array, ctx)?.into_array()))
163    }
164}
165
166#[derive(Clone, Debug)]
167pub struct FoR;
168
169impl FoR {
170    /// Construct a new FoR array from an encoded array and a reference scalar.
171    pub fn try_new(encoded: ArrayRef, reference: Scalar) -> VortexResult<FoRArray> {
172        vortex_ensure!(!reference.is_null(), "Reference value cannot be null");
173        let dtype = reference
174            .dtype()
175            .with_nullability(encoded.dtype().nullability());
176        let reference = reference.cast(&dtype)?;
177        let len = encoded.len();
178        let data = FoRData::try_new(reference)?;
179        let slots = smallvec![Some(encoded)];
180        Array::try_from_parts(ArrayParts::new(FoR, dtype, len, data).with_slots(slots))
181    }
182
183    /// Encode a primitive array using Frame of Reference encoding.
184    pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult<FoRArray> {
185        FoRData::encode(array, ctx)
186    }
187}
188
189fn validate_parts(
190    encoded_dtype: &DType,
191    encoded_len: usize,
192    reference: &Scalar,
193    dtype: &DType,
194    len: usize,
195) -> VortexResult<()> {
196    vortex_ensure!(dtype.is_int(), "FoR requires an integer dtype, got {dtype}");
197    vortex_ensure!(
198        reference.dtype() == dtype,
199        "FoR reference dtype mismatch: expected {dtype}, got {}",
200        reference.dtype()
201    );
202    vortex_ensure!(
203        encoded_dtype == dtype,
204        "FoR encoded dtype mismatch: expected {dtype}, got {}",
205        encoded_dtype
206    );
207    vortex_ensure!(
208        encoded_len == len,
209        "FoR encoded length mismatch: expected {len}, got {}",
210        encoded_len
211    );
212    Ok(())
213}
214
215#[cfg(test)]
216mod tests {
217    use vortex_array::scalar::ScalarValue;
218    use vortex_array::test_harness::check_metadata;
219
220    #[cfg_attr(miri, ignore)]
221    #[test]
222    fn test_for_metadata() {
223        let metadata: Vec<u8> = ScalarValue::to_proto_bytes(Some(&ScalarValue::from(i64::MAX)));
224        check_metadata("for.metadata", &metadata);
225    }
226}
vortex_fastlanes/for/vtable/mod.rs

vortex_fastlanes/for/vtable/
mod.rs
