vortex_btrblocks/schemes/string/
onpair.rs1use 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#[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 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
133fn 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
151const OFFSETS_DELTA_MIN_LEN: usize = 2048;
155
156fn 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}