From 78df2bf0e4c1aa75260e977fbb64e8987e6fd9cd Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 20 Jul 2026 12:31:16 -0400 Subject: [PATCH] Add query-aware compression cost experiment Signed-off-by: Connor Tsui --- .github/workflows/sql-benchmarks.yml | 13 +- vortex-bench/src/bin/data-gen.rs | 20 +- vortex-bench/src/conversions.rs | 116 +++++- vortex-btrblocks/Cargo.toml | 5 + vortex-btrblocks/benches/predicate_strings.rs | 267 +++++++++++++ vortex-btrblocks/src/builder.rs | 39 +- vortex-btrblocks/src/lib.rs | 12 +- vortex-btrblocks/src/predicate_string_cost.rs | 140 +++++++ .../schemes/string/scheme_selection_tests.rs | 27 +- vortex-compressor/src/cost/mod.rs | 5 + vortex-compressor/src/cost/workload.rs | 356 ++++++++++++++++++ vortex/src/lib.rs | 7 + 12 files changed, 983 insertions(+), 24 deletions(-) create mode 100644 vortex-btrblocks/benches/predicate_strings.rs create mode 100644 vortex-btrblocks/src/predicate_string_cost.rs create mode 100644 vortex-compressor/src/cost/workload.rs diff --git a/.github/workflows/sql-benchmarks.yml b/.github/workflows/sql-benchmarks.yml index 243db83aff8..bb7bb87125a 100644 --- a/.github/workflows/sql-benchmarks.yml +++ b/.github/workflows/sql-benchmarks.yml @@ -236,7 +236,8 @@ on: {"engine": "duckdb", "format": "vortex"}, {"engine": "duckdb", "format": "vortex-compact"} ], - "scale_factor": "100" + "scale_factor": "100", + "write_profile": "fineweb-queries" }, { "id": "fineweb-s3", @@ -261,7 +262,8 @@ on: {"engine": "duckdb", "format": "vortex"}, {"engine": "duckdb", "format": "vortex-compact"} ], - "scale_factor": "100" + "scale_factor": "100", + "write_profile": "fineweb-queries" }, { "id": "polarsignals", @@ -576,10 +578,12 @@ jobs: RUSTFLAGS: "-C target-cpu=native" run: | packages=(--bin data-gen --bin datafusion-bench --bin duckdb-bench) + features=() if [ "${{ inputs.mode }}" != "pr" ]; then packages+=(--bin lance-bench) + features+=(--features unstable_encodings) fi - cargo build "${packages[@]}" --profile release_debug --features unstable_encodings + cargo build "${packages[@]}" "${features[@]}" --profile release_debug - name: Generate data shell: bash @@ -588,7 +592,8 @@ jobs: run: | uv run --project bench-orchestrator vx-bench prepare-data "${{ matrix.subcommand }}" \ --formats-json '${{ toJSON(matrix.data_formats) }}' \ - ${{ matrix.scale_factor && format('--opt scale-factor={0}', matrix.scale_factor) || '' }} + ${{ matrix.scale_factor && format('--opt scale-factor={0}', matrix.scale_factor) || '' }} \ + ${{ matrix.write_profile && format('--opt write-profile={0}', matrix.write_profile) || '' }} - name: Setup AWS CLI if: inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false diff --git a/vortex-bench/src/bin/data-gen.rs b/vortex-bench/src/bin/data-gen.rs index 35c77d70c48..0ada3a34646 100644 --- a/vortex-bench/src/bin/data-gen.rs +++ b/vortex-bench/src/bin/data-gen.rs @@ -20,7 +20,8 @@ use vortex_bench::Format; use vortex_bench::LogFormat; use vortex_bench::Opt; use vortex_bench::Opts; -use vortex_bench::conversions::convert_parquet_directory_to_vortex; +use vortex_bench::conversions::WriteProfile; +use vortex_bench::conversions::convert_parquet_directory_to_vortex_with_profile; use vortex_bench::create_benchmark; use vortex_bench::generate_duckdb_registration_sql; use vortex_bench::setup_logging_and_tracing_with_format; @@ -58,6 +59,9 @@ async fn main() -> anyhow::Result<()> { setup_logging_and_tracing_with_format(args.verbose, args.tracing, args.log_format)?; let benchmark = create_benchmark(args.benchmark, &opts)?; + let write_profile = opts + .get_as::("write-profile") + .unwrap_or_default(); // Generate base Parquet data - this is the source for all other formats benchmark.generate_base_data().await?; @@ -74,7 +78,12 @@ async fn main() -> anyhow::Result<()> { .iter() .any(|f| matches!(f, Format::OnDiskVortex)) { - convert_parquet_directory_to_vortex(&base_path, CompactionStrategy::Default).await?; + convert_parquet_directory_to_vortex_with_profile( + &base_path, + CompactionStrategy::Default, + write_profile, + ) + .await?; } if args @@ -82,7 +91,12 @@ async fn main() -> anyhow::Result<()> { .iter() .any(|f| matches!(f, Format::VortexCompact)) { - convert_parquet_directory_to_vortex(&base_path, CompactionStrategy::Compact).await?; + convert_parquet_directory_to_vortex_with_profile( + &base_path, + CompactionStrategy::Compact, + write_profile, + ) + .await?; } if args diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 92064c25785..432a9d5a8a7 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -4,6 +4,7 @@ use std::fs; use std::path::Path; use std::path::PathBuf; +use std::str::FromStr; use std::sync::Arc; use futures::StreamExt; @@ -75,6 +76,30 @@ const MIN_CONCURRENCY: u64 = 1; /// Maximum number of concurrent conversion streams. This is somewhat arbitary. const MAX_CONCURRENCY: u64 = 16; +/// Compression objective used while generating Vortex benchmark files. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum WriteProfile { + /// Use the production size-oriented compressor. + #[default] + Default, + /// Use field-specific equality and `LIKE` profiles for the FineWeb query suite. + FinewebQueries, +} + +impl FromStr for WriteProfile { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + match value { + "default" => Ok(Self::Default), + "fineweb-queries" => Ok(Self::FinewebQueries), + _ => anyhow::bail!( + "unknown write profile {value:?}; expected default or fineweb-queries" + ), + } + } +} + /// Returns the available system memory in bytes. fn available_memory_bytes() -> u64 { System::new_all().available_memory() @@ -140,6 +165,22 @@ pub async fn convert_parquet_file_to_vortex( parquet_path: &Path, output_path: &Path, compaction: CompactionStrategy, +) -> anyhow::Result<()> { + convert_parquet_file_to_vortex_with_profile( + parquet_path, + output_path, + compaction, + WriteProfile::Default, + ) + .await +} + +/// Convert a single Parquet file to Vortex format with a benchmark-only write profile. +pub async fn convert_parquet_file_to_vortex_with_profile( + parquet_path: &Path, + output_path: &Path, + compaction: CompactionStrategy, + profile: WriteProfile, ) -> anyhow::Result<()> { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; @@ -157,7 +198,7 @@ pub async fn convert_parquet_file_to_vortex( .open(output_path) .await?; - write_options_for(compaction, &dtype, is_spatialbench(parquet_path)) + write_options_for(compaction, profile, &dtype, is_spatialbench(parquet_path)) .write( &mut output_file, ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), @@ -179,6 +220,7 @@ fn is_spatialbench(path: &Path) -> bool { /// unique, so the dictionary builder balloons memory (tens of GB) for zero gain. fn write_options_for( compaction: CompactionStrategy, + profile: WriteProfile, dtype: &DType, skip_binary_dict: bool, ) -> VortexWriteOptions { @@ -192,14 +234,30 @@ fn write_options_for( .collect(), _ => Vec::new(), }; - if binary_fields.is_empty() { + if binary_fields.is_empty() && profile == WriteProfile::Default { return compaction.apply_options(SESSION.write_options()); } - let mut builder = WriteStrategyBuilder::default(); - if matches!(compaction, CompactionStrategy::Compact) { - builder = - builder.with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()); + let compressor = compressor_builder(compaction); + let compressor = match profile { + WriteProfile::Default => compressor, + WriteProfile::FinewebQueries => compressor.with_like_strings(), + }; + let mut builder = WriteStrategyBuilder::default() + .with_btrblocks_builder(compressor) + // Layout dictionary eligibility is outside the scheme cost-model boundary. Preserve the + // production size-based probe so this experiment changes leaf encodings, not file shape. + .with_probe_compressor(compressor_builder(compaction).build()); + if profile == WriteProfile::FinewebQueries { + for name in ["dump", "language"] { + builder = builder.with_field_writer( + FieldPath::from_name(name), + profiled_layout( + compaction, + compressor_builder(compaction).with_equality_strings(), + ), + ); + } } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); @@ -207,14 +265,35 @@ fn write_options_for( SESSION.write_options().with_strategy(builder.build()) } -/// A chunked + compressed layout that skips dictionary encoding for opaque `Binary` blobs. -fn no_dict_layout() -> Arc { +fn compressor_builder(compaction: CompactionStrategy) -> BtrBlocksCompressorBuilder { + match compaction { + CompactionStrategy::Compact => BtrBlocksCompressorBuilder::default().with_compact(), + CompactionStrategy::Default => BtrBlocksCompressorBuilder::default(), + } +} + +fn profiled_layout( + compaction: CompactionStrategy, + compressor: BtrBlocksCompressorBuilder, +) -> Arc { + WriteStrategyBuilder::default() + .with_btrblocks_builder(compressor) + .with_probe_compressor(compressor_builder(compaction).build()) + .build() +} + +fn compressed_layout(builder: BtrBlocksCompressorBuilder) -> Arc { Arc::new(CompressingStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - BtrBlocksCompressorBuilder::default().build(), + builder.build(), )) } +/// A chunked + compressed layout that skips dictionary encoding for opaque `Binary` blobs. +fn no_dict_layout() -> Arc { + compressed_layout(BtrBlocksCompressorBuilder::default()) +} + /// Convert all Parquet files in a directory to Vortex format. /// /// This function reads Parquet files from `{input_path}/parquet/` and writes Vortex files to @@ -225,6 +304,16 @@ fn no_dict_layout() -> Arc { pub async fn convert_parquet_directory_to_vortex( input_path: &Path, compaction: CompactionStrategy, +) -> anyhow::Result<()> { + convert_parquet_directory_to_vortex_with_profile(input_path, compaction, WriteProfile::Default) + .await +} + +/// Convert all Parquet files in a directory using a benchmark-only write profile. +pub async fn convert_parquet_directory_to_vortex_with_profile( + input_path: &Path, + compaction: CompactionStrategy, + profile: WriteProfile, ) -> anyhow::Result<()> { let (format, dir_name) = match compaction { CompactionStrategy::Compact => (Format::VortexCompact, Format::VortexCompact.name()), @@ -264,8 +353,13 @@ pub async fn convert_parquet_directory_to_vortex( "Processing file '{filename}' with {:?} strategy", compaction ); - convert_parquet_file_to_vortex(&parquet_file_path, &vtx_file, compaction) - .await + convert_parquet_file_to_vortex_with_profile( + &parquet_file_path, + &vtx_file, + compaction, + profile, + ) + .await }) .await .expect("Failed to write Vortex file") diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 255d2fbd075..5972029d997 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -64,3 +64,8 @@ test = false name = "compress_listview" harness = false test = false + +[[bench]] +name = "predicate_strings" +harness = false +test = false diff --git a/vortex-btrblocks/benches/predicate_strings.rs b/vortex-btrblocks/benches/predicate_strings.rs new file mode 100644 index 00000000000..ea3180ffb25 --- /dev/null +++ b/vortex-btrblocks/benches/predicate_strings.rs @@ -0,0 +1,267 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! String size-versus-predicate-operation calibration benchmark. + +#![expect(clippy::unwrap_used)] + +use std::sync::Arc; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::scalar_fn::fns::like::Like; +use vortex_array::scalar_fn::fns::like::LikeOptions; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_btrblocks::ArrayAndStats; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::Cost; +use vortex_btrblocks::CostModel; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::SchemeId; +use vortex_btrblocks::SizeCost; +use vortex_btrblocks::schemes::string::FSSTScheme; +use vortex_btrblocks::schemes::string::StringDictScheme; +use vortex_compressor::cost::Candidate; +use vortex_session::VortexSession; + +const LEN: usize = 1 << 18; +const UNIQUE: usize = 3; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + vortex_fsst::initialize(&session); + session +}); + +static INPUT: LazyLock = LazyLock::new(|| { + VarBinViewArray::from_iter_str((0..LEN).map(|index| { + format!( + "tenant-{:06}-event-checkout-completed-region-us-east", + index % UNIQUE + ) + })) + .into_array() +}); + +static CANONICAL: LazyLock = LazyLock::new(|| LazyLock::force(&INPUT).clone()); + +static SIZE_WINNER: LazyLock = LazyLock::new(|| { + BtrBlocksCompressorBuilder::default() + .build() + .compress(&INPUT, &mut SESSION.create_execution_ctx()) + .unwrap() +}); + +static LIKE_WINNER: LazyLock = LazyLock::new(|| { + BtrBlocksCompressorBuilder::default() + .with_like_strings() + .build() + .compress(&INPUT, &mut SESSION.create_execution_ctx()) + .unwrap() +}); + +static FSST: LazyLock = LazyLock::new(|| compress_with_root(&INPUT, FSSTScheme.id())); +static DICT: LazyLock = + LazyLock::new(|| compress_with_root(&INPUT, StringDictScheme.id())); + +/// Forces one root family while preserving normal size-based descendant selection. +#[derive(Debug)] +struct ForceRootCost { + root: SchemeId, +} + +impl CostModel for ForceRootCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + if candidate.cascade().is_empty() && candidate.scheme_id() != self.root { + return None; + } + SizeCost.cost(candidate) + } + + fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost { + SizeCost.canonical_cost(data, n_values) + } +} + +fn compress_with_root(input: &ArrayRef, root: SchemeId) -> ArrayRef { + BtrBlocksCompressorBuilder::default() + .with_cost_model(Arc::new(ForceRootCost { root })) + .build() + .compress(input, &mut SESSION.create_execution_ctx()) + .unwrap() +} + +fn bench_decode(bencher: Bencher, array: &'static LazyLock) { + bencher + .counter(ItemsCount::new(LEN)) + .with_inputs(|| { + ( + LazyLock::force(array).clone(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(array, ctx)| { + divan::black_box(array.clone().execute::(ctx).unwrap()) + }); +} + +fn bench_eq(bencher: Bencher, array: &'static LazyLock) { + bencher + .counter(ItemsCount::new(LEN)) + .with_inputs(|| { + let lhs = LazyLock::force(array).clone(); + let rhs = + ConstantArray::new("tenant-131071-event-checkout-completed-region-us-east", LEN) + .into_array(); + (lhs, rhs, SESSION.create_execution_ctx()) + }) + .bench_refs(|(lhs, rhs, ctx)| { + divan::black_box( + lhs.clone() + .binary(rhs.clone(), Operator::Eq) + .unwrap() + .execute::(ctx) + .unwrap(), + ) + }); +} + +fn bench_like(bencher: Bencher, array: &'static LazyLock) { + bencher + .counter(ItemsCount::new(LEN)) + .with_inputs(|| { + let lhs = LazyLock::force(array).clone(); + let pattern = ConstantArray::new("%checkout-completed%", LEN).into_array(); + (lhs, pattern, SESSION.create_execution_ctx()) + }) + .bench_refs(|(lhs, pattern, ctx)| { + divan::black_box( + Like.try_new_array(LEN, LikeOptions::default(), [lhs.clone(), pattern.clone()]) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap(), + ) + }); +} + +#[divan::bench] +fn canonical_decode(bencher: Bencher) { + bench_decode(bencher, &CANONICAL); +} + +#[divan::bench] +fn canonical_eq(bencher: Bencher) { + bench_eq(bencher, &CANONICAL); +} + +#[divan::bench] +fn canonical_like(bencher: Bencher) { + bench_like(bencher, &CANONICAL); +} + +#[divan::bench] +fn fsst_decode(bencher: Bencher) { + bench_decode(bencher, &FSST); +} + +#[divan::bench] +fn fsst_eq(bencher: Bencher) { + bench_eq(bencher, &FSST); +} + +#[divan::bench] +fn fsst_like(bencher: Bencher) { + bench_like(bencher, &FSST); +} + +#[divan::bench] +fn dict_decode(bencher: Bencher) { + bench_decode(bencher, &DICT); +} + +#[divan::bench] +fn dict_eq(bencher: Bencher) { + bench_eq(bencher, &DICT); +} + +#[divan::bench] +fn dict_like(bencher: Bencher) { + bench_like(bencher, &DICT); +} + +#[divan::bench] +fn like_winner_decode(bencher: Bencher) { + bench_decode(bencher, &LIKE_WINNER); +} + +#[divan::bench] +fn like_winner_eq(bencher: Bencher) { + bench_eq(bencher, &LIKE_WINNER); +} + +#[divan::bench] +fn like_winner_like(bencher: Bencher) { + bench_like(bencher, &LIKE_WINNER); +} + +fn print_selection_sweep() { + const PROBE_LEN: usize = 1 << 14; + + for unique in [3, 32, 256, 2_048, 8_192] { + let input = VarBinViewArray::from_iter_str((0..PROBE_LEN).map(|index| { + format!( + "tenant-{:06}-event-checkout-completed-region-us-east", + index % unique + ) + })) + .into_array(); + let size = BtrBlocksCompressorBuilder::default() + .build() + .compress(&input, &mut SESSION.create_execution_ctx()) + .unwrap(); + let like = BtrBlocksCompressorBuilder::default() + .with_like_strings() + .build() + .compress(&input, &mut SESSION.create_execution_ctx()) + .unwrap(); + println!( + "unique={unique:5}: size={} bytes={:7}; like={} bytes={:7}", + size.encoding_id(), + size.nbytes(), + like.encoding_id(), + like.nbytes(), + ); + } +} + +fn main() { + print_selection_sweep(); + let input = LazyLock::force(&INPUT); + let size_winner = LazyLock::force(&SIZE_WINNER); + let like_winner = LazyLock::force(&LIKE_WINNER); + let fsst = LazyLock::force(&FSST); + let dict = LazyLock::force(&DICT); + println!( + "input={} bytes; size winner={} bytes={}; LIKE winner={} bytes={}; fsst={} dict={}", + input.nbytes(), + size_winner.encoding_id(), + size_winner.nbytes(), + like_winner.encoding_id(), + like_winner.nbytes(), + fsst.nbytes(), + dict.nbytes(), + ); + divan::main(); +} diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index e51fac9a9d5..97e65a5c132 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,10 +3,15 @@ //! Builder for configuring `BtrBlocksCompressor` instances. +use std::sync::Arc; + +use vortex_compressor::cost::CostModel; +use vortex_compressor::cost::SizeCost; use vortex_utils::aliases::hash_set::HashSet; use crate::BtrBlocksCompressor; use crate::CascadingCompressor; +use crate::PredicateStringCost; use crate::Scheme; use crate::SchemeExt; use crate::SchemeId; @@ -68,8 +73,9 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ /// Builder for creating configured [`BtrBlocksCompressor`] instances. /// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated -/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via +/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order and ranked by +/// [`SizeCost`]. Feature-gated schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added +/// explicitly via /// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the /// `zstd` feature is enabled. /// @@ -90,12 +96,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + cost_model: Arc, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + cost_model: Arc::new(SizeCost), } } } @@ -107,6 +115,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + cost_model: Arc::new(SizeCost), } } @@ -208,9 +217,33 @@ impl BtrBlocksCompressorBuilder { self } + /// Uses a custom model to rank compression candidates. + pub fn with_cost_model(mut self, cost_model: Arc) -> Self { + self.cost_model = cost_model; + self + } + + /// Optimizes string selection for a static mix of equality and `LIKE` predicates. + /// + /// This is an experimental, opt-in profile. All non-string candidate selection retains the + /// default size objective. + pub fn with_predicate_strings(self) -> Self { + self.with_cost_model(Arc::new(PredicateStringCost::default())) + } + + /// Optimizes string selection for scalar equality predicates. + pub fn with_equality_strings(self) -> Self { + self.with_cost_model(Arc::new(PredicateStringCost::equality())) + } + + /// Optimizes string selection for `LIKE` predicates and selective materialization. + pub fn with_like_strings(self) -> Self { + self.with_cost_model(Arc::new(PredicateStringCost::like())) + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + BtrBlocksCompressor(CascadingCompressor::new(self.schemes).with_cost_model(self.cost_model)) } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 37915b427b6..273d75836da 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -30,8 +30,8 @@ //! //! Each `Scheme` implementation declares whether it [`matches`](Scheme::matches) a given //! canonical form and, if so, estimates the compression ratio (often by compressing a ~1% -//! sample). There is no dynamic registry — the set of schemes is fixed at build time via -//! [`ALL_SCHEMES`]. +//! sample). The configured cost model ranks those estimates. There is no dynamic registry — the +//! set of schemes is fixed at build time via [`ALL_SCHEMES`]. //! //! Schemes can produce arrays that are themselves further compressed (e.g. FoR then BitPacking), //! up to [`MAX_CASCADE`] (3) layers deep. Descendant exclusion rules for of [`SchemeId`] prevents @@ -68,6 +68,7 @@ mod builder; mod canonical_compressor; +mod predicate_string_cost; /// Compression scheme implementations. pub mod schemes; @@ -76,8 +77,15 @@ pub mod schemes; pub use builder::ALL_SCHEMES; pub use builder::BtrBlocksCompressorBuilder; pub use canonical_compressor::BtrBlocksCompressor; +pub use predicate_string_cost::PredicateStringCost; pub use schemes::patches::compress_patches; pub use vortex_compressor::CascadingCompressor; +pub use vortex_compressor::cost::Cost; +pub use vortex_compressor::cost::CostModel; +pub use vortex_compressor::cost::OperationCosts; +pub use vortex_compressor::cost::OperationWeights; +pub use vortex_compressor::cost::SizeCost; +pub use vortex_compressor::cost::WorkloadCost; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; diff --git a/vortex-btrblocks/src/predicate_string_cost.rs b/vortex-btrblocks/src/predicate_string_cost.rs new file mode 100644 index 00000000000..569d5ae7dd5 --- /dev/null +++ b/vortex-btrblocks/src/predicate_string_cost.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A workload model for equality- and pattern-heavy string predicates. + +use vortex_array::arrays::Dict; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_compressor::cost::Candidate; +use vortex_compressor::cost::Cost; +use vortex_compressor::cost::CostModel; +use vortex_compressor::cost::OperationCosts; +use vortex_compressor::cost::OperationWeights; +use vortex_compressor::cost::SizeCost; +use vortex_compressor::cost::WorkloadCost; +use vortex_compressor::stats::ArrayAndStats; + +use crate::SchemeExt; +use crate::schemes::string::FSSTScheme; +use crate::schemes::string::StringDictScheme; + +/// Experimental cost model for string-heavy filtering workloads. +/// +/// The profile represents one scalar comparison, one `LIKE` evaluation, and a 10% chance of full +/// materialization per string value. Calibrations come from +/// `vortex-btrblocks/benches/predicate_strings.rs` and intentionally make FSST's specialized +/// pattern kernel compete with dictionary encoding's smaller representation. +/// +/// Non-string candidates retain [`SizeCost`] behavior, so the experiment does not globally alter +/// integer, float, or descendant selection. UTF-8 candidates without calibrations are rejected, +/// limiting the experiment to canonical, dictionary, and FSST representations. The default +/// compressor also remains size-oriented; callers opt into this model with +/// [`crate::BtrBlocksCompressorBuilder::with_predicate_strings`]. +#[derive(Debug, Clone)] +pub struct PredicateStringCost { + string_cost: WorkloadCost, +} + +impl Default for PredicateStringCost { + fn default() -> Self { + Self::new(OperationWeights { + full_decode: 0.1, + compare: 1.0, + like: 1.0, + }) + } +} + +impl PredicateStringCost { + /// Creates a calibrated UTF-8 cost model for the supplied operation mix. + /// + /// # Panics + /// + /// Panics if any operation weight is non-finite or negative. + pub fn new(weights: OperationWeights) -> Self { + let canonical = OperationCosts { + full_decode: 0.0, + compare: 2.0, + like: 4.8, + }; + + Self { + string_cost: WorkloadCost::new(20.0, weights, canonical, canonical) + .with_scheme_cost( + FSSTScheme.id(), + OperationCosts { + full_decode: 5.7, + compare: 1.25, + like: 6.7, + }, + ) + .with_scheme_cost( + StringDictScheme.id(), + OperationCosts { + // Conservative fallback when a dictionary estimate has no sample. + full_decode: 14.3, + compare: 2.47, + like: 22.0, + }, + ), + } + } + + /// Creates a profile for scalar equality predicates. + pub fn equality() -> Self { + Self::new(OperationWeights { + compare: 1.0, + ..Default::default() + }) + } + + /// Creates a profile for `LIKE` predicates followed by selective materialization. + pub fn like() -> Self { + Self::new(OperationWeights { + full_decode: 0.1, + like: 1.0, + ..Default::default() + }) + } +} + +impl CostModel for PredicateStringCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + if !candidate.array().dtype().is_utf8() { + return SizeCost.cost(candidate); + } + + let scheme_id = candidate.scheme_id(); + if scheme_id != StringDictScheme.id() && scheme_id != FSSTScheme.id() { + return None; + } + + if scheme_id == StringDictScheme.id() + && let Some(dict) = candidate + .sampled() + .and_then(|sample| sample.as_opt::()) + { + let distinct_fraction = dict.values().len() as f64 / dict.len() as f64; + return self.string_cost.cost_with_operations( + candidate, + OperationCosts { + // Dictionary decode and predicates pay a fixed codes scan plus work that + // scales with the dictionary cardinality. + full_decode: 8.7 + 5.6 * distinct_fraction, + compare: 0.27 + 2.2 * distinct_fraction, + like: 13.4 + 8.6 * distinct_fraction, + }, + ); + } + + self.string_cost.cost(candidate) + } + + fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost { + if data.array().dtype().is_utf8() { + self.string_cost.canonical_cost(data, n_values) + } else { + SizeCost.canonical_cost(data, n_values) + } + } +} diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index c1473146607..59c75af3484 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,6 +17,7 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::BtrBlocksCompressorBuilder; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -46,6 +47,31 @@ fn test_dict_compressed() -> VortexResult<()> { Ok(()) } +#[test] +fn query_profiles_split_low_cardinality_string_selection() -> VortexResult<()> { + let array = || { + VarBinViewArray::from_iter_str((0..1 << 14).map(|index| { + format!( + "tenant-{:06}-event-checkout-completed-region-us-east", + index % 3 + ) + })) + .into_array() + }; + let equality = BtrBlocksCompressorBuilder::default() + .with_equality_strings() + .build() + .compress(&array(), &mut SESSION.create_execution_ctx())?; + let like = BtrBlocksCompressorBuilder::default() + .with_like_strings() + .build() + .compress(&array(), &mut SESSION.create_execution_ctx())?; + + assert!(equality.is::()); + assert!(like.is::()); + Ok(()) +} + #[cfg(feature = "unstable_encodings")] #[test] fn test_onpair_in_default_scheme_list() { @@ -87,7 +113,6 @@ fn test_onpair_compressed() -> VortexResult<()> { /// FSST-only builder still produces an FSST array. #[test] fn test_fsst_in_default_scheme_list() -> VortexResult<()> { - use crate::BtrBlocksCompressorBuilder; use crate::SchemeExt; use crate::schemes::string::FSSTScheme; diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs index fda8c2c0c86..6ac3ad9eea3 100644 --- a/vortex-compressor/src/cost/mod.rs +++ b/vortex-compressor/src/cost/mod.rs @@ -47,5 +47,10 @@ mod model; pub use model::Cost; pub use model::CostModel; +mod workload; +pub use workload::OperationCosts; +pub use workload::OperationWeights; +pub use workload::WorkloadCost; + mod size; pub use size::SizeCost; diff --git a/vortex-compressor/src/cost/workload.rs b/vortex-compressor/src/cost/workload.rs new file mode 100644 index 00000000000..65d3c5eaa84 --- /dev/null +++ b/vortex-compressor/src/cost/workload.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Deterministic operation-aware workload costs. + +use vortex_utils::aliases::hash_map::HashMap; + +use crate::cost::Candidate; +use crate::cost::Cost; +use crate::cost::CostModel; +use crate::scheme::SchemeId; +use crate::stats::ArrayAndStats; + +/// Expected operation counts per value for a static workload. +/// +/// The weights need not sum to one. For example, a profile that expects one comparison and one +/// `LIKE` evaluation per value uses `compare = 1.0` and `like = 1.0`. +#[derive(Debug, Clone, Copy, Default)] +pub struct OperationWeights { + /// Expected full materializations per value. + pub full_decode: f64, + /// Expected scalar comparisons per value. + pub compare: f64, + /// Expected string pattern matches per value. + pub like: f64, +} + +/// Estimated nanoseconds per value for operations on one representation. +#[derive(Debug, Clone, Copy, Default)] +pub struct OperationCosts { + /// Nanoseconds per value to fully materialize the representation. + pub full_decode: f64, + /// Nanoseconds per value to compare the representation with a scalar. + pub compare: f64, + /// Nanoseconds per value to evaluate a string pattern against the representation. + pub like: f64, +} + +impl OperationCosts { + /// Returns the weighted cost for one value. + fn weighted(self, weights: OperationWeights) -> f64 { + self.full_decode * weights.full_decode + + self.compare * weights.compare + + self.like * weights.like + } + + /// Returns whether every cost is finite and non-negative. + fn is_valid(self) -> bool { + [self.full_decode, self.compare, self.like] + .into_iter() + .all(|cost| cost.is_finite() && cost >= 0.0) + } +} + +impl OperationWeights { + /// Returns whether every weight is finite and non-negative. + fn is_valid(self) -> bool { + [self.full_decode, self.compare, self.like] + .into_iter() + .all(|weight| weight.is_finite() && weight >= 0.0) + } +} + +/// Prices compression candidates for a fixed operation mix. +/// +/// A candidate is priced as: +/// +/// `estimated_bytes / effective_bandwidth + values * weighted_operation_ns_per_value` +/// +/// `effective_bandwidth` is bytes per nanosecond (numerically equal to decimal GB/s). Operation +/// costs are deterministic calibration data, never timings collected while writing a file. +/// Unknown schemes use the configured fallback costs. +/// +/// This model prices the scheme currently being selected. Because the compressor invokes the +/// model recursively, the same workload policy also applies to descendant selection. It does not +/// yet search multiple completed encoding trees or model cross-column query dependencies. +#[derive(Debug, Clone)] +pub struct WorkloadCost { + /// Assumed storage throughput in bytes per nanosecond. + effective_bandwidth: f64, + /// Expected operation counts per value. + weights: OperationWeights, + /// Operation costs for the canonical representation. + canonical_costs: OperationCosts, + /// Operation costs for schemes without a calibration. + fallback_costs: OperationCosts, + /// Per-scheme operation calibrations. + scheme_costs: HashMap, +} + +impl WorkloadCost { + /// Creates an operation-aware cost model. + /// + /// # Panics + /// + /// Panics if bandwidth is not finite and positive, or if any weight or operation cost is + /// non-finite or negative. + pub fn new( + effective_bandwidth: f64, + weights: OperationWeights, + canonical_costs: OperationCosts, + fallback_costs: OperationCosts, + ) -> Self { + assert!( + effective_bandwidth.is_finite() && effective_bandwidth > 0.0, + "effective bandwidth must be finite and positive" + ); + assert!( + (u64::MAX as f64 / effective_bandwidth).is_finite(), + "effective bandwidth is too small to produce finite costs" + ); + assert!( + weights.is_valid(), + "operation weights must be finite and non-negative" + ); + assert!( + canonical_costs.is_valid(), + "canonical operation costs must be finite and non-negative" + ); + assert!( + fallback_costs.is_valid(), + "fallback operation costs must be finite and non-negative" + ); + assert!( + canonical_costs.weighted(weights).is_finite() + && fallback_costs.weighted(weights).is_finite(), + "weighted operation costs must be finite" + ); + + Self { + effective_bandwidth, + weights, + canonical_costs, + fallback_costs, + scheme_costs: HashMap::default(), + } + } + + /// Sets the calibrated operation costs for one scheme. + /// + /// # Panics + /// + /// Panics if any cost is non-finite or negative. + pub fn with_scheme_cost(mut self, scheme: SchemeId, costs: OperationCosts) -> Self { + assert!( + costs.is_valid(), + "scheme operation costs must be finite and non-negative" + ); + assert!( + costs.weighted(self.weights).is_finite(), + "weighted scheme operation cost must be finite" + ); + self.scheme_costs.insert(scheme, costs); + self + } + + /// Returns the configured effective bandwidth in bytes per nanosecond. + pub fn effective_bandwidth(&self) -> f64 { + self.effective_bandwidth + } + + /// Prices a candidate with operation costs derived from candidate-specific evidence. + /// + /// This is useful when a representation's execution cost depends on its sampled shape, such + /// as dictionary predicates scaling with the number of dictionary values. + /// + /// # Panics + /// + /// Panics if any operation cost is non-finite or negative. + pub fn cost_with_operations( + &self, + candidate: &Candidate<'_>, + operation_costs: OperationCosts, + ) -> Option { + assert!( + operation_costs.is_valid(), + "operation costs must be finite and non-negative" + ); + let weighted_operation_cost = operation_costs.weighted(self.weights); + assert!( + weighted_operation_cost.is_finite(), + "weighted operation cost must be finite" + ); + + let ratio = candidate.estimate().estimated_compression_ratio()?; + if !ratio.is_finite() || ratio.is_subnormal() || ratio <= 1.0 { + return None; + } + + let estimated_nbytes = candidate.input_nbytes() as f64 / ratio; + let cost = estimated_nbytes / self.effective_bandwidth + + candidate.n_values() as f64 * weighted_operation_cost; + cost.is_finite().then(|| Cost::new(cost)) + } +} + +impl CostModel for WorkloadCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + let operation_costs = self + .scheme_costs + .get(&candidate.scheme_id()) + .copied() + .unwrap_or(self.fallback_costs); + self.cost_with_operations(candidate, operation_costs) + } + + fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost { + Cost::new( + data.array().nbytes() as f64 / self.effective_bandwidth + + n_values as f64 * self.canonical_costs.weighted(self.weights), + ) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + use crate::CascadingCompressor; + use crate::scheme::CandidateEstimate; + use crate::scheme::CompressorContext; + use crate::scheme::Scheme; + use crate::scheme::SchemeEvaluation; + use crate::scheme::SchemeExt; + use crate::stats::GenerateStatsOptions; + + #[derive(Debug)] + struct TestScheme; + + impl Scheme for TestScheme { + fn scheme_name(&self) -> &'static str { + "test.workload_cost" + } + + fn matches(&self, _canonical: &Canonical) -> bool { + true + } + + fn evaluate( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> SchemeEvaluation { + SchemeEvaluation::Skip + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper is never selected") + } + } + + fn test_data() -> ArrayAndStats { + ArrayAndStats::new( + PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(), + GenerateStatsOptions::default(), + ) + } + + #[test] + fn combines_io_and_weighted_operation_time() { + let data = test_data(); + let model = WorkloadCost::new( + 4.0, + OperationWeights { + compare: 1.0, + like: 0.5, + ..Default::default() + }, + OperationCosts::default(), + OperationCosts::default(), + ) + .with_scheme_cost( + TestScheme.id(), + OperationCosts { + compare: 0.5, + like: 1.0, + ..Default::default() + }, + ); + let candidate = Candidate::new( + &TestScheme, + CandidateEstimate::from_compression_ratio(2.0), + &data, + None, + &[], + ); + + // 16 input bytes / ratio 2 / 4 bytes/ns + 4 values * (0.5 + 0.5 * 1.0). + assert_eq!(model.cost(&candidate).map(Cost::value), Some(6.0)); + } + + #[test] + fn canonical_includes_operation_time() { + let data = test_data(); + let model = WorkloadCost::new( + 4.0, + OperationWeights { + like: 0.5, + ..Default::default() + }, + OperationCosts { + like: 2.0, + ..Default::default() + }, + OperationCosts::default(), + ); + + // 16 bytes / 4 bytes/ns + 4 values * 0.5 expected LIKEs * 2 ns/value. + assert_eq!(model.canonical_cost(&data, 4).value(), 8.0); + } + + #[test] + fn can_price_candidate_above_canonical() { + let data = test_data(); + let model = WorkloadCost::new( + 4.0, + OperationWeights { + like: 1.0, + ..Default::default() + }, + OperationCosts::default(), + OperationCosts { + like: 1.0, + ..Default::default() + }, + ); + let candidate = Candidate::new( + &TestScheme, + CandidateEstimate::from_compression_ratio(2.0), + &data, + None, + &[], + ); + + assert!( + model + .cost(&candidate) + .is_some_and(|cost| cost >= model.canonical_cost(&data, candidate.n_values())) + ); + } +} diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 2703020906c..08753d8c411 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -141,8 +141,15 @@ pub mod buffer { pub mod compressor { pub use vortex_btrblocks::BtrBlocksCompressor; pub use vortex_btrblocks::BtrBlocksCompressorBuilder; + pub use vortex_btrblocks::Cost; + pub use vortex_btrblocks::CostModel; + pub use vortex_btrblocks::OperationCosts; + pub use vortex_btrblocks::OperationWeights; + pub use vortex_btrblocks::PredicateStringCost; pub use vortex_btrblocks::Scheme; pub use vortex_btrblocks::SchemeId; + pub use vortex_btrblocks::SizeCost; + pub use vortex_btrblocks::WorkloadCost; } /// Logical Vortex data types.