diff --git a/docs/index.md b/docs/index.md index 6b06c40..0c59f0d 100644 --- a/docs/index.md +++ b/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: diff --git a/src/bm25.c b/src/bm25.c index 9824386..15a0736 100644 --- a/src/bm25.c +++ b/src/bm25.c @@ -20,6 +20,7 @@ #include #include "access/xact.h" +#include "catalog/namespace.h" #include "catalog/pg_type.h" #include "lib/stringinfo.h" #include "utils/array.h" @@ -27,6 +28,7 @@ #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 */ diff --git a/src/guc.c b/src/guc.c index 216ad25..295933d 100644 --- a/src/guc.c +++ b/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"); } diff --git a/src/pgedge_vectorizer.h b/src/pgedge_vectorizer.h index 30a32ba..586bf39 100644 --- a/src/pgedge_vectorizer.h +++ b/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 diff --git a/test/expected/hybrid_test.out b/test/expected/hybrid_test.out index 9e3da4f..5fe7af7 100644 --- a/test/expected/hybrid_test.out +++ b/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( diff --git a/test/sql/hybrid_test.sql b/test/sql/hybrid_test.sql index a5bb1be..4760028 100644 --- a/test/sql/hybrid_test.sql +++ b/test/sql/hybrid_test.sql @@ -472,6 +472,92 @@ SELECT pgedge_vectorizer.disable_vectorization( 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); + +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; + +-- 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; + +RESET pgedge_vectorizer.corpus_stats_cache_ttl; + +SELECT pgedge_vectorizer.disable_vectorization('cache_docs', 'body', true); +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; + +RESET pgedge_vectorizer.corpus_stats_cache_ttl; +DROP SCHEMA cache_tenant_a CASCADE; +DROP SCHEMA cache_tenant_b CASCADE; + --------------------------------------------------------------------------- -- Cleanup ---------------------------------------------------------------------------