Fix encoding discovery

bench-vortex/Cargo.toml

@@ -29,11 +29,15 @@ itertools = { workspace = true }
 reqwest = { workspace = true }
 parquet = { workspace = true }
 log = { workspace = true }
+simplelog = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
-simplelog = { workspace = true }
 
 [[bench]]
 name = "compress_benchmark"
+harness = false
+
+[[bench]]
+name = "random_access"
 harness = false

bench-vortex/benches/compress_benchmark.rs

@@ -1,6 +1,7 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 
-use bench_vortex::{compress_taxi_data, download_taxi_data};
+use bench_vortex::compress_taxi_data;
+use bench_vortex::taxi_data::download_taxi_data;
 
 fn enc_compress(c: &mut Criterion) {
     download_taxi_data();

bench-vortex/benches/random_access.rs

@@ -0,0 +1,20 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use itertools::Itertools;
+
+use bench_vortex::serde::{take_taxi_data, write_taxi_data};
+use vortex::array::ENCODINGS;
+
+fn random_access(c: &mut Criterion) {
+    let taxi_spiral = write_taxi_data();
+    let indices = [10, 11, 12, 13, 100_000, 3_000_000];
+    println!(
+        "ENCODINGS {:?}",
+        ENCODINGS.iter().map(|e| e.id()).collect_vec()
+    );
+    c.bench_function("random access", |b| {
+        b.iter(|| black_box(take_taxi_data(&taxi_spiral, &indices)))
+    });
+}
+
+criterion_group!(benches, random_access);
+criterion_main!(benches);

bench-vortex/src/lib.rs

@@ -1,18 +1,18 @@
 use arrow_array::RecordBatchReader;
-use std::collections::HashSet;
 use std::fs::{create_dir_all, File};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use itertools::Itertools;
-use log::{info, warn};
+use log::{info, warn, LevelFilter};
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use parquet::arrow::ProjectionMask;
+use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
 
 use vortex::array::chunked::ChunkedArray;
 use vortex::array::downcast::DowncastArrayBuiltin;
-use vortex::array::{Array, ArrayRef};
-use vortex::array::{EncodingId, IntoArray};
+use vortex::array::IntoArray;
+use vortex::array::{Array, ArrayRef, EncodingRef, ENCODINGS};
 use vortex::arrow::FromArrowType;
 use vortex::compress::{CompressConfig, CompressCtx};
 use vortex::formatter::display_tree;
@@ -24,50 +24,58 @@ use vortex_ree::REEEncoding;
 use vortex_roaring::RoaringBoolEncoding;
 use vortex_schema::DType;
 
-pub fn enumerate_arrays() -> Vec<EncodingId> {
+pub mod serde;
+pub mod taxi_data;
+
+pub fn idempotent(name: &str, f: impl FnOnce(&mut File)) -> PathBuf {
+    let path = Path::new(env!("CARGO_MANIFEST_DIR"))
+        .join("data")
+        .join(name);
+    if !path.exists() {
+        create_dir_all(path.parent().unwrap()).unwrap();
+        let mut file = File::create(&path).unwrap();
+        f(&mut file);
+    }
+    path.to_path_buf()
+}
+
+#[allow(dead_code)]
+fn setup_logger(level: LevelFilter) {
+    TermLogger::init(
+        level,
+        Config::default(),
+        TerminalMode::Mixed,
+        ColorChoice::Auto,
+    )
+    .unwrap();
+}
+
+pub fn enumerate_arrays() -> Vec<EncodingRef> {
+    println!("FOUND {:?}", ENCODINGS.iter().map(|e| e.id()).collect_vec());
     vec![
-        ALPEncoding::ID,
-        DictEncoding::ID,
-        BitPackedEncoding::ID,
-        FoREncoding::ID,
-        DateTimeEncoding::ID,
-        // DeltaEncoding::ID,
-        // FFoREncoding::ID,
-        REEEncoding::ID,
-        RoaringBoolEncoding::ID,
-        // RoaringIntEncoding::ID,
+        &ALPEncoding,
+        &DictEncoding,
+        &BitPackedEncoding,
+        &FoREncoding,
+        &DateTimeEncoding,
+        // DeltaEncoding,
+        // FFoREncoding,
+        &REEEncoding,
+        &RoaringBoolEncoding,
+        // RoaringIntEncoding,
         // Doesn't offer anything more than FoR really
-        // ZigZagEncoding::ID,
+        // ZigZagEncoding,
     ]
 }
 
 pub fn compress_ctx() -> CompressCtx {
-    let cfg = CompressConfig::new(HashSet::from_iter(enumerate_arrays()), HashSet::default());
+    let cfg = CompressConfig::new().with_enabled(enumerate_arrays());
     info!("Compression config {cfg:?}");
     CompressCtx::new(Arc::new(cfg))
 }
 
-pub fn download_taxi_data() -> PathBuf {
-    let download_path =
-        Path::new(env!("CARGO_MANIFEST_DIR")).join("data/yellow-tripdata-2023-11.parquet");
-    if download_path.exists() {
-        return download_path;
-    }
-
-    create_dir_all(download_path.parent().unwrap()).unwrap();
-    let mut download_file = File::create(&download_path).unwrap();
-    reqwest::blocking::get(
-        "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2023-11.parquet",
-    )
-    .unwrap()
-    .copy_to(&mut download_file)
-    .unwrap();
-
-    download_path
-}
-
 pub fn compress_taxi_data() -> ArrayRef {
-    let file = File::open(download_taxi_data()).unwrap();
+    let file = File::open(taxi_data::download_taxi_data()).unwrap();
     let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
     let _mask = ProjectionMask::roots(builder.parquet_schema(), [1]);
     let _no_datetime_mask = ProjectionMask::roots(
@@ -131,30 +139,19 @@ mod test {
     use arrow_array::{ArrayRef as ArrowArrayRef, StructArray as ArrowStructArray};
     use log::LevelFilter;
     use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
-    use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
 
     use vortex::array::ArrayRef;
     use vortex::compute::as_arrow::as_arrow;
     use vortex::encode::FromArrowArray;
     use vortex::serde::{ReadCtx, WriteCtx};
 
-    use crate::{compress_ctx, compress_taxi_data, download_taxi_data};
-
-    #[allow(dead_code)]
-    fn setup_logger(level: LevelFilter) {
-        TermLogger::init(
-            level,
-            Config::default(),
-            TerminalMode::Mixed,
-            ColorChoice::Auto,
-        )
-        .unwrap();
-    }
+    use crate::taxi_data::download_taxi_data;
+    use crate::{compress_ctx, compress_taxi_data, setup_logger};
 
     #[ignore]
     #[test]
     fn compression_ratio() {
-        setup_logger(LevelFilter::Info);
+        setup_logger(LevelFilter::Debug);
         _ = compress_taxi_data();
     }
 

bench-vortex/src/serde.rs

@@ -0,0 +1,87 @@
+use std::fs::File;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+
+use arrow_array::{ArrayRef as ArrowArrayRef, RecordBatchReader, StructArray as ArrowStructArray};
+use itertools::Itertools;
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+
+use vortex::array::chunked::ChunkedArray;
+use vortex::array::primitive::PrimitiveArray;
+use vortex::array::ArrayRef;
+use vortex::arrow::FromArrowType;
+use vortex::compute::take::take;
+use vortex::encode::FromArrowArray;
+use vortex::ptype::PType;
+use vortex::serde::{ReadCtx, WriteCtx};
+use vortex_schema::DType;
+
+use crate::taxi_data::download_taxi_data;
+use crate::{compress_ctx, idempotent};
+
+pub fn write_taxi_data() -> PathBuf {
+    idempotent("taxi.spiral", |write| {
+        let taxi_pq = File::open(download_taxi_data()).unwrap();
+        let builder = ParquetRecordBatchReaderBuilder::try_new(taxi_pq).unwrap();
+
+        // FIXME(ngates): the compressor should handle batch size.
+        let reader = builder.with_batch_size(65_536).build().unwrap();
+
+        let dtype = DType::from_arrow(reader.schema());
+        let ctx = compress_ctx();
+
+        let chunks = reader
+            .map(|batch_result| batch_result.unwrap())
+            .map(|record_batch| {
+                let struct_arrow: ArrowStructArray = record_batch.into();
+                let arrow_array: ArrowArrayRef = Arc::new(struct_arrow);
+                let vortex_array = ArrayRef::from_arrow(arrow_array.clone(), false);
+                ctx.compress(&vortex_array, None).unwrap()
+            })
+            .collect_vec();
+        let chunked = ChunkedArray::new(chunks, dtype.clone());
+
+        let mut write_ctx = WriteCtx::new(write);
+        write_ctx.dtype(&dtype).unwrap();
+        write_ctx.write(&chunked).unwrap();
+    })
+}
+
+pub fn take_taxi_data(path: &Path, indices: &[u64]) -> ArrayRef {
+    let chunked = {
+        let mut file = File::open(path).unwrap();
+        let dummy_dtype: DType = PType::U8.into();
+        let mut read_ctx = ReadCtx::new(&dummy_dtype, &mut file);
+        let dtype = read_ctx.dtype().unwrap();
+        read_ctx.with_schema(&dtype).read().unwrap()
+    };
+    take(&chunked, &PrimitiveArray::from(indices.to_vec())).unwrap()
+}
+
+#[cfg(test)]
+mod test {
+
+    use log::LevelFilter;
+    use simplelog::{ColorChoice, Config, TermLogger, TerminalMode};
+
+    use crate::serde::{take_taxi_data, write_taxi_data};
+
+    #[allow(dead_code)]
+    fn setup_logger(level: LevelFilter) {
+        TermLogger::init(
+            level,
+            Config::default(),
+            TerminalMode::Mixed,
+            ColorChoice::Auto,
+        )
+        .unwrap();
+    }
+
+    #[ignore]
+    #[test]
+    fn round_trip_serde() {
+        let taxi_spiral = write_taxi_data();
+        let rows = take_taxi_data(&taxi_spiral, &[10, 11, 12, 13, 100_000, 3_000_000]);
+        println!("TAKE TAXI DATA: {:?}", rows);
+    }
+}

bench-vortex/src/taxi_data.rs

@@ -0,0 +1,14 @@
+use std::path::PathBuf;
+
+use crate::idempotent;
+
+pub fn download_taxi_data() -> PathBuf {
+    idempotent("yellow-tripdata-2023-11.parquet", |file| {
+        reqwest::blocking::get(
+            "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2023-11.parquet",
+        )
+        .unwrap()
+        .copy_to(file)
+        .unwrap();
+    })
+}