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

vortex_array/arrays/varbin/vtable/
mod.rs
