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

vortex_array/arrays/varbin/vtable/
mod.rs
