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::array::Array;
21use crate::array::ArrayId;
22use crate::array::ArrayView;
23use crate::array::VTable;
24use crate::arrays::PrimitiveArray;
25use crate::arrays::varbin::VarBinArrayExt;
26use crate::arrays::varbin::VarBinArraySlotsExt;
27use crate::arrays::varbin::VarBinData;
28use crate::arrays::varbin::VarBinSlots;
29use crate::buffer::BufferHandle;
30use crate::builders::ArrayBuilder;
31use crate::builders::VarBinViewBuilder;
32use crate::dtype::DType;
33use crate::dtype::Nullability;
34use crate::dtype::PType;
35use crate::match_each_integer_ptype;
36use crate::match_each_varbin_builder;
37use crate::serde::ArrayChildren;
38use crate::validity::Validity;
39pub(crate) mod canonical;
40mod kernel;
41mod operations;
42mod validity;
43
44use canonical::varbin_to_canonical;
45use vortex_session::VortexSession;
46
47use crate::EqMode;
48use crate::arrays::varbin::compute::rules::PARENT_RULES;
49use crate::hash::ArrayEq;
50use crate::hash::ArrayHash;
51
52/// A [`VarBin`]-encoded Vortex array.
53pub type VarBinArray = Array<VarBin>;
54
55pub(crate) fn initialize(session: &VortexSession) {
56    kernel::initialize(session);
57}
58
59#[derive(Clone, prost::Message)]
60pub struct VarBinMetadata {
61    #[prost(enumeration = "PType", tag = "1")]
62    pub(crate) offsets_ptype: i32,
63}
64
65impl ArrayHash for VarBinData {
66    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
67        self.bytes().array_hash(state, accuracy);
68    }
69}
70
71impl ArrayEq for VarBinData {
72    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
73        self.bytes().array_eq(other.bytes(), accuracy)
74    }
75}
76
77impl VTable for VarBin {
78    type TypedArrayData = VarBinData;
79
80    type OperationsVTable = Self;
81    type ValidityVTable = Self;
82    fn id(&self) -> ArrayId {
83        static ID: CachedId = CachedId::new("vortex.varbin");
84        *ID
85    }
86
87    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
88        1
89    }
90
91    fn validate(
92        &self,
93        _data: &VarBinData,
94        dtype: &DType,
95        len: usize,
96        slots: &[Option<ArrayRef>],
97    ) -> VortexResult<()> {
98        vortex_ensure_eq!(slots.len(), VarBinSlots::COUNT);
99        let offsets = slots[VarBinSlots::OFFSETS]
100            .as_ref()
101            .vortex_expect("VarBinArray offsets slot");
102        vortex_ensure_eq!(
103            offsets.len().saturating_sub(1),
104            len,
105            "VarBinArray length does not match outer length",
106        );
107        vortex_ensure!(
108            matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
109            "VarBinArray dtype must be binary or utf8, got {dtype}"
110        );
111        Ok(())
112    }
113
114    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
115        match idx {
116            0 => array.bytes_handle().clone(),
117            _ => vortex_panic!("VarBinArray buffer index {idx} out of bounds"),
118        }
119    }
120
121    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
122        match idx {
123            0 => Some("bytes".to_string()),
124            _ => vortex_panic!("VarBinArray buffer_name index {idx} out of bounds"),
125        }
126    }
127
128    fn with_buffers(
129        &self,
130        array: ArrayView<'_, Self>,
131        buffers: &[BufferHandle],
132    ) -> VortexResult<ArrayParts<Self>> {
133        vortex_ensure_eq!(buffers.len(), 1);
134        let mut data = array.data().clone();
135        data.bytes = buffers[0].clone();
136        Ok(ArrayParts::new(
137            self.clone(),
138            array.dtype().clone(),
139            array.len(),
140            data,
141            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(
191            self.clone(),
192            dtype.clone(),
193            len,
194            data,
195            slots,
196        ))
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
