]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/api/search/token_assignment.py
Merge remote-tracking branch 'upstream/master'
[nominatim.git] / nominatim / api / search / token_assignment.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) 2023 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 """
8 Create query interpretations where each vertice in the query is assigned
9 a specific function (expressed as a token type).
10 """
11 from typing import Optional, List, Iterator
12 import dataclasses
13
14 import nominatim.api.search.query as qmod
15 from nominatim.api.logging import log
16
17 # pylint: disable=too-many-return-statements,too-many-branches
18
19 @dataclasses.dataclass
20 class TypedRange:
21     """ A token range for a specific type of tokens.
22     """
23     ttype: qmod.TokenType
24     trange: qmod.TokenRange
25
26
27 PENALTY_TOKENCHANGE = {
28     qmod.BreakType.START: 0.0,
29     qmod.BreakType.END: 0.0,
30     qmod.BreakType.PHRASE: 0.0,
31     qmod.BreakType.WORD: 0.1,
32     qmod.BreakType.PART: 0.2,
33     qmod.BreakType.TOKEN: 0.4
34 }
35
36 TypedRangeSeq = List[TypedRange]
37
38 @dataclasses.dataclass
39 class TokenAssignment: # pylint: disable=too-many-instance-attributes
40     """ Representation of a possible assignment of token types
41         to the tokens in a tokenized query.
42     """
43     penalty: float = 0.0
44     name: Optional[qmod.TokenRange] = None
45     address: List[qmod.TokenRange] = dataclasses.field(default_factory=list)
46     housenumber: Optional[qmod.TokenRange] = None
47     postcode: Optional[qmod.TokenRange] = None
48     country: Optional[qmod.TokenRange] = None
49     near_item: Optional[qmod.TokenRange] = None
50     qualifier: Optional[qmod.TokenRange] = None
51
52
53     @staticmethod
54     def from_ranges(ranges: TypedRangeSeq) -> 'TokenAssignment':
55         """ Create a new token assignment from a sequence of typed spans.
56         """
57         out = TokenAssignment()
58         for token in ranges:
59             if token.ttype == qmod.TokenType.PARTIAL:
60                 out.address.append(token.trange)
61             elif token.ttype == qmod.TokenType.HOUSENUMBER:
62                 out.housenumber = token.trange
63             elif token.ttype == qmod.TokenType.POSTCODE:
64                 out.postcode = token.trange
65             elif token.ttype == qmod.TokenType.COUNTRY:
66                 out.country = token.trange
67             elif token.ttype == qmod.TokenType.NEAR_ITEM:
68                 out.near_item = token.trange
69             elif token.ttype == qmod.TokenType.QUALIFIER:
70                 out.qualifier = token.trange
71         return out
72
73
74 class _TokenSequence:
75     """ Working state used to put together the token assignments.
76
77         Represents an intermediate state while traversing the tokenized
78         query.
79     """
80     def __init__(self, seq: TypedRangeSeq,
81                  direction: int = 0, penalty: float = 0.0) -> None:
82         self.seq = seq
83         self.direction = direction
84         self.penalty = penalty
85
86
87     def __str__(self) -> str:
88         seq = ''.join(f'[{r.trange.start} - {r.trange.end}: {r.ttype.name}]' for r in self.seq)
89         return f'{seq} (dir: {self.direction}, penalty: {self.penalty})'
90
91
92     @property
93     def end_pos(self) -> int:
94         """ Return the index of the global end of the current sequence.
95         """
96         return self.seq[-1].trange.end if self.seq else 0
97
98
99     def has_types(self, *ttypes: qmod.TokenType) -> bool:
100         """ Check if the current sequence contains any typed ranges of
101             the given types.
102         """
103         return any(s.ttype in ttypes for s in self.seq)
104
105
106     def is_final(self) -> bool:
107         """ Return true when the sequence cannot be extended by any
108             form of token anymore.
109         """
110         # Country and category must be the final term for left-to-right
111         return len(self.seq) > 1 and \
112                self.seq[-1].ttype in (qmod.TokenType.COUNTRY, qmod.TokenType.NEAR_ITEM)
113
114
115     def appendable(self, ttype: qmod.TokenType) -> Optional[int]:
116         """ Check if the give token type is appendable to the existing sequence.
117
118             Returns None if the token type is not appendable, otherwise the
119             new direction of the sequence after adding such a type. The
120             token is not added.
121         """
122         if ttype == qmod.TokenType.WORD:
123             return None
124
125         if not self.seq:
126             # Append unconditionally to the empty list
127             if ttype == qmod.TokenType.COUNTRY:
128                 return -1
129             if ttype in (qmod.TokenType.HOUSENUMBER, qmod.TokenType.QUALIFIER):
130                 return 1
131             return self.direction
132
133         # Name tokens are always acceptable and don't change direction
134         if ttype == qmod.TokenType.PARTIAL:
135             # qualifiers cannot appear in the middle of the query. They need
136             # to be near the next phrase.
137             if self.direction == -1 \
138                and any(t.ttype == qmod.TokenType.QUALIFIER for t in self.seq[:-1]):
139                 return None
140             return self.direction
141
142         # Other tokens may only appear once
143         if self.has_types(ttype):
144             return None
145
146         if ttype == qmod.TokenType.HOUSENUMBER:
147             if self.direction == 1:
148                 if len(self.seq) == 1 and self.seq[0].ttype == qmod.TokenType.QUALIFIER:
149                     return None
150                 if len(self.seq) > 2 \
151                    or self.has_types(qmod.TokenType.POSTCODE, qmod.TokenType.COUNTRY):
152                     return None # direction left-to-right: housenumber must come before anything
153             elif self.direction == -1 \
154                  or self.has_types(qmod.TokenType.POSTCODE, qmod.TokenType.COUNTRY):
155                 return -1 # force direction right-to-left if after other terms
156
157             return self.direction
158
159         if ttype == qmod.TokenType.POSTCODE:
160             if self.direction == -1:
161                 if self.has_types(qmod.TokenType.HOUSENUMBER, qmod.TokenType.QUALIFIER):
162                     return None
163                 return -1
164             if self.direction == 1:
165                 return None if self.has_types(qmod.TokenType.COUNTRY) else 1
166             if self.has_types(qmod.TokenType.HOUSENUMBER, qmod.TokenType.QUALIFIER):
167                 return 1
168             return self.direction
169
170         if ttype == qmod.TokenType.COUNTRY:
171             return None if self.direction == -1 else 1
172
173         if ttype == qmod.TokenType.NEAR_ITEM:
174             return self.direction
175
176         if ttype == qmod.TokenType.QUALIFIER:
177             if self.direction == 1:
178                 if (len(self.seq) == 1
179                     and self.seq[0].ttype in (qmod.TokenType.PARTIAL, qmod.TokenType.NEAR_ITEM)) \
180                    or (len(self.seq) == 2
181                        and self.seq[0].ttype == qmod.TokenType.NEAR_ITEM
182                        and self.seq[1].ttype == qmod.TokenType.PARTIAL):
183                     return 1
184                 return None
185             if self.direction == -1:
186                 return -1
187
188             tempseq = self.seq[1:] if self.seq[0].ttype == qmod.TokenType.NEAR_ITEM else self.seq
189             if len(tempseq) == 0:
190                 return 1
191             if len(tempseq) == 1 and self.seq[0].ttype == qmod.TokenType.HOUSENUMBER:
192                 return None
193             if len(tempseq) > 1 or self.has_types(qmod.TokenType.POSTCODE, qmod.TokenType.COUNTRY):
194                 return -1
195             return 0
196
197         return None
198
199
200     def advance(self, ttype: qmod.TokenType, end_pos: int,
201                 btype: qmod.BreakType) -> Optional['_TokenSequence']:
202         """ Return a new token sequence state with the given token type
203             extended.
204         """
205         newdir = self.appendable(ttype)
206         if newdir is None:
207             return None
208
209         if not self.seq:
210             newseq = [TypedRange(ttype, qmod.TokenRange(0, end_pos))]
211             new_penalty = 0.0
212         else:
213             last = self.seq[-1]
214             if btype != qmod.BreakType.PHRASE and last.ttype == ttype:
215                 # extend the existing range
216                 newseq = self.seq[:-1] + [TypedRange(ttype, last.trange.replace_end(end_pos))]
217                 new_penalty = 0.0
218             else:
219                 # start a new range
220                 newseq = list(self.seq) + [TypedRange(ttype,
221                                                       qmod.TokenRange(last.trange.end, end_pos))]
222                 new_penalty = PENALTY_TOKENCHANGE[btype]
223
224         return _TokenSequence(newseq, newdir, self.penalty + new_penalty)
225
226
227     def _adapt_penalty_from_priors(self, priors: int, new_dir: int) -> bool:
228         if priors == 2:
229             self.penalty += 1.0
230         elif priors > 2:
231             if self.direction == 0:
232                 self.direction = new_dir
233             else:
234                 return False
235
236         return True
237
238
239     def recheck_sequence(self) -> bool:
240         """ Check that the sequence is a fully valid token assignment
241             and adapt direction and penalties further if necessary.
242
243             This function catches some impossible assignments that need
244             forward context and can therefore not be excluded when building
245             the assignment.
246         """
247         # housenumbers may not be further than 2 words from the beginning.
248         # If there are two words in front, give it a penalty.
249         hnrpos = next((i for i, tr in enumerate(self.seq)
250                        if tr.ttype == qmod.TokenType.HOUSENUMBER),
251                       None)
252         if hnrpos is not None:
253             if self.direction != -1:
254                 priors = sum(1 for t in self.seq[:hnrpos] if t.ttype == qmod.TokenType.PARTIAL)
255                 if not self._adapt_penalty_from_priors(priors, -1):
256                     return False
257             if self.direction != 1:
258                 priors = sum(1 for t in self.seq[hnrpos+1:] if t.ttype == qmod.TokenType.PARTIAL)
259                 if not self._adapt_penalty_from_priors(priors, 1):
260                     return False
261             if any(t.ttype == qmod.TokenType.NEAR_ITEM for t in self.seq):
262                 self.penalty += 1.0
263
264         return True
265
266
267     def _get_assignments_postcode(self, base: TokenAssignment,
268                                   query_len: int)  -> Iterator[TokenAssignment]:
269         """ Yield possible assignments of Postcode searches with an
270             address component.
271         """
272         assert base.postcode is not None
273
274         if (base.postcode.start == 0 and self.direction != -1)\
275            or (base.postcode.end == query_len and self.direction != 1):
276             log().comment('postcode search')
277             # <address>,<postcode> should give preference to address search
278             if base.postcode.start == 0:
279                 penalty = self.penalty
280                 self.direction = -1 # name searches are only possible backwards
281             else:
282                 penalty = self.penalty + 0.1
283                 self.direction = 1 # name searches are only possible forwards
284             yield dataclasses.replace(base, penalty=penalty)
285
286
287     def _get_assignments_address_forward(self, base: TokenAssignment,
288                                          query: qmod.QueryStruct) -> Iterator[TokenAssignment]:
289         """ Yield possible assignments of address searches with
290             left-to-right reading.
291         """
292         first = base.address[0]
293
294         log().comment('first word = name')
295         yield dataclasses.replace(base, penalty=self.penalty,
296                                   name=first, address=base.address[1:])
297
298         # To paraphrase:
299         #  * if another name term comes after the first one and before the
300         #    housenumber
301         #  * a qualifier comes after the name
302         #  * the containing phrase is strictly typed
303         if (base.housenumber and first.end < base.housenumber.start)\
304            or (base.qualifier and base.qualifier > first)\
305            or (query.nodes[first.start].ptype != qmod.PhraseType.NONE):
306             return
307
308         penalty = self.penalty
309
310         # Penalty for:
311         #  * <name>, <street>, <housenumber> , ...
312         #  * queries that are comma-separated
313         if (base.housenumber and base.housenumber > first) or len(query.source) > 1:
314             penalty += 0.25
315
316         for i in range(first.start + 1, first.end):
317             name, addr = first.split(i)
318             log().comment(f'split first word = name ({i - first.start})')
319             yield dataclasses.replace(base, name=name, address=[addr] + base.address[1:],
320                                       penalty=penalty + PENALTY_TOKENCHANGE[query.nodes[i].btype])
321
322
323     def _get_assignments_address_backward(self, base: TokenAssignment,
324                                           query: qmod.QueryStruct) -> Iterator[TokenAssignment]:
325         """ Yield possible assignments of address searches with
326             right-to-left reading.
327         """
328         last = base.address[-1]
329
330         if self.direction == -1 or len(base.address) > 1:
331             log().comment('last word = name')
332             yield dataclasses.replace(base, penalty=self.penalty,
333                                       name=last, address=base.address[:-1])
334
335         # To paraphrase:
336         #  * if another name term comes before the last one and after the
337         #    housenumber
338         #  * a qualifier comes before the name
339         #  * the containing phrase is strictly typed
340         if (base.housenumber and last.start > base.housenumber.end)\
341            or (base.qualifier and base.qualifier < last)\
342            or (query.nodes[last.start].ptype != qmod.PhraseType.NONE):
343             return
344
345         penalty = self.penalty
346         if base.housenumber and base.housenumber < last:
347             penalty += 0.4
348         if len(query.source) > 1:
349             penalty += 0.25
350
351         for i in range(last.start + 1, last.end):
352             addr, name = last.split(i)
353             log().comment(f'split last word = name ({i - last.start})')
354             yield dataclasses.replace(base, name=name, address=base.address[:-1] + [addr],
355                                       penalty=penalty + PENALTY_TOKENCHANGE[query.nodes[i].btype])
356
357
358     def get_assignments(self, query: qmod.QueryStruct) -> Iterator[TokenAssignment]:
359         """ Yield possible assignments for the current sequence.
360
361             This function splits up general name assignments into name
362             and address and yields all possible variants of that.
363         """
364         base = TokenAssignment.from_ranges(self.seq)
365
366         num_addr_tokens = sum(t.end - t.start for t in base.address)
367         if num_addr_tokens > 50:
368             return
369
370         # Postcode search (postcode-only search is covered in next case)
371         if base.postcode is not None and base.address:
372             yield from self._get_assignments_postcode(base, query.num_token_slots())
373
374         # Postcode or country-only search
375         if not base.address:
376             if not base.housenumber and (base.postcode or base.country or base.near_item):
377                 log().comment('postcode/country search')
378                 yield dataclasses.replace(base, penalty=self.penalty)
379         else:
380             # <postcode>,<address> should give preference to postcode search
381             if base.postcode and base.postcode.start == 0:
382                 self.penalty += 0.1
383
384             # Right-to-left reading of the address
385             if self.direction != -1:
386                 yield from self._get_assignments_address_forward(base, query)
387
388             # Left-to-right reading of the address
389             if self.direction != 1:
390                 yield from self._get_assignments_address_backward(base, query)
391
392             # variant for special housenumber searches
393             if base.housenumber and not base.qualifier:
394                 yield dataclasses.replace(base, penalty=self.penalty)
395
396
397 def yield_token_assignments(query: qmod.QueryStruct) -> Iterator[TokenAssignment]:
398     """ Return possible word type assignments to word positions.
399
400         The assignments are computed from the concrete tokens listed
401         in the tokenized query.
402
403         The result includes the penalty for transitions from one word type to
404         another. It does not include penalties for transitions within a
405         type.
406     """
407     todo = [_TokenSequence([], direction=0 if query.source[0].ptype == qmod.PhraseType.NONE else 1)]
408
409     while todo:
410         state = todo.pop()
411         node = query.nodes[state.end_pos]
412
413         for tlist in node.starting:
414             newstate = state.advance(tlist.ttype, tlist.end, node.btype)
415             if newstate is not None:
416                 if newstate.end_pos == query.num_token_slots():
417                     if newstate.recheck_sequence():
418                         log().var_dump('Assignment', newstate)
419                         yield from newstate.get_assignments(query)
420                 elif not newstate.is_final():
421                     todo.append(newstate)