Research-Stack/4-Infrastructure/shim/hybrid_search.sql
allaun 88db8987b0 feat(infra): add hybrid arxiv search with pgvector HNSW + pg_trgm
- embed_arxiv.py: batch embedding of 700k arxiv abstracts using
  static-retrieval-mrl-en-v1 (1024-dim model)
- hybrid_search.sql: RRF merger of trigram + vector search, 22ms query time
- jaccard_hybrid.py: 48 cornfield concepts matched to arxiv papers

Results: 480 hybrid citations loaded (vector-only matches for abstract concepts)
Build: 3314 jobs, 0 errors (lake build Compiler)
2026-06-22 01:11:17 -05:00

67 lines
2.1 KiB
PL/PgSQL

-- hybrid_search.sql
-- Hybrid keyword + semantic search using pg_trgm + pgvector RRF
--
-- Usage:
-- SELECT * FROM hybrid_search('braid eigensolid convergence',
-- (SELECT embedding FROM arxiv_papers WHERE paper_id = 'some_id'), 10);
CREATE OR REPLACE FUNCTION hybrid_search(
query_text TEXT,
query_embedding vector(1024),
top_k INT DEFAULT 10
) RETURNS TABLE (
paper_id TEXT,
title TEXT,
trigram_rank BIGINT,
vector_rank BIGINT,
rrf_score DOUBLE PRECISION
) AS $$
WITH trigram_candidates AS (
SELECT p.paper_id, p.title,
ROW_NUMBER() OVER (ORDER BY similarity(p.title, query_text) DESC) AS trigram_rank
FROM arxiv_papers p
WHERE p.title % query_text
LIMIT 50
),
vector_candidates AS (
SELECT p.paper_id, p.title,
ROW_NUMBER() OVER (ORDER BY p.embedding <=> query_embedding) AS vector_rank
FROM arxiv_papers p
WHERE p.embedding IS NOT NULL
ORDER BY p.embedding <=> query_embedding
LIMIT 50
)
SELECT
COALESCE(t.paper_id, v.paper_id)::TEXT AS paper_id,
COALESCE(t.title, v.title)::TEXT AS title,
t.trigram_rank,
v.vector_rank,
(COALESCE(1.0 / (60 + t.trigram_rank), 0) +
COALESCE(1.0 / (60 + v.vector_rank), 0))::DOUBLE PRECISION AS rrf_score
FROM trigram_candidates t
FULL OUTER JOIN vector_candidates v ON t.paper_id = v.paper_id
ORDER BY rrf_score DESC
LIMIT top_k;
$$ LANGUAGE sql STABLE;
-- Variant: embed query text inline (for when we don't have a pre-computed embedding)
-- This uses a placeholder — actual embedding must be computed in Python
CREATE OR REPLACE FUNCTION hybrid_search_text(
query_text TEXT,
top_k INT DEFAULT 10
) RETURNS TABLE (
paper_id TEXT,
title TEXT,
trigram_rank BIGINT,
rrf_score DOUBLE PRECISION
) AS $$
SELECT
p.paper_id::TEXT,
p.title::TEXT,
ROW_NUMBER() OVER (ORDER BY similarity(p.title, query_text) DESC) AS trigram_rank,
similarity(p.title, query_text)::DOUBLE PRECISION AS rrf_score
FROM arxiv_papers p
WHERE p.title % query_text
ORDER BY similarity(p.title, query_text) DESC
LIMIT top_k;
$$ LANGUAGE sql STABLE;