include_rank parameter for Table.search

* Add `include_rank` parameter to `Table.search`
* Test for .search(include_rank)
* Docs for table.search(include_rank)

#628

Refs #480

---------

Co-authored-by: Simon Willison <[email protected]>

Author: liunux4odoo

docs/python-api.rst

@@ -2311,6 +2311,9 @@ The ``.search()`` method also accepts the following optional parameters:
 ``where_args`` dictionary
     Arguments to use for ``:param`` placeholders in the extra WHERE clause
 
+``include_rank`` bool
+    If set a ``rank`` column will be included with the BM25 ranking score - for FTS5 tables only.
+
 ``quote`` bool
     Apply :ref:`FTS quoting rules <python_api_quote_fts>` to the search query, disabling advanced query syntax in a way that avoids surprising errors.
 

sqlite_utils/db.py

@@ -2678,6 +2678,7 @@ def search(
         offset: Optional[int] = None,
         where: Optional[str] = None,
         where_args: Optional[Union[Iterable, dict]] = None,
+        include_rank: bool = False,
         quote: bool = False,
     ) -> Generator[dict, None, None]:
         """
@@ -2691,6 +2692,7 @@ def search(
         :param offset: Optional integer SQL offset.
         :param where: Extra SQL fragment for the WHERE clause
         :param where_args: Arguments to use for :param placeholders in the extra WHERE clause
+        :param include_rank: Select the search rank column in the final query
         :param quote: Apply quoting to disable any special characters in the search query
 
         See :ref:`python_api_fts_search`.
@@ -2710,6 +2712,7 @@ def search(
                 limit=limit,
                 offset=offset,
                 where=where,
+                include_rank=include_rank,
             ),
             args,
         )

tests/test_fts.py

@@ -1,6 +1,7 @@
 import pytest
 from sqlite_utils import Database
 from sqlite_utils.utils import sqlite3
+from unittest.mock import ANY
 
 search_records = [
     {
@@ -126,6 +127,32 @@ def test_search_where_args_disallows_query(fresh_db):
     )
 
 
+def test_search_include_rank(fresh_db):
+    table = fresh_db["t"]
+    table.insert_all(search_records)
+    table.enable_fts(["text", "country"], fts_version="FTS5")
+    results = list(table.search("are", include_rank=True))
+    assert results == [
+        {
+            "rowid": 1,
+            "text": "tanuki are running tricksters",
+            "country": "Japan",
+            "not_searchable": "foo",
+            "rank": ANY,
+        },
+        {
+            "rowid": 2,
+            "text": "racoons are biting trash pandas",
+            "country": "USA",
+            "not_searchable": "bar",
+            "rank": ANY,
+        },
+    ]
+    assert isinstance(results[0]["rank"], float)
+    assert isinstance(results[1]["rank"], float)
+    assert results[0]["rank"] < results[1]["rank"]
+
+
 def test_enable_fts_table_names_containing_spaces(fresh_db):
     table = fresh_db["test"]
     table.insert({"column with spaces": "in its name"})