Cache the BM25 corpus figures per backend instead of rescanning

docs/index.md

@@ -173,6 +173,7 @@ different use cases.
 | `pgedge_vectorizer.enable_hybrid` | `false` | -- | Enable BM25 sparse vectors alongside dense embeddings |
 | `pgedge_vectorizer.bm25_k1` | `1.2` | `0.0-3.0` | BM25 term-frequency saturation |
 | `pgedge_vectorizer.bm25_b` | `0.75` | `0.0-1.0` | BM25 document-length normalization |
+| `pgedge_vectorizer.corpus_stats_cache_ttl` | `60s` | `0-3600` | Seconds a backend reuses a chunk table's corpus size and mean document length before reading them again. Set to 0 to read them on every call |
 
 For more information or to download Vectorizer visit:
 

src/bm25.c

@@ -20,13 +20,15 @@
 #include <string.h>
 
 #include "access/xact.h"
+#include "catalog/namespace.h"
 #include "catalog/pg_type.h"
 #include "lib/stringinfo.h"
 #include "utils/array.h"
 #include "utils/builtins.h"
 #include "utils/errcodes.h"
 #include "utils/hsearch.h"
 #include "utils/memutils.h"
+#include "utils/timestamp.h"
 
 /* ----------------------------------------------------------------
  * English stopword list — must stay sorted (used with bsearch)
@@ -221,7 +223,30 @@ bm25_tokenize(const char *text, int *ntokens)
 }
 
 /*
- * bm25_corpus_stats — corpus size N and mean document length, in one pass.
+ * Cached corpus figures for one chunk table.
+ *
+ * Keyed on the relation's OID rather than the name it was reached by.  The
+ * name is resolved through search_path, so one string can mean different
+ * relations at different moments: a schema-per-tenant deployment gives every
+ * tenant a chunk table of the same name, and a pooled backend that switches
+ * search_path between requests would otherwise be handed whichever tenant's
+ * figures it saw first.  Dropping and recreating a chunk table under the same
+ * name has the same shape, and an OID changes in both cases.
+ *
+ * Key must be first for dynahash.
+ */
+typedef struct
+{
+	Oid			key;
+	int64		total_docs;
+	float8		avg_doc_len;
+	TimestampTz	read_at;
+} CorpusStatsEntry;
+
+static HTAB *corpus_stats_cache = NULL;
+
+/*
+ * bm25_corpus_stats_uncached — read N and the mean document length.
  *
  * Both come from the same scan: taken separately they made two passes over
  * the chunk table for every search.  Returns N, clamped to a minimum of 1 so
@@ -235,7 +260,7 @@ bm25_tokenize(const char *text, int *ntokens)
  * Caller must have an active SPI connection.
  */
 static int64
-bm25_corpus_stats(const char *chunk_table, float8 *avg_doc_len)
+bm25_corpus_stats_uncached(const char *chunk_table, float8 *avg_doc_len)
 {
 	char   *sql;
 	int     ret;
@@ -269,6 +294,93 @@ bm25_corpus_stats(const char *chunk_table, float8 *avg_doc_len)
 	return (total <= 0) ? 1 : total;
 }
 
+/*
+ * bm25_corpus_stats — corpus size N and mean document length, cached.
+ *
+ * The underlying read is an unindexable aggregate over the whole chunk table,
+ * and it was paid on every call: once per search, and once per queue item in
+ * the worker, so a batch of ten scanned the table ten times over.  At 300,000
+ * chunks that is around 33ms a time and it grows with the corpus.
+ *
+ * These two figures are inputs to a ranking heuristic rather than an account
+ * that has to balance, and during ingest they are a moving target in any
+ * case, since every chunk written changes them.  Holding a value for a few
+ * seconds therefore costs a little precision in a number that was never
+ * precise, and buys the removal of a scan from the hot path of every search.
+ *
+ * Cached per backend rather than shared, which suits both callers: a worker
+ * is a long-lived process working one database, and searches arrive on
+ * pooled connections that are also long-lived.  A short-lived backend that
+ * runs a single search pays the scan exactly as it did before.
+ *
+ * Set pgedge_vectorizer.corpus_stats_cache_ttl to 0 to read afresh every
+ * time, which restores the previous behaviour exactly.
+ *
+ * Caller must have an active SPI connection.
+ */
+static int64
+bm25_corpus_stats(const char *chunk_table, float8 *avg_doc_len)
+{
+	CorpusStatsEntry   *entry;
+	TimestampTz			now;
+	Oid					relid;
+	int64				total_docs;
+
+	if (pgedge_vectorizer_corpus_stats_cache_ttl <= 0)
+		return bm25_corpus_stats_uncached(chunk_table, avg_doc_len);
+
+	/*
+	 * Resolve the name the same way the query below will, through
+	 * search_path.  A name that resolves to nothing is left to the uncached
+	 * read, which raises the "relation does not exist" the caller expects
+	 * rather than inventing a different error here.
+	 */
+	relid = RelnameGetRelid(chunk_table);
+	if (!OidIsValid(relid))
+		return bm25_corpus_stats_uncached(chunk_table, avg_doc_len);
+
+	if (corpus_stats_cache == NULL)
+	{
+		HASHCTL		ctl;
+
+		memset(&ctl, 0, sizeof(ctl));
+		ctl.keysize = sizeof(Oid);
+		ctl.entrysize = sizeof(CorpusStatsEntry);
+		ctl.hcxt = TopMemoryContext;
+		corpus_stats_cache = hash_create("bm25_corpus_stats", 8, &ctl,
+										 HASH_ELEM | HASH_BLOBS |
+										 HASH_CONTEXT);
+	}
+
+	now = GetCurrentTimestamp();
+	entry = hash_search(corpus_stats_cache, &relid, HASH_FIND, NULL);
+
+	if (entry != NULL &&
+		!TimestampDifferenceExceeds(entry->read_at, now,
+									pgedge_vectorizer_corpus_stats_cache_ttl
+									* 1000))
+	{
+		*avg_doc_len = entry->avg_doc_len;
+		return entry->total_docs;
+	}
+
+	/*
+	 * Read before entering anything.  HASH_ENTER inserts an entry for the
+	 * caller to fill, and this read can throw, leaving one unfilled in
+	 * TopMemoryContext for a later call to read as figures -- dynahash.c warns
+	 * of exactly this.  The local is what sequences it: assigning the return
+	 * straight into the entry would put HASH_ENTER back before the call.
+	 */
+	total_docs = bm25_corpus_stats_uncached(chunk_table, avg_doc_len);
+
+	entry = hash_search(corpus_stats_cache, &relid, HASH_ENTER, NULL);
+	entry->total_docs = total_docs;
+	entry->avg_doc_len = *avg_doc_len;
+	entry->read_at = now;
+
+	return total_docs;
+}
+
 /*
  * read_idf_row — extract one IdfStat from a SPI result row
  */

src/guc.c

@@ -51,6 +51,7 @@ int pgedge_vectorizer_auto_cleanup_hours = 24;
 bool   pgedge_vectorizer_enable_hybrid = false;
 double pgedge_vectorizer_bm25_k1       = 1.2;
 double pgedge_vectorizer_bm25_b        = 0.75;
+int    pgedge_vectorizer_corpus_stats_cache_ttl = 60;
 
 /*
  * Initialize all GUC variables
@@ -312,5 +313,22 @@ pgedge_vectorizer_init_guc(void)
 		0,
 		NULL, NULL, NULL);
 
+	DefineCustomIntVariable(
+		"pgedge_vectorizer.corpus_stats_cache_ttl",
+		"Seconds a backend may reuse a chunk table's cached corpus figures",
+		"BM25 needs the corpus size and the mean document length, which "
+		"together cost one unindexable scan of the chunk table. Both feed a "
+		"ranking heuristic rather than an account that has to balance, so "
+		"each backend holds them for this long instead of reading them "
+		"again on every search and every queue item. Set to 0 to read them "
+		"afresh every time.",
+		&pgedge_vectorizer_corpus_stats_cache_ttl,
+		60,     /* default */
+		0,      /* min: 0 disables the cache */
+		3600,   /* max */
+		PGC_USERSET,
+		GUC_UNIT_S,
+		NULL, NULL, NULL);
+
 	elog(DEBUG1, "pgedge_vectorizer GUC variables initialized");
 }

src/pgedge_vectorizer.h

@@ -72,6 +72,7 @@ extern int pgedge_vectorizer_auto_cleanup_hours;
 extern bool   pgedge_vectorizer_enable_hybrid;
 extern double pgedge_vectorizer_bm25_k1;
 extern double pgedge_vectorizer_bm25_b;
+extern int    pgedge_vectorizer_corpus_stats_cache_ttl;
 
 /*
  * Chunking strategy enumeration

test/expected/hybrid_test.out

@@ -620,6 +620,110 @@ NOTICE:  Vectorization disabled and chunk table dropped: hybrid_multi_test_body_
 
 DROP TABLE hybrid_multi_test;
 ---------------------------------------------------------------------------
+-- Test 20: the corpus statistics cache does not change results
+---------------------------------------------------------------------------
+-- N and the mean document length feed a ranking heuristic, and reading them
+-- costs an unindexable scan of the chunk table on every call. Each backend
+-- therefore holds them for corpus_stats_cache_ttl seconds. Whatever the
+-- setting, the vector produced for a given corpus must be the same, so this
+-- compares the two paths directly rather than asserting on either alone.
+CREATE TABLE cache_docs (id BIGSERIAL PRIMARY KEY, body TEXT);
+INSERT INTO cache_docs (body)
+SELECT 'alpha beta gamma document number ' || g FROM generate_series(1, 25) g;
+SELECT pgedge_vectorizer.enable_vectorization(
+    'cache_docs', 'body', embedding_dimension => 3);
+NOTICE:  Using primary key column: id (bigint)
+NOTICE:  column "sparse_embedding" of relation "cache_docs_body_chunks" already exists, skipping
+NOTICE:  Vectorization enabled: cache_docs -> cache_docs_body_chunks
+NOTICE:  Strategy: token_based, chunk_size: 400, overlap: 50
+NOTICE:  Processing existing rows...
+NOTICE:  Processed 25 existing rows
+ enable_vectorization 
+----------------------
+
+(1 row)
+
+SET pgedge_vectorizer.corpus_stats_cache_ttl = 0;
+SELECT pgedge_vectorizer.bm25_query_vector('alpha beta', 'cache_docs_body_chunks')
+    AS uncached \gset
+SET pgedge_vectorizer.corpus_stats_cache_ttl = 60;
+SELECT pgedge_vectorizer.bm25_query_vector('alpha beta', 'cache_docs_body_chunks')
+    AS cached \gset
+SELECT :'uncached'::sparsevec = :'cached'::sparsevec AS cache_matches_uncached;
+ cache_matches_uncached 
+------------------------
+ t
+(1 row)
+
+-- A second table must not read the first one's cached figures.
+SELECT pgedge_vectorizer.bm25_avg_doc_len('cache_docs_body_chunks')
+    <> pgedge_vectorizer.bm25_avg_doc_len('hybrid_test_docs_content_chunks')
+    AS each_table_cached_separately;
+ each_table_cached_separately 
+------------------------------
+ t
+(1 row)
+
+RESET pgedge_vectorizer.corpus_stats_cache_ttl;
+SELECT pgedge_vectorizer.disable_vectorization('cache_docs', 'body', true);
+NOTICE:  Vectorization disabled and chunk table dropped: cache_docs_body_chunks
+ disable_vectorization 
+-----------------------
+
+(1 row)
+
+DROP TABLE cache_docs;
+---------------------------------------------------------------------------
+-- Test 21: the corpus statistics cache is keyed by relation, not by name
+---------------------------------------------------------------------------
+-- A chunk table is named unqualified and resolved through search_path, so one
+-- string means different relations in a schema-per-tenant deployment. Keyed by
+-- name, a pooled backend that switched tenants was handed whichever tenant's
+-- corpus it had seen first, and scored the second tenant's queries against the
+-- first one's statistics. The two corpora below differ only in document
+-- length, so the vectors must differ; keyed by name they came back identical.
+CREATE SCHEMA cache_tenant_a;
+CREATE SCHEMA cache_tenant_b;
+CREATE TABLE cache_tenant_a.t_chunks (id BIGSERIAL PRIMARY KEY, content TEXT,
+                                      token_count INT);
+CREATE TABLE cache_tenant_b.t_chunks (id BIGSERIAL PRIMARY KEY, content TEXT,
+                                      token_count INT);
+INSERT INTO cache_tenant_a.t_chunks (content, token_count)
+SELECT 'alpha', 10 FROM generate_series(1, 50);
+INSERT INTO cache_tenant_b.t_chunks (content, token_count)
+SELECT 'alpha', 900 FROM generate_series(1, 50);
+CREATE TABLE cache_tenant_a.t_chunks_idf_stats (term TEXT PRIMARY KEY,
+                                                doc_freq INT NOT NULL);
+CREATE TABLE cache_tenant_b.t_chunks_idf_stats (term TEXT PRIMARY KEY,
+                                                doc_freq INT NOT NULL);
+INSERT INTO cache_tenant_a.t_chunks_idf_stats VALUES ('alpha', 5);
+INSERT INTO cache_tenant_b.t_chunks_idf_stats VALUES ('alpha', 5);
+SET pgedge_vectorizer.corpus_stats_cache_ttl = 60;
+SET search_path = cache_tenant_a, public;
+SELECT pgedge_vectorizer.bm25_query_vector('alpha', 't_chunks') AS tenant_a \gset
+SET search_path = cache_tenant_b, public;
+SELECT pgedge_vectorizer.bm25_query_vector('alpha', 't_chunks') AS tenant_b \gset
+-- RESET rather than SET, so that whatever the session started with is what
+-- the rest of the file runs under. The comparison below needs no relation
+-- lookup of its own.
+RESET search_path;
+SELECT :'tenant_a'::sparsevec <> :'tenant_b'::sparsevec
+    AS same_name_different_schema_not_confused;
+ same_name_different_schema_not_confused 
+-----------------------------------------
+ t
+(1 row)
+
+RESET pgedge_vectorizer.corpus_stats_cache_ttl;
+DROP SCHEMA cache_tenant_a CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table cache_tenant_a.t_chunks
+drop cascades to table cache_tenant_a.t_chunks_idf_stats
+DROP SCHEMA cache_tenant_b CASCADE;
+NOTICE:  drop cascades to 2 other objects
+DETAIL:  drop cascades to table cache_tenant_b.t_chunks
+drop cascades to table cache_tenant_b.t_chunks_idf_stats
+---------------------------------------------------------------------------
 -- Cleanup
 ---------------------------------------------------------------------------
 SELECT pgedge_vectorizer.disable_vectorization(