add vector-search-bench crate

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

Cargo.lock

@@ -59,6 +59,7 @@ members = [
     "benchmarks/datafusion-bench",
     "benchmarks/duckdb-bench",
     "benchmarks/random-access-bench",
+    "benchmarks/vector-search-bench",
 ]
 exclude = ["java/testfiles", "wasm-test"]
 resolver = "2"

Cargo.toml

@@ -0,0 +1,40 @@
+[package]
+name = "vector-search-bench"
+description = "Vector similarity search benchmarks for Vortex on public embedding datasets"
+authors.workspace = true
+categories.workspace = true
+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
+publish = false
+
+[dependencies]
+anyhow = { workspace = true }
+arrow-array = { workspace = true }
+arrow-buffer = { workspace = true }
+arrow-schema = { workspace = true }
+clap = { workspace = true, features = ["derive"] }
+futures = { workspace = true }
+indicatif = { workspace = true }
+parquet = { workspace = true, features = ["async"] }
+rand = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+tabled = { workspace = true, features = ["std"] }
+tokio = { workspace = true, features = ["full"] }
+tracing = { workspace = true }
+vortex = { workspace = true, features = ["files", "tokio", "unstable_encodings"] }
+vortex-bench = { workspace = true, features = ["unstable_encodings"] }
+vortex-btrblocks = { workspace = true, features = ["unstable_encodings"] }
+vortex-tensor = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }
+
+[lints]
+workspace = true

benchmarks/vector-search-bench/Cargo.toml

@@ -0,0 +1,103 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Vector compression flavors exercised by the benchmark.
+//!
+//! Each [`VectorFlavor`] variant maps to a [`vortex::file::WriteStrategyBuilder`] configuration
+//! applied to the same input data.
+//!
+//! The benchmark writes one `.vortex` file per flavor per data file, then scans them all with the
+//! same query so the comparison is apples-to-apples with the Parquet files.
+//!
+//! Note that the handrolled `&[f32]` parquet baseline is **not** a flavor here.
+
+use clap::ValueEnum;
+use vortex::array::ArrayId;
+use vortex::array::scalar_fn::ScalarFnVTable;
+use vortex::file::ALLOWED_ENCODINGS;
+use vortex::file::VortexWriteOptions;
+use vortex::file::WriteOptionsSessionExt;
+use vortex::file::WriteStrategyBuilder;
+use vortex::session::VortexSession;
+use vortex::utils::aliases::hash_set::HashSet;
+use vortex_bench::Format;
+use vortex_btrblocks::BtrBlocksCompressorBuilder;
+use vortex_tensor::scalar_fns::l2_denorm::L2Denorm;
+use vortex_tensor::scalar_fns::sorf_transform::SorfTransform;
+
+/// Every [`VectorFlavor`] variant in CLI-help order.
+pub const ALL_VECTOR_FLAVORS: &[VectorFlavor] =
+    &[VectorFlavor::Uncompressed, VectorFlavor::TurboQuant];
+
+/// One write-side compression configuration we measure.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, ValueEnum)]
+pub enum VectorFlavor {
+    /// `BtrBlocksCompressorBuilder::empty()`
+    #[clap(name = "vortex-uncompressed")]
+    Uncompressed,
+    /// `BtrBlocksCompressorBuilder::default().with_turboquant()`.
+    #[clap(name = "vortex-turboquant")]
+    TurboQuant,
+    // TODO(connor): We will want to add `Default` here which is just the default compressor.
+}
+
+impl VectorFlavor {
+    /// Stable kebab-cased label used in CLI args and metric names.
+    pub fn label(&self) -> &'static str {
+        match self {
+            VectorFlavor::Uncompressed => "vortex-uncompressed",
+            VectorFlavor::TurboQuant => "vortex-turboquant",
+        }
+    }
+
+    /// The `target.format` value emitted on measurements for this flavor. Both flavors produce
+    /// `.vortex` files, so the compression label carries the flavor split.
+    pub fn as_format(&self) -> Format {
+        match self {
+            VectorFlavor::Uncompressed => Format::OnDiskVortex,
+            VectorFlavor::TurboQuant => Format::OnDiskVortex,
+        }
+    }
+
+    /// Subdirectory name under the per-dataset cache root used to store this flavor's `.vortex`
+    /// files.
+    pub fn dir_name(&self) -> &'static str {
+        match self {
+            VectorFlavor::Uncompressed => "vortex-uncompressed",
+            VectorFlavor::TurboQuant => "vortex-turboquant",
+        }
+    }
+
+    /// Build the [`vortex::file::WriteStrategyBuilder`]-backed write options for this flavor.
+    ///
+    /// TurboQuant produces `L2Denorm(SorfTransform(...))` which the default file
+    /// `ALLOWED_ENCODINGS` set rejects on normalization — we extend the allow-list with the two
+    /// scalar-fn array IDs the scheme actually emits.
+    pub fn create_write_options(&self, session: &VortexSession) -> VortexWriteOptions {
+        let strategy = match self {
+            VectorFlavor::Uncompressed => {
+                let compressor = BtrBlocksCompressorBuilder::empty().build();
+
+                WriteStrategyBuilder::default()
+                    .with_compressor(compressor)
+                    .build()
+            }
+            VectorFlavor::TurboQuant => {
+                let compressor = BtrBlocksCompressorBuilder::default()
+                    .with_turboquant()
+                    .build();
+
+                let mut allowed: HashSet<ArrayId> = ALLOWED_ENCODINGS.clone();
+                allowed.insert(L2Denorm.id());
+                allowed.insert(SorfTransform.id());
+
+                WriteStrategyBuilder::default()
+                    .with_compressor(compressor)
+                    .with_allow_encodings(allowed)
+                    .build()
+            }
+        };
+
+        session.write_options().with_strategy(strategy)
+    }
+}

benchmarks/vector-search-bench/src/compression.rs

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Cosine-similarity filter [`Expression`]s used by the file-scan path.
+//!
+//! We can easily build a cosine similarity filter by hand:
+//!
+//! ```text
+//! gt(
+//!     cosine_similarity(col("emb"), lit(query_scalar)),
+//!     lit(threshold),
+//! )
+//! ```
+//!
+//! The query is wrapped as `Scalar::extension::<Vector>(Scalar::fixed_size_list(F32, ...))` so
+//! [`CosineSimilarity`] can treat it as a single-row `Vector` value during evaluation.
+//!
+//! At scan time the literal expands into a `ConstantArray` whose row count matches the chunk batch
+//! size.
+
+use anyhow::Result;
+use vortex::array::expr::Expression;
+use vortex::array::expr::col;
+use vortex::array::expr::gt;
+use vortex::array::expr::lit;
+use vortex::array::extension::EmptyMetadata;
+use vortex::array::scalar::Scalar;
+use vortex::array::scalar_fn::EmptyOptions;
+use vortex::array::scalar_fn::ScalarFnVTableExt;
+use vortex::dtype::DType;
+use vortex::dtype::Nullability;
+use vortex::dtype::PType;
+use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity;
+use vortex_tensor::vector::Vector;
+
+/// Build the filter `cosine_similarity(emb, query) > threshold`.
+pub fn similarity_filter(query: &[f32], threshold: f32) -> Result<Expression> {
+    // Empty queries short-circuit to a literal `false`, so scans return no rows instead of trying
+    // to evaluate cosine similarity on a zero-dimensional vector.
+    if query.is_empty() {
+        return Ok(lit(false));
+    }
+
+    let query_lit = lit(query_scalar(query)?);
+    let cosine = CosineSimilarity.new_expr(EmptyOptions, [col("emb"), query_lit]);
+    Ok(gt(cosine, lit(threshold)))
+}
+
+/// Wrap a query vector as `Scalar::extension::<Vector>(Scalar::fixed_size_list(F32, ...))`.
+pub fn query_scalar(query: &[f32]) -> Result<Scalar> {
+    let children: Vec<Scalar> = query
+        .iter()
+        .map(|&v| Scalar::primitive(v, Nullability::NonNullable))
+        .collect();
+
+    let element_dtype = DType::Primitive(PType::F32, Nullability::NonNullable);
+    let fsl = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
+
+    Ok(Scalar::extension::<Vector>(EmptyMetadata, fsl))
+}
+
+/// Project just the `emb` column. Used by the throughput-only scan path.
+pub fn emb_projection() -> Expression {
+    col("emb")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn query_scalar_accepts_empty_query() {
+        let scalar = query_scalar(&[]).unwrap();
+        match scalar.dtype() {
+            DType::Extension(_) => {}
+            other => panic!("expected Extension, got {other}"),
+        }
+    }
+
+    #[test]
+    fn query_scalar_builds_extension_dtype() {
+        let scalar = query_scalar(&[1.0, 0.0, 0.0]).unwrap();
+        match scalar.dtype() {
+            DType::Extension(_) => {}
+            other => panic!("expected Extension, got {other}"),
+        }
+    }
+
+    #[test]
+    fn similarity_filter_uses_gt_operator() {
+        let expr = similarity_filter(&[1.0, 0.0, 0.0], 0.5).unwrap();
+        // Quick sanity check: the printed form contains the operator and the threshold so
+        // future refactors that change the structure get caught here.
+        let printed = format!("{expr:?}");
+        assert!(printed.contains("Gt") || printed.contains(">"), "{printed}");
+    }
+}