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_error::VortexResult;
10
11use crate::BtrBlocksCompressorBuilder;
12use crate::CascadingCompressor;
13
14/// The BtrBlocks-style compressor with all built-in schemes pre-registered.
15///
16/// This is a thin wrapper around [`CascadingCompressor`] that provides a default set of
17/// compression schemes via [`BtrBlocksCompressorBuilder`].
18///
19/// # Examples
20///
21/// ```rust
22/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt};
23/// use vortex_btrblocks::schemes::integer::IntDictScheme;
24///
25/// // Default compressor - all schemes allowed.
26/// let compressor = BtrBlocksCompressor::default();
27///
28/// // Remove specific schemes using the builder.
29/// let compressor = BtrBlocksCompressorBuilder::default()
30///     .exclude_schemes([IntDictScheme.id()])
31///     .build();
32/// ```
33#[derive(Clone)]
34pub struct BtrBlocksCompressor(
35    /// The underlying cascading compressor.
36    pub CascadingCompressor,
37);
38
39impl BtrBlocksCompressor {
40    /// Compresses an array using BtrBlocks-inspired compression.
41    pub fn compress(&self, array: &ArrayRef) -> VortexResult<ArrayRef> {
42        self.0.compress(array)
43    }
44}
45
46impl Deref for BtrBlocksCompressor {
47    type Target = CascadingCompressor;
48
49    fn deref(&self) -> &CascadingCompressor {
50        &self.0
51    }
52}
53
54impl Default for BtrBlocksCompressor {
55    fn default() -> Self {
56        BtrBlocksCompressorBuilder::default().build()
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use rstest::rstest;
63    use vortex_array::IntoArray;
64    use vortex_array::arrays::BoolArray;
65    use vortex_array::arrays::Constant;
66    use vortex_array::arrays::List;
67    use vortex_array::arrays::ListView;
68    use vortex_array::arrays::ListViewArray;
69    use vortex_array::assert_arrays_eq;
70    use vortex_array::validity::Validity;
71    use vortex_buffer::BitBuffer;
72    use vortex_buffer::buffer;
73    use vortex_error::VortexResult;
74
75    use crate::BtrBlocksCompressor;
76
77    #[rstest]
78    #[case::zctl(
79        unsafe {
80            ListViewArray::new_unchecked(
81                buffer![1i32, 2, 3, 4, 5].into_array(),
82                buffer![0i32, 3].into_array(),
83                buffer![3i32, 2].into_array(),
84                Validity::NonNullable,
85            ).with_zero_copy_to_list(true)
86        },
87        true,
88    )]
89    #[case::overlapping(
90        ListViewArray::new(
91            buffer![1i32, 2, 3].into_array(),
92            buffer![0i32, 0, 0].into_array(),
93            buffer![3i32, 3, 3].into_array(),
94            Validity::NonNullable,
95        ),
96        false,
97    )]
98    fn listview_compress_roundtrip(
99        #[case] input: ListViewArray,
100        #[case] expect_list: bool,
101    ) -> VortexResult<()> {
102        let array_ref = input.clone().into_array();
103        let result = BtrBlocksCompressor::default().compress(&array_ref)?;
104        if expect_list {
105            assert!(result.as_opt::<List>().is_some());
106        } else {
107            assert!(result.as_opt::<ListView>().is_some());
108        }
109        assert_arrays_eq!(result, input);
110        Ok(())
111    }
112
113    #[test]
114    fn test_constant_all_true() -> VortexResult<()> {
115        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable);
116        let btr = BtrBlocksCompressor::default();
117        let compressed = btr.compress(&array.clone().into_array())?;
118        assert!(compressed.is::<Constant>());
119        assert_arrays_eq!(compressed, array);
120        Ok(())
121    }
122
123    #[test]
124    fn test_constant_all_false() -> VortexResult<()> {
125        let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable);
126        let btr = BtrBlocksCompressor::default();
127        let compressed = btr.compress(&array.clone().into_array())?;
128        assert!(compressed.is::<Constant>());
129        assert_arrays_eq!(compressed, array);
130        Ok(())
131    }
132
133    #[test]
134    fn test_nullable_all_valid_compressed() -> VortexResult<()> {
135        let array = BoolArray::new(
136            BitBuffer::from(vec![true; 100]),
137            Validity::from(BitBuffer::from(vec![true; 100])),
138        );
139        let btr = BtrBlocksCompressor::default();
140        let compressed = btr.compress(&array.clone().into_array())?;
141        assert!(compressed.is::<Constant>());
142        assert_arrays_eq!(compressed, array);
143        Ok(())
144    }
145
146    #[test]
147    fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> {
148        let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0)));
149        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity);
150        let btr = BtrBlocksCompressor::default();
151        let compressed = btr.compress(&array.clone().into_array())?;
152        assert!(!compressed.is::<Constant>());
153        assert_arrays_eq!(compressed, array);
154        Ok(())
155    }
156
157    #[test]
158    fn test_mixed_not_constant() -> VortexResult<()> {
159        let array = BoolArray::new(
160            BitBuffer::from(vec![true, false, true, false, true]),
161            Validity::NonNullable,
162        );
163        let btr = BtrBlocksCompressor::default();
164        let compressed = btr.compress(&array.clone().into_array())?;
165        assert!(!compressed.is::<Constant>());
166        assert_arrays_eq!(compressed, array);
167        Ok(())
168    }
169}
vortex_btrblocks/canonical_compressor.rs

vortex_btrblocks/
canonical_compressor.rs
