vortex_array/arrays/varbin/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hasher;
5
6use prost::Message;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_ensure;
11use vortex_error::vortex_panic;
12use vortex_session::registry::CachedId;
13
14use crate::ArrayParts;
15use crate::ArrayRef;
16use crate::ExecutionCtx;
17use crate::ExecutionResult;
18use crate::IntoArray;
19use crate::array::Array;
20use crate::array::ArrayId;
21use crate::array::ArrayView;
22use crate::array::VTable;
23use crate::arrays::varbin::VarBinArraySlotsExt;
24use crate::arrays::varbin::VarBinData;
25use crate::arrays::varbin::VarBinSlots;
26use crate::buffer::BufferHandle;
27use crate::dtype::DType;
28use crate::dtype::Nullability;
29use crate::dtype::PType;
30use crate::serde::ArrayChildren;
31use crate::validity::Validity;
32mod canonical;
33mod kernel;
34mod operations;
35mod validity;
36
37use canonical::varbin_to_canonical;
38use vortex_session::VortexSession;
39
40use crate::EqMode;
41use crate::arrays::varbin::compute::rules::PARENT_RULES;
42use crate::hash::ArrayEq;
43use crate::hash::ArrayHash;
44
45/// A [`VarBin`]-encoded Vortex array.
46pub type VarBinArray = Array<VarBin>;
47
48pub(crate) fn initialize(session: &VortexSession) {
49    kernel::initialize(session);
50}
51
52#[derive(Clone, prost::Message)]
53pub struct VarBinMetadata {
54    #[prost(enumeration = "PType", tag = "1")]
55    pub(crate) offsets_ptype: i32,
56}
57
58impl ArrayHash for VarBinData {
59    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
60        self.bytes().array_hash(state, accuracy);
61    }
62}
63
64impl ArrayEq for VarBinData {
65    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
66        self.bytes().array_eq(other.bytes(), accuracy)
67    }
68}
69
70impl VTable for VarBin {
71    type TypedArrayData = VarBinData;
72
73    type OperationsVTable = Self;
74    type ValidityVTable = Self;
75    fn id(&self) -> ArrayId {
76        static ID: CachedId = CachedId::new("vortex.varbin");
77        *ID
78    }
79
80    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
81        1
82    }
83
84    fn validate(
85        &self,
86        _data: &VarBinData,
87        dtype: &DType,
88        len: usize,
89        slots: &[Option<ArrayRef>],
90    ) -> VortexResult<()> {
91        vortex_ensure!(
92            slots.len() == VarBinSlots::COUNT,
93            "VarBinArray expected {} slots, found {}",
94            VarBinSlots::COUNT,
95            slots.len()
96        );
97        let offsets = slots[VarBinSlots::OFFSETS]
98            .as_ref()
99            .vortex_expect("VarBinArray offsets slot");
100        vortex_ensure!(
101            offsets.len().saturating_sub(1) == len,
102            "VarBinArray length {} does not match outer length {}",
103            offsets.len().saturating_sub(1),
104            len
105        );
106        vortex_ensure!(
107            matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
108            "VarBinArray dtype must be binary or utf8, got {dtype}"
109        );
110        Ok(())
111    }
112
113    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
114        match idx {
115            0 => array.bytes_handle().clone(),
116            _ => vortex_panic!("VarBinArray buffer index {idx} out of bounds"),
117        }
118    }
119
120    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
121        match idx {
122            0 => Some("bytes".to_string()),
123            _ => vortex_panic!("VarBinArray buffer_name index {idx} out of bounds"),
124        }
125    }
126
127    fn with_buffers(
128        &self,
129        array: ArrayView<'_, Self>,
130        buffers: &[BufferHandle],
131    ) -> VortexResult<ArrayParts<Self>> {
132        vortex_ensure!(
133            buffers.len() == 1,
134            "Expected 1 buffer, got {}",
135            buffers.len()
136        );
137        let mut data = array.data().clone();
138        data.bytes = buffers[0].clone();
139        Ok(
140            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
141                .with_slots(array.slots().iter().cloned().collect()),
142        )
143    }
144
145    fn serialize(
146        array: ArrayView<'_, Self>,
147        _session: &VortexSession,
148    ) -> VortexResult<Option<Vec<u8>>> {
149        Ok(Some(
150            VarBinMetadata {
151                offsets_ptype: PType::try_from(array.offsets().dtype())
152                    .vortex_expect("Must be a valid PType") as i32,
153            }
154            .encode_to_vec(),
155        ))
156    }
157
158    fn deserialize(
159        &self,
160        dtype: &DType,
161        len: usize,
162        metadata: &[u8],
163        buffers: &[BufferHandle],
164        children: &dyn ArrayChildren,
165        _session: &VortexSession,
166    ) -> VortexResult<ArrayParts<Self>> {
167        let metadata = VarBinMetadata::decode(metadata)?;
168        let validity = if children.len() == 1 {
169            Validity::from(dtype.nullability())
170        } else if children.len() == 2 {
171            let validity = children.get(1, &Validity::DTYPE, len)?;
172            Validity::Array(validity)
173        } else {
174            vortex_bail!("Expected 1 or 2 children, got {}", children.len());
175        };
176
177        let offsets = children.get(
178            0,
179            &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable),
180            len + 1,
181        )?;
182
183        if buffers.len() != 1 {
184            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
185        }
186        let bytes = buffers[0].clone().try_to_host_sync()?;
187
188        let data = VarBinData::try_build(offsets.clone(), bytes, dtype.clone(), validity.clone())?;
189        let slots = VarBinData::make_slots(offsets, &validity, len);
190        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
191    }
192
193    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
194        VarBinSlots::NAMES[idx].to_string()
195    }
196
197    fn reduce_parent(
198        array: ArrayView<'_, Self>,
199        parent: &ArrayRef,
200        child_idx: usize,
201    ) -> VortexResult<Option<ArrayRef>> {
202        PARENT_RULES.evaluate(array, parent, child_idx)
203    }
204
205    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
206        Ok(ExecutionResult::done(
207            varbin_to_canonical(array.as_view(), ctx)?.into_array(),
208        ))
209    }
210}
211
212#[derive(Clone, Debug)]
213pub struct VarBin;
vortex_array/arrays/varbin/vtable/mod.rs

vortex_array/arrays/varbin/vtable/
mod.rs
