fix: don't recursively dictionary-encode dictionary codes

Signed-off-by: Alfonso Subiotto Marques <[email protected]>

Author: asubiotto

fuzz/fuzz_targets/array_ops.rs

@@ -27,7 +27,7 @@ fuzz_target!(|fuzz_action: FuzzArrayAction| -> Corpus {
     for (i, (action, expected)) in actions.into_iter().enumerate() {
         match action {
             Action::Compress => {
-                current_array = BtrBlocksCompressor
+                current_array = BtrBlocksCompressor::default()
                     .compress(current_array.to_canonical().vortex_unwrap().as_ref())
                     .vortex_unwrap();
                 assert_array_eq(&expected.array(), &current_array, i).unwrap();
@@ -56,7 +56,9 @@ fuzz_target!(|fuzz_action: FuzzArrayAction| -> Corpus {
                 ])
                 .contains(&current_array.encoding_id())
                 {
-                    sorted = BtrBlocksCompressor.compress(&sorted).vortex_unwrap();
+                    sorted = BtrBlocksCompressor::default()
+                        .compress(&sorted)
+                        .vortex_unwrap();
                 }
                 assert_search_sorted(sorted, s, side, expected.search(), i).unwrap()
             }

fuzz/src/array/mod.rs

@@ -128,7 +128,9 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
                     )
                     .into_array();
 
-                    let compressed = BtrBlocksCompressor.compress(&indices_array).vortex_unwrap();
+                    let compressed = BtrBlocksCompressor::default()
+                        .compress(&indices_array)
+                        .vortex_unwrap();
                     (
                         Action::Take(compressed),
                         ExpectedValue::Array(current_array.to_array()),

vortex-btrblocks/src/lib.rs

@@ -270,8 +270,10 @@ pub trait Compressor {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct BtrBlocksCompressor;
+#[derive(Default, Debug, Clone)]
+pub struct BtrBlocksCompressor {
+    pub exclude_int_dict_encoding: bool,
+}
 
 impl BtrBlocksCompressor {
     pub fn compress(&self, array: &dyn Array) -> VortexResult<ArrayRef> {
@@ -291,7 +293,11 @@ impl BtrBlocksCompressor {
             Canonical::Bool(bool_array) => Ok(bool_array.into_array()),
             Canonical::Primitive(primitive) => {
                 if primitive.ptype().is_int() {
-                    IntCompressor::compress(&primitive, false, MAX_CASCADE, &[])
+                    if self.exclude_int_dict_encoding {
+                        IntCompressor::compress_no_dict(&primitive, false, MAX_CASCADE, &[])
+                    } else {
+                        IntCompressor::compress(&primitive, false, MAX_CASCADE, &[])
+                    }
                 } else {
                     FloatCompressor::compress(&primitive, false, MAX_CASCADE, &[])
                 }

vortex-file/src/strategy.rs

@@ -76,7 +76,7 @@ impl WriteStrategyBuilder {
         let compressing = if let Some(ref compressor) = self.compressor {
             CompressingStrategy::new_opaque(buffered, compressor.clone(), executor.clone(), 16)
         } else {
-            CompressingStrategy::new_btrblocks(buffered, executor.clone(), 16)
+            CompressingStrategy::new_btrblocks(buffered, executor.clone(), 16, true)
         };
 
         // 4. prior to compression, coalesce up to a minimum size
@@ -97,7 +97,12 @@ impl WriteStrategyBuilder {
                 1,
             )
         } else {
-            CompressingStrategy::new_btrblocks(FlatLayoutStrategy::default(), executor.clone(), 1)
+            CompressingStrategy::new_btrblocks(
+                FlatLayoutStrategy::default(),
+                executor.clone(),
+                1,
+                false,
+            )
         };
 
         // 3. apply dict encoding or fallback

vortex-file/src/tests.rs

@@ -17,6 +17,7 @@ use vortex_array::stream::ArrayStreamExt;
 use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_buffer::{Buffer, ByteBufferMut, buffer};
+use vortex_dict::{DictEncoding, DictVTable};
 use vortex_dtype::PType::I32;
 use vortex_dtype::{DType, DecimalDType, Nullability, PType, StructFields};
 use vortex_error::VortexResult;
@@ -1183,3 +1184,26 @@ async fn test_into_tokio_array_stream() -> VortexResult<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_array_stream_no_double_dict_encode() -> VortexResult<()> {
+    let num_vals = 2048;
+    let mut values = Vec::<i64>::with_capacity(num_vals);
+    values.extend(iter::repeat_n(0, num_vals / 2));
+    values.extend(iter::repeat_n(1, num_vals / 2));
+
+    let array = PrimitiveArray::from_iter(values).into_array();
+    let buf = VortexWriteOptions::default().write_blocking(Vec::new(), array.to_array_stream())?;
+    let file = VortexOpenOptions::in_memory().open(buf)?;
+    let read_array = file.scan()?.into_array_iter()?.read_all()?;
+
+    let dict = read_array
+        .as_opt::<DictVTable>()
+        .expect("expected root to be dictionary");
+    assert_ne!(
+        dict.codes().encoding().id(),
+        DictEncoding.id(),
+        "dictionary codes should not be dictionary encoded"
+    );
+    Ok(())
+}

vortex-layout/src/layouts/compressed.rs

@@ -43,7 +43,7 @@ where
 
 impl CompressorPlugin for BtrBlocksCompressor {
     fn compress_chunk(&self, chunk: &dyn Array) -> VortexResult<ArrayRef> {
-        BtrBlocksCompressor::compress(self, chunk)
+        self.compress(chunk)
     }
 }
 
@@ -67,10 +67,20 @@ impl<S: LayoutStrategy> CompressingStrategy<S> {
     /// Create a new writer that uses the BtrBlocks-style cascading compressor to compress chunks.
     ///
     /// This provides a good balance between decoding speed and small file size.
-    pub fn new_btrblocks(child: S, executor: Arc<dyn TaskExecutor>, parallelism: usize) -> Self {
+    ///
+    /// Set `exclude_int_dict_encoding` to true to prevent dictionary encoding of integer arrays,
+    /// which is useful when compressing dictionary codes to avoid recursive dictionary encoding.
+    pub fn new_btrblocks(
+        child: S,
+        executor: Arc<dyn TaskExecutor>,
+        parallelism: usize,
+        exclude_int_dict_encoding: bool,
+    ) -> Self {
         Self {
             child,
-            compressor: Arc::new(BtrBlocksCompressor),
+            compressor: Arc::new(BtrBlocksCompressor {
+                exclude_int_dict_encoding,
+            }),
             executor,
             parallelism,
         }

vortex-layout/src/layouts/dict/writer/mod.rs

@@ -112,7 +112,7 @@ where
         let should_fallback = match first_chunk {
             None => true, // empty stream
             Some(chunk) => {
-                let compressed = BtrBlocksCompressor.compress(&chunk)?;
+                let compressed = BtrBlocksCompressor::default().compress(&chunk)?;
                 !compressed.is_encoding(DictEncoding.id())
             }
         };

vortex-python/src/compress.rs

@@ -50,6 +50,6 @@ pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
 ///    'vortex.alp(f64?, len=1000)'
 #[pyfunction]
 pub fn compress(array: PyArrayRef) -> PyResult<PyArrayRef> {
-    let compressed = BtrBlocksCompressor.compress(array.inner())?;
+    let compressed = BtrBlocksCompressor::default().compress(array.inner())?;
     Ok(PyArrayRef::from(compressed))
 }

vortex/src/lib.rs

@@ -91,7 +91,7 @@ mod test {
         let array = PrimitiveArray::new(buffer![42u64; 100_000], Validity::NonNullable);
 
         // You can compress an array in-memory with the BtrBlocks compressor
-        let compressed = BtrBlocksCompressor.compress(array.as_ref())?;
+        let compressed = BtrBlocksCompressor::default().compress(array.as_ref())?;
         println!(
             "BtrBlocks size: {} / {}",
             compressed.nbytes(),

wasm-test/src/main.rs

@@ -13,7 +13,7 @@ pub fn main() {
     // Extremely simple test of compression/decompression and a few compute functions.
     let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).to_array();
 
-    let compressed = BtrBlocksCompressor.compress(&array).unwrap();
+    let compressed = BtrBlocksCompressor::default().compress(&array).unwrap();
     println!("Compressed size: {}", compressed.len());
     println!("Tree view: {}", compressed.display_tree());
 }