Skip to content

add include_rank parameter to Table.search #628

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docs/python-api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2308,6 +2308,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.

Expand Down
3 changes: 3 additions & 0 deletions sqlite_utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -2641,6 +2641,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]:
"""
Expand All @@ -2654,6 +2655,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`.
Expand All @@ -2673,6 +2675,7 @@ def search(
limit=limit,
offset=offset,
where=where,
include_rank=include_rank,
),
args,
)
Expand Down
27 changes: 27 additions & 0 deletions tests/test_fts.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import pytest
from sqlite_utils import Database
from sqlite_utils.utils import sqlite3
from unittest.mock import ANY

search_records = [
{
Expand Down Expand Up @@ -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"})
Expand Down
Loading