vortex_onpair/
compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Train + compress entry points for the OnPair encoding.
5
6use onpair::Config;
7use onpair::Offset;
8use vortex_array::ArrayRef;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::arrays::ConstantArray;
12use vortex_array::arrays::VarBinViewArray;
13use vortex_array::arrays::varbinview::BinaryView;
14use vortex_array::buffer::BufferHandle;
15use vortex_array::validity::Validity;
16use vortex_buffer::Alignment;
17use vortex_buffer::Buffer;
18use vortex_buffer::BufferMut;
19use vortex_buffer::ByteBuffer;
20use vortex_buffer::ByteBufferMut;
21use vortex_error::VortexExpect;
22use vortex_error::VortexResult;
23use vortex_error::vortex_err;
24use vortex_mask::AllOr;
25
26use crate::OnPair;
27use crate::OnPairArray;
28
29/// Default OnPair training configuration: 12-bit codes ("dict-12").
30pub const DEFAULT_DICT12_CONFIG: Config = onpair::DEFAULT_CONFIG;
31
32fn onpair_compress_varbinview<O>(
33    array: VarBinViewArray,
34    config: Config,
35    ctx: &mut ExecutionCtx,
36) -> VortexResult<OnPairArray>
37where
38    O: Offset,
39{
40    let len = array.len();
41    let mask = array.validity()?.execute_mask(len, ctx)?;
42    if mask.all_false() {
43        return OnPair::try_new(
44            array.dtype().clone(),
45            BufferHandle::new_host(ByteBuffer::empty()),
46            ConstantArray::new(0, len).into_array(),
47            ConstantArray::new(0u16, len).into_array(),
48            ConstantArray::new(0u32, len + 1).into_array(),
49            ConstantArray::new(0i32, len).into_array(),
50            Validity::AllInvalid,
51            9,
52        );
53    }
54
55    let mut flat: Vec<u8> = Vec::with_capacity(len * 16);
56    let mut offsets: Vec<O> = Vec::with_capacity(len + 1);
57    let mut uncompressed_lengths: BufferMut<i32> = BufferMut::with_capacity(len);
58    offsets.push(O::from_usize(0));
59    let views = array.views();
60    let buffers = array
61        .data_buffers()
62        .as_ref()
63        .iter()
64        .map(|b| b.as_host())
65        .collect::<Vec<_>>();
66
67    match mask.bit_buffer() {
68        AllOr::All => {
69            for view in views {
70                let bytes = view_bytes(view, &buffers);
71                flat.extend_from_slice(bytes);
72                offsets.push(O::from_usize(flat.len()));
73                uncompressed_lengths
74                    .push(i32::try_from(view.len()).vortex_expect("must fit in i32"));
75            }
76        }
77        AllOr::None => {
78            unreachable!("all_false() should have been caught earlier");
79        }
80        AllOr::Some(validity) => {
81            for (view, valid) in views.iter().zip(validity.iter()) {
82                if valid {
83                    let bytes = view_bytes(view, &buffers);
84                    flat.extend_from_slice(bytes);
85                    offsets.push(O::from_usize(flat.len()));
86                    uncompressed_lengths
87                        .push(i32::try_from(view.len()).vortex_expect("must fit in i32"));
88                } else {
89                    offsets.push(O::from_usize(flat.len()));
90                    uncompressed_lengths.push(0);
91                }
92            }
93        }
94    }
95
96    let column = onpair::compress(&flat, &offsets, config)
97        .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?;
98    let bits = column.bits;
99    let dict_bytes = dict_bytes_to_buffer(column.dict_bytes);
100    let codes_offsets =
101        build_codes_offsets(&column.codes, &column.dict_offsets, &offsets)?.into_array();
102    let codes = Buffer::from(column.codes).into_array();
103    let dict_offsets = Buffer::from(column.dict_offsets).into_array();
104
105    let uncompressed_lengths = uncompressed_lengths.into_array();
106
107    OnPair::try_new(
108        array.dtype().clone(),
109        dict_bytes,
110        dict_offsets,
111        codes,
112        codes_offsets,
113        uncompressed_lengths,
114        array.validity()?,
115        bits,
116    )
117}
118
119fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] {
120    if view.is_inlined() {
121        view.as_inlined().value()
122    } else {
123        let view_ref = view.as_view();
124        &buffers[view_ref.buffer_index as usize][view_ref.as_range()]
125    }
126}
127
128/// Lift compressed dictionary bytes into the Vortex buffer slot.
129fn dict_bytes_to_buffer(dict_bytes: Vec<u8>) -> BufferHandle {
130    // Pad the dictionary blob with MAX_TOKEN_SIZE zero bytes so the
131    // over-copy decoder can issue a fixed 16-byte load for every token
132    // without risking an OOB read on the last entry.
133    //
134    // Align dict_bytes to 8 bytes so the segment that ultimately holds the
135    // OnPair tree starts at an 8-aligned in-memory address. Without this
136    // anchor, the per-buffer padding the serializer inserts is only
137    // *relative* to the segment start; if the segment lands at a u8-aligned
138    // heap address, downstream `PrimitiveArray<u32>::deserialize` panics
139    // with `Misaligned buffer cannot be used to build PrimitiveArray of u32`.
140    let mut padded = ByteBufferMut::with_capacity_aligned(
141        dict_bytes.len() + onpair::MAX_TOKEN_SIZE,
142        Alignment::new(8),
143    );
144    padded.extend_from_slice(&dict_bytes);
145    unsafe { padded.push_n_unchecked(0, dict_bytes.len() + onpair::MAX_TOKEN_SIZE - padded.len()) };
146    BufferHandle::new_host(padded.freeze())
147}
148
149/// Reconstruct the per-row `codes_offsets` from the flat `codes`, the
150/// dictionary `dict_offsets` (token byte lengths) and the per-row decoded byte
151/// boundaries. Returns `nrows + 1` cumulative code counts (`u32`).
152// TODO(joe): can we compute this while compressing the array, yes but a worse API.
153fn build_codes_offsets<O: Offset>(
154    codes: &[u16],
155    dict_offsets: &[u32],
156    row_byte_offsets: &[O],
157) -> VortexResult<Buffer<u32>> {
158    let nrows = row_byte_offsets.len() - 1;
159    let mut codes_offsets = BufferMut::with_capacity(nrows + 1);
160    codes_offsets.push(0u32);
161    let mut decoded_bytes: u64 = 0;
162    let mut code_idx: usize = 0;
163    for r in 0..nrows {
164        let target = row_byte_offsets[r + 1]
165            .to_usize()
166            .ok_or_else(|| vortex_err!("OnPair row byte offset does not fit usize"))?
167            as u64;
168        while decoded_bytes < target {
169            let code = codes[code_idx] as usize;
170            decoded_bytes += u64::from(dict_offsets[code + 1] - dict_offsets[code]);
171            code_idx += 1;
172        }
173        codes_offsets.push(
174            u32::try_from(code_idx)
175                .map_err(|_| vortex_err!("OnPair: code boundary {code_idx} does not fit u32"))?,
176        );
177    }
178    Ok(codes_offsets.freeze())
179}
180
181/// Compress any [`ArrayRef`] whose canonical form is a string array, by first
182/// canonicalising to `VarBinViewArray`.
183pub fn onpair_compress(
184    array: &ArrayRef,
185    config: Config,
186    ctx: &mut ExecutionCtx,
187) -> VortexResult<OnPairArray> {
188    let view = array.clone().execute::<VarBinViewArray>(ctx)?;
189    onpair_compress_varbinview::<u64>(view, config, ctx)
190}
vortex_onpair/compress.rs

vortex_onpair/
compress.rs
