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