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    #[cfg(feature = "zstd")]
67    use vortex_array::ArrayId;
68    #[cfg(feature = "zstd")]
69    use vortex_array::ArrayPlugin;
70    use vortex_array::IntoArray;
71    use vortex_array::VortexSessionExecute;
72    use vortex_array::arrays::BoolArray;
73    use vortex_array::arrays::Constant;
74    use vortex_array::arrays::Dict;
75    use vortex_array::arrays::List;
76    use vortex_array::arrays::ListView;
77    use vortex_array::arrays::ListViewArray;
78    use vortex_array::arrays::VarBinViewArray;
79    use vortex_array::assert_arrays_eq;
80    use vortex_array::dtype::DType;
81    use vortex_array::dtype::Nullability;
82    use vortex_array::validity::Validity;
83    use vortex_buffer::BitBuffer;
84    use vortex_buffer::buffer;
85    use vortex_error::VortexResult;
86    use vortex_session::VortexSession;
87    #[cfg(feature = "zstd")]
88    use vortex_utils::aliases::hash_set::HashSet;
89
90    use crate::BtrBlocksCompressor;
91    #[cfg(feature = "zstd")]
92    use crate::BtrBlocksCompressorBuilder;
93
94    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
95
96    #[rstest]
97    #[case::zctl(
98        unsafe {
99            ListViewArray::new_unchecked(
100                buffer![1i32, 2, 3, 4, 5].into_array(),
101                buffer![0i32, 3].into_array(),
102                buffer![3i32, 2].into_array(),
103                Validity::NonNullable,
104            ).with_zero_copy_to_list(true)
105        },
106        true,
107    )]
108    #[case::overlapping(
109        ListViewArray::new(
110            buffer![1i32, 2, 3].into_array(),
111            buffer![0i32, 0, 0].into_array(),
112            buffer![3i32, 3, 3].into_array(),
113            Validity::NonNullable,
114        ),
115        false,
116    )]
117    fn listview_compress_roundtrip(
118        #[case] input: ListViewArray,
119        #[case] expect_list: bool,
120    ) -> VortexResult<()> {
121        let mut ctx = SESSION.create_execution_ctx();
122        let array_ref = input.clone().into_array();
123        let result = BtrBlocksCompressor::default()
124            .compress(&array_ref, &mut SESSION.create_execution_ctx())?;
125        if expect_list {
126            assert!(result.as_opt::<List>().is_some());
127        } else {
128            assert!(result.as_opt::<ListView>().is_some());
129        }
130        assert_arrays_eq!(result, input, &mut ctx);
131        Ok(())
132    }
133
134    #[test]
135    fn test_constant_all_true() -> VortexResult<()> {
136        let mut ctx = SESSION.create_execution_ctx();
137        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable);
138        let btr = BtrBlocksCompressor::default();
139        let compressed = btr.compress(
140            &array.clone().into_array(),
141            &mut SESSION.create_execution_ctx(),
142        )?;
143        assert!(compressed.is::<Constant>());
144        assert_arrays_eq!(compressed, array, &mut ctx);
145        Ok(())
146    }
147
148    #[test]
149    fn test_constant_all_false() -> VortexResult<()> {
150        let mut ctx = SESSION.create_execution_ctx();
151        let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable);
152        let btr = BtrBlocksCompressor::default();
153        let compressed = btr.compress(
154            &array.clone().into_array(),
155            &mut SESSION.create_execution_ctx(),
156        )?;
157        assert!(compressed.is::<Constant>());
158        assert_arrays_eq!(compressed, array, &mut ctx);
159        Ok(())
160    }
161
162    #[test]
163    fn test_nullable_all_valid_compressed() -> VortexResult<()> {
164        let mut ctx = SESSION.create_execution_ctx();
165        let array = BoolArray::new(
166            BitBuffer::from(vec![true; 100]),
167            Validity::from(BitBuffer::from(vec![true; 100])),
168        );
169        let btr = BtrBlocksCompressor::default();
170        let compressed = btr.compress(
171            &array.clone().into_array(),
172            &mut SESSION.create_execution_ctx(),
173        )?;
174        assert!(compressed.is::<Constant>());
175        assert_arrays_eq!(compressed, array, &mut ctx);
176        Ok(())
177    }
178
179    #[test]
180    fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> {
181        let mut ctx = SESSION.create_execution_ctx();
182        let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0)));
183        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity);
184        let btr = BtrBlocksCompressor::default();
185        let compressed = btr.compress(
186            &array.clone().into_array(),
187            &mut SESSION.create_execution_ctx(),
188        )?;
189        assert!(!compressed.is::<Constant>());
190        assert_arrays_eq!(compressed, array, &mut ctx);
191        Ok(())
192    }
193
194    #[test]
195    fn test_mixed_not_constant() -> VortexResult<()> {
196        let mut ctx = SESSION.create_execution_ctx();
197        let array = BoolArray::new(
198            BitBuffer::from(vec![true, false, true, false, true]),
199            Validity::NonNullable,
200        );
201        let btr = BtrBlocksCompressor::default();
202        let compressed = btr.compress(
203            &array.clone().into_array(),
204            &mut SESSION.create_execution_ctx(),
205        )?;
206        assert!(!compressed.is::<Constant>());
207        assert_arrays_eq!(compressed, array, &mut ctx);
208        Ok(())
209    }
210
211    #[test]
212    fn test_binary_constant_compressed() -> VortexResult<()> {
213        let mut ctx = SESSION.create_execution_ctx();
214        let values = vec![Some(b"constant-bytes".as_slice()); 100];
215        let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable));
216        let btr = BtrBlocksCompressor::default();
217        let compressed = btr.compress(
218            &array.clone().into_array(),
219            &mut SESSION.create_execution_ctx(),
220        )?;
221        assert!(compressed.is::<Constant>());
222        assert_arrays_eq!(compressed, array, &mut ctx);
223        Ok(())
224    }
225
226    #[test]
227    fn test_binary_dict_compressed() -> VortexResult<()> {
228        let mut ctx = SESSION.create_execution_ctx();
229        let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"];
230        let values = (0..1000)
231            .map(|idx| Some(distinct_values[idx % distinct_values.len()]))
232            .collect::<Vec<_>>();
233        let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable));
234        let btr = BtrBlocksCompressor::default();
235        let compressed = btr.compress(
236            &array.clone().into_array(),
237            &mut SESSION.create_execution_ctx(),
238        )?;
239        assert!(compressed.is::<Dict>());
240        assert_arrays_eq!(compressed, array, &mut ctx);
241        Ok(())
242    }
243
244    #[cfg(feature = "zstd")]
245    #[test]
246    fn test_compact_binary_zstd_compressed() -> VortexResult<()> {
247        let values = (0..1024)
248            .map(|idx| {
249                let mut value = Vec::from(&b"common binary payload prefix "[..]);
250                value.extend_from_slice(&(idx as u32).to_le_bytes());
251                value.extend_from_slice(&[b'x'; 96]);
252                value
253            })
254            .collect::<Vec<_>>();
255        let array = VarBinViewArray::from_iter(
256            values.iter().map(|value| Some(value.as_slice())),
257            DType::Binary(Nullability::NonNullable),
258        );
259
260        let compressor = BtrBlocksCompressorBuilder::default().with_compact().build();
261        let mut ctx = SESSION.create_execution_ctx();
262        let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?;
263
264        assert!(
265            compressed.is::<vortex_zstd::Zstd>(),
266            "expected Zstd, got {}",
267            compressed.encoding_id()
268        );
269        assert_arrays_eq!(compressed, array, &mut ctx);
270        Ok(())
271    }
272
273    #[cfg(feature = "zstd")]
274    #[rstest]
275    #[case::array_level(vortex_zstd::Zstd.id())]
276    #[case::buffer_level(vortex_zstd::ZstdBuffers.id())]
277    fn test_cuda_compatible_binary_zstd_follows_editions(
278        #[case] allowed: ArrayId,
279    ) -> VortexResult<()> {
280        let values = (0..1024)
281            .map(|idx| {
282                let mut value = Vec::from(&b"common binary payload prefix "[..]);
283                value.extend_from_slice(&(idx as u32).to_le_bytes());
284                value.extend_from_slice(&[b'x'; 96]);
285                value
286            })
287            .collect::<Vec<_>>();
288        let array = VarBinViewArray::from_iter(
289            values.iter().map(|value| Some(value.as_slice())),
290            DType::Binary(Nullability::NonNullable),
291        );
292
293        // The CUDA preset carries both Zstd schemes; the edition filter decides which one
294        // survives.
295        let compressor = BtrBlocksCompressorBuilder::default()
296            .only_cuda_compatible()
297            .retain_allowed_encodings(&HashSet::from([allowed]))
298            .build();
299        let mut ctx = SESSION.create_execution_ctx();
300        let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?;
301
302        assert_eq!(compressed.encoding_id(), allowed);
303        assert_arrays_eq!(compressed, array, &mut ctx);
304        Ok(())
305    }
306}
vortex_btrblocks/canonical_compressor.rs

vortex_btrblocks/
canonical_compressor.rs
