From 617e6f070c85be66ea04c80b78dafd08c875f8c8 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sat, 1 Aug 2020 13:40:36 -0700 Subject: [PATCH] enable_fts(..., tokenize=X) parameter, refs #130 --- docs/python-api.rst | 8 ++++++++ sqlite_utils/db.py | 9 ++++++-- tests/test_fts.py | 50 ++++++++++++++++++++++++++++++++++++--------- 3 files changed, 55 insertions(+), 12 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index e6b1af6..9fc51f7 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -1135,6 +1135,14 @@ A better solution is to use database triggers. You can set up database triggers dogs.enable_fts(["name", "twitter"], fts_version="FTS4") +You can customize the tokenizer configured for the table using the ``tokenize=`` parameter. For example, to enable Porter stemming, where English words like "running" will match stemmed alternatives such as "run", use ``tokenize="porter"``: + +.. code-block:: python + + db["articles"].enable_fts(["headline", "body"], tokenize="porter") + +The SQLite documentation has more on `FTS5 tokenizers `__ and `FTS4 tokenizers `__. ``porter`` is a valid option for both. + To remove the FTS tables and triggers you created, use the ``disable_fts()`` table method: .. code-block:: python diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index ee26433..7d3a8fe 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -763,17 +763,22 @@ class Table(Queryable): ) self.db.add_foreign_keys([(self.name, column, other_table, other_column)]) - def enable_fts(self, columns, fts_version="FTS5", create_triggers=False): + def enable_fts( + self, columns, fts_version="FTS5", create_triggers=False, tokenize=None + ): "Enables FTS on the specified columns." sql = """ CREATE VIRTUAL TABLE [{table}_fts] USING {fts_version} ( - {columns}, + {columns},{tokenize} content=[{table}] ); """.format( table=self.name, columns=", ".join("[{}]".format(c) for c in columns), fts_version=fts_version, + tokenize="\n tokenize='{}',\n".format(tokenize) + if tokenize + else "", ) self.db.conn.executescript(sql) self.populate_fts(columns) diff --git a/tests/test_fts.py b/tests/test_fts.py index 7817477..bac4117 100644 --- a/tests/test_fts.py +++ b/tests/test_fts.py @@ -1,8 +1,16 @@ import pytest search_records = [ - {"text": "tanuki are tricksters", "country": "Japan", "not_searchable": "foo"}, - {"text": "racoons are trash pandas", "country": "USA", "not_searchable": "bar"}, + { + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + }, + { + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + }, ] @@ -19,8 +27,8 @@ def test_enable_fts(fresh_db): "searchable_fts_docsize", "searchable_fts_stat", ] == fresh_db.table_names() - assert [("tanuki are tricksters", "Japan", "foo")] == table.search("tanuki") - assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [("tanuki are running tricksters", "Japan", "foo")] == table.search("tanuki") + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") assert [] == table.search("bar") @@ -39,8 +47,8 @@ def test_enable_fts_escape_table_names(fresh_db): "http://example.com_fts_docsize", "http://example.com_fts_stat", ] == fresh_db.table_names() - assert [("tanuki are tricksters", "Japan", "foo")] == table.search("tanuki") - assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [("tanuki are running tricksters", "Japan", "foo")] == table.search("tanuki") + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") assert [] == table.search("bar") @@ -67,7 +75,7 @@ def test_populate_fts(fresh_db): assert [] == table.search("trash pandas") # Now run populate_fts to make this record available table.populate_fts(["text", "country"]) - assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") def test_populate_fts_escape_table_names(fresh_db): @@ -80,7 +88,29 @@ def test_populate_fts_escape_table_names(fresh_db): assert [] == table.search("trash pandas") # Now run populate_fts to make this record available table.populate_fts(["text", "country"]) - assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") + + +def test_fts_tokenize(fresh_db): + for fts_version in ("4", "5"): + table_name = "searchable_{}".format(fts_version) + table = fresh_db[table_name] + table.insert_all(search_records) + # Test without porter stemming + table.enable_fts( + ["text", "country"], fts_version="FTS{}".format(fts_version), + ) + assert [] == table.search("bite") + # Test WITH stemming + table.disable_fts() + table.enable_fts( + ["text", "country"], + fts_version="FTS{}".format(fts_version), + tokenize="porter", + ) + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search( + "bite" + ) def test_optimize_fts(fresh_db): @@ -103,10 +133,10 @@ def test_enable_fts_w_triggers(fresh_db): table = fresh_db["searchable"] table.insert(search_records[0]) table.enable_fts(["text", "country"], fts_version="FTS4", create_triggers=True) - assert [("tanuki are tricksters", "Japan", "foo")] == table.search("tanuki") + assert [("tanuki are running tricksters", "Japan", "foo")] == table.search("tanuki") table.insert(search_records[1]) # Triggers will auto-populate FTS virtual table, not need to call populate_fts() - assert [("racoons are trash pandas", "USA", "bar")] == table.search("usa") + assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") assert [] == table.search("bar")