yams/tests/benchmarks/vector_backend_engine_compare.cpp

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <filesystem>
#include <numeric>
#include <random>
#include <stdexcept>
#include <string>
#include <unordered_set>
#include <vector>

#include <spdlog/spdlog.h>

#include "bench_utils.h"
#include <yams/vector/sqlite_vec_backend.h>

using namespace yams::vector;

namespace {

struct Config {
    size_t corpus = 5000;
    size_t queries = 100;
    size_t dim = 128;
    size_t k = 10;
    unsigned int seed = 42;
    bool spq_sweep = false;
    bool filter_cell = false;
};

struct EngineResult {
    const char* name = "";
    double build_ms = 0.0;
    double mean_us = 0.0;
    double p95_us = 0.0;
    double p99_us = 0.0;
    double qps = 0.0;
    double recall = 0.0;
};

Config parseArgs(int argc, char* argv[]) {
    Config cfg;
    for (int i = 1; i < argc; ++i) {
        std::string_view arg(argv[i]);
        auto parse = [&](std::string_view prefix, auto& out) {
            if (arg.rfind(prefix, 0) == 0) {
                out = static_cast<std::decay_t<decltype(out)>>(
                    std::stoull(std::string(arg.substr(prefix.size()))));
                return true;
            }
            return false;
        };
        if (parse("--corpus=", cfg.corpus) || parse("--queries=", cfg.queries) ||
            parse("--dim=", cfg.dim) || parse("--k=", cfg.k) || parse("--seed=", cfg.seed)) {
            continue;
        }
        if (arg == "--spq-sweep") {
            cfg.spq_sweep = true;
            continue;
        }
        if (arg == "--filter-cell") {
            cfg.filter_cell = true;
            continue;
        }
        if (arg == "--help") {
            std::printf("Usage: %s [--corpus=N] [--queries=N] [--dim=N] [--k=N] [--seed=N] "
                        "[--spq-sweep] [--filter-cell]\n",
                        argv[0]);
            std::exit(0);
        }
    }
    if (cfg.queries == 0) {
        throw std::invalid_argument("--queries must be >= 1");
    }
    if (cfg.k == 0) {
        throw std::invalid_argument("--k must be >= 1");
    }
    return cfg;
}

std::vector<float> generateNormalizedVector(size_t dim, std::mt19937& rng) {
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    std::vector<float> vec(dim);
    float norm_sq = 0.0f;
    for (auto& value : vec) {
        value = dist(rng);
        norm_sq += value * value;
    }
    float norm = std::sqrt(norm_sq);
    if (norm > 0.0f) {
        for (auto& value : vec) {
            value /= norm;
        }
    }
    return vec;
}

std::vector<std::vector<float>> generateCorpus(size_t count, size_t dim, std::mt19937& rng) {
    std::vector<std::vector<float>> corpus;
    corpus.reserve(count);
    for (size_t i = 0; i < count; ++i) {
        corpus.push_back(generateNormalizedVector(dim, rng));
    }
    return corpus;
}

std::vector<size_t> bruteForceTopK(const std::vector<float>& query,
                                   const std::vector<std::vector<float>>& corpus, size_t k) {
    std::vector<std::pair<size_t, float>> scored;
    scored.reserve(corpus.size());
    for (size_t i = 0; i < corpus.size(); ++i) {
        float dot = 0.0f;
        for (size_t j = 0; j < query.size(); ++j) {
            dot += query[j] * corpus[i][j];
        }
        scored.emplace_back(i, dot);
    }

    const size_t top = std::min(k, scored.size());
    std::partial_sort(scored.begin(), scored.begin() + top, scored.end(),
                      [](const auto& lhs, const auto& rhs) { return lhs.second > rhs.second; });

    std::vector<size_t> ids;
    ids.reserve(top);
    for (size_t i = 0; i < top; ++i) {
        ids.push_back(scored[i].first);
    }
    return ids;
}

double percentile(std::vector<double> values, double p) {
    return yams::bench::percentileNearestRank(std::move(values), p);
}

EngineResult runEngine(VectorSearchEngine engine, const Config& cfg,
                       const std::vector<std::vector<float>>& corpus,
                       const std::vector<std::vector<float>>& queries,
                       const std::vector<std::vector<size_t>>& ground_truth,
                       size_t spq_subquantizers = 0, size_t spq_rerank_factor = 0) {
    SqliteVecBackend::Config backend_cfg;
    backend_cfg.embedding_dim = cfg.dim;
    backend_cfg.search_engine = engine;
    if (spq_subquantizers > 0) {
        backend_cfg.simeon_pq_subquantizers = spq_subquantizers;
    }
    if (spq_rerank_factor > 0) {
        backend_cfg.simeon_pq_rerank_factor = spq_rerank_factor;
    }

    SqliteVecBackend backend(backend_cfg);
    auto init = backend.initialize(":memory:");
    if (!init) {
        throw std::runtime_error(init.error().message);
    }
    auto create = backend.createTables(cfg.dim);
    if (!create) {
        throw std::runtime_error(create.error().message);
    }

    std::vector<VectorRecord> records;
    records.reserve(corpus.size());
    for (size_t i = 0; i < corpus.size(); ++i) {
        VectorRecord rec;
        rec.chunk_id = "chunk_" + std::to_string(i);
        rec.document_hash = "doc_" + std::to_string(i / 4);
        rec.embedding = corpus[i];
        rec.content = "benchmark content " + std::to_string(i);
        records.push_back(std::move(rec));
    }
    auto insert = backend.insertVectorsBatch(records);
    if (!insert) {
        throw std::runtime_error(insert.error().message);
    }

    const auto build_start = std::chrono::steady_clock::now();
    auto build = backend.buildIndex();
    if (!build) {
        throw std::runtime_error(build.error().message);
    }
    const auto build_end = std::chrono::steady_clock::now();

    // Exclude one-time lazy index initialization from steady-state query timing. Run the same
    // unmeasured query for every engine so Vec0, Simeon PQ, and exact scan enter the measured loop
    // from the same query-ready boundary.
    auto warmup = backend.searchSimilar(queries.front(), cfg.k, 0.0f, std::nullopt, {});
    if (!warmup) {
        throw std::runtime_error(warmup.error().message);
    }

    std::vector<double> latencies_us;
    latencies_us.reserve(queries.size());
    size_t total_hits = 0;

    for (size_t i = 0; i < queries.size(); ++i) {
        const auto start = std::chrono::steady_clock::now();
        auto result = backend.searchSimilar(queries[i], cfg.k, 0.0f, std::nullopt, {});
        const auto end = std::chrono::steady_clock::now();
        if (!result) {
            throw std::runtime_error(result.error().message);
        }

        latencies_us.push_back(std::chrono::duration<double, std::micro>(end - start).count());

        std::unordered_set<std::string> expected;
        for (size_t idx : ground_truth[i]) {
            expected.insert("chunk_" + std::to_string(idx));
        }
        for (const auto& rec : result.value()) {
            total_hits += expected.contains(rec.chunk_id) ? 1 : 0;
        }
    }

    const double total_us = std::accumulate(latencies_us.begin(), latencies_us.end(), 0.0);
    EngineResult out;
    switch (engine) {
        case VectorSearchEngine::SimeonPqAdc:
            out.name = "simeon-pq";
            break;
        case VectorSearchEngine::Vec0L2:
            out.name = "vec0-l2";
            break;
        case VectorSearchEngine::ExactScan:
            out.name = "exact-scan";
            break;
    }
    out.build_ms = std::chrono::duration<double, std::milli>(build_end - build_start).count();
    out.mean_us = total_us / static_cast<double>(latencies_us.size());
    out.p95_us = percentile(latencies_us, 0.95);
    out.p99_us = percentile(latencies_us, 0.99);
    out.qps = out.mean_us > 0.0 ? (1e6 / out.mean_us) : 0.0;
    out.recall = static_cast<double>(total_hits) / static_cast<double>(queries.size() * cfg.k);
    return out;
}

void printResult(const EngineResult& result, size_t k) {
    std::printf("%-12s build=%8.1f ms mean=%8.1f us p95=%8.1f us p99=%8.1f us qps=%8.0f "
                "recall@%zu=%5.1f%%\n",
                result.name, result.build_ms, result.mean_us, result.p95_us, result.p99_us,
                result.qps, k, result.recall * 100.0);
}

} // namespace

int main(int argc, char* argv[]) {
    spdlog::set_level(spdlog::level::warn);

    try {
        const Config cfg = parseArgs(argc, argv);
        std::mt19937 rng(cfg.seed);
        auto corpus = generateCorpus(cfg.corpus, cfg.dim, rng);
        auto queries = generateCorpus(cfg.queries, cfg.dim, rng);

        std::vector<std::vector<size_t>> ground_truth;
        ground_truth.reserve(queries.size());
        for (const auto& query : queries) {
            ground_truth.push_back(bruteForceTopK(query, corpus, cfg.k));
        }

        std::printf("============================================================\n");
        std::printf("YAMS Vector Backend Engine Compare\n");
        std::printf("============================================================\n");
        std::printf("corpus=%zu queries=%zu dim=%zu k=%zu seed=%u\n", cfg.corpus, cfg.queries,
                    cfg.dim, cfg.k, cfg.seed);
        std::printf("note: corpus/query vectors are normalized, so cosine and L2 rankings are "
                    "comparable\n\n");

        if (cfg.filter_cell) {
            std::printf("filtered-search cell: routed Simeon PQ/ADC versus exact allowed-set "
                        "control\n");
            SqliteVecBackend::Config backend_cfg;
            backend_cfg.embedding_dim = cfg.dim;
            backend_cfg.search_engine = VectorSearchEngine::SimeonPqAdc;
            SqliteVecBackend backend(backend_cfg);
            if (!backend.initialize(":memory:")) {
                throw std::runtime_error("init failed");
            }
            if (!backend.createTables(cfg.dim)) {
                throw std::runtime_error("createTables failed");
            }
            std::vector<VectorRecord> records;
            records.reserve(corpus.size());
            for (size_t i = 0; i < corpus.size(); ++i) {
                VectorRecord rec;
                rec.chunk_id = "chunk_" + std::to_string(i);
                rec.document_hash = "doc_" + std::to_string(i / 4);
                rec.embedding = corpus[i];
                rec.content = "benchmark content " + std::to_string(i);
                records.push_back(std::move(rec));
            }
            if (!backend.insertVectorsBatch(records)) {
                throw std::runtime_error("insert failed");
            }
            const auto buildStart = std::chrono::steady_clock::now();
            if (!backend.buildIndex()) {
                throw std::runtime_error("buildIndex failed");
            }
            const auto buildEnd = std::chrono::steady_clock::now();
            std::printf("index_build_ms=%.1f vectors=%zu documents=%zu chunks_per_document=4\n",
                        std::chrono::duration<double, std::milli>(buildEnd - buildStart).count(),
                        cfg.corpus, (cfg.corpus + 3) / 4);

            const size_t documentCount = (cfg.corpus + 3) / 4;
            const std::vector<size_t> candidateCounts = {2UL, 7UL, 16UL, 64UL, documentCount};
            for (const size_t nCandidates : candidateCounts) {
                const size_t candidateCount = std::min(nCandidates, documentCount);
                std::unordered_set<std::string> candidates;
                for (size_t i = 0; i < candidateCount; ++i) {
                    candidates.insert("doc_" + std::to_string(i));
                }

                VectorSearchDiagnostics warmupDiagnostics;
                auto warmup = backend.searchSimilarWithDiagnostics(
                    queries.front(), cfg.k, -1.0F, std::nullopt, candidates, {}, warmupDiagnostics);
                if (!warmup) {
                    throw std::runtime_error(warmup.error().message);
                }

                std::vector<double> annLatencies;
                std::vector<double> exactLatencies;
                annLatencies.reserve(queries.size());
                exactLatencies.reserve(queries.size());
                std::uint64_t lookupNs = 0;
                std::uint64_t projectionNs = 0;
                std::uint64_t lutNs = 0;
                std::uint64_t adcNs = 0;
                std::uint64_t selectionNs = 0;
                std::uint64_t materializationNs = 0;
                std::uint64_t rerankNs = 0;
                size_t candidateIndexPayloadBytes = 0;
                size_t recallHits = 0;
                size_t recallDenominator = 0;
                for (const auto& q : queries) {
                    VectorSearchDiagnostics annDiagnostics;
                    const auto annStart = std::chrono::steady_clock::now();
                    auto ann = backend.searchSimilarWithDiagnostics(q, cfg.k, -1.0F, std::nullopt,
                                                                    candidates, {}, annDiagnostics);
                    const auto annEnd = std::chrono::steady_clock::now();
                    if (!ann) {
                        throw std::runtime_error(ann.error().message);
                    }
                    VectorSearchDiagnostics exactDiagnostics;
                    const auto exactStart = std::chrono::steady_clock::now();
                    auto exact = backend.searchExactCandidatesWithDiagnostics(
                        q, cfg.k, -1.0F, candidates, exactDiagnostics);
                    const auto exactEnd = std::chrono::steady_clock::now();
                    if (!exact) {
                        throw std::runtime_error(exact.error().message);
                    }

                    annLatencies.push_back(
                        std::chrono::duration<double, std::micro>(annEnd - annStart).count());
                    exactLatencies.push_back(
                        std::chrono::duration<double, std::micro>(exactEnd - exactStart).count());
                    lookupNs += annDiagnostics.candidateLookupNanoseconds;
                    projectionNs += annDiagnostics.candidateProjectionNanoseconds;
                    lutNs += annDiagnostics.pqLutNanoseconds;
                    adcNs += annDiagnostics.adcScoringNanoseconds;
                    selectionNs += annDiagnostics.topKSelectionNanoseconds;
                    materializationNs += annDiagnostics.resultMaterializationNanoseconds;
                    rerankNs += annDiagnostics.exactRerankNanoseconds;
                    candidateIndexPayloadBytes = std::max(
                        candidateIndexPayloadBytes, annDiagnostics.candidateIndexPayloadBytes);

                    std::unordered_set<std::string> exactIds;
                    for (const auto& record : exact.value()) {
                        exactIds.insert(record.chunk_id);
                    }
                    recallDenominator += exactIds.size();
                    for (const auto& record : ann.value()) {
                        recallHits += exactIds.contains(record.chunk_id) ? 1 : 0;
                    }
                }
                const double divisor = static_cast<double>(queries.size());
                const double annMean =
                    std::accumulate(annLatencies.begin(), annLatencies.end(), 0.0) / divisor;
                const double exactMean =
                    std::accumulate(exactLatencies.begin(), exactLatencies.end(), 0.0) / divisor;
                const auto meanUs = [divisor](std::uint64_t nanoseconds) {
                    return static_cast<double>(nanoseconds) / (1000.0 * divisor);
                };
                const double recall =
                    recallDenominator == 0
                        ? 1.0
                        : static_cast<double>(recallHits) / static_cast<double>(recallDenominator);
                std::printf(
                    "allowed_docs=%2zu allowed_vectors<=%3zu ann_mean=%8.1f us ann_p95=%8.1f us "
                    "exact_mean=%8.1f us recall=%5.1f%% phases_us{lookup=%.1f projection=%.1f "
                    "lut=%.1f adc=%.1f select=%.1f materialize=%.1f rerank=%.1f} "
                    "candidate_index_payload_kib=%.1f\n",
                    candidateCount, candidateCount * 4, annMean, percentile(annLatencies, 0.95),
                    exactMean, recall * 100.0, meanUs(lookupNs), meanUs(projectionNs),
                    meanUs(lutNs), meanUs(adcNs), meanUs(selectionNs), meanUs(materializationNs),
                    meanUs(rerankNs), static_cast<double>(candidateIndexPayloadBytes) / 1024.0);
            }
            return 0;
        }

        if (cfg.spq_sweep) {
            std::printf("simeon-pq sweep: subquantizers x rerank_factor\n");
            for (size_t m : {8UL, 16UL, 32UL, 64UL}) {
                if (cfg.dim % m != 0) {
                    continue;
                }
                for (size_t rerank : {2UL, 4UL, 8UL}) {
                    auto result = runEngine(VectorSearchEngine::SimeonPqAdc, cfg, corpus, queries,
                                            ground_truth, m, rerank);
                    std::printf("m=%-3zu rerank=%zu ", m, rerank);
                    printResult(result, cfg.k);
                }
            }
            return 0;
        }

        const auto spq =
            runEngine(VectorSearchEngine::SimeonPqAdc, cfg, corpus, queries, ground_truth);
        const auto vec0 = runEngine(VectorSearchEngine::Vec0L2, cfg, corpus, queries, ground_truth);

        printResult(spq, cfg.k);
        printResult(vec0, cfg.k);
        return 0;
    } catch (const std::exception& e) {
        std::fprintf(stderr, "Benchmark failed: %s\n", e.what());
        return 1;
    }
}