src/nominatim_api/search/db_search_builder.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2024 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Conversion from token assignment to an abstract DB search.
   9 """
  10 from typing import Optional, List, Tuple, Iterator, Dict
  11 import heapq
  12
  13 from ..types import SearchDetails, DataLayer
  14 from .query import QueryStruct, Token, TokenType, TokenRange, BreakType
  15 from .token_assignment import TokenAssignment
  16 from . import db_search_fields as dbf
  17 from . import db_searches as dbs
  18 from . import db_search_lookups as lookups
  19
  20
  21 def wrap_near_search(categories: List[Tuple[str, str]],
  22                      search: dbs.AbstractSearch) -> dbs.NearSearch:
  23     """ Create a new search that wraps the given search in a search
  24         for near places of the given category.
  25     """
  26     return dbs.NearSearch(penalty=search.penalty,
  27                           categories=dbf.WeightedCategories(categories,
  28                                                             [0.0] * len(categories)),
  29                           search=search)
  30
  31
  32 def build_poi_search(category: List[Tuple[str, str]],
  33                      countries: Optional[List[str]]) -> dbs.PoiSearch:
  34     """ Create a new search for places by the given category, possibly
  35         constraint to the given countries.
  36     """
  37     if countries:
  38         ccs = dbf.WeightedStrings(countries, [0.0] * len(countries))
  39     else:
  40         ccs = dbf.WeightedStrings([], [])
  41
  42     class _PoiData(dbf.SearchData):
  43         penalty = 0.0
  44         qualifiers = dbf.WeightedCategories(category, [0.0] * len(category))
  45         countries=ccs
  46
  47     return dbs.PoiSearch(_PoiData())
  48
  49
  50 class SearchBuilder:
  51     """ Build the abstract search queries from token assignments.
  52     """
  53
  54     def __init__(self, query: QueryStruct, details: SearchDetails) -> None:
  55         self.query = query
  56         self.details = details
  57
  58
  59     @property
  60     def configured_for_country(self) -> bool:
  61         """ Return true if the search details are configured to
  62             allow countries in the result.
  63         """
  64         return self.details.min_rank <= 4 and self.details.max_rank >= 4 \
  65                and self.details.layer_enabled(DataLayer.ADDRESS)
  66
  67
  68     @property
  69     def configured_for_postcode(self) -> bool:
  70         """ Return true if the search details are configured to
  71             allow postcodes in the result.
  72         """
  73         return self.details.min_rank <= 5 and self.details.max_rank >= 11\
  74                and self.details.layer_enabled(DataLayer.ADDRESS)
  75
  76
  77     @property
  78     def configured_for_housenumbers(self) -> bool:
  79         """ Return true if the search details are configured to
  80             allow addresses in the result.
  81         """
  82         return self.details.max_rank >= 30 \
  83                and self.details.layer_enabled(DataLayer.ADDRESS)
  84
  85
  86     def build(self, assignment: TokenAssignment) -> Iterator[dbs.AbstractSearch]:
  87         """ Yield all possible abstract searches for the given token assignment.
  88         """
  89         sdata = self.get_search_data(assignment)
  90         if sdata is None:
  91             return
  92
  93         near_items = self.get_near_items(assignment)
  94         if near_items is not None and not near_items:
  95             return # impossible compbination of near items and category parameter
  96
  97         if assignment.name is None:
  98             if near_items and not sdata.postcodes:
  99                 sdata.qualifiers = near_items
 100                 near_items = None
 101                 builder = self.build_poi_search(sdata)
 102             elif assignment.housenumber:
 103                 hnr_tokens = self.query.get_tokens(assignment.housenumber,
 104                                                    TokenType.HOUSENUMBER)
 105                 builder = self.build_housenumber_search(sdata, hnr_tokens, assignment.address)
 106             else:
 107                 builder = self.build_special_search(sdata, assignment.address,
 108                                                     bool(near_items))
 109         else:
 110             builder = self.build_name_search(sdata, assignment.name, assignment.address,
 111                                              bool(near_items))
 112
 113         if near_items:
 114             penalty = min(near_items.penalties)
 115             near_items.penalties = [p - penalty for p in near_items.penalties]
 116             for search in builder:
 117                 search_penalty = search.penalty
 118                 search.penalty = 0.0
 119                 yield dbs.NearSearch(penalty + assignment.penalty + search_penalty,
 120                                      near_items, search)
 121         else:
 122             for search in builder:
 123                 search.penalty += assignment.penalty
 124                 yield search
 125
 126
 127     def build_poi_search(self, sdata: dbf.SearchData) -> Iterator[dbs.AbstractSearch]:
 128         """ Build abstract search query for a simple category search.
 129             This kind of search requires an additional geographic constraint.
 130         """
 131         if not sdata.housenumbers \
 132            and ((self.details.viewbox and self.details.bounded_viewbox) or self.details.near):
 133             yield dbs.PoiSearch(sdata)
 134
 135
 136     def build_special_search(self, sdata: dbf.SearchData,
 137                              address: List[TokenRange],
 138                              is_category: bool) -> Iterator[dbs.AbstractSearch]:
 139         """ Build abstract search queries for searches that do not involve
 140             a named place.
 141         """
 142         if sdata.qualifiers:
 143             # No special searches over qualifiers supported.
 144             return
 145
 146         if sdata.countries and not address and not sdata.postcodes \
 147            and self.configured_for_country:
 148             yield dbs.CountrySearch(sdata)
 149
 150         if sdata.postcodes and (is_category or self.configured_for_postcode):
 151             penalty = 0.0 if sdata.countries else 0.1
 152             if address:
 153                 sdata.lookups = [dbf.FieldLookup('nameaddress_vector',
 154                                                  [t.token for r in address
 155                                                   for t in self.query.get_partials_list(r)],
 156                                                  lookups.Restrict)]
 157                 penalty += 0.2
 158             yield dbs.PostcodeSearch(penalty, sdata)
 159
 160
 161     def build_housenumber_search(self, sdata: dbf.SearchData, hnrs: List[Token],
 162                                  address: List[TokenRange]) -> Iterator[dbs.AbstractSearch]:
 163         """ Build a simple address search for special entries where the
 164             housenumber is the main name token.
 165         """
 166         sdata.lookups = [dbf.FieldLookup('name_vector', [t.token for t in hnrs], lookups.LookupAny)]
 167         expected_count = sum(t.count for t in hnrs)
 168
 169         partials = {t.token: t.addr_count for trange in address
 170                        for t in self.query.get_partials_list(trange)}
 171
 172         if expected_count < 8000:
 173             sdata.lookups.append(dbf.FieldLookup('nameaddress_vector',
 174                                                  list(partials), lookups.Restrict))
 175         elif len(partials) != 1 or list(partials.values())[0] < 10000:
 176             sdata.lookups.append(dbf.FieldLookup('nameaddress_vector',
 177                                                  list(partials), lookups.LookupAll))
 178         else:
 179             addr_fulls = [t.token for t
 180                           in self.query.get_tokens(address[0], TokenType.WORD)]
 181             if len(addr_fulls) > 5:
 182                 return
 183             sdata.lookups.append(
 184                 dbf.FieldLookup('nameaddress_vector', addr_fulls, lookups.LookupAny))
 185
 186         sdata.housenumbers = dbf.WeightedStrings([], [])
 187         yield dbs.PlaceSearch(0.05, sdata, expected_count)
 188
 189
 190     def build_name_search(self, sdata: dbf.SearchData,
 191                           name: TokenRange, address: List[TokenRange],
 192                           is_category: bool) -> Iterator[dbs.AbstractSearch]:
 193         """ Build abstract search queries for simple name or address searches.
 194         """
 195         if is_category or not sdata.housenumbers or self.configured_for_housenumbers:
 196             ranking = self.get_name_ranking(name)
 197             name_penalty = ranking.normalize_penalty()
 198             if ranking.rankings:
 199                 sdata.rankings.append(ranking)
 200             for penalty, count, lookup in self.yield_lookups(name, address):
 201                 sdata.lookups = lookup
 202                 yield dbs.PlaceSearch(penalty + name_penalty, sdata, count)
 203
 204
 205     def yield_lookups(self, name: TokenRange, address: List[TokenRange])\
 206                           -> Iterator[Tuple[float, int, List[dbf.FieldLookup]]]:
 207         """ Yield all variants how the given name and address should best
 208             be searched for. This takes into account how frequent the terms
 209             are and tries to find a lookup that optimizes index use.
 210         """
 211         penalty = 0.0 # extra penalty
 212         name_partials = {t.token: t for t in self.query.get_partials_list(name)}
 213
 214         addr_partials = [t for r in address for t in self.query.get_partials_list(r)]
 215         addr_tokens = list({t.token for t in addr_partials})
 216
 217         partials_indexed = all(t.is_indexed for t in name_partials.values()) \
 218                            and all(t.is_indexed for t in addr_partials)
 219         exp_count = min(t.count for t in name_partials.values()) / (2**(len(name_partials) - 1))
 220
 221         if (len(name_partials) > 3 or exp_count < 8000) and partials_indexed:
 222             yield penalty, exp_count, dbf.lookup_by_names(list(name_partials.keys()), addr_tokens)
 223             return
 224
 225         addr_count = min(t.addr_count for t in addr_partials) if addr_partials else 30000
 226         # Partial term to frequent. Try looking up by rare full names first.
 227         name_fulls = self.query.get_tokens(name, TokenType.WORD)
 228         if name_fulls:
 229             fulls_count = sum(t.count for t in name_fulls)
 230             if partials_indexed:
 231                 penalty += 1.2 * sum(t.penalty for t in addr_partials if not t.is_indexed)
 232
 233             if fulls_count < 50000 or addr_count < 30000:
 234                 yield penalty,fulls_count / (2**len(addr_tokens)), \
 235                     self.get_full_name_ranking(name_fulls, addr_partials,
 236                                                fulls_count > 30000 / max(1, len(addr_tokens)))
 237
 238         # To catch remaining results, lookup by name and address
 239         # We only do this if there is a reasonable number of results expected.
 240         exp_count = exp_count / (2**len(addr_tokens)) if addr_tokens else exp_count
 241         if exp_count < 10000 and addr_count < 20000\
 242            and all(t.is_indexed for t in name_partials.values()):
 243             penalty += 0.35 * max(1 if name_fulls else 0.1,
 244                                   5 - len(name_partials) - len(addr_tokens))
 245             yield penalty, exp_count,\
 246                   self.get_name_address_ranking(list(name_partials.keys()), addr_partials)
 247
 248
 249     def get_name_address_ranking(self, name_tokens: List[int],
 250                                  addr_partials: List[Token]) -> List[dbf.FieldLookup]:
 251         """ Create a ranking expression looking up by name and address.
 252         """
 253         lookup = [dbf.FieldLookup('name_vector', name_tokens, lookups.LookupAll)]
 254
 255         addr_restrict_tokens = []
 256         addr_lookup_tokens = []
 257         for t in addr_partials:
 258             if t.is_indexed:
 259                 if t.addr_count > 20000:
 260                     addr_restrict_tokens.append(t.token)
 261                 else:
 262                     addr_lookup_tokens.append(t.token)
 263
 264         if addr_restrict_tokens:
 265             lookup.append(dbf.FieldLookup('nameaddress_vector',
 266                                           addr_restrict_tokens, lookups.Restrict))
 267         if addr_lookup_tokens:
 268             lookup.append(dbf.FieldLookup('nameaddress_vector',
 269                                           addr_lookup_tokens, lookups.LookupAll))
 270
 271         return lookup
 272
 273
 274     def get_full_name_ranking(self, name_fulls: List[Token], addr_partials: List[Token],
 275                               use_lookup: bool) -> List[dbf.FieldLookup]:
 276         """ Create a ranking expression with full name terms and
 277             additional address lookup. When 'use_lookup' is true, then
 278             address lookups will use the index, when the occurrences are not
 279             too many.
 280         """
 281         # At this point drop unindexed partials from the address.
 282         # This might yield wrong results, nothing we can do about that.
 283         if use_lookup:
 284             addr_restrict_tokens = []
 285             addr_lookup_tokens = []
 286             for t in addr_partials:
 287                 if t.is_indexed:
 288                     if t.addr_count > 20000:
 289                         addr_restrict_tokens.append(t.token)
 290                     else:
 291                         addr_lookup_tokens.append(t.token)
 292         else:
 293             addr_restrict_tokens = [t.token for t in addr_partials if t.is_indexed]
 294             addr_lookup_tokens = []
 295
 296         return dbf.lookup_by_any_name([t.token for t in name_fulls],
 297                                       addr_restrict_tokens, addr_lookup_tokens)
 298
 299
 300     def get_name_ranking(self, trange: TokenRange,
 301                          db_field: str = 'name_vector') -> dbf.FieldRanking:
 302         """ Create a ranking expression for a name term in the given range.
 303         """
 304         name_fulls = self.query.get_tokens(trange, TokenType.WORD)
 305         ranks = [dbf.RankedTokens(t.penalty, [t.token]) for t in name_fulls]
 306         ranks.sort(key=lambda r: r.penalty)
 307         # Fallback, sum of penalty for partials
 308         name_partials = self.query.get_partials_list(trange)
 309         default = sum(t.penalty for t in name_partials) + 0.2
 310         return dbf.FieldRanking(db_field, default, ranks)
 311
 312
 313     def get_addr_ranking(self, trange: TokenRange) -> dbf.FieldRanking:
 314         """ Create a list of ranking expressions for an address term
 315             for the given ranges.
 316         """
 317         todo: List[Tuple[int, int, dbf.RankedTokens]] = []
 318         heapq.heappush(todo, (0, trange.start, dbf.RankedTokens(0.0, [])))
 319         ranks: List[dbf.RankedTokens] = []
 320
 321         while todo: # pylint: disable=too-many-nested-blocks
 322             neglen, pos, rank = heapq.heappop(todo)
 323             for tlist in self.query.nodes[pos].starting:
 324                 if tlist.ttype in (TokenType.PARTIAL, TokenType.WORD):
 325                     if tlist.end < trange.end:
 326                         chgpenalty = PENALTY_WORDCHANGE[self.query.nodes[tlist.end].btype]
 327                         if tlist.ttype == TokenType.PARTIAL:
 328                             penalty = rank.penalty + chgpenalty \
 329                                       + max(t.penalty for t in tlist.tokens)
 330                             heapq.heappush(todo, (neglen - 1, tlist.end,
 331                                                   dbf.RankedTokens(penalty, rank.tokens)))
 332                         else:
 333                             for t in tlist.tokens:
 334                                 heapq.heappush(todo, (neglen - 1, tlist.end,
 335                                                       rank.with_token(t, chgpenalty)))
 336                     elif tlist.end == trange.end:
 337                         if tlist.ttype == TokenType.PARTIAL:
 338                             ranks.append(dbf.RankedTokens(rank.penalty
 339                                                           + max(t.penalty for t in tlist.tokens),
 340                                                           rank.tokens))
 341                         else:
 342                             ranks.extend(rank.with_token(t, 0.0) for t in tlist.tokens)
 343                         if len(ranks) >= 10:
 344                             # Too many variants, bail out and only add
 345                             # Worst-case Fallback: sum of penalty of partials
 346                             name_partials = self.query.get_partials_list(trange)
 347                             default = sum(t.penalty for t in name_partials) + 0.2
 348                             ranks.append(dbf.RankedTokens(rank.penalty + default, []))
 349                             # Bail out of outer loop
 350                             todo.clear()
 351                             break
 352
 353         ranks.sort(key=lambda r: len(r.tokens))
 354         default = ranks[0].penalty + 0.3
 355         del ranks[0]
 356         ranks.sort(key=lambda r: r.penalty)
 357
 358         return dbf.FieldRanking('nameaddress_vector', default, ranks)
 359
 360
 361     def get_search_data(self, assignment: TokenAssignment) -> Optional[dbf.SearchData]:
 362         """ Collect the tokens for the non-name search fields in the
 363             assignment.
 364         """
 365         sdata = dbf.SearchData()
 366         sdata.penalty = assignment.penalty
 367         if assignment.country:
 368             tokens = self.get_country_tokens(assignment.country)
 369             if not tokens:
 370                 return None
 371             sdata.set_strings('countries', tokens)
 372         elif self.details.countries:
 373             sdata.countries = dbf.WeightedStrings(self.details.countries,
 374                                                   [0.0] * len(self.details.countries))
 375         if assignment.housenumber:
 376             sdata.set_strings('housenumbers',
 377                               self.query.get_tokens(assignment.housenumber,
 378                                                     TokenType.HOUSENUMBER))
 379         if assignment.postcode:
 380             sdata.set_strings('postcodes',
 381                               self.query.get_tokens(assignment.postcode,
 382                                                     TokenType.POSTCODE))
 383         if assignment.qualifier:
 384             tokens = self.get_qualifier_tokens(assignment.qualifier)
 385             if not tokens:
 386                 return None
 387             sdata.set_qualifiers(tokens)
 388         elif self.details.categories:
 389             sdata.qualifiers = dbf.WeightedCategories(self.details.categories,
 390                                                       [0.0] * len(self.details.categories))
 391
 392         if assignment.address:
 393             if not assignment.name and assignment.housenumber:
 394                 # housenumber search: the first item needs to be handled like
 395                 # a name in ranking or penalties are not comparable with
 396                 # normal searches.
 397                 sdata.set_ranking([self.get_name_ranking(assignment.address[0],
 398                                                          db_field='nameaddress_vector')]
 399                                   + [self.get_addr_ranking(r) for r in assignment.address[1:]])
 400             else:
 401                 sdata.set_ranking([self.get_addr_ranking(r) for r in assignment.address])
 402         else:
 403             sdata.rankings = []
 404
 405         return sdata
 406
 407
 408     def get_country_tokens(self, trange: TokenRange) -> List[Token]:
 409         """ Return the list of country tokens for the given range,
 410             optionally filtered by the country list from the details
 411             parameters.
 412         """
 413         tokens = self.query.get_tokens(trange, TokenType.COUNTRY)
 414         if self.details.countries:
 415             tokens = [t for t in tokens if t.lookup_word in self.details.countries]
 416
 417         return tokens
 418
 419
 420     def get_qualifier_tokens(self, trange: TokenRange) -> List[Token]:
 421         """ Return the list of qualifier tokens for the given range,
 422             optionally filtered by the qualifier list from the details
 423             parameters.
 424         """
 425         tokens = self.query.get_tokens(trange, TokenType.QUALIFIER)
 426         if self.details.categories:
 427             tokens = [t for t in tokens if t.get_category() in self.details.categories]
 428
 429         return tokens
 430
 431
 432     def get_near_items(self, assignment: TokenAssignment) -> Optional[dbf.WeightedCategories]:
 433         """ Collect tokens for near items search or use the categories
 434             requested per parameter.
 435             Returns None if no category search is requested.
 436         """
 437         if assignment.near_item:
 438             tokens: Dict[Tuple[str, str], float] = {}
 439             for t in self.query.get_tokens(assignment.near_item, TokenType.NEAR_ITEM):
 440                 cat = t.get_category()
 441                 # The category of a near search will be that of near_item.
 442                 # Thus, if search is restricted to a category parameter,
 443                 # the two sets must intersect.
 444                 if (not self.details.categories or cat in self.details.categories)\
 445                    and t.penalty < tokens.get(cat, 1000.0):
 446                     tokens[cat] = t.penalty
 447             return dbf.WeightedCategories(list(tokens.keys()), list(tokens.values()))
 448
 449         return None
 450
 451
 452 PENALTY_WORDCHANGE = {
 453     BreakType.START: 0.0,
 454     BreakType.END: 0.0,
 455     BreakType.PHRASE: 0.0,
 456     BreakType.WORD: 0.1,
 457     BreakType.PART: 0.2,
 458     BreakType.TOKEN: 0.4
 459 }