From d68eaa9c6bb7d90dcbbfba07f4f0903d7975c43b Mon Sep 17 00:00:00 2001 From: Kazim Haider Date: Thu, 28 May 2026 23:27:23 +0500 Subject: [PATCH 1/4] Expose count_tokens() and add refresh_token_counts() helper Add a SQL declaration for the existing C count_tokens() function so users can call pgedge_vectorizer.count_tokens(text) directly to inspect approximate token counts (~4 chars/token, UTF-8 aware). Add refresh_token_counts(chunk_table regclass) to back-fill the token_count column on existing chunk rows without a full recreate_chunks() rebuild. Returns the number of rows updated. Includes regression tests covering approximation values, NULL/UTF-8 handling, the refresh path, and schema-qualified chunk tables. --- Makefile | 2 +- sql/pgedge_vectorizer--1.0--1.1.sql | 38 ++++++ sql/pgedge_vectorizer--1.1.sql | 38 ++++++ test/expected/count_tokens.out | 177 ++++++++++++++++++++++++++++ test/sql/count_tokens.sql | 117 ++++++++++++++++++ 5 files changed, 371 insertions(+), 1 deletion(-) create mode 100644 test/expected/count_tokens.out create mode 100644 test/sql/count_tokens.sql diff --git a/Makefile b/Makefile index aaf3ab9..847c56d 100644 --- a/Makefile +++ b/Makefile @@ -32,7 +32,7 @@ DATA = sql/$(EXTENSION)--$(EXTVERSION).sql \ sql/$(EXTENSION)--1.0-beta3--1.0.sql # Test configuration for pg_regress -REGRESS = setup chunking hybrid_chunking queue vectorization multi_column maintenance edge_cases providers worker cleanup embedding pk_types stale_embeddings hybrid_test +REGRESS = setup chunking hybrid_chunking queue vectorization multi_column maintenance edge_cases providers worker cleanup embedding pk_types stale_embeddings hybrid_test count_tokens REGRESS_OPTS = --inputdir=test --outputdir=test # Documentation files (if any) diff --git a/sql/pgedge_vectorizer--1.0--1.1.sql b/sql/pgedge_vectorizer--1.0--1.1.sql index 37c44c0..72b1b7c 100644 --- a/sql/pgedge_vectorizer--1.0--1.1.sql +++ b/sql/pgedge_vectorizer--1.0--1.1.sql @@ -129,6 +129,17 @@ LANGUAGE C STRICT; COMMENT ON FUNCTION pgedge_vectorizer.bm25_tokenize IS 'Tokenize text and return the non-stopword terms (useful for testing)'; +-- UTF-8 aware approximate token counter (C implementation) +CREATE OR REPLACE FUNCTION pgedge_vectorizer.count_tokens( + p_text TEXT +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_count_tokens_sql' +LANGUAGE C IMMUTABLE STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.count_tokens IS +'Approximate token count using UTF-8 character counting (~4 chars/token). ' +'Used internally by the chunking engine and for stored token_count values.'; + --------------------------------------------------------------------------- -- SQL Functions --------------------------------------------------------------------------- @@ -976,6 +987,33 @@ $$ LANGUAGE plpgsql; COMMENT ON FUNCTION pgedge_vectorizer.recreate_chunks IS 'Delete all chunks and recreate from source table (complete rebuild)'; +--------------------------------------------------------------------------- +-- refresh_token_counts() — recompute stored token_count for existing chunks +--------------------------------------------------------------------------- + +CREATE OR REPLACE FUNCTION pgedge_vectorizer.refresh_token_counts( + p_chunk_table REGCLASS +) RETURNS BIGINT AS $$ +DECLARE + rows_updated BIGINT; +BEGIN + EXECUTE format( + 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content)', + p_chunk_table + ); + GET DIAGNOSTICS rows_updated = ROW_COUNT; + RAISE NOTICE 'refresh_token_counts: updated % rows in %', + rows_updated, p_chunk_table::TEXT; + RETURN rows_updated; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.refresh_token_counts IS +'Recompute token_count for every row in the given chunk table using ' +'count_tokens(). Useful for back-filling accurate counts after adding ' +'new rows outside the normal trigger path, or after a schema change. ' +'Returns the number of rows updated.'; + -- Get configuration summary CREATE OR REPLACE FUNCTION pgedge_vectorizer.show_config() RETURNS TABLE ( diff --git a/sql/pgedge_vectorizer--1.1.sql b/sql/pgedge_vectorizer--1.1.sql index c4475dc..0216773 100644 --- a/sql/pgedge_vectorizer--1.1.sql +++ b/sql/pgedge_vectorizer--1.1.sql @@ -129,6 +129,17 @@ LANGUAGE C STRICT; COMMENT ON FUNCTION pgedge_vectorizer.bm25_tokenize IS 'Tokenize text and return the non-stopword terms (useful for testing)'; +-- UTF-8 aware approximate token counter (C implementation) +CREATE FUNCTION pgedge_vectorizer.count_tokens( + p_text TEXT +) RETURNS INT +AS 'MODULE_PATHNAME', 'pgedge_vectorizer_count_tokens_sql' +LANGUAGE C IMMUTABLE STRICT; + +COMMENT ON FUNCTION pgedge_vectorizer.count_tokens IS +'Approximate token count using UTF-8 character counting (~4 chars/token). ' +'Used internally by the chunking engine and for stored token_count values.'; + --------------------------------------------------------------------------- -- SQL Functions --------------------------------------------------------------------------- @@ -976,6 +987,33 @@ $$ LANGUAGE plpgsql; COMMENT ON FUNCTION pgedge_vectorizer.recreate_chunks IS 'Delete all chunks and recreate from source table (complete rebuild)'; +--------------------------------------------------------------------------- +-- refresh_token_counts() — recompute stored token_count for existing chunks +--------------------------------------------------------------------------- + +CREATE FUNCTION pgedge_vectorizer.refresh_token_counts( + p_chunk_table REGCLASS +) RETURNS BIGINT AS $$ +DECLARE + rows_updated BIGINT; +BEGIN + EXECUTE format( + 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content)', + p_chunk_table + ); + GET DIAGNOSTICS rows_updated = ROW_COUNT; + RAISE NOTICE 'refresh_token_counts: updated % rows in %', + rows_updated, p_chunk_table::TEXT; + RETURN rows_updated; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION pgedge_vectorizer.refresh_token_counts IS +'Recompute token_count for every row in the given chunk table using ' +'count_tokens(). Useful for back-filling accurate counts after adding ' +'new rows outside the normal trigger path, or after a schema change. ' +'Returns the number of rows updated.'; + -- Get configuration summary CREATE FUNCTION pgedge_vectorizer.show_config() RETURNS TABLE ( diff --git a/test/expected/count_tokens.out b/test/expected/count_tokens.out new file mode 100644 index 0000000..488c5d4 --- /dev/null +++ b/test/expected/count_tokens.out @@ -0,0 +1,177 @@ +-- count_tokens.sql +-- Regression tests for count_tokens() and refresh_token_counts(). +--------------------------------------------------------------------------- +-- Test 1: count_tokens() is registered +--------------------------------------------------------------------------- +SELECT proname +FROM pg_proc +WHERE proname = 'count_tokens' + AND pronamespace = ( + SELECT oid FROM pg_namespace WHERE nspname = 'pgedge_vectorizer' + ); + proname +-------------- + count_tokens +(1 row) + +--------------------------------------------------------------------------- +-- Test 2: count_tokens() approximation values +--------------------------------------------------------------------------- +-- 'hello world' = 11 chars, (11+3)/4 = 3 +SELECT pgedge_vectorizer.count_tokens('hello world') AS tokens; + tokens +-------- + 3 +(1 row) + +-- empty string returns 0 +SELECT pgedge_vectorizer.count_tokens('') AS tokens; + tokens +-------- + 0 +(1 row) + +-- 'test' = 4 chars, (4+3)/4 = 1 +SELECT pgedge_vectorizer.count_tokens('test') AS tokens; + tokens +-------- + 1 +(1 row) + +-- NULL returns NULL (STRICT function) +SELECT pgedge_vectorizer.count_tokens(NULL) IS NULL AS is_null; + is_null +--------- + t +(1 row) + +-- '你好世界' = 4 UTF-8 characters, not 12 bytes; (4+3)/4 = 1 +SELECT pgedge_vectorizer.count_tokens('你好世界') AS tokens; + tokens +-------- + 1 +(1 row) + +--------------------------------------------------------------------------- +-- Test 3: refresh_token_counts() is registered +--------------------------------------------------------------------------- +SELECT proname +FROM pg_proc +WHERE proname = 'refresh_token_counts' + AND pronamespace = ( + SELECT oid FROM pg_namespace WHERE nspname = 'pgedge_vectorizer' + ); + proname +---------------------- + refresh_token_counts +(1 row) + +--------------------------------------------------------------------------- +-- Test 4: refresh_token_counts() updates rows and returns the count +--------------------------------------------------------------------------- +CREATE TABLE count_tokens_refresh_test ( + id BIGSERIAL PRIMARY KEY, + content TEXT +); +INSERT INTO count_tokens_refresh_test (content) +VALUES ('Refresh test document one.'), + ('Refresh test document two.'); +SELECT pgedge_vectorizer.enable_vectorization( + 'count_tokens_refresh_test'::regclass, + 'content', + 'token_based', + 100, + 10, + 1536 +); +NOTICE: Using primary key column: id (bigint) +NOTICE: column "sparse_embedding" of relation "count_tokens_refresh_test_content_chunks" already exists, skipping +NOTICE: Vectorization enabled: count_tokens_refresh_test -> count_tokens_refresh_test_content_chunks +NOTICE: Strategy: token_based, chunk_size: 100, overlap: 10 +NOTICE: Processing existing rows... +NOTICE: Processed 2 existing rows + enable_vectorization +---------------------- + +(1 row) + +-- Force token_count to 0 to simulate stale values +UPDATE count_tokens_refresh_test_content_chunks SET token_count = 0; +SELECT COUNT(*) AS zeroed +FROM count_tokens_refresh_test_content_chunks +WHERE token_count = 0; + zeroed +-------- + 2 +(1 row) + +-- refresh_token_counts() should update all rows and return the count +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_refresh_test_content_chunks'::regclass +) >= 1 AS refresh_returned_count; +NOTICE: refresh_token_counts: updated 2 rows in count_tokens_refresh_test_content_chunks + refresh_returned_count +------------------------ + t +(1 row) + +-- All token_counts should now be positive +SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed +FROM count_tokens_refresh_test_content_chunks; + all_refreshed +--------------- + t +(1 row) + +-- Cleanup +SELECT pgedge_vectorizer.disable_vectorization( + 'count_tokens_refresh_test'::regclass, 'content', true +); +NOTICE: Vectorization disabled and chunk table dropped: count_tokens_refresh_test_content_chunks + disable_vectorization +----------------------- + +(1 row) + +DROP TABLE count_tokens_refresh_test; +--------------------------------------------------------------------------- +-- Test 5: refresh_token_counts() works with a schema-qualified table +--------------------------------------------------------------------------- +CREATE SCHEMA count_tokens_ns_test; +CREATE TABLE count_tokens_ns_test.ns_chunks ( + id BIGSERIAL PRIMARY KEY, + source_id BIGINT NOT NULL DEFAULT 1, + content TEXT NOT NULL, + token_count INT +); +INSERT INTO count_tokens_ns_test.ns_chunks (content) +VALUES ('Schema-qualified refresh test.'), + ('Another row for schema test.'); +UPDATE count_tokens_ns_test.ns_chunks SET token_count = 0; +SELECT COUNT(*) AS zeroed +FROM count_tokens_ns_test.ns_chunks +WHERE token_count = 0; + zeroed +-------- + 2 +(1 row) + +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_ns_test.ns_chunks'::regclass +) >= 1 AS refresh_returned_count; +NOTICE: refresh_token_counts: updated 2 rows in count_tokens_ns_test.ns_chunks + refresh_returned_count +------------------------ + t +(1 row) + +SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed +FROM count_tokens_ns_test.ns_chunks; + all_refreshed +--------------- + t +(1 row) + +-- Cleanup +DROP TABLE count_tokens_ns_test.ns_chunks; +DROP SCHEMA count_tokens_ns_test; diff --git a/test/sql/count_tokens.sql b/test/sql/count_tokens.sql new file mode 100644 index 0000000..c39508a --- /dev/null +++ b/test/sql/count_tokens.sql @@ -0,0 +1,117 @@ +-- count_tokens.sql +-- Regression tests for count_tokens() and refresh_token_counts(). + +--------------------------------------------------------------------------- +-- Test 1: count_tokens() is registered +--------------------------------------------------------------------------- +SELECT proname +FROM pg_proc +WHERE proname = 'count_tokens' + AND pronamespace = ( + SELECT oid FROM pg_namespace WHERE nspname = 'pgedge_vectorizer' + ); + +--------------------------------------------------------------------------- +-- Test 2: count_tokens() approximation values +--------------------------------------------------------------------------- + +-- 'hello world' = 11 chars, (11+3)/4 = 3 +SELECT pgedge_vectorizer.count_tokens('hello world') AS tokens; + +-- empty string returns 0 +SELECT pgedge_vectorizer.count_tokens('') AS tokens; + +-- 'test' = 4 chars, (4+3)/4 = 1 +SELECT pgedge_vectorizer.count_tokens('test') AS tokens; + +-- NULL returns NULL (STRICT function) +SELECT pgedge_vectorizer.count_tokens(NULL) IS NULL AS is_null; + +-- '你好世界' = 4 UTF-8 characters, not 12 bytes; (4+3)/4 = 1 +SELECT pgedge_vectorizer.count_tokens('你好世界') AS tokens; + +--------------------------------------------------------------------------- +-- Test 3: refresh_token_counts() is registered +--------------------------------------------------------------------------- +SELECT proname +FROM pg_proc +WHERE proname = 'refresh_token_counts' + AND pronamespace = ( + SELECT oid FROM pg_namespace WHERE nspname = 'pgedge_vectorizer' + ); + +--------------------------------------------------------------------------- +-- Test 4: refresh_token_counts() updates rows and returns the count +--------------------------------------------------------------------------- +CREATE TABLE count_tokens_refresh_test ( + id BIGSERIAL PRIMARY KEY, + content TEXT +); + +INSERT INTO count_tokens_refresh_test (content) +VALUES ('Refresh test document one.'), + ('Refresh test document two.'); + +SELECT pgedge_vectorizer.enable_vectorization( + 'count_tokens_refresh_test'::regclass, + 'content', + 'token_based', + 100, + 10, + 1536 +); + +-- Force token_count to 0 to simulate stale values +UPDATE count_tokens_refresh_test_content_chunks SET token_count = 0; + +SELECT COUNT(*) AS zeroed +FROM count_tokens_refresh_test_content_chunks +WHERE token_count = 0; + +-- refresh_token_counts() should update all rows and return the count +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_refresh_test_content_chunks'::regclass +) >= 1 AS refresh_returned_count; + +-- All token_counts should now be positive +SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed +FROM count_tokens_refresh_test_content_chunks; + +-- Cleanup +SELECT pgedge_vectorizer.disable_vectorization( + 'count_tokens_refresh_test'::regclass, 'content', true +); +DROP TABLE count_tokens_refresh_test; + +--------------------------------------------------------------------------- +-- Test 5: refresh_token_counts() works with a schema-qualified table +--------------------------------------------------------------------------- +CREATE SCHEMA count_tokens_ns_test; + +CREATE TABLE count_tokens_ns_test.ns_chunks ( + id BIGSERIAL PRIMARY KEY, + source_id BIGINT NOT NULL DEFAULT 1, + content TEXT NOT NULL, + token_count INT +); + +INSERT INTO count_tokens_ns_test.ns_chunks (content) +VALUES ('Schema-qualified refresh test.'), + ('Another row for schema test.'); + +UPDATE count_tokens_ns_test.ns_chunks SET token_count = 0; + +SELECT COUNT(*) AS zeroed +FROM count_tokens_ns_test.ns_chunks +WHERE token_count = 0; + +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_ns_test.ns_chunks'::regclass +) >= 1 AS refresh_returned_count; + +SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed +FROM count_tokens_ns_test.ns_chunks; + +-- Cleanup +DROP TABLE count_tokens_ns_test.ns_chunks; +DROP SCHEMA count_tokens_ns_test; From eeea965a72197f3c2ede7b64810bc4bc90ccd912 Mon Sep 17 00:00:00 2001 From: Kazim Haider Date: Fri, 29 May 2026 01:28:06 +0500 Subject: [PATCH 2/4] Use count_tokens() on all chunk write paths for consistent token_count The three PLpgSQL insert paths (enable_vectorization, vectorization_trigger, recreate_chunks) were using length(chunk_text) / 4 (floor division) while count_tokens() and refresh_token_counts() use ceiling division ((char_count + 3) / 4). This meant refresh_token_counts() would silently change values the trigger had just written for any chunk whose length is not divisible by 4. Replace all six occurrences (three per SQL file) with pgedge_vectorizer.count_tokens(chunk_text) so every write path uses the same formula. --- sql/pgedge_vectorizer--1.0--1.1.sql | 6 +++--- sql/pgedge_vectorizer--1.1.sql | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/sql/pgedge_vectorizer--1.0--1.1.sql b/sql/pgedge_vectorizer--1.0--1.1.sql index 72b1b7c..e3ae14a 100644 --- a/sql/pgedge_vectorizer--1.0--1.1.sql +++ b/sql/pgedge_vectorizer--1.0--1.1.sql @@ -355,7 +355,7 @@ BEGIN (sparse_embedding IS NULL) AS needs_sparse', chunk_table, chunk_table, chunk_table, chunk_table, chunk_table) USING row_record.pk_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id, needs_embedding, needs_sparse; -- Queue if dense or sparse work is needed. @@ -672,7 +672,7 @@ BEGIN VALUES ($1::%s, $2, $3, $4) RETURNING id', chunk_table, pk_type) USING source_id_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id; -- Queue for embedding @@ -967,7 +967,7 @@ BEGIN VALUES ($1::%s, $2, $3, $4) RETURNING id', chunk_table_name, pk_type) USING row_record.pk_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id; -- Queue for embedding diff --git a/sql/pgedge_vectorizer--1.1.sql b/sql/pgedge_vectorizer--1.1.sql index 0216773..3582083 100644 --- a/sql/pgedge_vectorizer--1.1.sql +++ b/sql/pgedge_vectorizer--1.1.sql @@ -355,7 +355,7 @@ BEGIN (sparse_embedding IS NULL) AS needs_sparse', chunk_table, chunk_table, chunk_table, chunk_table, chunk_table) USING row_record.pk_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id, needs_embedding, needs_sparse; -- Queue if dense or sparse work is needed. @@ -672,7 +672,7 @@ BEGIN VALUES ($1::%s, $2, $3, $4) RETURNING id', chunk_table, pk_type) USING source_id_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id; -- Queue for embedding @@ -967,7 +967,7 @@ BEGIN VALUES ($1::%s, $2, $3, $4) RETURNING id', chunk_table_name, pk_type) USING row_record.pk_val, i, chunk_text, - length(chunk_text) / 4 -- Approximate token count + pgedge_vectorizer.count_tokens(chunk_text) INTO chunk_id; -- Queue for embedding From eb41ada2504acdf8e5f0ec26f30dc0d83a86a3ed Mon Sep 17 00:00:00 2001 From: Kazim Haider Date: Fri, 29 May 2026 15:09:39 +0500 Subject: [PATCH 3/4] Skip no-op rewrites in refresh_token_counts() and tighten test assertions Add WHERE token_count IS DISTINCT FROM count_tokens(content) so the UPDATE only touches rows whose stored value actually differs. On large chunk tables this avoids unnecessary WAL writes and dead tuples. The return value now reflects rows that changed, not rows scanned. Tighten the two refresh test assertions from >= 1 to = COUNT(*) so a partial update cannot silently pass the regression test. --- sql/pgedge_vectorizer--1.0--1.1.sql | 4 +++- sql/pgedge_vectorizer--1.1.sql | 4 +++- test/expected/count_tokens.out | 10 +++++++--- test/sql/count_tokens.sql | 10 +++++++--- 4 files changed, 20 insertions(+), 8 deletions(-) diff --git a/sql/pgedge_vectorizer--1.0--1.1.sql b/sql/pgedge_vectorizer--1.0--1.1.sql index e3ae14a..b4aebfb 100644 --- a/sql/pgedge_vectorizer--1.0--1.1.sql +++ b/sql/pgedge_vectorizer--1.0--1.1.sql @@ -998,7 +998,9 @@ DECLARE rows_updated BIGINT; BEGIN EXECUTE format( - 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content)', + 'UPDATE %s + SET token_count = pgedge_vectorizer.count_tokens(content) + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content)', p_chunk_table ); GET DIAGNOSTICS rows_updated = ROW_COUNT; diff --git a/sql/pgedge_vectorizer--1.1.sql b/sql/pgedge_vectorizer--1.1.sql index 3582083..6e34ace 100644 --- a/sql/pgedge_vectorizer--1.1.sql +++ b/sql/pgedge_vectorizer--1.1.sql @@ -998,7 +998,9 @@ DECLARE rows_updated BIGINT; BEGIN EXECUTE format( - 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content)', + 'UPDATE %s + SET token_count = pgedge_vectorizer.count_tokens(content) + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content)', p_chunk_table ); GET DIAGNOSTICS rows_updated = ROW_COUNT; diff --git a/test/expected/count_tokens.out b/test/expected/count_tokens.out index 488c5d4..3c29442 100644 --- a/test/expected/count_tokens.out +++ b/test/expected/count_tokens.out @@ -105,10 +105,12 @@ WHERE token_count = 0; 2 (1 row) --- refresh_token_counts() should update all rows and return the count +-- refresh_token_counts() should update all rows and return the exact count SELECT pgedge_vectorizer.refresh_token_counts( 'count_tokens_refresh_test_content_chunks'::regclass -) >= 1 AS refresh_returned_count; +) = ( + SELECT COUNT(*) FROM count_tokens_refresh_test_content_chunks +) AS refresh_returned_count; NOTICE: refresh_token_counts: updated 2 rows in count_tokens_refresh_test_content_chunks refresh_returned_count ------------------------ @@ -158,7 +160,9 @@ WHERE token_count = 0; SELECT pgedge_vectorizer.refresh_token_counts( 'count_tokens_ns_test.ns_chunks'::regclass -) >= 1 AS refresh_returned_count; +) = ( + SELECT COUNT(*) FROM count_tokens_ns_test.ns_chunks +) AS refresh_returned_count; NOTICE: refresh_token_counts: updated 2 rows in count_tokens_ns_test.ns_chunks refresh_returned_count ------------------------ diff --git a/test/sql/count_tokens.sql b/test/sql/count_tokens.sql index c39508a..c6212ae 100644 --- a/test/sql/count_tokens.sql +++ b/test/sql/count_tokens.sql @@ -68,10 +68,12 @@ SELECT COUNT(*) AS zeroed FROM count_tokens_refresh_test_content_chunks WHERE token_count = 0; --- refresh_token_counts() should update all rows and return the count +-- refresh_token_counts() should update all rows and return the exact count SELECT pgedge_vectorizer.refresh_token_counts( 'count_tokens_refresh_test_content_chunks'::regclass -) >= 1 AS refresh_returned_count; +) = ( + SELECT COUNT(*) FROM count_tokens_refresh_test_content_chunks +) AS refresh_returned_count; -- All token_counts should now be positive SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed @@ -107,7 +109,9 @@ WHERE token_count = 0; SELECT pgedge_vectorizer.refresh_token_counts( 'count_tokens_ns_test.ns_chunks'::regclass -) >= 1 AS refresh_returned_count; +) = ( + SELECT COUNT(*) FROM count_tokens_ns_test.ns_chunks +) AS refresh_returned_count; SELECT COALESCE(BOOL_AND(token_count > 0), false) AS all_refreshed FROM count_tokens_ns_test.ns_chunks; From 74f9073a4de8bbd017db9242e620f96325f17a09 Mon Sep 17 00:00:00 2001 From: Kazim Haider Date: Mon, 8 Jun 2026 15:41:28 +0500 Subject: [PATCH 4/4] Add PG_FUNCTION_INFO_V1 wrapper for count_tokens SQL callable The SQL migration files declared pgedge_vectorizer_count_tokens_sql as a C symbol but the wrapper was never added to src/tokenizer.c, causing every call to pgedge_vectorizer.count_tokens() (and all three chunk-write paths that call it) to fail with "could not find function" at runtime. Co-Authored-By: Claude Sonnet 4.6 --- src/tokenizer.c | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/tokenizer.c b/src/tokenizer.c index fae2464..da0ffba 100644 --- a/src/tokenizer.c +++ b/src/tokenizer.c @@ -55,6 +55,15 @@ count_tokens(const char *text, const char *model) return token_estimate; } +PG_FUNCTION_INFO_V1(pgedge_vectorizer_count_tokens_sql); +Datum +pgedge_vectorizer_count_tokens_sql(PG_FUNCTION_ARGS) +{ + text *txt = PG_GETARG_TEXT_PP(0); + char *s = text_to_cstring(txt); + PG_RETURN_INT32(count_tokens(s, pgedge_vectorizer_model)); +} + /* * Tokenize text into token IDs *