Skip to content

Commit

Permalink
Update elasticsearch_backend.py, added support 'track_total_hits'
Browse files Browse the repository at this point in the history
Support track_total_hits Handling in ElasticsearchSearchBackend

- Implemented handling for `track_total_hits` in the Elasticsearch search backend.
- Supports `bool` values (True/False) and `int` values.
- Added a warning for invalid `track_total_hits` values.
- Ensured the parameter is only added to search kwargs when valid.
  • Loading branch information
nikolaysm authored Aug 9, 2024
1 parent d84327e commit 12675e8
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions haystack/backends/elasticsearch_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
FUZZY_MAX_EXPANSIONS,
FUZZY_MIN_SIM,
ID,
TRACK_TOTAL_HITS,
)
from haystack.exceptions import MissingDependency, MoreLikeThisError, SkipDocument
from haystack.inputs import Clean, Exact, PythonData, Raw
Expand Down Expand Up @@ -545,6 +546,28 @@ def build_search_kwargs(
if extra_kwargs:
kwargs.update(extra_kwargs)

# If TRACK_TOTAL_HITS is False, 0, or None, do not include the parameter
if TRACK_TOTAL_HITS:
# Define a mapping for the track_total_hits parameter
# - If TRACK_TOTAL_HITS is True (bool), map to "true" (string)
# - If TRACK_TOTAL_HITS is an integer, use its value directly
track_total_hits_mapper = {
bool: "true", # Maps boolean True to "true"
int: TRACK_TOTAL_HITS, # Maps integer to its value
}

Check failure on line 558 in haystack/backends/elasticsearch_backend.py

View workflow job for this annotation

GitHub Actions / ruff

Ruff (W293)

haystack/backends/elasticsearch_backend.py:558:1: W293 Blank line contains whitespace
# Get the mapped value based on the type of TRACK_TOTAL_HITS
# If the type is not in the mapper, fallback to False
track_total_hits = track_total_hits_mapper.get(type(TRACK_TOTAL_HITS), False)

Check failure on line 562 in haystack/backends/elasticsearch_backend.py

View workflow job for this annotation

GitHub Actions / ruff

Ruff (W293)

haystack/backends/elasticsearch_backend.py:562:1: W293 Blank line contains whitespace
# If a valid track_total_hits value is obtained, add it to search_kwargs
if track_total_hits:
search_kwargs['track_total_hits'] = track_total_hits

Check failure on line 565 in haystack/backends/elasticsearch_backend.py

View workflow job for this annotation

GitHub Actions / ruff

Ruff (F821)

haystack/backends/elasticsearch_backend.py:565:17: F821 Undefined name `search_kwargs`
else:
# Issue a warning if TRACK_TOTAL_HITS is not of type bool or int
warnings.warn(
"Wrong value of HAYSTACK_TRACK_TOTAL_HITS is provided. Valid options are `bool` or `int`."
)
return kwargs

@log_query
Expand Down

0 comments on commit 12675e8

Please sign in to comment.