280 lines
9.4 KiB
Python
280 lines
9.4 KiB
Python
|
"""
|
||
|
|
Elasticsearch search backend.
|
||
|
|
|
||
|
|
The query builder is a pure function (:func:`build_query`) that turns a
|
||
|
|
:class:`~apps.search.types.SearchRequest` into an Elasticsearch request body.
|
||
|
|
Keeping it free of I/O means the ranking logic is unit-testable without a
|
||
|
|
running cluster - see ``tests/test_query_builder.py``.
|
||
|
|
|
||
|
|
Ranking model
|
||
|
|
-------------
|
||
|
|
``must`` the keyword must match somewhere, with per-field boosts that encode
|
||
|
|
"a name hit is worth more than a summary hit", plus ``fuzziness:
|
||
|
|
AUTO`` so ``pyhton`` still finds ``python``.
|
||
|
|
``should`` an exact phrase match on the strong fields adds a large bonus, which
|
||
|
|
pushes "product manager" above documents that merely contain both
|
||
|
|
words far apart.
|
||
|
|
``filter`` structured constraints (skills, role, seniority, country, industry,
|
||
|
|
employer, years of experience). They live in the filter context, so
|
||
|
|
they are cached by Elasticsearch and contribute nothing to the score.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from django.conf import settings
|
||
|
|
from elasticsearch import ApiError, Elasticsearch, TransportError
|
||
|
|
|
||
|
|
from apps.core.exceptions import SearchBackendError
|
||
|
|
from apps.search.backends.base import SearchBackend
|
||
|
|
from apps.search.indexing import alias_name, get_client
|
||
|
|
from apps.search.types import (
|
||
|
|
FACET_FIELDS,
|
||
|
|
SORT_EXPERIENCE_ASC,
|
||
|
|
SORT_EXPERIENCE_DESC,
|
||
|
|
SORT_NAME,
|
||
|
|
FacetBucket,
|
||
|
|
SearchHit,
|
||
|
|
SearchRequest,
|
||
|
|
SearchResult,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
#: Field boosts for the free-text part of the query.
|
||
|
|
SEARCH_FIELDS = [
|
||
|
|
"full_name^6",
|
||
|
|
"job_title^4",
|
||
|
|
"skills^3.5",
|
||
|
|
"company^2.5",
|
||
|
|
"past_companies^1.5",
|
||
|
|
"industry^1.5",
|
||
|
|
"majors^1.5",
|
||
|
|
"education",
|
||
|
|
"location_name",
|
||
|
|
"headline^2",
|
||
|
|
"summary",
|
||
|
|
"interests",
|
||
|
|
]
|
||
|
|
|
||
|
|
#: Narrower field set for the phrase bonus - phrases in a summary are noise.
|
||
|
|
PHRASE_FIELDS = ["full_name^8", "job_title^6", "skills^5", "company^4", "headline^3"]
|
||
|
|
|
||
|
|
HIGHLIGHT_FIELDS = {
|
||
|
|
"full_name": {"number_of_fragments": 0},
|
||
|
|
"job_title": {"number_of_fragments": 0},
|
||
|
|
"skills": {"number_of_fragments": 3},
|
||
|
|
"summary": {"fragment_size": 160, "number_of_fragments": 1},
|
||
|
|
}
|
||
|
|
|
||
|
|
#: Maps an API filter to the indexed keyword field it constrains.
|
||
|
|
FILTER_FIELDS = {
|
||
|
|
"skills": "skills.keyword",
|
||
|
|
"roles": "job_title_role",
|
||
|
|
"seniority": "seniority_levels",
|
||
|
|
"countries": "location_country",
|
||
|
|
"industries": "industry.keyword",
|
||
|
|
"companies": "company.keyword",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def build_query(request: SearchRequest) -> dict[str, Any]:
|
||
|
|
"""Translate a search request into an Elasticsearch request body."""
|
||
|
|
must: list[dict[str, Any]] = []
|
||
|
|
should: list[dict[str, Any]] = []
|
||
|
|
|
||
|
|
if request.q:
|
||
|
|
must.append(
|
||
|
|
{
|
||
|
|
"multi_match": {
|
||
|
|
"query": request.q,
|
||
|
|
"type": "best_fields",
|
||
|
|
"fields": SEARCH_FIELDS,
|
||
|
|
"operator": "or",
|
||
|
|
# One-word queries must match; longer queries tolerate a
|
||
|
|
# missing term so "senior python django" still returns
|
||
|
|
# people who only list two of the three.
|
||
|
|
"minimum_should_match": "2<70%",
|
||
|
|
"fuzziness": "AUTO",
|
||
|
|
"prefix_length": 1,
|
||
|
|
"max_expansions": 30,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
)
|
||
|
|
should.append(
|
||
|
|
{
|
||
|
|
"multi_match": {
|
||
|
|
"query": request.q,
|
||
|
|
"type": "phrase",
|
||
|
|
"fields": PHRASE_FIELDS,
|
||
|
|
"slop": 1,
|
||
|
|
"boost": 4,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
must.append({"match_all": {}})
|
||
|
|
|
||
|
|
body: dict[str, Any] = {
|
||
|
|
"query": {
|
||
|
|
"bool": {
|
||
|
|
"must": must,
|
||
|
|
"should": should,
|
||
|
|
"filter": _build_filters(request),
|
||
|
|
}
|
||
|
|
},
|
||
|
|
"from": request.offset,
|
||
|
|
"size": request.page_size,
|
||
|
|
"track_total_hits": True,
|
||
|
|
"sort": _build_sort(request),
|
||
|
|
}
|
||
|
|
|
||
|
|
if request.q:
|
||
|
|
body["highlight"] = {
|
||
|
|
"pre_tags": ["<mark>"],
|
||
|
|
"post_tags": ["</mark>"],
|
||
|
|
"fields": HIGHLIGHT_FIELDS,
|
||
|
|
}
|
||
|
|
|
||
|
|
if request.include_facets:
|
||
|
|
body["aggs"] = _build_aggregations()
|
||
|
|
|
||
|
|
return body
|
||
|
|
|
||
|
|
|
||
|
|
def _build_filters(request: SearchRequest) -> list[dict[str, Any]]:
|
||
|
|
filters: list[dict[str, Any]] = []
|
||
|
|
|
||
|
|
for attribute, field_name in FILTER_FIELDS.items():
|
||
|
|
values = getattr(request, attribute)
|
||
|
|
if values:
|
||
|
|
# ``terms`` is OR within one filter; separate filters are ANDed,
|
||
|
|
# which matches how users read the UI: "python OR django" skills,
|
||
|
|
# AND located in Germany.
|
||
|
|
filters.append({"terms": {field_name: list(values)}})
|
||
|
|
|
||
|
|
if request.min_experience is not None or request.max_experience is not None:
|
||
|
|
bounds: dict[str, int] = {}
|
||
|
|
if request.min_experience is not None:
|
||
|
|
bounds["gte"] = request.min_experience
|
||
|
|
if request.max_experience is not None:
|
||
|
|
bounds["lte"] = request.max_experience
|
||
|
|
filters.append({"range": {"years_experience": bounds}})
|
||
|
|
|
||
|
|
return filters
|
||
|
|
|
||
|
|
|
||
|
|
def _build_sort(request: SearchRequest) -> list[Any]:
|
||
|
|
if request.sort == SORT_NAME:
|
||
|
|
return [{"full_name.keyword": "asc"}, {"id": "asc"}]
|
||
|
|
if request.sort == SORT_EXPERIENCE_DESC:
|
||
|
|
return [{"years_experience": {"order": "desc", "missing": "_last"}}, "_score"]
|
||
|
|
if request.sort == SORT_EXPERIENCE_ASC:
|
||
|
|
return [{"years_experience": {"order": "asc", "missing": "_last"}}, "_score"]
|
||
|
|
return ["_score", {"id": "asc"}] # deterministic tie-break
|
||
|
|
|
||
|
|
|
||
|
|
def _build_aggregations() -> dict[str, Any]:
|
||
|
|
size = settings.SEARCH["FACET_SIZE"]
|
||
|
|
return {
|
||
|
|
name: {"terms": {"field": field_name, "size": size, "order": {"_count": "desc"}}}
|
||
|
|
for name, field_name in FACET_FIELDS.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class ElasticsearchBackend(SearchBackend):
|
||
|
|
name = "elasticsearch"
|
||
|
|
|
||
|
|
def __init__(self, client: Elasticsearch | None = None) -> None:
|
||
|
|
self._client = client or get_client()
|
||
|
|
self._index = alias_name()
|
||
|
|
|
||
|
|
# -- SearchBackend ------------------------------------------------------
|
||
|
|
def search(self, request: SearchRequest) -> SearchResult:
|
||
|
|
body = build_query(request)
|
||
|
|
started = time.perf_counter()
|
||
|
|
try:
|
||
|
|
response = self._client.search(index=self._index, **body)
|
||
|
|
except (ApiError, TransportError) as exc:
|
||
|
|
logger.exception("search_backend_error", extra={"event": "search_backend_error"})
|
||
|
|
raise SearchBackendError() from exc
|
||
|
|
|
||
|
|
took_ms = int((time.perf_counter() - started) * 1000)
|
||
|
|
hits = [self._to_hit(raw) for raw in response["hits"]["hits"]]
|
||
|
|
return SearchResult(
|
||
|
|
total=response["hits"]["total"]["value"],
|
||
|
|
hits=hits,
|
||
|
|
facets=self._to_facets(response.get("aggregations", {})),
|
||
|
|
took_ms=took_ms,
|
||
|
|
backend=self.name,
|
||
|
|
)
|
||
|
|
|
||
|
|
def suggest(self, prefix: str, limit: int = 8) -> list[str]:
|
||
|
|
if not prefix:
|
||
|
|
return []
|
||
|
|
try:
|
||
|
|
response = self._client.search(
|
||
|
|
index=self._index,
|
||
|
|
size=limit,
|
||
|
|
query={
|
||
|
|
"bool": {
|
||
|
|
"should": [
|
||
|
|
{"match": {"full_name.autocomplete": {"query": prefix, "boost": 2}}},
|
||
|
|
{"match": {"skills.autocomplete": prefix}},
|
||
|
|
],
|
||
|
|
"minimum_should_match": 1,
|
||
|
|
}
|
||
|
|
},
|
||
|
|
_source=["full_name", "job_title"],
|
||
|
|
)
|
||
|
|
except (ApiError, TransportError) as exc:
|
||
|
|
raise SearchBackendError() from exc
|
||
|
|
|
||
|
|
seen: list[str] = []
|
||
|
|
for raw in response["hits"]["hits"]:
|
||
|
|
value = raw["_source"].get("full_name", "")
|
||
|
|
if value and value not in seen:
|
||
|
|
seen.append(value)
|
||
|
|
return seen[:limit]
|
||
|
|
|
||
|
|
def ping(self) -> bool:
|
||
|
|
if not self._client.ping():
|
||
|
|
raise SearchBackendError("Elasticsearch did not answer a ping.")
|
||
|
|
return True
|
||
|
|
|
||
|
|
# -- Mapping helpers ----------------------------------------------------
|
||
|
|
@staticmethod
|
||
|
|
def _to_hit(raw: dict[str, Any]) -> SearchHit:
|
||
|
|
source = raw["_source"]
|
||
|
|
return SearchHit(
|
||
|
|
id=source["id"],
|
||
|
|
score=raw.get("_score") or 0.0,
|
||
|
|
full_name=source.get("full_name", ""),
|
||
|
|
headline=source.get("headline", ""),
|
||
|
|
job_title=source.get("job_title", ""),
|
||
|
|
job_title_role=source.get("job_title_role", ""),
|
||
|
|
seniority_levels=source.get("seniority_levels", []),
|
||
|
|
company=source.get("company", ""),
|
||
|
|
industry=source.get("industry", ""),
|
||
|
|
location_name=source.get("location_name", ""),
|
||
|
|
location_country=source.get("location_country", ""),
|
||
|
|
years_experience=source.get("years_experience"),
|
||
|
|
linkedin_url=source.get("linkedin_url", ""),
|
||
|
|
skills=source.get("skills", [])[:12],
|
||
|
|
highlight=raw.get("highlight", {}),
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _to_facets(aggregations: dict[str, Any]) -> dict[str, list[FacetBucket]]:
|
||
|
|
return {
|
||
|
|
name: [
|
||
|
|
FacetBucket(key=str(bucket["key"]), count=bucket["doc_count"])
|
||
|
|
for bucket in payload.get("buckets", [])
|
||
|
|
if bucket["key"] != ""
|
||
|
|
]
|
||
|
|
for name, payload in aggregations.items()
|
||
|
|
}
|