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

PostgreSQL – pgvector – Keep the ivfflat index when you give results a weight

by

In a vector search, you often want some documents to count more than others. An official page should come before an old PDF, an up-to-date page before a stale archive. The easy way seems to be putting that weight straight into the sort.

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

The result is right, but PostgreSQL stops using the pgvector ivfflat index. On a table of 500,000 chunks, each search took 3.7 seconds instead of 150 ms.

Why a weight breaks the pgvector ivfflat index

A pgvector index, ivfflat or HNSW, can only do one thing: sort by distance, from nearest to farthest. PostgreSQL only uses it when the ORDER BY is directly on the distance (embedding <=> :q), ascending, with a LIMIT.

If you do any math around the distance, PostgreSQL can no longer use the index. Even a plain 1 - distance is enough to block it. It then computes the score of every row in the table, and sorts them all. EXPLAIN shows it right away:

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

Let the index sort, then apply the weight

The fix: let the index sort by distance, but ask it for more rows than you need. Then apply the weight, on that small batch only.

SELECT id, weight * (1 - distance) AS score
FROM (
  SELECT id, weight, embedding <=> :q AS distance
  FROM documents
  ORDER BY embedding <=> :q  -- sort by distance only: the ivfflat index works
  LIMIT 30 * 6                 -- 6 times more rows than needed
) candidates
ORDER BY score DESC
LIMIT 30;

The subquery does go through the index (Index Scan using documents_embedding_idx). The weighted sort now only covers 180 rows. On the 500,000 chunks, the search is back to 150-200 ms, about 25 times faster.

ivfflat.probes and hnsw.ef_search: when the index returns fewer rows than asked

Watch out: a pgvector index, ivfflat or HNSW, does not always return the 180 rows you asked for. Sometimes it returns fewer, with no error and no warning.

With ivfflat, ivfflat.probes says how many lists the index reads. The default is 1. In a test with 1,000 lists of about 50 rows and probes = 1, the query only returned 51 rows instead of 180. The pgvector docs suggest starting from the square root of the number of lists, so about 32 for 1,000 lists.

With HNSW, without iterative scans, the number of rows returned is capped by hnsw.ef_search, which defaults to 40. In the same test, the query only returned 40 rows, four times fewer than planned, and nothing tells you.

BEGIN;
SET LOCAL ivfflat.probes = 32;   -- ivfflat: about the square root of lists
SET LOCAL hnsw.ef_search = 200;  -- HNSW: at least as much as the LIMIT
-- ... the query with the subquery ...
COMMIT;

SET LOCAL only lasts for the current transaction, hence the BEGIN. The higher the value, the more of the true nearest neighbours the index usually finds, but the slower the search.

Since pgvector 0.8, both indexes can also keep searching to try to get enough rows. They still stop at a limit (hnsw.max_scan_tuples, ivfflat.max_probes):

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

With relaxed_order, the rows can come back slightly out of order by distance. That matters little here: they are mostly used to fill the batch, which is then sorted again with the weight.

How many extra rows to take

This method is not perfect, and it can lose documents in two places. The index itself is approximate: it can miss a few close neighbours. And a document a bit far away, but with a big weight, can stay outside the batch.

These numbers come from 500,000 chunks, with weights going up to x3. The reference is the true top 30, from the slow query: it reads the whole table and computes the weighted score for all 500,000 rows. With 3 times more rows, you get back 79% of it. With 6 times more, you get 86%. These numbers count both losses together, and the first result was always the same.

In this test, going from 90 to 180 rows did not change the response time: about 150 ms either way, because walking the index was what cost the most. With other settings or a much bigger batch, that may not hold.

The stronger the weights, the more extra rows you need. If almost all weights are 1, the weight barely changes the ranking.

One detail about the math: it multiplies the similarity by the weight. If the similarity is negative, a weight of 3 pushes the document further down instead of up. A minimum similarity threshold avoids that case.

If the second sort happens in code rather than in SQL, mind the direction. The distance sorts from smallest to largest, the score from largest to smallest. If you keep the query’s direction, you keep the worst rows of the batch instead of the best.

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

Leave a comment