feat: add ChunkedCompressor which compresses chunk n+1 like chunk n

vortex-sampling-compressor/Cargo.toml

@@ -16,6 +16,7 @@ readme = { workspace = true }
 [dependencies]
 arbitrary = { workspace = true, optional = true }
 fsst-rs = { workspace = true }
+itertools = { workspace = true }
 lazy_static = { workspace = true }
 log = { workspace = true }
 rand = { workspace = true }

vortex-sampling-compressor/src/compressors/chunked.rs

@@ -0,0 +1,142 @@
+use std::any::Any;
+use std::collections::HashSet;
+use std::sync::Arc;
+
+use log::warn;
+use vortex::array::{Chunked, ChunkedArray};
+use vortex::encoding::EncodingRef;
+use vortex::{Array, ArrayDType, ArrayDef, IntoArray};
+use vortex_error::{vortex_bail, VortexResult};
+
+use super::EncoderMetadata;
+use crate::compressors::{CompressedArray, CompressionTree, EncodingCompressor};
+use crate::SamplingCompressor;
+
+#[derive(Debug)]
+pub struct ChunkedCompressor;
+
+pub struct ChunkedCompressorMetadata(Option<f32>);
+
+impl EncoderMetadata for ChunkedCompressorMetadata {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
+impl EncodingCompressor for ChunkedCompressor {
+    fn id(&self) -> &str {
+        Chunked::ID.as_ref()
+    }
+
+    fn cost(&self) -> u8 {
+        0
+    }
+
+    fn can_compress(&self, array: &Array) -> Option<&dyn EncodingCompressor> {
+        ChunkedArray::try_from(array)
+            .ok()
+            .map(|_| self as &dyn EncodingCompressor)
+    }
+
+    fn compress<'a>(
+        &'a self,
+        array: &Array,
+        like: Option<CompressionTree<'a>>,
+        ctx: SamplingCompressor<'a>,
+    ) -> VortexResult<CompressedArray<'a>> {
+        let chunked_array = ChunkedArray::try_from(array)?;
+        let like_and_ratio = like_into_parts(like)?;
+        self.compress_chunked(&chunked_array, like_and_ratio, ctx)
+    }
+
+    fn used_encodings(&self) -> HashSet<EncodingRef> {
+        HashSet::from([])
+    }
+}
+
+fn like_into_parts(
+    tree: Option<CompressionTree<'_>>,
+) -> VortexResult<Option<(CompressionTree<'_>, f32)>> {
+    let (_, mut children, metadata) = match tree {
+        None => return Ok(None),
+        Some(tree) => tree.into_parts(),
+    };
+
+    let Some(target_ratio) = metadata else {
+        vortex_bail!("chunked array compression tree must have metadata")
+    };
+
+    let Some(ChunkedCompressorMetadata(target_ratio)) =
+        target_ratio.as_ref().as_any().downcast_ref()
+    else {
+        vortex_bail!(
+            "chunked array compression
+        tree must ChunkedCompressorMetadata"
+        )
+    };
+
+    if children.len() != 1 {
+        vortex_bail!("chunked array compression tree must have one child")
+    }
+
+    let child = children.remove(0);
+
+    match (child, target_ratio) {
+        (None, None) => Ok(None),
+        (Some(child), Some(ratio)) => Ok(Some((child, *ratio))),
+        (..) => vortex_bail!("chunked array compression tree must have a child iff it has a ratio"),
+    }
+}
+
+impl ChunkedCompressor {
+    fn compress_chunked<'a>(
+        &'a self,
+        array: &ChunkedArray,
+        mut previous: Option<(CompressionTree<'a>, f32)>,
+        ctx: SamplingCompressor<'a>,
+    ) -> VortexResult<CompressedArray<'a>> {
+        let less_chunked = array.rechunk(
+            ctx.options().target_block_bytesize,
+            ctx.options().target_block_size,
+        )?;
+        let mut compressed_chunks = Vec::with_capacity(less_chunked.nchunks());
+        for (index, chunk) in less_chunked.chunks().enumerate() {
+            let like = previous.as_ref().map(|(like, _)| like);
+            let (compressed_chunk, tree) = ctx
+                .named(&format!("chunk-{}", index))
+                .compress(&chunk, like)?
+                .into_parts();
+
+            let ratio = (compressed_chunk.nbytes() as f32) / (chunk.nbytes() as f32);
+            let exceeded_target_ratio = previous
+                .as_ref()
+                .map(|(_, target_ratio)| ratio > target_ratio * 1.2)
+                .unwrap_or(false);
+
+            if ratio > 1.0 || exceeded_target_ratio {
+                warn!("unsatisfactory ratio {} {:?}", ratio, previous);
+                let (compressed_chunk, tree) = ctx.compress_array(&chunk)?.into_parts();
+                let new_ratio = (compressed_chunk.nbytes() as f32) / (chunk.nbytes() as f32);
+                previous = tree.map(|tree| (tree, new_ratio));
+                compressed_chunks.push(compressed_chunk);
+            } else {
+                previous = previous.or_else(|| tree.map(|tree| (tree, ratio)));
+                compressed_chunks.push(compressed_chunk);
+            }
+        }
+
+        let (child, ratio) = match previous {
+            Some((child, ratio)) => (Some(child), Some(ratio)),
+            None => (None, None),
+        };
+
+        Ok(CompressedArray::new(
+            ChunkedArray::try_new(compressed_chunks, array.dtype().clone())?.into_array(),
+            Some(CompressionTree::new_with_metadata(
+                self,
+                vec![child],
+                Arc::new(ChunkedCompressorMetadata(ratio)),
+            )),
+        ))
+    }
+}

vortex-sampling-compressor/src/compressors/fsst.rs

@@ -63,7 +63,7 @@ impl EncodingCompressor for FSSTCompressor {
         // between 2-3x depending on the text quality.
         //
         // It's not worth running a full compression step unless the array is large enough.
-        if array.nbytes() < 10 * FSST_SYMTAB_MAX_SIZE {
+        if array.nbytes() < 5 * FSST_SYMTAB_MAX_SIZE {
             return Ok(CompressedArray::uncompressed(array.clone()));
         }
 

vortex-sampling-compressor/src/compressors/mod.rs

@@ -13,6 +13,7 @@ use crate::SamplingCompressor;
 pub mod alp;
 pub mod alp_rd;
 pub mod bitpacked;
+pub mod chunked;
 pub mod constant;
 pub mod date_time_parts;
 pub mod delta;
@@ -23,6 +24,7 @@ pub mod roaring_bool;
 pub mod roaring_int;
 pub mod runend;
 pub mod sparse;
+pub mod struct_;
 pub mod zigzag;
 
 pub trait EncodingCompressor: Sync + Send + Debug {
@@ -162,6 +164,17 @@ impl<'a> CompressionTree<'a> {
             .filter_map(|child| child.as_ref().map(|c| c.num_descendants() + 1))
             .sum::<usize>()
     }
+
+    #[allow(clippy::type_complexity)]
+    pub fn into_parts(
+        self,
+    ) -> (
+        &'a dyn EncodingCompressor,
+        Vec<Option<CompressionTree<'a>>>,
+        Option<Arc<dyn EncoderMetadata>>,
+    ) {
+        (self.compressor, self.children, self.metadata)
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -199,6 +212,7 @@ impl<'a> CompressedArray<'a> {
         self.path
     }
 
+    #[inline]
     pub fn into_parts(self) -> (Array, Option<CompressionTree<'a>>) {
         (self.array, self.path)
     }

vortex-sampling-compressor/src/compressors/struct_.rs

@@ -0,0 +1,66 @@
+use std::collections::HashSet;
+
+use itertools::Itertools;
+use vortex::array::{Struct, StructArray};
+use vortex::encoding::EncodingRef;
+use vortex::variants::StructArrayTrait;
+use vortex::{Array, ArrayDef, IntoArray};
+use vortex_error::VortexResult;
+
+use crate::compressors::{CompressedArray, CompressionTree, EncodingCompressor};
+use crate::SamplingCompressor;
+
+#[derive(Debug)]
+pub struct StructCompressor;
+
+impl EncodingCompressor for StructCompressor {
+    fn id(&self) -> &str {
+        Struct::ID.as_ref()
+    }
+
+    fn cost(&self) -> u8 {
+        0
+    }
+
+    fn can_compress(&self, array: &Array) -> Option<&dyn EncodingCompressor> {
+        StructArray::try_from(array)
+            .ok()
+            .map(|_| self as &dyn EncodingCompressor)
+    }
+
+    fn compress<'a>(
+        &'a self,
+        array: &Array,
+        like: Option<CompressionTree<'a>>,
+        ctx: SamplingCompressor<'a>,
+    ) -> VortexResult<CompressedArray<'a>> {
+        let array = StructArray::try_from(array)?;
+        let compressed_validity = ctx.compress_validity(array.validity())?;
+
+        let children_trees = match like {
+            Some(tree) => tree.children,
+            None => vec![None; array.nfields()],
+        };
+
+        let (arrays, trees) = array
+            .children()
+            .zip(children_trees)
+            .map(|(array, like)| ctx.compress(&array, like.as_ref()))
+            .process_results(|iter| iter.map(|x| (x.array, x.path)).unzip())?;
+
+        Ok(CompressedArray::new(
+            StructArray::try_new(
+                array.names().clone(),
+                arrays,
+                array.len(),
+                compressed_validity,
+            )?
+            .into_array(),
+            Some(CompressionTree::new(self, trees)),
+        ))
+    }
+
+    fn used_encodings(&self) -> HashSet<EncodingRef> {
+        HashSet::from([])
+    }
+}