vortex_onpair/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8
9use prost::Message as _;
10use vortex_array::Array;
11use vortex_array::ArrayEq;
12use vortex_array::ArrayHash;
13use vortex_array::ArrayId;
14use vortex_array::ArrayParts;
15use vortex_array::ArrayRef;
16use vortex_array::ArrayView;
17use vortex_array::Canonical;
18use vortex_array::EqMode;
19use vortex_array::ExecutionCtx;
20use vortex_array::ExecutionResult;
21use vortex_array::IntoArray;
22use vortex_array::array_slots;
23use vortex_array::buffer::BufferHandle;
24use vortex_array::builders::ArrayBuilder;
25use vortex_array::builders::VarBinViewBuilder;
26use vortex_array::dtype::DType;
27use vortex_array::dtype::Nullability;
28use vortex_array::dtype::PType;
29use vortex_array::serde::ArrayChildren;
30use vortex_array::validity::Validity;
31use vortex_array::vtable::VTable;
32use vortex_array::vtable::ValidityVTable;
33use vortex_array::vtable::child_to_validity;
34use vortex_array::vtable::validity_to_child;
35use vortex_buffer::ByteBuffer;
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::canonical::canonicalize_onpair;
45use crate::canonical::onpair_decode_views;
46use crate::rules::RULES;
47
48/// An [`OnPair`]-encoded Vortex array.
49pub type OnPairArray = Array<OnPair>;
50
51/// Wire-format metadata persisted alongside the OnPair buffer + slot children.
52///
53/// On disk the layout is FSST-shape:
54///
55/// * Buffer 0 — `dict_bytes`: the dictionary blob built by the OnPair trainer,
56///   padded with `onpair::MAX_TOKEN_SIZE` trailing zero
57///   bytes so the over-copy decoder can read 16 bytes past the last token.
58/// * Slots — see [`OnPairSlots`].
59///
60/// The four integer slot children flow through the standard `compress_child`
61/// pipeline (see `vortex-btrblocks::schemes::string::OnPairScheme`), so any
62/// encoding registered with the compressor can re-encode them — exactly the
63/// same shape as FSST's `codes` `VarBinArray`.
64#[derive(Clone, prost::Message)]
65pub struct OnPairMetadata {
66    /// Width of the per-row primitive `uncompressed_lengths` child.
67    #[prost(enumeration = "PType", tag = "1")]
68    pub uncompressed_lengths_ptype: i32,
69    /// Bits-per-token the column was compressed with (9..=16). Every value
70    /// in the `codes` child only uses its low `bits` bits.
71    #[prost(uint32, tag = "2")]
72    pub bits: u32,
73    /// Number of dictionary tokens. `dict_offsets` has length `dict_size + 1`.
74    /// Bounded by `2^bits ≤ 2^16 = 65_536`, so `u32` is comfortably wide.
75    #[prost(uint32, tag = "3")]
76    pub dict_size: u32,
77    /// Total number of tokens across all rows. `codes` has this length;
78    /// `codes_offsets.last() == total_tokens`.
79    #[prost(uint64, tag = "4")]
80    pub total_tokens: u64,
81    /// PType of the `dict_offsets` slot child (defaults to U32, may be
82    /// narrowed to U16/U8 by the cascading compressor when values fit).
83    #[prost(enumeration = "PType", tag = "5")]
84    pub dict_offsets_ptype: i32,
85    /// PType of the `codes` slot child (typically U16, may be narrowed to U8
86    /// when `bits <= 8`).
87    #[prost(enumeration = "PType", tag = "6")]
88    pub codes_ptype: i32,
89    /// PType of the `codes_offsets` slot child.
90    #[prost(enumeration = "PType", tag = "7")]
91    pub codes_offsets_ptype: i32,
92}
93
94impl OnPairMetadata {
95    pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult<PType> {
96        PType::try_from(self.uncompressed_lengths_ptype)
97            .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype))
98    }
99}
100
101#[array_slots(OnPair)]
102pub struct OnPairSlots {
103    /// `PrimitiveArray<u32>`, length `dict_size + 1`. Cascading compressor may
104    /// narrow the ptype to U16/U8.
105    pub dict_offsets: ArrayRef,
106    /// `PrimitiveArray<u16>`. Each value only uses its low `bits` bits;
107    /// downstream `FastLanes::BitPacking` losslessly shrinks the child to
108    /// exactly `bits`-bit codes on disk.
109    pub codes: ArrayRef,
110    /// `PrimitiveArray<u32>`, length `num_rows + 1`. FoR / RunEnd / etc. apply
111    /// naturally via the cascading compressor.
112    pub codes_offsets: ArrayRef,
113    /// Integer `PrimitiveArray`, length `num_rows`. Used to size the canonical
114    /// output buffer.
115    pub uncompressed_lengths: ArrayRef,
116    /// Optional validity child for the outer string column.
117    pub validity: Option<ArrayRef>,
118}
119
120/// Inner data for an OnPair-encoded array.
121///
122/// Holds only the dictionary blob (buffer 0). Every other piece —
123/// `dict_offsets`, the per-token `codes`, the per-row `codes_offsets`, the
124/// per-row `uncompressed_lengths`, and the optional validity child — is a
125/// Vortex slot child so it can be re-encoded by the cascading compressor.
126#[derive(Clone)]
127pub struct OnPairData {
128    /// The dictionary blob (buffer 0).
129    ///
130    /// INVARIANT: this buffer must be over-padded past its logical end
131    /// (`dict_offsets.last()`) by the decoder's fixed token read width,
132    /// `onpair::MAX_TOKEN_SIZE`. The over-copy decoder reads
133    /// every dictionary entry with one fixed-width load and then advances the
134    /// cursor by the token's true length, so the load for the final, shortest
135    /// token over-reads past the logical end of the dictionary. This is the
136    /// same over-read the decoder accounts for on the final few codes; the
137    /// trailing padding absorbs it so that any entry can be read in bounds.
138    /// `onpair_compress` establishes this padding (see `parts_to_children`);
139    /// the over-copy decoder lives in the `onpair` crate.
140    dict_bytes: BufferHandle,
141    bits: u32,
142    len: usize,
143}
144
145impl OnPairData {
146    pub fn new(dict_bytes: BufferHandle, bits: u32, len: usize) -> Self {
147        Self {
148            dict_bytes,
149            bits,
150            len,
151        }
152    }
153
154    pub fn len(&self) -> usize {
155        self.len
156    }
157
158    pub fn is_empty(&self) -> bool {
159        self.len == 0
160    }
161
162    pub fn bits(&self) -> u32 {
163        self.bits
164    }
165
166    pub fn dict_bytes(&self) -> &ByteBuffer {
167        self.dict_bytes.as_host()
168    }
169
170    pub fn dict_bytes_handle(&self) -> &BufferHandle {
171        &self.dict_bytes
172    }
173}
174
175impl Display for OnPairData {
176    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
177        write!(
178            f,
179            "len: {}, bits: {}, dict_bytes_len: {}",
180            self.len,
181            self.bits,
182            self.dict_bytes.len()
183        )
184    }
185}
186
187impl Debug for OnPairData {
188    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
189        f.debug_struct("OnPairData")
190            .field("len", &self.len)
191            .field("bits", &self.bits)
192            .field("dict_bytes_len", &self.dict_bytes.len())
193            .finish()
194    }
195}
196
197impl ArrayHash for OnPairData {
198    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {
199        self.dict_bytes.as_host().array_hash(state, accuracy);
200        state.write_u32(self.bits);
201    }
202}
203
204impl ArrayEq for OnPairData {
205    fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool {
206        self.bits == other.bits
207            && self
208                .dict_bytes
209                .as_host()
210                .array_eq(other.dict_bytes.as_host(), accuracy)
211    }
212}
213
214/// Zero-sized VTable marker for the OnPair encoding.
215#[derive(Clone, Debug)]
216pub struct OnPair;
217
218impl OnPair {
219    /// Build an [`OnPairArray`] from already-materialised parts.
220    #[expect(clippy::too_many_arguments, reason = "every child is a real input")]
221    pub fn try_new(
222        dtype: DType,
223        dict_bytes: BufferHandle,
224        dict_offsets: ArrayRef,
225        codes: ArrayRef,
226        codes_offsets: ArrayRef,
227        uncompressed_lengths: ArrayRef,
228        validity: Validity,
229        bits: u32,
230    ) -> VortexResult<OnPairArray> {
231        validate_parts(
232            &dtype,
233            &dict_offsets,
234            &codes,
235            &codes_offsets,
236            &uncompressed_lengths,
237            bits,
238        )?;
239        let len = uncompressed_lengths.len();
240        let data = OnPairData::new(dict_bytes, bits, len);
241        let slots = OnPairSlots {
242            dict_offsets,
243            codes,
244            codes_offsets,
245            uncompressed_lengths,
246            validity: validity_to_child(&validity, len),
247        }
248        .into_slots();
249        Ok(unsafe {
250            Array::from_parts_unchecked(ArrayParts::new(OnPair, dtype, len, data).with_slots(slots))
251        })
252    }
253
254    #[expect(clippy::too_many_arguments, reason = "every child is a real input")]
255    pub(crate) unsafe fn new_unchecked(
256        dtype: DType,
257        dict_bytes: BufferHandle,
258        dict_offsets: ArrayRef,
259        codes: ArrayRef,
260        codes_offsets: ArrayRef,
261        uncompressed_lengths: ArrayRef,
262        validity: Validity,
263        bits: u32,
264    ) -> OnPairArray {
265        let len = uncompressed_lengths.len();
266        let data = OnPairData::new(dict_bytes, bits, len);
267        let slots = OnPairSlots {
268            dict_offsets,
269            codes,
270            codes_offsets,
271            uncompressed_lengths,
272            validity: validity_to_child(&validity, len),
273        }
274        .into_slots();
275        unsafe {
276            Array::from_parts_unchecked(ArrayParts::new(OnPair, dtype, len, data).with_slots(slots))
277        }
278    }
279}
280
281fn validate_parts(
282    dtype: &DType,
283    dict_offsets: &ArrayRef,
284    codes: &ArrayRef,
285    codes_offsets: &ArrayRef,
286    uncompressed_lengths: &ArrayRef,
287    bits: u32,
288) -> VortexResult<()> {
289    vortex_ensure!(
290        matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
291        "OnPair arrays must be Binary or Utf8, found {dtype}"
292    );
293    vortex_ensure!((9..=16).contains(&bits), "bits {bits} out of range [9, 16]");
294
295    if !dict_offsets.dtype().is_int() || dict_offsets.dtype().is_nullable() {
296        vortex_bail!(InvalidArgument: "dict_offsets must be non-nullable integer");
297    }
298    if !codes.dtype().is_int() || codes.dtype().is_nullable() {
299        vortex_bail!(InvalidArgument: "codes must be non-nullable integer");
300    }
301    if !codes_offsets.dtype().is_int() || codes_offsets.dtype().is_nullable() {
302        vortex_bail!(InvalidArgument: "codes_offsets must be non-nullable integer");
303    }
304    if !uncompressed_lengths.dtype().is_int() || uncompressed_lengths.dtype().is_nullable() {
305        vortex_bail!(InvalidArgument: "uncompressed_lengths must be non-nullable integer");
306    }
307    if codes_offsets.len() != uncompressed_lengths.len() + 1 {
308        vortex_bail!(InvalidArgument:
309            "codes_offsets.len ({}) != uncompressed_lengths.len + 1 ({})",
310            codes_offsets.len(),
311            uncompressed_lengths.len() + 1
312        );
313    }
314    Ok(())
315}
316
317impl VTable for OnPair {
318    type TypedArrayData = OnPairData;
319    type OperationsVTable = Self;
320    type ValidityVTable = Self;
321
322    fn id(&self) -> ArrayId {
323        static ID: CachedId = CachedId::new("vortex.onpair");
324        *ID
325    }
326
327    fn validate(
328        &self,
329        data: &Self::TypedArrayData,
330        dtype: &DType,
331        len: usize,
332        slots: &[Option<ArrayRef>],
333    ) -> VortexResult<()> {
334        let s = OnPairSlotsView::from_slots(slots);
335        validate_parts(
336            dtype,
337            s.dict_offsets,
338            s.codes,
339            s.codes_offsets,
340            s.uncompressed_lengths,
341            data.bits,
342        )?;
343        if s.uncompressed_lengths.len() != len {
344            vortex_bail!(InvalidArgument: "uncompressed_lengths must have same len as outer array");
345        }
346        if data.len != len {
347            vortex_bail!(InvalidArgument: "OnPairData len {} != outer len {}", data.len, len);
348        }
349        Ok(())
350    }
351
352    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
353        1
354    }
355
356    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
357        match idx {
358            0 => array.dict_bytes_handle().clone(),
359            _ => vortex_panic!("OnPairArray buffer index {idx} out of bounds"),
360        }
361    }
362
363    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
364        match idx {
365            0 => Some("dict_bytes".to_string()),
366            _ => vortex_panic!("OnPairArray buffer_name index {idx} out of bounds"),
367        }
368    }
369
370    fn with_buffers(
371        &self,
372        array: ArrayView<'_, Self>,
373        buffers: &[BufferHandle],
374    ) -> VortexResult<ArrayParts<Self>> {
375        vortex_ensure!(
376            buffers.len() == 1,
377            "Expected 1 buffer, got {}",
378            buffers.len()
379        );
380        let mut data = array.data().clone();
381        data.dict_bytes = buffers[0].clone();
382        Ok(
383            ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
384                .with_slots(array.slots().iter().cloned().collect()),
385        )
386    }
387
388    fn serialize(
389        array: ArrayView<'_, Self>,
390        _session: &VortexSession,
391    ) -> VortexResult<Option<Vec<u8>>> {
392        let dict_size = u32::try_from(array.dict_offsets().len().saturating_sub(1))
393            .map_err(|_| vortex_err!("OnPair dict_size exceeds u32"))?;
394        let total_tokens = array.codes().len() as u64;
395        Ok(Some(
396            OnPairMetadata {
397                uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(),
398                bits: array.bits(),
399                dict_size,
400                total_tokens,
401                dict_offsets_ptype: array.dict_offsets().dtype().as_ptype().into(),
402                codes_ptype: array.codes().dtype().as_ptype().into(),
403                codes_offsets_ptype: array.codes_offsets().dtype().as_ptype().into(),
404            }
405            .encode_to_vec(),
406        ))
407    }
408
409    fn deserialize(
410        &self,
411        dtype: &DType,
412        len: usize,
413        metadata: &[u8],
414        buffers: &[BufferHandle],
415        children: &dyn ArrayChildren,
416        _session: &VortexSession,
417    ) -> VortexResult<ArrayParts<Self>> {
418        if buffers.len() != 1 {
419            vortex_bail!(InvalidArgument: "Expected 1 buffer, got {}", buffers.len());
420        }
421        let metadata = OnPairMetadata::decode(metadata)?;
422        let uncompressed_ptype = metadata.get_uncompressed_lengths_ptype()?;
423
424        // Slot children. We pass `usize::MAX` for slots whose length we
425        // don't know up front (`dict_offsets` and `codes`). `codes_offsets`
426        // has known length `len + 1`.
427        let dict_offsets_len = metadata.dict_size as usize + 1;
428        let total_tokens = usize::try_from(metadata.total_tokens)
429            .map_err(|_| vortex_err!("total_tokens {} overflows usize", metadata.total_tokens))?;
430        // The cascading compressor may have narrowed any of these integer
431        // children to a tighter ptype; the recorded ptype tells the framework
432        // exactly which dtype to materialise as.
433        let dict_offsets_ptype = PType::try_from(metadata.dict_offsets_ptype).map_err(|_| {
434            vortex_err!("invalid dict_offsets_ptype {}", metadata.dict_offsets_ptype)
435        })?;
436        let codes_ptype = PType::try_from(metadata.codes_ptype)
437            .map_err(|_| vortex_err!("invalid codes_ptype {}", metadata.codes_ptype))?;
438        let codes_offsets_ptype = PType::try_from(metadata.codes_offsets_ptype).map_err(|_| {
439            vortex_err!(
440                "invalid codes_offsets_ptype {}",
441                metadata.codes_offsets_ptype
442            )
443        })?;
444        let dict_offsets = children.get(
445            0,
446            &DType::Primitive(dict_offsets_ptype, Nullability::NonNullable),
447            dict_offsets_len,
448        )?;
449        let codes = children.get(
450            1,
451            &DType::Primitive(codes_ptype, Nullability::NonNullable),
452            total_tokens,
453        )?;
454        let codes_offsets = children.get(
455            2,
456            &DType::Primitive(codes_offsets_ptype, Nullability::NonNullable),
457            len + 1,
458        )?;
459        let uncompressed_lengths = children.get(
460            3,
461            &DType::Primitive(uncompressed_ptype, Nullability::NonNullable),
462            len,
463        )?;
464        let validity = match children.len() {
465            4 => Validity::from(dtype.nullability()),
466            5 => Validity::Array(children.get(4, &Validity::DTYPE, len)?),
467            other => vortex_bail!(InvalidArgument: "Expected 4 or 5 children, got {other}"),
468        };
469
470        let data = OnPairData::new(buffers[0].clone(), metadata.bits, len);
471        let slots = OnPairSlots {
472            dict_offsets,
473            codes,
474            codes_offsets,
475            uncompressed_lengths,
476            validity: validity_to_child(&validity, len),
477        }
478        .into_slots();
479        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
480    }
481
482    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
483        OnPairSlots::NAMES[idx].to_string()
484    }
485
486    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
487        canonicalize_onpair(array.as_view(), ctx).map(ExecutionResult::done)
488    }
489
490    fn append_to_builder(
491        array: ArrayView<'_, Self>,
492        builder: &mut dyn ArrayBuilder,
493        ctx: &mut ExecutionCtx,
494    ) -> VortexResult<()> {
495        let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
496            return array
497                .array()
498                .clone()
499                .execute::<Canonical>(ctx)?
500                .into_array()
501                .append_to_builder(builder, ctx);
502        };
503
504        let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress());
505        let (buffers, views) = onpair_decode_views(array, next_buffer_index, ctx)?;
506        builder.push_buffer_and_adjusted_views(
507            &buffers,
508            &views,
509            array
510                .array()
511                .validity()?
512                .execute_mask(array.array().len(), ctx)?,
513        );
514        Ok(())
515    }
516
517    fn reduce_parent(
518        array: ArrayView<'_, Self>,
519        parent: &ArrayRef,
520        child_idx: usize,
521    ) -> VortexResult<Option<ArrayRef>> {
522        RULES.evaluate(array, parent, child_idx)
523    }
524}
525
526impl ValidityVTable<OnPair> for OnPair {
527    fn validity(array: ArrayView<'_, OnPair>) -> VortexResult<Validity> {
528        Ok(child_to_validity(
529            array.slots()[OnPairSlots::VALIDITY].as_ref(),
530            array.dtype().nullability(),
531        ))
532    }
533}
534
535/// Convenience methods on top of the macro-generated [`OnPairArraySlotsExt`].
536pub trait OnPairArrayExt: OnPairArraySlotsExt {
537    fn array_validity(&self) -> Validity {
538        child_to_validity(
539            self.as_ref().slots()[OnPairSlots::VALIDITY].as_ref(),
540            self.as_ref().dtype().nullability(),
541        )
542    }
543}
544
545impl<T: OnPairArraySlotsExt> OnPairArrayExt for T {}
vortex_onpair/array.rs

vortex_onpair/
array.rs
