Skip to main content

vortex_btrblocks/schemes/string/
onpair.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! OnPair short-string compression (dict-12).
5
6use vortex_array::ArrayRef;
7use vortex_array::Canonical;
8use vortex_array::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::arrays::PrimitiveArray;
11use vortex_array::arrays::primitive::PrimitiveArrayExt;
12use vortex_compressor::scheme::CompressionEstimate;
13use vortex_compressor::scheme::DeferredEstimate;
14use vortex_compressor::scheme::SchemeId;
15use vortex_error::VortexResult;
16use vortex_onpair::DEFAULT_DICT12_CONFIG;
17use vortex_onpair::OnPair;
18use vortex_onpair::OnPairArrayExt;
19use vortex_onpair::OnPairArraySlotsExt;
20use vortex_onpair::onpair_compress;
21
22use crate::ArrayAndStats;
23use crate::CascadingCompressor;
24use crate::CompressorContext;
25use crate::Scheme;
26use crate::SchemeExt;
27use crate::schemes::integer::try_compress_delta;
28
29/// OnPair short-string compression (dict-12).
30///
31/// A default string-fragmentation scheme (alongside [`super::FSSTScheme`]) —
32/// targets large columns of short-to-medium strings with high lexical
33/// overlap, like URLs or log lines. Uses a learned dictionary of frequent
34/// adjacent substrings (built by the OnPair trainer at compress time) and
35/// 12-bit token codes stored as a u16 child, with offsets /
36/// uncompressed-lengths flowing through the cascading compressor like any
37/// other primitive children.
38#[derive(Debug, Copy, Clone, PartialEq, Eq)]
39pub struct OnPairScheme;
40
41impl Scheme for OnPairScheme {
42    fn scheme_name(&self) -> &'static str {
43        "vortex.string.onpair"
44    }
45
46    fn matches(&self, canonical: &Canonical) -> bool {
47        canonical.dtype().is_utf8()
48    }
49
50    /// 4 primitive slot children flow through the cascading compressor:
51    /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 →
52    /// usually `FastLanes::BitPacked` after scheme selection),
53    /// `codes_offsets` (u32 → `FoR`), `uncompressed_lengths` (i32 → narrow
54    /// + `FoR`). Validity stays untouched.
55    fn num_children(&self) -> usize {
56        4
57    }
58
59    fn expected_compression_ratio(
60        &self,
61        _data: &ArrayAndStats,
62        _compress_ctx: CompressorContext,
63        _exec_ctx: &mut ExecutionCtx,
64    ) -> CompressionEstimate {
65        CompressionEstimate::Deferred(DeferredEstimate::Sample)
66    }
67
68    fn compress(
69        &self,
70        compressor: &CascadingCompressor,
71        data: &ArrayAndStats,
72        compress_ctx: CompressorContext,
73        exec_ctx: &mut ExecutionCtx,
74    ) -> VortexResult<ArrayRef> {
75        let utf8 = data.array_as_varbinview().into_owned();
76        let encoded = onpair_compress(utf8.as_array(), DEFAULT_DICT12_CONFIG, exec_ctx)?;
77        let Some(onpair_array) = encoded.as_opt::<OnPair>() else {
78            return Ok(encoded);
79        };
80
81        let dict_offsets = compress_offsets_child(
82            compressor,
83            onpair_array.dict_offsets(),
84            &compress_ctx,
85            self.id(),
86            0,
87            exec_ctx,
88        )?;
89        let codes = compress_primitive_child(
90            compressor,
91            onpair_array.codes(),
92            &compress_ctx,
93            self.id(),
94            1,
95            exec_ctx,
96        )?;
97        let codes_offsets = compress_offsets_child(
98            compressor,
99            onpair_array.codes_offsets(),
100            &compress_ctx,
101            self.id(),
102            2,
103            exec_ctx,
104        )?;
105        let uncompressed_lengths = compress_primitive_child(
106            compressor,
107            onpair_array.uncompressed_lengths(),
108            &compress_ctx,
109            self.id(),
110            3,
111            exec_ctx,
112        )?;
113
114        Ok(OnPair::try_new(
115            onpair_array.dtype().clone(),
116            onpair_array.dict_bytes_handle().clone(),
117            dict_offsets,
118            codes,
119            codes_offsets,
120            uncompressed_lengths,
121            onpair_array.array_validity(),
122        )?
123        .into_array())
124    }
125}
126
127/// Narrow a primitive child to its tightest int type, then forward it to
128/// the cascading compressor.
129fn compress_primitive_child(
130    compressor: &CascadingCompressor,
131    child: &ArrayRef,
132    compress_ctx: &CompressorContext,
133    scheme_id: SchemeId,
134    child_idx: usize,
135    exec_ctx: &mut ExecutionCtx,
136) -> VortexResult<ArrayRef> {
137    let narrowed = child
138        .clone()
139        .execute::<PrimitiveArray>(exec_ctx)?
140        .narrow(exec_ctx)?
141        .into_array();
142    compressor.compress_child(&narrowed, compress_ctx, scheme_id, child_idx, exec_ctx)
143}
144
145/// Minimum child length before delta is even attempted. Delta carries fixed
146/// overhead (a separate `bases` array plus FastLanes' 1024-element lane
147/// packing), so on short children it can only lose.
148const OFFSETS_DELTA_MIN_LEN: usize = 2048;
149
150/// Compress a monotonic offsets child. For children of at least
151/// [`OFFSETS_DELTA_MIN_LEN`] it tries both the normal cascading path and a
152/// delta path and keeps whichever produces fewer bytes; shorter children
153/// skip delta entirely. `dict_offsets` and `codes_offsets` are cumulative
154/// (monotonic), so delta (per-entry deltas) usually packs much tighter than
155/// FoR+bitpacking over the full range.
156fn compress_offsets_child(
157    compressor: &CascadingCompressor,
158    child: &ArrayRef,
159    compress_ctx: &CompressorContext,
160    scheme_id: SchemeId,
161    child_idx: usize,
162    exec_ctx: &mut ExecutionCtx,
163) -> VortexResult<ArrayRef> {
164    let narrowed = child
165        .clone()
166        .execute::<PrimitiveArray>(exec_ctx)?
167        .narrow(exec_ctx)?
168        .into_array();
169    let plain =
170        compressor.compress_child(&narrowed, compress_ctx, scheme_id, child_idx, exec_ctx)?;
171    if narrowed.len() < OFFSETS_DELTA_MIN_LEN {
172        return Ok(plain);
173    }
174    let delta = try_compress_delta(
175        compressor,
176        &narrowed,
177        compress_ctx,
178        scheme_id,
179        child_idx,
180        exec_ctx,
181    )?;
182    if delta.nbytes() < plain.nbytes() {
183        Ok(delta)
184    } else {
185        Ok(plain)
186    }
187}