farine qui passe à travers un tamis au-dessus d'un bol

PostgreSQL – pgvector – Garder l’index ivfflat quand on donne un poids aux résultats

Dans une recherche vectorielle, on veut souvent que certains documents comptent plus que d’autres. Une page officielle doit passer avant un vieux PDF, une page à jour avant une archive périmée. Le plus simple semble de mettre ce poids directement dans le tri.

SELECT id, weight * (1 - (embedding <=> :q)) AS score
FROM documents
ORDER BY weight * (1 - (embedding <=> :q)) DESC
LIMIT 30;

Le résultat est bon, mais PostgreSQL n’utilise plus l’index ivfflat de pgvector. Sur une table de 500 000 chunks, chaque recherche prenait 3,7 secondes au lieu de 150 ms.

Pourquoi un poids casse l’index ivfflat de pgvector

Un index pgvector, ivfflat comme HNSW, ne sait faire qu’une chose : trier par distance, du plus proche au plus loin. PostgreSQL ne s’en sert que si le ORDER BY porte directement sur la distance (embedding <=> :q), en ordre croissant, avec un LIMIT.

Si on fait un calcul autour de la distance, PostgreSQL ne peut plus utiliser l’index. Même un simple 1 - distance suffit à le bloquer. Il calcule alors le score de toutes les lignes de la table, puis il les trie. On le voit tout de suite avec EXPLAIN :

Limit
  ->  Sort
        Sort Key: ((weight * (1 - (embedding <=> '[...]'::vector))) DESC
        ->  Seq Scan on documents

Laisser l’index trier, puis appliquer le poids

La solution : on laisse l’index trier par distance, mais on lui demande plus de lignes qu’il n’en faut. On applique le poids ensuite, sur ce petit lot seulement.

SELECT id, weight * (1 - distance) AS score
FROM (
  SELECT id, weight, embedding <=> :q AS distance
  FROM documents
  ORDER BY embedding <=> :q  -- tri par distance seule : l'index ivfflat marche
  LIMIT 30 * 6                 -- 6 fois plus de lignes que nécessaire
) candidates
ORDER BY score DESC
LIMIT 30;

La sous-requête passe bien par l’index (Index Scan using documents_embedding_idx). Le tri avec le poids ne porte plus que sur 180 lignes. Sur les 500 000 chunks, la recherche retombe à 150-200 ms, environ 25 fois plus vite.

ivfflat.probes et hnsw.ef_search : quand l’index rend moins de lignes que demandé

Attention : un index pgvector, ivfflat ou HNSW, ne rend pas toujours les 180 lignes demandées. Il en rend parfois moins, sans erreur et sans prévenir.

Avec ivfflat, ivfflat.probes dit combien de listes l’index va lire. Par défaut, c’est 1. Dans un test avec 1 000 listes d’environ 50 lignes et probes = 1, la requête n’a rendu que 51 lignes au lieu de 180. La doc de pgvector conseille de partir de la racine carrée du nombre de listes, soit environ 32 pour 1 000 listes.

Avec HNSW, sans parcours itératif, le nombre de lignes rendues est limité par hnsw.ef_search, qui vaut 40 par défaut. Dans le même test, la requête n’a rendu que 40 lignes, quatre fois moins que prévu, et rien ne le signale.

BEGIN;
SET LOCAL ivfflat.probes = 32;   -- ivfflat : environ la racine carrée de lists
SET LOCAL hnsw.ef_search = 200;  -- HNSW : au moins autant que le LIMIT
-- ... la requête avec la sous-requête ...
COMMIT;

SET LOCAL ne vaut que dans la transaction en cours, d’où le BEGIN. Plus la valeur est haute, plus l’index retrouve en général les vrais voisins les plus proches, mais plus la recherche est lente.

Depuis pgvector 0.8, les deux index peuvent aussi continuer à chercher pour essayer d’avoir assez de lignes. Ils s’arrêtent quand même à une limite (hnsw.max_scan_tuples, ivfflat.max_probes) :

SET LOCAL hnsw.iterative_scan = relaxed_order;
-- ou
SET LOCAL ivfflat.iterative_scan = relaxed_order;

Avec relaxed_order, les lignes peuvent arriver un peu dans le désordre par rapport à leur distance. Ici, ça gêne peu : elles servent surtout à remplir le lot, qui est ensuite retrié avec le poids.

Combien de lignes prendre en plus

La méthode n’est pas parfaite, et elle peut perdre des documents à deux endroits. L’index lui-même est approximatif : il peut rater quelques voisins proches. Et un document un peu loin, mais avec un gros poids, peut rester en dehors du lot.

La mesure a été faite sur 500 000 chunks, avec des poids qui vont jusqu’à x3. La référence est le vrai top 30, celui de la requête lente : elle lit toute la table et calcule le score avec le poids pour chacune des 500 000 lignes. Avec 3 fois plus de lignes, on en retrouve 79 %. Avec 6 fois plus, on en retrouve 86 %. Ces chiffres comptent les deux pertes à la fois, et le premier résultat était toujours le même.

Dans ce test, passer de 90 à 180 lignes ne changeait pas le temps de réponse : environ 150 ms dans les deux cas, parce que c’est surtout le parcours de l’index qui coûtait. Avec d’autres réglages ou un lot beaucoup plus gros, ce ne sera pas forcément vrai.

Plus les poids sont forts, plus il faut prendre de lignes en plus. Si presque tous les poids valent 1, le poids change peu le classement.

Un détail sur le calcul : il multiplie la similarité par le poids. Si la similarité est négative, un poids de 3 enfonce le document au lieu de le remonter. Un seuil minimum sur la similarité évite ce cas.

Si le deuxième tri se fait dans le code et pas en SQL, attention au sens. La distance se trie du plus petit au plus grand, le score du plus grand au plus petit. Si on garde le sens de la requête, on garde les pires lignes du lot au lieu des meilleures.

rows.sort((a, b) => b.score - a.score).slice(0, limit);

Laisser un commentaire