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;
29
30/// OnPair short-string compression (dict-12).
31///
32/// A default string-fragmentation scheme (alongside [`super::FSSTScheme`]) —
33/// targets large columns of short-to-medium strings with high lexical
34/// overlap, like URLs or log lines. Uses a learned dictionary of frequent
35/// adjacent substrings (built by the OnPair trainer at compress time) and
36/// 12-bit token codes stored as a u16 child, with offsets /
37/// uncompressed-lengths flowing through the cascading compressor like any
38/// other primitive children.
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40pub struct OnPairScheme;
41
42impl Scheme for OnPairScheme {
43    fn scheme_name(&self) -> &'static str {
44        "vortex.string.onpair"
45    }
46
47    fn matches(&self, canonical: &Canonical) -> bool {
48        canonical.dtype().is_utf8()
49    }
50
51    fn produced_encodings(&self) -> Vec<ArrayId> {
52        vec![OnPair.id()]
53    }
54
55    /// 4 primitive slot children flow through the cascading compressor:
56    /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 →
57    /// usually `FastLanes::BitPacked` after scheme selection),
58    /// `codes_offsets` (u32 → `FoR`), `uncompressed_lengths` (i32 → narrow
59    /// + `FoR`). Validity stays untouched.
60    fn num_children(&self) -> usize {
61        4
62    }
63
64    fn expected_compression_ratio(
65        &self,
66        _data: &ArrayAndStats,
67        _compress_ctx: CompressorContext,
68        _exec_ctx: &mut ExecutionCtx,
69    ) -> CompressionEstimate {
70        CompressionEstimate::Deferred(DeferredEstimate::Sample)
71    }
72
73    fn compress(
74        &self,
75        compressor: &CascadingCompressor,
76        data: &ArrayAndStats,
77        compress_ctx: CompressorContext,
78        exec_ctx: &mut ExecutionCtx,
79    ) -> VortexResult<ArrayRef> {
80        let utf8 = data.array_as_varbinview().into_owned();
81        let encoded = onpair_compress(utf8.as_array(), DEFAULT_CONFIG, exec_ctx)?;
82        let Some(onpair_array) = encoded.as_opt::<OnPair>() else {
83            return Ok(encoded);
84        };
85
86        let dict_offsets = compress_primitive_child(
87            compressor,
88            onpair_array.dict_offsets(),
89            &compress_ctx,
90            self.id(),
91            0,
92            exec_ctx,
93        )?;
94        let codes = compress_primitive_child(
95            compressor,
96            onpair_array.codes(),
97            &compress_ctx,
98            self.id(),
99            1,
100            exec_ctx,
101        )?;
102        let codes_offsets = compress_primitive_child(
103            compressor,
104            onpair_array.codes_offsets(),
105            &compress_ctx,
106            self.id(),
107            2,
108            exec_ctx,
109        )?;
110        let uncompressed_lengths = compress_primitive_child(
111            compressor,
112            onpair_array.uncompressed_lengths(),
113            &compress_ctx,
114            self.id(),
115            3,
116            exec_ctx,
117        )?;
118
119        Ok(OnPair::try_new_with_data(
120            onpair_array.dtype().clone(),
121            onpair_array.data().clone(),
122            dict_offsets,
123            codes,
124            codes_offsets,
125            uncompressed_lengths,
126            onpair_array.array_validity(),
127        )?
128        .into_array())
129    }
130}
131
132/// Narrow a primitive child to its tightest int type, then forward it to
133/// the cascading compressor.
134fn compress_primitive_child(
135    compressor: &CascadingCompressor,
136    child: &ArrayRef,
137    compress_ctx: &CompressorContext,
138    scheme_id: SchemeId,
139    child_idx: usize,
140    exec_ctx: &mut ExecutionCtx,
141) -> VortexResult<ArrayRef> {
142    let narrowed = child
143        .clone()
144        .execute::<PrimitiveArray>(exec_ctx)?
145        .narrow(exec_ctx)?
146        .into_array();
147    compressor.compress_child(&narrowed, compress_ctx, scheme_id, child_idx, exec_ctx)
148}
149
150// TODO(joe): re-apply Delta to the monotone `dict_offsets` / `codes_offsets` children once the
151// compressor can say whether a scheme is eligible for a given child; applying it by hand
152// bypassed the exclusion rules.
vortex_btrblocks/schemes/string/onpair.rs

vortex_btrblocks/schemes/string/
onpair.rs
