Vortex layout strategy (#2217)

Just a quick hack together of a layout strategy that includes
size+row-based chunking and the sampling compressor.

Author: gatesn

Cargo.lock

@@ -63,7 +63,6 @@ tokio = { workspace = true, features = ["full"] }
 uuid = { workspace = true, features = ["v4"] }
 vortex = { workspace = true, features = ["object_store", "parquet"] }
 vortex-datafusion = { workspace = true }
-vortex-mask = { workspace = true }
 xshell = { workspace = true }
 
 [dev-dependencies]

bench-vortex/Cargo.toml

@@ -7,6 +7,7 @@ use vortex::array::PrimitiveArray;
 use vortex::compute::filter;
 use vortex::dtype::{DType, Nullability, PType};
 use vortex::encodings::alp::{ALPArray, ALPEncoding};
+use vortex::mask::Mask;
 use vortex::sampling_compressor::compressors::alp::ALPCompressor;
 use vortex::sampling_compressor::compressors::bitpacked::{
     BitPackedCompressor, BITPACK_NO_PATCHES, BITPACK_WITH_PATCHES,
@@ -16,7 +17,6 @@ use vortex::sampling_compressor::compressors::EncodingCompressor;
 use vortex::sampling_compressor::SamplingCompressor;
 use vortex::variants::PrimitiveArrayTrait;
 use vortex::{Array, Encoding, IntoArray, IntoCanonical};
-use vortex_mask::Mask;
 
 // criterion benchmark setup:
 fn bench_sel_vec(c: &mut Criterion) {

bench-vortex/benches/sel_vec.rs

@@ -102,16 +102,9 @@ fn main() {
     });
 
     let formats = if args.only_vortex {
-        vec![Format::OnDiskVortex {
-            enable_compression: true,
-        }]
+        vec![Format::OnDiskVortex]
     } else {
-        vec![
-            Format::Parquet,
-            Format::OnDiskVortex {
-                enable_compression: true,
-            },
-        ]
+        vec![Format::Parquet, Format::OnDiskVortex]
     };
 
     let queries = match args.queries.clone() {
@@ -140,9 +133,7 @@ fn main() {
                 .await
                 .unwrap()
             }),
-            Format::OnDiskVortex {
-                enable_compression: true,
-            } => {
+            Format::OnDiskVortex => {
                 runtime.block_on(async {
                     clickbench::register_vortex_files(
                         context.clone(),

bench-vortex/src/bin/clickbench.rs

@@ -87,20 +87,9 @@ async fn bench_main(
 
     // The formats to run against (vs the baseline)
     let formats = if only_vortex {
-        vec![
-            Format::Arrow,
-            Format::OnDiskVortex {
-                enable_compression: true,
-            },
-        ]
+        vec![Format::Arrow, Format::OnDiskVortex]
     } else {
-        vec![
-            Format::Arrow,
-            Format::Parquet,
-            Format::OnDiskVortex {
-                enable_compression: true,
-            },
-        ]
+        vec![Format::Arrow, Format::Parquet, Format::OnDiskVortex]
     };
 
     // Load datasets

bench-vortex/src/bin/tpch_benchmark.rs

@@ -9,14 +9,12 @@ use datafusion::datasource::listing::{
 use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use futures::{stream, StreamExt, TryStreamExt};
 use tokio::fs::{create_dir_all, OpenOptions};
-use vortex::aliases::hash_map::HashMap;
-use vortex::array::{ChunkedArray, StructArray};
+use vortex::arrow::FromArrowType;
 use vortex::dtype::DType;
-use vortex::error::vortex_err;
+use vortex::error::{vortex_err, VortexError};
 use vortex::file::{VortexWriteOptions, VORTEX_FILE_EXTENSION};
-use vortex::sampling_compressor::SamplingCompressor;
-use vortex::variants::StructArrayTrait;
-use vortex::{Array, IntoArray, IntoArrayVariant};
+use vortex::stream::ArrayStreamAdapter;
+use vortex::Array;
 use vortex_datafusion::persistent::VortexFormat;
 
 use crate::{idempotent_async, CTX};
@@ -156,7 +154,6 @@ pub async fn register_vortex_files(
                 .join(format!("hits_{idx}.parquet"));
             let output_path = vortex_dir.join(format!("hits_{idx}.{VORTEX_FILE_EXTENSION}"));
             let session = session.clone();
-            let schema = schema.clone();
 
             tokio::spawn(async move {
                 let output_path = output_path.clone();
@@ -168,49 +165,17 @@ pub async fn register_vortex_files(
                             ParquetReadOptions::default(),
                         )
                         .await?
-                        .collect()
+                        .execute_stream()
                         .await?;
 
-                    // Create a ChunkedArray from the set of chunks.
-                    let sts = record_batches
-                        .into_iter()
-                        .map(Array::try_from)
-                        .map(|a| a.unwrap().into_struct().unwrap())
-                        .collect::<Vec<_>>();
-
-                    let mut arrays_map: HashMap<Arc<str>, Vec<Array>> = HashMap::default();
-                    let mut types_map: HashMap<Arc<str>, DType> = HashMap::default();
-
-                    for st in sts.into_iter() {
-                        let struct_dtype = st.dtype().as_struct().unwrap();
-                        let names = struct_dtype.names().iter();
-                        let types = struct_dtype.fields();
-
-                        for (field_name, field_type) in names.zip(types) {
-                            let val = arrays_map.entry(field_name.clone()).or_default();
-                            val.push(st.maybe_null_field_by_name(field_name.as_ref()).unwrap());
-
-                            types_map.insert(field_name.clone(), field_type.clone());
-                        }
-                    }
-
-                    let fields = schema
-                        .fields()
-                        .iter()
-                        .map(|field| {
-                            let name: Arc<str> = field.name().as_str().into();
-                            let dtype = types_map[&name].clone();
-                            let chunks = arrays_map.remove(&name).unwrap();
-                            let chunked_child = ChunkedArray::try_new(chunks, dtype).unwrap();
-
-                            (name, chunked_child.into_array())
-                        })
-                        .collect::<Vec<_>>();
-
-                    let data = StructArray::from_fields(&fields)?.into_array();
-
-                    let compressor = SamplingCompressor::default();
-                    let data = compressor.compress(&data, None)?.into_array();
+                    // Convert the arrow schema to a Vortex DType
+                    let array_stream = ArrayStreamAdapter::new(
+                        // TODO(ngates): or should we use the provided schema?
+                        DType::from_arrow(record_batches.schema()),
+                        record_batches.map(|batch| {
+                            batch.map_err(VortexError::from).and_then(Array::try_from)
+                        }),
+                    );
 
                     let f = OpenOptions::new()
                         .write(true)
@@ -219,9 +184,7 @@ pub async fn register_vortex_files(
                         .open(&vtx_file)
                         .await?;
 
-                    VortexWriteOptions::default()
-                        .write(f, data.into_array_stream())
-                        .await?;
+                    VortexWriteOptions::default().write(f, array_stream).await?;
 
                     anyhow::Ok(())
                 })

bench-vortex/src/clickbench.rs

@@ -94,7 +94,6 @@ pub fn decompress_bz2(input_path: PathBuf, output_path: PathBuf) -> PathBuf {
 pub trait BenchmarkDataset {
     fn as_uncompressed(&self);
     fn to_vortex_array(&self) -> VortexResult<Array>;
-    fn compress_to_vortex(&self) -> VortexResult<()>;
     fn write_as_parquet(&self);
     fn write_as_vortex(&self) -> impl Future<Output = ()>;
     fn list_files(&self, file_type: FileType) -> Vec<PathBuf>;

bench-vortex/src/data_downloads.rs

@@ -48,10 +48,6 @@ pub mod taxi_data;
 pub mod tpch;
 pub mod vortex_utils;
 
-// Sizes match default compressor configuration
-const TARGET_BLOCK_BYTESIZE: usize = 16 * (1 << 20);
-const TARGET_BLOCK_SIZE: usize = 64 * (1 << 10);
-
 #[macro_export]
 macro_rules! feature_flagged_allocator {
     () => {
@@ -81,7 +77,7 @@ pub enum Format {
     Arrow,
     Parquet,
     InMemoryVortex,
-    OnDiskVortex { enable_compression: bool },
+    OnDiskVortex,
 }
 
 impl std::fmt::Display for Format {
@@ -93,8 +89,8 @@ impl std::fmt::Display for Format {
             Format::InMemoryVortex => {
                 write!(f, "in_memory_vortex")
             }
-            Format::OnDiskVortex { enable_compression } => {
-                write!(f, "on_disk_vortex(compressed={enable_compression})")
+            Format::OnDiskVortex => {
+                write!(f, "on_disk_vortex(compressed=true)")
             }
         }
     }
@@ -107,12 +103,7 @@ impl Format {
             Format::Arrow => "arrow".to_string(),
             Format::Parquet => "parquet".to_string(),
             Format::InMemoryVortex => "vortex-in-memory".to_string(),
-            Format::OnDiskVortex { enable_compression } => if *enable_compression {
-                "vortex-file-compressed"
-            } else {
-                "vortex-file-uncompressed"
-            }
-            .to_string(),
+            Format::OnDiskVortex => "vortex-file-compressed".to_string(),
         }
     }
 }
@@ -404,65 +395,3 @@ pub fn generate_struct_of_list_of_ints_array(
         DType::Struct(struct_dtype.clone(), Nullability::NonNullable),
     )
 }
-
-#[cfg(test)]
-mod test {
-    use std::fs::File;
-    use std::ops::Deref;
-    use std::sync::Arc;
-
-    use arrow_array::{ArrayRef as ArrowArrayRef, StructArray as ArrowStructArray};
-    use log::LevelFilter;
-    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
-    use vortex::arrow::{FromArrowArray, IntoArrowArray};
-    use vortex::compress::CompressionStrategy;
-    use vortex::sampling_compressor::SamplingCompressor;
-    use vortex::Array;
-
-    use crate::taxi_data::taxi_data_parquet;
-    use crate::{compress_taxi_data, setup_logger};
-
-    #[ignore]
-    #[test]
-    fn compression_ratio() {
-        setup_logger(LevelFilter::Debug);
-        _ = compress_taxi_data();
-    }
-
-    #[ignore]
-    #[test]
-    fn round_trip_arrow() {
-        let file = File::open(taxi_data_parquet()).unwrap();
-        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
-        let reader = builder.with_limit(1).build().unwrap();
-
-        for record_batch in reader.map(|batch_result| batch_result.unwrap()) {
-            let struct_arrow: ArrowStructArray = record_batch.into();
-            let arrow_array: ArrowArrayRef = Arc::new(struct_arrow);
-            let vortex_array = Array::from_arrow(arrow_array.clone(), false);
-            let vortex_as_arrow = vortex_array.into_arrow_preferred().unwrap();
-            assert_eq!(vortex_as_arrow.deref(), arrow_array.deref());
-        }
-    }
-
-    // Ignoring since Struct arrays don't currently support equality.
-    // https://github.com/apache/arrow-rs/issues/5199
-    #[ignore]
-    #[test]
-    fn round_trip_arrow_compressed() {
-        let file = File::open(taxi_data_parquet()).unwrap();
-        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
-        let reader = builder.with_limit(1).build().unwrap();
-        let compressor: &dyn CompressionStrategy = &SamplingCompressor::default();
-
-        for record_batch in reader.map(|batch_result| batch_result.unwrap()) {
-            let struct_arrow: ArrowStructArray = record_batch.into();
-            let arrow_array: ArrowArrayRef = Arc::new(struct_arrow);
-            let vortex_array = Array::from_arrow(arrow_array.clone(), false);
-
-            let compressed = compressor.compress(&vortex_array).unwrap();
-            let compressed_as_arrow = compressed.into_arrow_preferred().unwrap();
-            assert_eq!(compressed_as_arrow.deref(), arrow_array.deref());
-        }
-    }
-}

bench-vortex/src/lib.rs

@@ -12,14 +12,13 @@ use reqwest::Url;
 use tokio::fs::File;
 use vortex::aliases::hash_map::HashMap;
 use vortex::array::ChunkedArray;
-use vortex::error::VortexResult;
+use vortex::error::{VortexExpect, VortexResult};
 use vortex::{Array, IntoArray};
 
 use crate::data_downloads::{decompress_bz2, download_data, BenchmarkDataset, FileType};
 use crate::public_bi_data::PBIDataset::*;
 use crate::reader::{
-    compress_parquet_to_vortex, open_vortex, read_parquet_to_vortex, rewrite_parquet_as_vortex,
-    write_csv_as_parquet,
+    open_vortex, read_parquet_to_vortex, rewrite_parquet_as_vortex, write_csv_as_parquet,
 };
 use crate::{idempotent, IdempotentPath};
 
@@ -561,34 +560,13 @@ impl BenchmarkDataset for BenchmarkDatasets {
         let arrays = self
             .list_files(FileType::Parquet)
             .iter()
-            .map(|f| read_parquet_to_vortex(f.as_path()))
+            .flat_map(|f| {
+                read_parquet_to_vortex(f.as_path()).vortex_expect("Failed to read parquet")
+            })
             .collect::<VortexResult<Vec<_>>>()?;
-        assert!(!arrays.is_empty());
-        let dtype = arrays[0].dtype().clone();
-        ChunkedArray::try_new(
-            arrays.iter().flat_map(|x| x.chunks()).collect::<Vec<_>>(),
-            dtype,
-        )
-        .map(|x| x.into_array())
-    }
-
-    fn compress_to_vortex(&self) -> VortexResult<()> {
-        self.write_as_parquet();
-        for f in self.list_files(FileType::Parquet) {
-            info!(
-                "Compressing and writing {} to vortex",
-                f.to_str().unwrap_or("None")
-            );
-            let from_vortex = compress_parquet_to_vortex(f.as_path())?;
-            let vx_size = from_vortex.nbytes();
 
-            info!(
-                "Vortex size: {}, {}B",
-                format_size(vx_size as u64, DECIMAL),
-                vx_size
-            );
-        }
-        Ok(())
+        assert!(!arrays.is_empty());
+        Ok(ChunkedArray::from_iter(arrays).into_array())
     }
 
     fn write_as_parquet(&self) {