feat(btrblocks): compress binary arrays with VarBin offsets and FSST

vortex-btrblocks/src/builder.rs

@@ -61,6 +61,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
     // Binary schemes.
     ////////////////////////////////////////////////////////////////////////////////////////////////
     &binary::BinaryDictScheme,
+    &binary::VarBinScheme,
     // Decimal schemes.
     &decimal::DecimalScheme,
     // Temporal schemes.

vortex-btrblocks/src/schemes/binary/mod.rs

@@ -3,12 +3,14 @@
 
 //! Binary compression schemes.
 
+mod varbin;
 #[cfg(feature = "zstd")]
 mod zstd;
 #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
 mod zstd_buffers;
 
 // Re-export builtin schemes from vortex-compressor.
+pub use varbin::VarBinScheme;
 pub use vortex_compressor::builtins::BinaryDictScheme;
 #[cfg(feature = "zstd")]
 pub use zstd::ZstdScheme;

vortex-btrblocks/src/schemes/binary/varbin.rs

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Offset-based storage for binary arrays.
+//!
+//! Canonical binary arrays are [`VarBinViewArray`](vortex_array::arrays::VarBinViewArray), which
+//! spends a fixed 16 bytes per element on an opaque views buffer that no scheme can compress.
+//! Re-encoding as [`VarBinArray`] replaces that buffer with an offsets child array, which the
+//! cascading compressor can then compress with the ordinary integer schemes. For fixed-width
+//! values the offsets are a constant-stride sequence and collapse to nothing.
+
+use vortex_array::ArrayId;
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::VTable;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::VarBin;
+use vortex_array::arrays::VarBinArray;
+use vortex_array::arrays::primitive::PrimitiveArrayExt;
+use vortex_array::arrays::varbin::VarBinArraySlotsExt;
+use vortex_array::builders::VarBinBuilder;
+use vortex_compressor::scheme::CompressionEstimate;
+use vortex_compressor::scheme::DeferredEstimate;
+use vortex_compressor::scheme::SchemeExt;
+use vortex_error::VortexResult;
+
+use crate::ArrayAndStats;
+use crate::CascadingCompressor;
+use crate::CompressorContext;
+use crate::Scheme;
+
+/// Offset-based (rather than view-based) storage for binary arrays.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct VarBinScheme;
+
+impl Scheme for VarBinScheme {
+    fn scheme_name(&self) -> &'static str {
+        "vortex.binary.varbin"
+    }
+
+    fn matches(&self, canonical: &Canonical) -> bool {
+        canonical.dtype().is_binary()
+    }
+
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![VarBin.id()]
+    }
+
+    fn num_children(&self) -> usize {
+        1
+    }
+
+    fn expected_compression_ratio(
+        &self,
+        _data: &ArrayAndStats,
+        _compress_ctx: CompressorContext,
+        _exec_ctx: &mut ExecutionCtx,
+    ) -> CompressionEstimate {
+        CompressionEstimate::Deferred(DeferredEstimate::Sample)
+    }
+
+    fn compress(
+        &self,
+        compressor: &CascadingCompressor,
+        data: &ArrayAndStats,
+        compress_ctx: CompressorContext,
+        exec_ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        // `append_to_builder` resolves the views slice and data buffers once and appends
+        // borrowed slices into a single pre-sized allocation. Iterating the array per element
+        // instead would clone a buffer handle and allocate for every value.
+        let array = data.array();
+        let mut builder = VarBinBuilder::<u64>::with_capacity(array.dtype().clone(), array.len());
+        array.append_to_builder(&mut builder, exec_ctx)?;
+        let varbin = builder.finish_into_varbin();
+
+        let offsets = varbin
+            .offsets()
+            .clone()
+            .execute::<PrimitiveArray>(exec_ctx)?
+            .narrow(exec_ctx)?
+            .into_array();
+        let compressed_offsets =
+            compressor.compress_child(&offsets, &compress_ctx, self.id(), 0, exec_ctx)?;
+
+        Ok(VarBinArray::try_new(
+            compressed_offsets,
+            varbin.bytes().clone(),
+            varbin.dtype().clone(),
+            varbin.validity()?,
+        )?
+        .into_array())
+    }
+}

vortex-btrblocks/src/schemes/string/fsst.rs

@@ -47,7 +47,7 @@ impl Scheme for FSSTScheme {
     }
 
     fn matches(&self, canonical: &Canonical) -> bool {
-        canonical.dtype().is_utf8()
+        canonical.dtype().is_utf8() || canonical.dtype().is_binary()
     }
 
     fn produced_encodings(&self) -> Vec<ArrayId> {

vortex-btrblocks/tests/snapshots/golden__compact__binary_low_cardinality.snap

@@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: binary, len=16384, nbytes=315856
-root: vortex.dict(binary, len=16384) nbytes=6199
+root: vortex.dict(binary, len=16384) nbytes=6196
   metadata: all_values_referenced: true
   codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144
     metadata: bit_width: 3, offset: 0
-  values: vortex.zstd(binary, len=5) nbytes=55
-    metadata: nrows: 5, slice: 0..5
+  values: vortex.varbin(binary, len=5) nbytes=52
+    metadata: 
+    offsets: vortex.primitive(u8, len=6) nbytes=6
+      metadata: ptype: u8

vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap

@@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: binary, len=16384, nbytes=315856
-root: vortex.dict(binary, len=16384) nbytes=6240
+root: vortex.dict(binary, len=16384) nbytes=6196
   metadata: all_values_referenced: true
   codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144
     metadata: bit_width: 3, offset: 0
-  values: vortex.varbinview(binary, len=5) nbytes=96
-    metadata:
+  values: vortex.varbin(binary, len=5) nbytes=52
+    metadata: 
+    offsets: vortex.primitive(u8, len=6) nbytes=6
+      metadata: ptype: u8

vortex-btrblocks/tests/snapshots/golden__unstable__binary_low_cardinality.snap

@@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: binary, len=16384, nbytes=315856
-root: vortex.dict(binary, len=16384) nbytes=6240
+root: vortex.dict(binary, len=16384) nbytes=6196
   metadata: all_values_referenced: true
   codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144
     metadata: bit_width: 3, offset: 0
-  values: vortex.varbinview(binary, len=5) nbytes=96
-    metadata:
+  values: vortex.varbin(binary, len=5) nbytes=52
+    metadata: 
+    offsets: vortex.primitive(u8, len=6) nbytes=6
+      metadata: ptype: u8

vortex-btrblocks/tests/varbin_scheme.rs

@@ -0,0 +1,169 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Measures `VarBinScheme` against the same compressor with the scheme excluded.
+
+#![allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)]
+
+use std::sync::LazyLock;
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::VarBinViewArray;
+use vortex_array::assert_arrays_eq;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_btrblocks::BtrBlocksCompressorBuilder;
+use vortex_btrblocks::SchemeExt;
+use vortex_btrblocks::schemes::binary::VarBinScheme;
+#[cfg(feature = "unstable_encodings")]
+use vortex_btrblocks::schemes::string::OnPairScheme;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
+
+const N: usize = 100_000;
+
+fn lcg(state: &mut u64) -> u64 {
+    *state = state
+        .wrapping_mul(6364136223846793005)
+        .wrapping_add(1442695040888963407);
+    *state
+}
+
+fn cases() -> Vec<(&'static str, ArrayRef)> {
+    let mut s = 42u64;
+    let random: Vec<Vec<u8>> = (0..N)
+        .map(|_| (0..16).map(|_| (lcg(&mut s) >> 33) as u8).collect())
+        .collect();
+    let prefixed: Vec<Vec<u8>> = (0..N)
+        .map(|i| format!("PREFIX_{i:09}").into_bytes())
+        .collect();
+    let mut s2 = 7u64;
+    let wide: Vec<Vec<u8>> = (0..N)
+        .map(|_| (0..256).map(|_| (lcg(&mut s2) >> 33) as u8).collect())
+        .collect();
+
+    let nullable = VarBinViewArray::from_iter(
+        (0..N).map(|i| (i % 7 != 0).then(|| prefixed[i].as_slice())),
+        DType::Binary(Nullability::Nullable),
+    )
+    .into_array();
+
+    vec![
+        ("nulls every 7th", nullable),
+        (
+            "random 16B (hash)",
+            VarBinViewArray::from_iter_bin(random.iter().map(|v| v.as_slice())).into_array(),
+        ),
+        (
+            "shared prefix",
+            VarBinViewArray::from_iter_bin(prefixed.iter().map(|v| v.as_slice())).into_array(),
+        ),
+        (
+            "random 256B",
+            VarBinViewArray::from_iter_bin(wide.iter().map(|v| v.as_slice())).into_array(),
+        ),
+    ]
+}
+
+#[test]
+fn varbin_scheme_shrinks_binary() -> VortexResult<()> {
+    let with = BtrBlocksCompressorBuilder::default().build();
+    let without = BtrBlocksCompressorBuilder::default()
+        .exclude_schemes([VarBinScheme.id()])
+        .build();
+
+    println!(
+        "{:<20}{:>12}{:>14}{:>14}{:>9}",
+        "case", "input", "without", "with", "ratio"
+    );
+    for (name, array) in cases() {
+        let a = {
+            let mut ctx = SESSION.create_execution_ctx();
+            without.compress(&array, &mut ctx)?.nbytes()
+        };
+        let b = {
+            let mut ctx = SESSION.create_execution_ctx();
+            with.compress(&array, &mut ctx)?.nbytes()
+        };
+        println!(
+            "{:<20}{:>12}{:>14}{:>14}{:>9.2}",
+            name,
+            array.nbytes(),
+            a,
+            b,
+            b as f64 / a as f64
+        );
+
+        assert!(
+            b <= a,
+            "{name}: enabling VarBinScheme grew the output, {a} -> {b}"
+        );
+
+        let mut ctx = SESSION.create_execution_ctx();
+        let compressed = with.compress(&array, &mut ctx)?;
+        let decoded = compressed
+            .execute::<VarBinViewArray>(&mut ctx)?
+            .into_array();
+        assert_arrays_eq!(&array, &decoded, &mut ctx);
+    }
+    Ok(())
+}
+
+/// Same bytes, two dtypes: FSST compresses bytes rather than codepoints, so the dtype must not
+/// change the result. `OnPairScheme` only matches utf8 and would otherwise win the utf8 column.
+#[test]
+fn fsst_versus_varbin_on_identical_bytes() -> VortexResult<()> {
+    let builder = BtrBlocksCompressorBuilder::default();
+    #[cfg(feature = "unstable_encodings")]
+    let builder = builder.exclude_schemes([OnPairScheme.id()]);
+    let compressor = builder.build();
+    let mut seed = 99u64;
+
+    let shared_prefix: Vec<String> = (0..N).map(|i| format!("PREFIX_{i:09}")).collect();
+    let hex_random: Vec<String> = (0..N)
+        .map(|_| {
+            (0..16)
+                .map(|_| format!("{:02x}", (lcg(&mut seed) >> 33) as u8))
+                .collect()
+        })
+        .collect();
+
+    println!(
+        "{:<20}{:>14}{:>14}{:>9}",
+        "case", "binary", "utf8(fsst)", "fsst/bin"
+    );
+    for (name, vals) in [
+        ("shared prefix", &shared_prefix),
+        ("hex random", &hex_random),
+    ] {
+        let as_bin = VarBinViewArray::from_iter_bin(vals.iter().map(|v| v.as_bytes())).into_array();
+        let as_str = VarBinViewArray::from_iter_str(vals.iter().map(|v| v.as_str())).into_array();
+
+        let bin_bytes = {
+            let mut ctx = SESSION.create_execution_ctx();
+            compressor.compress(&as_bin, &mut ctx)?.nbytes()
+        };
+        let utf8_bytes = {
+            let mut ctx = SESSION.create_execution_ctx();
+            compressor.compress(&as_str, &mut ctx)?.nbytes()
+        };
+        println!(
+            "{:<20}{:>14}{:>14}{:>9.2}",
+            name,
+            bin_bytes,
+            utf8_bytes,
+            utf8_bytes as f64 / bin_bytes as f64
+        );
+
+        // FSST compresses bytes, not codepoints, so the dtype must not change the result.
+        assert_eq!(
+            bin_bytes, utf8_bytes,
+            "{name}: binary and utf8 must compress identically"
+        );
+    }
+    Ok(())
+}