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::varbin::VarBinArrayExt;
24use crate::arrays::varbin::VarBinData;
25use crate::arrays::varbin::array::NUM_SLOTS;
26use crate::arrays::varbin::array::OFFSETS_SLOT;
27use crate::arrays::varbin::array::SLOT_NAMES;
28use crate::buffer::BufferHandle;
29use crate::dtype::DType;
30use crate::dtype::Nullability;
31use crate::dtype::PType;
32use crate::serde::ArrayChildren;
33use crate::validity::Validity;
34mod canonical;
35mod kernel;
36mod operations;
37mod validity;
38
39use canonical::varbin_to_canonical;
40use vortex_session::VortexSession;
41
42use crate::EqMode;
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
50pub(crate) fn initialize(session: &VortexSession) {
51    kernel::initialize(session);
52}
53
54#[derive(Clone, prost::Message)]
55pub struct VarBinMetadata {
56    #[prost(enumeration = "PType", tag = "1")]
57    pub(crate) offsets_ptype: i32,
58}
59
60impl ArrayHash for VarBinData {
61    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
62        self.bytes().array_hash(state, accuracy);
63    }
64}
65
66impl ArrayEq for VarBinData {
67    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
68        self.bytes().array_eq(other.bytes(), accuracy)
69    }
70}
71
72impl VTable for VarBin {
73    type TypedArrayData = VarBinData;
74
75    type OperationsVTable = Self;
76    type ValidityVTable = Self;
77    fn id(&self) -> ArrayId {
78        static ID: CachedId = CachedId::new("vortex.varbin");
79        *ID
80    }
81
82    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
83        1
84    }
85
86    fn validate(
87        &self,
88        _data: &VarBinData,
89        dtype: &DType,
90        len: usize,
91        slots: &[Option<ArrayRef>],
92    ) -> VortexResult<()> {
93        vortex_ensure!(
94            slots.len() == NUM_SLOTS,
95            "VarBinArray expected {NUM_SLOTS} slots, found {}",
96            slots.len()
97        );
98        let offsets = slots[OFFSETS_SLOT]
99            .as_ref()
100            .vortex_expect("VarBinArray offsets slot");
101        vortex_ensure!(
102            offsets.len().saturating_sub(1) == len,
103            "VarBinArray length {} does not match outer length {}",
104            offsets.len().saturating_sub(1),
105            len
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!(
134            buffers.len() == 1,
135            "Expected 1 buffer, got {}",
136            buffers.len()
137        );
138        let mut data = array.data().clone();
139        data.bytes = buffers[0].clone();
140        Ok(
141            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
142                .with_slots(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(offsets.clone(), bytes, dtype.clone(), validity.clone())?;
190        let slots = VarBinData::make_slots(offsets, &validity, len);
191        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
192    }
193
194    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
195        SLOT_NAMES[idx].to_string()
196    }
197
198    fn reduce_parent(
199        array: ArrayView<'_, Self>,
200        parent: &ArrayRef,
201        child_idx: usize,
202    ) -> VortexResult<Option<ArrayRef>> {
203        PARENT_RULES.evaluate(array, parent, child_idx)
204    }
205
206    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
207        Ok(ExecutionResult::done(
208            varbin_to_canonical(array.as_view(), ctx)?.into_array(),
209        ))
210    }
211}
212
213#[derive(Clone, Debug)]
214pub struct VarBin;
vortex_array/arrays/varbin/vtable/mod.rs

vortex_array/arrays/varbin/vtable/
mod.rs
