diff --git a/Cargo.lock b/Cargo.lock index 0010a4e6cec..67b69a74d43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,9 +76,9 @@ checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" [[package]] name = "alp" -version = "0.0.2" +version = "0.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41ceb950ed9f82837662e3f3c470403987f32e3a25e5b741c8f17b3b0cca394d" +checksum = "639486620c15f839f53206c75be0d1a5a79c743c3e2a097d5e9e4b67fdb966fe" dependencies = [ "fastlanes", "itertools 0.15.0", @@ -1586,6 +1586,7 @@ dependencies = [ "vortex", "vortex-arrow", "vortex-bench", + "vortex-btrblocks", "vortex-cuda", ] @@ -7287,6 +7288,7 @@ name = "random-access-bench" version = "0.1.0" dependencies = [ "anyhow", + "async-trait", "clap", "indicatif", "lance-bench", @@ -7294,6 +7296,7 @@ dependencies = [ "rand_distr 0.6.0", "tabled", "tokio", + "vortex", "vortex-bench", ] @@ -9531,6 +9534,7 @@ dependencies = [ "vortex-alp", "vortex-array", "vortex-arrow", + "vortex-block-residual", "vortex-btrblocks", "vortex-buffer", "vortex-bytebool", @@ -9542,6 +9546,7 @@ dependencies = [ "vortex-fastlanes", "vortex-file", "vortex-flatbuffers", + "vortex-float-quant", "vortex-fsst", "vortex-io", "vortex-ipc", @@ -9729,6 +9734,7 @@ dependencies = [ "uuid", "vortex", "vortex-arrow", + "vortex-btrblocks", "vortex-spatial", "vortex-tensor", "wkb", @@ -9743,6 +9749,18 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "vortex-block-residual" +version = "0.1.0" +dependencies = [ + "fastlanes", + "rstest", + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-session", +] + [[package]] name = "vortex-btrblocks" version = "0.1.0" @@ -9760,12 +9778,14 @@ dependencies = [ "vortex-alp", "vortex-array", "vortex-arrow", + "vortex-block-residual", "vortex-buffer", "vortex-compressor", "vortex-datetime-parts", "vortex-decimal-byte-parts", "vortex-error", "vortex-fastlanes", + "vortex-float-quant", "vortex-fsst", "vortex-mask", "vortex-onpair", @@ -10148,6 +10168,7 @@ dependencies = [ "url", "vortex-alp", "vortex-array", + "vortex-block-residual", "vortex-btrblocks", "vortex-buffer", "vortex-bytebool", @@ -10157,6 +10178,7 @@ dependencies = [ "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", + "vortex-float-quant", "vortex-fsst", "vortex-io", "vortex-layout", @@ -10184,6 +10206,18 @@ dependencies = [ "vortex-error", ] +[[package]] +name = "vortex-float-quant" +version = "0.1.0" +dependencies = [ + "vortex-array", + "vortex-buffer", + "vortex-error", + "vortex-fastlanes", + "vortex-session", + "vortex-utils", +] + [[package]] name = "vortex-fsst" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 9abe9e379a2..c62cce728a8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,8 @@ members = [ "encodings/bytebool", "encodings/parquet-variant", "encodings/onpair", + "encodings/block-residual", + "encodings/float-quant", # Benchmarks "benchmarks/bench-support", "benchmarks/lance-bench", @@ -97,7 +99,7 @@ rust-version = "1.95" version = "0.1.0" [workspace.dependencies] -alp = "0.0.2" +alp = "0.0.3" anyhow = "1.0.100" arbitrary = "1.3.2" arc-swap = "1.9" @@ -297,6 +299,7 @@ vortex = { version = "0.1.0", path = "./vortex" } vortex-alp = { version = "0.1.0", path = "./encodings/alp", default-features = false } vortex-array = { version = "0.1.0", path = "./vortex-array", default-features = false } vortex-arrow = { version = "0.1.0", path = "./vortex-arrow", default-features = false } +vortex-block-residual = { version = "0.1.0", path = "./encodings/block-residual", default-features = false } vortex-btrblocks = { version = "0.1.0", path = "./vortex-btrblocks", default-features = false } vortex-buffer = { version = "0.1.0", path = "./vortex-buffer", default-features = false } vortex-bytebool = { version = "0.1.0", path = "./encodings/bytebool", default-features = false } @@ -311,6 +314,7 @@ vortex-error = { version = "0.1.0", path = "./vortex-error", default-features = vortex-fastlanes = { version = "0.1.0", path = "./encodings/fastlanes", default-features = false } vortex-file = { version = "0.1.0", path = "./vortex-file", default-features = false } vortex-flatbuffers = { version = "0.1.0", path = "./vortex-flatbuffers", default-features = false } +vortex-float-quant = { version = "0.1.0", path = "./encodings/float-quant", default-features = false } vortex-fsst = { version = "0.1.0", path = "./encodings/fsst", default-features = false } vortex-io = { version = "0.1.0", path = "./vortex-io", default-features = false } vortex-ipc = { version = "0.1.0", path = "./vortex-ipc", default-features = false } diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 7b1dadc3209..a887bd5948a 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -35,6 +35,7 @@ tracing = { workspace = true } vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } +vortex-btrblocks = { workspace = true } vortex-cuda = { workspace = true, optional = true } [features] diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 01b3e6e7d82..39c773e760f 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,6 +15,26 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format cargo run -p compress-bench --profile release_debug ``` +Compare the default compressor, Compact, and Parquet with Zstd: + +```bash +cargo run -p compress-bench --profile release_debug -- \ + --formats vortex,vortex-compact,parquet +``` + +Compare a numeric scheme bundle with the same command and one of these values: + +- `--vortex-numeric-bundle prior-default` +- `--vortex-numeric-bundle block-residual` +- `--vortex-numeric-bundle current-default` + +Add a local Parquet file with `--parquet-path`. Use `--datasets` to select its file stem. + +```bash +cargo run -p compress-bench --profile release_debug -- \ + --parquet-path /tmp/input.parquet --datasets '^input$' +``` + ## GPU decompression `--gpu-decompress` is opt-in, requires the `cuda` feature, and restricts the suite to the diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index cbf0378f344..5a515effa0a 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -24,6 +24,7 @@ use vortex_bench::Engine; use vortex_bench::Format; use vortex_bench::LogFormat; use vortex_bench::Target; +use vortex_bench::VortexNumericBundle; use vortex_bench::compress::CompressMeasurements; use vortex_bench::compress::CompressOp; use vortex_bench::compress::Compressor; @@ -32,6 +33,9 @@ use vortex_bench::compress::benchmark_decompress; use vortex_bench::compress::calculate_ratios; use vortex_bench::create_output_writer; use vortex_bench::datasets::Dataset; +use vortex_bench::datasets::feature_vectors::GloveEmbeddingsData; +use vortex_bench::datasets::feature_vectors::OpenAiEmbeddingsData; +use vortex_bench::datasets::local_parquet::LocalParquetData; use vortex_bench::datasets::struct_list_of_ints::StructListOfInts; use vortex_bench::datasets::taxi_data::TaxiData; use vortex_bench::datasets::tpch_l_comment::TPCHLCommentCanonical; @@ -73,6 +77,9 @@ struct Args { ops: Vec, #[arg(long)] datasets: Option, + /// Add a local Parquet file to the benchmark suite. + #[arg(long)] + parquet_path: Vec, /// Print the dataset names that would run, one per line, and exit. /// Knowledge of datasets lies only in this binary so we need /// orchestrator to know whan queries to run one by one. @@ -111,6 +118,9 @@ struct Args { ingest_output: Option, #[arg(long)] tracing: bool, + /// Select the numeric scheme bundle for Vortex compression. + #[arg(long, value_enum, default_value_t)] + vortex_numeric_bundle: VortexNumericBundle, /// Format for the primary stderr log sink. `text` is the default human-readable format; /// `json` emits one JSON object per event, suitable for piping into `jq`. #[arg(long, value_enum, default_value_t = LogFormat::Text)] @@ -149,6 +159,7 @@ async fn main() -> anyhow::Result<()> { run_compress( args.iterations, args.datasets.map(|d| Regex::new(&d)).transpose()?, + args.parquet_path, formats, ops, mode, @@ -156,6 +167,7 @@ async fn main() -> anyhow::Result<()> { args.display_format, args.output_path, args.ingest_output, + args.vortex_numeric_bundle, ) .await } @@ -182,13 +194,24 @@ impl BenchMode { } /// Get a compressor for the given format. -fn get_compressor(format: Format, mode: BenchMode) -> Box { +fn get_compressor( + format: Format, + mode: BenchMode, + vortex_numeric_bundle: VortexNumericBundle, +) -> Box { if let BenchMode::Gpu(options) = mode { return gpu_compressor(format, options); } match format { - Format::OnDiskVortex => Box::new(VortexCompressor), + Format::OnDiskVortex => Box::new(VortexCompressor::new( + Format::OnDiskVortex, + vortex_numeric_bundle, + )), + Format::VortexCompact => Box::new(VortexCompressor::new( + Format::VortexCompact, + VortexNumericBundle::CurrentDefault, + )), Format::Parquet => Box::new(ParquetCompressor::new()), #[cfg(feature = "lance")] Format::Lance => Box::new(LanceCompressor), @@ -209,6 +232,7 @@ const DOC_PATH: &str = "benchmarks/compress-bench/README.md"; async fn run_compress( iterations: usize, datasets_filter: Option, + parquet_paths: Vec, formats: Vec, ops: Vec, mode: BenchMode, @@ -216,6 +240,7 @@ async fn run_compress( display_format: DisplayFormat, output_path: Option, ingest_output: Option, + vortex_numeric_bundle: VortexNumericBundle, ) -> anyhow::Result<()> { let targets = formats .iter() @@ -238,6 +263,10 @@ async fn run_compress( // Some(READ_PROJECTION_COLUMNS), // ), ]; + let local_parquet = parquet_paths + .into_iter() + .map(LocalParquetData::try_new) + .collect::>>()?; // Datasets run in GPU mode. Add one only after a `--gpu-verify` run has confirmed its CUDA // decode end to end; a dataset here that cannot decode fails the benchmark job. Between them @@ -262,6 +291,8 @@ async fn run_compress( let all_datasets: Vec<&dyn Dataset> = [ &TaxiData as &dyn Dataset, + &GloveEmbeddingsData, + &OpenAiEmbeddingsData, PBI_DATASETS.get(Arade), PBI_DATASETS.get(Bimbo), PBI_DATASETS.get(CMSprovider), @@ -273,13 +304,14 @@ async fn run_compress( // Hatred, // panic in fsst_compress_iter // TableroSistemaPenal, // Unexpected type error // YaleLanguages, // 4th column looks like integer but also contains Y - &TPCHLCommentChunked, + &TPCHLCommentChunked as &dyn Dataset, &TPCHLCommentCanonical, &DownloadableDataset::RPlace, &DownloadableDataset::AirQuality, ] .into_iter() .chain(structlistofints.iter().map(|d| d as &dyn Dataset)) + .chain(local_parquet.iter().map(|d| d as &dyn Dataset)) .collect(); let datasets: Vec<&dyn Dataset> = if mode.is_gpu() { @@ -316,9 +348,16 @@ async fn run_compress( let survey_all = mode.is_gpu(); let mut failures: Vec<(String, anyhow::Error)> = Vec::new(); - for dataset_handle in datasets.into_iter() { - let run = - run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, mode); + for dataset_handle in datasets { + let run = run_benchmark_for_dataset( + &progress, + &formats, + &ops, + iterations, + dataset_handle, + mode, + vortex_numeric_bundle, + ); // Missing CUDA kernel support surfaces as a panic rather than an error, so the survey // has to catch those too or the first unsupported dataset ends the run. @@ -403,6 +442,7 @@ async fn run_benchmark_for_dataset( iterations: usize, dataset_handle: &dyn Dataset, mode: BenchMode, + vortex_numeric_bundle: VortexNumericBundle, ) -> anyhow::Result<(CompressMeasurements, Vec)> { let bench_name = dataset_handle.name(); // A GPU decode and a host decode of the same dataset would otherwise publish the same @@ -427,8 +467,7 @@ async fn run_benchmark_for_dataset( let mut v3_records: Vec = Vec::new(); for format in formats { - let compressor = get_compressor(*format, mode); - + let compressor = get_compressor(*format, mode, vortex_numeric_bundle); for op in ops { let time = match op { CompressOp::Compress => { diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 20b4b9f1402..545284ef176 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -17,21 +17,48 @@ use vortex::dtype::FieldNames; use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexWriteOptions; use vortex::file::WriteOptionsSessionExt; use vortex_arrow::ArrowSessionExt; +use vortex_bench::CompactionStrategy; use vortex_bench::Format; use vortex_bench::SESSION; +use vortex_bench::VortexNumericBundle; use vortex_bench::compress::Compressor; use vortex_bench::compress::read_projection; use vortex_bench::conversions::parquet_to_vortex_chunks; /// Compressor implementation for Vortex format. -pub struct VortexCompressor; +pub struct VortexCompressor { + format: Format, + numeric_bundle: VortexNumericBundle, +} + +impl VortexCompressor { + pub fn new(format: Format, numeric_bundle: VortexNumericBundle) -> Self { + assert!(matches!( + format, + Format::OnDiskVortex | Format::VortexCompact + )); + Self { + format, + numeric_bundle, + } + } + + fn write_options(&self) -> VortexWriteOptions { + let options = SESSION.write_options(); + if self.format == Format::VortexCompact { + return CompactionStrategy::Compact.apply_options(options); + } + self.numeric_bundle.apply_options(options) + } +} #[async_trait] impl Compressor for VortexCompressor { fn format(&self) -> Format { - Format::OnDiskVortex + self.format } async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)> { @@ -41,8 +68,7 @@ impl Compressor for VortexCompressor { let mut buf = Vec::new(); let start = Instant::now(); let mut cursor = Cursor::new(&mut buf); - SESSION - .write_options() + self.write_options() .write(&mut cursor, uncompressed.into_array().to_array_stream()) .await?; let elapsed = start.elapsed(); @@ -55,8 +81,7 @@ impl Compressor for VortexCompressor { let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; let mut buf = Vec::new(); let mut cursor = Cursor::new(&mut buf); - SESSION - .write_options() + self.write_options() .write(&mut cursor, uncompressed.into_array().to_array_stream()) .await?; diff --git a/benchmarks/random-access-bench/Cargo.toml b/benchmarks/random-access-bench/Cargo.toml index c45c1ac2bff..37633685d0a 100644 --- a/benchmarks/random-access-bench/Cargo.toml +++ b/benchmarks/random-access-bench/Cargo.toml @@ -16,6 +16,7 @@ publish = false [dependencies] anyhow = { workspace = true } +async-trait = { workspace = true } clap = { workspace = true, features = ["derive"] } indicatif = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } @@ -23,6 +24,7 @@ rand = { workspace = true } rand_distr = { workspace = true } tabled = { workspace = true } tokio = { workspace = true, features = ["full"] } +vortex = { workspace = true } vortex-bench = { workspace = true } [features] diff --git a/benchmarks/random-access-bench/README.md b/benchmarks/random-access-bench/README.md index 14949c85fcb..d9faaaf1dbe 100644 --- a/benchmarks/random-access-bench/README.md +++ b/benchmarks/random-access-bench/README.md @@ -20,3 +20,20 @@ reopening the file per lookup. CI drives the full matrix via ```bash cargo run -p random-access-bench --profile release_debug --features lance ``` + +Compare random access for a numeric scheme bundle with one of these values: + +- `--vortex-numeric-bundle prior-default` +- `--vortex-numeric-bundle block-residual` +- `--vortex-numeric-bundle current-default` + +Each bundle uses a separate Vortex file. Existing files remain available for repeated runs. + +Add a local Parquet file with `--parquet-path`. A local path replaces the default dataset list. + +```bash +cargo run -p random-access-bench --profile release_debug -- \ + --formats vortex --parquet-path /tmp/input.parquet +``` + +If `--datasets` also selects built-in datasets, the benchmark runs both groups. diff --git a/benchmarks/random-access-bench/src/main.rs b/benchmarks/random-access-bench/src/main.rs index 11ef81edd3c..59d3748efc8 100644 --- a/benchmarks/random-access-bench/src/main.rs +++ b/benchmarks/random-access-bench/src/main.rs @@ -4,13 +4,19 @@ use std::path::PathBuf; use anyhow::Result; +use async_trait::async_trait; use clap::Parser; use clap::ValueEnum; use random_access_bench::AccessPattern; use random_access_bench::OpenMode; use random_access_bench::RunConfig; +use vortex::file::WriteOptionsSessionExt; use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::VortexNumericBundle; +use vortex_bench::conversions::write_parquet_as_vortex_with_options; use vortex_bench::datasets::feature_vectors::FeatureVectorsData; +use vortex_bench::datasets::local_parquet::LocalParquetData; use vortex_bench::datasets::nested_lists::NestedListsData; use vortex_bench::datasets::nested_structs::NestedStructsData; use vortex_bench::datasets::taxi_data::TaxiData; @@ -42,6 +48,34 @@ impl DatasetArg { } } +struct NumericBundleDataset { + dataset: Box, + bundle: VortexNumericBundle, +} + +#[async_trait] +impl BenchDataset for NumericBundleDataset { + fn name(&self) -> &str { + self.dataset.name() + } + + fn row_count(&self) -> u64 { + self.dataset.row_count() + } + + async fn path(&self, format: Format) -> Result { + if format != Format::OnDiskVortex { + return self.dataset.path(format).await; + } + + let parquet_path = self.dataset.path(Format::Parquet).await?; + let name = self.dataset.name(); + let path = format!("random_access/{name}/{name}-{}.vortex", self.bundle.name()); + let options = self.bundle.apply_options(SESSION.write_options()); + write_parquet_as_vortex_with_options(parquet_path, &path, options).await + } +} + #[derive(Parser, Debug)] #[command(version, about, long_about = None)] struct Args { @@ -67,13 +101,11 @@ struct Args { #[arg(long = "ingest-jsonl")] ingest_output: Option, /// Which datasets to benchmark random access on. - #[arg( - long, - value_delimiter = ',', - value_enum, - default_values_t = vec![DatasetArg::Taxi, DatasetArg::FeatureVectors, DatasetArg::NestedLists, DatasetArg::NestedStructs] - )] + #[arg(long, value_delimiter = ',', value_enum)] datasets: Vec, + /// Add a local Parquet file to the benchmark suite. + #[arg(long)] + parquet_path: Vec, /// Which access patterns to benchmark. #[arg( long, @@ -85,6 +117,9 @@ struct Args { /// Whether to reopen the file on each iteration, use a cached handle, or run both. #[arg(long, value_enum, default_value_t = OpenMode::Both)] open_mode: OpenMode, + /// Select the numeric scheme bundle for Vortex files. + #[arg(long, value_enum, default_value_t)] + vortex_numeric_bundle: VortexNumericBundle, } #[tokio::main] @@ -92,11 +127,37 @@ async fn main() -> Result<()> { let args = Args::parse(); setup_logging_and_tracing(args.verbose, args.tracing)?; + let dataset_args = if args.datasets.is_empty() && args.parquet_path.is_empty() { + vec![ + DatasetArg::Taxi, + DatasetArg::FeatureVectors, + DatasetArg::NestedLists, + DatasetArg::NestedStructs, + ] + } else { + args.datasets + }; + let mut datasets = dataset_args + .into_iter() + .map(DatasetArg::into_dataset) + .collect::>(); + datasets.extend( + args.parquet_path + .into_iter() + .map(LocalParquetData::try_new) + .map(|dataset| dataset.map(|dataset| Box::new(dataset) as Box)) + .collect::>>()?, + ); + let run_config = RunConfig { - datasets: args - .datasets + datasets: datasets .into_iter() - .map(DatasetArg::into_dataset) + .map(|dataset| { + Box::new(NumericBundleDataset { + dataset, + bundle: args.vortex_numeric_bundle, + }) as Box + }) .collect(), formats: args.formats, patterns: args.patterns, diff --git a/encodings/block-residual/Cargo.toml b/encodings/block-residual/Cargo.toml new file mode 100644 index 00000000000..e0bb52093d0 --- /dev/null +++ b/encodings/block-residual/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "vortex-block-residual" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex ordered-float and block-residual array encodings" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[lints] +workspace = true + +[dependencies] +fastlanes = { workspace = true } +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-session = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/encodings/block-residual/src/block_residual_array.rs b/encodings/block-residual/src/block_residual_array.rs new file mode 100644 index 00000000000..5e5babd5cde --- /dev/null +++ b/encodings/block-residual/src/block_residual_array.rs @@ -0,0 +1,1656 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; +use std::ops::Range; + +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::validity_to_child; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::BlockResidualCodec; +use crate::BlockResidualParts; +use crate::codec::BlockResidualCodecEstimate; +use crate::codec::ResidualWord; +use crate::codec::packed_words_as_native; +use crate::codec::read_wide_bits; + +const BLOCK_LEN: usize = 1024; +const METADATA_VERSION: u8 = 2; +const METADATA_LEN: usize = 41; + +/// Ordered unsigned integers with one reference and packed residuals per block. +pub type BlockResidualArray = Array; + +#[array_slots(BlockResidual)] +pub struct BlockResidualSlots { + #[slot(0)] + pub validity: Option, +} + +#[derive(Clone, Debug)] +pub struct BlockResidualData { + unsliced_len: usize, + slice_start: usize, + slice_stop: usize, + payload: ByteBuffer, + bases: Buffer, + residual_widths: Buffer, + high_widths: Buffer, + residual_starts: Buffer, + patch_starts: Buffer, + high_starts: Buffer, + residual_words: Buffer, + patch_positions: Buffer, + patch_highs: Buffer, +} + +impl Display for BlockResidualData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "blocks: {}, slice: {}..{}", + self.unsliced_len.div_ceil(BLOCK_LEN), + self.slice_start, + self.slice_stop + ) + } +} + +impl ArrayHash for BlockResidualData { + fn array_hash(&self, state: &mut H, accuracy: EqMode) { + self.unsliced_len.hash(state); + self.slice_start.hash(state); + self.slice_stop.hash(state); + self.payload.array_hash(state, accuracy); + } +} + +impl ArrayEq for BlockResidualData { + fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { + self.unsliced_len == other.unsliced_len + && self.slice_start == other.slice_start + && self.slice_stop == other.slice_stop + && self.payload.array_eq(&other.payload, accuracy) + } +} + +#[derive(Clone, Debug)] +pub struct BlockResidual; + +/// Exact encoded size and patch count without materialized payloads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct BlockResidualEstimate { + nbytes: u64, + patch_count: usize, +} + +impl BlockResidualEstimate { + pub(crate) fn try_new( + encoded_nbytes: usize, + validity_nbytes: u64, + patch_count: usize, + ) -> VortexResult { + let nbytes = u64::try_from(encoded_nbytes)? + .checked_add(validity_nbytes) + .ok_or_else(|| vortex_error::vortex_err!("BlockResidual estimate size overflow"))?; + Ok(Self { + nbytes, + patch_count, + }) + } + + /// Return the estimated physical bytes. + pub fn nbytes(self) -> u64 { + self.nbytes + } + + /// Return the estimated patch count. + pub fn patch_count(self) -> usize { + self.patch_count + } +} + +impl VTable for BlockResidual { + type TypedArrayData = BlockResidualData; + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.block_residual"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let slots = BlockResidualSlotsView::from_slots(slots); + let validity = child_to_validity(slots.validity, dtype.nullability()); + data.validate(dtype, len, slots, &validity) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 1 + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + match idx { + 0 => BufferHandle::new_host(array.payload.clone()), + _ => vortex_panic!("BlockResidualArray buffer index {idx} out of bounds"), + } + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + match idx { + 0 => Some("payload".to_string()), + _ => vortex_panic!("BlockResidualArray buffer_name {idx} out of bounds"), + } + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!(buffers.len() == 1, "BlockResidualArray expects one buffer"); + let mut data = array.data().clone(); + data.replace_payload(&buffers[0])?; + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + BlockResidualMetadata::from_data(array.data())?.encode(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = BlockResidualMetadata::decode(metadata)?; + let unsliced_len = usize::try_from(metadata.unsliced_len)?; + let slice_start = usize::try_from(metadata.slice_start)?; + let slice_stop = slice_start + .checked_add(len) + .ok_or_else(|| vortex_error::vortex_err!("block residual slice length overflows"))?; + let residual_word_count = usize::try_from(metadata.residual_word_count)?; + let patch_count = usize::try_from(metadata.patch_count)?; + let patch_high_count = usize::try_from(metadata.patch_high_count)?; + vortex_ensure!(buffers.len() == 1, "BlockResidualArray expects one buffer"); + let validity = match children.len() { + 0 => Validity::from(dtype.nullability()), + 1 => Validity::Array(children.get(0, &Validity::DTYPE, unsliced_len)?), + count => vortex_bail!("BlockResidualArray expects zero or one child, got {count}"), + }; + let slots = BlockResidualSlots { + validity: validity_to_child(&validity, unsliced_len), + } + .into_slots(); + let data = BlockResidualData::try_new( + unsliced_len, + slice_start, + slice_stop, + residual_word_count, + patch_count, + patch_high_count, + host_payload(&buffers[0])?, + )?; + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + BlockResidualSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + decompress_array(array.as_view(), ctx)?.into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for BlockResidual { + fn scalar_at( + array: ArrayView<'_, BlockResidual>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !array.as_ref().is_valid(index, ctx)? { + return Ok(Scalar::null(array.dtype().clone())); + } + let value = scalar_from_array(array, index, ctx)?; + let nullability = array.dtype().nullability(); + Ok(match array.dtype().as_ptype() { + PType::U8 => Scalar::primitive(u8::try_from(value)?, nullability), + PType::U16 => Scalar::primitive(u16::try_from(value)?, nullability), + PType::U32 => Scalar::primitive(u32::try_from(value)?, nullability), + PType::U64 => Scalar::primitive(value, nullability), + PType::I8 => Scalar::primitive( + i8::from_le_bytes([(u8::try_from(value)? ^ (1_u8 << 7))]), + nullability, + ), + PType::I16 => Scalar::primitive( + i16::from_le_bytes((u16::try_from(value)? ^ (1_u16 << 15)).to_le_bytes()), + nullability, + ), + PType::I32 => Scalar::primitive( + i32::from_le_bytes((u32::try_from(value)? ^ (1_u32 << 31)).to_le_bytes()), + nullability, + ), + PType::I64 => Scalar::primitive( + i64::from_le_bytes((value ^ (1_u64 << 63)).to_le_bytes()), + nullability, + ), + ptype => vortex_bail!("BlockResidual scalar access does not support {ptype}"), + }) + } +} + +impl ValidityVTable for BlockResidual { + fn validity(array: ArrayView<'_, BlockResidual>) -> VortexResult { + array + .unsliced_validity() + .slice(array.data().slice_start..array.data().slice_stop) + } +} + +impl SliceReduce for BlockResidual { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + let data = array.data().slice(range); + let parts = ArrayParts::new(BlockResidual, array.dtype().clone(), data.len(), data) + .with_slots(array.slots().iter().cloned().collect()); + // SAFETY: The source array is valid. The slice only narrows its logical bounds. + Ok(Some( + unsafe { Array::from_parts_unchecked(parts) }.into_array(), + )) + } +} + +static RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(BlockResidual))]); + +pub(crate) trait BlockResidualArrayExt: TypedArrayRef { + fn unsliced_validity(&self) -> Validity { + child_to_validity( + self.as_ref().slots()[BlockResidualSlots::VALIDITY].as_ref(), + self.as_ref().dtype().nullability(), + ) + } + + /// Return the reference value for each block. + fn bases(&self) -> &[u64] { + &self.bases + } + + /// Return the packed residual width for each block. + fn residual_widths(&self) -> &[u8] { + &self.residual_widths + } + + /// Return the packed patch width for each block. + fn high_widths(&self) -> &[u8] { + &self.high_widths + } + + /// Return the residual payload offsets. + fn residual_starts(&self) -> &[u32] { + &self.residual_starts + } + + /// Return the patch position offsets. + fn patch_starts(&self) -> &[u32] { + &self.patch_starts + } + + /// Return the patch high-bit offsets. + fn high_starts(&self) -> &[u32] { + &self.high_starts + } + + /// Return the packed residual payload. + fn residual_words(&self) -> &[u64] { + &self.residual_words + } + + /// Return the patch positions. + fn patch_positions(&self) -> &[u16] { + &self.patch_positions + } + + /// Return the packed patch high bits. + fn patch_highs(&self) -> &[u8] { + &self.patch_highs + } +} + +impl> BlockResidualArrayExt for T {} + +impl BlockResidual { + /// Estimate the exact encoded size without materializing packed payloads. + pub fn estimate_primitive( + array: ArrayView<'_, Primitive>, + ) -> VortexResult { + vortex_ensure!( + array.ptype().is_int(), + "BlockResidual requires integer values" + ); + let BlockResidualCodecEstimate { + encoded_nbytes, + patch_count, + } = match array.ptype() { + PType::U8 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), u64::from) + } + PType::U16 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), u64::from) + } + PType::U32 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), u64::from) + } + PType::U64 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| value) + } + PType::I8 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + u64::from((value as u8) ^ (1_u8 << 7)) + }) + } + PType::I16 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + u64::from((value as u16) ^ (1_u16 << 15)) + }) + } + PType::I32 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + u64::from((value as u32) ^ (1_u32 << 31)) + }) + } + PType::I64 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + (value as u64) ^ (1_u64 << 63) + }) + } + ptype => vortex_bail!("BlockResidual does not support {ptype}"), + }; + let validity_nbytes = validity_to_child(&array.validity()?, array.len()) + .map(|validity| validity.nbytes()) + .unwrap_or(0); + BlockResidualEstimate::try_new(encoded_nbytes, validity_nbytes, patch_count) + } + + /// Encode an integer array in independent blocks. + pub fn from_primitive(array: ArrayView<'_, Primitive>) -> VortexResult { + vortex_ensure!( + array.ptype().is_int(), + "BlockResidual requires integer values" + ); + let validity = array.validity()?; + let values = ordered_values(array)?; + let parts = BlockResidualCodec::encode_with_word_width( + &values, + u8::try_from(array.ptype().bit_width())?, + )? + .into_parts()?; + Self::try_new(parts, validity, array.ptype()) + } + + fn try_new( + parts: BlockResidualParts, + validity: Validity, + ptype: PType, + ) -> VortexResult { + let payload = payload_from_parts(&parts)?; + let data = BlockResidualData::try_new( + parts.len, + 0, + parts.len, + parts.residual_words.len(), + parts.patch_positions.len(), + parts.patch_highs.len(), + payload, + )?; + let slots = BlockResidualSlots { + validity: validity_to_child(&validity, data.unsliced_len), + } + .into_slots(); + Array::try_from_parts( + ArrayParts::new( + BlockResidual, + DType::Primitive(ptype, validity.nullability()), + data.unsliced_len, + data, + ) + .with_slots(slots), + ) + } +} + +fn ordered_values(array: ArrayView<'_, Primitive>) -> VortexResult> { + Ok(match array.ptype() { + PType::U8 => array + .as_slice::() + .iter() + .map(|&value| u64::from(value)) + .collect(), + PType::U16 => array + .as_slice::() + .iter() + .map(|&value| u64::from(value)) + .collect(), + PType::U32 => array + .as_slice::() + .iter() + .map(|&value| u64::from(value)) + .collect(), + PType::U64 => array.as_slice::().to_vec(), + PType::I8 => array + .as_slice::() + .iter() + .map(|&value| u64::from((value as u8) ^ (1_u8 << 7))) + .collect(), + PType::I16 => array + .as_slice::() + .iter() + .map(|&value| u64::from((value as u16) ^ (1_u16 << 15))) + .collect(), + PType::I32 => array + .as_slice::() + .iter() + .map(|&value| u64::from((value as u32) ^ (1_u32 << 31))) + .collect(), + PType::I64 => array + .as_slice::() + .iter() + .map(|&value| (value as u64) ^ (1_u64 << 63)) + .collect(), + ptype => vortex_bail!("BlockResidual does not support {ptype}"), + }) +} + +impl BlockResidualData { + fn try_new( + unsliced_len: usize, + slice_start: usize, + slice_stop: usize, + residual_word_count: usize, + patch_count: usize, + patch_high_count: usize, + payload: ByteBuffer, + ) -> VortexResult { + let block_count = unsliced_len.div_ceil(BLOCK_LEN); + let mut offset = 0; + let bases = take_payload(&payload, &mut offset, block_count, "bases")?; + let residual_words = + take_payload(&payload, &mut offset, residual_word_count, "residual words")?; + let residual_starts = + take_payload(&payload, &mut offset, block_count + 1, "residual starts")?; + let patch_starts = take_payload(&payload, &mut offset, block_count + 1, "patch starts")?; + let high_starts = take_payload(&payload, &mut offset, block_count + 1, "high starts")?; + let patch_positions = take_payload(&payload, &mut offset, patch_count, "patch positions")?; + let residual_widths = take_payload(&payload, &mut offset, block_count, "residual widths")?; + let high_widths = take_payload(&payload, &mut offset, block_count, "high widths")?; + let patch_highs = take_payload(&payload, &mut offset, patch_high_count, "patch highs")?; + vortex_ensure!( + offset == payload.len(), + "block residual payload contains trailing bytes" + ); + Ok(Self { + unsliced_len, + slice_start, + slice_stop, + payload, + bases, + residual_widths, + high_widths, + residual_starts, + patch_starts, + high_starts, + residual_words, + patch_positions, + patch_highs, + }) + } + + fn validate( + &self, + dtype: &DType, + len: usize, + _slots: BlockResidualSlotsView<'_>, + validity: &Validity, + ) -> VortexResult<()> { + vortex_ensure!( + dtype.is_int(), + "BlockResidualArray requires an integer dtype" + ); + vortex_ensure!( + self.slice_start <= self.slice_stop && self.slice_stop <= self.unsliced_len, + "block residual slice exceeds its source length" + ); + vortex_ensure!(len == self.len(), "block residual slice length is invalid"); + let block_count = self.unsliced_len.div_ceil(BLOCK_LEN); + vortex_ensure!( + self.bases.len() == block_count + && self.residual_widths.len() == block_count + && self.high_widths.len() == block_count, + "block residual block tables have invalid lengths" + ); + vortex_ensure!( + self.residual_starts.len() == block_count + 1 + && self.patch_starts.len() == block_count + 1 + && self.high_starts.len() == block_count + 1, + "block residual offset tables have invalid lengths" + ); + validate_offset_table( + &self.residual_starts, + block_count, + self.residual_words.len(), + "residual", + )?; + validate_offset_table( + &self.patch_starts, + block_count, + self.patch_positions.len(), + "patch", + )?; + validate_offset_table( + &self.high_starts, + block_count, + self.patch_highs.len(), + "patch high", + )?; + let logical_width = dtype.as_ptype().bit_width(); + let maximum = if logical_width == 64 { + u64::MAX + } else { + (1_u64 << logical_width) - 1 + }; + for block_index in 0..block_count { + vortex_ensure!( + self.bases[block_index] <= maximum, + "block residual base exceeds its logical type" + ); + let residual_width = self.residual_widths[block_index]; + let high_width = self.high_widths[block_index]; + vortex_ensure!( + usize::from(residual_width) <= logical_width + && usize::from(high_width) <= logical_width + && usize::from(residual_width) + usize::from(high_width) <= logical_width, + "block residual bit widths are invalid" + ); + let residual_range = payload_range( + &self.residual_starts, + block_index, + self.residual_words.len(), + "residual", + )?; + vortex_ensure!( + residual_range.len() == BLOCK_LEN * usize::from(residual_width) / 64, + "block residual word count is invalid" + ); + let patch_range = payload_range( + &self.patch_starts, + block_index, + self.patch_positions.len(), + "patch", + )?; + let high_range = payload_range( + &self.high_starts, + block_index, + self.patch_highs.len(), + "patch high", + )?; + let positions = &self.patch_positions[patch_range]; + validate_patch_header( + residual_width, + high_width, + positions.len(), + high_range.len(), + )?; + let block_start = block_index * BLOCK_LEN; + let block_len = (self.unsliced_len - block_start).min(BLOCK_LEN); + let mut previous = None; + for &position in positions { + validate_patch_position(block_len, previous, position)?; + previous = Some(position); + } + } + if let Some(validity_len) = validity.maybe_len() { + vortex_ensure!( + validity_len == self.unsliced_len, + "block residual validity length is invalid" + ); + } + Ok(()) + } + + fn len(&self) -> usize { + self.slice_stop - self.slice_start + } + + fn slice(&self, range: Range) -> Self { + Self { + slice_start: self.slice_start + range.start, + slice_stop: self.slice_start + range.end, + ..self.clone() + } + } + + fn replace_payload(&mut self, buffer: &BufferHandle) -> VortexResult<()> { + *self = Self::try_new( + self.unsliced_len, + self.slice_start, + self.slice_stop, + self.residual_words.len(), + self.patch_positions.len(), + self.patch_highs.len(), + host_payload(buffer)?, + )?; + Ok(()) + } +} + +fn host_payload(buffer: &BufferHandle) -> VortexResult { + buffer + .clone() + .ensure_aligned(Alignment::of::())? + .try_into_host_sync() +} + +fn take_payload( + payload: &ByteBuffer, + offset: &mut usize, + len: usize, + name: &str, +) -> VortexResult> { + let nbytes = len + .checked_mul(size_of::()) + .ok_or_else(|| vortex_error::vortex_err!("block residual {name} size overflows"))?; + let stop = offset + .checked_add(nbytes) + .ok_or_else(|| vortex_error::vortex_err!("block residual {name} offset overflows"))?; + vortex_ensure!( + stop <= payload.len(), + "block residual {name} exceeds the payload" + ); + let bytes = payload.slice_with_alignment(*offset..stop, Alignment::of::()); + *offset = stop; + Ok(Buffer::from_byte_buffer(bytes)) +} + +fn payload_from_parts(parts: &BlockResidualParts) -> VortexResult { + let total_nbytes = [ + size_of_val(parts.bases.as_slice()), + size_of_val(parts.residual_words.as_slice()), + size_of_val(parts.residual_starts.as_slice()), + size_of_val(parts.patch_starts.as_slice()), + size_of_val(parts.high_starts.as_slice()), + size_of_val(parts.patch_positions.as_slice()), + size_of_val(parts.residual_widths.as_slice()), + size_of_val(parts.high_widths.as_slice()), + size_of_val(parts.patch_highs.as_slice()), + ] + .into_iter() + .try_fold(0_usize, |total, nbytes| total.checked_add(nbytes)) + .ok_or_else(|| vortex_error::vortex_err!("block residual payload size overflows"))?; + let mut payload = ByteBufferMut::with_capacity_aligned(total_nbytes, Alignment::of::()); + append_native(&mut payload, &parts.bases); + append_native(&mut payload, &parts.residual_words); + append_native(&mut payload, &parts.residual_starts); + append_native(&mut payload, &parts.patch_starts); + append_native(&mut payload, &parts.high_starts); + append_native(&mut payload, &parts.patch_positions); + append_native(&mut payload, &parts.residual_widths); + append_native(&mut payload, &parts.high_widths); + append_native(&mut payload, &parts.patch_highs); + Ok(payload.freeze()) +} + +fn append_native(payload: &mut ByteBufferMut, values: &[T]) { + // SAFETY: NativePType values contain no padding and permit every initialized bit pattern. + let bytes = + unsafe { std::slice::from_raw_parts(values.as_ptr().cast::(), size_of_val(values)) }; + payload.extend_from_slice(bytes); +} + +fn decode_array_values( + array: ArrayView<'_, BlockResidual>, + _ctx: &mut ExecutionCtx, + mut transform: impl FnMut(U) -> T, +) -> VortexResult { + let bases = array.bases(); + let residual_widths = array.residual_widths(); + let high_widths = array.high_widths(); + let residual_starts = array.residual_starts(); + let patch_starts = array.patch_starts(); + let high_starts = array.high_starts(); + let residual_words = array.residual_words(); + let patch_positions = array.patch_positions(); + let patch_highs = array.patch_highs(); + let logical_range = array.data().slice_start..array.data().slice_stop; + let mut direct_values = + DIRECT_OUTPUT.then(|| BufferMut::::with_capacity(logical_range.len())); + let mut transformed_values = + (!DIRECT_OUTPUT).then(|| BufferMut::::with_capacity(logical_range.len())); + let mut residuals = [U::default(); BLOCK_LEN]; + + let first_block = logical_range.start / BLOCK_LEN; + let last_block = logical_range.end.div_ceil(BLOCK_LEN); + for block_index in first_block..last_block { + let block_start = block_index * BLOCK_LEN; + let block_len = (array.data().unsliced_len - block_start).min(BLOCK_LEN); + let block_stop = block_start + block_len; + if block_stop <= logical_range.start || block_start >= logical_range.end { + continue; + } + + let residual_width = residual_widths[block_index]; + let high_width = high_widths[block_index]; + let base = ::from_u64(bases[block_index]); + vortex_ensure!( + residual_width <= U::BITS + && high_width <= U::BITS + && u16::from(residual_width) + u16::from(high_width) <= u16::from(U::BITS), + "block residual bit widths are invalid" + ); + let residual_payload = payload_range( + residual_starts, + block_index, + residual_words.len(), + "residual", + )?; + vortex_ensure!( + residual_payload.len() == BLOCK_LEN * usize::from(residual_width) / 64, + "block residual word count is invalid" + ); + let patch_payload = + payload_range(patch_starts, block_index, patch_positions.len(), "patch")?; + let high_payload = + payload_range(high_starts, block_index, patch_highs.len(), "patch high")?; + let positions = &patch_positions[patch_payload]; + validate_patch_header( + residual_width, + high_width, + positions.len(), + high_payload.len(), + )?; + let highs = &patch_highs[high_payload]; + let local_start = logical_range.start.saturating_sub(block_start); + let local_stop = (logical_range.end - block_start).min(block_len); + + if residual_width == 0 { + let output_start = decoded_len(direct_values.as_ref(), transformed_values.as_ref()); + append_repeated_decoded::( + &mut direct_values, + &mut transformed_values, + base, + local_stop - local_start, + &mut transform, + ); + let mut previous_position = None; + for (patch_index, &position) in positions.iter().enumerate() { + validate_patch_position(block_len, previous_position, position)?; + previous_position = Some(position); + let position = usize::from(position); + if position < local_start || position >= local_stop { + continue; + } + // SAFETY: The payload includes fifteen readable padding bytes. + let high = unsafe { + read_wide_bits(highs, patch_index * usize::from(high_width), high_width) + }; + set_decoded::( + &mut direct_values, + &mut transformed_values, + output_start + position - local_start, + base.wrapping_add(::from_u64(high)), + &mut transform, + ); + } + continue; + } + + let packed = packed_words_as_native::(&residual_words[residual_payload]); + if positions.is_empty() && DIRECT_OUTPUT && local_start == 0 && local_stop == BLOCK_LEN { + // SAFETY: The encoder writes one complete FastLanes chunk, and the output has capacity. + unsafe { + append_unpacked_add( + direct_values + .as_mut() + .vortex_expect("direct BlockResidual output is present"), + usize::from(residual_width), + packed, + base, + ); + } + continue; + } + // SAFETY: The encoder writes one complete FastLanes chunk for each block. + unsafe { + if positions.is_empty() && DIRECT_OUTPUT { + U::unpack_add(usize::from(residual_width), packed, base, &mut residuals); + } else { + U::unchecked_unpack(usize::from(residual_width), packed, &mut residuals); + } + } + if positions.is_empty() && DIRECT_OUTPUT { + append_decoded::( + &mut direct_values, + &mut transformed_values, + &residuals[local_start..local_stop], + &mut transform, + ); + continue; + } + let mut previous_position = None; + for (patch_index, &position) in positions.iter().enumerate() { + validate_patch_position(block_len, previous_position, position)?; + previous_position = Some(position); + // SAFETY: The payload includes fifteen readable padding bytes. + let high = + unsafe { read_wide_bits(highs, patch_index * usize::from(high_width), high_width) }; + residuals[usize::from(position)].apply_high(high, residual_width); + } + + append_residuals::( + &mut direct_values, + &mut transformed_values, + &mut residuals[local_start..local_stop], + base, + &mut transform, + ); + } + let validity = array.validity()?; + if DIRECT_OUTPUT { + Ok(PrimitiveArray::new( + direct_values + .vortex_expect("direct BlockResidual output is present") + .freeze(), + validity, + )) + } else { + Ok(PrimitiveArray::new( + transformed_values + .vortex_expect("transformed BlockResidual output is present") + .freeze(), + validity, + )) + } +} + +fn decoded_len( + direct: Option<&BufferMut>, + transformed: Option<&BufferMut>, +) -> usize { + direct.map_or_else(|| transformed.map_or(0, BufferMut::len), BufferMut::len) +} + +fn append_repeated_decoded( + direct: &mut Option>, + transformed: &mut Option>, + value: U, + count: usize, + transform: &mut impl FnMut(U) -> T, +) { + if DIRECT_OUTPUT { + append_repeated( + direct + .as_mut() + .vortex_expect("direct BlockResidual output is present"), + value, + count, + ); + } else { + let output = transformed + .as_mut() + .vortex_expect("transformed BlockResidual output is present"); + let output_len = output.len(); + for destination in &mut output.spare_capacity_mut()[..count] { + destination.write(transform(value)); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + count) }; + } +} + +fn append_decoded( + direct: &mut Option>, + transformed: &mut Option>, + values: &[U], + transform: &mut impl FnMut(U) -> T, +) { + if DIRECT_OUTPUT { + append_values( + direct + .as_mut() + .vortex_expect("direct BlockResidual output is present"), + values, + ); + } else { + let output = transformed + .as_mut() + .vortex_expect("transformed BlockResidual output is present"); + let output_len = output.len(); + for (destination, &value) in output.spare_capacity_mut()[..values.len()] + .iter_mut() + .zip(values) + { + destination.write(transform(value)); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + values.len()) }; + } +} + +fn append_residuals( + direct: &mut Option>, + transformed: &mut Option>, + residuals: &mut [U], + base: U, + transform: &mut impl FnMut(U) -> T, +) { + if DIRECT_OUTPUT { + for residual in residuals.iter_mut() { + *residual = residual.wrapping_add(base); + } + append_values( + direct + .as_mut() + .vortex_expect("direct BlockResidual output is present"), + residuals, + ); + } else { + let output = transformed + .as_mut() + .vortex_expect("transformed BlockResidual output is present"); + let output_len = output.len(); + for (destination, &residual) in output.spare_capacity_mut()[..residuals.len()] + .iter_mut() + .zip(residuals.iter()) + { + destination.write(transform(residual.wrapping_add(base))); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + residuals.len()) }; + } +} + +fn set_decoded( + direct: &mut Option>, + transformed: &mut Option>, + index: usize, + value: U, + transform: &mut impl FnMut(U) -> T, +) { + if DIRECT_OUTPUT { + direct + .as_mut() + .vortex_expect("direct BlockResidual output is present")[index] = value; + } else { + transformed + .as_mut() + .vortex_expect("transformed BlockResidual output is present")[index] = transform(value); + } +} + +fn append_repeated(output: &mut BufferMut, value: T, count: usize) { + let output_len = output.len(); + for destination in &mut output.spare_capacity_mut()[..count] { + destination.write(value); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + count) }; +} + +fn append_values(output: &mut BufferMut, values: &[T]) { + let output_len = output.len(); + for (destination, &value) in output.spare_capacity_mut()[..values.len()] + .iter_mut() + .zip(values) + { + destination.write(value); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + values.len()) }; +} + +unsafe fn append_unpacked_add( + output: &mut BufferMut, + bit_width: usize, + packed: &[T], + base: T, +) { + let output_len = output.len(); + let destination = output.spare_capacity_mut()[..BLOCK_LEN].as_mut_ptr(); + // SAFETY: The caller guarantees capacity. FastLanes initializes one complete output chunk. + let destination = unsafe { std::slice::from_raw_parts_mut(destination.cast::(), BLOCK_LEN) }; + // SAFETY: The caller provides one complete FastLanes input and output chunk. + unsafe { T::unpack_add(bit_width, packed, base, destination) }; + // SAFETY: FastLanes initialized each new output value. + unsafe { output.set_len(output_len + BLOCK_LEN) }; +} + +fn decompress_array( + array: ArrayView<'_, BlockResidual>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match array.dtype().as_ptype() { + PType::U8 => decode_array_values::(array, ctx, |value| value), + PType::U16 => decode_array_values::(array, ctx, |value| value), + PType::U32 => decode_array_values::(array, ctx, |value| value), + PType::U64 => decode_array_values::(array, ctx, |value| value), + PType::I8 => { + decode_array_values::(array, ctx, |value| (value ^ (1_u8 << 7)) as i8) + } + PType::I16 => decode_array_values::(array, ctx, |value| { + (value ^ (1_u16 << 15)) as i16 + }), + PType::I32 => decode_array_values::(array, ctx, |value| { + (value ^ (1_u32 << 31)) as i32 + }), + PType::I64 => decode_array_values::(array, ctx, |value| { + (value ^ (1_u64 << 63)) as i64 + }), + ptype => vortex_bail!("BlockResidual decode does not support {ptype}"), + } +} + +pub(crate) fn decompress_ordered_f32( + array: ArrayView<'_, BlockResidual>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + decode_array_values::(array, ctx, |ordered| { + let bits = if ordered & (1_u32 << 31) == 0 { + !ordered + } else { + ordered ^ (1_u32 << 31) + }; + f32::from_bits(bits) + }) +} + +pub(crate) fn decompress_ordered_f16( + array: ArrayView<'_, BlockResidual>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + decode_array_values::(array, ctx, |ordered| { + let bits = if ordered & (1_u16 << 15) == 0 { + !ordered + } else { + ordered ^ (1_u16 << 15) + }; + f16::from_bits(bits) + }) +} + +pub(crate) fn decompress_ordered_f64( + array: ArrayView<'_, BlockResidual>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + decode_array_values::(array, ctx, |ordered| { + let bits = if ordered & (1_u64 << 63) == 0 { + !ordered + } else { + ordered ^ (1_u64 << 63) + }; + f64::from_bits(bits) + }) +} + +fn scalar_from_array( + array: ArrayView<'_, BlockResidual>, + index: usize, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + let source_index = array.data().slice_start + index; + let block_index = source_index / BLOCK_LEN; + let index_in_block = source_index % BLOCK_LEN; + let residual_width = array.residual_widths()[block_index]; + let high_width = array.high_widths()[block_index]; + let logical_width = array.dtype().as_ptype().bit_width(); + vortex_ensure!( + usize::from(residual_width) <= logical_width + && usize::from(high_width) <= logical_width + && usize::from(residual_width) + usize::from(high_width) <= logical_width, + "block residual bit widths are invalid" + ); + let residual_words = array.residual_words(); + let residual_payload = payload_range( + array.residual_starts(), + block_index, + residual_words.len(), + "residual", + )?; + vortex_ensure!( + residual_payload.len() == BLOCK_LEN * usize::from(residual_width) / 64, + "block residual word count is invalid" + ); + let mut residual = match logical_width { + 8 => unpack_single_residual::( + residual_width, + &residual_words[residual_payload], + index_in_block, + ), + 16 => unpack_single_residual::( + residual_width, + &residual_words[residual_payload], + index_in_block, + ), + 32 => unpack_single_residual::( + residual_width, + &residual_words[residual_payload], + index_in_block, + ), + 64 => unpack_single_residual::( + residual_width, + &residual_words[residual_payload], + index_in_block, + ), + _ => vortex_bail!("block residual logical bit width is invalid"), + }; + + let positions = array.patch_positions(); + let patch_payload = payload_range(array.patch_starts(), block_index, positions.len(), "patch")?; + let block_positions = &positions[patch_payload]; + let highs = array.patch_highs(); + let high_payload = payload_range(array.high_starts(), block_index, highs.len(), "patch high")?; + validate_patch_header( + residual_width, + high_width, + block_positions.len(), + high_payload.len(), + )?; + if let Ok(patch_index) = block_positions.binary_search(&u16::try_from(index_in_block)?) { + let high_payload = &highs[high_payload]; + // SAFETY: The payload includes fifteen readable padding bytes. + let high = unsafe { + read_wide_bits( + high_payload, + patch_index * usize::from(high_width), + high_width, + ) + }; + residual |= high << residual_width; + } + Ok(array.bases()[block_index].wrapping_add(residual)) +} + +fn unpack_single_residual(width: u8, packed_words: &[u64], index: usize) -> u64 { + if width == 0 { + return 0; + } + let packed = packed_words_as_native::(packed_words); + // SAFETY: The encoder writes one complete FastLanes chunk for each block. + unsafe { T::unchecked_unpack_single(usize::from(width), packed, index).to_u64() } +} + +fn validate_patch_header( + residual_width: u8, + high_width: u8, + patch_count: usize, + high_payload_len: usize, +) -> VortexResult<()> { + vortex_ensure!( + patch_count == 0 || (high_width > 0 && residual_width < 64), + "block residual patches require nonzero high bits" + ); + let expected_high_len = if patch_count == 0 { + 0 + } else { + (patch_count * usize::from(high_width)).div_ceil(8) + 15 + }; + vortex_ensure!( + high_payload_len == expected_high_len, + "block residual patch high payload is invalid" + ); + Ok(()) +} + +#[inline(always)] +fn validate_patch_position( + block_len: usize, + previous_position: Option, + position: u16, +) -> VortexResult<()> { + if usize::from(position) < block_len + && previous_position.is_none_or(|previous| previous < position) + { + Ok(()) + } else { + invalid_patch_position() + } +} + +#[cold] +#[inline(never)] +fn invalid_patch_position() -> VortexResult<()> { + vortex_bail!("block residual patch positions are invalid") +} + +fn validate_offset_table( + starts: &[u32], + block_count: usize, + payload_len: usize, + name: &str, +) -> VortexResult<()> { + vortex_ensure!( + starts.len() == block_count + 1, + "block residual {name} offsets have an invalid length" + ); + vortex_ensure!( + starts.first() == Some(&0) + && starts.last().copied().map(usize::try_from).transpose()? == Some(payload_len), + "block residual {name} offsets do not cover the payload" + ); + vortex_ensure!( + starts.windows(2).all(|window| window[0] <= window[1]), + "block residual {name} offsets are not ordered" + ); + Ok(()) +} + +fn payload_range( + starts: &[u32], + block_index: usize, + payload_len: usize, + name: &str, +) -> VortexResult> { + let start = usize::try_from( + *starts + .get(block_index) + .ok_or_else(|| vortex_error::vortex_err!("block residual {name} start is missing"))?, + )?; + let stop = usize::try_from( + *starts + .get(block_index + 1) + .ok_or_else(|| vortex_error::vortex_err!("block residual {name} stop is missing"))?, + )?; + vortex_ensure!( + start <= stop && stop <= payload_len, + "block residual {name} offsets are invalid" + ); + Ok(start..stop) +} + +#[derive(Clone, Copy)] +struct BlockResidualMetadata { + unsliced_len: u64, + slice_start: u64, + residual_word_count: u64, + patch_count: u64, + patch_high_count: u64, +} + +impl BlockResidualMetadata { + fn from_data(data: &BlockResidualData) -> VortexResult { + Ok(Self { + unsliced_len: u64::try_from(data.unsliced_len)?, + slice_start: u64::try_from(data.slice_start)?, + residual_word_count: u64::try_from(data.residual_words.len())?, + patch_count: u64::try_from(data.patch_positions.len())?, + patch_high_count: u64::try_from(data.patch_highs.len())?, + }) + } + + fn encode(self) -> Vec { + let mut bytes = Vec::with_capacity(METADATA_LEN); + bytes.push(METADATA_VERSION); + for value in [ + self.unsliced_len, + self.slice_start, + self.residual_word_count, + self.patch_count, + self.patch_high_count, + ] { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } + + fn decode(bytes: &[u8]) -> VortexResult { + vortex_ensure!( + bytes.len() == METADATA_LEN, + "BlockResidualArray metadata requires {METADATA_LEN} bytes" + ); + vortex_ensure!( + bytes[0] == METADATA_VERSION, + "unsupported BlockResidualArray metadata version {}", + bytes[0] + ); + let read = |offset: usize| { + u64::from_le_bytes([ + bytes[offset], + bytes[offset + 1], + bytes[offset + 2], + bytes[offset + 3], + bytes[offset + 4], + bytes[offset + 5], + bytes[offset + 6], + bytes[offset + 7], + ]) + }; + Ok(Self { + unsliced_len: read(1), + slice_start: read(9), + residual_word_count: read(17), + patch_count: read(25), + patch_high_count: read(33), + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::dtype::NativePType; + use vortex_array::dtype::PType; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBufferMut; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::BlockResidual; + use super::BlockResidualArrayExt; + use crate::BlockResidualCodec; + + #[test] + fn roundtrip_and_scalar_access() -> VortexResult<()> { + let values = (0..4_099) + .map(|index| Ok(1_000_000_u64 + u64::try_from(index * index)?)) + .collect::>>()?; + let primitive = PrimitiveArray::from_iter(values.clone()); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(encoded.clone(), primitive.into_array(), &mut ctx); + for index in [0, 1, 1_023, 1_024, 4_098] { + let scalar = encoded.execute_scalar(index, &mut ctx)?; + assert_eq!( + scalar.as_primitive().typed_value::(), + Some(values[index]) + ); + } + Ok(()) + } + + #[test] + fn signed_roundtrip_and_scalar_access() -> VortexResult<()> { + let values = (0..2_050) + .map(|index| match index { + 0 => i64::MIN, + 1_023 => -1, + 1_024 => 0, + 2_049 => i64::MAX, + _ => (index as i64 - 1_025) * 1_000_003, + }) + .collect::>(); + let primitive = PrimitiveArray::from_iter(values.clone()); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + assert_arrays_eq!(encoded.clone(), primitive.into_array(), &mut ctx); + for index in [0, 1, 1_023, 1_024, 2_049] { + let scalar = encoded.execute_scalar(index, &mut ctx)?; + assert_eq!( + scalar.as_primitive().typed_value::(), + Some(values[index]) + ); + } + Ok(()) + } + + #[rstest] + #[case(vec![0_u8, 1, u8::MAX])] + #[case(vec![0_u16, 1, u16::MAX])] + #[case(vec![0_u32, 1, u32::MAX])] + #[case(vec![0_u64, 1, u64::MAX])] + #[case(vec![i8::MIN, -1, 0, 1, i8::MAX])] + #[case(vec![i16::MIN, -1, 0, 1, i16::MAX])] + #[case(vec![i32::MIN, -1, 0, 1, i32::MAX])] + #[case(vec![i64::MIN, -1, 0, 1, i64::MAX])] + fn integer_ptype_roundtrip(#[case] values: Vec) -> VortexResult<()> + where + T: NativePType + Copy, + { + let primitive = PrimitiveArray::from_iter(values); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + assert_arrays_eq!(encoded, primitive, &mut ctx); + Ok(()) + } + + #[test] + fn rejects_components_outside_logical_width() -> VortexResult<()> { + let mut parts = + BlockResidualCodec::encode_with_word_width(&[0_u64, 1, 2], 64)?.into_parts()?; + parts.bases[0] = u64::from(u8::MAX) + 1; + assert!(BlockResidual::try_new(parts, Validity::NonNullable, PType::U8).is_err()); + Ok(()) + } + + #[test] + fn rejects_invalid_component_offsets() -> VortexResult<()> { + let mut parts = + BlockResidualCodec::encode_with_word_width(&[0_u64, 1, 2], 64)?.into_parts()?; + parts.residual_starts[0] = 1; + assert!(BlockResidual::try_new(parts, Validity::NonNullable, PType::U64).is_err()); + Ok(()) + } + + #[test] + fn rejects_non_integer_input() { + let primitive = PrimitiveArray::from_iter([0.0_f32, 1.0, 2.0]); + assert!(BlockResidual::from_primitive(primitive.as_view()).is_err()); + } + + #[test] + fn u32_direct_decode_roundtrip() -> VortexResult<()> { + let primitive = PrimitiveArray::from_iter((0..2_050_u32).map(|index| { + let block = index / 1_024; + block * 1_000_000 + (index * 7_919) % 1_024 + })); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + assert_arrays_eq!(encoded, primitive.clone().into_array(), &mut ctx); + assert_arrays_eq!( + encoded.into_array().slice(1_023..1_026)?, + primitive.into_array().slice(1_023..1_026)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn nullable_slice_and_scalar_access() -> VortexResult<()> { + let values = (0..2_050) + .map(|index| Ok(u64::try_from(index * index)?)) + .collect::>>()?; + let validity = Validity::from_iter((0..values.len()).map(|index| index != 1_024)); + let primitive = PrimitiveArray::new(Buffer::from(values), validity); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + assert!(encoded.execute_scalar(1_024, &mut ctx)?.is_null()); + let sliced = encoded.into_array().slice(1_023..1_026)?; + let expected = primitive.into_array().slice(1_023..1_026)?; + assert_arrays_eq!(sliced, expected, &mut ctx); + Ok(()) + } + + #[test] + fn zero_width_patched_block_roundtrip() -> VortexResult<()> { + let mut values = vec![42_u32; 2_050]; + values[1_023] = u32::MAX; + let validity = Validity::from_iter((0..values.len()).map(|index| index != 1_024)); + let primitive = PrimitiveArray::new(Buffer::from(values.clone()), validity); + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + assert_eq!(encoded.residual_widths()[0], 0); + assert_eq!( + encoded + .execute_scalar(1_023, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(u32::MAX) + ); + assert!(encoded.execute_scalar(1_024, &mut ctx)?.is_null()); + assert_arrays_eq!( + encoded.into_array().slice(1_022..1_025)?, + primitive.into_array().slice(1_022..1_025)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn estimate_matches_materialized_size() -> VortexResult<()> { + let mut values = vec![42_u32; 2_050]; + values[1_023] = u32::MAX; + values[2_049] = u32::MAX - 1; + let validity = Validity::from_iter((0..values.len()).map(|index| index != 1_024)); + let primitive = PrimitiveArray::new(Buffer::from(values), validity); + let estimate = BlockResidual::estimate_primitive(primitive.as_view())?; + let encoded = BlockResidual::from_primitive(primitive.as_view())?; + + assert_eq!(estimate.nbytes(), encoded.nbytes()); + assert_eq!(estimate.patch_count(), encoded.patch_positions().len()); + Ok(()) + } + + #[test] + fn nullable_slice_serialization_roundtrip() -> VortexResult<()> { + let values = (0..2_050) + .map(|index| Ok(u64::try_from(index * index)?)) + .collect::>>()?; + let validity = Validity::from_iter((0..values.len()).map(|index| index != 1_024)); + let primitive = PrimitiveArray::new(Buffer::from(values), validity); + let sliced = BlockResidual::from_primitive(primitive.as_view())? + .into_array() + .slice(1_023..1_026)?; + let expected = primitive.into_array().slice(1_023..1_026)?; + let dtype = sliced.dtype().clone(); + let len = sliced.len(); + let array_context = ArrayContext::empty(); + let session = array_session(); + crate::initialize(&session); + let serialized = + sliced.serialize(&array_context, &session, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in serialized { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_context.to_ids()), + &session, + )?; + assert!(decoded.is::()); + assert_arrays_eq!(decoded, expected, &mut session.create_execution_ctx()); + Ok(()) + } + + #[test] + fn conformance() -> VortexResult<()> { + let mut values = vec![42_i16; 2_050]; + values[1_023] = i16::MAX; + let primitive = PrimitiveArray::new( + Buffer::from(values), + Validity::from_iter((0..2_050).map(|index| index != 1_024)), + ); + let array = BlockResidual::from_primitive(primitive.as_view())?.into_array(); + let session = array_session(); + crate::initialize(&session); + test_array_consistency(&array, &mut session.create_execution_ctx()); + Ok(()) + } +} diff --git a/encodings/block-residual/src/codec.rs b/encodings/block-residual/src/codec.rs new file mode 100644 index 00000000000..b99846b0776 --- /dev/null +++ b/encodings/block-residual/src/codec.rs @@ -0,0 +1,443 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use fastlanes::BitPacking; +use fastlanes::FoR as FastLanesFoR; +use vortex_error::VortexResult; + +const CHUNK_LEN: usize = 1024; +const HIGH_PADDING: usize = 15; +const PATCH_DECODE_PENALTY_BITS: usize = 16; +const SERIALIZED_BLOCK_METADATA_BYTES: usize = 12; + +/// Block-local residual codec for ordered unsigned latents. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct BlockResidualCodec { + len: usize, + blocks: Vec, +} + +/// Serialized children for the one-reference block residual codec. +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct BlockResidualParts { + pub(crate) len: usize, + pub(crate) bases: Vec, + pub(crate) residual_widths: Vec, + pub(crate) high_widths: Vec, + pub(crate) residual_starts: Vec, + pub(crate) patch_starts: Vec, + pub(crate) high_starts: Vec, + pub(crate) residual_words: Vec, + pub(crate) patch_positions: Vec, + pub(crate) patch_highs: Vec, +} + +pub(crate) struct BlockResidualCodecEstimate { + pub encoded_nbytes: usize, + pub patch_count: usize, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +struct BlockResidualBlock { + base: u64, + residual_width: u8, + high_width: u8, + residuals: Vec, + patch_positions: Vec, + patch_highs: Vec, +} + +impl BlockResidualCodec { + pub(crate) fn encode_with_word_width(values: &[u64], word_width: u8) -> VortexResult { + vortex_error::vortex_ensure!( + matches!(word_width, 8 | 16 | 32 | 64), + "block residual word width is invalid" + ); + let blocks = values + .chunks(CHUNK_LEN) + .map(|block| encode_block(block, word_width)) + .collect::>>()?; + Ok(Self { + len: values.len(), + blocks, + }) + } + + pub(crate) fn estimate_transformed( + values: &[T], + transform: impl Fn(T) -> u64 + Copy, + ) -> BlockResidualCodecEstimate { + let mut encoded_nbytes = 3 * size_of::(); + let mut total_patch_count = 0; + for values in values.chunks(CHUNK_LEN) { + let plan = estimate_block(values, transform); + let residual_nbytes = CHUNK_LEN * usize::from(plan.residual_width) / 8; + let patch_high_nbytes = if plan.patch_count == 0 { + 0 + } else { + (plan.patch_count * usize::from(plan.high_width)).div_ceil(8) + HIGH_PADDING + }; + encoded_nbytes += size_of::() + + 2 * size_of::() + + 3 * size_of::() + + residual_nbytes + + plan.patch_count * size_of::() + + patch_high_nbytes; + total_patch_count += plan.patch_count; + } + BlockResidualCodecEstimate { + encoded_nbytes, + patch_count: total_patch_count, + } + } + + pub(crate) fn into_parts(self) -> VortexResult { + let mut parts = BlockResidualParts { + len: self.len, + bases: Vec::with_capacity(self.blocks.len()), + residual_widths: Vec::with_capacity(self.blocks.len()), + high_widths: Vec::with_capacity(self.blocks.len()), + residual_starts: Vec::with_capacity(self.blocks.len() + 1), + patch_starts: Vec::with_capacity(self.blocks.len() + 1), + high_starts: Vec::with_capacity(self.blocks.len() + 1), + residual_words: Vec::new(), + patch_positions: Vec::new(), + patch_highs: Vec::new(), + }; + parts.residual_starts.push(0); + parts.patch_starts.push(0); + parts.high_starts.push(0); + for block in self.blocks { + parts.bases.push(block.base); + parts.residual_widths.push(block.residual_width); + parts.high_widths.push(block.high_width); + parts.residual_words.extend(block.residuals); + parts.patch_positions.extend(block.patch_positions); + parts.patch_highs.extend(block.patch_highs); + parts + .residual_starts + .push(u32::try_from(parts.residual_words.len())?); + parts + .patch_starts + .push(u32::try_from(parts.patch_positions.len())?); + parts + .high_starts + .push(u32::try_from(parts.patch_highs.len())?); + } + Ok(parts) + } +} +fn encode_block(values: &[u64], word_width: u8) -> VortexResult { + let base = values.iter().copied().min().unwrap_or(0); + let mut residuals = Vec::with_capacity(CHUNK_LEN); + let mut width_counts = [0usize; 65]; + let mut maximum_width = 0u8; + for &value in values { + let residual = value - base; + let width = bit_width(residual); + residuals.push(residual); + width_counts[usize::from(width)] += 1; + maximum_width = maximum_width.max(width); + } + residuals.resize(CHUNK_LEN, 0); + + let width_plan = choose_width(&width_counts, maximum_width, values.len()); + + materialize_block( + values, + BlockPlan { + base, + residual_width: width_plan.residual_width, + high_width: width_plan.high_width, + residuals, + patch_count: width_plan.patch_count, + }, + word_width, + ) +} + +struct BlockWidthPlan { + residual_width: u8, + high_width: u8, + patch_count: usize, +} + +fn estimate_block(values: &[T], transform: impl Fn(T) -> u64) -> BlockWidthPlan { + let base = values.iter().copied().map(&transform).min().unwrap_or(0); + let mut width_counts = [0usize; 65]; + let mut maximum_width = 0u8; + for &value in values { + let width = bit_width(transform(value) - base); + width_counts[usize::from(width)] += 1; + maximum_width = maximum_width.max(width); + } + choose_width(&width_counts, maximum_width, values.len()) +} + +fn choose_width( + width_counts: &[usize; 65], + maximum_width: u8, + value_count: usize, +) -> BlockWidthPlan { + let mut patch_count = value_count; + let mut best = (usize::MAX, maximum_width, 0u8, 0usize); + for residual_width in 0..=maximum_width { + patch_count -= width_counts[usize::from(residual_width)]; + let high_width = if patch_count == 0 { + 0 + } else { + maximum_width - residual_width + }; + let cost_bits = usize::from(residual_width) * CHUNK_LEN + + patch_count * (u16::BITS as usize + usize::from(high_width)) + + patch_count * PATCH_DECODE_PENALTY_BITS + + u64::BITS as usize + + SERIALIZED_BLOCK_METADATA_BYTES * 8 + + usize::from(patch_count > 0) * HIGH_PADDING * 8; + if cost_bits < best.0 { + best = (cost_bits, residual_width, high_width, patch_count); + } + } + BlockWidthPlan { + residual_width: best.1, + high_width: best.2, + patch_count: best.3, + } +} + +struct BlockPlan { + base: u64, + residual_width: u8, + high_width: u8, + residuals: Vec, + patch_count: usize, +} + +fn materialize_block( + values: &[u64], + plan: BlockPlan, + word_width: u8, +) -> VortexResult { + let residual_mask = low_mask(plan.residual_width); + let low_residuals = plan + .residuals + .iter() + .map(|&residual| residual & residual_mask) + .collect::>(); + let residuals = fast_pack(&low_residuals, plan.residual_width, word_width); + let mut patch_positions = Vec::with_capacity(plan.patch_count); + let mut patch_highs = BitWriter::with_capacity(plan.patch_count * 8); + if plan.high_width > 0 { + for (position, &residual) in plan.residuals[..values.len()].iter().enumerate() { + let high = residual >> plan.residual_width; + if high != 0 { + patch_positions.push(u16::try_from(position)?); + patch_highs.write(high, plan.high_width); + } + } + } + let patch_highs = if patch_positions.is_empty() { + Vec::new() + } else { + let mut encoded = patch_highs.finish(); + encoded.extend_from_slice(&[0; HIGH_PADDING]); + encoded + }; + + Ok(BlockResidualBlock { + base: plan.base, + residual_width: plan.residual_width, + high_width: plan.high_width, + residuals, + patch_positions, + patch_highs, + }) +} + +fn fast_pack(values: &[u64], width: u8, word_width: u8) -> Vec { + if width == 0 { + return Vec::new(); + } + match word_width { + 8 => fast_pack_native::(values, width), + 16 => fast_pack_native::(values, width), + 32 => fast_pack_native::(values, width), + 64 => fast_pack_native::(values, width), + _ => unreachable!("validated block residual word width"), + } +} + +fn fast_pack_native(values: &[u64], width: u8) -> Vec { + let mut packed_words = vec![0u64; CHUNK_LEN * usize::from(width) / u64::BITS as usize]; + let unpacked = values.iter().copied().map(T::from_u64).collect::>(); + let packed_native = packed_words_as_native_mut::(&mut packed_words); + // SAFETY: Both slices have the exact lengths required for one FastLanes chunk. + unsafe { T::unchecked_pack(usize::from(width), &unpacked, packed_native) }; + packed_words +} + +pub(crate) trait ResidualWord: BitPacking + FastLanesFoR + Copy + Default { + const BITS: u8; + + fn from_u64(value: u64) -> Self; + + fn to_u64(self) -> u64; + + fn wrapping_add(self, other: Self) -> Self; + + fn apply_high(&mut self, high: u64, shift: u8); + + unsafe fn unpack_add(bit_width: usize, packed: &[Self], base: Self, output: &mut [Self]); +} + +macro_rules! impl_residual_word { + ($T:ty, $bits:literal) => { + impl ResidualWord for $T { + const BITS: u8 = $bits; + + #[allow(clippy::cast_possible_truncation)] + fn from_u64(value: u64) -> Self { + value as $T + } + + fn to_u64(self) -> u64 { + u64::from(self) + } + + fn wrapping_add(self, other: Self) -> Self { + self.wrapping_add(other) + } + + #[allow(clippy::cast_possible_truncation)] + fn apply_high(&mut self, high: u64, shift: u8) { + *self |= (high as $T) << shift; + } + + unsafe fn unpack_add( + bit_width: usize, + packed: &[Self], + base: Self, + output: &mut [Self], + ) { + // SAFETY: The caller provides one complete FastLanes input and output chunk. + unsafe { FastLanesFoR::unchecked_unfor_pack(bit_width, packed, base, output) }; + } + } + }; +} + +impl_residual_word!(u8, 8); +impl_residual_word!(u16, 16); +impl_residual_word!(u32, 32); +impl_residual_word!(u64, 64); + +pub(crate) fn packed_words_as_native(words: &[u64]) -> &[T] { + // SAFETY: Unsigned integer types permit every bit pattern. A `u64` slice has sufficient + // alignment, and packed FastLanes payloads always contain a whole number of bytes. + let (prefix, native, suffix) = unsafe { words.align_to::() }; + debug_assert!(prefix.is_empty() && suffix.is_empty()); + native +} + +fn packed_words_as_native_mut(words: &mut [u64]) -> &mut [T] { + // SAFETY: Unsigned integer types permit every bit pattern. A `u64` slice has sufficient + // alignment, and packed FastLanes payloads always contain a whole number of bytes. + let (prefix, native, suffix) = unsafe { words.align_to_mut::() }; + debug_assert!(prefix.is_empty() && suffix.is_empty()); + native +} + +fn bit_width(value: u64) -> u8 { + u8::try_from(u64::BITS - value.leading_zeros()).unwrap_or(64) +} + +fn low_mask(bits: u8) -> u64 { + match bits { + 0 => 0, + 64 => u64::MAX, + _ => (1_u64 << bits) - 1, + } +} + +pub(crate) unsafe fn read_wide_bits(bytes: &[u8], bit_position: usize, width: u8) -> u64 { + let byte_position = bit_position / 8; + let bits_past_byte = bit_position % 8; + // SAFETY: The caller provides fifteen readable padding bytes. + let first = unsafe { read_u64_unaligned(bytes.as_ptr().add(byte_position)) }; + if width <= 57 { + (first >> bits_past_byte) & low_mask(width) + } else { + // SAFETY: The caller provides fifteen readable padding bytes. + let second = unsafe { read_u64_unaligned(bytes.as_ptr().add(byte_position + 7)) }; + let processed = 56 - bits_past_byte; + ((first >> bits_past_byte) | (second << processed)) & low_mask(width) + } +} + +unsafe fn read_u64_unaligned(pointer: *const u8) -> u64 { + // SAFETY: The caller provides eight readable bytes at the pointer. + u64::from_le(unsafe { pointer.cast::().read_unaligned() }) +} + +struct BitWriter { + bytes: Vec, + pending: u64, + pending_bits: u8, +} + +impl BitWriter { + fn with_capacity(capacity: usize) -> Self { + Self { + bytes: Vec::with_capacity(capacity), + pending: 0, + pending_bits: 0, + } + } + + fn write(&mut self, value: u64, width: u8) { + if width == 0 { + return; + } + let available = 64 - self.pending_bits; + self.pending |= value << self.pending_bits; + if width < available { + self.pending_bits += width; + return; + } + self.bytes.extend_from_slice(&self.pending.to_le_bytes()); + let remaining = width - available; + self.pending = if remaining == 0 { + 0 + } else { + value >> available + }; + self.pending_bits = remaining; + } + + fn finish(mut self) -> Vec { + if self.pending_bits > 0 { + let byte_count = usize::from(self.pending_bits).div_ceil(8); + self.bytes + .extend_from_slice(&self.pending.to_le_bytes()[..byte_count]); + } + self.bytes + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::BlockResidualCodec; + + #[test] + fn patch_penalty_prefers_dense_residuals() -> VortexResult<()> { + let mut values = vec![0_u64; 1_024]; + values[..307].fill(4_095); + let parts = BlockResidualCodec::encode_with_word_width(&values, 64)?.into_parts()?; + + assert_eq!(parts.residual_widths, [12]); + assert!(parts.patch_positions.is_empty()); + Ok(()) + } +} diff --git a/encodings/block-residual/src/lib.rs b/encodings/block-residual/src/lib.rs new file mode 100644 index 00000000000..b4298389790 --- /dev/null +++ b/encodings/block-residual/src/lib.rs @@ -0,0 +1,21 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Ordered-float and block-residual array encodings. + +mod block_residual_array; +mod codec; +mod ordered_float_array; + +pub use block_residual_array::*; +pub(crate) use codec::BlockResidualCodec; +pub(crate) use codec::BlockResidualParts; +pub use ordered_float_array::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Register the ordered-float and block-residual encodings in one session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(BlockResidual); + session.arrays().register(OrderedFloat); +} diff --git a/encodings/block-residual/src/ordered_float_array.rs b/encodings/block-residual/src/ordered_float_array.rs new file mode 100644 index 00000000000..3d293205f6c --- /dev/null +++ b/encodings/block-residual/src/ordered_float_array.rs @@ -0,0 +1,663 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; +use std::ops::Range; + +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::BlockResidual; +use crate::BlockResidualCodec; +use crate::BlockResidualEstimate; +use crate::block_residual_array::decompress_ordered_f16; +use crate::block_residual_array::decompress_ordered_f32; +use crate::block_residual_array::decompress_ordered_f64; +use crate::codec::BlockResidualCodecEstimate; + +/// IEEE floats mapped to unsigned integers that preserve numeric order. +pub type OrderedFloatArray = Array; + +#[array_slots(OrderedFloat)] +pub struct OrderedFloatSlots { + /// Ordered unsigned float bits. + #[slot(0)] + pub encoded: ArrayRef, +} + +#[derive(Clone, Debug, Default)] +pub struct OrderedFloatData; + +impl Display for OrderedFloatData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayHash for OrderedFloatData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl ArrayEq for OrderedFloatData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +#[derive(Clone, Debug)] +pub struct OrderedFloat; + +impl VTable for OrderedFloat { + type TypedArrayData = OrderedFloatData; + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.ordered_float"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let ptype = PType::try_from(dtype)?; + vortex_ensure!( + matches!(ptype, PType::F16 | PType::F32 | PType::F64), + "OrderedFloatArray requires f16, f32, or f64" + ); + let encoded = OrderedFloatSlotsView::from_slots(slots).encoded; + let expected = DType::Primitive(ordered_ptype(ptype)?, dtype.nullability()); + vortex_ensure!( + encoded.dtype() == &expected, + "OrderedFloatArray expected child dtype {expected}, got {}", + encoded.dtype() + ); + vortex_ensure!( + encoded.len() == len, + "OrderedFloatArray child length differs" + ); + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("OrderedFloatArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("OrderedFloatArray buffer_name {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(Vec::new())) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!(buffers.is_empty(), "OrderedFloatArray expects no buffers"); + vortex_ensure!( + metadata.is_empty(), + "OrderedFloatArray metadata must be empty" + ); + vortex_ensure!(children.len() == 1, "OrderedFloatArray requires one child"); + let ptype = PType::try_from(dtype)?; + let child_dtype = DType::Primitive(ordered_ptype(ptype)?, dtype.nullability()); + let encoded = children.get(0, &child_dtype, len)?; + Ok( + ArrayParts::new(self.clone(), dtype.clone(), len, OrderedFloatData) + .with_slots(OrderedFloatSlots { encoded }.into_slots()), + ) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + OrderedFloatSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let decoded = if let Some(block_residual) = array.encoded().as_typed::() { + match array.dtype().as_ptype() { + PType::F16 => decompress_ordered_f16(block_residual, ctx)?, + PType::F32 => decompress_ordered_f32(block_residual, ctx)?, + PType::F64 => decompress_ordered_f64(block_residual, ctx)?, + ptype => vortex_bail!("unsupported OrderedFloat ptype {ptype}"), + } + } else { + decode_primitive(array.as_view(), ctx)? + }; + Ok(ExecutionResult::done(decoded.into_array())) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for OrderedFloat { + fn scalar_at( + array: ArrayView<'_, OrderedFloat>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = array.encoded().execute_scalar(index, ctx)?; + if scalar.is_null() { + return Ok(Scalar::null(array.dtype().clone())); + } + Ok(match array.dtype().as_ptype() { + PType::F16 => Scalar::primitive( + f16::from_bits(unordered_u16( + scalar + .as_primitive() + .typed_value::() + .vortex_expect("validated ordered float scalar"), + )), + array.dtype().nullability(), + ), + PType::F32 => Scalar::primitive( + f32::from_bits(unordered_u32( + scalar + .as_primitive() + .typed_value::() + .vortex_expect("validated ordered float scalar"), + )), + array.dtype().nullability(), + ), + PType::F64 => Scalar::primitive( + f64::from_bits(unordered_u64( + scalar + .as_primitive() + .typed_value::() + .vortex_expect("validated ordered float scalar"), + )), + array.dtype().nullability(), + ), + ptype => vortex_panic!("unsupported OrderedFloat ptype {ptype}"), + }) + } +} + +impl ValidityChild for OrderedFloat { + fn validity_child(array: ArrayView<'_, OrderedFloat>) -> ArrayRef { + array.encoded().clone() + } +} + +impl SliceReduce for OrderedFloat { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + Ok(Some( + OrderedFloat::try_new(array.encoded().slice(range)?, array.dtype().as_ptype())? + .into_array(), + )) + } +} + +static RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(OrderedFloat))]); + +pub trait OrderedFloatArrayExt: TypedArrayRef + OrderedFloatArraySlotsExt {} + +impl> OrderedFloatArrayExt for T {} + +impl OrderedFloat { + /// Estimate BlockResidual bytes for ordered float bits without materialized integer values. + pub fn estimate_block_residual( + array: ArrayView<'_, Primitive>, + ) -> VortexResult { + let BlockResidualCodecEstimate { + encoded_nbytes, + patch_count, + } = match array.ptype() { + PType::F16 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + u64::from(ordered_u16(value.to_bits())) + }) + } + PType::F32 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + u64::from(ordered_u32(value.to_bits())) + }) + } + PType::F64 => { + BlockResidualCodec::estimate_transformed(array.as_slice::(), |value| { + ordered_u64(value.to_bits()) + }) + } + ptype => vortex_bail!("OrderedFloat requires f16, f32, or f64, got {ptype}"), + }; + let validity_nbytes = validity_to_child(&array.validity()?, array.len()) + .map(|validity| validity.nbytes()) + .unwrap_or(0); + BlockResidualEstimate::try_new(encoded_nbytes, validity_nbytes, patch_count) + } + + /// Construct an ordered float array from an unsigned child. + pub fn try_new(encoded: ArrayRef, float_ptype: PType) -> VortexResult { + vortex_ensure!( + matches!(float_ptype, PType::F16 | PType::F32 | PType::F64), + "OrderedFloat requires f16, f32, or f64" + ); + let dtype = DType::Primitive(float_ptype, encoded.dtype().nullability()); + let len = encoded.len(); + Array::try_from_parts( + ArrayParts::new(OrderedFloat, dtype, len, OrderedFloatData) + .with_slots(OrderedFloatSlots { encoded }.into_slots()), + ) + } + + /// Map canonical floats to ordered unsigned integer bits. + pub fn from_primitive(array: ArrayView<'_, Primitive>) -> VortexResult { + let validity = array.validity()?; + match array.ptype() { + PType::F16 => Self::try_new( + PrimitiveArray::new( + Buffer::from( + array + .as_slice::() + .iter() + .map(|value| ordered_u16(value.to_bits())) + .collect::>(), + ), + validity, + ) + .into_array(), + PType::F16, + ), + PType::F32 => Self::try_new( + PrimitiveArray::new( + Buffer::from( + array + .as_slice::() + .iter() + .map(|value| ordered_u32(value.to_bits())) + .collect::>(), + ), + validity, + ) + .into_array(), + PType::F32, + ), + PType::F64 => Self::try_new( + PrimitiveArray::new( + Buffer::from( + array + .as_slice::() + .iter() + .map(|value| ordered_u64(value.to_bits())) + .collect::>(), + ), + validity, + ) + .into_array(), + PType::F64, + ), + ptype => vortex_bail!("OrderedFloat requires f16, f32, or f64, got {ptype}"), + } + } +} + +fn decode_primitive( + array: ArrayView<'_, OrderedFloat>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let encoded = array.encoded().clone().execute::(ctx)?; + let validity = encoded.validity()?; + Ok(match array.dtype().as_ptype() { + PType::F16 => PrimitiveArray::new( + encoded + .into_buffer::() + .map_each_in_place(|value| f16::from_bits(unordered_u16(value))) + .freeze(), + validity, + ), + PType::F32 => PrimitiveArray::new( + encoded + .into_buffer::() + .map_each_in_place(|value| f32::from_bits(unordered_u32(value))) + .freeze(), + validity, + ), + PType::F64 => PrimitiveArray::new( + encoded + .into_buffer::() + .map_each_in_place(|value| f64::from_bits(unordered_u64(value))) + .freeze(), + validity, + ), + ptype => vortex_panic!("unsupported OrderedFloat ptype {ptype}"), + }) +} + +fn ordered_ptype(ptype: PType) -> VortexResult { + match ptype { + PType::F16 => Ok(PType::U16), + PType::F32 => Ok(PType::U32), + PType::F64 => Ok(PType::U64), + _ => vortex_bail!("OrderedFloat requires f16, f32, or f64, got {ptype}"), + } +} + +fn ordered_u16(bits: u16) -> u16 { + if bits & (1_u16 << 15) == 0 { + bits ^ (1_u16 << 15) + } else { + !bits + } +} + +fn unordered_u16(value: u16) -> u16 { + if value & (1_u16 << 15) == 0 { + !value + } else { + value ^ (1_u16 << 15) + } +} + +fn ordered_u32(bits: u32) -> u32 { + if bits & (1_u32 << 31) == 0 { + bits ^ (1_u32 << 31) + } else { + !bits + } +} + +fn unordered_u32(value: u32) -> u32 { + if value & (1_u32 << 31) == 0 { + !value + } else { + value ^ (1_u32 << 31) + } +} + +fn ordered_u64(bits: u64) -> u64 { + if bits & (1_u64 << 63) == 0 { + bits ^ (1_u64 << 63) + } else { + !bits + } +} + +fn unordered_u64(value: u64) -> u64 { + if value & (1_u64 << 63) == 0 { + !value + } else { + value ^ (1_u64 << 63) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::dtype::PType; + use vortex_array::dtype::half::f16; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBufferMut; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + + use super::OrderedFloat; + use super::OrderedFloatArraySlotsExt; + use crate::BlockResidual; + use crate::BlockResidualArrayExt; + + #[test] + fn roundtrip_f16_special_values() -> VortexResult<()> { + let values = [ + f16::from_bits(0xfc00), + f16::from_f32(-1.0), + f16::NEG_ZERO, + f16::ZERO, + f16::from_f32(1.0), + f16::INFINITY, + f16::from_bits(0x7e42), + ]; + let primitive = PrimitiveArray::from_iter(values); + let estimate = OrderedFloat::estimate_block_residual(primitive.as_view())?; + let ordered = OrderedFloat::from_primitive(primitive.as_view())?; + let residuals = BlockResidual::from_primitive(ordered.encoded().as_::())?; + assert_eq!(estimate.nbytes(), residuals.nbytes()); + assert_eq!(estimate.patch_count(), residuals.patch_positions().len()); + let encoded = OrderedFloat::try_new(residuals.into_array(), PType::F16)?; + let session = array_session(); + crate::initialize(&session); + let decoded = encoded + .into_array() + .execute::(&mut session.create_execution_ctx())?; + + assert_eq!( + decoded + .as_slice::() + .iter() + .map(|value| value.to_bits()) + .collect::>(), + values.map(f16::to_bits) + ); + Ok(()) + } + + #[test] + fn roundtrip_special_values() -> VortexResult<()> { + let primitive = PrimitiveArray::from_iter([ + f64::NEG_INFINITY, + -1.0, + -0.0, + 0.0, + 1.0, + f64::INFINITY, + f64::from_bits(0x7ff8_0000_0000_0042), + ]); + let estimate = OrderedFloat::estimate_block_residual(primitive.as_view())?; + let encoded = OrderedFloat::from_primitive(primitive.as_view())?; + let residuals = BlockResidual::from_primitive(encoded.encoded().as_::())?; + assert_eq!(estimate.nbytes(), residuals.nbytes()); + assert_eq!(estimate.patch_count(), residuals.patch_positions().len()); + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(encoded, primitive.into_array(), &mut ctx); + Ok(()) + } + + #[test] + fn roundtrip_f32_special_values() -> VortexResult<()> { + let values = [ + f32::NEG_INFINITY, + -1.0, + -0.0, + 0.0, + 1.0, + f32::INFINITY, + f32::from_bits(0x7fc0_0042), + ]; + let primitive = PrimitiveArray::from_iter(values); + let encoded = OrderedFloat::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let decoded = encoded + .into_array() + .execute::(&mut session.create_execution_ctx())?; + + assert_eq!( + decoded + .as_slice::() + .iter() + .map(|value| value.to_bits()) + .collect::>(), + values.map(f32::to_bits) + ); + Ok(()) + } + + #[test] + fn rejects_invalid_logical_and_child_types() { + let u32_child = PrimitiveArray::from_iter([0_u32, 1, 2]).into_array(); + assert!(OrderedFloat::try_new(u32_child.clone(), PType::F64).is_err()); + assert!(OrderedFloat::try_new(u32_child, PType::U32).is_err()); + + let u64_child = PrimitiveArray::from_iter([0_u64, 1, 2]).into_array(); + assert!(OrderedFloat::try_new(u64_child, PType::F32).is_err()); + } + + #[test] + fn roundtrip_f32_block_residual() -> VortexResult<()> { + let values = (0..2_050) + .scan(1_000.0_f32, |value, index| { + *value += ((index * 7_919 % 101) as f32 - 50.0) * 0.0001; + Some(*value) + }) + .collect::>(); + let primitive = PrimitiveArray::from_iter(values.clone()); + let estimate = OrderedFloat::estimate_block_residual(primitive.as_view())?; + let ordered = OrderedFloat::from_primitive(primitive.as_view())?; + let residuals = BlockResidual::from_primitive(ordered.encoded().as_::())?; + assert_eq!(estimate.nbytes(), residuals.nbytes()); + assert_eq!(estimate.patch_count(), residuals.patch_positions().len()); + let encoded = OrderedFloat::try_new(residuals.into_array(), PType::F32)?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + assert_arrays_eq!(encoded, primitive, &mut ctx); + assert_eq!( + encoded + .execute_scalar(1_024, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(values[1_024]) + ); + assert_arrays_eq!( + encoded.into_array().slice(1_023..1_026)?, + primitive.into_array().slice(1_023..1_026)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn nullable_serialized_slice_roundtrip() -> VortexResult<()> { + let primitive = PrimitiveArray::new( + Buffer::from(vec![f64::NEG_INFINITY, -0.0, 0.0, 42.25, f64::INFINITY]), + Validity::from_iter([true, false, true, true, true]), + ); + let encoded = OrderedFloat::from_primitive(primitive.as_view())?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + assert!(encoded.execute_scalar(1, &mut ctx)?.is_null()); + + let sliced = encoded.into_array().slice(1..5)?; + let expected = primitive.into_array().slice(1..5)?; + let dtype = sliced.dtype().clone(); + let len = sliced.len(); + let array_context = ArrayContext::empty(); + let serialized = + sliced.serialize(&array_context, &session, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in serialized { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_context.to_ids()), + &session, + )?; + + assert!(decoded.is::()); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn conformance() -> VortexResult<()> { + let primitive = PrimitiveArray::new( + Buffer::from(vec![f32::NEG_INFINITY, -0.0, 0.0, 42.25, f32::INFINITY]), + Validity::from_iter([true, false, true, true, true]), + ); + let direct = OrderedFloat::from_primitive(primitive.as_view())?; + let residuals = BlockResidual::from_primitive(direct.encoded().as_::())?; + let fused = OrderedFloat::try_new(residuals.into_array(), PType::F32)?; + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + for array in [direct.into_array(), fused.into_array()] { + test_array_consistency(&array, &mut ctx); + } + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index a393db6ecc8..cb85537f303 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -194,6 +194,133 @@ pub fn bitpack_primitive(array: &[T], bit_width: u8 output.freeze() } +/// Maps and bit-packs primitive values without a full intermediate buffer. +/// +/// The mapped values must fit in `bit_width` bits. This function does not create patches. +pub fn bitpack_primitive_map(array: &[S], bit_width: u8, mut map: F) -> Buffer +where + T: NativePType + BitPacking, + F: FnMut(&S) -> T, +{ + if bit_width == 0 { + return Buffer::::empty(); + } + + let bit_width = usize::from(bit_width); + let num_chunks = array.len().div_ceil(1024); + let num_full_chunks = array.len() / 1024; + let packed_len = 128 * bit_width / size_of::(); + let mut output = BufferMut::::with_capacity(num_chunks * packed_len); + let mut mapped = [T::zero(); 1024]; + + for chunk_index in 0..num_full_chunks { + let start = chunk_index * 1024; + mapped + .iter_mut() + .zip(&array[start..start + 1024]) + .for_each(|(output, input)| *output = map(input)); + let output_len = output.len(); + // SAFETY: The output has capacity for one packed vector and both slices have exact sizes. + unsafe { + output.set_len(output_len + packed_len); + BitPacking::unchecked_pack(bit_width, &mapped, &mut output[output_len..][..packed_len]); + } + } + + if num_chunks != num_full_chunks { + let start = num_full_chunks * 1024; + let last_chunk_len = array.len() - start; + mapped[..last_chunk_len] + .iter_mut() + .zip(&array[start..]) + .for_each(|(output, input)| *output = map(input)); + mapped[last_chunk_len..].fill(T::zero()); + + let output_len = output.len(); + // SAFETY: The output has capacity for one packed vector and both slices have exact sizes. + unsafe { + output.set_len(output_len + packed_len); + BitPacking::unchecked_pack(bit_width, &mapped, &mut output[output_len..][..packed_len]); + } + } + + output.freeze() +} + +/// Maps each primitive value to two values and bit-packs both outputs. +/// +/// Each mapped value must fit in its corresponding bit width. This function does not create +/// patches. +pub fn bitpack_primitive_map_pair( + array: &[S], + left_bit_width: u8, + right_bit_width: u8, + mut map: F, +) -> (Buffer, Buffer) +where + T: NativePType + BitPacking, + F: FnMut(&S) -> (T, T), +{ + if left_bit_width == 0 && right_bit_width == 0 { + return (Buffer::empty(), Buffer::empty()); + } + + let left_bit_width = usize::from(left_bit_width); + let right_bit_width = usize::from(right_bit_width); + let num_chunks = array.len().div_ceil(1024); + let num_full_chunks = array.len() / 1024; + let left_packed_len = 128 * left_bit_width / size_of::(); + let right_packed_len = 128 * right_bit_width / size_of::(); + let mut left_output = BufferMut::::with_capacity(num_chunks * left_packed_len); + let mut right_output = BufferMut::::with_capacity(num_chunks * right_packed_len); + let mut left_mapped = [T::zero(); 1024]; + let mut right_mapped = [T::zero(); 1024]; + + for chunk_index in 0..num_full_chunks { + let start = chunk_index * 1024; + left_mapped + .iter_mut() + .zip(&mut right_mapped) + .zip(&array[start..start + 1024]) + .for_each(|((left, right), input)| (*left, *right) = map(input)); + append_packed_chunk(&mut left_output, &left_mapped, left_bit_width); + append_packed_chunk(&mut right_output, &right_mapped, right_bit_width); + } + + if num_chunks != num_full_chunks { + let start = num_full_chunks * 1024; + let last_chunk_len = array.len() - start; + left_mapped[..last_chunk_len] + .iter_mut() + .zip(&mut right_mapped[..last_chunk_len]) + .zip(&array[start..]) + .for_each(|((left, right), input)| (*left, *right) = map(input)); + left_mapped[last_chunk_len..].fill(T::zero()); + right_mapped[last_chunk_len..].fill(T::zero()); + append_packed_chunk(&mut left_output, &left_mapped, left_bit_width); + append_packed_chunk(&mut right_output, &right_mapped, right_bit_width); + } + + (left_output.freeze(), right_output.freeze()) +} + +fn append_packed_chunk( + output: &mut BufferMut, + values: &[T; 1024], + bit_width: usize, +) { + if bit_width == 0 { + return; + } + let packed_len = 128 * bit_width / size_of::(); + let output_len = output.len(); + // SAFETY: The output has capacity for one packed vector and both slices have exact sizes. + unsafe { + output.set_len(output_len + packed_len); + BitPacking::unchecked_pack(bit_width, values, &mut output[output_len..][..packed_len]); + } +} + pub fn gather_patches( parray: &PrimitiveArray, bit_width: u8, @@ -464,6 +591,41 @@ mod test { ); } + #[test] + fn bitpack_primitive_map_matches_materialized_input() { + for len in [0, 1, 1023, 1024, 1025, 4097] { + let input = (0..len).map(|value| value as u32).collect::>(); + let mapped = input + .iter() + .map(|value| value.wrapping_mul(31) & 0x3ff) + .collect::>(); + assert_eq!( + bitpack_primitive_map(&input, 10, |value| value.wrapping_mul(31) & 0x3ff), + bitpack_primitive(&mapped, 10), + ); + } + } + + #[test] + fn bitpack_primitive_map_pair_matches_materialized_inputs() { + for len in [0, 1, 1023, 1024, 1025, 4097] { + let input = (0..len).map(|value| value as u32).collect::>(); + let left = input + .iter() + .map(|value| value.wrapping_mul(31) & 0x3ff) + .collect::>(); + let right = input + .iter() + .map(|value| value.wrapping_mul(7) & 0x7) + .collect::>(); + let actual = bitpack_primitive_map_pair(&input, 10, 3, |value| { + (value.wrapping_mul(31) & 0x3ff, value.wrapping_mul(7) & 0x7) + }); + assert_eq!(actual.0, bitpack_primitive(&left, 10)); + assert_eq!(actual.1, bitpack_primitive(&right, 3)); + } + } + #[test] fn null_patches() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 692e7dcdd7f..60141cde5e3 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -17,8 +17,11 @@ use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::Patches; use vortex_array::scalar::Scalar; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::BitPacked; use crate::BitPackedArrayExt; @@ -45,6 +48,114 @@ pub fn unpack_primitive_array( Ok(builder.finish_into_primitive()) } +/// Unpacks one bit-packed array and maps each value into the output primitive array. +pub fn unpack_map( + array: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, + map: M, +) -> VortexResult +where + F: BitPackedUnpack, + T: NativePType, + M: Fn(F) -> T, +{ + let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); + unpack_map_into_builder::(array, &mut builder, ctx, map)?; + Ok(builder.finish_into_primitive()) +} + +/// Unpacks two aligned bit-packed arrays and maps each value pair into one output buffer. +/// +/// This path requires equal lengths and offsets. Neither input can contain patches. +pub fn unpack_pair_map( + left: ArrayView<'_, BitPacked>, + right: ArrayView<'_, BitPacked>, + mut map: F, +) -> VortexResult> +where + T: BitPackedUnpack + BitPacking, + U: NativePType, + F: FnMut(T, T) -> U, +{ + vortex_ensure!(left.len() == right.len(), "bit-packed pair length differs"); + vortex_ensure!( + left.offset() == right.offset(), + "bit-packed pair offset differs" + ); + vortex_ensure!( + left.dtype().as_ptype() == T::PTYPE && right.dtype().as_ptype() == T::PTYPE, + "bit-packed pair requires {} inputs", + T::PTYPE + ); + vortex_ensure!( + left.patches().is_none() && right.patches().is_none(), + "bit-packed pair mapping does not support patches" + ); + if left.is_empty() { + return Ok(Buffer::empty()); + } + + let len = left.len(); + let offset = usize::from(left.offset()); + let num_chunks = (offset + len).div_ceil(1024); + let left_width = usize::from(left.bit_width()); + let right_width = usize::from(right.bit_width()); + let left_packed_len = 128 * left_width / size_of::(); + let right_packed_len = 128 * right_width / size_of::(); + let left_packed = left.packed_slice::(); + let right_packed = right.packed_slice::(); + let mut left_values = [T::default(); 1024]; + let mut right_values = [T::default(); 1024]; + let mut output = BufferMut::with_capacity(len); + + for chunk_index in 0..num_chunks { + if left_width == 0 { + left_values.fill(T::default()); + } else { + let start = chunk_index * left_packed_len; + // SAFETY: BitPacked validation guarantees one packed FastLanes block here. + unsafe { + BitPacking::unchecked_unpack( + left_width, + &left_packed[start..start + left_packed_len], + &mut left_values, + ); + } + } + if right_width == 0 { + right_values.fill(T::default()); + } else { + let start = chunk_index * right_packed_len; + // SAFETY: BitPacked validation guarantees one packed FastLanes block here. + unsafe { + BitPacking::unchecked_unpack( + right_width, + &right_packed[start..start + right_packed_len], + &mut right_values, + ); + } + } + + let start = if chunk_index == 0 { offset } else { 0 }; + let logical_start = chunk_index * 1024; + let end = (offset + len - logical_start).min(1024); + let output_len = output.len(); + let mapped_len = end - start; + for (output, (&left, &right)) in output.spare_capacity_mut()[..mapped_len].iter_mut().zip( + left_values[start..end] + .iter() + .zip(&right_values[start..end]), + ) { + output.write(map(left, right)); + } + // SAFETY: The loop initialized each new output value. + unsafe { output.set_len(output_len + mapped_len) }; + } + + debug_assert_eq!(output.len(), len); + Ok(output.freeze()) +} + /// Unpack a bit-packed array directly into a same-typed `PrimitiveBuilder`. /// /// This is the fast path for ordinary decompression: full FastLanes chunks are unpacked straight @@ -203,6 +314,7 @@ pub fn count_exceptions(bit_width: u8, bit_width_freq: &[usize]) -> usize { #[cfg(test)] mod tests { + use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; @@ -228,6 +340,65 @@ mod tests { use crate::BitPackedData; use crate::bitpack_compress::bitpack_encode; + #[test] + fn unpack_map_preserves_slice_and_validity() -> VortexResult<()> { + let values = (0_u32..3073) + .map(|value| (value % 17 != 0).then_some(value & 0x3ff)) + .collect::>(); + let input = PrimitiveArray::from_option_iter(values.iter().copied()); + let mut ctx = SESSION.create_execution_ctx(); + let packed = bitpack_encode(&input, 10, None, &mut ctx)?.into_array(); + + for range in [0..3073, 3..2051, 1023..2050] { + let packed = packed.slice(range.clone())?; + let actual = unpack_map::(packed.as_::(), &mut ctx, |value| { + u64::from(value) * 17 + })?; + let expected = PrimitiveArray::from_option_iter( + values[range] + .iter() + .map(|value| value.map(|value| u64::from(value) * 17)), + ); + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) + } + + #[test] + fn unpack_pair_map_matches_sliced_inputs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let left_values = PrimitiveArray::from_iter((0_u32..3073).map(|value| value & 0x3ff)); + let right_values = + PrimitiveArray::from_iter((0_u32..3073).map(|value| value.wrapping_mul(7) & 0x7f)); + let left = bitpack_encode(&left_values, 10, None, &mut ctx)?.into_array(); + let right = bitpack_encode(&right_values, 7, None, &mut ctx)?.into_array(); + + for range in [ + Range { + start: 0, + end: 3073, + }, + 3..2051, + 1023..2050, + ] { + let left = left.slice(range.clone())?; + let right = right.slice(range.clone())?; + let actual = unpack_pair_map::( + left.as_::(), + right.as_::(), + |left, right| u64::from(left) * 17 + u64::from(right), + )?; + let expected = left_values.as_slice::()[range.clone()] + .iter() + .copied() + .zip(right_values.as_slice::()[range].iter().copied()) + .map(|(left, right)| u64::from(left) * 17 + u64::from(right)) + .collect::>(); + assert_eq!(actual, expected); + } + Ok(()) + } + fn encode(array: &PrimitiveArray, bit_width: u8) -> BitPackedArray { bitpack_encode(array, bit_width, None, &mut SESSION.create_execution_ctx()).unwrap() } diff --git a/encodings/float-quant/Cargo.toml b/encodings/float-quant/Cargo.toml new file mode 100644 index 00000000000..c3592f20840 --- /dev/null +++ b/encodings/float-quant/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "vortex-float-quant" +authors = { workspace = true } +categories = { workspace = true } +description = "Vortex lossless float transform array encodings" +edition = { workspace = true } +homepage = { workspace = true } +include = { workspace = true } +keywords = { workspace = true } +license = { workspace = true } +readme = { workspace = true } +repository = { workspace = true } +rust-version = { workspace = true } +version = { workspace = true } + +[dependencies] +vortex-array = { workspace = true } +vortex-buffer = { workspace = true } +vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } +vortex-session = { workspace = true } +vortex-utils = { workspace = true } + +[dev-dependencies] +vortex-array = { workspace = true, features = ["_test-harness"] } + +[lints] +workspace = true diff --git a/encodings/float-quant/src/array.rs b/encodings/float-quant/src/array.rs new file mode 100644 index 00000000000..c25c49c28f4 --- /dev/null +++ b/encodings/float-quant/src/array.rs @@ -0,0 +1,1429 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; + +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability::NonNullable; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; +use vortex_array::scalar::Scalar; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FoR; +use vortex_fastlanes::FoRArrayExt; +use vortex_fastlanes::FoRArraySlotsExt; +use vortex_fastlanes::bitpack_decompress::unpack_map; +use vortex_fastlanes::bitpack_decompress::unpack_pair_map; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::rules::RULES; + +const METADATA_VERSION: u8 = 1; +const METADATA_LEN: usize = 2; + +/// A lossless float split with quantized high bits and low-bit adjustments. +pub type FloatQuantArray = Array; + +#[array_slots(FloatQuant)] +pub struct FloatQuantSlots { + /// Ordered float bits after the low `k` bits are removed. + #[slot(0)] + pub primary: ArrayRef, + /// Sign-normalized low `k` bits. + #[slot(1)] + pub secondary: Option, +} + +#[derive(Clone, Debug)] +pub struct FloatQuantData { + pub(crate) k: u8, +} + +impl Display for FloatQuantData { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "k: {}", self.k) + } +} + +impl ArrayHash for FloatQuantData { + fn array_hash(&self, state: &mut H, _accuracy: EqMode) { + self.k.hash(state); + } +} + +impl ArrayEq for FloatQuantData { + fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool { + self.k == other.k + } +} + +#[derive(Clone, Debug)] +pub struct FloatQuant; + +impl VTable for FloatQuant { + type TypedArrayData = FloatQuantData; + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.float_quant"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let ptype = PType::try_from(dtype)?; + let latent_ptype = latent_ptype(ptype)?; + let precision_bits = precision_bits(ptype)?; + vortex_ensure!( + data.k > 0 && data.k <= precision_bits, + "FloatQuant k {} exceeds {ptype} precision {precision_bits}", + data.k + ); + + let slots = FloatQuantSlotsView::from_slots(slots); + let expected_primary = DType::Primitive(latent_ptype, dtype.nullability()); + vortex_ensure!( + slots.primary.dtype() == &expected_primary, + "expected primary dtype {expected_primary}, got {}", + slots.primary.dtype() + ); + vortex_ensure!( + slots.primary.len() == len, + "FloatQuant primary length differs" + ); + if let Some(secondary) = slots.secondary { + let expected_secondary = DType::Primitive(latent_ptype, NonNullable); + vortex_ensure!( + secondary.dtype() == &expected_secondary, + "expected secondary dtype {expected_secondary}, got {}", + secondary.dtype() + ); + vortex_ensure!( + secondary.len() == len, + "FloatQuant secondary length differs" + ); + } + Ok(()) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("FloatQuantArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("FloatQuantArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some(vec![METADATA_VERSION, array.data().k])) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_ensure!(buffers.is_empty(), "FloatQuant expects no buffers"); + vortex_ensure!( + metadata.len() == METADATA_LEN, + "FloatQuant metadata requires {METADATA_LEN} bytes" + ); + vortex_ensure!( + metadata[0] == METADATA_VERSION, + "unsupported FloatQuant metadata version {}", + metadata[0] + ); + vortex_ensure!( + matches!(children.len(), 1 | 2), + "FloatQuant requires one or two children" + ); + + let ptype = PType::try_from(dtype)?; + let latent_ptype = latent_ptype(ptype)?; + let primary_dtype = DType::Primitive(latent_ptype, dtype.nullability()); + let primary = children.get(0, &primary_dtype, len)?; + let secondary = if children.len() == 2 { + let secondary_dtype = DType::Primitive(latent_ptype, NonNullable); + Some(children.get(1, &secondary_dtype, len)?) + } else { + None + }; + let slots = FloatQuantSlots { primary, secondary }.into_slots(); + Ok(ArrayParts::new( + self.clone(), + dtype.clone(), + len, + FloatQuantData { k: metadata[1] }, + ) + .with_slots(slots)) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + FloatQuantSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ExecutionResult::done( + decode(array.as_view(), ctx)?.into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl OperationsVTable for FloatQuant { + fn scalar_at( + array: ArrayView<'_, FloatQuant>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primary = array.primary().execute_scalar(index, ctx)?; + if primary.is_null() { + return Ok(Scalar::null(array.dtype().clone())); + } + let k = array.data().k; + Ok(match PType::try_from(array.dtype())? { + PType::F16 => Scalar::primitive( + join_f16( + primary + .as_primitive() + .typed_value::() + .vortex_expect("validated primary scalar"), + array + .secondary() + .map(|secondary| { + secondary + .execute_scalar(index, ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| { + vortex_error::vortex_err!("validated secondary scalar is null") + }) + }) + .transpose()? + .unwrap_or(0), + k, + ), + array.dtype().nullability(), + ), + PType::F32 => Scalar::primitive( + join_f32( + primary + .as_primitive() + .typed_value::() + .vortex_expect("validated primary scalar"), + array + .secondary() + .map(|secondary| { + secondary + .execute_scalar(index, ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| { + vortex_error::vortex_err!("validated secondary scalar is null") + }) + }) + .transpose()? + .unwrap_or(0), + k, + ), + array.dtype().nullability(), + ), + PType::F64 => Scalar::primitive( + join_f64( + primary + .as_primitive() + .typed_value::() + .vortex_expect("validated primary scalar"), + array + .secondary() + .map(|secondary| { + secondary + .execute_scalar(index, ctx)? + .as_primitive() + .typed_value::() + .ok_or_else(|| { + vortex_error::vortex_err!("validated secondary scalar is null") + }) + }) + .transpose()? + .unwrap_or(0), + k, + ), + array.dtype().nullability(), + ), + ptype => vortex_panic!("unsupported FloatQuant ptype {ptype}"), + }) + } +} + +impl ValidityChild for FloatQuant { + fn validity_child(array: ArrayView<'_, FloatQuant>) -> ArrayRef { + array.primary().clone() + } +} + +pub trait FloatQuantArrayExt: TypedArrayRef + FloatQuantArraySlotsExt { + /// Return the number of split low bits. + fn k(&self) -> u8 { + self.deref().k + } +} + +impl> FloatQuantArrayExt for T {} + +impl FloatQuant { + /// Construct a float quantization array from one or two latent children. + pub fn try_new( + primary: ArrayRef, + secondary: Option, + float_ptype: PType, + k: u8, + ) -> VortexResult { + let dtype = DType::Primitive(float_ptype, primary.dtype().nullability()); + let len = primary.len(); + let slots = FloatQuantSlots { primary, secondary }.into_slots(); + Array::try_from_parts( + ArrayParts::new(FloatQuant, dtype, len, FloatQuantData { k }).with_slots(slots), + ) + } + + /// Split a canonical float array into two unsigned latent children. + pub fn from_primitive(array: ArrayView<'_, Primitive>, k: u8) -> VortexResult { + let validity = array.validity()?; + match array.ptype() { + PType::F16 => { + let (primary, secondary) = split_f16(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + Some( + PrimitiveArray::new(Buffer::from(secondary), NonNullable.into()) + .into_array(), + ), + PType::F16, + k, + ) + } + PType::F32 => { + let (primary, secondary) = split_f32(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + Some( + PrimitiveArray::new(Buffer::from(secondary), NonNullable.into()) + .into_array(), + ), + PType::F32, + k, + ) + } + PType::F64 => { + let (primary, secondary) = split_f64(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + Some( + PrimitiveArray::new(Buffer::from(secondary), NonNullable.into()) + .into_array(), + ), + PType::F64, + k, + ) + } + ptype => vortex_bail!("FloatQuant requires f16, f32, or f64, got {ptype}"), + } + } + + /// Split floats whose lowest `k` fraction bits are zero. + pub fn from_primitive_constant_secondary( + array: ArrayView<'_, Primitive>, + k: u8, + ) -> VortexResult { + let validity = array.validity()?; + match array.ptype() { + PType::F16 => { + let primary = split_primary_f16(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + None, + PType::F16, + k, + ) + } + PType::F32 => { + let primary = split_primary_f32(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + None, + PType::F32, + k, + ) + } + PType::F64 => { + let primary = split_primary_f64(array.as_slice::(), k)?; + Self::try_new( + PrimitiveArray::new(Buffer::from(primary), validity).into_array(), + None, + PType::F64, + k, + ) + } + ptype => vortex_bail!("FloatQuant requires f16, f32, or f64, got {ptype}"), + } + } + + /// Split a constant-secondary float array into frame-of-reference primary values. + pub fn primary_for_primitive( + array: ArrayView<'_, Primitive>, + k: u8, + primary_min: u64, + ) -> VortexResult { + let validity = array.validity()?; + match array.ptype() { + PType::F16 => { + let primary_min = u16::try_from(primary_min)?; + let primary = split_primary_for_f16(array.as_slice::(), k, primary_min)?; + Ok(PrimitiveArray::new(Buffer::from(primary), validity)) + } + PType::F32 => { + let primary_min = u32::try_from(primary_min)?; + let primary = split_primary_for_f32(array.as_slice::(), k, primary_min)?; + Ok(PrimitiveArray::new(Buffer::from(primary), validity)) + } + PType::F64 => { + let primary = split_primary_for_f64(array.as_slice::(), k, primary_min)?; + Ok(PrimitiveArray::new(Buffer::from(primary), validity)) + } + ptype => vortex_bail!("FloatQuant requires f16, f32, or f64, got {ptype}"), + } + } +} + +/// Compression facts derived during FloatQuant split selection. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FloatQuantAnalysis { + /// Selected low-bit width. + pub k: u8, + /// Bit width of the frame-of-reference primary values. + pub primary_bit_width: u8, + /// Minimum primary value before frame-of-reference subtraction. + pub primary_min: u64, + /// Bit width of the secondary values. Zero identifies an implicit-zero child. + pub secondary_bit_width: u8, +} + +/// Analyze a canonical float array for a FloatQuant split. +pub fn analyze_float_quant(array: ArrayView<'_, Primitive>) -> Option { + match array.ptype() { + PType::F16 => analyze_f16(array.as_slice::()), + PType::F32 => analyze_f32(array.as_slice::()), + PType::F64 => analyze_f64(array.as_slice::()), + _ => None, + } +} + +fn analyze_bits( + low_bits_or: u64, + precision_bits: u8, + len: usize, + primary_min: u64, + primary_max: u64, +) -> Option { + if len == 0 { + return None; + } + + let mut best = None; + for k in 1..=precision_bits { + let low_mask = (1_u64 << k) - 1; + let secondary_bit_width = + u8::try_from(u64::BITS - (low_bits_or & low_mask).leading_zeros()).unwrap_or(u8::MAX); + if k - secondary_bit_width < 2 { + continue; + } + + let shifted_min = primary_min >> k; + let shifted_max = primary_max >> k; + let primary_bit_width = + u8::try_from(u64::BITS - (shifted_max - shifted_min).leading_zeros()) + .unwrap_or(u8::MAX); + let total_bit_width = primary_bit_width + secondary_bit_width; + let candidate = ( + total_bit_width, + secondary_bit_width, + k, + primary_bit_width, + shifted_min, + ); + if best.is_none_or(|current| candidate < current) { + best = Some(candidate); + } + } + + let (_, secondary_bit_width, k, primary_bit_width, primary_min) = best?; + Some(FloatQuantAnalysis { + k, + primary_bit_width, + primary_min, + secondary_bit_width, + }) +} + +fn analyze_f16(values: &[f16]) -> Option { + let mut minimum = u16::MAX; + let mut maximum = u16::MIN; + let mut low_bits_or = 0_u16; + for value in values { + let bits = value.to_bits(); + let ordered = ordered_u16(bits); + minimum = minimum.min(ordered); + maximum = maximum.max(ordered); + low_bits_or |= bits; + } + analyze_bits( + u64::from(low_bits_or), + 10, + values.len(), + u64::from(minimum), + u64::from(maximum), + ) +} + +fn analyze_f32(values: &[f32]) -> Option { + let mut minimum = u32::MAX; + let mut maximum = u32::MIN; + let mut low_bits_or = 0_u32; + for value in values { + let bits = value.to_bits(); + let ordered = ordered_u32(bits); + minimum = minimum.min(ordered); + maximum = maximum.max(ordered); + low_bits_or |= bits; + } + analyze_bits( + u64::from(low_bits_or), + 23, + values.len(), + u64::from(minimum), + u64::from(maximum), + ) +} + +fn analyze_f64(values: &[f64]) -> Option { + let mut minimum = u64::MAX; + let mut maximum = u64::MIN; + let mut low_bits_or = 0_u64; + for value in values { + let bits = value.to_bits(); + let ordered = ordered_u64(bits); + minimum = minimum.min(ordered); + maximum = maximum.max(ordered); + low_bits_or |= bits; + } + analyze_bits(low_bits_or, 52, values.len(), minimum, maximum) +} + +fn latent_ptype(ptype: PType) -> VortexResult { + match ptype { + PType::F16 => Ok(PType::U16), + PType::F32 => Ok(PType::U32), + PType::F64 => Ok(PType::U64), + _ => vortex_bail!("FloatQuant requires f16, f32, or f64, got {ptype}"), + } +} + +fn precision_bits(ptype: PType) -> VortexResult { + match ptype { + PType::F16 => Ok(10), + PType::F32 => Ok(23), + PType::F64 => Ok(52), + _ => vortex_bail!("FloatQuant requires f16, f32, or f64, got {ptype}"), + } +} + +fn ordered_u16(bits: u16) -> u16 { + if bits & (1_u16 << 15) == 0 { + bits ^ (1_u16 << 15) + } else { + !bits + } +} + +fn ordered_u32(bits: u32) -> u32 { + if bits & (1_u32 << 31) == 0 { + bits ^ (1_u32 << 31) + } else { + !bits + } +} + +fn ordered_u64(bits: u64) -> u64 { + if bits & (1_u64 << 63) == 0 { + bits ^ (1_u64 << 63) + } else { + !bits + } +} + +fn split_f16(values: &[f16], k: u8) -> VortexResult<(Vec, Vec)> { + vortex_ensure!(k > 0 && k <= 10, "FloatQuant f16 k must be in 1..=10"); + let low_mask = (1_u16 << k) - 1; + let mut primary = Vec::with_capacity(values.len()); + let mut secondary = Vec::with_capacity(values.len()); + for &value in values { + let bits = value.to_bits(); + let ordered = ordered_u16(bits); + primary.push(ordered >> k); + let low = ordered & low_mask; + secondary.push(if bits & (1_u16 << 15) == 0 { + low + } else { + low_mask - low + }); + } + Ok((primary, secondary)) +} + +fn split_f32(values: &[f32], k: u8) -> VortexResult<(Vec, Vec)> { + vortex_ensure!(k > 0 && k <= 23, "FloatQuant f32 k must be in 1..=23"); + let low_mask = (1_u32 << k) - 1; + let mut primary = Vec::with_capacity(values.len()); + let mut secondary = Vec::with_capacity(values.len()); + for &value in values { + let bits = value.to_bits(); + let ordered = ordered_u32(bits); + primary.push(ordered >> k); + let low = ordered & low_mask; + secondary.push(if bits & (1_u32 << 31) == 0 { + low + } else { + low_mask - low + }); + } + Ok((primary, secondary)) +} + +fn split_f64(values: &[f64], k: u8) -> VortexResult<(Vec, Vec)> { + vortex_ensure!(k > 0 && k <= 52, "FloatQuant f64 k must be in 1..=52"); + let low_mask = (1_u64 << k) - 1; + let mut primary = Vec::with_capacity(values.len()); + let mut secondary = Vec::with_capacity(values.len()); + for &value in values { + let bits = value.to_bits(); + let ordered = ordered_u64(bits); + primary.push(ordered >> k); + let low = ordered & low_mask; + secondary.push(if bits & (1_u64 << 63) == 0 { + low + } else { + low_mask - low + }); + } + Ok((primary, secondary)) +} + +fn split_primary_f32(values: &[f32], k: u8) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 23, "FloatQuant f32 k must be in 1..=23"); + let low_mask = (1_u32 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + Ok(values + .iter() + .map(|value| ordered_u32(value.to_bits()) >> k) + .collect()) +} + +fn split_primary_f16(values: &[f16], k: u8) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 10, "FloatQuant f16 k must be in 1..=10"); + let low_mask = (1_u16 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + Ok(values + .iter() + .map(|value| ordered_u16(value.to_bits()) >> k) + .collect()) +} + +fn split_primary_f64(values: &[f64], k: u8) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 52, "FloatQuant f64 k must be in 1..=52"); + let low_mask = (1_u64 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + Ok(values + .iter() + .map(|value| ordered_u64(value.to_bits()) >> k) + .collect()) +} + +fn split_primary_for_f32(values: &[f32], k: u8, primary_min: u32) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 23, "FloatQuant f32 k must be in 1..=23"); + let low_mask = (1_u32 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + values + .iter() + .map(|value| { + (ordered_u32(value.to_bits()) >> k) + .checked_sub(primary_min) + .ok_or_else(|| vortex_error::vortex_err!("FloatQuant primary minimum is invalid")) + }) + .collect::>>() +} + +fn split_primary_for_f16(values: &[f16], k: u8, primary_min: u16) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 10, "FloatQuant f16 k must be in 1..=10"); + let low_mask = (1_u16 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + values + .iter() + .map(|value| { + (ordered_u16(value.to_bits()) >> k) + .checked_sub(primary_min) + .ok_or_else(|| vortex_error::vortex_err!("FloatQuant primary minimum is invalid")) + }) + .collect::>>() +} + +fn split_primary_for_f64(values: &[f64], k: u8, primary_min: u64) -> VortexResult> { + vortex_ensure!(k > 0 && k <= 52, "FloatQuant f64 k must be in 1..=52"); + let low_mask = (1_u64 << k) - 1; + vortex_ensure!( + values.iter().all(|value| value.to_bits() & low_mask == 0), + "FloatQuant constant secondary requires zero low bits" + ); + values + .iter() + .map(|value| { + (ordered_u64(value.to_bits()) >> k) + .checked_sub(primary_min) + .ok_or_else(|| vortex_error::vortex_err!("FloatQuant primary minimum is invalid")) + }) + .collect::>>() +} + +fn join_f16(primary: u16, secondary: u16, k: u8) -> f16 { + let low_mask = (1_u16 << k) - 1; + let sign_cutoff = (1_u16 << 15) >> k; + let low = if primary >= sign_cutoff { + secondary + } else { + low_mask.wrapping_sub(secondary) + }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u16 << 15) == 0 { + !ordered + } else { + ordered ^ (1_u16 << 15) + }; + f16::from_bits(bits) +} + +fn join_f32(primary: u32, secondary: u32, k: u8) -> f32 { + let low_mask = (1_u32 << k) - 1; + let sign_cutoff = (1_u32 << 31) >> k; + let low = if primary >= sign_cutoff { + secondary + } else { + low_mask.wrapping_sub(secondary) + }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u32 << 31) == 0 { + !ordered + } else { + ordered ^ (1_u32 << 31) + }; + f32::from_bits(bits) +} + +fn join_f64(primary: u64, secondary: u64, k: u8) -> f64 { + let low_mask = (1_u64 << k) - 1; + let sign_cutoff = (1_u64 << 63) >> k; + let low = if primary >= sign_cutoff { + secondary + } else { + low_mask.wrapping_sub(secondary) + }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u64 << 63) == 0 { + !ordered + } else { + ordered ^ (1_u64 << 63) + }; + f64::from_bits(bits) +} + +fn join_zero_f32(primary: u32, k: u8) -> f32 { + let low_mask = (1_u32 << k) - 1; + let sign_cutoff = (1_u32 << 31) >> k; + let low = if primary >= sign_cutoff { 0 } else { low_mask }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u32 << 31) == 0 { + !ordered + } else { + ordered ^ (1_u32 << 31) + }; + f32::from_bits(bits) +} + +fn join_zero_f16(primary: u16, k: u8) -> f16 { + let low_mask = (1_u16 << k) - 1; + let sign_cutoff = (1_u16 << 15) >> k; + let low = if primary >= sign_cutoff { 0 } else { low_mask }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u16 << 15) == 0 { + !ordered + } else { + ordered ^ (1_u16 << 15) + }; + f16::from_bits(bits) +} + +fn join_zero_f64(primary: u64, k: u8) -> f64 { + let low_mask = (1_u64 << k) - 1; + let sign_cutoff = (1_u64 << 63) >> k; + let low = if primary >= sign_cutoff { 0 } else { low_mask }; + let ordered = (primary << k).wrapping_add(low); + let bits = if ordered & (1_u64 << 63) == 0 { + !ordered + } else { + ordered ^ (1_u64 << 63) + }; + f64::from_bits(bits) +} + +fn decode( + array: ArrayView<'_, FloatQuant>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if array.dtype().as_ptype() == PType::F64 + && let Some(decoded) = decode_fastlanes_zero(array, ctx)? + { + return Ok(decoded); + } + if let Some(decoded) = decode_fastlanes_pair(array)? { + return Ok(decoded); + } + + let primary = array.primary().clone().execute::(ctx)?; + let validity = primary.validity()?; + let k = array.data().k; + let Some(secondary) = array.secondary() else { + return Ok(match PType::try_from(array.dtype())? { + PType::F16 => PrimitiveArray::new( + primary + .into_buffer::() + .map_each_in_place(|primary| join_zero_f16(primary, k)) + .freeze(), + validity, + ), + PType::F32 => PrimitiveArray::new( + primary + .into_buffer::() + .map_each_in_place(|primary| join_zero_f32(primary, k)) + .freeze(), + validity, + ), + PType::F64 => PrimitiveArray::new( + primary + .into_buffer::() + .map_each_in_place(|primary| join_zero_f64(primary, k)) + .freeze(), + validity, + ), + ptype => vortex_panic!("unsupported FloatQuant ptype {ptype}"), + }); + }; + let secondary = secondary.clone().execute::(ctx)?; + Ok(match PType::try_from(array.dtype())? { + PType::F16 => { + let secondary_values = secondary.as_slice::(); + let mut index = 0; + let values = primary + .into_buffer::() + .map_each_in_place(|primary| { + let value = join_f16(primary, secondary_values[index], k); + index += 1; + value + }) + .freeze(); + PrimitiveArray::new(values, validity) + } + PType::F32 => { + let secondary_values = secondary.as_slice::(); + let mut index = 0; + let values = primary + .into_buffer::() + .map_each_in_place(|primary| { + let value = join_f32(primary, secondary_values[index], k); + index += 1; + value + }) + .freeze(); + PrimitiveArray::new(values, validity) + } + PType::F64 => { + let secondary_values = secondary.as_slice::(); + let mut index = 0; + let values = primary + .into_buffer::() + .map_each_in_place(|primary| { + let value = join_f64(primary, secondary_values[index], k); + index += 1; + value + }) + .freeze(); + PrimitiveArray::new(values, validity) + } + ptype => vortex_panic!("unsupported FloatQuant ptype {ptype}"), + }) +} + +fn decode_fastlanes_zero( + array: ArrayView<'_, FloatQuant>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if array.secondary().is_some() { + return Ok(None); + } + let Some(primary_for) = array.primary().as_opt::() else { + return Ok(None); + }; + let Some(primary) = primary_for.encoded().as_opt::() else { + return Ok(None); + }; + + let k = array.data().k; + let reference = primary_for + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("validated f64 primary reference"); + Ok(Some(unpack_map::(primary, ctx, |primary| { + join_zero_f64(primary.wrapping_add(reference), k) + })?)) +} + +fn decode_fastlanes_pair(array: ArrayView<'_, FloatQuant>) -> VortexResult> { + let Some(secondary) = array + .secondary() + .and_then(|child| child.as_opt::()) + else { + return Ok(None); + }; + let Some(primary_for) = array.primary().as_opt::() else { + return Ok(None); + }; + let Some(primary) = primary_for.encoded().as_opt::() else { + return Ok(None); + }; + if primary.patches().is_some() + || secondary.patches().is_some() + || primary.offset() != secondary.offset() + { + return Ok(None); + } + + let validity = primary.validity()?; + let k = array.data().k; + Ok(Some(match PType::try_from(array.dtype())? { + PType::F16 => { + let reference = primary_for + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("validated f16 primary reference"); + PrimitiveArray::new( + unpack_pair_map::(primary, secondary, |primary, secondary| { + join_f16(primary.wrapping_add(reference), secondary, k) + })?, + validity, + ) + } + PType::F32 => { + let reference = primary_for + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("validated f32 primary reference"); + PrimitiveArray::new( + unpack_pair_map::(primary, secondary, |primary, secondary| { + join_f32(primary.wrapping_add(reference), secondary, k) + })?, + validity, + ) + } + PType::F64 => { + let reference = primary_for + .reference_scalar() + .as_primitive() + .typed_value::() + .vortex_expect("validated f64 primary reference"); + PrimitiveArray::new( + unpack_pair_map::(primary, secondary, |primary, secondary| { + join_f64(primary.wrapping_add(reference), secondary, k) + })?, + validity, + ) + } + ptype => vortex_panic!("unsupported FloatQuant ptype {ptype}"), + })) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::assert_arrays_eq; + use vortex_array::assert_nth_scalar; + use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::validity::Validity; + use vortex_buffer::ByteBufferMut; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + use vortex_session::registry::ReadContext; + + use super::*; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session(); + crate::initialize(&session); + session + }); + + #[test] + fn f16_bit_patterns_roundtrip() -> VortexResult<()> { + let values = [ + f16::from_bits(0xfc00), + f16::from_f32(-1.5), + f16::NEG_ZERO, + f16::ZERO, + f16::from_f32(1.5), + f16::INFINITY, + f16::from_bits(0x7e34), + f16::from_bits(0xfe56), + ]; + let array = PrimitiveArray::from_iter(values); + let encoded = FloatQuant::from_primitive(array.as_view(), 5)?; + let decoded = encoded + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + assert_eq!( + decoded + .as_slice::() + .iter() + .map(|value| value.to_bits()) + .collect::>(), + values.map(f16::to_bits) + ); + Ok(()) + } + + #[test] + fn fixed_tree_analysis_accounts_for_secondary_width() -> VortexResult<()> { + let values = PrimitiveArray::from_iter((0_u32..4096).map(|index| { + let value = f64::from(f32::from_bits(0x3f80_0000 | index.wrapping_mul(7_919))); + if index % 10 == 0 { + f64::from_bits(value.to_bits() | 1) + } else { + value + } + })); + let analysis = analyze_float_quant(values.as_view()).vortex_expect("FloatQuant input"); + assert_eq!(analysis.k, 29); + assert_eq!(analysis.secondary_bit_width, 1); + + let general = + PrimitiveArray::from_iter((0_u32..4096).map(|index| { + f32::from_bits(0x3f80_0000 | (index.wrapping_mul(7_919) & 0x007f_ffff)) + })); + assert_eq!(analyze_float_quant(general.as_view()), None); + Ok(()) + } + + #[test] + fn float_bit_patterns_roundtrip() -> VortexResult<()> { + let values = [ + f64::NEG_INFINITY, + -1.5, + -0.0, + 0.0, + 1.5, + f64::INFINITY, + f64::from_bits(0x7ff8_0000_0000_1234), + f64::from_bits(0xfff8_0000_0000_5678), + ]; + let array = PrimitiveArray::from_iter(values); + let encoded = FloatQuant::from_primitive(array.as_view(), 29)?; + let decoded = encoded + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + assert_eq!( + decoded + .as_slice::() + .iter() + .map(|value| value.to_bits()) + .collect::>(), + values + .iter() + .map(|value| value.to_bits()) + .collect::>() + ); + Ok(()) + } + + #[test] + fn f32_bit_patterns_roundtrip() -> VortexResult<()> { + let values = [ + f32::NEG_INFINITY, + -1.5, + -0.0, + 0.0, + 1.5, + f32::INFINITY, + f32::from_bits(0x7fc0_1234), + f32::from_bits(0xffc0_5678), + ]; + let array = PrimitiveArray::from_iter(values); + let encoded = FloatQuant::from_primitive(array.as_view(), 8)?; + let decoded = encoded + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + assert_eq!( + decoded + .as_slice::() + .iter() + .map(|value| value.to_bits()) + .collect::>(), + values.map(f32::to_bits) + ); + Ok(()) + } + + #[test] + fn rejects_invalid_split_shapes() { + let primary = PrimitiveArray::from_iter([0_u32, 1, 2]).into_array(); + assert!(FloatQuant::try_new(primary.clone(), None, PType::F32, 0).is_err()); + assert!(FloatQuant::try_new(primary.clone(), None, PType::F32, 24).is_err()); + assert!(FloatQuant::try_new(primary.clone(), None, PType::U32, 8).is_err()); + + let wrong_ptype = PrimitiveArray::from_iter([0_u64, 1, 2]).into_array(); + assert!(FloatQuant::try_new(primary.clone(), Some(wrong_ptype), PType::F32, 8).is_err()); + let nullable = PrimitiveArray::from_option_iter([Some(0_u32), None, Some(2)]).into_array(); + assert!(FloatQuant::try_new(primary, Some(nullable), PType::F32, 8).is_err()); + } + + #[test] + fn nullable_slice_and_scalar_access() -> VortexResult<()> { + let array = PrimitiveArray::new( + Buffer::from(vec![1.25_f32, 0.0, -0.0, 42.5, -10.0]), + Validity::from_iter([true, false, true, true, false]), + ); + let encoded = FloatQuant::from_primitive(array.as_view(), 8)?; + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(encoded, array, &mut ctx); + assert_nth_scalar!(encoded, 3, 42.5_f32, &mut ctx); + assert!(encoded.execute_scalar(1, &mut ctx)?.is_null()); + + let sliced = encoded.into_array().slice(1..4)?; + assert!(sliced.is::()); + assert_arrays_eq!(sliced, array.into_array().slice(1..4)?, &mut ctx); + Ok(()) + } + + #[test] + fn implicit_zero_secondary_roundtrip() -> VortexResult<()> { + let original = PrimitiveArray::from_option_iter([ + Some(f64::from(-10.5_f32)), + None, + Some(f64::from(-0.0_f32)), + Some(f64::from(42.25_f32)), + ]); + let encoded = FloatQuant::from_primitive_constant_secondary(original.as_view(), 29)?; + assert!(encoded.secondary().is_none()); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(encoded, original, &mut ctx); + assert_nth_scalar!(encoded, 3, 42.25_f64, &mut ctx); + + let sliced = encoded.into_array().slice(1..4)?; + let expected = original.into_array().slice(1..4)?; + assert_arrays_eq!(sliced, expected, &mut ctx); + + let dtype = sliced.dtype().clone(); + let len = sliced.len(); + let array_context = ArrayContext::empty(); + let serialized = + sliced.serialize(&array_context, &SESSION, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in serialized { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_context.to_ids()), + &SESSION, + )?; + assert!(decoded.as_::().secondary().is_none()); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn implicit_zero_secondary_rejects_nonzero_low_bits() { + let f32_values = PrimitiveArray::from_iter([f32::from_bits(1.0_f32.to_bits() | 1)]); + assert!(FloatQuant::from_primitive_constant_secondary(f32_values.as_view(), 8).is_err()); + assert!(FloatQuant::primary_for_primitive(f32_values.as_view(), 8, 0).is_err()); + + let f64_values = PrimitiveArray::from_iter([f64::from_bits(1.0_f64.to_bits() | 1)]); + assert!(FloatQuant::from_primitive_constant_secondary(f64_values.as_view(), 29).is_err()); + assert!(FloatQuant::primary_for_primitive(f64_values.as_view(), 29, 0).is_err()); + } + + #[test] + fn primary_for_primitive_rejects_invalid_reference() { + let values = PrimitiveArray::from_iter([1.0_f32, 2.0]); + assert!( + FloatQuant::primary_for_primitive(values.as_view(), 1, u64::from(u32::MAX)).is_err() + ); + } + + #[test] + fn fastlanes_zero_decode_roundtrip_and_slice() -> VortexResult<()> { + let original = PrimitiveArray::from_option_iter((0_u32..4097).map(|index| { + (index % 17 != 0).then(|| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + f64::from(f32::from_bits(0x3f80_0000 | mantissa)) + }) + })); + let analysis = analyze_float_quant(original.as_view()).vortex_expect("FloatQuant input"); + assert_eq!(analysis.secondary_bit_width, 0); + let primary = FloatQuant::primary_for_primitive( + original.as_view(), + analysis.k, + analysis.primary_min, + )?; + // SAFETY: The analysis computes the exact primary width. + let primary = unsafe { + vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked( + primary, + analysis.primary_bit_width, + )? + }; + let primary = FoR::try_new(primary.into_array(), Scalar::from(analysis.primary_min))?; + let encoded = + FloatQuant::try_new(primary.into_array(), None, PType::F64, analysis.k)?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(encoded, original, &mut ctx); + assert_arrays_eq!( + encoded.slice(3..2051)?, + original.into_array().slice(3..2051)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn fastlanes_pair_decode_roundtrip_and_slice() -> VortexResult<()> { + let original = PrimitiveArray::from_iter((0_u32..4097).map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + let value = f64::from(f32::from_bits(0x3f80_0000 | mantissa)); + if index % 10 == 0 { + f64::from_bits(value.to_bits() | 1) + } else { + value + } + })); + let analysis = analyze_float_quant(original.as_view()).vortex_expect("FloatQuant input"); + assert_eq!(analysis.secondary_bit_width, 1); + let split = FloatQuant::from_primitive(original.as_view(), analysis.k)?; + let primary = split + .primary() + .clone() + .execute::(&mut SESSION.create_execution_ctx())?; + let secondary = split + .secondary() + .vortex_expect("nonzero secondary") + .clone() + .execute::(&mut SESSION.create_execution_ctx())?; + let biased_primary = PrimitiveArray::from_iter( + primary + .as_slice::() + .iter() + .map(|value| value - analysis.primary_min), + ); + // SAFETY: The analysis computes the exact primary width. + let primary = unsafe { + vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked( + biased_primary, + analysis.primary_bit_width, + )? + }; + let primary = FoR::try_new(primary.into_array(), Scalar::from(analysis.primary_min))?; + // SAFETY: The analysis computes the exact secondary width. + let secondary = unsafe { + vortex_fastlanes::bitpack_compress::bitpack_encode_unchecked( + secondary, + analysis.secondary_bit_width, + )? + }; + let encoded = FloatQuant::try_new( + primary.into_array(), + Some(secondary.into_array()), + PType::F64, + analysis.k, + )? + .into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(encoded, original, &mut ctx); + assert_arrays_eq!( + encoded.slice(3..2051)?, + original.into_array().slice(3..2051)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn serialization_roundtrip() -> VortexResult<()> { + let original = PrimitiveArray::from_option_iter([ + Some(f64::NEG_INFINITY), + None, + Some(-0.0), + Some(42.25), + Some(f64::from_bits(0x7ff8_0000_0000_1234)), + ]); + let encoded = FloatQuant::from_primitive(original.as_view(), 29)?; + let sliced = encoded.into_array().slice(1..5)?; + let dtype = sliced.dtype().clone(); + let len = sliced.len(); + let array_context = ArrayContext::empty(); + let serialized = + sliced.serialize(&array_context, &SESSION, &SerializeOptions::default())?; + let mut bytes = ByteBufferMut::empty(); + for buffer in serialized { + bytes.extend_from_slice(buffer.as_ref()); + } + + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + &dtype, + len, + &ReadContext::new(array_context.to_ids()), + &SESSION, + )?; + assert!(decoded.is::()); + assert_arrays_eq!( + decoded, + original.into_array().slice(1..5)?, + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + #[test] + fn conformance() -> VortexResult<()> { + let explicit = PrimitiveArray::from_option_iter([ + Some(f32::NEG_INFINITY), + None, + Some(-0.0), + Some(0.0), + Some(42.25), + Some(f32::INFINITY), + ]); + let explicit = FloatQuant::from_primitive(explicit.as_view(), 8)?.into_array(); + let implicit = PrimitiveArray::from_option_iter([ + Some(f64::from(-10.5_f32)), + None, + Some(f64::from(-0.0_f32)), + Some(f64::from(42.25_f32)), + ]); + let implicit = + FloatQuant::from_primitive_constant_secondary(implicit.as_view(), 29)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + for array in [explicit, implicit] { + test_array_consistency(&array, &mut ctx); + } + Ok(()) + } +} diff --git a/encodings/float-quant/src/lib.rs b/encodings/float-quant/src/lib.rs new file mode 100644 index 00000000000..6efde148ff0 --- /dev/null +++ b/encodings/float-quant/src/lib.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lossless float quantization as two integer child arrays. + +mod array; +mod rules; +mod slice; + +pub use array::*; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +/// Register the float quantization encoding in one session. +pub fn initialize(session: &VortexSession) { + session.arrays().register(FloatQuant); +} diff --git a/encodings/float-quant/src/rules.rs b/encodings/float-quant/src/rules.rs new file mode 100644 index 00000000000..5faedb68866 --- /dev/null +++ b/encodings/float-quant/src/rules.rs @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleSet; + +use crate::FloatQuant; + +pub(crate) static RULES: ParentRuleSet = + ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(FloatQuant))]); diff --git a/encodings/float-quant/src/slice.rs b/encodings/float-quant/src/slice.rs new file mode 100644 index 00000000000..b46d982d0d5 --- /dev/null +++ b/encodings/float-quant/src/slice.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::dtype::PType; +use vortex_error::VortexResult; + +use crate::FloatQuant; +use crate::FloatQuantArraySlotsExt; + +impl SliceReduce for FloatQuant { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + Ok(Some( + FloatQuant::try_new( + array.primary().slice(range.clone())?, + array + .secondary() + .map(|secondary| secondary.slice(range)) + .transpose()?, + PType::try_from(array.dtype())?, + array.data().k, + )? + .into_array(), + )) + } +} diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index 965a80c6259..4569cdf765f 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -28,6 +28,7 @@ vortex = { workspace = true, features = [ "zstd", ] } vortex-arrow = { workspace = true } +vortex-btrblocks = { workspace = true } vortex-spatial = { workspace = true } vortex-tensor = { workspace = true } # TODO(connor): In the future, this might be inside vortex. diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index f9c4ee63734..4303237ebcd 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -338,11 +338,22 @@ pub async fn write_parquet_as_vortex( parquet_path: PathBuf, vortex_path: &str, compaction: CompactionStrategy, +) -> anyhow::Result { + let write_options = compaction.apply_options(SESSION.write_options()); + write_parquet_as_vortex_with_options(parquet_path, vortex_path, write_options).await +} + +/// Convert a Parquet file to Vortex with explicit write options. +/// +/// The function skips conversion when the output file already exists. +pub async fn write_parquet_as_vortex_with_options( + parquet_path: PathBuf, + vortex_path: &str, + write_options: VortexWriteOptions, ) -> anyhow::Result { idempotent_async(vortex_path, |output_fname| async move { let mut output_file = File::create(&output_fname).await?; let data = parquet_to_vortex_chunks(parquet_path).await?; - let write_options = compaction.apply_options(SESSION.write_options()); write_options .write(&mut output_file, data.into_array().to_array_stream()) .await?; diff --git a/vortex-bench/src/datasets/feature_vectors.rs b/vortex-bench/src/datasets/feature_vectors.rs index e09f29f49ec..2ae9e5b7701 100644 --- a/vortex-bench/src/datasets/feature_vectors.rs +++ b/vortex-bench/src/datasets/feature_vectors.rs @@ -18,13 +18,20 @@ use parquet::arrow::ArrowWriter; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; use crate::CompactionStrategy; use crate::Format; +use crate::conversions::parquet_to_vortex_chunks; use crate::conversions::write_parquet_as_vortex; +use crate::datasets::Dataset; use crate::idempotent_async; use crate::random_access::BenchDataset; use crate::random_access::data_path; +use crate::vector_dataset::TrainLayout; +use crate::vector_dataset::VectorDataset; +use crate::vector_dataset::download; /// Dataset identifier used for data path generation. pub const DATASET: &str = "feature_vectors"; @@ -34,6 +41,56 @@ pub const ROW_COUNT: usize = 1_000_000; pub struct FeatureVectorsData; +/// A real f32 embedding dataset from the GloVe vector benchmark corpus. +pub struct GloveEmbeddingsData; + +/// A real f64 embedding dataset from the OpenAI-on-C4 vector benchmark corpus. +pub struct OpenAiEmbeddingsData; + +#[async_trait] +impl Dataset for GloveEmbeddingsData { + fn name(&self) -> &str { + "glove-embeddings-100k" + } + + async fn to_vortex_array(&self, _ctx: &mut ExecutionCtx) -> Result { + Ok(parquet_to_vortex_chunks(self.to_parquet_path().await?) + .await? + .into()) + } + + async fn to_parquet_path(&self) -> Result { + let paths = download(VectorDataset::GloveSmall100k, TrainLayout::Single).await?; + paths + .train_files + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("GloVe embedding train file is missing")) + } +} + +#[async_trait] +impl Dataset for OpenAiEmbeddingsData { + fn name(&self) -> &str { + "openai-c4-embeddings-50k" + } + + async fn to_vortex_array(&self, _ctx: &mut ExecutionCtx) -> Result { + Ok(parquet_to_vortex_chunks(self.to_parquet_path().await?) + .await? + .into()) + } + + async fn to_parquet_path(&self) -> Result { + let paths = download(VectorDataset::OpenaiSmall50k, TrainLayout::Single).await?; + paths + .train_files + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("OpenAI embedding train file is missing")) + } +} + #[async_trait] impl BenchDataset for FeatureVectorsData { fn name(&self) -> &str { diff --git a/vortex-bench/src/datasets/local_parquet.rs b/vortex-bench/src/datasets/local_parquet.rs new file mode 100644 index 00000000000..fb3c8c983a6 --- /dev/null +++ b/vortex-bench/src/datasets/local_parquet.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fs::File; +use std::path::PathBuf; + +use anyhow::Result; +use anyhow::bail; +use async_trait::async_trait; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; + +use crate::Format; +use crate::conversions::parquet_to_vortex_chunks; +use crate::datasets::Dataset; +use crate::random_access::BenchDataset; + +/// A local Parquet file for compression and random-access benchmarks. +pub struct LocalParquetData { + name: String, + path: PathBuf, + row_count: u64, +} + +impl LocalParquetData { + /// Read the dataset name and row count from a local Parquet file. + pub fn try_new(path: PathBuf) -> Result { + let name = path + .file_stem() + .and_then(|value| value.to_str()) + .ok_or_else(|| anyhow::anyhow!("local Parquet path has no valid file stem"))? + .to_string(); + let reader = ParquetRecordBatchReaderBuilder::try_new(File::open(&path)?)?; + let row_count = u64::try_from(reader.metadata().file_metadata().num_rows())?; + Ok(Self { + name, + path, + row_count, + }) + } +} + +#[async_trait] +impl Dataset for LocalParquetData { + fn name(&self) -> &str { + &self.name + } + + async fn to_vortex_array(&self, _ctx: &mut ExecutionCtx) -> Result { + Ok(parquet_to_vortex_chunks(self.path.clone()).await?.into()) + } + + async fn to_parquet_path(&self) -> Result { + Ok(self.path.clone()) + } +} + +#[async_trait] +impl BenchDataset for LocalParquetData { + fn name(&self) -> &str { + &self.name + } + + fn row_count(&self) -> u64 { + self.row_count + } + + async fn path(&self, format: Format) -> Result { + if format == Format::Parquet { + return Ok(self.path.clone()); + } + bail!("local Parquet dataset does not provide {format}") + } +} diff --git a/vortex-bench/src/datasets/mod.rs b/vortex-bench/src/datasets/mod.rs index 7eb341edbb7..b53ed965324 100644 --- a/vortex-bench/src/datasets/mod.rs +++ b/vortex-bench/src/datasets/mod.rs @@ -14,6 +14,7 @@ use crate::clickbench::Flavor; pub mod data_downloads; pub mod feature_vectors; +pub mod local_parquet; pub mod nested_lists; pub mod nested_structs; pub mod struct_list_of_ints; diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 286d9d41bcc..bb5d8619d5f 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -34,6 +34,10 @@ use vortex::error::vortex_err; use vortex::file::VortexWriteOptions; use vortex::file::WriteStrategyBuilder; use vortex::utils::aliases::hash_map::HashMap; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::float::FloatQuantScheme; +use vortex_btrblocks::schemes::float::OrderedBlockResidualScheme; +use vortex_btrblocks::schemes::integer::BlockResidualScheme; use crate::spatialbench::SpatialBenchBenchmark; use crate::vortex_queries::VortexBenchmark; @@ -70,6 +74,8 @@ pub use datasets::BenchmarkDataset; pub use output::BenchmarkOutput; pub use output::create_output_writer; use vortex::VortexSessionDefault; +use vortex::editions::EditionSessionExt; +use vortex::editions::PREVIEW_2026_08_0; pub use vortex::error::vortex_panic; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; @@ -80,6 +86,11 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; pub static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::default().with_tokio(); + session + .enable_edition(PREVIEW_2026_08_0) + .unwrap_or_else(|error| { + vortex_panic!("numeric benchmark edition is not registered: {error}") + }); vortex_spatial::initialize(&session); session }); @@ -239,6 +250,49 @@ pub enum CompactionStrategy { Default, } +/// Numeric scheme bundles for compression benchmarks. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum VortexNumericBundle { + /// Exclude the new numeric schemes. + PriorDefault, + /// Add the integer and ordered-float BlockResidual schemes. + BlockResidual, + /// Use the current Default compressor. + #[default] + CurrentDefault, +} + +impl VortexNumericBundle { + /// Return the stable CLI and file-name component for this bundle. + pub fn name(self) -> &'static str { + match self { + Self::PriorDefault => "prior-default", + Self::BlockResidual => "block-residual", + Self::CurrentDefault => "current-default", + } + } + + /// Apply this numeric bundle to Vortex write options. + pub fn apply_options(self, options: VortexWriteOptions) -> VortexWriteOptions { + let compressor = match self { + Self::PriorDefault => BtrBlocksCompressorBuilder::default().exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]), + Self::BlockResidual => { + BtrBlocksCompressorBuilder::default().exclude_schemes([FloatQuantScheme.id()]) + } + Self::CurrentDefault => BtrBlocksCompressorBuilder::default(), + }; + options.with_strategy( + WriteStrategyBuilder::default() + .with_btrblocks_builder(compressor) + .build(), + ) + } +} + impl CompactionStrategy { pub fn apply_options(&self, options: VortexWriteOptions) -> VortexWriteOptions { match self { diff --git a/vortex-bench/src/measurements.rs b/vortex-bench/src/measurements.rs index d6d4ad85b32..905f120a06a 100644 --- a/vortex-bench/src/measurements.rs +++ b/vortex-bench/src/measurements.rs @@ -351,11 +351,11 @@ pub struct CompressionTimingMeasurement { impl ToJson for CompressionTimingMeasurement { fn to_json(&self) -> serde_json::Value { let (name, engine) = match self.format { - Format::OnDiskVortex => (self.name.to_string(), Engine::Vortex), + Format::OnDiskVortex | Format::VortexCompact => (self.name.to_string(), Engine::Vortex), Format::Parquet => (format!("parquet_rs-zstd {}", self.name), Engine::Vortex), Format::Lance => (format!("lance {}", self.name), Engine::Vortex), _ => vortex_panic!( - "CompressionTimingMeasurement only supports vortex, lance, and parquet formats" + "CompressionTimingMeasurement only supports Vortex, Lance, and Parquet formats" ), }; diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 4e22f042adf..38da11310c0 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -22,12 +22,14 @@ pco = { workspace = true, optional = true } rand = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } +vortex-block-residual = { workspace = true } vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } +vortex-float-quant = { workspace = true } vortex-fsst = { workspace = true } vortex-onpair = { workspace = true, optional = true } vortex-pco = { workspace = true, optional = true } diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6f38e29cd86..cb34d38a4e6 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -31,6 +31,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. &integer::ZigZagScheme, &integer::BitPackingScheme, + &integer::BlockResidualScheme, &integer::SparseScheme, &integer::IntDictScheme, &integer::RunEndScheme, @@ -44,6 +45,8 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// &float::ALPScheme, &float::ALPRDScheme, + &float::FloatQuantScheme, + &float::OrderedBlockResidualScheme, &float::FloatDictScheme, &float::NullDominatedSparseScheme, &float::FloatRLEScheme, @@ -170,9 +173,12 @@ impl BtrBlocksCompressorBuilder { allow(unused_mut) )] let mut excluded: Vec = vec![ + integer::BlockResidualScheme.id(), integer::SparseScheme.id(), integer::IntRLEScheme.id(), float::ALPRDScheme.id(), + float::FloatQuantScheme.id(), + float::OrderedBlockResidualScheme.id(), float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), string::StringDictScheme.id(), @@ -219,6 +225,7 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use rstest::rstest; use vortex_array::VTable; use vortex_fastlanes::FoR; @@ -233,7 +240,29 @@ mod tests { #[test] fn default_includes_all_schemes() { let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + assert_eq!( + builder + .schemes + .iter() + .map(|scheme| scheme.id()) + .collect::>(), + ALL_SCHEMES + .iter() + .map(|scheme| scheme.id()) + .collect::>() + ); + } + + #[test] + fn float_quant_can_be_excluded() { + let builder = + BtrBlocksCompressorBuilder::default().exclude_schemes([float::FloatQuantScheme.id()]); + assert!( + !builder + .schemes + .iter() + .any(|scheme| scheme.id() == float::FloatQuantScheme.id()) + ); } #[test] @@ -257,14 +286,18 @@ mod tests { assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } - #[test] - fn cuda_compatible_excludes_alprd() { + #[rstest] + #[case(float::ALPRDScheme.id())] + #[case(float::FloatQuantScheme.id())] + #[case(float::OrderedBlockResidualScheme.id())] + #[case(integer::BlockResidualScheme.id())] + fn cuda_compatible_excludes_non_cuda_schemes(#[case] scheme_id: SchemeId) { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); assert!( !builder .schemes .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) + .any(|scheme| scheme.id() == scheme_id) ); } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 1ca05c86b4e..59064ab3432 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -80,6 +80,11 @@ pub use builder::ALL_SCHEMES; pub use builder::BtrBlocksCompressorBuilder; pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::match_each_native_ptype; pub use vortex_compressor::CascadingCompressor; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; @@ -92,3 +97,24 @@ pub use vortex_compressor::stats::FloatStats; pub use vortex_compressor::stats::GenerateStatsOptions; pub use vortex_compressor::stats::IntegerStats; pub use vortex_compressor::stats::StringStats; +use vortex_error::VortexResult; + +fn normalize_null_values( + array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = array.validity()?; + let primitive = array.into_owned(); + if validity.definitely_no_nulls() { + return Ok(primitive); + } + + match_each_native_ptype!(primitive.ptype(), |T| { + primitive.map_each_with_validity::( + ctx, + |(value, valid)| { + if valid { value } else { T::default() } + }, + ) + }) +} diff --git a/vortex-btrblocks/src/schemes/float/float_quant.rs b/vortex-btrblocks/src/schemes/float/float_quant.rs new file mode 100644 index 00000000000..2edc224e7ab --- /dev/null +++ b/vortex-btrblocks/src/schemes/float/float_quant.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lossless float quantization with a fixed frame-of-reference child. + +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::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::PType; +use vortex_array::dtype::half::f16; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_array::vtable::validity_to_child; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::FL_CHUNK_SIZE; +use vortex_fastlanes::FoR; +use vortex_fastlanes::bitpack_compress::bitpack_primitive_map; +use vortex_fastlanes::bitpack_compress::bitpack_primitive_map_pair; +use vortex_float_quant::FloatQuant; +use vortex_float_quant::FloatQuantAnalysis; +use vortex_float_quant::analyze_float_quant; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::normalize_null_values; +use crate::schemes::sample_primitive_one_percent; + +// Food needs a factor above 1.078 to prevent its sample from displacing a smaller ALP tree. +const SELECTION_COST_FACTOR: f64 = 1.10; + +/// FloatQuant split with a fixed frame-of-reference primary child. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct FloatQuantScheme; + +impl Scheme for FloatQuantScheme { + fn scheme_name(&self) -> &'static str { + "vortex.float.float_quant" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_float() + } + + fn produced_encodings(&self) -> Vec { + vec![FloatQuant.id(), FoR.id(), BitPacked.id()] + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + if compress_ctx.finished_cascading() { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, data, _best_so_far, _compress_ctx, exec_ctx| { + let sample = sample_primitive_one_percent(data.array_as_primitive(), exec_ctx)?; + estimate_float_quant_sample(&sample) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let source = data.array_as_primitive(); + let primitive = normalize_null_values(source, exec_ctx)?; + let Some(analysis) = analyze_float_quant(primitive.as_view()) else { + return Ok(source.array().clone()); + }; + encode_float_quant(primitive.as_view(), analysis) + } +} + +fn estimate_float_quant_sample(sample: &PrimitiveArray) -> VortexResult { + let Some(analysis) = analyze_float_quant(sample.as_view()) else { + return Ok(EstimateVerdict::Skip); + }; + // A constant sample does not prove that the full array is constant. + if analysis.primary_bit_width == 0 && analysis.secondary_bit_width == 0 { + return Ok(EstimateVerdict::Skip); + } + + let before_nbytes = sample.nbytes(); + let after_nbytes = estimate_float_quant_nbytes(sample, analysis)?; + if after_nbytes == 0 || after_nbytes >= before_nbytes { + return Ok(EstimateVerdict::Skip); + } + + let adjusted_ratio = before_nbytes as f64 / after_nbytes as f64 / SELECTION_COST_FACTOR; + if adjusted_ratio <= 1.0 { + return Ok(EstimateVerdict::Skip); + } + + Ok(EstimateVerdict::Ratio(adjusted_ratio)) +} + +fn estimate_float_quant_nbytes( + sample: &PrimitiveArray, + analysis: FloatQuantAnalysis, +) -> VortexResult { + let packed_chunks = u64::try_from(sample.len().div_ceil(FL_CHUNK_SIZE))?; + let bytes_per_bit = u64::try_from(FL_CHUNK_SIZE / 8)?; + let packed_bit_width = + u64::from(analysis.primary_bit_width) + u64::from(analysis.secondary_bit_width); + let packed_nbytes = packed_chunks * bytes_per_bit * packed_bit_width; + let validity_nbytes = validity_to_child(&sample.validity()?, sample.len()) + .map(|validity| validity.nbytes()) + .unwrap_or(0); + Ok(packed_nbytes + validity_nbytes) +} + +fn encode_float_quant( + primitive: vortex_array::ArrayView<'_, Primitive>, + analysis: FloatQuantAnalysis, +) -> VortexResult { + // The ordered transform complements the low bits of negative values. Decode complements the + // secondary bits again, so both signs store the original low bits. + let (primary_packed, secondary_packed, latent_ptype, reference) = match primitive.ptype() { + PType::F16 => { + let primary_min = u16::try_from(analysis.primary_min)?; + let values = primitive.as_slice::(); + let (primary, secondary) = if analysis.secondary_bit_width == 0 { + ( + bitpack_primitive_map(values, analysis.primary_bit_width, |value| { + (ordered_u16(value.to_bits()) >> analysis.k) - primary_min + }), + None, + ) + } else { + let low_mask = (1_u16 << analysis.k) - 1; + let (primary, secondary) = bitpack_primitive_map_pair( + values, + analysis.primary_bit_width, + analysis.secondary_bit_width, + |value| { + let bits = value.to_bits(); + ( + (ordered_u16(bits) >> analysis.k) - primary_min, + bits & low_mask, + ) + }, + ); + (primary, Some(secondary)) + }; + ( + primary.into_byte_buffer(), + secondary.map(|packed| packed.into_byte_buffer()), + PType::U16, + Scalar::from(primary_min), + ) + } + PType::F32 => { + let primary_min = u32::try_from(analysis.primary_min)?; + let values = primitive.as_slice::(); + let (primary, secondary) = if analysis.secondary_bit_width == 0 { + ( + bitpack_primitive_map(values, analysis.primary_bit_width, |value| { + (ordered_u32(value.to_bits()) >> analysis.k) - primary_min + }), + None, + ) + } else { + let low_mask = (1_u32 << analysis.k) - 1; + let (primary, secondary) = bitpack_primitive_map_pair( + values, + analysis.primary_bit_width, + analysis.secondary_bit_width, + |value| { + let bits = value.to_bits(); + ( + (ordered_u32(bits) >> analysis.k) - primary_min, + bits & low_mask, + ) + }, + ); + (primary, Some(secondary)) + }; + ( + primary.into_byte_buffer(), + secondary.map(|packed| packed.into_byte_buffer()), + PType::U32, + Scalar::from(primary_min), + ) + } + PType::F64 => { + let values = primitive.as_slice::(); + let (primary, secondary) = if analysis.secondary_bit_width == 0 { + ( + bitpack_primitive_map(values, analysis.primary_bit_width, |value| { + (ordered_u64(value.to_bits()) >> analysis.k) - analysis.primary_min + }), + None, + ) + } else { + let low_mask = (1_u64 << analysis.k) - 1; + let (primary, secondary) = bitpack_primitive_map_pair( + values, + analysis.primary_bit_width, + analysis.secondary_bit_width, + |value| { + let bits = value.to_bits(); + ( + (ordered_u64(bits) >> analysis.k) - analysis.primary_min, + bits & low_mask, + ) + }, + ); + (primary, Some(secondary)) + }; + ( + primary.into_byte_buffer(), + secondary.map(|packed| packed.into_byte_buffer()), + PType::U64, + Scalar::from(analysis.primary_min), + ) + } + _ => unreachable!(), + }; + let compressed_primary = BitPacked::try_new( + BufferHandle::new_host(primary_packed), + latent_ptype, + primitive.validity()?, + None, + analysis.primary_bit_width, + primitive.len(), + 0, + )? + .into_array(); + let compressed_primary = FoR::try_new(compressed_primary, reference)?.into_array(); + let compressed_secondary = secondary_packed + .map(|packed| { + BitPacked::try_new( + BufferHandle::new_host(packed), + latent_ptype, + Validity::NonNullable, + None, + analysis.secondary_bit_width, + primitive.len(), + 0, + ) + .map(IntoArray::into_array) + }) + .transpose()?; + Ok(FloatQuant::try_new( + compressed_primary, + compressed_secondary, + primitive.ptype(), + analysis.k, + )? + .into_array()) +} + +#[inline] +fn ordered_u16(bits: u16) -> u16 { + if bits & (1_u16 << 15) == 0 { + bits ^ (1_u16 << 15) + } else { + !bits + } +} + +#[inline] +fn ordered_u32(bits: u32) -> u32 { + if bits & (1_u32 << 31) == 0 { + bits ^ (1_u32 << 31) + } else { + !bits + } +} + +#[inline] +fn ordered_u64(bits: u64) -> u64 { + if bits & (1_u64 << 63) == 0 { + bits ^ (1_u64 << 63) + } else { + !bits + } +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::half::f16; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::analyze_float_quant; + use super::encode_float_quant; + use super::estimate_float_quant_nbytes; + use super::estimate_float_quant_sample; + + #[test] + fn constant_sample_is_not_evidence_for_float_quant() -> VortexResult<()> { + let sample = PrimitiveArray::from_iter(vec![1.0_f64; 1_024]); + + assert!(matches!( + estimate_float_quant_sample(&sample)?, + EstimateVerdict::Skip + )); + Ok(()) + } + + #[test] + fn estimated_nbytes_matches_encoded_tree() -> VortexResult<()> { + let f16_values = PrimitiveArray::from_iter((0..2_050).map(|index| { + let high_bits = u16::try_from(index).unwrap_or_default().wrapping_mul(17) & 0x03f0; + let low_bit = u16::from(index % 10 == 0); + f16::from_bits(0x3c00 | high_bits | low_bit) + })); + let f32_values = PrimitiveArray::from_option_iter((0..2_050).map(|index| { + let high_bits = (index as u32).wrapping_mul(7_919) & 0x007f_ff00; + let low_bit = u32::from(index % 10 == 0); + (index % 17 != 0).then_some(f32::from_bits(0x3f80_0000 | high_bits | low_bit)) + })); + let f64_values = PrimitiveArray::from_iter((0..2_050).map(|index| { + let high_bits = ((index as u64).wrapping_mul(7_919) << 29) & 0x000f_ffff_ffff_ff00; + let low_bit = u64::from(index % 10 == 0); + f64::from_bits(0x3ff0_0000_0000_0000 | high_bits | low_bit) + })); + + for values in [f16_values, f32_values, f64_values] { + let analysis = analyze_float_quant(values.as_view()) + .ok_or_else(|| vortex_err!("FloatQuant test input did not produce an analysis"))?; + let expected = encode_float_quant(values.as_view(), analysis)?.nbytes(); + assert_eq!(estimate_float_quant_nbytes(&values, analysis)?, expected); + } + Ok(()) + } + + #[test] + fn near_miss_sample_is_rejected() -> VortexResult<()> { + let f32_values = PrimitiveArray::from_iter((0..2_048).map(|index| { + let scrambled = (index as u32).wrapping_mul(2_654_435_761); + let sign = (scrambled & 1) << 31; + let exponent = ((scrambled >> 1) % 254 + 1) << 23; + let mantissa = scrambled & 0x007f_fffc; + f32::from_bits(sign | exponent | mantissa) + })); + let f64_values = PrimitiveArray::from_iter((0..2_048).map(|index| { + let scrambled = (index as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15); + let sign = (scrambled & 1) << 63; + let exponent = ((scrambled >> 1) % 2_046 + 1) << 52; + let mantissa = scrambled & 0x000f_ffff_ffff_fffc; + f64::from_bits(sign | exponent | mantissa) + })); + + for values in [f32_values, f64_values] { + assert!(analyze_float_quant(values.as_view()).is_some()); + assert!(matches!( + estimate_float_quant_sample(&values)?, + EstimateVerdict::Skip + )); + } + Ok(()) + } + + #[test] + fn nonzero_secondary_round_trips_negative_values() -> VortexResult<()> { + let values = PrimitiveArray::from_iter((0_u32..2_050).flat_map(|index| { + let high_bits = index.wrapping_mul(7_919) & 0x007f_ff00; + let low_bits = index % 7; + let positive = f32::from_bits(0x3f80_0000 | high_bits | low_bits); + [positive, -positive] + })); + let analysis = analyze_float_quant(values.as_view()) + .ok_or_else(|| vortex_err!("FloatQuant test input did not produce an analysis"))?; + assert!(analysis.secondary_bit_width > 0); + let encoded = encode_float_quant(values.as_view(), analysis)?; + assert_arrays_eq!( + encoded, + values.into_array(), + &mut array_session().create_execution_ctx() + ); + Ok(()) + } +} diff --git a/vortex-btrblocks/src/schemes/float/mod.rs b/vortex-btrblocks/src/schemes/float/mod.rs index 1301184ac0c..e2f8aecb398 100644 --- a/vortex-btrblocks/src/schemes/float/mod.rs +++ b/vortex-btrblocks/src/schemes/float/mod.rs @@ -5,6 +5,8 @@ mod alp; mod alprd; +mod float_quant; +mod ordered_block_residual; mod rle; mod sparse; @@ -13,6 +15,8 @@ mod pco; pub use alp::ALPScheme; pub use alprd::ALPRDScheme; +pub use float_quant::FloatQuantScheme; +pub use ordered_block_residual::OrderedBlockResidualScheme; #[cfg(feature = "pco")] pub use pco::PcoScheme; pub use rle::FloatRLEScheme; diff --git a/vortex-btrblocks/src/schemes/float/ordered_block_residual.rs b/vortex-btrblocks/src/schemes/float/ordered_block_residual.rs new file mode 100644 index 00000000000..4fbfe07aed8 --- /dev/null +++ b/vortex-btrblocks/src/schemes/float/ordered_block_residual.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Ordered float bits with one block-local reference and packed residuals. + +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::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::match_each_float_ptype; +use vortex_block_residual::BlockResidual; +use vortex_block_residual::OrderedFloat; +use vortex_block_residual::OrderedFloatArraySlotsExt; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::normalize_null_values; +use crate::schemes::integer::patch_adjusted_estimate_nbytes; +use crate::schemes::sample_primitive_blocks; + +const BLOCK_LEN: usize = 1024; +const ESTIMATE_BLOCKS: usize = 8; +const MIN_COMPRESSION_RATIO: f64 = 1.05; +const DECODE_COST_FACTOR: f64 = 1.02; + +/// Compress floats as block-local residuals of ordered IEEE bits. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct OrderedBlockResidualScheme; + +impl Scheme for OrderedBlockResidualScheme { + fn scheme_name(&self) -> &'static str { + "vortex.float.ordered_block_residual" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_float() + } + + fn produced_encodings(&self) -> Vec { + vec![OrderedFloat.id(), BlockResidual.id()] + } + + fn num_children(&self) -> usize { + 1 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + if compress_ctx.finished_cascading() { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, data, _best_so_far, _compress_ctx, exec_ctx| { + let sample = locality_sample(data.array_as_primitive(), exec_ctx)?; + let sample = normalize_null_values(sample.as_view(), exec_ctx)?; + let before_nbytes = sample.nbytes(); + let estimate = OrderedFloat::estimate_block_residual(sample.as_view())?; + let after_nbytes = patch_adjusted_estimate_nbytes(estimate, sample.len()); + if after_nbytes == 0 { + return Ok(EstimateVerdict::Skip); + } + + let ratio = before_nbytes as f64 / after_nbytes as f64; + if ratio < MIN_COMPRESSION_RATIO { + return Ok(EstimateVerdict::Skip); + } + let adjusted_ratio = ratio / DECODE_COST_FACTOR; + if adjusted_ratio < MIN_COMPRESSION_RATIO { + return Ok(EstimateVerdict::Skip); + } + Ok(EstimateVerdict::Ratio(adjusted_ratio)) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primitive = normalize_null_values(data.array_as_primitive(), exec_ctx)?; + let ordered = OrderedFloat::from_primitive(primitive.as_view())?; + let ordered_values = ordered.encoded().as_::(); + let residuals = BlockResidual::from_primitive(ordered_values)?; + Ok(OrderedFloat::try_new(residuals.into_array(), primitive.ptype())?.into_array()) + } +} + +fn locality_sample( + primitive: vortex_array::ArrayView<'_, Primitive>, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = primitive + .validity()? + .execute_mask(primitive.len(), exec_ctx)?; + let full_blocks = primitive.len() / BLOCK_LEN; + + if full_blocks <= ESTIMATE_BLOCKS { + return primitive + .array() + .clone() + .execute::(exec_ctx); + } + + let sample_blocks = ESTIMATE_BLOCKS.min(full_blocks); + Ok(match_each_float_ptype!(primitive.ptype(), |T| { + sample_primitive_blocks( + primitive.as_slice::(), + validity.all_true(), + |index| validity.value(index), + full_blocks, + sample_blocks, + BLOCK_LEN, + ) + })) +} diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 4c7a5f85fa6..7237bb31ae5 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -3,26 +3,78 @@ //! Tests to verify that each float compression scheme produces the expected encoding. +use std::f64::consts::TAU; use std::sync::LazyLock; use vortex_alp::ALP; +use vortex_array::ArrayEq; +use vortex_array::EqMode; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Constant; use vortex_array::arrays::Dict; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::PrimitiveBuilder; use vortex_array::dtype::Nullability; +use vortex_array::dtype::half::f16; use vortex_array::validity::Validity; +use vortex_block_residual::BlockResidual; +use vortex_block_residual::OrderedFloat; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_float_quant::FloatQuant; +use vortex_float_quant::FloatQuantArraySlotsExt; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +#[cfg(feature = "unstable_encodings")] +use crate::BtrBlocksCompressorBuilder; +use crate::CascadingCompressor; +#[cfg(feature = "unstable_encodings")] +use crate::SchemeExt; +use crate::schemes::float::FloatQuantScheme; +#[cfg(feature = "unstable_encodings")] +use crate::schemes::integer::DeltaScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +#[test] +fn test_quantized_f16_uses_float_quant() -> VortexResult<()> { + let values = (0_u16..16_384) + .map(|index| f16::from_bits(0x3c00 | (index.wrapping_mul(7_919) & 0x03f0))) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_f16_secondary_uses_float_quant() -> VortexResult<()> { + let values = (0_u16..16_384) + .map(|index| { + let high_mantissa = index.wrapping_mul(7_919) & 0x03f0; + f16::from_bits(0x3c00 | high_mantissa | (index & 1)) + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressor = CascadingCompressor::new(vec![&FloatQuantScheme]); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!(compressed.is::()); + assert!(compressed.as_::().secondary().is_some()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + #[test] fn test_constant_compressed() -> VortexResult<()> { let values: Vec = vec![42.5; 100]; @@ -66,8 +118,202 @@ fn test_null_dominated_compressed() -> VortexResult<()> { builder.append_nulls(95); let array = builder.finish_into_primitive(); let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; - // Verify the compressed array preserves values. - assert_eq!(compressed.len(), 100); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = btr.compress(&array.clone().into_array(), &mut ctx)?; + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_widened_f32_uses_float_quant() -> VortexResult<()> { + let values = (0u32..16_384) + .map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + f64::from(f32::from_bits(0x3f80_0000 | mantissa)) + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + assert!(compressed.is::()); + assert!(compressed.as_::().secondary().is_none()); + Ok(()) +} + +#[test] +fn test_nonzero_secondary_uses_float_quant() -> VortexResult<()> { + let values = (0u32..65_536) + .map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + let value = f64::from(f32::from_bits(0x3f80_0000 | mantissa)); + if index % 10 == 0 { + f64::from_bits(value.to_bits() | 1) + } else { + value + } + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + + assert!(compressed.is::()); + let float_quant = compressed.as_::(); + let secondary = float_quant + .secondary() + .ok_or_else(|| vortex_error::vortex_err!("missing nonzero FloatQuant secondary"))? + .as_::(); + assert_eq!(secondary.bit_width(), 1); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_float_quant_ignores_null_payloads() -> VortexResult<()> { + let values = (0u32..16_384) + .map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + f64::from(f32::from_bits(0x3f80_0000 | mantissa)) + }) + .collect::>(); + let validity = Validity::from_iter((0..values.len()).map(|index| index % 17 != 0)); + let mut alternate = values.clone(); + for index in (0..alternate.len()).step_by(17) { + alternate[index] = f64::from_bits(0x3ff0_0000_0000_0001 + index as u64); + } + let first = PrimitiveArray::new(Buffer::copy_from(&values), validity.clone()).into_array(); + let second = PrimitiveArray::new(Buffer::copy_from(&alternate), validity).into_array(); + let compressor = BtrBlocksCompressor::default(); + let first = compressor.compress(&first, &mut SESSION.create_execution_ctx())?; + let second = compressor.compress(&second, &mut SESSION.create_execution_ctx())?; + + assert!(first.is::()); + assert!(second.is::()); + assert!(first.array_eq(&second, EqMode::Value)); + Ok(()) +} + +#[test] +fn test_f32_does_not_use_float_quant() -> VortexResult<()> { + let values = (0u32..16_384) + .map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + f32::from_bits(0x3f80_0000 | mantissa) + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + assert!(!compressed.is::()); + Ok(()) +} + +#[test] +fn test_quantized_f32_uses_float_quant() -> VortexResult<()> { + let values = (0_u32..65_536) + .map(|index| { + let mantissa = (index.wrapping_mul(7_919) & 0x7fff) << 8; + f32::from_bits(0x3f80_0000 | mantissa) + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + + assert!(compressed.is::()); + assert!(compressed.as_::().secondary().is_none()); + Ok(()) +} + +#[test] +fn test_repeated_f64_prefers_existing_scheme() -> VortexResult<()> { + let values = (0u32..16_384) + .map(|index| f64::from(index % 8)) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + assert!(!compressed.is::()); + Ok(()) +} + +#[test] +fn test_random_walk_uses_ordered_block_residual() -> VortexResult<()> { + fn uniform(state: &mut u64) -> f64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + ((*state >> 11) as f64 + 0.5) / (1_u64 << 53) as f64 + } + + let mut state = 0x4d59_5df4_d0f3_3173_u64; + let mut value = 0.0_f64; + let values = (0usize..65_536) + .map(|_| { + let radius = (-2.0 * uniform(&mut state).ln()).sqrt(); + let normal = radius * (TAU * uniform(&mut state)).cos(); + value += normal * 0.01; + value + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + assert!(compressed.is::()); + assert!(compressed.children()[0].is::()); + Ok(()) +} + +#[test] +fn test_f32_random_walk_uses_ordered_block_residual() -> VortexResult<()> { + let values = (0_u32..65_536) + .map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(7_919) % 1_024; + f32::from_bits(0x3f80_0000 + (block * 0x1_0000) + residual) + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values).into_array(); + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; + + assert!( + compressed.is::(), + "expected OrderedFloat, got tree:\n{}", + compressed.display_tree() + ); + assert!(compressed.children()[0].is::()); + Ok(()) +} + +#[test] +fn test_block_residual_composes_with_alp() -> VortexResult<()> { + let values = (0..65_536_usize).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + (block * 1_000_000 + residual) as f64 + }); + let array = PrimitiveArray::from_iter(values).into_array(); + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + let compressed = compressor.compress(&array, &mut SESSION.create_execution_ctx())?; + + assert!( + compressed.is::(), + "expected ALP, got tree:\n{}", + compressed.display_tree() + ); + assert!( + compressed + .children() + .iter() + .any(|child| child.is::()), + "expected a BlockResidual child:\n{}", + compressed.display_tree() + ); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 2d0a04542cc..2d6e2f9900c 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -42,15 +42,10 @@ fn test_compress() -> VortexResult<()> { } let array = values.into_array(); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; + let compressed = + BtrBlocksCompressor::default().compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); - - let display = compressed - .display_as(DisplayOptions::MetadataOnly) - .to_string() - .to_lowercase(); - assert_eq!(display, "vortex.dict(f32, len=1024)"); + assert_arrays_eq!(compressed, array, &mut SESSION.create_execution_ctx()); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/block_residual.rs b/vortex-btrblocks/src/schemes/integer/block_residual.rs new file mode 100644 index 00000000000..c3780dd75f9 --- /dev/null +++ b/vortex-btrblocks/src/schemes/integer/block_residual.rs @@ -0,0 +1,185 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Integer compression with one block-local reference and packed residuals. + +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::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::match_each_integer_ptype; +use vortex_block_residual::BlockResidual; +use vortex_block_residual::BlockResidualEstimate; +use vortex_compressor::builtins::BinaryDictScheme; +use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; + +use super::ZigZagScheme; +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::SchemeExt; +use crate::normalize_null_values; +use crate::schemes::sample_primitive_blocks; + +const BLOCK_LEN: usize = 1024; +const ESTIMATE_BLOCKS: usize = 8; +const MIN_COMPRESSION_RATIO: f64 = 1.05; +// The weakest measured 8-bit gain increased file-access latency by 29 to 55 percent. +// Require about 12 percent estimated savings for 8-bit values. +const EIGHT_BIT_ACCESS_COST_FACTOR: f64 = 1.12; +// The 25-percent patch sweep needs about 80 cost bits per patch to reject its slow tree. +// This quadratic slope reaches that cost at 25 percent while it keeps sparse patches cheap. +const PATCH_DENSITY_COST_BITS: u64 = 320; + +/// Compress integers with one reference and packed residuals per 1,024-value block. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct BlockResidualScheme; + +impl Scheme for BlockResidualScheme { + fn scheme_name(&self) -> &'static str { + "vortex.int.block_residual" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() + } + + fn produced_encodings(&self) -> Vec { + vec![BlockResidual.id()] + } + + fn ancestor_exclusions(&self) -> Vec { + vec![ + AncestorExclusion { + ancestor: StringDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: BinaryDictScheme.id(), + children: ChildSelection::One(1), + }, + AncestorExclusion { + ancestor: ZigZagScheme.id(), + children: ChildSelection::One(0), + }, + ] + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // A single block cannot amortize a block-local reference against FoR. + if data.array().len() <= BLOCK_LEN + || compress_ctx.finished_cascading() + || compress_ctx.is_sample() + { + return CompressionEstimate::Verdict(EstimateVerdict::Skip); + } + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + |_compressor, data, _best_so_far, _compress_ctx, exec_ctx| { + let sample = locality_sample(data.array_as_primitive(), exec_ctx)?; + let sample = normalize_null_values(sample.as_view(), exec_ctx)?; + let before_nbytes = sample.nbytes(); + let estimate = BlockResidual::estimate_primitive(sample.as_view())?; + let after_nbytes = patch_adjusted_estimate_nbytes(estimate, sample.len()); + if after_nbytes == 0 { + return Ok(EstimateVerdict::Skip); + } + + let ratio = before_nbytes as f64 / after_nbytes as f64; + if ratio < MIN_COMPRESSION_RATIO { + return Ok(EstimateVerdict::Skip); + } + let adjusted_ratio = if sample.ptype().byte_width() == 1 { + ratio / EIGHT_BIT_ACCESS_COST_FACTOR + } else { + ratio + }; + Ok(EstimateVerdict::Ratio(adjusted_ratio)) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primitive = normalize_null_values(data.array_as_primitive(), exec_ctx)?; + Ok(BlockResidual::from_primitive(primitive.as_view())?.into_array()) + } +} + +pub(crate) fn patch_adjusted_estimate_nbytes(estimate: BlockResidualEstimate, len: usize) -> u64 { + estimate + .nbytes() + .saturating_add(patch_density_cost_bytes(len, estimate.patch_count())) +} + +fn patch_density_cost_bytes(len: usize, patch_count: usize) -> u64 { + if len == 0 { + return 0; + } + let patch_count = u64::try_from(patch_count).unwrap_or(u64::MAX); + patch_count + .saturating_mul(patch_count) + .saturating_mul(PATCH_DENSITY_COST_BITS) + .div_ceil(u64::try_from(len).unwrap_or(u64::MAX).saturating_mul(8)) +} + +fn locality_sample( + primitive: vortex_array::ArrayView<'_, Primitive>, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = primitive + .validity()? + .execute_mask(primitive.len(), exec_ctx)?; + let full_blocks = primitive.len() / BLOCK_LEN; + + if full_blocks <= ESTIMATE_BLOCKS { + return primitive + .array() + .clone() + .execute::(exec_ctx); + } + + let sample_blocks = ESTIMATE_BLOCKS.min(full_blocks); + Ok(match_each_integer_ptype!(primitive.ptype(), |T| { + sample_primitive_blocks( + primitive.as_slice::(), + validity.all_true(), + |index| validity.value(index), + full_blocks, + sample_blocks, + BLOCK_LEN, + ) + })) +} + +#[cfg(test)] +mod tests { + use super::patch_density_cost_bytes; + + #[test] + fn patch_density_cost_is_nonlinear() { + assert_eq!(patch_density_cost_bytes(1_024, 0), 0); + assert_eq!(patch_density_cost_bytes(1_024, 102), 407); + assert_eq!(patch_density_cost_bytes(1_024, 256), 2_560); + } +} diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 3aae2ae5601..43ff8bbb6cd 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -4,6 +4,7 @@ //! Integer compression schemes. mod bitpacking; +mod block_residual; #[cfg(feature = "unstable_encodings")] mod delta; mod for_; @@ -17,6 +18,8 @@ mod zigzag; mod pco; pub use bitpacking::BitPackingScheme; +pub use block_residual::BlockResidualScheme; +pub(crate) use block_residual::patch_adjusted_estimate_nbytes; #[cfg(feature = "unstable_encodings")] pub use delta::DeltaScheme; pub use for_::FoRScheme; diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index e4227a472ec..f94ed226a4a 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -9,6 +9,8 @@ use std::sync::LazyLock; use rand::Rng; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::ArrayEq; +use vortex_array::EqMode; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Constant; @@ -18,6 +20,7 @@ use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; use vortex_array::validity::Validity; +use vortex_block_residual::BlockResidual; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_fastlanes::BitPacked; @@ -28,6 +31,12 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; +#[cfg(feature = "unstable_encodings")] +use crate::BtrBlocksCompressorBuilder; +#[cfg(feature = "unstable_encodings")] +use crate::SchemeExt; +#[cfg(feature = "unstable_encodings")] +use crate::schemes::integer::DeltaScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] @@ -50,6 +59,264 @@ fn test_for_compressed() -> VortexResult<()> { Ok(()) } +#[test] +fn test_block_residual_compressed() -> VortexResult<()> { + let values = (0..8_192) + .map(|index| { + let block = index / 1_024; + let residual = (index * 2_654_435_761_usize) % 1_024; + (block as i64 - 4) * 1_000_000_000_000 + residual as i64 + }) + .collect::>(); + let array = PrimitiveArray::from_iter(values); + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + let compressed = + compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!( + compressed.is::(), + "expected BlockResidual, got tree:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +#[test] +fn test_block_residual_ignores_null_payloads() -> VortexResult<()> { + let values = (0usize..8_192) + .map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + (block as i64 - 4) * 1_000_000_000_000 + residual as i64 + }) + .collect::>(); + let validity = Validity::from_iter((0..values.len()).map(|index| index % 17 != 0)); + let mut alternate = values.clone(); + for index in (0..alternate.len()).step_by(17) { + alternate[index] = i64::MAX - index as i64; + } + let first = PrimitiveArray::new(Buffer::copy_from(&values), validity.clone()).into_array(); + let second = PrimitiveArray::new(Buffer::copy_from(&alternate), validity).into_array(); + let compressor = BtrBlocksCompressor::default(); + let first = compressor.compress(&first, &mut SESSION.create_execution_ctx())?; + let second = compressor.compress(&second, &mut SESSION.create_execution_ctx())?; + + assert!(first.is::()); + assert!(second.is::()); + assert!(first.array_eq(&second, EqMode::Value)); + Ok(()) +} + +#[test] +fn test_block_residual_compresses_16_bit_integers() -> VortexResult<()> { + let signed_values = (0..8_192) + .map(|index| { + let block = index / 1_024; + let residual = (index * 2_654_435_761_usize) % 32; + Ok(i16::try_from(block * 1_000 + residual)?) + }) + .collect::>>()?; + let unsigned_values = signed_values + .iter() + .copied() + .map(u16::try_from) + .collect::, _>>()?; + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + + for array in [ + PrimitiveArray::from_iter(signed_values), + PrimitiveArray::from_iter(unsigned_values), + ] { + let compressed = + compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + contains_block_residual(&compressed), + "BlockResidual must encode this 16-bit input:\n{}", + compressed.display_tree() + ); + } + Ok(()) +} + +#[test] +fn test_block_residual_compresses_8_bit_integers() -> VortexResult<()> { + let unsigned_values = (0..16_384) + .map(|index| { + let block = (index / 1_024) % 32; + let residual = (index * 2_654_435_761_usize) % 8; + u8::try_from(block * 8 + residual) + }) + .collect::, _>>()?; + let signed_values = unsigned_values + .iter() + .copied() + .map(|value| i8::from_le_bytes([value])) + .collect::>(); + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + + for array in [ + PrimitiveArray::from_iter(signed_values), + PrimitiveArray::from_iter(unsigned_values), + ] { + let compressed = + compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + contains_block_residual(&compressed), + "BlockResidual must encode this 8-bit input:\n{}", + compressed.display_tree() + ); + } + Ok(()) +} + +#[test] +fn test_block_residual_rejects_weak_8_bit_gain() -> VortexResult<()> { + let unsigned_values = (0..16_384) + .map(|index| { + let block = (index / 1_024) % 2; + let residual = (index * 2_654_435_761_usize) % 128; + u8::try_from(block * 128 + residual) + }) + .collect::, _>>()?; + let signed_values = unsigned_values + .iter() + .copied() + .map(|value| i8::from_le_bytes([value])) + .collect::>(); + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + + for array in [ + PrimitiveArray::from_iter(signed_values), + PrimitiveArray::from_iter(unsigned_values), + ] { + let compressed = + compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + !contains_block_residual(&compressed), + "BlockResidual must reject this weak 8-bit gain:\n{}", + compressed.display_tree() + ); + } + Ok(()) +} + +#[test] +fn test_block_residual_rejects_uniform_8_bit_integers() -> VortexResult<()> { + let unsigned_values = (0..16_384) + .map(|index| u8::try_from((index * 2_654_435_761_usize) % 256)) + .collect::, _>>()?; + let signed_values = unsigned_values + .iter() + .copied() + .map(|value| i8::from_le_bytes([value])) + .collect::>(); + #[cfg(not(feature = "unstable_encodings"))] + let compressor = BtrBlocksCompressor::default(); + #[cfg(feature = "unstable_encodings")] + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DeltaScheme::default().id()]) + .build(); + + for array in [ + PrimitiveArray::from_iter(signed_values), + PrimitiveArray::from_iter(unsigned_values), + ] { + let compressed = + compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + assert!( + !contains_block_residual(&compressed), + "BlockResidual must reject this uniform 8-bit input:\n{}", + compressed.display_tree() + ); + } + Ok(()) +} + +#[test] +fn test_block_residual_rejects_dense_patches() -> VortexResult<()> { + let values = (0..8_192_u32).map(|index| if index % 4 == 0 { u32::MAX - index } else { 42 }); + let array = PrimitiveArray::from_iter(values); + let compressed = BtrBlocksCompressor::default() + .compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!( + !contains_block_residual(&compressed), + "dense patches must not select BlockResidual:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +#[test] +fn test_block_residual_composes_with_sparse() -> VortexResult<()> { + let values = (0..65_536_usize).map(|index| { + if index % 16 == 0 { + let value_index = index / 16; + let block = value_index / 1_024; + let residual = value_index.wrapping_mul(2_654_435_761) % 1_024; + block as u64 * 1_000_000_000_000 + residual as u64 + } else { + 42 + } + }); + let array = PrimitiveArray::from_iter(values); + let compressed = BtrBlocksCompressor::default() + .compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!(compressed.is::()); + assert!( + contains_block_residual(&compressed), + "expected a BlockResidual child:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +#[test] +fn test_block_residual_composes_with_runend() -> VortexResult<()> { + let values = (0..65_536_usize).map(|index| { + let value_index = index / 16; + let block = value_index / 1_024; + let residual = value_index.wrapping_mul(2_654_435_761) % 1_024; + block as u64 * 1_000_000_000_000 + residual as u64 + }); + let array = PrimitiveArray::from_iter(values); + let compressed = BtrBlocksCompressor::default() + .compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; + + assert!(compressed.is::()); + assert!( + contains_block_residual(&compressed), + "expected a BlockResidual child:\n{}", + compressed.display_tree() + ); + Ok(()) +} + +fn contains_block_residual(array: &vortex_array::ArrayRef) -> bool { + array.is::() || array.children().iter().any(contains_block_residual) +} + #[test] fn test_bitpacking_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| i % 16).collect(); diff --git a/vortex-btrblocks/src/schemes/mod.rs b/vortex-btrblocks/src/schemes/mod.rs index a0e9b042a66..2aa2b5ad03f 100644 --- a/vortex-btrblocks/src/schemes/mod.rs +++ b/vortex-btrblocks/src/schemes/mod.rs @@ -13,6 +13,14 @@ pub mod temporal; pub(crate) mod patches; +use std::ops::Range; + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_native_ptype; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -21,9 +29,96 @@ use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::SchemeExt; +use vortex_error::VortexResult; +use crate::normalize_null_values; use crate::schemes::integer::SparseScheme; +const SAMPLE_BLOCK_LEN: usize = 64; +const MIN_SAMPLE_BLOCKS: usize = 16; +const SAMPLE_BLOCK_MULTIPLE: usize = 16; + +fn sample_primitive_one_percent( + primitive: ArrayView<'_, Primitive>, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let sample_blocks = (primitive.len() / 100 / SAMPLE_BLOCK_LEN) + .next_multiple_of(SAMPLE_BLOCK_MULTIPLE) + .max(MIN_SAMPLE_BLOCKS); + let sample_len = sample_blocks * SAMPLE_BLOCK_LEN; + if primitive.len() <= sample_len { + return normalize_null_values(primitive, exec_ctx); + } + + let ranges = one_percent_sample_ranges(primitive.len(), sample_blocks); + let validity = primitive.validity()?; + if validity.definitely_no_nulls() { + return Ok(match_each_native_ptype!(primitive.ptype(), |T| { + let values = primitive.as_slice::(); + let mut sample = Vec::with_capacity(sample_len); + for range in ranges { + sample.extend_from_slice(&values[range]); + } + PrimitiveArray::from_iter(sample) + })); + } + + let validity = validity.execute_mask(primitive.len(), exec_ctx)?; + Ok(match_each_native_ptype!(primitive.ptype(), |T| { + let values = primitive.as_slice::(); + PrimitiveArray::from_option_iter( + ranges + .into_iter() + .flatten() + .map(|index| validity.value(index).then_some(values[index])), + ) + })) +} + +fn one_percent_sample_ranges(len: usize, sample_blocks: usize) -> Vec> { + let partition_len = len / sample_blocks; + let long_partitions = len % sample_blocks; + let mut partition_start = 0; + (0..sample_blocks) + .map(|partition_index| { + let current_partition_len = + partition_len + usize::from(partition_index < long_partitions); + let start = partition_start + (current_partition_len - SAMPLE_BLOCK_LEN) / 2; + partition_start += current_partition_len; + start..start + SAMPLE_BLOCK_LEN + }) + .collect() +} + +fn sample_primitive_blocks( + values: &[T], + all_valid: bool, + is_valid: impl Fn(usize) -> bool, + full_blocks: usize, + sample_blocks: usize, + block_len: usize, +) -> PrimitiveArray { + if all_valid { + let mut sample = Vec::with_capacity(sample_blocks * block_len); + for sample_index in 0..sample_blocks { + let block_index = sample_index * full_blocks / sample_blocks; + let start = block_index * block_len; + sample.extend_from_slice(&values[start..start + block_len]); + } + PrimitiveArray::from_iter(sample) + } else { + let mut sample = Vec::with_capacity(sample_blocks * block_len); + for sample_index in 0..sample_blocks { + let block_index = sample_index * full_blocks / sample_blocks; + let start = block_index * block_len; + sample.extend( + (start..start + block_len).map(|index| is_valid(index).then_some(values[index])), + ); + } + PrimitiveArray::from_option_iter(sample) + } +} + /// Shared descendant exclusion rules for RLE schemes. /// /// RLE indices (child 1) and offsets (child 2) are monotonically increasing positions with all diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index fea1ae10b47..60aa85a429b 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -219,7 +219,7 @@ fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { /// Q6-style predicate over the quantity column: `l_quantity < 24`. /// -/// The column compresses to `decimal_byte_parts -> dict -> bitpacked/sequence`. +/// The column compresses to `decimal_byte_parts -> dict -> bitpacked/primitive`. fn quantity_predicate(column: ArrayRef, len: usize) -> VortexResult { let cutoff = Scalar::decimal( DecimalValue::I128(2400), @@ -361,6 +361,11 @@ fn trace_scan_like_on_compressed_comment() -> VortexResult<()> { insta::assert_snapshot!(executed.trace.to_string(), @" execute_until target=AnyCanonical root=vortex.like(bool, len=4096) iter 0 current=vortex.like(bool, len=4096) builder_active=false + execute_until target=AnyCanonical root=vortex.block_residual(u16, len=4097) + iter 0 current=vortex.block_residual(u16, len=4097) builder_active=false + Done array=vortex.primitive(u16, len=4097) + iter 1 current=vortex.primitive(u16, len=4097) builder_active=false + return output=vortex.primitive(u16, len=4097) child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.like(bool, len=4096) child=vortex.fsst(utf8, len=4096) -> vortex.bool(bool, len=4096) iter 1 current=vortex.bool(bool, len=4096) builder_active=false return output=vortex.bool(bool, len=4096) diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap index 4800aae289d..a29fba7b5a4 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap @@ -3,7 +3,5 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: u64, len=16384, nbytes=131072 -root: fastlanes.for(u64, len=16384) nbytes=49152 - metadata: reference: 1700000001036u64 - encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 +root: vortex.block_residual(u64, len=16384) nbytes=41324 + metadata: blocks: 16, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap index c9554add05a..e5f7bbb6eee 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: list(i32), len=4066, nbytes=81804 -root: vortex.list(list(i32), len=4066) nbytes=11146 +root: vortex.list(list(i32), len=4066) nbytes=10304 metadata: elements: vortex.runend(i32, len=16384) nbytes=3968 metadata: offset: 0 @@ -15,11 +15,5 @@ root: vortex.list(list(i32), len=4066) nbytes=11146 metadata: reference: -49931i32 encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 metadata: bit_width: 17, offset: 0 - offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 - metadata: bit_width: 14, offset: 0 - patch_indices: vortex.primitive(u16, len=1) nbytes=2 - metadata: ptype: u16 - patch_values: vortex.constant(u16, len=1) nbytes=4 - metadata: scalar: 16384u16 - patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 - metadata: ptype: u8 + offsets: vortex.block_residual(u16, len=4067) nbytes=6336 + metadata: blocks: 4, slice: 0..4067 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap index 327f050b0d7..ddbd6548e12 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.fsst(utf8, len=16384) nbytes=151382 +root: vortex.fsst(utf8, len=16384) nbytes=141400 metadata: len: 16384, nsymbols: 223 uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 metadata: fill_value: 24u8 @@ -11,5 +11,5 @@ root: vortex.fsst(utf8, len=16384) nbytes=151382 metadata: ptype: u16 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 - codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 - metadata: bit_width: 17, offset: 0 + codes_offsets: vortex.block_residual(u32, len=16385) nbytes=27010 + metadata: blocks: 17, slice: 0..16385 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap index 6f38e8e6f5d..bc08ade5b96 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap @@ -3,9 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 -root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67584 +root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=59756 metadata: - storage: fastlanes.for(i64, len=16384) nbytes=67584 - metadata: reference: 1700000000891673i64 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67584 - metadata: bit_width: 33, offset: 0 + storage: vortex.block_residual(i64, len=16384) nbytes=59756 + metadata: blocks: 16, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap index c9554add05a..e5f7bbb6eee 100644 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__list_of_int_runs.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: list(i32), len=4066, nbytes=81804 -root: vortex.list(list(i32), len=4066) nbytes=11146 +root: vortex.list(list(i32), len=4066) nbytes=10304 metadata: elements: vortex.runend(i32, len=16384) nbytes=3968 metadata: offset: 0 @@ -15,11 +15,5 @@ root: vortex.list(list(i32), len=4066) nbytes=11146 metadata: reference: -49931i32 encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 metadata: bit_width: 17, offset: 0 - offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 - metadata: bit_width: 14, offset: 0 - patch_indices: vortex.primitive(u16, len=1) nbytes=2 - metadata: ptype: u16 - patch_values: vortex.constant(u16, len=1) nbytes=4 - metadata: scalar: 16384u16 - patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 - metadata: ptype: u8 + offsets: vortex.block_residual(u16, len=4067) nbytes=6336 + metadata: blocks: 4, slice: 0..4067 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap index 327f050b0d7..ddbd6548e12 100644 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.fsst(utf8, len=16384) nbytes=151382 +root: vortex.fsst(utf8, len=16384) nbytes=141400 metadata: len: 16384, nsymbols: 223 uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 metadata: fill_value: 24u8 @@ -11,5 +11,5 @@ root: vortex.fsst(utf8, len=16384) nbytes=151382 metadata: ptype: u16 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 - codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 - metadata: bit_width: 17, offset: 0 + codes_offsets: vortex.block_residual(u32, len=16385) nbytes=27010 + metadata: blocks: 17, slice: 0..16385 diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index bddaff61a02..44c7331b3d4 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -33,6 +33,7 @@ tracing = { workspace = true } url = { workspace = true } vortex-alp = { workspace = true } vortex-array = { workspace = true } +vortex-block-residual = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-bytebool = { workspace = true } @@ -43,6 +44,7 @@ vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["file"] } +vortex-float-quant = { workspace = true } vortex-fsst = { workspace = true } vortex-io = { workspace = true } vortex-layout = { workspace = true } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index d66aaa54a9b..abb9d423b9f 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -189,6 +189,8 @@ pub fn register_default_encodings(session: &VortexSession) { vortex_datetime_parts::initialize(session); vortex_decimal_byte_parts::initialize(session); vortex_fastlanes::initialize(session); + vortex_float_quant::initialize(session); + vortex_block_residual::initialize(session); vortex_runend::initialize(session); vortex_sequence::initialize(session); vortex_sparse::initialize(session); diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index fc3f9e420f8..34589d8444c 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -26,6 +26,7 @@ workspace = true vortex-alp = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } +vortex-block-residual = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-bytebool = { workspace = true } @@ -37,6 +38,7 @@ vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-file = { workspace = true, optional = true } vortex-flatbuffers = { workspace = true } +vortex-float-quant = { workspace = true } vortex-fsst = { workspace = true } vortex-io = { workspace = true } vortex-ipc = { workspace = true } diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index 2cc7bba0481..54aa0805ac3 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -2,8 +2,11 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] use std::sync::LazyLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use divan::Bencher; #[cfg(not(codspeed))] @@ -17,19 +20,33 @@ use vortex::VortexSessionDefault; use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::arrays::Primitive; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::builders::dict::dict_encode; use vortex::array::builtins::ArrayBuiltins; use vortex::array::dtype::Nullability; +use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::dtype::NativePType; use vortex::dtype::PType; +use vortex::dtype::half::f16; use vortex::encodings::alp::RDEncoder; use vortex::encodings::alp::RDEncoderExt; use vortex::encodings::alp::alp_encode; +use vortex::encodings::block_residual::BlockResidual; +use vortex::encodings::block_residual::OrderedFloat; +use vortex::encodings::block_residual::OrderedFloatArraySlotsExt; +use vortex::encodings::fastlanes::BitPacked; use vortex::encodings::fastlanes::Delta; use vortex::encodings::fastlanes::DeltaData; use vortex::encodings::fastlanes::FoR; +use vortex::encodings::fastlanes::FoRArrayExt; +use vortex::encodings::fastlanes::FoRArraySlotsExt; +use vortex::encodings::fastlanes::bitpack_compress::bitpack_encode_unchecked; use vortex::encodings::fastlanes::delta_compress; +use vortex::encodings::float_quant::FloatQuant; +use vortex::encodings::float_quant::FloatQuantArraySlotsExt; +use vortex::encodings::float_quant::analyze_float_quant; use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::encodings::pco::Pco; @@ -38,7 +55,13 @@ use vortex::encodings::sequence::sequence_encode; use vortex::encodings::zigzag::zigzag_encode; use vortex::encodings::zstd::Zstd; use vortex::encodings::zstd::ZstdData; +use vortex::scalar::Scalar; use vortex_array::VortexSessionExecute; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::float::FloatQuantScheme; +use vortex_btrblocks::schemes::float::OrderedBlockResidualScheme; +use vortex_btrblocks::schemes::integer::BlockResidualScheme; use vortex_error::VortexResult; use vortex_sequence::Sequence; use vortex_session::VortexSession; @@ -57,6 +80,8 @@ fn main() { // smaller inputs because they are much slower per element. const NUM_VALUES: u64 = 4096; const PCO_NUM_VALUES: u64 = 1024; +const PCO_COMPRESSION_LEVEL: usize = 8; +const PCO_VALUES_PER_PAGE: usize = 8192; #[cfg(feature = "zstd")] const ZSTD_NUM_VALUES: u64 = 128; const STRING_NUM_VALUES: usize = 2048; @@ -82,6 +107,12 @@ fn canonicalize(array: impl IntoArray, ctx: &mut ExecutionCtx) -> VortexResult(ctx) } +fn bench_compressor(bencher: Bencher, array: PrimitiveArray, compressor: BtrBlocksCompressor) { + with_byte_counter(bencher, array.nbytes()) + .with_inputs(|| (array.clone().into_array(), SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| compressor.compress(&array, &mut ctx).unwrap()); +} + // Setup functions fn setup_primitive_arrays(len: u64) -> (PrimitiveArray, PrimitiveArray, PrimitiveArray) { let mut ctx = SESSION.create_execution_ctx(); @@ -104,6 +135,351 @@ fn setup_primitive_arrays(len: u64) -> (PrimitiveArray, PrimitiveArray, Primitiv (uint_array, int_array, float_array) } +fn setup_widened_f32_array() -> PrimitiveArray { + let mut rng = StdRng::seed_from_u64(1); + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let trend = (index % 10_000) as f32 * 0.001; + f64::from(trend + rng.random_range(-1.0_f32..1.0)) + })) +} + +fn setup_quantized_f32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let mantissa = (index.wrapping_mul(7_919) as u32 & 0x7fff) << 8; + f32::from_bits(0x3f80_0000 | mantissa) + })) +} + +fn setup_quantized_f16_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let mantissa = (index.wrapping_mul(7_919) as u16) & 0x03f0; + f16::from_bits(0x3c00 | mantissa) + })) +} + +fn setup_general_f16_array() -> PrimitiveArray { + PrimitiveArray::from_iter( + (0..NUM_VALUES).map(|index| f16::from_bits(index.wrapping_mul(7_919) as u16)), + ) +} + +fn setup_general_f32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let mantissa = index.wrapping_mul(7_919) as u32 & 0x007f_ffff; + f32::from_bits(0x3f80_0000 | mantissa) + })) +} + +fn setup_general_f64_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let mantissa = index.wrapping_mul(0x9e37_79b9_7f4a_7c15) & 0x000f_ffff_ffff_ffff; + f64::from_bits(0x3ff0_0000_0000_0000 | mantissa) + })) +} + +fn setup_float_quant_near_miss_f32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let scrambled = (index as u32).wrapping_mul(2_654_435_761); + let sign = (scrambled & 1) << 31; + let exponent = ((scrambled >> 1) % 254 + 1) << 23; + let mantissa = scrambled & 0x007f_fffc; + f32::from_bits(sign | exponent | mantissa) + })) +} + +fn setup_float_quant_near_miss_f64_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let scrambled = index.wrapping_mul(0x9e37_79b9_7f4a_7c15); + let sign = (scrambled & 1) << 63; + let exponent = ((scrambled >> 1) % 2_046 + 1) << 52; + let mantissa = scrambled & 0x000f_ffff_ffff_fffc; + f64::from_bits(sign | exponent | mantissa) + })) +} + +fn setup_nonzero_secondary_array() -> PrimitiveArray { + let widened = setup_widened_f32_array(); + PrimitiveArray::from_iter( + widened + .as_slice::() + .iter() + .enumerate() + .map(|(index, value)| { + if index % 10 == 0 { + f64::from_bits(value.to_bits() | 1) + } else { + *value + } + }), + ) +} + +fn setup_nonzero_secondary_f16_array() -> PrimitiveArray { + let quantized = setup_quantized_f16_array(); + PrimitiveArray::from_iter(quantized.as_slice::().iter().enumerate().map( + |(index, value)| { + if index % 10 == 0 { + f16::from_bits(value.to_bits() | 1) + } else { + *value + } + }, + )) +} + +fn setup_nonzero_secondary_f32_array() -> PrimitiveArray { + let quantized = setup_quantized_f32_array(); + PrimitiveArray::from_iter(quantized.as_slice::().iter().enumerate().map( + |(index, value)| { + if index % 10 == 0 { + f32::from_bits(value.to_bits() | 1) + } else { + *value + } + }, + )) +} + +fn setup_secondary_width_array(width: u8) -> PrimitiveArray { + let widened = setup_widened_f32_array(); + let low_mask = (1_u64 << width) - 1; + PrimitiveArray::from_iter( + widened + .as_slice::() + .iter() + .enumerate() + .map(|(index, value)| { + let low = (index as u64).wrapping_mul(2_654_435_761) & low_mask; + f64::from_bits(value.to_bits() | low) + }), + ) +} + +fn setup_random_walk_array() -> PrimitiveArray { + let mut rng = StdRng::seed_from_u64(2); + let mut value = 1_000.0_f64; + PrimitiveArray::from_iter((0..NUM_VALUES).map(|_| { + value += rng.random_range(-0.01_f64..0.01); + value + })) +} + +fn setup_block_local_u64_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + block * 1_000_000 + residual + })) +} + +fn setup_block_local_u32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + (block * 1_000_000 + residual) as u32 + })) +} + +fn setup_block_local_i32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + (block as i32 - 1_000) * 1_000_000 + residual as i32 + })) +} + +fn setup_patch_density_u32_array(stride: u64) -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + if index % stride == 0 { + u32::MAX - index as u32 + } else { + 42 + } + })) +} + +fn setup_ordered_f32_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(7_919) % 1_024; + f32::from_bits(0x3f80_0000 + (block as u32 * 0x1_0000) + residual as u32) + })) +} + +fn setup_ordered_f16_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = (index / 1_024) % 8; + let residual = index.wrapping_mul(7_919) % 64; + f16::from_bits(0x3c00 + (block * 64 + residual) as u16) + })) +} + +fn setup_block_local_i16_array() -> PrimitiveArray { + PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = (index / 1_024) % 128; + let residual = index.wrapping_mul(2_654_435_761) % 128; + (block * 128 + residual) as i16 + })) +} + +fn setup_block_local_integer_array() -> PrimitiveArray { + match T::PTYPE { + PType::U8 => PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = (index / 1_024) % 16; + let residual = index.wrapping_mul(2_654_435_761) % 16; + (block * 16 + residual) as u8 + })), + PType::U16 => PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = (index / 1_024) % 256; + let residual = index.wrapping_mul(2_654_435_761) % 256; + (block * 256 + residual) as u16 + })), + PType::U32 => setup_block_local_u32_array(), + PType::U64 => setup_block_local_u64_array(), + PType::I8 => PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = (index / 1_024) % 16; + let residual = index.wrapping_mul(2_654_435_761) % 16; + ((block * 16 + residual) as i16 - 128) as i8 + })), + PType::I16 => setup_block_local_i16_array(), + PType::I32 => setup_block_local_i32_array(), + PType::I64 => PrimitiveArray::from_iter((0..NUM_VALUES).map(|index| { + let block = index / 1_024; + let residual = index.wrapping_mul(2_654_435_761) % 1_024; + (block as i64 - 1_000) * 1_000_000_000_000 + residual as i64 + })), + ptype => unreachable!("unsupported block residual benchmark type {ptype}"), + } +} + +fn encode_for_bitpacked_tree(array: &PrimitiveArray, bit_width: u8) -> vortex::array::ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let encoded = FoR::encode(array.clone(), &mut ctx).unwrap(); + let bitpacked = BitPacked::encode(encoded.encoded(), bit_width, &mut ctx).unwrap(); + FoR::try_new(bitpacked.into_array(), encoded.reference_scalar().clone()) + .unwrap() + .into_array() +} + +fn ordered_values(array: &PrimitiveArray) -> PrimitiveArray { + let ordered = OrderedFloat::from_primitive(array.as_view()).unwrap(); + ordered + .encoded() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap() +} + +fn encode_ordered_block_residual(array: &PrimitiveArray) -> vortex::array::ArrayRef { + let ordered = OrderedFloat::from_primitive(array.as_view()).unwrap(); + let residuals = BlockResidual::from_primitive(ordered.encoded().as_::()).unwrap(); + OrderedFloat::try_new(residuals.into_array(), array.ptype()) + .unwrap() + .into_array() +} + +fn encode_float_quant_tree(array: &PrimitiveArray) -> vortex::array::ArrayRef { + let analysis = analyze_float_quant(array.as_view()).unwrap(); + assert_eq!(analysis.secondary_bit_width, 0); + let primary = + FloatQuant::primary_for_primitive(array.as_view(), analysis.k, analysis.primary_min) + .unwrap(); + // SAFETY: The analysis computes this width from the exact primary range. + let primary = unsafe { bitpack_encode_unchecked(primary, analysis.primary_bit_width) }.unwrap(); + let reference = if array.ptype() == PType::F32 { + Scalar::from(u32::try_from(analysis.primary_min).unwrap()) + } else { + debug_assert_eq!(array.ptype(), PType::F64); + Scalar::from(analysis.primary_min) + }; + let primary = FoR::try_new(primary.into_array(), reference).unwrap(); + FloatQuant::try_new(primary.into_array(), None, array.ptype(), analysis.k) + .unwrap() + .into_array() +} + +fn encode_float_quant_nonzero_secondary_tree(array: &PrimitiveArray) -> vortex::array::ArrayRef { + let analysis = analyze_float_quant(array.as_view()).unwrap(); + assert_ne!(analysis.secondary_bit_width, 0); + let split = FloatQuant::from_primitive(array.as_view(), analysis.k).unwrap(); + let primary = split + .primary() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + let secondary = split + .secondary() + .unwrap() + .clone() + .execute::(&mut SESSION.create_execution_ctx()) + .unwrap(); + let biased_primary = PrimitiveArray::from_iter( + primary + .as_slice::() + .iter() + .map(|value| value - analysis.primary_min), + ); + // SAFETY: The analysis computes this width from the exact primary range. + let primary = unsafe { bitpack_encode_unchecked(biased_primary, analysis.primary_bit_width) } + .unwrap() + .into_array(); + let primary = FoR::try_new(primary, Scalar::from(analysis.primary_min)) + .unwrap() + .into_array(); + // SAFETY: The analysis computes the exact secondary width. + let secondary = unsafe { bitpack_encode_unchecked(secondary, analysis.secondary_bit_width) } + .unwrap() + .into_array(); + FloatQuant::try_new(primary, Some(secondary), PType::F64, analysis.k) + .unwrap() + .into_array() +} + +fn encode_float_quant_scheme_tree(array: &PrimitiveArray) -> vortex::array::ArrayRef { + BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build() + .compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + ) + .unwrap() +} + +fn encode_float_quant_nonzero_secondary_scheme_tree( + array: &PrimitiveArray, +) -> vortex::array::ArrayRef { + let encoded = encode_float_quant_scheme_tree(array); + let float_quant = encoded.as_::(); + assert!(float_quant.secondary().is_some()); + encoded +} + +fn encode_prior_default(array: &PrimitiveArray) -> vortex::array::ArrayRef { + BtrBlocksCompressorBuilder::default() + .exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]) + .build() + .compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + ) + .unwrap() +} + +fn encode_proposed_default(array: &PrimitiveArray) -> vortex::array::ArrayRef { + BtrBlocksCompressorBuilder::default() + .build() + .compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + ) + .unwrap() +} + #[expect(clippy::cast_possible_truncation)] fn gen_varbin_words(len: usize, uniqueness: f64) -> Vec { let mut rng = StdRng::seed_from_u64(0); @@ -329,6 +705,1290 @@ fn bench_alp_rd_decompress_f64(bencher: Bencher) { .bench_refs(|(a, ctx)| canonicalize((**a).clone(), ctx)); } +#[divan::bench(name = "ordered_float_compress_f64")] +fn bench_ordered_float_compress_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| OrderedFloat::from_primitive(array.as_view()).unwrap()); +} + +#[divan::bench(name = "ordered_float_decompress_f64")] +fn bench_ordered_float_decompress_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let encoded = OrderedFloat::from_primitive(float_array.as_view()) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_float_scalar_at_f64")] +fn bench_ordered_float_scalar_at_f64(bencher: Bencher) { + let encoded = OrderedFloat::from_primitive(setup_random_walk_array().as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_float_compress_f16")] +fn bench_ordered_float_compress_f16(bencher: Bencher) { + let float_array = setup_ordered_f16_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| &float_array) + .bench_refs(|array| OrderedFloat::from_primitive(array.as_view()).unwrap()); +} + +#[divan::bench(name = "ordered_float_decompress_f16")] +fn bench_ordered_float_decompress_f16(bencher: Bencher) { + let encoded = OrderedFloat::from_primitive(setup_ordered_f16_array().as_view()) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_float_scalar_at_f16")] +fn bench_ordered_float_scalar_at_f16(bencher: Bencher) { + let encoded = OrderedFloat::from_primitive(setup_ordered_f16_array().as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_float_compress_f32")] +fn bench_ordered_float_compress_f32(bencher: Bencher) { + let float_array = setup_ordered_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| OrderedFloat::from_primitive(array.as_view()).unwrap()); +} + +#[divan::bench(name = "ordered_float_decompress_f32")] +fn bench_ordered_float_decompress_f32(bencher: Bencher) { + let encoded = OrderedFloat::from_primitive(setup_ordered_f32_array().as_view()) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_float_scalar_at_f32")] +fn bench_ordered_float_scalar_at_f32(bencher: Bencher) { + let encoded = OrderedFloat::from_primitive(setup_ordered_f32_array().as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "block_residual_compress_u64")] +fn bench_block_residual_compress_u64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let ordered = ordered_values(&float_array); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &ordered) + .bench_refs(|array| BlockResidual::from_primitive(array.as_view()).unwrap()); +} + +#[divan::bench(name = "block_residual_decompress_u64")] +fn bench_block_residual_decompress_u64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let ordered = ordered_values(&float_array); + let encoded = BlockResidual::from_primitive(ordered.as_view()) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(types = [u8, u16, u32, u64, i8, i16, i32, i64])] +fn block_local_block_residual_compress(bencher: Bencher) { + let array = setup_block_local_integer_array::(); + let byte_width = u64::try_from(T::PTYPE.byte_width()).unwrap(); + + with_byte_counter(bencher, NUM_VALUES * byte_width) + .with_inputs(|| &array) + .bench_refs(|array| BlockResidual::from_primitive(array.as_view()).unwrap()); +} + +#[divan::bench(types = [u8, u16, u32, u64, i8, i16, i32, i64])] +fn block_local_block_residual_decompress(bencher: Bencher) { + let encoded = BlockResidual::from_primitive(setup_block_local_integer_array::().as_view()) + .unwrap() + .into_array(); + let byte_width = u64::try_from(T::PTYPE.byte_width()).unwrap(); + + with_byte_counter(bencher, NUM_VALUES * byte_width) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(types = [u8, u16, u32, u64, i8, i16, i32, i64])] +fn block_local_block_residual_scalar_at(bencher: Bencher) { + let encoded = BlockResidual::from_primitive(setup_block_local_integer_array::().as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "block_residual_slice_patched_u32")] +fn bench_block_residual_slice_patched_u32(bencher: Bencher) { + const SLICE_LEN: usize = 100; + + let encoded = BlockResidual::from_primitive(setup_patch_density_u32_array(16).as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + let start = next_index.fetch_add(2_654_435_761, Ordering::Relaxed) + % (encoded.len() - SLICE_LEN); + (&encoded, start) + }) + .bench_values(|(array, start)| array.slice(start..start + SLICE_LEN).unwrap()); +} + +#[divan::bench(name = "block_local_for_bitpacked_compress_u64")] +fn bench_block_local_for_bitpacked_compress_u64(bencher: Bencher) { + let array = setup_block_local_u64_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &array) + .bench_refs(|array| encode_for_bitpacked_tree(array, 31)); +} + +#[divan::bench(name = "block_local_for_bitpacked_decompress_u64")] +fn bench_block_local_for_bitpacked_decompress_u64(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_u64_array(), 31); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "block_local_for_bitpacked_scalar_at_u64")] +fn bench_block_local_for_bitpacked_scalar_at_u64(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_u64_array(), 31); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "block_local_for_bitpacked_compress_u32")] +fn bench_block_local_for_bitpacked_compress_u32(bencher: Bencher) { + let array = setup_block_local_u32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &array) + .bench_refs(|array| encode_for_bitpacked_tree(array, 31)); +} + +#[divan::bench(name = "block_local_for_bitpacked_decompress_u32")] +fn bench_block_local_for_bitpacked_decompress_u32(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_u32_array(), 31); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "block_local_for_bitpacked_scalar_at_u32")] +fn bench_block_local_for_bitpacked_scalar_at_u32(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_u32_array(), 31); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "block_local_for_bitpacked_compress_i32")] +fn bench_block_local_for_bitpacked_compress_i32(bencher: Bencher) { + let array = setup_block_local_i32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &array) + .bench_refs(|array| encode_for_bitpacked_tree(array, 31)); +} + +#[divan::bench(name = "block_local_for_bitpacked_decompress_i32")] +fn bench_block_local_for_bitpacked_decompress_i32(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_i32_array(), 31); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "block_local_for_bitpacked_scalar_at_i32")] +fn bench_block_local_for_bitpacked_scalar_at_i32(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_i32_array(), 31); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_block_residual_decompress_u32(bencher: Bencher, stride: u64) { + let encoded = BlockResidual::from_primitive(setup_patch_density_u32_array(stride).as_view()) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_block_residual_scalar_at_u32(bencher: Bencher, stride: u64) { + let encoded = BlockResidual::from_primitive(setup_patch_density_u32_array(stride).as_view()) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_prior_default_compress_u32(bencher: Bencher, stride: u64) { + let input = setup_patch_density_u32_array(stride); + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([BlockResidualScheme.id()]) + .build(); + bench_compressor(bencher, input, compressor); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_default_compress_u32(bencher: Bencher, stride: u64) { + bench_compressor( + bencher, + setup_patch_density_u32_array(stride), + BtrBlocksCompressorBuilder::default().build(), + ); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_prior_default_decompress_u32(bencher: Bencher, stride: u64) { + let encoded = encode_prior_default(&setup_patch_density_u32_array(stride)); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(args = [256, 64, 16, 4, 1])] +fn patch_density_default_decompress_u32(bencher: Bencher, stride: u64) { + let encoded = encode_proposed_default(&setup_patch_density_u32_array(stride)); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "block_local_for_bitpacked_compress_i16")] +fn bench_block_local_for_bitpacked_compress_i16(bencher: Bencher) { + let array = setup_block_local_i16_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| &array) + .bench_refs(|array| encode_for_bitpacked_tree(array, 14)); +} + +#[divan::bench(name = "block_local_for_bitpacked_decompress_i16")] +fn bench_block_local_for_bitpacked_decompress_i16(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_i16_array(), 14); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "block_local_for_bitpacked_scalar_at_i16")] +fn bench_block_local_for_bitpacked_scalar_at_i16(bencher: Bencher) { + let encoded = encode_for_bitpacked_tree(&setup_block_local_i16_array(), 14); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "block_local_pcodec_compress_i16")] +fn bench_block_local_pcodec_compress_i16(bencher: Bencher) { + let array = setup_block_local_i16_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + Pco::from_primitive( + array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + ctx, + ) + .unwrap() + }); +} + +#[divan::bench(name = "block_local_pcodec_decompress_i16")] +fn bench_block_local_pcodec_decompress_i16(bencher: Bencher) { + let array = setup_block_local_i16_array(); + let compressed = Pco::from_primitive( + array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + &mut SESSION.create_execution_ctx(), + ) + .unwrap(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_block_residual_compress_f64")] +fn bench_ordered_block_residual_compress_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_ordered_block_residual(array)); +} + +#[divan::bench(name = "ordered_block_residual_decompress_f64")] +fn bench_ordered_block_residual_decompress_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let encoded = encode_ordered_block_residual(&float_array); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_block_residual_scalar_at_f64")] +fn bench_ordered_block_residual_scalar_at_f64(bencher: Bencher) { + let encoded = encode_ordered_block_residual(&setup_random_walk_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_block_residual_prior_default_scalar_at_f64")] +fn bench_ordered_block_residual_prior_default_scalar_at_f64(bencher: Bencher) { + let encoded = encode_prior_default(&setup_random_walk_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_block_residual_compress_f16")] +fn bench_ordered_block_residual_compress_f16(bencher: Bencher) { + let float_array = setup_ordered_f16_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_ordered_block_residual(array)); +} + +#[divan::bench(name = "ordered_block_residual_decompress_f16")] +fn bench_ordered_block_residual_decompress_f16(bencher: Bencher) { + let encoded = encode_ordered_block_residual(&setup_ordered_f16_array()); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_block_residual_scalar_at_f16")] +fn bench_ordered_block_residual_scalar_at_f16(bencher: Bencher) { + let encoded = encode_ordered_block_residual(&setup_ordered_f16_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_block_residual_compress_f32")] +fn bench_ordered_block_residual_compress_f32(bencher: Bencher) { + let float_array = setup_ordered_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_ordered_block_residual(array)); +} + +#[divan::bench(name = "ordered_block_residual_decompress_f32")] +fn bench_ordered_block_residual_decompress_f32(bencher: Bencher) { + let encoded = encode_ordered_block_residual(&setup_ordered_f32_array()); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "ordered_block_residual_scalar_at_f32")] +fn bench_ordered_block_residual_scalar_at_f32(bencher: Bencher) { + let encoded = encode_ordered_block_residual(&setup_ordered_f32_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "ordered_block_residual_scheme_compress_f64")] +fn bench_ordered_block_residual_scheme_compress_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&OrderedBlockResidualScheme) + .build(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| { + ( + float_array.clone().into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_values(|(array, mut ctx)| compressor.compress(&array, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_split_compress_f16")] +fn bench_float_quant_split_compress_f16(bencher: Bencher) { + let float_array = setup_quantized_f16_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| &float_array) + .bench_refs(|array| { + FloatQuant::from_primitive_constant_secondary(array.as_view(), 4).unwrap() + }); +} + +#[divan::bench(name = "float_quant_split_decompress_f16")] +fn bench_float_quant_split_decompress_f16(bencher: Bencher) { + let encoded = + FloatQuant::from_primitive_constant_secondary(setup_quantized_f16_array().as_view(), 4) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_split_scalar_at_f16")] +fn bench_float_quant_split_scalar_at_f16(bencher: Bencher) { + let encoded = + FloatQuant::from_primitive_constant_secondary(setup_quantized_f16_array().as_view(), 4) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_scheme_compress_f16")] +fn bench_float_quant_scheme_compress_f16(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, setup_quantized_f16_array(), compressor); +} + +#[divan::bench(name = "float_quant_tree_decompress_f16")] +fn bench_float_quant_tree_decompress_f16(bencher: Bencher) { + let encoded = encode_float_quant_scheme_tree(&setup_quantized_f16_array()); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_tree_scalar_at_f16")] +fn bench_float_quant_tree_scalar_at_f16(bencher: Bencher) { + let encoded = encode_float_quant_scheme_tree(&setup_quantized_f16_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_prior_default_compress_f16")] +fn bench_float_quant_prior_default_compress_f16(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]) + .build(); + bench_compressor(bencher, setup_quantized_f16_array(), compressor); +} + +#[divan::bench(name = "float_quant_proposed_default_compress_f16")] +fn bench_float_quant_proposed_default_compress_f16(bencher: Bencher) { + bench_compressor( + bencher, + setup_quantized_f16_array(), + BtrBlocksCompressorBuilder::default().build(), + ); +} + +#[divan::bench(name = "float_quant_prior_default_decompress_f16")] +fn bench_float_quant_prior_default_decompress_f16(bencher: Bencher) { + let encoded = encode_prior_default(&setup_quantized_f16_array()); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_proposed_default_decompress_f16")] +fn bench_float_quant_proposed_default_decompress_f16(bencher: Bencher) { + let encoded = encode_proposed_default(&setup_quantized_f16_array()); + + with_byte_counter(bencher, NUM_VALUES * 2) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_prior_default_reject_f16")] +fn bench_float_quant_prior_default_reject_f16(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]) + .build(); + bench_compressor(bencher, setup_general_f16_array(), compressor); +} + +#[divan::bench(name = "float_quant_proposed_default_reject_f16")] +fn bench_float_quant_proposed_default_reject_f16(bencher: Bencher) { + bench_compressor( + bencher, + setup_general_f16_array(), + BtrBlocksCompressorBuilder::default().build(), + ); +} + +macro_rules! float_quant_rejection_benches { + ($scheme_only:ident, $without_scheme:ident, $with_scheme:ident, $setup:ident) => { + #[divan::bench] + fn $scheme_only(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, $setup(), compressor); + } + + #[divan::bench] + fn $without_scheme(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([FloatQuantScheme.id()]) + .build(); + bench_compressor(bencher, $setup(), compressor); + } + + #[divan::bench] + fn $with_scheme(bencher: Bencher) { + bench_compressor( + bencher, + $setup(), + BtrBlocksCompressorBuilder::default().build(), + ); + } + }; +} + +float_quant_rejection_benches!( + float_quant_scheme_reject_f16, + float_quant_default_without_scheme_reject_f16, + float_quant_default_with_scheme_reject_f16, + setup_general_f16_array +); + +float_quant_rejection_benches!( + float_quant_scheme_reject_f32, + float_quant_default_without_scheme_reject_f32, + float_quant_default_with_scheme_reject_f32, + setup_general_f32_array +); + +float_quant_rejection_benches!( + float_quant_scheme_reject_f64, + float_quant_default_without_scheme_reject_f64, + float_quant_default_with_scheme_reject_f64, + setup_general_f64_array +); + +float_quant_rejection_benches!( + float_quant_scheme_near_miss_f32, + float_quant_default_without_scheme_near_miss_f32, + float_quant_default_with_scheme_near_miss_f32, + setup_float_quant_near_miss_f32_array +); + +float_quant_rejection_benches!( + float_quant_scheme_near_miss_f64, + float_quant_default_without_scheme_near_miss_f64, + float_quant_default_with_scheme_near_miss_f64, + setup_float_quant_near_miss_f64_array +); + +#[divan::bench(name = "float_quant_split_compress_f64")] +fn bench_float_quant_split_compress_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| { + FloatQuant::from_primitive_constant_secondary(array.as_view(), 29).unwrap() + }); +} + +#[divan::bench(name = "float_quant_split_decompress_f64")] +fn bench_float_quant_split_decompress_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + let encoded = FloatQuant::from_primitive_constant_secondary(float_array.as_view(), 29) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_split_scalar_at_f64")] +fn bench_float_quant_split_scalar_at_f64(bencher: Bencher) { + let encoded = + FloatQuant::from_primitive_constant_secondary(setup_widened_f32_array().as_view(), 29) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_materialized_tree_compress_f64")] +fn bench_float_quant_tree_compress_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_float_quant_tree(array)); +} + +#[divan::bench(name = "float_quant_scheme_compress_f64")] +fn bench_float_quant_scheme_compress_f64(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, setup_widened_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_tree_decompress_f64")] +fn bench_float_quant_tree_decompress_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + let encoded = encode_float_quant_tree(&float_array); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_tree_scalar_at_f64")] +fn bench_float_quant_tree_scalar_at_f64(bencher: Bencher) { + let encoded = encode_float_quant_tree(&setup_widened_f32_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_split_compress_f32")] +fn bench_float_quant_split_compress_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| { + FloatQuant::from_primitive_constant_secondary(array.as_view(), k).unwrap() + }); +} + +#[divan::bench(name = "float_quant_split_decompress_f32")] +fn bench_float_quant_split_decompress_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive_constant_secondary(float_array.as_view(), k) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_split_scalar_at_f32")] +fn bench_float_quant_split_scalar_at_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive_constant_secondary(float_array.as_view(), k) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_materialized_tree_compress_f32")] +fn bench_float_quant_tree_compress_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_float_quant_tree(array)); +} + +#[divan::bench(name = "float_quant_alp_rd_compress_f32")] +fn bench_float_quant_alp_rd_compress_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| { + let encoder = RDEncoder::new(array.as_slice::()); + encoder.encode(array.as_view()) + }); +} + +#[divan::bench(name = "float_quant_tree_decompress_f32")] +fn bench_float_quant_tree_decompress_f32(bencher: Bencher) { + let encoded = encode_float_quant_tree(&setup_quantized_f32_array()); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_tree_scalar_at_f32")] +fn bench_float_quant_tree_scalar_at_f32(bencher: Bencher) { + let encoded = encode_float_quant_tree(&setup_quantized_f32_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_analyze_f32")] +fn bench_float_quant_analyze_f32(bencher: Bencher) { + let float_array = setup_quantized_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 4) + .with_inputs(|| &float_array) + .bench_refs(|array| analyze_float_quant(array.as_view()).unwrap()); +} + +#[divan::bench(name = "float_quant_scheme_compress_f32")] +fn bench_float_quant_scheme_compress_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, setup_quantized_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_prior_default_compress_f32")] +fn bench_float_quant_prior_default_compress_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]) + .build(); + bench_compressor(bencher, setup_quantized_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_default_compress_f32")] +fn bench_float_quant_default_compress_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([OrderedBlockResidualScheme.id(), BlockResidualScheme.id()]) + .build(); + bench_compressor(bencher, setup_quantized_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_proposed_default_compress_f32")] +fn bench_float_quant_proposed_default_compress_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default().build(); + bench_compressor(bencher, setup_quantized_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_prior_default_reject_f32")] +fn bench_float_quant_prior_default_reject_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([ + FloatQuantScheme.id(), + OrderedBlockResidualScheme.id(), + BlockResidualScheme.id(), + ]) + .build(); + bench_compressor(bencher, setup_general_f32_array(), compressor); +} + +#[divan::bench(name = "float_quant_proposed_default_reject_f32")] +fn bench_float_quant_proposed_default_reject_f32(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::default().build(); + bench_compressor(bencher, setup_general_f32_array(), compressor); +} + +macro_rules! float_quant_nonzero_secondary_benches { + ( + $split_compress:ident, + $split_decompress:ident, + $split_scalar:ident, + $scheme_compress:ident, + $tree_decompress:ident, + $tree_scalar:ident, + $setup:ident, + $byte_width:expr + ) => { + #[divan::bench] + fn $split_compress(bencher: Bencher) { + let float_array = $setup(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + + with_byte_counter(bencher, NUM_VALUES * $byte_width) + .with_inputs(|| &float_array) + .bench_refs(|array| FloatQuant::from_primitive(array.as_view(), k).unwrap()); + } + + #[divan::bench] + fn $split_decompress(bencher: Bencher) { + let float_array = $setup(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive(float_array.as_view(), k) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * $byte_width) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); + } + + #[divan::bench] + fn $split_scalar(bencher: Bencher) { + let float_array = $setup(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive(float_array.as_view(), k) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| { + array.execute_scalar(index, &mut ctx).unwrap() + }); + } + + #[divan::bench] + fn $scheme_compress(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, $setup(), compressor); + } + + #[divan::bench] + fn $tree_decompress(bencher: Bencher) { + let encoded = encode_float_quant_nonzero_secondary_scheme_tree(&$setup()); + + with_byte_counter(bencher, NUM_VALUES * $byte_width) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); + } + + #[divan::bench] + fn $tree_scalar(bencher: Bencher) { + let encoded = encode_float_quant_nonzero_secondary_scheme_tree(&$setup()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| { + array.execute_scalar(index, &mut ctx).unwrap() + }); + } + }; +} + +float_quant_nonzero_secondary_benches!( + float_quant_nonzero_secondary_split_compress_f16, + float_quant_nonzero_secondary_split_decompress_f16, + float_quant_nonzero_secondary_split_scalar_at_f16, + float_quant_nonzero_secondary_scheme_compress_f16, + float_quant_nonzero_secondary_tree_decompress_f16, + float_quant_nonzero_secondary_tree_scalar_at_f16, + setup_nonzero_secondary_f16_array, + 2 +); + +float_quant_nonzero_secondary_benches!( + float_quant_nonzero_secondary_split_compress_f32, + float_quant_nonzero_secondary_split_decompress_f32, + float_quant_nonzero_secondary_split_scalar_at_f32, + float_quant_nonzero_secondary_scheme_compress_f32, + float_quant_nonzero_secondary_tree_decompress_f32, + float_quant_nonzero_secondary_tree_scalar_at_f32, + setup_nonzero_secondary_f32_array, + 4 +); + +#[divan::bench(name = "float_quant_nonzero_secondary_split_compress_f64")] +fn bench_float_quant_nonzero_secondary_split_compress_f64(bencher: Bencher) { + let float_array = setup_nonzero_secondary_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| FloatQuant::from_primitive(array.as_view(), k).unwrap()); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_split_decompress_f64")] +fn bench_float_quant_nonzero_secondary_split_decompress_f64(bencher: Bencher) { + let float_array = setup_nonzero_secondary_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive(float_array.as_view(), k) + .unwrap() + .into_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_split_scalar_at_f64")] +fn bench_float_quant_nonzero_secondary_split_scalar_at_f64(bencher: Bencher) { + let float_array = setup_nonzero_secondary_array(); + let k = analyze_float_quant(float_array.as_view()).unwrap().k; + let encoded = FloatQuant::from_primitive(float_array.as_view(), k) + .unwrap() + .into_array(); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_tree_compress_f64")] +fn bench_float_quant_nonzero_secondary_tree_compress_f64(bencher: Bencher) { + let float_array = setup_nonzero_secondary_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_float_quant_nonzero_secondary_tree(array)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_scheme_compress_f64")] +fn bench_float_quant_nonzero_secondary_scheme_compress_f64(bencher: Bencher) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&FloatQuantScheme) + .build(); + bench_compressor(bencher, setup_nonzero_secondary_array(), compressor); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_default_compress_f64")] +fn bench_float_quant_nonzero_secondary_default_compress_f64(bencher: Bencher) { + bench_compressor( + bencher, + setup_nonzero_secondary_array(), + BtrBlocksCompressorBuilder::default().build(), + ); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_tree_decompress_f64")] +fn bench_float_quant_nonzero_secondary_tree_decompress_f64(bencher: Bencher) { + let encoded = encode_float_quant_nonzero_secondary_tree(&setup_nonzero_secondary_array()); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_default_decompress_f64")] +fn bench_float_quant_nonzero_secondary_default_decompress_f64(bencher: Bencher) { + let input = setup_nonzero_secondary_array().into_array(); + let encoded = BtrBlocksCompressorBuilder::default() + .build() + .compress(&input, &mut SESSION.create_execution_ctx()) + .unwrap(); + assert!(encoded.is::()); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(args = [1, 4, 8, 16])] +fn float_quant_secondary_width_decompress_f64(bencher: Bencher, width: u8) { + let encoded = encode_float_quant_nonzero_secondary_tree(&setup_secondary_width_array(width)); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_tree_scalar_at_f64")] +fn bench_float_quant_nonzero_secondary_tree_scalar_at_f64(bencher: Bencher) { + let encoded = encode_float_quant_nonzero_secondary_tree(&setup_nonzero_secondary_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_prior_default_compress_f64")] +fn bench_float_quant_nonzero_secondary_prior_default_compress_f64(bencher: Bencher) { + let float_array = setup_nonzero_secondary_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| &float_array) + .bench_refs(|array| encode_prior_default(array)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_prior_default_decompress_f64")] +fn bench_float_quant_nonzero_secondary_prior_default_decompress_f64(bencher: Bencher) { + let encoded = encode_prior_default(&setup_nonzero_secondary_array()); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "float_quant_nonzero_secondary_prior_default_scalar_at_f64")] +fn bench_float_quant_nonzero_secondary_prior_default_scalar_at_f64(bencher: Bencher) { + let encoded = encode_prior_default(&setup_nonzero_secondary_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + +#[divan::bench(name = "float_quant_prior_default_scalar_at_f64")] +fn bench_float_quant_prior_default_scalar_at_f64(bencher: Bencher) { + let encoded = encode_prior_default(&setup_widened_f32_array()); + let next_index = AtomicUsize::new(0); + + bencher + .with_inputs(|| { + ( + &encoded, + SESSION.create_execution_ctx(), + next_index.fetch_add(2_654_435_761, Ordering::Relaxed) % encoded.len(), + ) + }) + .bench_values(|(array, mut ctx, index)| array.execute_scalar(index, &mut ctx).unwrap()); +} + #[divan::bench(name = "pcodec_compress_f64")] fn bench_pcodec_compress_f64(bencher: Bencher) { let (_, _, float_array) = setup_primitive_arrays(PCO_NUM_VALUES); @@ -354,6 +2014,72 @@ fn bench_pcodec_decompress_f64(bencher: Bencher) { .bench_refs(|(a, ctx)| canonicalize((**a).clone(), ctx)); } +#[divan::bench(name = "pcodec_compress_widened_f32_f64")] +fn bench_pcodec_compress_widened_f32_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&float_array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + Pco::from_primitive( + array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + ctx, + ) + .unwrap() + }); +} + +#[divan::bench(name = "pcodec_decompress_widened_f32_f64")] +fn bench_pcodec_decompress_widened_f32_f64(bencher: Bencher) { + let float_array = setup_widened_f32_array(); + let compressed = Pco::from_primitive( + float_array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + &mut SESSION.create_execution_ctx(), + ) + .unwrap(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + +#[divan::bench(name = "pcodec_compress_random_walk_f64")] +fn bench_pcodec_compress_random_walk_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&float_array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| { + Pco::from_primitive( + array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + ctx, + ) + .unwrap() + }); +} + +#[divan::bench(name = "pcodec_decompress_random_walk_f64")] +fn bench_pcodec_decompress_random_walk_f64(bencher: Bencher) { + let float_array = setup_random_walk_array(); + let compressed = Pco::from_primitive( + float_array.as_view(), + PCO_COMPRESSION_LEVEL, + PCO_VALUES_PER_PAGE, + &mut SESSION.create_execution_ctx(), + ) + .unwrap(); + + with_byte_counter(bencher, NUM_VALUES * 8) + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| canonicalize((**array).clone(), ctx)); +} + #[cfg(feature = "zstd")] #[divan::bench(name = "zstd_compress_u32")] fn bench_zstd_compress_u32(bencher: Bencher) { diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index a6dd8ee7fe9..1037bf097fd 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -46,13 +46,14 @@ pub use self::preview::PREVIEW_2025_05_0; pub use self::preview::PREVIEW_2026_02_0; pub use self::preview::PREVIEW_2026_04_0; pub use self::preview::PREVIEW_2026_06_0; +pub use self::preview::PREVIEW_2026_08_0; /// The `core` edition enabled for writing by the default Vortex session. pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_08_1; /// The `preview` edition enabled for writing by the default Vortex session when the /// `unstable_encodings` feature is selected. -pub const DEFAULT_PREVIEW_EDITION: EditionId = PREVIEW_2026_06_0; +pub const DEFAULT_PREVIEW_EDITION: EditionId = PREVIEW_2026_08_0; /// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ @@ -67,6 +68,7 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &preview::v2026_02::DECLARATION, &preview::v2026_04::DECLARATION, &preview::v2026_06::DECLARATION, + &preview::v2026_08::DECLARATION, ]; /// Register the Vortex edition declarations with the session's [`EditionSession`]. diff --git a/vortex/src/editions/preview/mod.rs b/vortex/src/editions/preview/mod.rs index cba3ba1dc6c..07f07c5431e 100644 --- a/vortex/src/editions/preview/mod.rs +++ b/vortex/src/editions/preview/mod.rs @@ -10,8 +10,10 @@ pub mod v2025_05; pub mod v2026_02; pub mod v2026_04; pub mod v2026_06; +pub mod v2026_08; pub use v2025_05::PREVIEW_2025_05_0; pub use v2026_02::PREVIEW_2026_02_0; pub use v2026_04::PREVIEW_2026_04_0; pub use v2026_06::PREVIEW_2026_06_0; +pub use v2026_08::PREVIEW_2026_08_0; diff --git a/vortex/src/editions/preview/v2026_08.rs b/vortex/src/editions/preview/v2026_08.rs new file mode 100644 index 00000000000..97d92410520 --- /dev/null +++ b/vortex/src/editions/preview/v2026_08.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The August 2026 `preview` encoding cohort. + +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionMember; + +/// The August 2026 draft edition of the `preview` family. +pub const PREVIEW_2026_08_0: EditionId = EditionId::new("preview", 2026, 8, 0); + +/// The declaration of [`PREVIEW_2026_08_0`] and the encodings that join the family at it. +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: PREVIEW_2026_08_0, + min_vortex_version: None, + }, + added: &[ + EditionMember::array(&"vortex.block_residual"), + EditionMember::array(&"vortex.float_quant"), + EditionMember::array(&"vortex.ordered_float"), + ], +}; diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 4ba621baf7a..53ac1149ee3 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::f64::consts::TAU; use std::sync::Arc; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; @@ -16,6 +18,8 @@ use vortex_array::dtype::PType; use vortex_array::field_path; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; +use vortex_block_residual::BlockResidual; +use vortex_block_residual::OrderedFloat; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::ByteBufferMut; use vortex_edition::ComponentKind; @@ -33,6 +37,7 @@ use vortex_error::vortex_err; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; +use vortex_float_quant::FloatQuant; use vortex_io::session::RuntimeSession; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::compressed::CompressingStrategy; @@ -52,6 +57,7 @@ use super::DEFAULT_CORE_EDITION; use super::DEFAULT_PREVIEW_EDITION; use super::EDITION_DECLARATIONS; use super::PREVIEW_2026_06_0; +use super::PREVIEW_2026_08_0; fn session() -> Result { let session = EditionSession::empty(); @@ -188,6 +194,29 @@ fn core_2026_08_3_adds_variants() { ); } +#[test] +fn preview_2026_08_adds_numeric_arrays() { + let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); + assert!( + session + .find(&PREVIEW_2026_08_0) + .unwrap_or_else(|| panic!("{PREVIEW_2026_08_0} is not registered")) + .is_draft() + ); + let arrays = session.components_in(&PREVIEW_2026_08_0, ComponentKind::Array); + for id in [ + "vortex.block_residual", + "vortex.float_quant", + "vortex.ordered_float", + ] { + assert!( + arrays + .iter() + .any(|inclusion| inclusion.component_id.as_str() == id) + ); + } +} + #[test] fn encodings_in_editions_unions_families() { let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); @@ -690,6 +719,93 @@ async fn default_strategy_round_trip_uses_only_enabled_encodings() -> VortexResu .await } +#[tokio::test] +async fn numeric_draft_writer_round_trips_float_quant() -> VortexResult<()> { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + session + .enable_edition(PREVIEW_2026_08_0) + .map_err(|error| vortex_err!("{error}"))?; + let values = (0u32..65_536) + .map(|index| { + let mantissa = index.wrapping_mul(7_919) & 0x007f_ffff; + f64::from(f32::from_bits(0x3f80_0000 | mantissa)) + }) + .collect::>(); + let expected = PrimitiveArray::from_iter(values.clone()).into_array(); + let buffer = write_with(&session, expected.clone()).await?; + let actual = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert!( + actual + .depth_first_traversal() + .any(|array| array.is::()) + ); + let decoded = actual.execute::(&mut session.create_execution_ctx())?; + assert_eq!(decoded.as_slice::(), values); + Ok(()) +} + +#[tokio::test] +async fn numeric_draft_writer_round_trips_ordered_block_residual() -> VortexResult<()> { + use crate::VortexSessionDefault; + + fn uniform(state: &mut u64) -> f64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + ((*state >> 11) as f64 + 0.5) / (1_u64 << 53) as f64 + } + + let session = VortexSession::default(); + session + .enable_edition(PREVIEW_2026_08_0) + .map_err(|error| vortex_err!("{error}"))?; + let mut state = 0x4d59_5df4_d0f3_3173_u64; + let mut value = 0.0_f64; + let values = (0..65_536) + .map(|_| { + let radius = (-2.0 * uniform(&mut state).ln()).sqrt(); + let normal = radius * (TAU * uniform(&mut state)).cos(); + value += normal * 0.01; + value + }) + .collect::>(); + let expected = PrimitiveArray::from_iter(values.clone()).into_array(); + let buffer = write_with(&session, expected).await?; + let actual = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let encoding_ids = actual + .depth_first_traversal() + .map(|array| array.encoding_id().to_string()) + .collect::>(); + assert!( + actual + .depth_first_traversal() + .any(|array| array.is::()), + "written tree does not contain OrderedFloat: {encoding_ids:?}" + ); + assert!( + actual + .depth_first_traversal() + .any(|array| array.is::()) + ); + let decoded = actual.execute::(&mut session.create_execution_ctx())?; + assert_eq!(decoded.as_slice::(), values); + Ok(()) +} + #[tokio::test] async fn replacement_default_builder_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { let session = writer_test_session()?; diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1d1ab1252ac..214041ac889 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -260,6 +260,11 @@ pub mod encodings { pub use vortex_fastlanes::*; } + /// Lossless float quantization with integer child arrays. + pub mod float_quant { + pub use vortex_float_quant::*; + } + /// Fast Static Symbol Table string encoding. pub mod fsst { pub use vortex_fsst::*; @@ -270,6 +275,11 @@ pub mod encodings { pub use vortex_pco::*; } + /// Ordered-float and block-residual numeric encodings. + pub mod block_residual { + pub use vortex_block_residual::*; + } + /// Arrow-compatible run-end encoding. pub mod runend { pub use vortex_runend::*;