From a6485bf682307543db3b64e1fafdb41383675b41 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 5 Jun 2026 22:53:44 +0000 Subject: [PATCH 1/2] fix(postgres): index-backed RRF hybrid search to avoid full seqscan The hybrid search query wrapped the vector distance operator in a scalar similarity expression and sorted on the alias (ORDER BY similarity DESC). That blinds the planner: pgvector's HNSW/DiskANN index can only serve a bare "ORDER BY embedding <=> $vec" path, so the wrapped form degraded to a full sequential scan over every row. On multi-million-row collections this exceeds the statement timeout and returns an empty result set (mudler/LocalAI#10186). Rewrite the query using the Reciprocal Rank Fusion pattern recommended by pgvector and Timescale (pg_textsearch + pgvectorscale): each arm retrieves its top-N candidates via a bare operator (index-served) and assigns a rank; the arms are combined with a FULL OUTER JOIN and scored by weighted RRF (weight/(60+rank)), then joined back to the table by primary key to fetch payloads. RRF fuses by rank rather than raw score, so it no longer mixes BM25's unbounded scores with cosine similarity's [0,1] range. The bm25 and vector weights are preserved as per-arm RRF multipliers. Add a planning regression test that asserts, via EXPLAIN, that the vector ordering is served by the index ("Order By: (embedding <=>") and that the documents table is never sequentially scanned. Signed-off-by: Ettore Di Giacinto --- rag/engine/postgres.go | 87 ++++++++++++++---- rag/engine/postgres_planning_test.go | 127 +++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 17 deletions(-) create mode 100644 rag/engine/postgres_planning_test.go diff --git a/rag/engine/postgres.go b/rag/engine/postgres.go index 61d799d..43d6901 100644 --- a/rag/engine/postgres.go +++ b/rag/engine/postgres.go @@ -698,6 +698,74 @@ func (p *PostgresDB) GetBySource(source string) ([]types.Result, error) { return results, nil } +// hybridCandidateMultiplier controls how many candidates each retrieval arm +// (vector and BM25) pulls before fusion: max(limit*multiplier, floor). A larger +// pool trades latency for recall; 10x with a floor of 100 keeps recall high for +// typical small result limits while staying index-bound. +const ( + hybridCandidateMultiplier = 10 + hybridCandidateFloor = 100 + // rrfK is the Reciprocal Rank Fusion smoothing constant. 60 is the value used + // across the IR literature and by pgvector/Timescale's reference hybrid-search + // examples; it damps the influence of the very top ranks so the two arms blend + // smoothly. + rrfK = 60 +) + +// buildHybridSearchQuery returns the SQL for hybrid (BM25 + vector) search over +// the documents table. Bind parameters: $1 query text, $2 BM25 weight, +// $3 query embedding (::vector), $4 vector weight, $5 result limit. +// +// It follows the canonical Reciprocal Rank Fusion pattern recommended by both +// pgvector and Timescale (pg_textsearch + pgvectorscale). Each arm retrieves its +// top-N candidates via a *bare* operator - "ORDER BY embedding <=> $vec" and +// "ORDER BY full_text <@> to_bm25query(...)" - which pgvector's HNSW/DiskANN and +// the BM25 index serve directly, and assigns a rank. The arms are combined with a +// FULL OUTER JOIN and scored by weighted RRF (sum of weight/(rrfK+rank)); the +// final id list is joined back to the table by primary key to fetch payloads. +// +// RRF fuses by *rank*, not raw score, which avoids mixing BM25's unbounded scores +// with cosine similarity's [0,1] range. The previous query sorted on a wrapped +// scalar similarity expression in a single stage, which blinded the planner into +// a full sequential scan over every row and exceeded the statement timeout on +// multi-million-row collections (LocalAI issue #10186). $2/$4 weight each arm +// (equal by default), so an arm can be biased without breaking the index path. +func buildHybridSearchQuery(tableName string) string { + candidatePool := fmt.Sprintf("GREATEST($5 * %d, %d)", hybridCandidateMultiplier, hybridCandidateFloor) + return fmt.Sprintf(` + WITH bm25_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY full_text <@> to_bm25query($1, 'idx_%[1]s_bm25')) AS rank + FROM %[1]s + ORDER BY full_text <@> to_bm25query($1, 'idx_%[1]s_bm25') + LIMIT %[2]s + ), + vector_results AS ( + SELECT id, ROW_NUMBER() OVER (ORDER BY embedding <=> $3::vector) AS rank + FROM %[1]s + WHERE embedding IS NOT NULL + ORDER BY embedding <=> $3::vector + LIMIT %[2]s + ), + fused AS ( + SELECT + COALESCE(b.id, v.id) AS id, + COALESCE($2 / (%[3]d + b.rank), 0) + COALESCE($4 / (%[3]d + v.rank), 0) AS similarity + FROM bm25_results b + FULL OUTER JOIN vector_results v ON b.id = v.id + ) + SELECT + d.id::text, + COALESCE(d.title, '') as title, + d.content, + d.metadata, + f.similarity + FROM fused f + JOIN %[1]s d ON d.id = f.id + ORDER BY f.similarity DESC + LIMIT $5 + `, tableName, candidatePool, rrfK) +} + func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error) { ctx := context.Background() @@ -708,23 +776,8 @@ func (p *PostgresDB) Search(s string, similarEntries int) ([]types.Result, error } queryEmbeddingStr := formatVector(queryEmbedding) - // Build hybrid search query - // Combine BM25 score and vector similarity - query := fmt.Sprintf(` - SELECT - id::text, - COALESCE(title, '') as title, - content, - metadata, - ( - COALESCE(-(full_text <@> to_bm25query($1, 'idx_%s_bm25')), 0) * $2 + - COALESCE((1 - (embedding <=> $3::vector)), 0) * $4 - ) as similarity - FROM %s - WHERE embedding IS NOT NULL - ORDER BY similarity DESC - LIMIT $5 - `, p.tableName, p.tableName) + // Build hybrid search query (BM25 + vector similarity) + query := buildHybridSearchQuery(p.tableName) rows, err := p.pool.Query(ctx, query, s, p.bm25Weight, queryEmbeddingStr, p.vectorWeight, similarEntries) if err != nil { diff --git a/rag/engine/postgres_planning_test.go b/rag/engine/postgres_planning_test.go new file mode 100644 index 0000000..a01c2cf --- /dev/null +++ b/rag/engine/postgres_planning_test.go @@ -0,0 +1,127 @@ +package engine + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/jackc/pgx/v5/pgxpool" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Regression test for LocalAI issue #10186: the hybrid search query wrapped the +// vector distance operator in a scalar similarity expression and sorted on the +// alias (ORDER BY similarity DESC). That blinds the planner: pgvector's +// HNSW/DiskANN index can only serve a bare "ORDER BY embedding <=> $vec" path, +// so the wrapped form degrades to a full sequential scan over every row and, on +// multi-million-row tables, blows past the statement timeout. +// +// This is a query-planning test, so it only needs a PostgreSQL with pg_textsearch +// + pgvector/vectorscale (the docker-compose stack). It builds the schema via the +// real setupDatabase() and asserts, through EXPLAIN, that the vector ordering is +// served by the index rather than a full scan. No embedding model is required. +var _ = Describe("hybrid search query planning (LocalAI issue #10186)", func() { + var ( + ctx context.Context + pool *pgxpool.Pool + tableName string + queryVec string + ) + + BeforeEach(func() { + ctx = context.Background() + + databaseURL := os.Getenv("POSTGRES_TEST_URL") + if databaseURL == "" { + databaseURL = "postgresql://localrecall:localrecall@localhost:5432/localrecall?sslmode=disable" + } + + var err error + pool, err = pgxpool.New(ctx, databaseURL) + Expect(err).ToNot(HaveOccurred()) + Expect(pool.Ping(ctx)).To(Succeed(), + "PostgreSQL with pg_textsearch + pgvector must be reachable for the planning test") + + const dims = 8 + collectionName := "plan10186" + p := &PostgresDB{ + pool: pool, + collectionName: collectionName, + tableName: sanitizeTableName(collectionName), + embeddingDims: dims, + bm25Weight: 0.5, + vectorWeight: 0.5, + } + tableName = p.tableName + + // Start from a clean slate so setupDatabase()'s CREATE TABLE IF NOT EXISTS + // builds the table at the dimensions this test expects. + _, err = pool.Exec(ctx, "DROP TABLE IF EXISTS "+tableName) + Expect(err).ToNot(HaveOccurred()) + Expect(p.setupDatabase()).To(Succeed()) + + // Seed enough rows that an index path is the cheap plan. Random vectors are + // inserted directly: a planning test needs row shape, not real embeddings. + _, err = pool.Exec(ctx, fmt.Sprintf(` + INSERT INTO %s (title, content, embedding) + SELECT 'title '||g, 'content number '||g, + ('['||random()||','||random()||','||random()||','||random()||','|| + random()||','||random()||','||random()||','||random()||']')::vector + FROM generate_series(1, 2000) g + `, tableName)) + Expect(err).ToNot(HaveOccurred()) + _, err = pool.Exec(ctx, "ANALYZE "+tableName) + Expect(err).ToNot(HaveOccurred()) + + queryVec = "[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8]" + }) + + AfterEach(func() { + if pool != nil { + _, _ = pool.Exec(ctx, "DROP TABLE IF EXISTS "+tableName) + pool.Close() + } + }) + + It("serves the vector ordering from the index instead of a full sequential scan", func() { + query := buildHybridSearchQuery(tableName) + + // Disable plain sequential scans so the planner is forced onto an index + // path *if the query is index-compatible*. The buggy wrapped-scalar ORDER + // BY cannot be served by the vector index and still falls back to a full + // (disabled, high-cost) scan, which this test catches. + tx, err := pool.Begin(ctx) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = tx.Rollback(ctx) }() + + _, err = tx.Exec(ctx, "SET LOCAL enable_seqscan = off") + Expect(err).ToNot(HaveOccurred()) + + rows, err := tx.Query(ctx, "EXPLAIN "+query, "content", 0.5, queryVec, 0.5, 5) + Expect(err).ToNot(HaveOccurred()) + + var plan strings.Builder + for rows.Next() { + var line string + Expect(rows.Scan(&line)).To(Succeed()) + plan.WriteString(line) + plan.WriteByte('\n') + } + Expect(rows.Err()).ToNot(HaveOccurred()) + planText := plan.String() + AddReportEntry("EXPLAIN plan", planText) + + // pgvector/vectorscale emit an index "Order By:" condition on the distance + // operator only when the index actually serves the nearest-neighbour + // ordering. The buggy query produces a "Sort Key:" on the wrapped scalar + // instead, and never this line. + Expect(planText).To(ContainSubstring("Order By: (embedding <=>"), + "the vector nearest-neighbour ordering must be served by the index") + + // And the documents table must never be sequentially scanned for a search. + Expect(planText).ToNot(ContainSubstring("Seq Scan on "+tableName), + "the documents table must not be sequentially scanned") + }) +}) From 84120c8fa86667ccad23ec5499e2e1a24a28326f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 6 Jun 2026 00:02:52 +0000 Subject: [PATCH 2/2] ci: wait for the embedding model to load before running tests docker-compose starts LocalAI with the embedding model on its command line, but the model is downloaded and loaded lazily on first request. wait-localai only checked the HTTP health endpoints, which return OK as soon as the server is up - so the test suites raced the model download and intermittently failed in BeforeEach/Store with `model "granite-embedding-107m-multilingual" not found` (404), across both the postgres and chromem engines. Warm the model as part of wait-localai: after the server is healthy, POST to /v1/embeddings and retry until it answers (up to 600s, since the first call triggers the download). This makes the model a hard readiness gate for every suite that goes through start-test-services, removing the race at its source rather than retrying each individual embedding call. Signed-off-by: Ettore Di Giacinto --- Makefile | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/Makefile b/Makefile index 40eaf25..98a7d2f 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,13 @@ export CGO_ENABLED?=0 IMAGE?=quay.io/mudler/localrecall:latest +# Embedding model the test suites exercise. docker-compose starts LocalAI with +# this model, but it is downloaded/loaded lazily on first request - so the test +# harness must wait for it to actually answer before running specs (see +# wait-localai), otherwise the first embedding call races the download and the +# whole suite fails with "model not found". +EMBEDDING_MODEL?=granite-embedding-107m-multilingual + print-version: @echo "Version: ${VERSION}" @@ -55,6 +62,23 @@ wait-localai: docker compose logs localai | tail -20; \ exit 1; \ fi + @echo "Warming embedding model '$(EMBEDDING_MODEL)' (downloaded lazily on first use)..." + @timeout=600; \ + while [ $$timeout -gt 0 ]; do \ + if curl -fsS http://localhost:8081/v1/embeddings \ + -H 'Content-Type: application/json' \ + -d '{"model":"$(EMBEDDING_MODEL)","input":"warmup"}' >/dev/null 2>&1; then \ + echo "Embedding model '$(EMBEDDING_MODEL)' is ready"; \ + break; \ + fi; \ + sleep 3; \ + timeout=$$((timeout - 3)); \ + done; \ + if [ $$timeout -le 0 ]; then \ + echo "Error: embedding model '$(EMBEDDING_MODEL)' did not load in time"; \ + docker compose logs localai | tail -30; \ + exit 1; \ + fi # Start all test services (LocalAI and PostgreSQL from docker-compose) start-test-services: