clean up pluggable compressing some more

Signed-off-by: Will Manning <will@willmanning.io>

Author: lwwmanning

vortex-btrblocks/public-api.lock

@@ -616,6 +616,12 @@ impl core::clone::Clone for vortex_btrblocks::BtrBlocksCompressorBuilder
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::clone(&self) -> vortex_btrblocks::BtrBlocksCompressorBuilder
 
+impl core::cmp::Eq for vortex_btrblocks::BtrBlocksCompressorBuilder
+
+impl core::cmp::PartialEq for vortex_btrblocks::BtrBlocksCompressorBuilder
+
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::eq(&self, other: &vortex_btrblocks::BtrBlocksCompressorBuilder) -> bool
+
 impl core::default::Default for vortex_btrblocks::BtrBlocksCompressorBuilder
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::default() -> Self
@@ -624,6 +630,8 @@ impl core::fmt::Debug for vortex_btrblocks::BtrBlocksCompressorBuilder
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
+impl core::marker::StructuralPartialEq for vortex_btrblocks::BtrBlocksCompressorBuilder
+
 pub const vortex_btrblocks::ALL_SCHEMES: &[&dyn vortex_compressor::scheme::Scheme]
 
 pub fn vortex_btrblocks::compress_patches(patches: &vortex_array::patches::Patches) -> vortex_error::VortexResult<vortex_array::patches::Patches>

vortex-btrblocks/src/builder.rs

@@ -79,8 +79,6 @@ pub fn default_excluded() -> HashSet<SchemeId> {
     excluded.insert(string::ZstdScheme.id());
     #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
     excluded.insert(string::ZstdBuffersScheme.id());
-    #[cfg(feature = "unstable_encodings")]
-    excluded.insert(turboquant::scheme::TURBOQUANT_SCHEME.id());
     excluded
 }
 
@@ -109,7 +107,7 @@ pub fn default_excluded() -> HashSet<SchemeId> {
 ///     .include([IntDictScheme.id()])
 ///     .build();
 /// ```
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct BtrBlocksCompressorBuilder {
     schemes: HashSet<&'static dyn Scheme>,
 }

vortex-file/src/strategy.rs

@@ -55,6 +55,7 @@ use vortex_pco::Pco;
 use vortex_runend::RunEnd;
 use vortex_sequence::Sequence;
 use vortex_sparse::Sparse;
+#[cfg(feature = "unstable_encodings")]
 use vortex_tensor::encodings::turboquant::TurboQuant;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_zigzag::ZigZag;
@@ -113,6 +114,7 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
     session.register(RunEnd);
     session.register(Sequence);
     session.register(Sparse);
+    #[cfg(feature = "unstable_encodings")]
     session.register(TurboQuant);
     session.register(ZigZag);
 
@@ -135,7 +137,7 @@ pub struct WriteStrategyBuilder {
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
     allow_encodings: Option<ArrayRegistry>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
-    builder: Option<BtrBlocksCompressorBuilder>,
+    builder: BtrBlocksCompressorBuilder,
 }
 
 impl Default for WriteStrategyBuilder {
@@ -148,7 +150,7 @@ impl Default for WriteStrategyBuilder {
             field_writers: HashMap::new(),
             allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
             flat_strategy: None,
-            builder: None,
+            builder: BtrBlocksCompressorBuilder::default(),
         }
     }
 }
@@ -203,8 +205,7 @@ impl WriteStrategyBuilder {
     /// GPU decompression. Without it, strings use interleaved Zstd compression.
     #[cfg(feature = "zstd")]
     pub fn with_cuda_compatible_encodings(mut self) -> Self {
-        let mut builder = self.builder.unwrap_or_default();
-        builder = builder.exclude([
+        self.builder = self.builder.exclude([
             integer::SparseScheme.id(),
             integer::RLE_INTEGER_SCHEME.id(),
             float::RLE_FLOAT_SCHEME.id(),
@@ -215,14 +216,13 @@ impl WriteStrategyBuilder {
 
         #[cfg(feature = "unstable_encodings")]
         {
-            builder = builder.include([string::ZstdBuffersScheme.id()]);
+            self.builder = self.builder.include([string::ZstdBuffersScheme.id()]);
         }
         #[cfg(not(feature = "unstable_encodings"))]
         {
-            builder = builder.include([string::ZstdScheme.id()]);
+            self.builder = self.builder.include([string::ZstdScheme.id()]);
         }
 
-        self.builder = Some(builder);
         self
     }
 
@@ -233,13 +233,11 @@ impl WriteStrategyBuilder {
     /// especially for floating-point heavy datasets.
     #[cfg(feature = "zstd")]
     pub fn with_compact_encodings(mut self) -> Self {
-        let mut builder = self.builder.unwrap_or_default();
-        builder = builder.include([
+        self.builder = self.builder.include([
             string::ZstdScheme.id(),
             integer::PcoScheme.id(),
             float::PcoScheme.id(),
         ]);
-        self.builder = Some(builder);
         self
     }
 
@@ -254,15 +252,17 @@ impl WriteStrategyBuilder {
     /// compressor is used with TurboQuant added.
     #[cfg(feature = "unstable_encodings")]
     pub fn with_vector_quantization(mut self) -> Self {
-        let mut builder = self.builder.unwrap_or_default();
-        builder = builder.include([turboquant::scheme::TURBOQUANT_SCHEME.id()]);
-        self.builder = Some(builder);
+        use vortex_tensor::encodings::turboquant::scheme::TURBOQUANT_SCHEME;
+        self.builder = self.builder.with_scheme(&TURBOQUANT_SCHEME);
         self
     }
 
     /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides
     /// applied.
     pub fn build(self) -> Arc<dyn LayoutStrategy> {
+        use vortex_btrblocks::SchemeExt as _;
+        use vortex_btrblocks::schemes::integer::IntDictScheme;
+
         let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = self.flat_strategy {
             flat
         } else if let Some(allow_encodings) = self.allow_encodings {
@@ -275,19 +275,28 @@ impl WriteStrategyBuilder {
         let chunked = ChunkedLayoutStrategy::new(flat.clone());
         // 6. buffer chunks so they end up with closer segment ids physically
         let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
-        // 5. compress each chunk
-        if self.builder.is_some() && self.compressor.is_some() {
-            vortex_panic!("Cannot configure both a custom compressor and custom builder schemes");
-        }
 
-        let compressing = if let Some(ref compressor) = self.compressor {
-            CompressingStrategy::new_opaque(buffered, compressor.clone())
-        } else if let Some(ref builder) = self.builder {
-            CompressingStrategy::new_opaque(buffered, builder.build())
+        // 5. compress each chunk
+        // Build separate compressors for data (excludes IntDict to avoid recursive dict encoding)
+        // and stats/dict values (includes IntDict).
+        let (data_compressor, stats_compressor): (
+            Arc<dyn CompressorPlugin>,
+            Arc<dyn CompressorPlugin>,
+        ) = if let Some(compressor) = self.compressor {
+            if self.builder != BtrBlocksCompressorBuilder::default() {
+                vortex_panic!(
+                    "Cannot configure both a custom compressor and custom builder schemes"
+                );
+            }
+            (compressor.clone(), compressor)
         } else {
-            CompressingStrategy::new_btrblocks(buffered, true)
+            let stats = Arc::new(self.builder.clone().build());
+            let data = Arc::new(self.builder.exclude([IntDictScheme.id()]).build());
+            (data, stats)
         };
 
+        let compressing = CompressingStrategy::new(buffered, data_compressor);
+
         // 4. prior to compression, coalesce up to a minimum size
         let coalescing = RepartitionStrategy::new(
             compressing,
@@ -306,11 +315,7 @@ impl WriteStrategyBuilder {
         );
 
         // 2.1. | 3.1. compress stats tables and dict values.
-        let compress_then_flat = if let Some(ref compressor) = compressor {
-            CompressingStrategy::new_opaque(flat, compressor.clone())
-        } else {
-            CompressingStrategy::new_btrblocks(flat, false)
-        };
+        let compress_then_flat = CompressingStrategy::new(flat, stats_compressor);
 
         // 3. apply dict encoding or fallback
         let dict = DictStrategy::new(

vortex-file/tests/test_write_table.rs

@@ -20,6 +20,7 @@ use vortex_array::field_path;
 use vortex_array::scalar_fn::session::ScalarFnSession;
 use vortex_array::session::ArraySession;
 use vortex_array::validity::Validity;
+use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_buffer::ByteBuffer;
 use vortex_file::OpenOptionsSessionExt;
 use vortex_file::WriteOptionsSessionExt;
@@ -67,9 +68,9 @@ async fn test_file_roundtrip() {
 
     // Create a writer which by default uses the BtrBlocks compressor for a.compressed, but leaves
     // the b and the a.raw columns uncompressed.
-    let default_strategy = Arc::new(CompressingStrategy::new_btrblocks(
+    let default_strategy = Arc::new(CompressingStrategy::new(
         FlatLayoutStrategy::default(),
-        false,
+        BtrBlocksCompressor::default(),
     ));
 
     let writer = Arc::new(

vortex-layout/public-api.lock

@@ -168,9 +168,7 @@ pub struct vortex_layout::layouts::compressed::CompressingStrategy
 
 impl vortex_layout::layouts::compressed::CompressingStrategy
 
-pub fn vortex_layout::layouts::compressed::CompressingStrategy::new_btrblocks<S: vortex_layout::LayoutStrategy>(child: S, exclude_int_dict_encoding: bool) -> Self
-
-pub fn vortex_layout::layouts::compressed::CompressingStrategy::new_opaque<S: vortex_layout::LayoutStrategy, C: vortex_layout::layouts::compressed::CompressorPlugin>(child: S, compressor: C) -> Self
+pub fn vortex_layout::layouts::compressed::CompressingStrategy::new<S: vortex_layout::LayoutStrategy, C: vortex_layout::layouts::compressed::CompressorPlugin>(child: S, compressor: C) -> Self
 
 pub fn vortex_layout::layouts::compressed::CompressingStrategy::with_concurrency(self, concurrency: usize) -> Self
 

vortex-layout/src/layouts/compressed.rs

@@ -10,9 +10,6 @@ use vortex_array::ArrayRef;
 use vortex_array::DynArray;
 use vortex_array::expr::stats::Stat;
 use vortex_btrblocks::BtrBlocksCompressor;
-use vortex_btrblocks::BtrBlocksCompressorBuilder;
-use vortex_btrblocks::SchemeExt;
-use vortex_btrblocks::schemes::integer::IntDictScheme;
 use vortex_error::VortexResult;
 use vortex_io::runtime::Handle;
 
@@ -61,32 +58,11 @@ pub struct CompressingStrategy {
 }
 
 impl CompressingStrategy {
-    /// 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.
-    ///
-    /// 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<S: LayoutStrategy>(child: S, exclude_int_dict_encoding: bool) -> Self {
-        let compressor = if exclude_int_dict_encoding {
-            BtrBlocksCompressorBuilder::default()
-                .exclude([IntDictScheme.id()])
-                .build()
-        } else {
-            BtrBlocksCompressor::default()
-        };
-        Self::new(child, Arc::new(compressor))
-    }
-
-    /// Create a new compressor from a plugin interface.
-    pub fn new_opaque<S: LayoutStrategy, C: CompressorPlugin>(child: S, compressor: C) -> Self {
-        Self::new(child, Arc::new(compressor))
-    }
-
-    fn new<S: LayoutStrategy>(child: S, compressor: Arc<dyn CompressorPlugin>) -> Self {
+    /// Create a new compressing strategy that wraps a child strategy with a compressor plugin.
+    pub fn new<S: LayoutStrategy, C: CompressorPlugin>(child: S, compressor: C) -> Self {
         Self {
             child: Arc::new(child),
-            compressor,
+            compressor: Arc::new(compressor),
             concurrency: std::thread::available_parallelism()
                 .map(|v| v.get())
                 .unwrap_or(1),

vortex-layout/src/layouts/table.rs

@@ -85,12 +85,13 @@ impl TableStrategy {
     /// ```ignore
     /// # use std::sync::Arc;
     /// # use vortex_array::dtype::{field_path, Field, FieldPath};
+    /// # use vortex_btrblocks::BtrBlocksCompressor;
     /// # use vortex_layout::layouts::compressed::CompressingStrategy;
     /// # use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
     /// # use vortex_layout::layouts::table::TableStrategy;
     ///
     /// // A strategy for compressing data using the balanced BtrBlocks compressor.
-    /// let compress = CompressingStrategy::new_btrblocks(FlatLayoutStrategy::default(), true);
+    /// let compress = CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default());
     ///
     /// // Our combined strategy uses no compression for validity buffers, BtrBlocks compression
     /// // for most columns, and stores a nested binary column uncompressed (flat) because it