Skip to main content

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;
5
6use vortex_array::ArrayRef;
7use vortex_array::ExecutionCtx;
8use vortex_array::IntoArray;
9use vortex_array::buffer::BufferHandle;
10use vortex_array::dtype::DType;
11use vortex_array::scalar::Scalar;
12use vortex_array::scalar::ScalarValue;
13use vortex_array::serde::ArrayChildren;
14use vortex_array::vtable;
15use vortex_array::vtable::ArrayId;
16use vortex_array::vtable::VTable;
17use vortex_array::vtable::ValidityVTableFromChild;
18use vortex_error::VortexResult;
19use vortex_error::vortex_bail;
20use vortex_error::vortex_ensure;
21use vortex_session::VortexSession;
22
23use crate::FoRArray;
24use crate::r#for::array::for_decompress::decompress;
25use crate::r#for::vtable::kernels::PARENT_KERNELS;
26use crate::r#for::vtable::rules::PARENT_RULES;
27
28mod array;
29mod kernels;
30mod operations;
31mod rules;
32mod slice;
33mod validity;
34mod visitor;
35
36vtable!(FoR);
37
38impl VTable for FoRVTable {
39    type Array = FoRArray;
40
41    type Metadata = Scalar;
42
43    type ArrayVTable = Self;
44    type OperationsVTable = Self;
45    type ValidityVTable = ValidityVTableFromChild;
46    type VisitorVTable = Self;
47
48    fn id(_array: &Self::Array) -> ArrayId {
49        Self::ID
50    }
51
52    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
53        // FoRArray children order (from visit_children):
54        // 1. encoded
55
56        vortex_ensure!(
57            children.len() == 1,
58            "Expected 1 child for FoR encoding, got {}",
59            children.len()
60        );
61
62        array.encoded = children[0].clone();
63
64        Ok(())
65    }
66
67    fn metadata(array: &FoRArray) -> VortexResult<Self::Metadata> {
68        Ok(array.reference_scalar().clone())
69    }
70
71    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
72        // Note that we **only** serialize the optional scalar value (not including the dtype).
73        Ok(Some(ScalarValue::to_proto_bytes(metadata.value())))
74    }
75
76    fn deserialize(
77        bytes: &[u8],
78        dtype: &DType,
79        _len: usize,
80        _buffers: &[BufferHandle],
81        _session: &VortexSession,
82    ) -> VortexResult<Self::Metadata> {
83        let scalar_value = ScalarValue::from_proto_bytes(bytes, dtype)?;
84        Scalar::try_new(dtype.clone(), scalar_value)
85    }
86
87    fn build(
88        dtype: &DType,
89        len: usize,
90        metadata: &Self::Metadata,
91        _buffers: &[BufferHandle],
92        children: &dyn ArrayChildren,
93    ) -> VortexResult<FoRArray> {
94        if children.len() != 1 {
95            vortex_bail!(
96                "Expected 1 child for FoR encoding, found {}",
97                children.len()
98            )
99        }
100
101        let encoded = children.get(0, dtype, len)?;
102
103        FoRArray::try_new(encoded, metadata.clone())
104    }
105
106    fn reduce_parent(
107        array: &Self::Array,
108        parent: &ArrayRef,
109        child_idx: usize,
110    ) -> VortexResult<Option<ArrayRef>> {
111        PARENT_RULES.evaluate(array, parent, child_idx)
112    }
113
114    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
115        Ok(decompress(array, ctx)?.into_array())
116    }
117
118    fn execute_parent(
119        array: &Self::Array,
120        parent: &ArrayRef,
121        child_idx: usize,
122        ctx: &mut ExecutionCtx,
123    ) -> VortexResult<Option<ArrayRef>> {
124        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
125    }
126}
127
128#[derive(Debug)]
129pub struct FoRVTable;
130
131impl FoRVTable {
132    pub const ID: ArrayId = ArrayId::new_ref("fastlanes.for");
133}