gpu compatible write strategy, move compact strategy to use btrblocks with zstd and pco

fuzz/Cargo.toml

@@ -21,7 +21,7 @@ cargo-fuzz = true
 default = ["native"]
 native = ["libfuzzer-sys", "zstd", "vortex-file", "vortex/files"]
 wasmfuzz = []
-zstd = ["vortex/zstd"]
+zstd = ["vortex/zstd", "vortex-btrblocks/zstd", "vortex-btrblocks/pco"]
 cuda = ["vortex-cuda", "vortex/cuda", "tokio"]
 
 [dependencies]

fuzz/fuzz_targets/file_io.rs

@@ -28,7 +28,6 @@ use vortex_fuzz::CompressorStrategy;
 use vortex_fuzz::FuzzFileAction;
 use vortex_fuzz::RUNTIME;
 use vortex_fuzz::SESSION;
-use vortex_layout::layouts::compact::CompactCompressor;
 use vortex_utils::aliases::DefaultHashBuilder;
 use vortex_utils::aliases::hash_set::HashSet;
 
@@ -62,7 +61,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus {
         CompressorStrategy::Default => SESSION.write_options(),
         CompressorStrategy::Compact => {
             let strategy = WriteStrategyBuilder::default()
-                .with_compressor(CompactCompressor::default())
+                .with_compact_encodings()
                 .build();
             SESSION.write_options().with_strategy(strategy)
         }

fuzz/src/array/mod.rs

@@ -517,15 +517,22 @@ fn random_action_from_list(
 /// Compress an array using the given strategy.
 #[cfg(feature = "zstd")]
 pub fn compress_array(array: &dyn Array, strategy: CompressorStrategy) -> ArrayRef {
-    use vortex_layout::layouts::compact::CompactCompressor;
+    use vortex_btrblocks::BtrBlocksCompressorBuilder;
+    use vortex_btrblocks::FloatCode;
+    use vortex_btrblocks::IntCode;
+    use vortex_btrblocks::StringCode;
 
     match strategy {
         CompressorStrategy::Default => BtrBlocksCompressor::default()
             .compress(array)
             .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"),
-        CompressorStrategy::Compact => CompactCompressor::default()
+        CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default()
+            .include_string([StringCode::Zstd])
+            .include_int([IntCode::Pco])
+            .include_float([FloatCode::Pco])
+            .build()
             .compress(array)
-            .vortex_expect("CompactCompressor compress should succeed in fuzz test"),
+            .vortex_expect("Compact compress should succeed in fuzz test"),
     }
 }
 

vortex-bench/src/lib.rs

@@ -30,7 +30,6 @@ use vortex::error::VortexExpect;
 use vortex::error::vortex_err;
 use vortex::file::VortexWriteOptions;
 use vortex::file::WriteStrategyBuilder;
-use vortex::layout::layouts::compact::CompactCompressor;
 use vortex::utils::aliases::hash_map::HashMap;
 
 pub mod benchmark;
@@ -215,7 +214,7 @@ impl CompactionStrategy {
         match self {
             CompactionStrategy::Compact => options.with_strategy(
                 WriteStrategyBuilder::default()
-                    .with_compressor(CompactCompressor::default())
+                    .with_compact_encodings()
                     .build(),
             ),
             CompactionStrategy::Default => options,

vortex-btrblocks/Cargo.toml

@@ -18,6 +18,7 @@ enum-iterator = { workspace = true }
 getrandom_v03 = { workspace = true }
 itertools = { workspace = true }
 num-traits = { workspace = true }
+pco = { workspace = true, optional = true }
 rand = { workspace = true }
 rustc-hash = { workspace = true }
 tracing = { workspace = true }
@@ -31,12 +32,14 @@ vortex-error = { workspace = true }
 vortex-fastlanes = { workspace = true }
 vortex-fsst = { workspace = true }
 vortex-mask = { workspace = true }
+vortex-pco = { workspace = true, optional = true }
 vortex-runend = { workspace = true }
 vortex-scalar = { workspace = true }
 vortex-sequence = { workspace = true }
 vortex-sparse = { workspace = true }
 vortex-utils = { workspace = true }
 vortex-zigzag = { workspace = true }
+vortex-zstd = { workspace = true, optional = true }
 
 [dev-dependencies]
 divan = { workspace = true }
@@ -47,6 +50,8 @@ vortex-array = { workspace = true, features = ["_test-harness"] }
 [features]
 # This feature enabled unstable encodings for which we don't guarantee stability.
 unstable_encodings = []
+pco = ["dep:pco", "dep:vortex-pco"]
+zstd = ["dep:vortex-zstd"]
 
 [lints]
 workspace = true

vortex-btrblocks/src/builder.rs

@@ -51,9 +51,21 @@ pub struct BtrBlocksCompressorBuilder {
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
-            int_schemes: ALL_INT_SCHEMES.iter().copied().collect(),
-            float_schemes: ALL_FLOAT_SCHEMES.iter().copied().collect(),
-            string_schemes: ALL_STRING_SCHEMES.iter().copied().collect(),
+            int_schemes: ALL_INT_SCHEMES
+                .iter()
+                .copied()
+                .filter(|s| s.code() != IntCode::Pco)
+                .collect(),
+            float_schemes: ALL_FLOAT_SCHEMES
+                .iter()
+                .copied()
+                .filter(|s| s.code() != FloatCode::Pco)
+                .collect(),
+            string_schemes: ALL_STRING_SCHEMES
+                .iter()
+                .copied()
+                .filter(|s| s.code() != StringCode::Zstd)
+                .collect(),
         }
     }
 }

vortex-btrblocks/src/compressor/float/mod.rs

@@ -77,6 +77,8 @@ pub const ALL_FLOAT_SCHEMES: &[&dyn FloatScheme] = &[
     &DictScheme,
     &NullDominated,
     &RLE_FLOAT_SCHEME,
+    #[cfg(feature = "pco")]
+    &PcoScheme,
 ];
 
 /// [`Compressor`] for floating-point numbers.
@@ -142,6 +144,8 @@ pub enum FloatCode {
     Rle,
     /// Sparse encoding for null-dominated arrays.
     Sparse,
+    /// Pco (pcodec) compression for floats.
+    Pco,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -162,6 +166,11 @@ struct DictScheme;
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct NullDominated;
 
+/// Pco (pcodec) compression for floats.
+#[cfg(feature = "pco")]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct PcoScheme;
+
 /// Configuration for float RLE compression.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct FloatRLEConfig;
@@ -520,6 +529,31 @@ impl Scheme for NullDominated {
     }
 }
 
+#[cfg(feature = "pco")]
+impl Scheme for PcoScheme {
+    type StatsType = FloatStats;
+    type CodeType = FloatCode;
+
+    fn code(&self) -> FloatCode {
+        FloatCode::Pco
+    }
+
+    fn compress(
+        &self,
+        _compressor: &BtrBlocksCompressor,
+        stats: &Self::StatsType,
+        _ctx: CompressorContext,
+        _excludes: &[FloatCode],
+    ) -> VortexResult<ArrayRef> {
+        Ok(vortex_pco::PcoArray::from_primitive(
+            stats.source(),
+            pco::DEFAULT_COMPRESSION_LEVEL,
+            8192,
+        )?
+        .into_array())
+    }
+}
+
 #[cfg(test)]
 mod tests {
 

vortex-btrblocks/src/compressor/integer/mod.rs

@@ -62,6 +62,8 @@ pub const ALL_INT_SCHEMES: &[&dyn IntegerScheme] = &[
     &RunEndScheme,
     &SequenceScheme,
     &RLE_INTEGER_SCHEME,
+    #[cfg(feature = "pco")]
+    &PcoScheme,
 ];
 
 /// [`Compressor`] for signed and unsigned integers.
@@ -156,6 +158,8 @@ pub enum IntCode {
     Sequence,
     /// RLE encoding - generic run-length encoding.
     Rle,
+    /// Pco (pcodec) compression for integers.
+    Pco,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -188,6 +192,11 @@ pub struct RunEndScheme;
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct SequenceScheme;
 
+/// Pco (pcodec) compression for integers.
+#[cfg(feature = "pco")]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct PcoScheme;
+
 /// Threshold for the average run length in an array before we consider run-end encoding.
 const RUN_END_THRESHOLD: u32 = 4;
 
@@ -818,6 +827,49 @@ impl Scheme for SequenceScheme {
     }
 }
 
+#[cfg(feature = "pco")]
+impl Scheme for PcoScheme {
+    type StatsType = IntegerStats;
+    type CodeType = IntCode;
+
+    fn code(&self) -> IntCode {
+        IntCode::Pco
+    }
+
+    fn expected_compression_ratio(
+        &self,
+        compressor: &BtrBlocksCompressor,
+        stats: &Self::StatsType,
+        ctx: CompressorContext,
+        excludes: &[IntCode],
+    ) -> VortexResult<f64> {
+        // Pco does not support I8 or U8.
+        if matches!(
+            stats.src.ptype(),
+            vortex_dtype::PType::I8 | vortex_dtype::PType::U8
+        ) {
+            return Ok(0.0);
+        }
+
+        self.estimate_compression_ratio_with_sampling(compressor, stats, ctx, excludes)
+    }
+
+    fn compress(
+        &self,
+        _compressor: &BtrBlocksCompressor,
+        stats: &Self::StatsType,
+        _ctx: CompressorContext,
+        _excludes: &[IntCode],
+    ) -> VortexResult<ArrayRef> {
+        Ok(vortex_pco::PcoArray::from_primitive(
+            stats.source(),
+            pco::DEFAULT_COMPRESSION_LEVEL,
+            8192,
+        )?
+        .into_array())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::iter;

vortex-btrblocks/src/compressor/string.rs

@@ -124,6 +124,8 @@ pub const ALL_STRING_SCHEMES: &[&dyn StringScheme] = &[
     &FSSTScheme,
     &ConstantScheme,
     &NullDominated,
+    #[cfg(feature = "zstd")]
+    &ZstdScheme,
 ];
 
 /// [`Compressor`] for strings.
@@ -209,6 +211,11 @@ pub struct ConstantScheme;
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct NullDominated;
 
+/// Zstd compression without dictionaries (nvCOMP compatible).
+#[cfg(feature = "zstd")]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct ZstdScheme;
+
 /// Unique identifier for string compression schemes.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Sequence, Ord, PartialOrd)]
 pub enum StringCode {
@@ -222,6 +229,8 @@ pub enum StringCode {
     Constant,
     /// Sparse encoding for null-dominated arrays.
     Sparse,
+    /// Zstd compression without dictionaries.
+    Zstd,
 }
 
 impl Scheme for UncompressedScheme {
@@ -502,6 +511,30 @@ impl Scheme for NullDominated {
     }
 }
 
+#[cfg(feature = "zstd")]
+impl Scheme for ZstdScheme {
+    type StatsType = StringStats;
+    type CodeType = StringCode;
+
+    fn code(&self) -> StringCode {
+        StringCode::Zstd
+    }
+
+    fn compress(
+        &self,
+        _compressor: &BtrBlocksCompressor,
+        stats: &Self::StatsType,
+        _ctx: CompressorContext,
+        _excludes: &[StringCode],
+    ) -> VortexResult<ArrayRef> {
+        let compacted = stats.source().compact_buffers()?;
+        Ok(
+            vortex_zstd::ZstdArray::from_var_bin_view_without_dict(&compacted, 3, 8192)?
+                .into_array(),
+        )
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::arrays::VarBinViewArray;