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_ensure_eq;
12use vortex_error::vortex_panic;
13use vortex_session::registry::CachedId;
14
15use crate::ArrayParts;
16use crate::ArrayRef;
17use crate::ExecutionCtx;
18use crate::ExecutionResult;
19use crate::IntoArray;
20use crate::VortexSessionExecute;
21use crate::array::Array;
22use crate::array::ArrayId;
23use crate::array::ArrayView;
24use crate::array::VTable;
25use crate::arrays::PrimitiveArray;
26use crate::arrays::varbin::VarBinArrayExt;
27use crate::arrays::varbin::VarBinArraySlotsExt;
28use crate::arrays::varbin::VarBinData;
29use crate::arrays::varbin::VarBinSlots;
30use crate::buffer::BufferHandle;
31use crate::builders::ArrayBuilder;
32use crate::builders::VarBinViewBuilder;
33use crate::dtype::DType;
34use crate::dtype::Nullability;
35use crate::dtype::PType;
36use crate::match_each_integer_ptype;
37use crate::match_each_varbin_builder;
38use crate::serde::ArrayChildren;
39use crate::validity::Validity;
40pub(crate) mod canonical;
41mod kernel;
42mod operations;
43mod validity;
44
45use canonical::varbin_to_canonical;
46use vortex_session::VortexSession;
47
48use crate::EqMode;
49use crate::arrays::varbin::compute::rules::PARENT_RULES;
50use crate::hash::ArrayEq;
51use crate::hash::ArrayHash;
52
53/// A [`VarBin`]-encoded Vortex array.
54pub type VarBinArray = Array<VarBin>;
55
56pub(crate) fn initialize(session: &VortexSession) {
57    kernel::initialize(session);
58}
59
60#[derive(Clone, prost::Message)]
61pub struct VarBinMetadata {
62    #[prost(enumeration = "PType", tag = "1")]
63    pub(crate) offsets_ptype: i32,
64}
65
66impl ArrayHash for VarBinData {
67    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
68        self.bytes().array_hash(state, accuracy);
69    }
70}
71
72impl ArrayEq for VarBinData {
73    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
74        self.bytes().array_eq(other.bytes(), accuracy)
75    }
76}
77
78impl VTable for VarBin {
79    type TypedArrayData = VarBinData;
80
81    type OperationsVTable = Self;
82    type ValidityVTable = Self;
83    fn id(&self) -> ArrayId {
84        static ID: CachedId = CachedId::new("vortex.varbin");
85        *ID
86    }
87
88    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
89        1
90    }
91
92    fn validate(
93        &self,
94        _data: &VarBinData,
95        dtype: &DType,
96        len: usize,
97        slots: &[Option<ArrayRef>],
98    ) -> VortexResult<()> {
99        vortex_ensure_eq!(slots.len(), VarBinSlots::COUNT);
100        let offsets = slots[VarBinSlots::OFFSETS]
101            .as_ref()
102            .vortex_expect("VarBinArray offsets slot");
103        vortex_ensure_eq!(
104            offsets.len().saturating_sub(1),
105            len,
106            "VarBinArray length does not match outer length",
107        );
108        vortex_ensure!(
109            matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
110            "VarBinArray dtype must be binary or utf8, got {dtype}"
111        );
112        Ok(())
113    }
114
115    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
116        match idx {
117            0 => array.bytes_handle().clone(),
118            _ => vortex_panic!("VarBinArray buffer index {idx} out of bounds"),
119        }
120    }
121
122    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
123        match idx {
124            0 => Some("bytes".to_string()),
125            _ => vortex_panic!("VarBinArray buffer_name index {idx} out of bounds"),
126        }
127    }
128
129    fn with_buffers(
130        &self,
131        array: ArrayView<'_, Self>,
132        buffers: &[BufferHandle],
133    ) -> VortexResult<ArrayParts<Self>> {
134        vortex_ensure_eq!(buffers.len(), 1);
135        let mut data = array.data().clone();
136        data.bytes = buffers[0].clone();
137        Ok(ArrayParts::new(
138            self.clone(),
139            array.dtype().clone(),
140            array.len(),
141            data,
142            array.slots().iter().cloned().collect(),
143        ))
144    }
145
146    fn serialize(
147        array: ArrayView<'_, Self>,
148        _session: &VortexSession,
149    ) -> VortexResult<Option<Vec<u8>>> {
150        Ok(Some(
151            VarBinMetadata {
152                offsets_ptype: PType::try_from(array.offsets().dtype())
153                    .vortex_expect("Must be a valid PType") as i32,
154            }
155            .encode_to_vec(),
156        ))
157    }
158
159    fn deserialize(
160        &self,
161        dtype: &DType,
162        len: usize,
163        metadata: &[u8],
164        buffers: &[BufferHandle],
165        children: &dyn ArrayChildren,
166        session: &VortexSession,
167    ) -> VortexResult<ArrayParts<Self>> {
168        let metadata = VarBinMetadata::decode(metadata)?;
169        let validity = if children.len() == 1 {
170            Validity::from(dtype.nullability())
171        } else if children.len() == 2 {
172            let validity = children.get(1, &Validity::DTYPE, len)?;
173            Validity::Array(validity)
174        } else {
175            vortex_bail!("Expected 1 or 2 children, got {}", children.len());
176        };
177
178        let offsets = children.get(
179            0,
180            &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable),
181            len + 1,
182        )?;
183
184        if buffers.len() != 1 {
185            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
186        }
187        let bytes = buffers[0].clone().try_to_host_sync()?;
188
189        let data = VarBinData::try_build(
190            offsets.clone(),
191            bytes,
192            dtype.clone(),
193            validity.clone(),
194            &mut session.create_execution_ctx(),
195        )?;
196        let slots = VarBinData::make_slots(offsets, &validity, len);
197        Ok(ArrayParts::new(
198            self.clone(),
199            dtype.clone(),
200            len,
201            data,
202            slots,
203        ))
204    }
205
206    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
207        VarBinSlots::NAMES[idx].to_string()
208    }
209
210    fn reduce_parent(
211        array: ArrayView<'_, Self>,
212        parent: &ArrayRef,
213        child_idx: usize,
214    ) -> VortexResult<Option<ArrayRef>> {
215        PARENT_RULES.evaluate(array, parent, child_idx)
216    }
217
218    fn append_to_builder(
219        array: ArrayView<'_, Self>,
220        builder: &mut dyn ArrayBuilder,
221        ctx: &mut ExecutionCtx,
222    ) -> VortexResult<()> {
223        if let Some(result) =
224            match_each_varbin_builder!(builder, |builder| builder.append_varbin(array, ctx))
225        {
226            return result;
227        }
228
229        // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four
230        // `VarBinBuilder` widths above, and `VarBinViewBuilder` below.
231        let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
232            vortex_bail!("append_to_builder for VarBin requires a variable-binary builder")
233        };
234        append_to_varbinview(array, builder, ctx)
235    }
236
237    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
238        Ok(ExecutionResult::done(
239            varbin_to_canonical(array.as_view(), ctx)?.into_array(),
240        ))
241    }
242}
243
244/// Hands the value bytes to `builder` as a data buffer with views built over them.
245///
246/// Canonicalizing first would build the same views, then pay for them twice more: once to wrap
247/// them in a `VarBinViewArray` the builder immediately unwraps, and once for
248/// `append_varbinview_array` to rewrite every view so its buffer index is rebased onto the
249/// builder's. Handing the heap and offsets to the builder instead makes the whole append one view
250/// per row with no byte copy — the builder adopts the referenced range of the heap as it is. That
251/// range is fully covered by the new views, so this stays valid for a compacting builder too.
252fn append_to_varbinview(
253    array: ArrayView<'_, VarBin>,
254    builder: &mut VarBinViewBuilder,
255    ctx: &mut ExecutionCtx,
256) -> VortexResult<()> {
257    let len = array.as_ref().len();
258    let validity = array.varbin_validity().execute_mask(len, ctx)?;
259
260    let parts = array.into_owned().into_data_parts();
261    let offsets = parts.offsets.execute::<PrimitiveArray>(ctx)?;
262    match_each_integer_ptype!(offsets.ptype(), |P| {
263        builder.append_buffer_with_offsets(
264            parts.bytes.unwrap_host(),
265            offsets.as_slice::<P>(),
266            &validity,
267        )
268    });
269    Ok(())
270}
271
272#[derive(Clone, Debug)]
273pub struct VarBin;
vortex_array/arrays/varbin/vtable/mod.rs

vortex_array/arrays/varbin/vtable/
mod.rs
