llm_kernel/embedding/

pgvector.rs

//! pgvector `AsyncVectorIndex` — PostgreSQL + the `pgvector` extension.
//!
//! `PgVectorIndex` implements [`AsyncVectorIndex`] over a `sqlx::PgPool`,
//! mirroring `qdrant`/`elastic` (async remote vector backend). Vectors live in
//! a `{table}(id BIGINT PK, vec <type>)` relation; search uses the cosine
//! distance operator `<=>`. Requires `CREATE EXTENSION vector` in the target DB
//! (the table + HNSW index are created automatically by [`PgVectorIndex::new`]).
//!
//! Two precisions: [`new`](PgVectorIndex::new) → `vector` (float32, dim ≤ 2000),
//! [`new_halfvec`](PgVectorIndex::new_halfvec) → `halfvec` (float16, ~half RAM,
//! dim ≤ 4000; needs pgvector ≥ 0.6).

use async_trait::async_trait;
use sqlx::PgPool;
use sqlx::QueryBuilder;
use sqlx::postgres::PgPoolOptions;

use crate::embedding::sparse::SparseVector;
use crate::embedding::vector_index::SearchHit;
use crate::error::{KernelError, Result};

/// 검색 결과 행(id + cosine 유사도) — 튜플 대신 구조체로 sqlx `FromRow` 안정 매핑.
#[derive(sqlx::FromRow)]
struct ScoreRow {
    id: i64,
    score: f64,
}

/// Row filter for `search_where` on [`PgVectorIndex`] and [`PgSparseVectorIndex`]:
/// a SQL boolean expression evaluated per candidate row, so the allowed set never
/// leaves the database. Prefer it to `search_filtered` when the allowed IDs are
/// many: an allowlist is fetched client-side and bound as one array per query.
///
/// `predicate` is interpolated verbatim, so it must be **trusted SQL written by
/// the caller, never user input**. It sees the vector row as `v` (`v.id`,
/// `v.vec`) and the text values of `binds` as `$3`, `$4`, ... (`$1` is the query
/// vector, `$2` the limit). Example: `EXISTS (SELECT 1 FROM chunks c WHERE
/// c.id = v.id AND c.category = $3)` with `binds: &["statute"]`.
#[derive(Debug, Clone, Copy)]
pub struct SqlFilter<'a> {
    /// Trusted SQL boolean expression over the vector row `v`.
    pub predicate: &'a str,
    /// Text values bound as `$3`, `$4`, ... in `predicate`.
    pub binds: &'a [&'a str],
}

/// Shared body of the `search_where` methods. `sql` uses `$1` (query literal)
/// and `$2` (limit); `filter.binds` follow as `$3..`.
async fn fetch_where(
    pool: &PgPool,
    sql: &str,
    q: String,
    k: usize,
    filter: &SqlFilter<'_>,
    what: &str,
) -> Result<Vec<SearchHit>> {
    let mut query = sqlx::query_as::<_, ScoreRow>(sql).bind(q).bind(k as i64);
    for b in filter.binds {
        query = query.bind(*b);
    }
    let rows = query
        .fetch_all(pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("{what}: {e}")))?;
    Ok(rows
        .into_iter()
        .map(|r| SearchHit {
            id: r.id as u64,
            score: r.score as f32,
        })
        .collect())
}

/// f32 슬라이스 → pgvector 문자열 리터럴 `[1,2,3]` (text → vector 입력 캐스트).
/// `pgvector::Vector`의 sqlx `Type` 바인드가 의존 환경에 따라 충돌해 문자열로 회피.
fn vec_literal(v: &[f32]) -> String {
    let mut s = String::from("[");
    for (i, f) in v.iter().enumerate() {
        if i > 0 {
            s.push(',');
        }
        s.push_str(&f.to_string());
    }
    s.push(']');
    s
}

/// Storage precision + HNSW tuning for [`PgVectorIndex`].
///
/// All fields default to pgvector's own defaults, so `PgVectorOpts::default()`
/// is equivalent to [`PgVectorIndex::new`].
#[derive(Debug, Clone, Default)]
pub struct PgVectorOpts {
    /// Store `halfvec` (float16) instead of `vector` (float32) — ~half the RAM.
    pub half: bool,
    /// HNSW `m` (edges per node). `None` → pgvector default (16). Higher `m`
    /// raises recall and index size.
    pub hnsw_m: Option<u32>,
    /// HNSW `ef_construction` (build-time candidate list). `None` → default
    /// (64). Higher values build a better graph, more slowly.
    pub hnsw_ef_construction: Option<u32>,
    /// `hnsw.ef_search` (query-time candidate list), applied to every pooled
    /// connection. `None` → default (40). Raise it for higher recall on large
    /// indexes; it is the main recall/latency knob at query time.
    pub hnsw_ef_search: Option<u32>,
}

/// PostgreSQL vector index backed by the `pgvector` extension.
///
/// All operations are async over a shared `PgPool` (MVCC, connection-pooled).
/// The relation named `table` holds `(id BIGINT PK, vec <type>(N))` rows where
/// `<type>` is `vector` (float32) or `halfvec` (float16) per the `half` flag.
pub struct PgVectorIndex {
    pool: PgPool,
    table: String,
    dim: usize,
    /// `true` → `halfvec` (float16, ~half RAM, recall 손실 ~0). `false` → `vector` (float32).
    half: bool,
    hnsw_m: Option<u32>,
    hnsw_ef_construction: Option<u32>,
}

impl PgVectorIndex {
    /// Connect to `url` (libpq connstring / `postgresql://…`), create the
    /// vector table + HNSW cosine index if missing, and return a ready index.
    ///
    /// `dim` is enforced by the fixed `vector(dim)` column; vectors whose
    /// length differs from `dim` are rejected by pgvector on insert. Uses
    /// full-precision `vector` (float32). For half-precision see
    /// [`new_halfvec`](Self::new_halfvec).
    pub async fn new(url: &str, table: &str, dim: usize) -> Result<Self> {
        Self::new_with_opts(url, table, dim, PgVectorOpts::default()).await
    }

    /// Half-precision variant — stores `halfvec` (float16): **~half the RAM** of
    /// `vector` with negligible recall loss for cosine similarity. Requires the
    /// `pgvector` extension ≥ 0.6 (`halfvec` type + `halfvec_cosine_ops`).
    /// `dim` ≤ 4000 (vs 2000 for `vector`).
    pub async fn new_halfvec(url: &str, table: &str, dim: usize) -> Result<Self> {
        Self::new_with_opts(
            url,
            table,
            dim,
            PgVectorOpts {
                half: true,
                ..Default::default()
            },
        )
        .await
    }

    /// Full form — pick storage precision and HNSW tuning via [`PgVectorOpts`].
    ///
    /// `hnsw_m` / `hnsw_ef_construction` are applied to the `CREATE INDEX` (they
    /// only affect a *newly* created index — an existing one keeps its build
    /// parameters). `hnsw_ef_search` is applied to every pooled connection and
    /// therefore takes effect immediately for queries.
    pub async fn new_with_opts(
        url: &str,
        table: &str,
        dim: usize,
        opts: PgVectorOpts,
    ) -> Result<Self> {
        validate_table_name(table)?;
        let mut pool_opts = PgPoolOptions::new().max_connections(8);
        if let Some(ef) = opts.hnsw_ef_search {
            // `ef` is a `u32`, so the interpolation has no injection surface.
            pool_opts = pool_opts.after_connect(move |conn, _meta| {
                Box::pin(async move {
                    sqlx::query(&format!("SET hnsw.ef_search = {ef}"))
                        .execute(conn)
                        .await?;
                    Ok(())
                })
            });
        }
        let pool = pool_opts
            .connect(url)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector connect: {e}")))?;
        let idx = Self {
            pool,
            table: table.to_string(),
            dim,
            half: opts.half,
            hnsw_m: opts.hnsw_m,
            hnsw_ef_construction: opts.hnsw_ef_construction,
        };
        idx.init_schema().await?;
        Ok(idx)
    }

    /// SQL vector type name for the active precision (`vector` | `halfvec`).
    fn sql_type(&self) -> &'static str {
        if self.half { "halfvec" } else { "vector" }
    }

    /// HNSW ops class for the active precision.
    fn ops_class(&self) -> &'static str {
        if self.half {
            "halfvec_cosine_ops"
        } else {
            "vector_cosine_ops"
        }
    }

    /// `CREATE TABLE IF NOT EXISTS {table} (id BIGINT PK, vec {vector|halfvec}(N))`
    /// + matching HNSW cosine index, per the active precision. Idempotent.
    async fn init_schema(&self) -> Result<()> {
        // Identifier is caller-controlled (not user input at runtime) and
        // validated in `new_with_opts`, so `format!` is acceptable here. PG cannot
        // bind identifiers. Callers pass a fixed, validated table name.
        let ty = self.sql_type();
        let ops = self.ops_class();
        let with = hnsw_with_clause(self.hnsw_m, self.hnsw_ef_construction);
        sqlx::query(&format!(
            "CREATE TABLE IF NOT EXISTS {} (id BIGINT PRIMARY KEY, vec {}({}) NOT NULL)",
            self.table, ty, self.dim
        ))
        .execute(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector create table: {e}")))?;
        sqlx::query(&format!(
            "CREATE INDEX IF NOT EXISTS idx_{}_vec ON {} USING hnsw (vec {}){}",
            self.table, self.table, ops, with
        ))
        .execute(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector hnsw index: {e}")))?;
        Ok(())
    }

    /// Underlying connection pool.
    ///
    /// Callers that need **cross-table transactional consistency** — e.g. pruning
    /// a law's chunks plus their vectors atomically in one transaction — can
    /// `pool().begin()` and run their own DML on the same pool the index uses.
    /// This preserves the table-name encapsulation while letting the caller
    /// coordinate a multi-statement transaction.
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Remove vectors by external IDs **within a caller-provided transaction**.
    ///
    /// Enables atomic cross-table deletes: begin a tx on [`pool`](Self::pool),
    /// delete related rows in other relations (chunks, edges, …), call this with
    /// the same `&mut PgConnection`, then `commit()`. A failure anywhere rolls
    /// back the whole set — no orphaned vectors or chunks. The table name stays
    /// encapsulated (`format!` over the validated identifier).
    pub async fn remove_in_tx(&self, tx: &mut sqlx::PgConnection, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
            .bind(&ids)
            .execute(tx)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector remove_in_tx: {e}")))?;
        Ok(())
    }

    /// Top-`k` by cosine similarity among rows matching `filter`, evaluated
    /// inside PostgreSQL (see [`SqlFilter`]). With pgvector >= 0.8
    /// `hnsw.iterative_scan` enabled the HNSW scan continues until `k` rows pass;
    /// without it a selective predicate can return fewer than `k`.
    pub async fn search_where(
        &self,
        query: &[f32],
        k: usize,
        filter: &SqlFilter<'_>,
    ) -> Result<Vec<SearchHit>> {
        let ty = self.sql_type();
        let sql = format!(
            "SELECT v.id, 1 - (v.vec <=> $1::{ty}) AS score FROM {} v WHERE ({}) \
             ORDER BY v.vec <=> $1::{ty}, v.id LIMIT $2",
            self.table, filter.predicate
        );
        fetch_where(
            &self.pool,
            &sql,
            vec_literal(query),
            k,
            filter,
            "pgvector search_where",
        )
        .await
    }
}

#[async_trait]
impl crate::embedding::AsyncVectorIndex for PgVectorIndex {
    async fn add(&self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
        if vectors.len() != ids.len() {
            return Err(KernelError::Embedding(format!(
                "vectors.len() ({}) must equal ids.len() ({})",
                vectors.len(),
                ids.len()
            )));
        }
        if vectors.is_empty() {
            return Ok(());
        }
        // Map u64 → i64 up front: pgvector stores BIGINT (i64), so values
        // above i64::MAX cannot be represented and must be rejected rather
        // than silently wrapped.
        let pg_ids: Vec<i64> = ids.iter().map(|&id| to_pg_id(id)).collect::<Result<_>>()?;

        // Each row binds 2 params (id, vec literal). PostgreSQL caps bound
        // parameters at u16::MAX (65535), so a single INSERT over ~32k vectors
        // overflows it. Chunk well under the limit; each chunk is its own
        // round trip but shares one prepared-statement shape.
        const ROWS_PER_CHUNK: usize = 16_000;
        for chunk in vectors
            .chunks(ROWS_PER_CHUNK)
            .zip(pg_ids.chunks(ROWS_PER_CHUNK))
        {
            let (chunk_vecs, chunk_ids) = chunk;
            let mut q = QueryBuilder::new("INSERT INTO ");
            q.push(self.table.as_str());
            q.push(" (id, vec) VALUES ");
            // pgvector `vec` 컬럼은 text 리터럴 입력 시 `::vector` 캐스트가 필수다.
            // `push_values` 는 값별 캐스트를 붙일 수 없어 수동으로 VALUES 튜플을 조립한다
            // (캐스트 누락 시 "column vec is of type vector but expression is of type text").
            for (i, (v, &id)) in chunk_vecs.iter().zip(chunk_ids.iter()).enumerate() {
                if i > 0 {
                    q.push(", ");
                }
                q.push("(");
                q.push_bind(id);
                q.push(", ");
                q.push_bind(vec_literal(v));
                q.push("::");
                q.push(self.sql_type());
                q.push(")");
            }
            q.push(" ON CONFLICT (id) DO UPDATE SET vec = EXCLUDED.vec");
            q.build()
                .execute(&self.pool)
                .await
                .map_err(|e| KernelError::Embedding(format!("pgvector add: {e}")))?;
        }
        Ok(())
    }

    async fn remove(&self, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
            .bind(&ids)
            .execute(&self.pool)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector remove: {e}")))?;
        Ok(())
    }

    async fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchHit>> {
        let q = vec_literal(query);
        let ty = self.sql_type();
        // cosine distance <=> : 0 (동일) .. 2 (반대). score = 1 - distance.
        // `ORDER BY ..., id` keeps ties deterministic so the per-branch rank order
        // that RRF feeds on is stable across branches (pgvector's own tie order
        // is otherwise index-scan dependent and unstable).
        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
            "SELECT id, 1 - (vec <=> $1::{ty}) AS score FROM {} \
             ORDER BY vec <=> $1::{ty}, id LIMIT $2",
            self.table
        ))
        .bind(q)
        .bind(k as i64)
        .fetch_all(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector search: {e}")))?;
        Ok(rows
            .into_iter()
            .map(|r| SearchHit {
                id: r.id as u64,
                score: r.score as f32,
            })
            .collect())
    }

    async fn search_filtered(
        &self,
        query: &[f32],
        k: usize,
        allowlist: &[u64],
    ) -> Result<Vec<SearchHit>> {
        if allowlist.is_empty() {
            return Ok(Vec::new());
        }
        let q = vec_literal(query);
        let ty = self.sql_type();
        let allow: Vec<i64> = allowlist
            .iter()
            .map(|&i| to_pg_id(i))
            .collect::<Result<_>>()?;
        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
            "SELECT id, 1 - (vec <=> $1::{ty}) AS score FROM {} WHERE id = ANY($2) \
             ORDER BY vec <=> $1::{ty}, id LIMIT $3",
            self.table
        ))
        .bind(q)
        .bind(&allow)
        .bind(k as i64)
        .fetch_all(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector search_filtered: {e}")))?;
        Ok(rows
            .into_iter()
            .map(|r| SearchHit {
                id: r.id as u64,
                score: r.score as f32,
            })
            .collect())
    }

    async fn len(&self) -> Result<usize> {
        let n: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {}", self.table))
            .fetch_one(&self.pool)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector len: {e}")))?;
        Ok(n as usize)
    }

    fn dim(&self) -> usize {
        self.dim
    }
}

/// `SparseVector` → pgvector `sparsevec` literal `{i:v,…}/dim`.
///
/// pgvector's text form uses **1-based** indices while models emit 0-based
/// vocabulary positions, so every index is shifted by one on the way out.
///
/// Weights use `f32::to_string`, which yields the shortest string that
/// round-trips back to the same `f32` (Rust's `Display` guarantee). pgvector
/// parses the literal as `float8` (f64), so no extra precision is lost in
/// transit beyond what `f32` already discards. The same holds for dense
/// [`vec_literal`].
fn sparse_literal(v: &SparseVector, dim: usize) -> String {
    let mut s = String::from("{");
    for (n, (idx, val)) in v.indices().iter().zip(v.values()).enumerate() {
        if n > 0 {
            s.push(',');
        }
        s.push_str(&(u64::from(*idx) + 1).to_string());
        s.push(':');
        s.push_str(&val.to_string());
    }
    s.push('}');
    s.push('/');
    s.push_str(&dim.to_string());
    s
}

/// PostgreSQL **sparse** vector index — the lexical half of hybrid retrieval.
///
/// Stores `sparsevec(N)` rows (`N` = vocabulary size, e.g. 250 002 for BGE-M3's
/// XLM-R vocabulary) and ranks by inner product, which is how learned-lexical
/// weights (BGE-M3 sparse, SPLADE) are scored. Deliberately *not* an
/// [`AsyncVectorIndex`](crate::embedding::AsyncVectorIndex) — that trait speaks
/// dense `Vec<f32>`. Pair it with a [`PgVectorIndex`] and combine the two
/// result lists with [`Fusion`](crate::embedding::Fusion).
///
/// Requires the `pgvector` extension ≥ 0.7 (`sparsevec` + `sparsevec_ip_ops`).
/// pgvector caps the number of non-zero elements it will index, so prune long
/// lexical vectors with [`SparseVector::prune_top_k`] before inserting.
pub struct PgSparseVectorIndex {
    pool: PgPool,
    table: String,
    dim: usize,
}

impl PgSparseVectorIndex {
    /// Connect and create the `sparsevec` table + HNSW inner-product index.
    ///
    /// `dim` is the vocabulary size and is enforced by the `sparsevec(dim)`
    /// column; vectors carrying an index ≥ `dim` are rejected on insert.
    pub async fn new(url: &str, table: &str, dim: usize) -> Result<Self> {
        validate_table_name(table)?;
        let pool = PgPoolOptions::new()
            .max_connections(8)
            .connect(url)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector sparse connect: {e}")))?;
        let idx = Self {
            pool,
            table: table.to_string(),
            dim,
        };
        sqlx::query(&format!(
            "CREATE TABLE IF NOT EXISTS {} (id BIGINT PRIMARY KEY, vec sparsevec({}) NOT NULL)",
            idx.table, idx.dim
        ))
        .execute(&idx.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector sparse create table: {e}")))?;
        sqlx::query(&format!(
            "CREATE INDEX IF NOT EXISTS idx_{}_vec ON {} USING hnsw (vec sparsevec_ip_ops)",
            idx.table, idx.table
        ))
        .execute(&idx.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector sparse hnsw index: {e}")))?;
        Ok(idx)
    }

    /// Vocabulary size of the indexed vectors.
    pub fn dim(&self) -> usize {
        self.dim
    }

    /// Underlying connection pool — see [`PgVectorIndex::pool`].
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Upsert sparse vectors by external ID.
    pub async fn add(&self, vectors: &[SparseVector], ids: &[u64]) -> Result<()> {
        if vectors.len() != ids.len() {
            return Err(KernelError::Embedding(format!(
                "vectors.len() ({}) must equal ids.len() ({})",
                vectors.len(),
                ids.len()
            )));
        }
        if vectors.is_empty() {
            return Ok(());
        }
        let pg_ids: Vec<i64> = ids.iter().map(|&id| to_pg_id(id)).collect::<Result<_>>()?;
        let mut q = QueryBuilder::new("INSERT INTO ");
        q.push(self.table.as_str());
        q.push(" (id, vec) VALUES ");
        for (i, (v, &id)) in vectors.iter().zip(pg_ids.iter()).enumerate() {
            if i > 0 {
                q.push(", ");
            }
            q.push("(");
            q.push_bind(id);
            q.push(", ");
            q.push_bind(sparse_literal(v, self.dim));
            q.push("::sparsevec)");
        }
        q.push(" ON CONFLICT (id) DO UPDATE SET vec = EXCLUDED.vec");
        q.build()
            .execute(&self.pool)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector sparse add: {e}")))?;
        Ok(())
    }

    /// Top-`k` by inner product. An empty query matches nothing by definition
    /// (its inner product with every row is zero), so it short-circuits.
    pub async fn search(&self, query: &SparseVector, k: usize) -> Result<Vec<SearchHit>> {
        if query.is_empty() {
            return Ok(Vec::new());
        }
        let q = sparse_literal(query, self.dim);
        // `<#>` yields the *negative* inner product, so ascending order is most
        // similar first; negate it back into a positive similarity score. The
        // `id` tie-break mirrors the dense index so RRF sees stable ranks.
        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
            "SELECT id, -(vec <#> $1::sparsevec) AS score FROM {} \
             ORDER BY vec <#> $1::sparsevec, id LIMIT $2",
            self.table
        ))
        .bind(q)
        .bind(k as i64)
        .fetch_all(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector sparse search: {e}")))?;
        Ok(rows
            .into_iter()
            .map(|r| SearchHit {
                id: r.id as u64,
                score: r.score as f32,
            })
            .collect())
    }

    /// Top-`k` by inner product, restricted to `allowlist`.
    pub async fn search_filtered(
        &self,
        query: &SparseVector,
        k: usize,
        allowlist: &[u64],
    ) -> Result<Vec<SearchHit>> {
        if query.is_empty() || allowlist.is_empty() {
            return Ok(Vec::new());
        }
        let q = sparse_literal(query, self.dim);
        let allow: Vec<i64> = allowlist
            .iter()
            .map(|&i| to_pg_id(i))
            .collect::<Result<_>>()?;
        let rows: Vec<ScoreRow> = sqlx::query_as(&format!(
            "SELECT id, -(vec <#> $1::sparsevec) AS score FROM {} WHERE id = ANY($2) \
             ORDER BY vec <#> $1::sparsevec, id LIMIT $3",
            self.table
        ))
        .bind(q)
        .bind(&allow)
        .bind(k as i64)
        .fetch_all(&self.pool)
        .await
        .map_err(|e| KernelError::Embedding(format!("pgvector sparse search_filtered: {e}")))?;
        Ok(rows
            .into_iter()
            .map(|r| SearchHit {
                id: r.id as u64,
                score: r.score as f32,
            })
            .collect())
    }

    /// Top-`k` by inner product among rows matching `filter`, evaluated inside
    /// PostgreSQL (see [`SqlFilter`]). An empty query matches nothing.
    pub async fn search_where(
        &self,
        query: &SparseVector,
        k: usize,
        filter: &SqlFilter<'_>,
    ) -> Result<Vec<SearchHit>> {
        if query.is_empty() {
            return Ok(Vec::new());
        }
        let sql = format!(
            "SELECT v.id, -(v.vec <#> $1::sparsevec) AS score FROM {} v WHERE ({}) \
             ORDER BY v.vec <#> $1::sparsevec, v.id LIMIT $2",
            self.table, filter.predicate
        );
        fetch_where(
            &self.pool,
            &sql,
            sparse_literal(query, self.dim),
            k,
            filter,
            "pgvector sparse search_where",
        )
        .await
    }

    /// Remove vectors by external ID.
    pub async fn remove(&self, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
            .bind(&ids)
            .execute(&self.pool)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector sparse remove: {e}")))?;
        Ok(())
    }

    /// Remove within a caller-provided transaction — see
    /// [`PgVectorIndex::remove_in_tx`]. Lets a dense index, a sparse index and
    /// the caller's own tables be pruned atomically together.
    pub async fn remove_in_tx(&self, tx: &mut sqlx::PgConnection, ids: &[u64]) -> Result<()> {
        if ids.is_empty() {
            return Ok(());
        }
        let ids: Vec<i64> = ids.iter().map(|&i| to_pg_id(i)).collect::<Result<_>>()?;
        sqlx::query(&format!("DELETE FROM {} WHERE id = ANY($1)", self.table))
            .bind(&ids)
            .execute(tx)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector sparse remove_in_tx: {e}")))?;
        Ok(())
    }

    /// Number of stored vectors.
    pub async fn len(&self) -> Result<usize> {
        let n: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {}", self.table))
            .fetch_one(&self.pool)
            .await
            .map_err(|e| KernelError::Embedding(format!("pgvector sparse len: {e}")))?;
        Ok(n as usize)
    }

    /// Whether the index holds no vectors.
    pub async fn is_empty(&self) -> Result<bool> {
        Ok(self.len().await? == 0)
    }
}

/// ` WITH (m = …, ef_construction = …)` for the HNSW index, or `""` when both
/// are unset (pgvector defaults: `m = 16`, `ef_construction = 64`). Values are
/// `u32`, so the interpolation has no injection surface.
fn hnsw_with_clause(m: Option<u32>, ef_construction: Option<u32>) -> String {
    let mut parts: Vec<String> = Vec::new();
    if let Some(m) = m {
        parts.push(format!("m = {m}"));
    }
    if let Some(ef) = ef_construction {
        parts.push(format!("ef_construction = {ef}"));
    }
    if parts.is_empty() {
        String::new()
    } else {
        format!(" WITH ({})", parts.join(", "))
    }
}

/// `u64` external ID → PG `BIGINT` (`i64`). IDs exceeding `i64::MAX` cannot
/// be stored in a BIGINT column — reject rather than silently wrap.
fn to_pg_id(id: u64) -> Result<i64> {
    i64::try_from(id).map_err(|_| KernelError::Embedding(format!("id {id} exceeds BIGINT range")))
}

/// Validate that `table` is a plain, safe SQL identifier (ASCII alphanumeric +
/// `_`, starting with a letter or `_`). It is interpolated into DDL/DML via
/// `format!` (PG cannot bind identifiers), so reject anything that could break
/// out of the identifier context.
fn validate_table_name(table: &str) -> Result<()> {
    let mut chars = table.chars();
    let first_ok = matches!(chars.next(), Some(c) if c.is_ascii_alphabetic() || c == '_');
    let valid = first_ok && chars.all(|c| c.is_ascii_alphanumeric() || c == '_');
    if !valid {
        return Err(KernelError::Embedding(format!(
            "invalid table identifier: {table:?}"
        )));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::embedding::AsyncVectorIndex;

    /// `LLMKERNEL_PG_URL` 미설정 시 자동 skip (graph-pg pg.rs 패턴).
    fn pg_url() -> Option<String> {
        std::env::var("LLMKERNEL_PG_URL").ok()
    }

    /// add → search → search_filtered → remove 라운드트립. HNSW는 대량 데이터에서
    /// 정확도가 보장되지만 소규모 테스트에선 정확 매칭이 간헐적일 수 있어
    /// 여기선 id/회수 위주로 검증.
    #[tokio::test]
    async fn roundtrip_add_search_remove() {
        let Some(url) = pg_url() else {
            eprintln!("skip pgvector test: LLMKERNEL_PG_URL unset");
            return;
        };
        let table = format!("lk_test_{}", line!());
        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");

        let vecs = vec![
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
        ];
        let ids = vec![10u64, 20, 30];
        idx.add(&vecs, &ids).await.expect("add");
        assert_eq!(idx.len().await.unwrap(), 3);

        // nearest to [1,0,0] → id 10 우선
        let hits = idx.search(&[1.0, 0.0, 0.0], 1).await.unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].id, 10);

        // filtered: allow [20,30] → 10 제외
        let hits = idx
            .search_filtered(&[1.0, 0.0, 0.0], 1, &[20, 30])
            .await
            .unwrap();
        assert_eq!(hits.len(), 1);
        assert_ne!(hits[0].id, 10);

        idx.remove(&[10]).await.unwrap();
        assert_eq!(idx.len().await.unwrap(), 2);

        // cleanup
        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
            .execute(&idx.pool)
            .await
            .ok();
    }

    /// halfvec(float16) 변형 라운드트립 — `new_halfvec` 로 halfvec 컬럼/인덱스 생성 후
    /// add/search/remove 가 float32 경로와 동일하게 동작하는지. pgvector ≥ 0.6 필요.
    #[tokio::test]
    async fn roundtrip_halfvec() {
        let Some(url) = pg_url() else {
            eprintln!("skip pgvector halfvec test: LLMKERNEL_PG_URL unset");
            return;
        };
        let table = format!("lk_test_hv_{}", line!());
        let idx = PgVectorIndex::new_halfvec(&url, &table, 3)
            .await
            .expect("new_halfvec");

        let vecs = vec![
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
        ];
        let ids = vec![10u64, 20, 30];
        idx.add(&vecs, &ids).await.expect("add");
        assert_eq!(idx.len().await.unwrap(), 3);

        // nearest to [1,0,0] → id 10 (halfvec cosine 도 동일 순위)
        let hits = idx.search(&[1.0, 0.0, 0.0], 1).await.unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].id, 10);

        idx.remove(&[10]).await.unwrap();
        assert_eq!(idx.len().await.unwrap(), 2);

        // cleanup
        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
            .execute(&idx.pool)
            .await
            .ok();
    }

    /// `remove_in_tx`: 같은 풀에서 시작한 트랜잭션 내 삭제가 원자 반영되는지 검증.
    /// klr prune(chunks + vectors 단일 tx)의 전제 — 커밋 전 롤백 시 벡터 보존 확인.
    #[tokio::test]
    async fn remove_in_tx_atomic_delete() {
        let Some(url) = pg_url() else {
            eprintln!("skip pgvector test: LLMKERNEL_PG_URL unset");
            return;
        };
        let table = format!("lk_txtx_{}", line!());
        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");
        idx.add(&[vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]], &[11, 12])
            .await
            .expect("add");
        assert_eq!(idx.len().await.unwrap(), 2);

        // 커밋 경로: tx 내 remove → commit → 삭제 반영
        let mut tx = idx.pool().begin().await.expect("begin tx");
        idx.remove_in_tx(&mut tx, &[11])
            .await
            .expect("remove_in_tx");
        tx.commit().await.expect("commit");
        assert_eq!(idx.len().await.unwrap(), 1);

        // 롤백 경로: tx 내 remove → rollback → 벡터 보존(원자성)
        let mut tx2 = idx.pool().begin().await.expect("begin tx2");
        idx.remove_in_tx(&mut tx2, &[12])
            .await
            .expect("remove_in_tx2");
        tx2.rollback().await.expect("rollback");
        assert_eq!(idx.len().await.unwrap(), 1);

        // cleanup
        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
            .execute(idx.pool())
            .await
            .ok();
    }

    #[test]
    fn rejects_invalid_table_name() {
        // valid identifiers accepted
        assert!(validate_table_name("lk_test_1").is_ok());
        assert!(validate_table_name("_vec").is_ok());
        // rejected — would break out of identifier context
        assert!(validate_table_name("").is_err());
        assert!(validate_table_name("1bad").is_err());
        assert!(validate_table_name("rm; DROP").is_err());
        assert!(validate_table_name("weird\"name").is_err());
        assert!(validate_table_name("sch.tbl").is_err());
    }

    /// pgvector's `sparsevec` text form is 1-based; model indices are 0-based.
    #[test]
    fn sparse_literal_is_one_based() {
        let sv = SparseVector::new(vec![0, 2], vec![0.5, 0.25]).unwrap();
        assert_eq!(sparse_literal(&sv, 5), "{1:0.5,3:0.25}/5");
    }

    #[test]
    fn sparse_literal_handles_empty() {
        assert_eq!(sparse_literal(&SparseVector::default(), 5), "{}/5");
    }

    /// Sparse index roundtrip: inner-product ranking, empty-query
    /// short-circuit, and delete. Requires pgvector ≥ 0.7 for `sparsevec`.
    #[tokio::test]
    async fn roundtrip_sparse() {
        let Some(url) = pg_url() else {
            eprintln!("skip pgvector sparse test: LLMKERNEL_PG_URL unset");
            return;
        };
        let table = format!("lk_test_sp_{}", line!());
        let idx = PgSparseVectorIndex::new(&url, &table, 10)
            .await
            .expect("new sparse");

        let a = SparseVector::new(vec![0, 3], vec![1.0, 0.5]).unwrap();
        let b = SparseVector::new(vec![3, 7], vec![0.5, 1.0]).unwrap();
        idx.add(&[a.clone(), b], &[10, 20]).await.expect("add");
        assert_eq!(idx.len().await.unwrap(), 2);

        // Querying with `a` itself must rank id 10 first (largest dot product).
        let hits = idx.search(&a, 1).await.unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].id, 10);

        // An all-zero query has zero inner product with everything.
        let empty = idx.search(&SparseVector::default(), 5).await.unwrap();
        assert!(empty.is_empty());

        idx.remove(&[10]).await.unwrap();
        assert_eq!(idx.len().await.unwrap(), 1);

        sqlx::query(&format!("DROP TABLE IF EXISTS {}", table))
            .execute(idx.pool())
            .await
            .ok();
    }

    /// `search_where`: the predicate joins a side table inside PG, and text binds
    /// start at `$3`. Dense and sparse indexes share the same filter.
    #[tokio::test]
    async fn search_where_joins_side_table() {
        let Some(url) = pg_url() else {
            eprintln!("skip pgvector search_where test: LLMKERNEL_PG_URL unset");
            return;
        };
        let table = format!("lk_test_sw_{}", line!());
        let sp_table = format!("{table}_sp");
        let tags = format!("{table}_tags");
        let idx = PgVectorIndex::new(&url, &table, 3).await.expect("new");
        let sp = PgSparseVectorIndex::new(&url, &sp_table, 10)
            .await
            .expect("new sparse");
        idx.add(
            &[
                vec![1.0, 0.0, 0.0],
                vec![0.9, 0.1, 0.0],
                vec![0.0, 1.0, 0.0],
            ],
            &[10, 20, 30],
        )
        .await
        .expect("add");
        let a = SparseVector::new(vec![0, 3], vec![1.0, 0.5]).unwrap();
        let b = SparseVector::new(vec![0, 7], vec![0.5, 1.0]).unwrap();
        sp.add(&[a.clone(), b], &[10, 20])
            .await
            .expect("add sparse");
        sqlx::query(&format!(
            "CREATE TABLE {tags} (id BIGINT PRIMARY KEY, tag TEXT NOT NULL)"
        ))
        .execute(idx.pool())
        .await
        .expect("create tags");
        sqlx::query(&format!(
            "INSERT INTO {tags} VALUES (10, 'x'), (20, 'y'), (30, 'y')"
        ))
        .execute(idx.pool())
        .await
        .expect("insert tags");

        let pred = format!("EXISTS (SELECT 1 FROM {tags} t WHERE t.id = v.id AND t.tag = $3)");
        let filter = SqlFilter {
            predicate: &pred,
            binds: &["y"],
        };
        // Nearest to [1,0,0] is 10, but it is tagged x: 20 must come first.
        let hits = idx
            .search_where(&[1.0, 0.0, 0.0], 2, &filter)
            .await
            .unwrap();
        assert_eq!(hits.iter().map(|h| h.id).collect::<Vec<_>>(), vec![20, 30]);
        let hits = sp.search_where(&a, 5, &filter).await.unwrap();
        assert_eq!(hits.iter().map(|h| h.id).collect::<Vec<_>>(), vec![20]);
        assert!(
            sp.search_where(&SparseVector::default(), 5, &filter)
                .await
                .unwrap()
                .is_empty()
        );

        for t in [&tags, &table, &sp_table] {
            sqlx::query(&format!("DROP TABLE IF EXISTS {t}"))
                .execute(idx.pool())
                .await
                .ok();
        }
    }

    #[test]
    fn hnsw_with_clause_variants() {
        assert_eq!(hnsw_with_clause(None, None), "");
        assert_eq!(hnsw_with_clause(Some(32), None), " WITH (m = 32)");
        assert_eq!(
            hnsw_with_clause(None, Some(200)),
            " WITH (ef_construction = 200)"
        );
        assert_eq!(
            hnsw_with_clause(Some(32), Some(200)),
            " WITH (m = 32, ef_construction = 200)"
        );
    }

    /// `PgVectorOpts::default()` must stay behaviourally identical to `new`
    /// (full precision, pgvector's own HNSW defaults).
    #[test]
    fn default_opts_are_full_precision_and_untuned() {
        let o = PgVectorOpts::default();
        assert!(!o.half);
        assert!(o.hnsw_m.is_none());
        assert!(o.hnsw_ef_construction.is_none());
        assert!(o.hnsw_ef_search.is_none());
    }

    #[test]
    fn rejects_overflowing_id() {
        assert_eq!(to_pg_id(0).unwrap(), 0);
        assert_eq!(to_pg_id(42).unwrap(), 42);
        assert_eq!(to_pg_id(i64::MAX as u64).unwrap(), i64::MAX);
        assert!(to_pg_id((i64::MAX as u64) + 1).is_err());
        assert!(to_pg_id(u64::MAX).is_err());
    }
}