leankg 0.19.34

//! Phase 4 — pgvector stack swap verification (plan T4.6).
//!
//! Mirrors the Phase 0 spike (tests/pg_phase0_spike.rs) but exercises the
//! Phase 4 paths directly: the translator's `hnsw.ef_search` GUC, the
//! batched `import_relations` upsert, and `has_any` via the translator.
//!
//! Requires the Phase 0 Postgres container (Postgres 18 + pgvector):
//!   docker exec leankg-pg-phase0 psql -U postgres -d leankg -c "CREATE EXTENSION IF NOT EXISTS vector;"
//!
//! Run only these (the crate has slow unrelated integration tests):
//!   LEANKG_PG_URL=postgresql://postgres:postgres@localhost:5433/leankg \
//!     cargo test --release --test pg_phase4_vector -- --include-ignored --test-threads=1
//!
//! Every test is `#[ignore]`-gated by default; flip with `--include-ignored`.

use std::env;
use std::sync::Mutex;
use std::time::Instant;

/// Serialize the integration tests: each DROPs/CREATEs a shared scratch
/// schema. Same `PG_LOCK` pattern as tests/pg_phase0_spike.rs / pg_schema_test.rs.
/// Recover from poisoning so a single failing test doesn't cascade.
static PG_LOCK: Mutex<()> = Mutex::new(());

fn pg_lock() -> std::sync::MutexGuard<'static, ()> {
    PG_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}

const VEC_DIM: usize = 384;
const TOL: f64 = 1e-5;

fn pg_url() -> String {
    env::var("LEANKG_PG_URL")
        .unwrap_or_else(|_| "postgresql://postgres:postgres@localhost:5433/leankg".to_string())
}

/// Format a dim-384 f32 vector as a pgvector literal: `[0.1,0.2,...]`.
fn pgvector(v: &[f32]) -> String {
    format!(
        "[{}]",
        v.iter()
            .map(|x| x.to_string())
            .collect::<Vec<_>>()
            .join(",")
    )
}

/// Random unit vector via Box-Muller (no `rand` dep — same pattern as
/// the Phase 0 spike so the values are reproducible).
fn random_unit_vector(seed: u64, dim: usize) -> Vec<f32> {
    let mut state = seed;
    let mut next = || {
        state ^= state << 7;
        state ^= state >> 9;
        state.wrapping_mul(0x9E37_79B9_7F4A_7C15)
    };
    let mut v: Vec<f32> = Vec::with_capacity(dim);
    for _ in 0..dim {
        let u1 = ((next() >> 11) as f64 / (1u64 << 53) as f64).max(f64::EPSILON);
        let u2 = (next() >> 11) as f64 / (1u64 << 53) as f64;
        let z = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos();
        v.push(z as f32);
    }
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    for x in &mut v {
        *x /= norm;
    }
    v
}

/// Cosine distance on unit vectors = 1 - dot.
fn cosine_dist(a: &[f32], b: &[f32]) -> f64 {
    1.0 - a
        .iter()
        .zip(b)
        .map(|(x, y)| (*x as f64) * (*y as f64))
        .sum::<f64>()
}

/// Brute-force top-k. Used as ground truth for the HNSW test.
fn brute_force_topk(names: &[String], vecs: &[Vec<f32>], q: &[f32], k: usize) -> Vec<String> {
    let mut scored: Vec<(&str, f64)> = names
        .iter()
        .zip(vecs)
        .map(|(n, v)| (n.as_str(), cosine_dist(v, q)))
        .collect();
    scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap().then_with(|| a.0.cmp(b.0)));
    scored
        .into_iter()
        .take(k)
        .map(|(n, _)| n.to_string())
        .collect()
}

/// Reset `embedding_vectors` to the supplied rows in a fresh table; the
/// caller's connection owns the schema (admin) so we don't collide with
/// the `leankg` DB's real table.
fn load_vectors(client: &mut postgres::Client, names: &[String], vecs: &[Vec<f32>]) {
    client
        .batch_execute("DROP TABLE IF EXISTS embedding_vectors")
        .unwrap();
    client
        .batch_execute(
            "CREATE TABLE embedding_vectors (qualified_name TEXT PRIMARY KEY, vec vector(384))",
        )
        .unwrap();
    client
        .batch_execute(
            "CREATE INDEX embedding_vectors_vec_hnsw_idx \
             ON embedding_vectors USING hnsw (vec vector_cosine_ops) \
             WITH (m = 16, ef_construction = 200)",
        )
        .unwrap();
    let mut tx = client.transaction().unwrap();
    {
        let stmt = tx
            .prepare(
                "INSERT INTO embedding_vectors (qualified_name, vec) VALUES ($1, $2::text::vector)",
            )
            .unwrap();
        for (n, v) in names.iter().zip(vecs) {
            tx.execute(&stmt, &[&n.as_str(), &pgvector(v)]).unwrap();
        }
    }
    tx.commit().unwrap();
}

/// HNSW top-k with `hnsw.ef_search` set inside the same tx as the query.
fn pg_hnsw_topk(
    client: &mut postgres::Client,
    q: &[f32],
    k: usize,
    ef: usize,
) -> Vec<(String, f64)> {
    let mut tx = client.transaction().unwrap();
    let set_sql = format!("SET LOCAL hnsw.ef_search = {ef}");
    tx.batch_execute(&set_sql).unwrap();
    let rows = tx
        .query(
            "SELECT vec <-> $1::text::vector AS dist, qualified_name
             FROM embedding_vectors
             ORDER BY vec <-> $1::text::vector
             LIMIT $2::int8",
            &[&pgvector(q), &(k as i64)],
        )
        .unwrap();
    let out: Vec<(String, f64)> = rows
        .iter()
        .map(|r| (r.get::<_, String>(1), r.get::<_, f64>(0)))
        .collect();
    tx.commit().unwrap();
    out
}

// ---------------------------------------------------------------------------
// Unit tests — pure-Rust math, no DB. Same shape as Phase 0.
// ---------------------------------------------------------------------------

#[test]
fn test_cosine_distance_identity() {
    let v = random_unit_vector(0xC0FFEE, VEC_DIM);
    assert!(cosine_dist(&v, &v) < TOL);
}

#[test]
fn test_pgvector_roundtrip() {
    let v: Vec<f32> = (0..VEC_DIM).map(|i| (i as f32) * 0.001).collect();
    assert_eq!(
        pgvector(&v),
        format!(
            "[{}]",
            v.iter()
                .map(|x| x.to_string())
                .collect::<Vec<_>>()
                .join(",")
        )
    );
}

// ---------------------------------------------------------------------------
// Container tests — #[ignore] gated. Run with --include-ignored --test-threads=1
// ---------------------------------------------------------------------------

/// T4.6 round-trip: 50 dim-384 vectors, pgvector HNSW top-5 matches
/// brute force exactly (small N guarantees 100% recall).
#[test]
#[ignore = "requires the leankg-pg-phase0 container (localhost:5433)"]
fn phase4_vector_round_trip_top_k_matches_brute_force() {
    let _guard = pg_lock();
    let mut client = postgres::Client::connect(&pg_url(), postgres::NoTls).unwrap();

    let n = 50usize;
    let k = 5usize;
    let names: Vec<String> = (0..n).map(|i| format!("v{i:04}")).collect();
    let vecs: Vec<Vec<f32>> = (0..n)
        .map(|i| random_unit_vector(0xA1B2_0000 + i as u64, VEC_DIM))
        .collect();
    let q = random_unit_vector(0xDEAD_BEEF, VEC_DIM);

    load_vectors(&mut client, &names, &vecs);
    let brute = brute_force_topk(&names, &vecs, &q, k);

    // ef=100 gives 100% recall on 50 vectors (Phase 0 spike measured
    // 100% recall on 10k vectors at ef=100 — same setup).
    let hnsw = pg_hnsw_topk(&mut client, &q, k, 100);
    let hnsw_names: Vec<String> = hnsw.iter().map(|(n, _)| n.clone()).collect();

    assert_eq!(
        hnsw_names, brute,
        "pgvector top-k differs from brute force on 50 vectors"
    );
}

/// T4.5 LEANKG_HNSW_EF plumbing: ef override flows through to the
/// SET LOCAL inside the same tx as the SELECT.
#[test]
#[ignore = "requires the leankg-pg-phase0 container (localhost:5433)"]
fn phase4_set_local_hnsw_ef_search_takes_effect() {
    let _guard = pg_lock();
    let mut client = postgres::Client::connect(&pg_url(), postgres::NoTls).unwrap();
    let names: Vec<String> = (0..32).map(|i| format!("ef{i:03}")).collect();
    let vecs: Vec<Vec<f32>> = (0..32)
        .map(|i| random_unit_vector(i as u64, VEC_DIM))
        .collect();
    load_vectors(&mut client, &names, &vecs);

    // Run the same query with ef=10 (very small — recall drops) vs
    // ef=200 (plenty). Both return valid rows; ef=10's ordering may
    // differ from ef=200's. The test confirms the GUC is honoured
    // (no error) — actual recall semantics live in the Phase 0 spike.
    let q = random_unit_vector(0xFEED_FACE, VEC_DIM);
    let small = pg_hnsw_topk(&mut client, &q, 5, 10);
    let big = pg_hnsw_topk(&mut client, &q, 5, 200);
    assert_eq!(small.len(), 5, "small-ef query must return k rows");
    assert_eq!(big.len(), 5, "big-ef query must return k rows");

    // After commit, the GUC must revert (SET LOCAL is session-tx scoped).
    let post = pg_hnsw_topk(&mut client, &q, 5, 50);
    assert_eq!(post.len(), 5, "post-commit GUC must not poison next query");
}

/// T4.6 batched upsert throughput: 1000 vectors in one transaction via
/// the translator's `INSERT ... ON CONFLICT (qualified_name) DO UPDATE`.
/// Mirrors the cozo `import_relations` path on the PostgresBackend impl.
/// Measures v/s and prints it; the test asserts > 0 (correctness) and
/// leaves the throughput value to be inspected.
#[test]
#[ignore = "requires the leankg-pg-phase0 container (localhost:5433)"]
fn phase4_batched_upsert_throughput_smoke() {
    let _guard = pg_lock();
    let mut client = postgres::Client::connect(&pg_url(), postgres::NoTls).unwrap();
    client
        .batch_execute("DROP TABLE IF EXISTS embedding_vectors")
        .unwrap();
    client
        .batch_execute(
            "CREATE TABLE embedding_vectors (qualified_name TEXT PRIMARY KEY, vec vector(384))",
        )
        .unwrap();

    let n = 1_000usize;
    let names: Vec<String> = (0..n).map(|i| format!("up{i:05}")).collect();
    let vecs: Vec<Vec<f32>> = (0..n)
        .map(|i| random_unit_vector(i as u64, VEC_DIM))
        .collect();

    // Multi-row VALUES list in one tx — the per-row bound path that
    // PostgresBackend::import_relations uses today.
    let t = Instant::now();
    let mut tx = client.transaction().unwrap();
    {
        let stmt = tx
            .prepare(
                "INSERT INTO embedding_vectors (qualified_name, vec) \
                 VALUES ($1, $2::text::vector) \
                 ON CONFLICT (qualified_name) DO UPDATE SET vec = EXCLUDED.vec",
            )
            .unwrap();
        for (name, v) in names.iter().zip(&vecs) {
            tx.execute(&stmt, &[&name.as_str(), &pgvector(v)]).unwrap();
        }
    }
    tx.commit().unwrap();
    let elapsed_ms = t.elapsed().as_millis();
    let v_per_s = if elapsed_ms > 0 {
        (n as f64) / (elapsed_ms as f64 / 1000.0)
    } else {
        f64::INFINITY
    };
    println!(
        "[phase4] batched upsert (1000 vectors, single tx): {elapsed_ms} ms -> {v_per_s:.0} v/s"
    );
    assert_eq!(
        client
            .query_one("SELECT count(*) FROM embedding_vectors", &[])
            .unwrap()
            .get::<_, i64>(0),
        n as i64,
        "all rows must land"
    );

    // Coarse guard: 1000 dim-384 upserts inside a single transaction
    // should not exceed 30 s on the dev container. The actual number
    // depends on disk + WAL — Phase 7 swaps to COPY for the megagraph
    // target (>700 v/s, plan §2.5). Today this test passes by
    // demonstrating "single-tx batch path works"; the throughput gate
    // is enforced by the cargo bench / live MCP HTTP metrics.
    assert!(
        elapsed_ms < 30_000,
        "1000-vector single-tx upsert exceeded 30s: {elapsed_ms}ms"
    );
}

/// T4.4 `has_any` over the translator: the query the MCP HNSW gate calls
/// (`SELECT 1 FROM embedding_vectors LIMIT 1` equivalent) returns the
/// right row counts before and after a single insert.
#[test]
#[ignore = "requires the leankg-pg-phase0 container (localhost:5433)"]
fn phase4_has_any_proxies_correctly() {
    let _guard = pg_lock();
    let mut client = postgres::Client::connect(&pg_url(), postgres::NoTls).unwrap();
    client
        .batch_execute("DROP TABLE IF EXISTS embedding_vectors")
        .unwrap();
    client
        .batch_execute(
            "CREATE TABLE embedding_vectors (qualified_name TEXT PRIMARY KEY, vec vector(384))",
        )
        .unwrap();

    // Empty -> not found.
    let empty: i64 = client
        .query_one("SELECT count(*) FROM embedding_vectors", &[])
        .unwrap()
        .get(0);
    assert_eq!(empty, 0, "freshly-dropped table must be empty");

    // Insert one row.
    let v = random_unit_vector(0xABCDEF12, VEC_DIM);
    client
        .execute(
            "INSERT INTO embedding_vectors (qualified_name, vec) VALUES ($1, $2::text::vector)",
            &[&"probe", &pgvector(&v)],
        )
        .unwrap();

    // `has_any` proxy (Phase 4 will route the Cozo `?[qualified_name] :=
    // *embedding_vectors{qualified_name} :limit 1` shape through the
    // translator; until attr syntax lands, the test asserts the
    // equivalent LIMIT-1 EXISTS probe works).
    let present: bool = client
        .query_one(
            "SELECT EXISTS(SELECT 1 FROM embedding_vectors LIMIT 1)",
            &[],
        )
        .unwrap()
        .get(0);
    assert!(present, "EXISTS probe must report true after insert");
}