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::List;
71    use vortex_array::arrays::ListView;
72    use vortex_array::arrays::ListViewArray;
73    use vortex_array::assert_arrays_eq;
74    use vortex_array::session::ArraySession;
75    use vortex_array::validity::Validity;
76    use vortex_buffer::BitBuffer;
77    use vortex_buffer::buffer;
78    use vortex_error::VortexResult;
79    use vortex_session::VortexSession;
80
81    use crate::BtrBlocksCompressor;
82
83    static SESSION: LazyLock<VortexSession> =
84        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
85
86    #[rstest]
87    #[case::zctl(
88        unsafe {
89            ListViewArray::new_unchecked(
90                buffer![1i32, 2, 3, 4, 5].into_array(),
91                buffer![0i32, 3].into_array(),
92                buffer![3i32, 2].into_array(),
93                Validity::NonNullable,
94            ).with_zero_copy_to_list(true)
95        },
96        true,
97    )]
98    #[case::overlapping(
99        ListViewArray::new(
100            buffer![1i32, 2, 3].into_array(),
101            buffer![0i32, 0, 0].into_array(),
102            buffer![3i32, 3, 3].into_array(),
103            Validity::NonNullable,
104        ),
105        false,
106    )]
107    fn listview_compress_roundtrip(
108        #[case] input: ListViewArray,
109        #[case] expect_list: bool,
110    ) -> VortexResult<()> {
111        let array_ref = input.clone().into_array();
112        let result = BtrBlocksCompressor::default()
113            .compress(&array_ref, &mut SESSION.create_execution_ctx())?;
114        if expect_list {
115            assert!(result.as_opt::<List>().is_some());
116        } else {
117            assert!(result.as_opt::<ListView>().is_some());
118        }
119        assert_arrays_eq!(result, input);
120        Ok(())
121    }
122
123    #[test]
124    fn test_constant_all_true() -> VortexResult<()> {
125        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable);
126        let btr = BtrBlocksCompressor::default();
127        let compressed = btr.compress(
128            &array.clone().into_array(),
129            &mut SESSION.create_execution_ctx(),
130        )?;
131        assert!(compressed.is::<Constant>());
132        assert_arrays_eq!(compressed, array);
133        Ok(())
134    }
135
136    #[test]
137    fn test_constant_all_false() -> VortexResult<()> {
138        let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable);
139        let btr = BtrBlocksCompressor::default();
140        let compressed = btr.compress(
141            &array.clone().into_array(),
142            &mut SESSION.create_execution_ctx(),
143        )?;
144        assert!(compressed.is::<Constant>());
145        assert_arrays_eq!(compressed, array);
146        Ok(())
147    }
148
149    #[test]
150    fn test_nullable_all_valid_compressed() -> VortexResult<()> {
151        let array = BoolArray::new(
152            BitBuffer::from(vec![true; 100]),
153            Validity::from(BitBuffer::from(vec![true; 100])),
154        );
155        let btr = BtrBlocksCompressor::default();
156        let compressed = btr.compress(
157            &array.clone().into_array(),
158            &mut SESSION.create_execution_ctx(),
159        )?;
160        assert!(compressed.is::<Constant>());
161        assert_arrays_eq!(compressed, array);
162        Ok(())
163    }
164
165    #[test]
166    fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> {
167        let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0)));
168        let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity);
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);
176        Ok(())
177    }
178
179    #[test]
180    fn test_mixed_not_constant() -> VortexResult<()> {
181        let array = BoolArray::new(
182            BitBuffer::from(vec![true, false, true, false, true]),
183            Validity::NonNullable,
184        );
185        let btr = BtrBlocksCompressor::default();
186        let compressed = btr.compress(
187            &array.clone().into_array(),
188            &mut SESSION.create_execution_ctx(),
189        )?;
190        assert!(!compressed.is::<Constant>());
191        assert_arrays_eq!(compressed, array);
192        Ok(())
193    }
194}
vortex_btrblocks/canonical_compressor.rs

vortex_btrblocks/
canonical_compressor.rs
