Hybrid søgning i Postgres uden pg_search

Supabase tilbyder ikke pg_search, BM25-udvidelsen fra ParadeDB. Den er blevet efterspurgt siden 2023 i en GitHub-diskussion. Indtil da kan man bygge en solid hybrid søgning med de udvidelser, Supabase allerede har. Det har Lovguiden gjort i omkring to år. Her gennemgås den nuværende version, inklusive det, vi gjorde forkert første gang.

Opsætningen

• Ca. 1,45 millioner tekststykker fra afgørelser, 510.000 fra lovgivning og 150.000 sager
• Supabase med vector 0.8, rum 1.3, pg_trgm og unaccent
• 1024-dimensionelle embeddings fra Isaacus kanon-2-embedder
• Én søgning er ét databasekald efterfulgt af ét rerank-kald i applikationen

Hver søgefunktion har tre spor, der kører på den samme forespørgsel. Et ordspor finder dokumenter, der indeholder ordene. Et vektorspor finder dokumenter, der betyder det samme. Et trigram-spor fanger navne, sagsnumre og stavefejl. De tre rangerede lister flettes med Reciprocal Rank Fusion (RRF), og de 30 bedste sendes til en reranker.

Indeksér det, folk søger efter, ikke hele dokumentet

Denne ændring gjorde mere for relevansen end nogen rangeringsfunktion. En afgørelse kan fylde 30 sider, og det meste er procedure. Søger man direkte i den tekst, rammer man meget støj. Derfor skriver en sprogmodel for hver sag en kort liste over formuleringer, en bruger faktisk ville søge på, ud fra sagens resumé og resultat i stedet for hele teksten. Vi kalder det søgeintentionen.

Søgeintentionen får den højeste vægt i en genereret tsvector-kolonne, mens titel og sagsnummer får lavere vægte. Fordi kolonnen er genereret, kan den aldrig komme ud af trit med teksten, den bygger på. De samme formuleringer lægges også ind i den tekst, der laves embedding af, så vektorsporet får glæde af dem.

ALTER TABLE court_cases ADD COLUMN search_intent text;

ALTER TABLE court_cases ADD COLUMN search_intent_tsv tsvector
  GENERATED ALWAYS AS (
    setweight(to_tsvector('danish', coalesce(search_intent, '')), 'A') ||
    setweight(to_tsvector('danish', coalesce(ai_title, '')),      'B') ||
    setweight(to_tsvector('danish', coalesce(original_id, '')),   'C')
  ) STORED;

CREATE INDEX ON court_cases USING gin (search_intent_tsv);

Det afbøder også en kendt svaghed ved Postgres' fuldtekstsøgning. Både plainto_tsquery og websearch_to_tsquery kræver, at alle ord i forespørgslen matcher, så et spørgsmål på fem ord giver ofte nul træf. Når omskrivninger og synonymer allerede står i den indekserede tekst, matcher langt flere spørgsmål. Vi beholder websearch_to_tsquery, fordi citerede fraser og udelukkede ord er nyttige i juridisk søgning. Når ordsporet stadig intet finder, bærer de to andre spor resultatet.

Ordspor: RUM i stedet for GIN

Det sædvanlige råd er et GIN-indeks og ts_rank_cd. GIN finder træffene hurtigt, men ts_rank_cd skal læse hver træffende rækkes tsvector fra tabellen for at give den en score. På en bred forespørgsel med 267.000 træf tog det 88 millisekunder at finde dem. At score dem ramte vores timeout på 44 sekunder. Det hjælper ikke at flytte vinduesfunktionen uden for LIMIT, fordi scoren stadig skal beregnes for alle træf, før man ved, hvilke der er de bedste.

Et RUM-indeks gemmer ordenes positioner i selve indekset, så rangeringen sker, mens indekset læses. Operatoren <=> returnerer en afstand, hvor lavere er bedre.

CREATE EXTENSION IF NOT EXISTS rum;

CREATE INDEX idx_vdv_rum_chunk_text_tsv
  ON verdict_doc_vec USING rum (chunk_text_tsv rum_tsvector_ops);

-- <=> is a distance (lower is better) and is served from the index
SELECT id, row_number() OVER (ORDER BY chunk_text_tsv <=> ts_query) AS rank
FROM verdict_doc_vec
WHERE chunk_text_tsv @@ ts_query
ORDER BY chunk_text_tsv <=> ts_query
LIMIT candidate_pool;

Rangeringen er stadig ikke ægte BM25, men den er nu billig nok til en million rækker. RUM-indekser er større end GIN og langsommere at bygge. Vores tog 10 til 20 minutter på 1,1 millioner rækker og blokerede skrivninger imens. På lovgivningstabellen bruger vi et partielt RUM-indeks, der kun dækker gældende versioner, så det holdes mindre.

Vektorspor: HNSW på halfvec

Embeddings gemmes med fuld præcision, men HNSW-indekserne bygges på en cast til halv præcision. Det sparede 15 til 20 GB indeks, uden målbart tab af træfsikkerhed ved hnsw.ef_search = 100.

-- the column stays vector(1024); only the index is half precision
CREATE INDEX idx_vdv_hnsw_cosine_half
  ON verdict_doc_vec USING hnsw ((embedding::halfvec(1024)) halfvec_cosine_ops)
  WITH (m = 16, ef_construction = 200);

-- the query must use the exact same expression
SELECT id FROM verdict_doc_vec
ORDER BY embedding::halfvec(1024) <=> $1::halfvec(1024)
LIMIT candidate_pool;

Forespørgslen skal bruge præcis den samme cast som indekset. Ellers ignorerer Postgres indekset uden at sige det og læser hele tabellen. Der kommer ingen fejl, bare en langsom forespørgsel.

Filtreret vektorsøgning kræver en konstant

Brugere søger tit i én domstol eller ét nævn ad gangen. De største kilder har deres egne partielle HNSW-indekser. I en PL/pgSQL-funktion kan en betingelse som type = p_source_filter blive planlagt uden kendskab til værdien, og så kan Postgres ikke bruge et partielt indeks med betingelsen type = 'landsskatteretten'. Den falder tilbage til det globale indeks og filtrerer bagefter. Fordi HNSW returnerer et fast antal nærmeste naboer, før filteret anvendes, kan der blive meget få resultater tilbage.

CREATE INDEX idx_vdv_hnsw_landsskat_half
  ON verdict_doc_vec USING hnsw ((embedding::halfvec(1024)) halfvec_cosine_ops)
  WITH (m = 16, ef_construction = 200)
  WHERE type = 'landsskatteretten';

-- a) inline the value so the planner can match the index predicate
RETURN QUERY EXECUTE format($q$
  SELECT ... FROM verdict_doc_vec vdv
  WHERE vdv.type = %L
  ORDER BY vdv.embedding::halfvec(1024) <=> $1
  LIMIT $2
$q$, p_source_filter) USING v_halfvec, candidate_pool;

-- b) or plan every call with its real values
CREATE FUNCTION ... SET plan_cache_mode TO 'force_custom_plan' AS $$ ... $$;

Når der er filtre, hæver vi også hnsw.ef_search til 200, så grafsøgningen kigger længere, før der filtreres. pgvector 0.8 har iterative indeksscanninger, som løser det samme problem på en anden måde.

Trigram-spor: navne, stavefejl og sammensatte ord

Sammensatte ord som parkeringsbøde og hundeloven, personnavne og sagsnumre er netop dér, hvor både ordstammer og embeddings fejler. Et trigram-indeks på en kopi af titel og søgeintention uden accenter og med små bogstaver klarer dem. Funktionen unaccent er ikke markeret som immutable, så den skal pakkes ind, før den kan bruges i en genereret kolonne.

CREATE FUNCTION immutable_unaccent(text) RETURNS text
  LANGUAGE sql IMMUTABLE PARALLEL SAFE
  AS $$ SELECT extensions.unaccent('extensions.unaccent'::regdictionary, $1) $$;

ALTER TABLE court_cases ADD COLUMN fuzzy_text text
  GENERATED ALWAYS AS (
    immutable_unaccent(lower(coalesce(ai_title, '') || ' ' || coalesce(search_intent, '')))
  ) STORED;

CREATE INDEX ON court_cases USING gist (fuzzy_text gist_trgm_ops (siglen = 256));

-- nearest neighbours by word-similarity distance, then a similarity floor
SELECT id, word_similarity(q, fuzzy_text) AS sim
FROM court_cases
ORDER BY q <<-> fuzzy_text
LIMIT candidate_pool;   -- keep sim >= 0.5, rank by sim

Vi kører kun dette spor for forespørgsler på højst tre ord. For længere spørgsmål bliver nabosøgningen langsom og tilføjer næsten intet.

Fletning: vægtet RRF

RRF giver hvert dokument 1 / (placering + k) for hver liste, det optræder i, og lægger dem sammen. Den bruger kun placeringer, så det er ligegyldigt, at en vektorafstand og en tekstscore ligger på helt forskellige skalaer.

CREATE FUNCTION rrf_score(rank bigint, rrf_k integer DEFAULT 50)
RETURNS numeric LANGUAGE sql IMMUTABLE PARALLEL SAFE
AS $$ SELECT COALESCE(1.0 / ($1 + $2), 0.0) $$;

-- inside the search function
combined AS (
  SELECT *, 0.5 AS w FROM vec_arm
  UNION ALL SELECT *, 0.5 AS w FROM lex_arm
  UNION ALL SELECT *, 0.4 AS w FROM fuzzy_arm
)
SELECT id, sum(w * rrf_score(rank, 50)) AS score
FROM combined
GROUP BY id
ORDER BY score DESC
LIMIT p_limit;

Nogle detaljer betyder mere, end man skulle tro:

• Hvert spor returnerer mindst 60 kandidater eller det dobbelte af den endelige grænse, ikke kun den endelige grænse. Ellers får et dokument, der ligger nr. 31 i alle tre spor, aldrig chancen for at vinde fletningen.
• Sporene har vægte. På afgørelser tæller vektor og ord lige meget og trigram lidt mindre. På lovgivning tæller ord lidt mere, fordi folk søger på præcise begreber.
• Faglige justeringer kommer efter fletningen. Love rangeres over kommentarer, nyere afgørelser får et lille løft på højst 18 procent, og en forespørgsel med § 12 sætter rækker for § 12 øverst i ordsporet.
• Lige scorer skal have en afgører. ts_rank_cd giver mange ens scorer, så sortering på score alene kan give forskellige rækker for den samme forespørgsel. Tilføj rækkens id som sidste sorteringsnøgle.

Rerank i applikationen

Databasen returnerer ca. 30 kandidater, og en cross-encoder-reranker (Isaacus kanon-2-reranker) sætter dem i den endelige rækkefølge. Arbejdsdelingen fungerer godt: RRF er god til at få de rigtige dokumenter med blandt de 30, og rerankeren er god til at ordne dem. At reranke 30 dokumenter går hurtigt. At reranke 3.000 ville ikke.

Faldgruber

• SECURITY DEFINER omgår row-level security. Vores søgefunktioner bruger det, fordi data er offentlige. For brugerdata skal man beholde standarden SECURITY INVOKER, så gælder RLS for søgningen præcis som alle andre steder.
• SET LOCAL work_mem og SET LOCAL hnsw.ef_search i funktionen gælder kun for det ene kald, hvilket er præcis det, man vil have bag PostgREST.
• Genererede kolonner, der kalder en SQL-indpakning som immutable_unaccent, fik vores opgradering af Postgres-hovedversion på Supabase til at fejle. Vi måtte fjerne dem før opgraderingen og tilføje dem igen bagefter.
• At bygge HNSW- og RUM-indekser uden CONCURRENTLY blokerer skrivninger. Sæt indlæsningen på pause, eller byg samtidigt uden for en transaktion.
• Tjek, hvor tit hvert spor faktisk returnerer rækker. Et ordspor, der stille og roligt er tomt, ser fint ud i resultaterne, fordi vektorsporet stadig returnerer noget.

Når pg_search kommer

Ægte BM25 ville stadig være bedre end RUM-rangering, og den dag Supabase udgiver pg_search, skifter vi ordsporet ud. Alt andet forbliver det samme: søgeintentionen, vektor- og trigram-sporene, fletningen og rerankeren. Indtil da kører det i den database, vi allerede har, med row-level security intakt og uden en ekstra søgetjeneste, der skal holdes synkroniseret.