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

vortex_array/arrays/varbin/vtable/
mod.rs
