vortex_fastlanes/bitpacking/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hash;
5use std::hash::Hasher;
6
7use prost::Message;
8use vortex_array::Array;
9use vortex_array::ArrayEq;
10use vortex_array::ArrayHash;
11use vortex_array::ArrayId;
12use vortex_array::ArrayParts;
13use vortex_array::ArrayRef;
14use vortex_array::ArraySlots;
15use vortex_array::ArrayView;
16use vortex_array::EqMode;
17use vortex_array::ExecutionCtx;
18use vortex_array::ExecutionResult;
19use vortex_array::IntoArray;
20use vortex_array::buffer::BufferHandle;
21use vortex_array::builders::ArrayBuilder;
22use vortex_array::dtype::DType;
23use vortex_array::dtype::PType;
24use vortex_array::match_each_integer_ptype;
25use vortex_array::patches::Patches;
26use vortex_array::patches::PatchesData;
27use vortex_array::patches::PatchesMetadata;
28use vortex_array::require_patches;
29use vortex_array::require_validity;
30use vortex_array::serde::ArrayChildren;
31use vortex_array::validity::Validity;
32use vortex_array::vtable::VTable;
33use vortex_array::vtable::child_to_validity;
34use vortex_array::vtable::validity_to_child;
35use vortex_error::VortexExpect;
36use vortex_error::VortexResult;
37use vortex_error::vortex_bail;
38use vortex_error::vortex_ensure;
39use vortex_error::vortex_err;
40use vortex_error::vortex_panic;
41use vortex_session::VortexSession;
42use vortex_session::registry::CachedId;
43
44use crate::BitPackedArrayExt;
45use crate::BitPackedData;
46use crate::BitPackedDataParts;
47use crate::bitpack_decompress::unpack_array;
48use crate::bitpack_decompress::unpack_into_primitive_builder;
49use crate::bitpacking::array::BitPackedSlots;
50use crate::bitpacking::array::BitPackedSlotsView;
51use crate::bitpacking::array::PATCH_SLOTS;
52use crate::bitpacking::vtable::rules::RULES;
53mod kernels;
54mod operations;
55mod rules;
56mod validity;
57
58/// A [`BitPacked`]-encoded Vortex array.
59pub type BitPackedArray = Array<BitPacked>;
60
61pub(crate) fn initialize(session: &VortexSession) {
62    kernels::initialize(session);
63}
64
65#[derive(Clone, prost::Message)]
66pub struct BitPackedMetadata {
67    #[prost(uint32, tag = "1")]
68    pub(crate) bit_width: u32,
69    #[prost(uint32, tag = "2")]
70    pub(crate) offset: u32, // must be <1024
71    #[prost(message, optional, tag = "3")]
72    pub(crate) patches: Option<PatchesMetadata>,
73}
74
75impl ArrayHash for BitPackedData {
76    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
77        self.offset.hash(state);
78        self.bit_width.hash(state);
79        self.packed.array_hash(state, accuracy);
80        self.patches_data.hash(state);
81    }
82}
83
84impl ArrayEq for BitPackedData {
85    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
86        self.offset == other.offset
87            && self.bit_width == other.bit_width
88            && self.packed.array_eq(&other.packed, accuracy)
89            && self.patches_data == other.patches_data
90    }
91}
92
93impl VTable for BitPacked {
94    type TypedArrayData = BitPackedData;
95
96    type OperationsVTable = Self;
97    type ValidityVTable = Self;
98
99    fn id(&self) -> ArrayId {
100        static ID: CachedId = CachedId::new("fastlanes.bitpacked");
101        *ID
102    }
103
104    fn validate(
105        &self,
106        data: &Self::TypedArrayData,
107        dtype: &DType,
108        len: usize,
109        slots: &[Option<ArrayRef>],
110    ) -> VortexResult<()> {
111        let bp_slots = BitPackedSlotsView::from_slots(slots);
112
113        let validity = child_to_validity(bp_slots.validity_child, dtype.nullability());
114        let patches =
115            PatchesData::patches_from_slots(data.patches_data.as_ref(), len, slots, PATCH_SLOTS);
116        BitPackedData::validate(
117            &data.packed,
118            dtype.as_ptype(),
119            &validity,
120            patches.as_ref(),
121            data.bit_width,
122            len,
123            data.offset,
124        )
125    }
126
127    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
128        1
129    }
130
131    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
132        match idx {
133            0 => array.packed().clone(),
134            _ => vortex_panic!("BitPackedArray buffer index {idx} out of bounds"),
135        }
136    }
137
138    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
139        match idx {
140            0 => Some("packed".to_string()),
141            _ => None,
142        }
143    }
144
145    fn with_buffers(
146        &self,
147        array: ArrayView<'_, Self>,
148        buffers: &[BufferHandle],
149    ) -> VortexResult<ArrayParts<Self>> {
150        vortex_ensure!(
151            buffers.len() == 1,
152            "Expected 1 buffer, got {}",
153            buffers.len()
154        );
155        let mut data = array.data().clone();
156        data.packed = buffers[0].clone();
157        Ok(
158            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
159                .with_slots(array.slots().iter().cloned().collect()),
160        )
161    }
162
163    fn serialize(
164        array: ArrayView<'_, Self>,
165        _session: &VortexSession,
166    ) -> VortexResult<Option<Vec<u8>>> {
167        Ok(Some(
168            BitPackedMetadata {
169                bit_width: array.bit_width() as u32,
170                offset: array.offset() as u32,
171                patches: array
172                    .patches()
173                    .map(|p| p.to_metadata(array.len(), array.dtype()))
174                    .transpose()?,
175            }
176            .encode_to_vec(),
177        ))
178    }
179
180    fn deserialize(
181        &self,
182        dtype: &DType,
183        len: usize,
184        metadata: &[u8],
185        buffers: &[BufferHandle],
186        children: &dyn ArrayChildren,
187        _session: &VortexSession,
188    ) -> VortexResult<ArrayParts<Self>> {
189        let metadata = BitPackedMetadata::decode(metadata)?;
190        if buffers.len() != 1 {
191            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
192        }
193        let packed = buffers[0].clone();
194
195        let load_validity = |child_idx: usize| {
196            if children.len() == child_idx {
197                Ok(Validity::from(dtype.nullability()))
198            } else if children.len() == child_idx + 1 {
199                let validity = children.get(child_idx, &Validity::DTYPE, len)?;
200                Ok(Validity::Array(validity))
201            } else {
202                vortex_bail!(
203                    "Expected {} or {} children, got {}",
204                    child_idx,
205                    child_idx + 1,
206                    children.len()
207                );
208            }
209        };
210
211        let validity_idx = match &metadata.patches {
212            None => 0,
213            Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3,
214            Some(_) => 2,
215        };
216
217        let validity = load_validity(validity_idx)?;
218
219        let patches = metadata
220            .patches
221            .map(|p| {
222                let indices = children.get(0, &p.indices_dtype()?, p.len()?)?;
223                let values = children.get(1, dtype, p.len()?)?;
224                let chunk_offsets = p
225                    .chunk_offsets_dtype()?
226                    .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize))
227                    .transpose()?;
228
229                Patches::new(len, p.offset()?, indices, values, chunk_offsets)
230            })
231            .transpose()?;
232
233        let slots = {
234            let mut s = ArraySlots::with_capacity(4);
235            PatchesData::push_slots(&mut s, patches.as_ref());
236            s.push(validity_to_child(&validity, len));
237            s
238        };
239        let data = BitPackedData::try_new(
240            packed,
241            patches,
242            u8::try_from(metadata.bit_width).map_err(|_| {
243                vortex_err!(
244                    "BitPackedMetadata bit_width {} does not fit in u8",
245                    metadata.bit_width
246                )
247            })?,
248            u16::try_from(metadata.offset).map_err(|_| {
249                vortex_err!(
250                    "BitPackedMetadata offset {} does not fit in u16",
251                    metadata.offset
252                )
253            })?,
254        )?;
255        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
256    }
257
258    fn append_to_builder(
259        array: ArrayView<'_, Self>,
260        builder: &mut dyn ArrayBuilder,
261        ctx: &mut ExecutionCtx,
262    ) -> VortexResult<()> {
263        match_each_integer_ptype!(array.dtype().as_ptype(), |T| {
264            unpack_into_primitive_builder::<T>(
265                array,
266                builder
267                    .as_any_mut()
268                    .downcast_mut()
269                    .vortex_expect("bit packed array must canonicalize into a primitive array"),
270                ctx,
271            )
272        })
273    }
274
275    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
276        BitPackedSlots::NAMES[idx].to_string()
277    }
278
279    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
280        require_patches!(
281            array,
282            BitPackedSlots::PATCH_INDICES,
283            BitPackedSlots::PATCH_VALUES,
284            BitPackedSlots::PATCH_CHUNK_OFFSETS
285        );
286        require_validity!(array, BitPackedSlots::VALIDITY_CHILD);
287
288        Ok(ExecutionResult::done(
289            unpack_array(array.as_view(), ctx)?.into_array(),
290        ))
291    }
292
293    fn reduce_parent(
294        array: ArrayView<'_, Self>,
295        parent: &ArrayRef,
296        child_idx: usize,
297    ) -> VortexResult<Option<ArrayRef>> {
298        RULES.evaluate(array, parent, child_idx)
299    }
300}
301
302#[derive(Clone, Debug)]
303pub struct BitPacked;
304
305impl BitPacked {
306    pub fn try_new(
307        packed: BufferHandle,
308        ptype: PType,
309        validity: Validity,
310        patches: Option<Patches>,
311        bit_width: u8,
312        len: usize,
313        offset: u16,
314    ) -> VortexResult<BitPackedArray> {
315        let dtype = DType::Primitive(ptype, validity.nullability());
316        let slots = {
317            let mut s = ArraySlots::with_capacity(4);
318            PatchesData::push_slots(&mut s, patches.as_ref());
319            s.push(validity_to_child(&validity, len));
320            s
321        };
322        let data = BitPackedData::try_new(packed, patches, bit_width, offset)?;
323        Array::try_from_parts(ArrayParts::new(BitPacked, dtype, len, data).with_slots(slots))
324    }
325
326    pub fn into_parts(array: BitPackedArray) -> BitPackedDataParts {
327        let len = array.len();
328        let patches = array.patches();
329        let validity = array.validity().vortex_expect("BitPacked validity");
330        let data = array.into_data();
331        BitPackedDataParts {
332            offset: data.offset,
333            bit_width: data.bit_width,
334            len,
335            packed: data.packed,
336            patches,
337            validity,
338        }
339    }
340
341    /// Encode an array into a bitpacked representation with the given bit width.
342    pub fn encode(
343        array: &ArrayRef,
344        bit_width: u8,
345        ctx: &mut ExecutionCtx,
346    ) -> VortexResult<BitPackedArray> {
347        BitPackedData::encode(array, bit_width, ctx)
348    }
349}
vortex_fastlanes/bitpacking/vtable/mod.rs

vortex_fastlanes/bitpacking/vtable/
mod.rs
