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..b4aebfb 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 --------------------------------------------------------------------------- @@ -344,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. @@ -661,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 @@ -956,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 @@ -976,6 +987,35 @@ $$ 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) + WHERE token_count IS DISTINCT FROM 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..6e34ace 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 --------------------------------------------------------------------------- @@ -344,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. @@ -661,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 @@ -956,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 @@ -976,6 +987,35 @@ $$ 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) + WHERE token_count IS DISTINCT FROM 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/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 * diff --git a/test/expected/count_tokens.out b/test/expected/count_tokens.out new file mode 100644 index 0000000..3c29442 --- /dev/null +++ b/test/expected/count_tokens.out @@ -0,0 +1,181 @@ +-- 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 exact count +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_refresh_test_content_chunks'::regclass +) = ( + 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 +------------------------ + 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 +) = ( + 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 +------------------------ + 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..c6212ae --- /dev/null +++ b/test/sql/count_tokens.sql @@ -0,0 +1,121 @@ +-- 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 exact count +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_refresh_test_content_chunks'::regclass +) = ( + 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 +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 +) = ( + 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; + +-- Cleanup +DROP TABLE count_tokens_ns_test.ns_chunks; +DROP SCHEMA count_tokens_ns_test;