diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 6a41932..f1a1e74 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -11,6 +11,7 @@ import json import os import pathlib import re +from sqlite_fts4 import rank_bm25 import sys import textwrap import uuid @@ -189,6 +190,9 @@ class Database: else: register(fn) + def register_fts4_bm25(self): + self.register_function(rank_bm25, deterministic=True) + def execute(self, sql, parameters=None): if self._tracer: self._tracer(sql, parameters) @@ -1330,52 +1334,51 @@ class Table(Queryable): assert fts_table, "Full-text search is not configured for table '{}'".format( self.name ) - if self.db[fts_table].virtual_table_using == "FTS5": - sql = textwrap.dedent( - """ - with {original} as ( - select - rowid, - {columns} - from [{dbtable}] - ) + virtual_table_using = self.db[fts_table].virtual_table_using + sql = textwrap.dedent( + """ + with {original} as ( select - {original}.*, - [{fts}].rank as {rank} - from - [{original}] - join [{fts}] on [{original}].rowid = [{fts}].rowid - where - [{fts}] match :query - order by - {order} - {limit} - """ - ).strip() + rowid, + {columns} + from [{dbtable}] + ) + select + {original}.*, + {rank_implementation} as {rank} + from + [{original}] + join [{fts_table}] on [{original}].rowid = [{fts_table}].rowid + where + [{fts_table}] match :query + order by + {order} + {limit} + """ + ).strip() + if virtual_table_using == "FTS5": + rank_implementation = "[{}].rank".format(fts_table) else: - if order == rank or order is None: - order = "rowid" - sql = textwrap.dedent( - """ - select * from "{dbtable}" where rowid in ( - select rowid from [{fts}] - where [{fts}] match :query + self.db.register_fts4_bm25() + rank_implementation = "-rank_bm25(matchinfo([{}], 'pcnalx'))".format( + fts_table ) - order by {order} - """ - ).strip() return sql.format( dbtable=self.name, original=original, columns=columns_sql, rank=rank, - fts=fts_table, + rank_implementation=rank_implementation, + fts_table=fts_table, order=order or "{} desc".format(rank), limit="limit {}".format(limit) if limit else "", ).strip() def search(self, q, order=None): - return self.db.execute(self.search_sql(order=order), {"query": q}).fetchall() + cursor = self.db.execute(self.search_sql(order=order), {"query": q}) + columns = [c[0] for c in cursor.description] + for row in cursor: + yield dict(zip(columns, row)) def value_or_default(self, key, value): return self._defaults[key] if value is DEFAULT else value diff --git a/tests/test_cli.py b/tests/test_cli.py index 7e1365d..f9b2509 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -510,14 +510,14 @@ def test_rebuild_fts(db_path, tables): ["c1", "c2", "c3"], fts_version="FTS5", create_triggers=True ) # Search should work - assert db["fts4_table"].search("verb1") - assert db["fts5_table"].search("verb1") + assert list(db["fts4_table"].search("verb1")) + assert list(db["fts5_table"].search("verb1")) # Deleting _fts_segments to break FTS4 with db.conn: db["fts4_table_fts_segments"].delete_where() # Now this should error: with pytest.raises(sqlite3.DatabaseError): - db["fts4_table"].search("verb1") + list(db["fts4_table"].search("verb1")) # Replicate docsize error from this issue for FTS5 # https://github.com/simonw/sqlite-utils/issues/149 assert db["fts5_table_fts_docsize"].count == 10000 @@ -530,10 +530,10 @@ def test_rebuild_fts(db_path, tables): assert 0 == result.exit_code fixed_tables = tables or ["fts4_table", "fts5_table"] if "fts4_table" in fixed_tables: - assert db["fts4_table"].search("verb1") + assert list(db["fts4_table"].search("verb1")) else: with pytest.raises(sqlite3.DatabaseError): - db["fts4_table"].search("verb1") + list(db["fts4_table"].search("verb1")) if "fts5_table" in fixed_tables: assert db["fts5_table_fts_docsize"].count == 10000 else: diff --git a/tests/test_fts.py b/tests/test_fts.py index 9ae692f..0b0c106 100644 --- a/tests/test_fts.py +++ b/tests/test_fts.py @@ -29,9 +29,25 @@ def test_enable_fts(fresh_db): "searchable_fts_docsize", "searchable_fts_stat", ] == fresh_db.table_names() - 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") + assert [ + { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + "rank": 0.0, + } + ] == list(table.search("tanuki")) + assert [ + { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + "rank": 0.0, + } + ] == list(table.search("usa")) + assert [] == list(table.search("bar")) def test_enable_fts_escape_table_names(fresh_db): @@ -49,9 +65,25 @@ 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 running tricksters", "Japan", "foo")] == table.search("tanuki") - assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") - assert [] == table.search("bar") + assert [ + { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + "rank": 0.0, + } + ] == list(table.search("tanuki")) + assert [ + { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + "rank": 0.0, + } + ] == list(table.search("usa")) + assert [] == list(table.search("bar")) def test_enable_fts_table_names_containing_spaces(fresh_db): @@ -72,12 +104,21 @@ def test_populate_fts(fresh_db): table = fresh_db["populatable"] table.insert(search_records[0]) table.enable_fts(["text", "country"], fts_version="FTS4") - assert [] == table.search("trash pandas") + assert [] == list(table.search("trash pandas")) table.insert(search_records[1]) - assert [] == table.search("trash pandas") + assert [] == list(table.search("trash pandas")) # Now run populate_fts to make this record available table.populate_fts(["text", "country"]) - assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") + rows = list(table.search("usa")) + assert [ + { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + "rank": 0.5108256237659907, + } + ] == rows def test_populate_fts_escape_table_names(fresh_db): @@ -85,12 +126,20 @@ def test_populate_fts_escape_table_names(fresh_db): table = fresh_db["http://example.com"] table.insert(search_records[0]) table.enable_fts(["text", "country"], fts_version="FTS4") - assert [] == table.search("trash pandas") + assert [] == list(table.search("trash pandas")) table.insert(search_records[1]) - assert [] == table.search("trash pandas") + assert [] == list(table.search("trash pandas")) # Now run populate_fts to make this record available table.populate_fts(["text", "country"]) - assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") + assert [ + { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + "rank": 0.5108256237659907, + } + ] == list(table.search("usa")) def test_fts_tokenize(fresh_db): @@ -103,7 +152,7 @@ def test_fts_tokenize(fresh_db): ["text", "country"], fts_version="FTS{}".format(fts_version), ) - assert [] == table.search("bite") + assert [] == list(table.search("bite")) # Test WITH stemming table.disable_fts() table.enable_fts( @@ -111,9 +160,14 @@ def test_fts_tokenize(fresh_db): fts_version="FTS{}".format(fts_version), tokenize="porter", ) - assert [("racoons are biting trash pandas", "USA", "bar")] == table.search( - "bite", order="rowid" - ) + rows = list(table.search("bite", order="rowid")) + assert len(rows) == 1 + assert { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + }.items() <= rows[0].items() def test_optimize_fts(fresh_db): @@ -132,15 +186,34 @@ def test_optimize_fts(fresh_db): fresh_db[table_name].optimize() -def test_enable_fts_w_triggers(fresh_db): +def test_enable_fts_with_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 running tricksters", "Japan", "foo")] == table.search("tanuki") + rows1 = list(table.search("tanuki")) + assert len(rows1) == 1 + assert rows1 == [ + { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + "rank": 0.0, + } + ] table.insert(search_records[1]) # Triggers will auto-populate FTS virtual table, not need to call populate_fts() - assert [("racoons are biting trash pandas", "USA", "bar")] == table.search("usa") - assert [] == table.search("bar") + rows2 = list(table.search("usa")) + assert rows2 == [ + { + "rowid": 2, + "text": "racoons are biting trash pandas", + "country": "USA", + "not_searchable": "bar", + "rank": 0.0, + } + ] + assert [] == list(table.search("bar")) @pytest.mark.parametrize("create_triggers", [True, False]) @@ -183,15 +256,29 @@ def test_rebuild_fts(fresh_db, table_to_fix): table.insert(search_records[0]) table.enable_fts(["text", "country"]) # Run a search - assert [("tanuki are running tricksters", "Japan", "foo")] == table.search("tanuki") + rows = list(table.search("tanuki")) + assert len(rows) == 1 + assert { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + }.items() <= rows[0].items() # Delete from searchable_fts_data fresh_db["searchable_fts_data"].delete_where() # This should have broken the index with pytest.raises(sqlite3.DatabaseError): - table.search("tanuki") + list(table.search("tanuki")) # Running rebuild_fts() should fix it fresh_db[table_to_fix].rebuild_fts() - assert [("tanuki are running tricksters", "Japan", "foo")] == table.search("tanuki") + rows2 = list(table.search("tanuki")) + assert len(rows2) == 1 + assert { + "rowid": 1, + "text": "tanuki are running tricksters", + "country": "Japan", + "not_searchable": "foo", + }.items() <= rows2[0].items() @pytest.mark.parametrize("invalid_table", ["does_not_exist", "not_searchable"]) diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 6a58ec0..dab9afb 100644 --- a/tests/test_tracer.py +++ b/tests/test_tracer.py @@ -27,10 +27,6 @@ def test_tracer(): None, ), ("select name from sqlite_master where type = 'view'", None), - ( - 'select * from "dogs" where rowid in (\n select rowid from [dogs_fts]\n where [dogs_fts] match :search\n)\norder by rowid', - ("Cleopaws",), - ), ] @@ -46,17 +42,25 @@ def test_with_tracer(): assert len(collected) == 0 with db.tracer(tracer): - db["dogs"].search("Cleopaws") + list(db["dogs"].search("Cleopaws")) - assert len(collected) == 2 + assert len(collected) == 7 assert collected == [ ("select name from sqlite_master where type = 'view'", None), + ("select name from sqlite_master where type = 'table'", None), + ("PRAGMA table_info([dogs])", None), ( - 'select * from "dogs" where rowid in (\n select rowid from [dogs_fts]\n where [dogs_fts] match :search\n)\norder by rowid', - ("Cleopaws",), + "SELECT name FROM sqlite_master\n WHERE rootpage = 0\n AND (\n sql LIKE '%VIRTUAL TABLE%USING FTS%content=%dogs%'\n OR (\n tbl_name = \"dogs\"\n AND sql LIKE '%VIRTUAL TABLE%USING FTS%'\n )\n )", + None, + ), + ("select name from sqlite_master where type = 'view'", None), + ("select sql from sqlite_master where name = ?", ("dogs_fts",)), + ( + "with original as (\n select\n rowid,\n *\n from [dogs]\n)\nselect\n original.*,\n [dogs_fts].rank as rank\nfrom\n [original]\n join [dogs_fts] on [original].rowid = [dogs_fts].rowid\nwhere\n [dogs_fts] match :query\norder by\n rank desc", + {"query": "Cleopaws"}, ), ] # Outside the with block collected should not be appended to db["dogs"].insert({"name": "Cleopaws"}) - assert len(collected) == 2 + assert len(collected) == 7