vortex_btrblocks/
canonical_compressor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! BtrBlocks-specific compressor wrapping the generic [`CascadingCompressor`].
5
6use std::ops::Deref;
7
8use vortex_array::ArrayRef;
9use vortex_array::ExecutionCtx;
10use vortex_error::VortexResult;
11
12use crate::BtrBlocksCompressorBuilder;
13use crate::CascadingCompressor;
14
15/// The BtrBlocks-style compressor with all built-in schemes pre-registered.
16///
17/// This is a thin wrapper around [`CascadingCompressor`] that provides a default set of
18/// compression schemes via [`BtrBlocksCompressorBuilder`].
19///
20/// # Examples
21///
22/// ```rust
23/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt};
24/// use vortex_btrblocks::schemes::integer::IntDictScheme;
25///
26/// // Default compressor - all schemes allowed.
27/// let compressor = BtrBlocksCompressor::default();
28///
29/// // Remove specific schemes using the builder.
30/// let compressor = BtrBlocksCompressorBuilder::default()
31///     .exclude_schemes([IntDictScheme.id()])
32///     .build();
33/// ```
34#[derive(Clone)]
35pub struct BtrBlocksCompressor(
36    /// The underlying cascading compressor.
37    pub CascadingCompressor,
38);
39
40impl BtrBlocksCompressor {
41    /// Compresses an array using BtrBlocks-inspired compression.
42    pub fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
43        self.0.compress(array, ctx)
44    }
45}
46
47impl Deref for BtrBlocksCompressor {
48    type Target = CascadingCompressor;
49
50    fn deref(&self) -> &CascadingCompressor {
51        &self.0
52    }
53}
54
55impl Default for BtrBlocksCompressor {
56    fn default() -> Self {
57        BtrBlocksCompressorBuilder::default().build()
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use std::sync::LazyLock;
64
65    use rstest::rstest;
66    use vortex_array::IntoArray;
67    use vortex_array::VortexSessionExecute;
68    use vortex_array::arrays::BoolArray;
69    use vortex_array::arrays::Constant;
70    use vortex_array::arrays::Dict;
71    use vortex_array::arrays::List;
72    use vortex_array::arrays::ListView;
73    use vortex_array::arrays::ListViewArray;
74    use vortex_array::arrays::VarBinViewArray;
75    use vortex_array::assert_arrays_eq;
76    use vortex_array::dtype::DType;
77    use vortex_array::dtype::Nullability;
78    use vortex_array::validity::Validity;
79    use vortex_buffer::BitBuffer;
80    use vortex_buffer::buffer;
81    use vortex_error::VortexResult;
82    use vortex_session::VortexSession;
83
84    use crate::BtrBlocksCompressor;
85    #[cfg(feature = "zstd")]
86    use crate::BtrBlocksCompressorBuilder;
87
88    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
89
90    #[rstest]
91    #[case::zctl(
92        unsafe {
93            ListViewArray::new_unchecked(
94                buffer![1i32, 2, 3, 4, 5].into_array(),
95                buffer![0i32, 3].into_array(),
96                buffer![3i32, 2].into_array(),
97                Validity::NonNullable,
98            ).with_zero_copy_to_list(true)
99        },
100        true,
101    )]
102    #[case::overlapping(
103        ListViewArray::new(
104            buffer![1i32, 2, 3].into_array(),
105            buffer![0i32, 0, 0].into_array(),
106            buffer![3i32, 3, 3].into_array(),
107            Validity::NonNullable,
108        ),
109        false,
110    )]
111    fn listview_compress_roundtrip(
112        #[case] input: ListViewArray,
113        #[case] expect_list: bool,
114    ) -> VortexResult<()> {
115        let mut ctx = SESSION.create_execution_ctx();
116        let array_ref = input.clone().into_array();
117        let result = BtrBlocksCompressor::default()
118            .compress(&array_ref, &mut SESSION.create_execution_ctx())?;
119        if expect_list {
120            assert!(result.as_opt::<List>().is_some());
121        } else {
122            assert!(result.as_opt::<ListView>().is_some());
123        }
124        assert_arrays_eq!(result, input, &mut ctx);
125        Ok(())
126    }
127
128    #[test]
129    fn test_constant_all_true() -> VortexResult<()> {
130        let mut ctx = SESSION.create_execution_ctx();
131        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable);
132        let btr = BtrBlocksCompressor::default();
133        let compressed = btr.compress(
134            &array.clone().into_array(),
135            &mut SESSION.create_execution_ctx(),
136        )?;
137        assert!(compressed.is::<Constant>());
138        assert_arrays_eq!(compressed, array, &mut ctx);
139        Ok(())
140    }
141
142    #[test]
143    fn test_constant_all_false() -> VortexResult<()> {
144        let mut ctx = SESSION.create_execution_ctx();
145        let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable);
146        let btr = BtrBlocksCompressor::default();
147        let compressed = btr.compress(
148            &array.clone().into_array(),
149            &mut SESSION.create_execution_ctx(),
150        )?;
151        assert!(compressed.is::<Constant>());
152        assert_arrays_eq!(compressed, array, &mut ctx);
153        Ok(())
154    }
155
156    #[test]
157    fn test_nullable_all_valid_compressed() -> VortexResult<()> {
158        let mut ctx = SESSION.create_execution_ctx();
159        let array = BoolArray::new(
160            BitBuffer::from(vec![true; 100]),
161            Validity::from(BitBuffer::from(vec![true; 100])),
162        );
163        let btr = BtrBlocksCompressor::default();
164        let compressed = btr.compress(
165            &array.clone().into_array(),
166            &mut SESSION.create_execution_ctx(),
167        )?;
168        assert!(compressed.is::<Constant>());
169        assert_arrays_eq!(compressed, array, &mut ctx);
170        Ok(())
171    }
172
173    #[test]
174    fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> {
175        let mut ctx = SESSION.create_execution_ctx();
176        let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0)));
177        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity);
178        let btr = BtrBlocksCompressor::default();
179        let compressed = btr.compress(
180            &array.clone().into_array(),
181            &mut SESSION.create_execution_ctx(),
182        )?;
183        assert!(!compressed.is::<Constant>());
184        assert_arrays_eq!(compressed, array, &mut ctx);
185        Ok(())
186    }
187
188    #[test]
189    fn test_mixed_not_constant() -> VortexResult<()> {
190        let mut ctx = SESSION.create_execution_ctx();
191        let array = BoolArray::new(
192            BitBuffer::from(vec![true, false, true, false, true]),
193            Validity::NonNullable,
194        );
195        let btr = BtrBlocksCompressor::default();
196        let compressed = btr.compress(
197            &array.clone().into_array(),
198            &mut SESSION.create_execution_ctx(),
199        )?;
200        assert!(!compressed.is::<Constant>());
201        assert_arrays_eq!(compressed, array, &mut ctx);
202        Ok(())
203    }
204
205    #[test]
206    fn test_binary_constant_compressed() -> VortexResult<()> {
207        let mut ctx = SESSION.create_execution_ctx();
208        let values = vec![Some(b"constant-bytes".as_slice()); 100];
209        let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable));
210        let btr = BtrBlocksCompressor::default();
211        let compressed = btr.compress(
212            &array.clone().into_array(),
213            &mut SESSION.create_execution_ctx(),
214        )?;
215        assert!(compressed.is::<Constant>());
216        assert_arrays_eq!(compressed, array, &mut ctx);
217        Ok(())
218    }
219
220    #[test]
221    fn test_binary_dict_compressed() -> VortexResult<()> {
222        let mut ctx = SESSION.create_execution_ctx();
223        let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"];
224        let values = (0..1000)
225            .map(|idx| Some(distinct_values[idx % distinct_values.len()]))
226            .collect::<Vec<_>>();
227        let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable));
228        let btr = BtrBlocksCompressor::default();
229        let compressed = btr.compress(
230            &array.clone().into_array(),
231            &mut SESSION.create_execution_ctx(),
232        )?;
233        assert!(compressed.is::<Dict>());
234        assert_arrays_eq!(compressed, array, &mut ctx);
235        Ok(())
236    }
237
238    #[cfg(feature = "zstd")]
239    #[test]
240    fn test_compact_binary_zstd_compressed() -> VortexResult<()> {
241        let values = (0..1024)
242            .map(|idx| {
243                let mut value = Vec::from(&b"common binary payload prefix "[..]);
244                value.extend_from_slice(&(idx as u32).to_le_bytes());
245                value.extend_from_slice(&[b'x'; 96]);
246                value
247            })
248            .collect::<Vec<_>>();
249        let array = VarBinViewArray::from_iter(
250            values.iter().map(|value| Some(value.as_slice())),
251            DType::Binary(Nullability::NonNullable),
252        );
253
254        let compressor = BtrBlocksCompressorBuilder::default().with_compact().build();
255        let mut ctx = SESSION.create_execution_ctx();
256        let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?;
257
258        assert!(
259            compressed.is::<vortex_zstd::Zstd>(),
260            "expected Zstd, got {}",
261            compressed.encoding_id()
262        );
263        assert_arrays_eq!(compressed, array, &mut ctx);
264        Ok(())
265    }
266
267    #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
268    #[test]
269    fn test_cuda_compatible_binary_zstd_buffers_compressed() -> VortexResult<()> {
270        let values = (0..1024)
271            .map(|idx| {
272                let mut value = Vec::from(&b"common binary payload prefix "[..]);
273                value.extend_from_slice(&(idx as u32).to_le_bytes());
274                value.extend_from_slice(&[b'x'; 96]);
275                value
276            })
277            .collect::<Vec<_>>();
278        let array = VarBinViewArray::from_iter(
279            values.iter().map(|value| Some(value.as_slice())),
280            DType::Binary(Nullability::NonNullable),
281        );
282
283        let compressor = BtrBlocksCompressorBuilder::default()
284            .only_cuda_compatible()
285            .build();
286        let mut ctx = SESSION.create_execution_ctx();
287        let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?;
288
289        assert!(
290            compressed.is::<vortex_zstd::ZstdBuffers>(),
291            "expected ZstdBuffers, got {}",
292            compressed.encoding_id()
293        );
294        assert_arrays_eq!(compressed, array, &mut ctx);
295        Ok(())
296    }
297}
vortex_btrblocks/canonical_compressor.rs

vortex_btrblocks/
canonical_compressor.rs
