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::fmt::Formatter;
6
7use vortex_array::ArrayRef;
8use vortex_array::DeserializeMetadata;
9use vortex_array::SerializeMetadata;
10use vortex_array::buffer::BufferHandle;
11use vortex_array::serde::ArrayChildren;
12use vortex_array::vtable;
13use vortex_array::vtable::ArrayId;
14use vortex_array::vtable::ArrayVTable;
15use vortex_array::vtable::ArrayVTableExt;
16use vortex_array::vtable::NotSupported;
17use vortex_array::vtable::VTable;
18use vortex_array::vtable::ValidityVTableFromChild;
19use vortex_dtype::DType;
20use vortex_error::VortexResult;
21use vortex_error::vortex_bail;
22use vortex_error::vortex_ensure;
23use vortex_scalar::Scalar;
24use vortex_scalar::ScalarValue;
25
26use crate::FoRArray;
27use crate::r#for::vtable::rules::PARENT_RULES;
28
29mod array;
30mod canonical;
31mod encode;
32mod operations;
33mod rules;
34mod validity;
35mod visitor;
36
37vtable!(FoR);
38
39impl VTable for FoRVTable {
40    type Array = FoRArray;
41
42    type Metadata = ScalarValueMetadata;
43
44    type ArrayVTable = Self;
45    type CanonicalVTable = Self;
46    type OperationsVTable = Self;
47    type ValidityVTable = ValidityVTableFromChild;
48    type VisitorVTable = Self;
49    type ComputeVTable = NotSupported;
50    type EncodeVTable = Self;
51
52    fn id(&self) -> ArrayId {
53        ArrayId::new_ref("fastlanes.for")
54    }
55
56    fn encoding(_array: &Self::Array) -> ArrayVTable {
57        FoRVTable.as_vtable()
58    }
59
60    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
61        // FoRArray children order (from visit_children):
62        // 1. encoded
63
64        vortex_ensure!(
65            children.len() == 1,
66            "Expected 1 child for FoR encoding, got {}",
67            children.len()
68        );
69
70        array.encoded = children[0].clone();
71
72        Ok(())
73    }
74
75    fn metadata(array: &FoRArray) -> VortexResult<Self::Metadata> {
76        Ok(ScalarValueMetadata(
77            array.reference_scalar().value().clone(),
78        ))
79    }
80
81    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
82        Ok(Some(metadata.serialize()))
83    }
84
85    fn deserialize(buffer: &[u8]) -> VortexResult<Self::Metadata> {
86        ScalarValueMetadata::deserialize(buffer)
87    }
88
89    fn build(
90        &self,
91        dtype: &DType,
92        len: usize,
93        metadata: &Self::Metadata,
94        _buffers: &[BufferHandle],
95        children: &dyn ArrayChildren,
96    ) -> VortexResult<FoRArray> {
97        if children.len() != 1 {
98            vortex_bail!(
99                "Expected 1 child for FoR encoding, found {}",
100                children.len()
101            )
102        }
103
104        let encoded = children.get(0, dtype, len)?;
105        let reference = Scalar::new(dtype.clone(), metadata.0.clone());
106
107        FoRArray::try_new(encoded, reference)
108    }
109
110    fn reduce_parent(
111        array: &Self::Array,
112        parent: &ArrayRef,
113        child_idx: usize,
114    ) -> VortexResult<Option<ArrayRef>> {
115        PARENT_RULES.evaluate(array, parent, child_idx)
116    }
117}
118
119#[derive(Debug)]
120pub struct FoRVTable;
121
122#[derive(Clone)]
123pub struct ScalarValueMetadata(pub ScalarValue);
124
125impl SerializeMetadata for ScalarValueMetadata {
126    fn serialize(self) -> Vec<u8> {
127        self.0.to_protobytes()
128    }
129}
130
131impl DeserializeMetadata for ScalarValueMetadata {
132    type Output = ScalarValueMetadata;
133
134    fn deserialize(metadata: &[u8]) -> VortexResult<Self::Output> {
135        let scalar_value = ScalarValue::from_protobytes(metadata)?;
136        Ok(ScalarValueMetadata(scalar_value))
137    }
138}
139
140impl Debug for ScalarValueMetadata {
141    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142        write!(f, "{}", &self.0)
143    }
144}