]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/api/search/icu_tokenizer.py
enable all API tests for sqlite and port missing features
[nominatim.git] / nominatim / api / search / icu_tokenizer.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 Implementation of query analysis for the ICU tokenizer.
9 """
10 from typing import Tuple, Dict, List, Optional, NamedTuple, Iterator, Any, cast
11 from copy import copy
12 from collections import defaultdict
13 import dataclasses
14 import difflib
15
16 from icu import Transliterator
17
18 import sqlalchemy as sa
19
20 from nominatim.typing import SaRow
21 from nominatim.api.connection import SearchConnection
22 from nominatim.api.logging import log
23 from nominatim.api.search import query as qmod
24 from nominatim.api.search.query_analyzer_factory import AbstractQueryAnalyzer
25 from nominatim.db.sqlalchemy_types import Json
26
27
28 DB_TO_TOKEN_TYPE = {
29     'W': qmod.TokenType.WORD,
30     'w': qmod.TokenType.PARTIAL,
31     'H': qmod.TokenType.HOUSENUMBER,
32     'P': qmod.TokenType.POSTCODE,
33     'C': qmod.TokenType.COUNTRY
34 }
35
36
37 class QueryPart(NamedTuple):
38     """ Normalized and transliterated form of a single term in the query.
39         When the term came out of a split during the transliteration,
40         the normalized string is the full word before transliteration.
41         The word number keeps track of the word before transliteration
42         and can be used to identify partial transliterated terms.
43     """
44     token: str
45     normalized: str
46     word_number: int
47
48
49 QueryParts = List[QueryPart]
50 WordDict = Dict[str, List[qmod.TokenRange]]
51
52 def yield_words(terms: List[QueryPart], start: int) -> Iterator[Tuple[str, qmod.TokenRange]]:
53     """ Return all combinations of words in the terms list after the
54         given position.
55     """
56     total = len(terms)
57     for first in range(start, total):
58         word = terms[first].token
59         yield word, qmod.TokenRange(first, first + 1)
60         for last in range(first + 1, min(first + 20, total)):
61             word = ' '.join((word, terms[last].token))
62             yield word, qmod.TokenRange(first, last + 1)
63
64
65 @dataclasses.dataclass
66 class ICUToken(qmod.Token):
67     """ Specialised token for ICU tokenizer.
68     """
69     word_token: str
70     info: Optional[Dict[str, Any]]
71
72     def get_category(self) -> Tuple[str, str]:
73         assert self.info
74         return self.info.get('class', ''), self.info.get('type', '')
75
76
77     def rematch(self, norm: str) -> None:
78         """ Check how well the token matches the given normalized string
79             and add a penalty, if necessary.
80         """
81         if not self.lookup_word:
82             return
83
84         seq = difflib.SequenceMatcher(a=self.lookup_word, b=norm)
85         distance = 0
86         for tag, afrom, ato, bfrom, bto in seq.get_opcodes():
87             if tag in ('delete', 'insert') and (afrom == 0 or ato == len(self.lookup_word)):
88                 distance += 1
89             elif tag == 'replace':
90                 distance += max((ato-afrom), (bto-bfrom))
91             elif tag != 'equal':
92                 distance += abs((ato-afrom) - (bto-bfrom))
93         self.penalty += (distance/len(self.lookup_word))
94
95
96     @staticmethod
97     def from_db_row(row: SaRow) -> 'ICUToken':
98         """ Create a ICUToken from the row of the word table.
99         """
100         count = 1 if row.info is None else row.info.get('count', 1)
101
102         penalty = 0.0
103         if row.type == 'w':
104             penalty = 0.3
105         elif row.type == 'W':
106             if len(row.word_token) == 1 and row.word_token == row.word:
107                 penalty = 0.2 if row.word.isdigit() else 0.3
108         elif row.type == 'H':
109             penalty = sum(0.1 for c in row.word_token if c != ' ' and not c.isdigit())
110             if all(not c.isdigit() for c in row.word_token):
111                 penalty += 0.2 * (len(row.word_token) - 1)
112         elif row.type == 'C':
113             if len(row.word_token) == 1:
114                 penalty = 0.3
115
116         if row.info is None:
117             lookup_word = row.word
118         else:
119             lookup_word = row.info.get('lookup', row.word)
120         if lookup_word:
121             lookup_word = lookup_word.split('@', 1)[0]
122         else:
123             lookup_word = row.word_token
124
125         return ICUToken(penalty=penalty, token=row.word_id, count=count,
126                         lookup_word=lookup_word, is_indexed=True,
127                         word_token=row.word_token, info=row.info)
128
129
130
131 class ICUQueryAnalyzer(AbstractQueryAnalyzer):
132     """ Converter for query strings into a tokenized query
133         using the tokens created by a ICU tokenizer.
134     """
135
136     def __init__(self, conn: SearchConnection) -> None:
137         self.conn = conn
138
139
140     async def setup(self) -> None:
141         """ Set up static data structures needed for the analysis.
142         """
143         async def _make_normalizer() -> Any:
144             rules = await self.conn.get_property('tokenizer_import_normalisation')
145             return Transliterator.createFromRules("normalization", rules)
146
147         self.normalizer = await self.conn.get_cached_value('ICUTOK', 'normalizer',
148                                                            _make_normalizer)
149
150         async def _make_transliterator() -> Any:
151             rules = await self.conn.get_property('tokenizer_import_transliteration')
152             return Transliterator.createFromRules("transliteration", rules)
153
154         self.transliterator = await self.conn.get_cached_value('ICUTOK', 'transliterator',
155                                                                _make_transliterator)
156
157         if 'word' not in self.conn.t.meta.tables:
158             sa.Table('word', self.conn.t.meta,
159                      sa.Column('word_id', sa.Integer),
160                      sa.Column('word_token', sa.Text, nullable=False),
161                      sa.Column('type', sa.Text, nullable=False),
162                      sa.Column('word', sa.Text),
163                      sa.Column('info', Json))
164
165
166     async def analyze_query(self, phrases: List[qmod.Phrase]) -> qmod.QueryStruct:
167         """ Analyze the given list of phrases and return the
168             tokenized query.
169         """
170         log().section('Analyze query (using ICU tokenizer)')
171         normalized = list(filter(lambda p: p.text,
172                                  (qmod.Phrase(p.ptype, self.normalize_text(p.text))
173                                   for p in phrases)))
174         query = qmod.QueryStruct(normalized)
175         log().var_dump('Normalized query', query.source)
176         if not query.source:
177             return query
178
179         parts, words = self.split_query(query)
180         log().var_dump('Transliterated query', lambda: _dump_transliterated(query, parts))
181
182         for row in await self.lookup_in_db(list(words.keys())):
183             for trange in words[row.word_token]:
184                 token = ICUToken.from_db_row(row)
185                 if row.type == 'S':
186                     if row.info['op'] in ('in', 'near'):
187                         if trange.start == 0:
188                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
189                     else:
190                         query.add_token(trange, qmod.TokenType.QUALIFIER, token)
191                         if trange.start == 0 or trange.end == query.num_token_slots():
192                             token = copy(token)
193                             token.penalty += 0.1 * (query.num_token_slots())
194                             query.add_token(trange, qmod.TokenType.NEAR_ITEM, token)
195                 else:
196                     query.add_token(trange, DB_TO_TOKEN_TYPE[row.type], token)
197
198         self.add_extra_tokens(query, parts)
199         self.rerank_tokens(query, parts)
200
201         log().table_dump('Word tokens', _dump_word_tokens(query))
202
203         return query
204
205
206     def normalize_text(self, text: str) -> str:
207         """ Bring the given text into a normalized form. That is the
208             standardized form search will work with. All information removed
209             at this stage is inevitably lost.
210         """
211         return cast(str, self.normalizer.transliterate(text))
212
213
214     def split_query(self, query: qmod.QueryStruct) -> Tuple[QueryParts, WordDict]:
215         """ Transliterate the phrases and split them into tokens.
216
217             Returns the list of transliterated tokens together with their
218             normalized form and a dictionary of words for lookup together
219             with their position.
220         """
221         parts: QueryParts = []
222         phrase_start = 0
223         words = defaultdict(list)
224         wordnr = 0
225         for phrase in query.source:
226             query.nodes[-1].ptype = phrase.ptype
227             for word in phrase.text.split(' '):
228                 trans = self.transliterator.transliterate(word)
229                 if trans:
230                     for term in trans.split(' '):
231                         if term:
232                             parts.append(QueryPart(term, word, wordnr))
233                             query.add_node(qmod.BreakType.TOKEN, phrase.ptype)
234                     query.nodes[-1].btype = qmod.BreakType.WORD
235                 wordnr += 1
236             query.nodes[-1].btype = qmod.BreakType.PHRASE
237
238             for word, wrange in yield_words(parts, phrase_start):
239                 words[word].append(wrange)
240
241             phrase_start = len(parts)
242         query.nodes[-1].btype = qmod.BreakType.END
243
244         return parts, words
245
246
247     async def lookup_in_db(self, words: List[str]) -> 'sa.Result[Any]':
248         """ Return the token information from the database for the
249             given word tokens.
250         """
251         t = self.conn.t.meta.tables['word']
252         return await self.conn.execute(t.select().where(t.c.word_token.in_(words)))
253
254
255     def add_extra_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
256         """ Add tokens to query that are not saved in the database.
257         """
258         for part, node, i in zip(parts, query.nodes, range(1000)):
259             if len(part.token) <= 4 and part[0].isdigit()\
260                and not node.has_tokens(i+1, qmod.TokenType.HOUSENUMBER):
261                 query.add_token(qmod.TokenRange(i, i+1), qmod.TokenType.HOUSENUMBER,
262                                 ICUToken(0.5, 0, 1, part.token, True, part.token, None))
263
264
265     def rerank_tokens(self, query: qmod.QueryStruct, parts: QueryParts) -> None:
266         """ Add penalties to tokens that depend on presence of other token.
267         """
268         for i, node, tlist in query.iter_token_lists():
269             if tlist.ttype == qmod.TokenType.POSTCODE:
270                 for repl in node.starting:
271                     if repl.end == tlist.end and repl.ttype != qmod.TokenType.POSTCODE \
272                        and (repl.ttype != qmod.TokenType.HOUSENUMBER
273                             or len(tlist.tokens[0].lookup_word) > 4):
274                         repl.add_penalty(0.39)
275             elif tlist.ttype == qmod.TokenType.HOUSENUMBER \
276                  and len(tlist.tokens[0].lookup_word) <= 3:
277                 if any(c.isdigit() for c in tlist.tokens[0].lookup_word):
278                     for repl in node.starting:
279                         if repl.end == tlist.end and repl.ttype != qmod.TokenType.HOUSENUMBER:
280                             repl.add_penalty(0.5 - tlist.tokens[0].penalty)
281             elif tlist.ttype not in (qmod.TokenType.COUNTRY, qmod.TokenType.PARTIAL):
282                 norm = parts[i].normalized
283                 for j in range(i + 1, tlist.end):
284                     if parts[j - 1].word_number != parts[j].word_number:
285                         norm += '  ' + parts[j].normalized
286                 for token in tlist.tokens:
287                     cast(ICUToken, token).rematch(norm)
288
289
290 def _dump_transliterated(query: qmod.QueryStruct, parts: QueryParts) -> str:
291     out = query.nodes[0].btype.value
292     for node, part in zip(query.nodes[1:], parts):
293         out += part.token + node.btype.value
294     return out
295
296
297 def _dump_word_tokens(query: qmod.QueryStruct) -> Iterator[List[Any]]:
298     yield ['type', 'token', 'word_token', 'lookup_word', 'penalty', 'count', 'info']
299     for node in query.nodes:
300         for tlist in node.starting:
301             for token in tlist.tokens:
302                 t = cast(ICUToken, token)
303                 yield [tlist.ttype.name, t.token, t.word_token or '',
304                        t.lookup_word or '', t.penalty, t.count, t.info]
305
306
307 async def create_query_analyzer(conn: SearchConnection) -> AbstractQueryAnalyzer:
308     """ Create and set up a new query analyzer for a database based
309         on the ICU tokenizer.
310     """
311     out = ICUQueryAnalyzer(conn)
312     await out.setup()
313
314     return out