ivf_rq_bench

ivf_rq_bench.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! benchmark for batched streaming search PR (#7642):
//! 1M x 1024d f32 vectors, IVF_RQ with default index params (1024 partitions
//! = sqrt(1M), num_bits=1, fast rotation), index prewarmed, plain ANN queries
//! at concurrency 1 and (num_cores - 2).
//!
//! ```bash
//! cargo build --release --example ivf_rq_bench
//! ./target/release/examples/ivf_rq_bench /path/to/dataset-dir   # builds on first run
//! ```

use std::sync::Arc;
use std::time::{Duration, Instant};

use arrow_array::Float32Array;
use arrow_array::types::{Float32Type, Int32Type};
use futures::{StreamExt, TryStreamExt, stream};
use lance::Dataset;
use lance::index::DatasetIndexExt;
use lance::index::vector::VectorIndexParams;
use lance_datagen::{BatchCount, Dimension, RowCount, array, gen_batch};
use lance_index::IndexType;
use lance_linalg::distance::MetricType;

const DIM: u32 = 1024;
const NUM_ROWS: u64 = 1_000_000;
const ROWS_PER_BATCH: u64 = 10_000;
const NUM_PARTITIONS: usize = 1024; // sqrt(1M): the default heuristic
const RQ_NUM_BITS: u8 = 1; // RQBuildParams::default()
const NUM_QUERIES: usize = 1024;
const WARMUP_QUERIES: usize = 32;
const K: usize = 10;

async fn build_dataset(uri: &str) {
    let start = Instant::now();
    eprintln!("building dataset at {uri} ({NUM_ROWS} x {DIM}d) ...");
    let data = gen_batch()
        .col("i", array::step::<Int32Type>())
        .col("vec", array::rand_vec::<Float32Type>(Dimension::from(DIM)))
        .into_reader_rows(
            RowCount::from(ROWS_PER_BATCH),
            BatchCount::from((NUM_ROWS / ROWS_PER_BATCH) as u32),
        );
    let mut dataset = Dataset::write(data, uri, None).await.unwrap();
    eprintln!("data written in {:.0?}, building IVF_RQ index ...", start.elapsed());
    let start = Instant::now();
    let params = VectorIndexParams::ivf_rq(NUM_PARTITIONS, RQ_NUM_BITS, MetricType::L2);
    dataset
        .create_index(&["vec"], IndexType::Vector, None, &params, true)
        .await
        .unwrap();
    eprintln!("index built in {:.0?}", start.elapsed());
}

fn make_query(seed: usize) -> Float32Array {
    let mut state = seed as u64 ^ 0x9E37_79B9_7F4A_7C15;
    Float32Array::from(
        (0..DIM)
            .map(|_| {
                state = state
                    .wrapping_mul(6364136223846793005)
                    .wrapping_add(1442695040888963407);
                ((state >> 33) as f32) / (u32::MAX as f32)
            })
            .collect::<Vec<_>>(),
    )
}

async fn run_query(dataset: Arc<Dataset>, seed: usize) -> Duration {
    let query = make_query(seed);
    let start = Instant::now();
    let mut scan = dataset.scan();
    scan.nearest("vec", &query, K).unwrap();
    let batches = scan
        .try_into_stream()
        .await
        .unwrap()
        .try_collect::<Vec<_>>()
        .await
        .unwrap();
    let elapsed = start.elapsed();
    let num_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
    assert_eq!(num_rows, K, "expected k results");
    elapsed
}

async fn bench_concurrency(dataset: &Arc<Dataset>, concurrency: usize) {
    // Warmup at this concurrency.
    stream::iter(0..WARMUP_QUERIES)
        .map(|i| run_query(dataset.clone(), i))
        .buffer_unordered(concurrency)
        .collect::<Vec<_>>()
        .await;

    let wall = Instant::now();
    let mut latencies = stream::iter(0..NUM_QUERIES)
        .map(|i| run_query(dataset.clone(), WARMUP_QUERIES + i))
        .buffer_unordered(concurrency)
        .collect::<Vec<_>>()
        .await;
    let wall = wall.elapsed();

    latencies.sort_unstable();
    let pct = |p: f64| latencies[((latencies.len() as f64 * p) as usize).min(latencies.len() - 1)];
    println!(
        "concurrency={concurrency} queries={NUM_QUERIES} wall={:.2}s qps={:.1} \
         p50={:.2}ms p90={:.2}ms p99={:.2}ms",
        wall.as_secs_f64(),
        NUM_QUERIES as f64 / wall.as_secs_f64(),
        pct(0.50).as_secs_f64() * 1e3,
        pct(0.90).as_secs_f64() * 1e3,
        pct(0.99).as_secs_f64() * 1e3,
    );
}

#[tokio::main]
async fn main() {
    let uri = std::env::args()
        .nth(1)
        .expect("usage: ivf_rq_bench <dataset_dir>");
    if !std::path::Path::new(&format!("{uri}/_versions")).exists() {
        build_dataset(&uri).await;
    }

    let dataset = Arc::new(Dataset::open(&uri).await.unwrap());
    let indices = dataset.load_indices().await.unwrap();
    let index_name = indices.first().expect("index must exist").name.clone();
    let start = Instant::now();
    dataset.prewarm_index(&index_name).await.unwrap();
    eprintln!("prewarmed index '{index_name}' in {:.0?}", start.elapsed());

    let num_cores = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    let high_concurrency = num_cores.saturating_sub(2).max(2);
    for concurrency in [1, high_concurrency] {
        bench_concurrency(&dataset, concurrency).await;
    }
}