vortex_btrblocks/schemes/string/
onpair.rs1use 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#[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 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
127fn 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
145const OFFSETS_DELTA_MIN_LEN: usize = 2048;
149
150fn 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}