feat: Use previous chunks to inform compressions decisions (#2724)

Followup of #2723, re-encoding array children according to a known
well-compressed chunk from the `BtrBlockCompressor`.

Current
[benchmarks](#2724 (comment))
seem to show that compress throughput is up, with decompression slightly
slower and query times are up by up to 10%.
Two interesting points are that wide columned arrays are basically
always slower to compress but decompression held-up better. Arade is
also the one dataset where compressed files shrunk significantly,
regardless of throughput.
I think these tradeoffs should be configurable, and I assume that
increasing the tolerance (currently hard coded at 20%) for compression
ratio drift will allow users to control this knob in a way that better
fits their usecases.

Other benchmarks (because there are so many commits/comments on this
PR):
- [Clickbench on
NVME](#2724 (comment))
- [TPCH on
NVME](#2724 (comment))
- [Random
Access](#2724 (comment))
- [TPC-H on
S3](#2724 (comment))
- Note that it doesn't actually generate new files here, so the results
aren't useful.

---------

Co-authored-by: Nicholas Gates <[email protected]>
Co-authored-by: Robert Kruszewski <[email protected]>

Author: 3 people

bench-vortex/src/bin/compress.rs

@@ -118,9 +118,9 @@ fn compress(
                     let vx_array =
                         runtime.block_on(async { dataset_handle.to_vortex_array().await });
                     ChunkedArray::from_iter(vx_array.as_::<ChunkedArray>().chunks().iter().map(
-                        |c| {
-                            let mut builder = builder_with_capacity(c.dtype(), c.len());
-                            c.append_to_builder(builder.as_mut()).unwrap();
+                        |chunk| {
+                            let mut builder = builder_with_capacity(chunk.dtype(), chunk.len());
+                            chunk.append_to_builder(builder.as_mut()).unwrap();
                             builder.finish()
                         },
                     ))

bench-vortex/src/bin/notimplemented.rs

@@ -21,9 +21,10 @@ use vortex::encodings::fastlanes::{BitPackedArray, DeltaArray, FoRArray};
 use vortex::encodings::fsst::{fsst_compress, fsst_train_compressor};
 use vortex::encodings::runend::RunEndArray;
 use vortex::encodings::sparse::SparseArray;
-use vortex::encodings::zigzag::ZigZagArray;
+use vortex::encodings::zigzag::ZigZagEncoding;
 use vortex::scalar::Scalar;
 use vortex::validity::Validity;
+use vortex::vtable::EncodingVTable;
 use vortex::{Array, ArrayRef, IntoArray};
 
 fn fsst_array() -> ArrayRef {
@@ -140,7 +141,12 @@ fn enc_impls() -> Vec<ArrayRef> {
         .into_array(),
         varbin_array(),
         varbinview_array(),
-        ZigZagArray::encode(&buffer![-1, 1, -9, 9].into_array())
+        ZigZagEncoding
+            .encode(
+                &buffer![-1, 1, -9, 9].into_array().to_canonical().unwrap(),
+                None,
+            )
+            .unwrap()
             .unwrap()
             .into_array(),
     ]

bench-vortex/src/clickbench.rs

@@ -15,6 +15,7 @@ use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use reqwest::IntoUrl;
 use reqwest::blocking::Response;
 use tokio::fs::{OpenOptions, create_dir_all};
+use tokio::io::AsyncWriteExt;
 use tracing::{info, warn};
 use url::Url;
 use vortex::TryIntoArray;
@@ -208,7 +209,11 @@ pub async fn convert_parquet_to_vortex(
                         .open(&vtx_file)
                         .await?;
 
-                    VortexWriteOptions::default().write(f, array_stream).await?;
+                    VortexWriteOptions::default()
+                        .write(f, array_stream)
+                        .await?
+                        .flush()
+                        .await?;
 
                     anyhow::Ok(())
                 })

bench-vortex/src/compress/bench.rs

@@ -56,6 +56,7 @@ pub fn benchmark_compress<F>(
 where
     F: Fn() -> ArrayRef,
 {
+    tracing::info!("Running {bench_name} benchmark");
     let uncompressed = make_uncompressed();
     let uncompressed_size = uncompressed.nbytes();
     let compressed_size = AtomicU64::default();

bench-vortex/src/datasets/public_bi_data.rs

@@ -13,6 +13,7 @@ use humansize::{DECIMAL, format_size};
 use log::{debug, info};
 use reqwest::Url;
 use tokio::fs::File;
+use tokio::io::AsyncWriteExt;
 use vortex::aliases::hash_map::HashMap;
 use vortex::arrays::ChunkedArray;
 use vortex::error::{VortexExpect, VortexResult, vortex_err};
@@ -521,12 +522,16 @@ impl PBIDataset {
             let compressed = idempotent_async(
                 &self.path_for_file_type(output_fname, FileType::Vortex),
                 |output_path| async {
-                    VortexWriteOptions::default()
+                    let mut f = VortexWriteOptions::default()
                         .write(
                             File::create(output_path).await.unwrap(),
                             parquet_to_vortex(f).await.unwrap(),
                         )
-                        .await
+                        .await?;
+
+                    f.flush().await?;
+
+                    VortexResult::Ok(())
                 },
             )
             .await

bench-vortex/src/datasets/taxi_data.rs

@@ -2,6 +2,7 @@ use std::path::PathBuf;
 
 use async_trait::async_trait;
 use tokio::fs::File;
+use tokio::io::AsyncWriteExt;
 use vortex::ArrayRef;
 use vortex::error::VortexError;
 use vortex::file::{VortexOpenOptions, VortexWriteOptions};
@@ -53,6 +54,8 @@ pub async fn taxi_data_vortex() -> PathBuf {
                 output_file,
                 parquet_to_vortex(taxi_data_parquet()).await.unwrap(),
             )
+            .await?
+            .flush()
             .await?;
         Ok::<PathBuf, VortexError>(buf)
     })

encodings/alp/benches/alp_compress.rs

@@ -58,7 +58,7 @@ fn compress_alp<T: ALPFloat + NativePType>(bencher: Bencher, args: (usize, f64,
     bencher
         .with_inputs(|| (values.clone(), validity.clone()))
         .bench_values(|(values, validity)| {
-            alp_encode(&PrimitiveArray::new(values, validity)).unwrap()
+            alp_encode(&PrimitiveArray::new(values, validity), None).unwrap()
         })
 }
 
@@ -80,7 +80,7 @@ fn decompress_alp<T: ALPFloat + NativePType>(bencher: Bencher, args: (usize, f64
         Validity::NonNullable
     };
     let values = values.freeze();
-    let array = alp_encode(&PrimitiveArray::new(values, validity)).unwrap();
+    let array = alp_encode(&PrimitiveArray::new(values, validity), None).unwrap();
     bencher
         .with_inputs(|| array.clone())
         .bench_values(|array| array.to_canonical().unwrap());

encodings/alp/src/alp/array.rs

@@ -1,27 +1,26 @@
 use std::fmt::Debug;
 
-use vortex_array::arrays::PrimitiveArray;
 use vortex_array::patches::Patches;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::VTableRef;
 use vortex_array::{
-    Array, ArrayCanonicalImpl, ArrayExt, ArrayImpl, ArrayRef, ArrayStatisticsImpl,
-    ArrayValidityImpl, ArrayVariantsImpl, Canonical, Encoding, SerdeMetadata,
+    Array, ArrayCanonicalImpl, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl,
+    ArrayVariantsImpl, Canonical, Encoding, SerdeMetadata,
 };
 use vortex_dtype::{DType, PType};
 use vortex_error::{VortexResult, vortex_bail};
 use vortex_mask::Mask;
 
 use crate::alp::serde::ALPMetadata;
-use crate::alp::{Exponents, alp_encode, decompress};
+use crate::alp::{Exponents, decompress};
 
 #[derive(Clone, Debug)]
 pub struct ALPArray {
-    dtype: DType,
     encoded: ArrayRef,
-    exponents: Exponents,
     patches: Option<Patches>,
+    dtype: DType,
+    exponents: Exponents,
     stats_set: ArrayStats,
 }
 
@@ -52,14 +51,6 @@ impl ALPArray {
         })
     }
 
-    pub fn encode(array: ArrayRef) -> VortexResult<ArrayRef> {
-        if let Some(parray) = array.as_opt::<PrimitiveArray>() {
-            Ok(alp_encode(parray)?.into_array())
-        } else {
-            vortex_bail!("ALP can only encode primitive arrays");
-        }
-    }
-
     pub fn encoded(&self) -> &ArrayRef {
         &self.encoded
     }
@@ -88,6 +79,18 @@ impl ArrayImpl for ALPArray {
     fn _vtable(&self) -> VTableRef {
         VTableRef::new_ref(&ALPEncoding)
     }
+
+    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
+        let encoded = children[0].clone();
+
+        let patches = self.patches().map(|existing| {
+            let indices = children[1].clone();
+            let values = children[2].clone();
+            Patches::new(existing.array_len(), existing.offset(), indices, values)
+        });
+
+        ALPArray::try_new(encoded, self.exponents(), patches)
+    }
 }
 
 impl ArrayCanonicalImpl for ALPArray {
@@ -135,29 +138,3 @@ impl ArrayVariantsImpl for ALPArray {
 }
 
 impl PrimitiveArrayTrait for ALPArray {}
-
-#[cfg(test)]
-mod tests {
-    use vortex_array::SerdeMetadata;
-    use vortex_array::patches::PatchesMetadata;
-    use vortex_array::test_harness::check_metadata;
-    use vortex_dtype::PType;
-
-    use crate::Exponents;
-    use crate::alp::serde::ALPMetadata;
-
-    #[cfg_attr(miri, ignore)]
-    #[test]
-    fn test_alp_metadata() {
-        check_metadata(
-            "alp.metadata",
-            SerdeMetadata(ALPMetadata {
-                patches: Some(PatchesMetadata::new(usize::MAX, usize::MAX, PType::U64)),
-                exponents: Exponents {
-                    e: u8::MAX,
-                    f: u8::MAX,
-                },
-            }),
-        );
-    }
-}

encodings/alp/src/alp/compress.rs

@@ -28,24 +28,20 @@ macro_rules! match_each_alp_float_ptype {
     })
 }
 
-pub fn alp_encode(parray: &PrimitiveArray) -> VortexResult<ALPArray> {
-    let (exponents, encoded, patches) = alp_encode_components(parray)?;
-    ALPArray::try_new(encoded, exponents, patches)
-}
-
-pub fn alp_encode_components(
-    parray: &PrimitiveArray,
-) -> VortexResult<(Exponents, ArrayRef, Option<Patches>)> {
-    match parray.ptype() {
-        PType::F32 => alp_encode_components_typed::<f32>(parray),
-        PType::F64 => alp_encode_components_typed::<f64>(parray),
+pub fn alp_encode(parray: &PrimitiveArray, exponents: Option<Exponents>) -> VortexResult<ALPArray> {
+    let (exponents, encoded, patches) = match parray.ptype() {
+        PType::F32 => alp_encode_components_typed::<f32>(parray, exponents)?,
+        PType::F64 => alp_encode_components_typed::<f64>(parray, exponents)?,
         _ => vortex_bail!("ALP can only encode f32 and f64"),
-    }
+    };
+
+    ALPArray::try_new(encoded, exponents, patches)
 }
 
 #[allow(clippy::cast_possible_truncation)]
 fn alp_encode_components_typed<T>(
     values: &PrimitiveArray,
+    exponents: Option<Exponents>,
 ) -> VortexResult<(Exponents, ArrayRef, Option<Patches>)>
 where
     T: ALPFloat + NativePType,
@@ -55,7 +51,7 @@ where
     let values_slice = values.as_slice::<T>();
 
     let (exponents, encoded, exceptional_positions, exceptional_values) =
-        T::encode(values_slice, None);
+        T::encode(values_slice, exponents);
 
     let encoded_array = PrimitiveArray::new(encoded, values.validity().clone()).into_array();
 
@@ -132,7 +128,7 @@ mod tests {
     #[test]
     fn test_compress() {
         let array = PrimitiveArray::new(buffer![1.234f32; 1025], Validity::NonNullable);
-        let encoded = alp_encode(&array).unwrap();
+        let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
         assert_eq!(
             encoded.encoded().to_primitive().unwrap().as_slice::<i32>(),
@@ -147,7 +143,7 @@ mod tests {
     #[test]
     fn test_nullable_compress() {
         let array = PrimitiveArray::from_option_iter([None, Some(1.234f32), None]);
-        let encoded = alp_encode(&array).unwrap();
+        let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
         assert_eq!(
             encoded.encoded().to_primitive().unwrap().as_slice::<i32>(),
@@ -165,7 +161,7 @@ mod tests {
     fn test_patched_compress() {
         let values = buffer![1.234f64, 2.718, f64::consts::PI, 4.0];
         let array = PrimitiveArray::new(values.clone(), Validity::NonNullable);
-        let encoded = alp_encode(&array).unwrap();
+        let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_some());
         assert_eq!(
             encoded.encoded().to_primitive().unwrap().as_slice::<i64>(),
@@ -182,7 +178,7 @@ mod tests {
     fn test_compress_ignores_invalid_exceptional_values() {
         let values = buffer![1.234f64, 2.718, f64::consts::PI, 4.0];
         let array = PrimitiveArray::new(values, Validity::from_iter([true, true, false, true]));
-        let encoded = alp_encode(&array).unwrap();
+        let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_none());
         assert_eq!(
             encoded.encoded().to_primitive().unwrap().as_slice::<i64>(),
@@ -216,7 +212,7 @@ mod tests {
             Some(4.0),
             None,
         ]);
-        let encoded = alp_encode(&array).unwrap();
+        let encoded = alp_encode(&array, None).unwrap();
         assert!(encoded.patches().is_some());
 
         assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 });
@@ -236,7 +232,7 @@ mod tests {
     #[test]
     fn roundtrips_close_fractional() {
         let original = PrimitiveArray::from_iter([195.26274f32, 195.27837, -48.815685]);
-        let alp_arr = alp_encode(&original).unwrap();
+        let alp_arr = alp_encode(&original, None).unwrap();
         let decompressed = alp_arr.to_primitive().unwrap();
         assert_eq!(original.as_slice::<f32>(), decompressed.as_slice::<f32>());
     }
@@ -247,7 +243,7 @@ mod tests {
             Buffer::from_iter([195.26274f64, f64::consts::PI, -48.815685]),
             Validity::AllInvalid,
         );
-        let alp_arr = alp_encode(&original).unwrap();
+        let alp_arr = alp_encode(&original, None).unwrap();
         let decompressed = alp_arr.to_primitive().unwrap();
         assert_eq!(
             // The second and third values become exceptions and are replaced
@@ -275,7 +271,7 @@ mod tests {
             buffer![0.0f32, -0.0, f32::NAN, f32::NEG_INFINITY, f32::INFINITY],
             Validity::NonNullable,
         );
-        let encoded = alp_encode(&original).unwrap();
+        let encoded = alp_encode(&original, None).unwrap();
         let decoded = encoded.to_primitive().unwrap();
         for idx in 0..original.len() {
             let decoded_val = decoded.as_slice::<f32>()[idx];