Merge pull request #3292 from lonvia/faster-country-search

[nominatim.git] / nominatim / api / lookup.py

diff --git a/nominatim/api/lookup.py b/nominatim/api/lookup.py
index 5c54da4a1823209a9dcede95a7f3fd730b957a4d..402b85316853173967c13c73f7f0a30aabc97c92 100644 (file)
--- a/nominatim/api/lookup.py
+++ b/nominatim/api/lookup.py
@@ -12,7 +12,7 @@ import datetime as dt
 
 import sqlalchemy as sa
 
 
 import sqlalchemy as sa
 

-from nominatim.typing import SaColumn, SaLabel, SaRow
+from nominatim.typing import SaColumn, SaRow, SaSelect

 from nominatim.api.connection import SearchConnection
 import nominatim.api.types as ntyp
 import nominatim.api.results as nres
 from nominatim.api.connection import SearchConnection
 import nominatim.api.types as ntyp
 import nominatim.api.results as nres


@@ -20,23 +20,12 @@ from nominatim.api.logging import log
 
 RowFunc = Callable[[Optional[SaRow], Type[nres.BaseResultT]], Optional[nres.BaseResultT]]
 
 
 RowFunc = Callable[[Optional[SaRow], Type[nres.BaseResultT]], Optional[nres.BaseResultT]]
 

-def _select_column_geometry(column: SaColumn,
-                            geometry_output: ntyp.GeometryFormat) -> SaLabel:
-    """ Create the appropriate column expression for selecting a
-        geometry for the details response.
-    """
-    if geometry_output & ntyp.GeometryFormat.GEOJSON:
-        return sa.literal_column(f"""
-                  ST_AsGeoJSON(CASE WHEN ST_NPoints({column.name}) > 5000
-                               THEN ST_SimplifyPreserveTopology({column.name}, 0.0001)
-                               ELSE {column.name} END)
-                  """).label('geometry_geojson')
+GeomFunc = Callable[[SaSelect, SaColumn], SaSelect]

 
 

-    return sa.func.ST_GeometryType(column).label('geometry_type')

 
 
 async def find_in_placex(conn: SearchConnection, place: ntyp.PlaceRef,
 
 
 async def find_in_placex(conn: SearchConnection, place: ntyp.PlaceRef,

-                         details: ntyp.LookupDetails) -> Optional[SaRow]:
+                         add_geometries: GeomFunc) -> Optional[SaRow]:

     """ Search for the given place in the placex table and return the
         base information.
     """
     """ Search for the given place in the placex table and return the
         base information.
     """


@@ -49,8 +38,8 @@ async def find_in_placex(conn: SearchConnection, place: ntyp.PlaceRef,
                     t.c.importance, t.c.wikipedia, t.c.indexed_date,
                     t.c.parent_place_id, t.c.rank_address, t.c.rank_search,
                     t.c.linked_place_id,
                     t.c.importance, t.c.wikipedia, t.c.indexed_date,
                     t.c.parent_place_id, t.c.rank_address, t.c.rank_search,
                     t.c.linked_place_id,

-                    t.c.centroid,
-                    _select_column_geometry(t.c.geometry, details.geometry_output))
+                    t.c.geometry.ST_Expand(0).label('bbox'),
+                    t.c.centroid)

 
     if isinstance(place, ntyp.PlaceID):
         sql = sql.where(t.c.place_id == place.place_id)
 
     if isinstance(place, ntyp.PlaceID):
         sql = sql.where(t.c.place_id == place.place_id)


@@ -65,11 +54,11 @@ async def find_in_placex(conn: SearchConnection, place: ntyp.PlaceRef,
     else:
         return None
 
     else:
         return None
 

-    return (await conn.execute(sql)).one_or_none()
+    return (await conn.execute(add_geometries(sql, t.c.geometry))).one_or_none()

 
 
 async def find_in_osmline(conn: SearchConnection, place: ntyp.PlaceRef,
 
 
 async def find_in_osmline(conn: SearchConnection, place: ntyp.PlaceRef,

-                          details: ntyp.LookupDetails) -> Optional[SaRow]:
+                          add_geometries: GeomFunc) -> Optional[SaRow]:

     """ Search for the given place in the osmline table and return the
         base information.
     """
     """ Search for the given place in the osmline table and return the
         base information.
     """


@@ -78,8 +67,7 @@ async def find_in_osmline(conn: SearchConnection, place: ntyp.PlaceRef,
     sql = sa.select(t.c.place_id, t.c.osm_id, t.c.parent_place_id,
                     t.c.indexed_date, t.c.startnumber, t.c.endnumber,
                     t.c.step, t.c.address, t.c.postcode, t.c.country_code,
     sql = sa.select(t.c.place_id, t.c.osm_id, t.c.parent_place_id,
                     t.c.indexed_date, t.c.startnumber, t.c.endnumber,
                     t.c.step, t.c.address, t.c.postcode, t.c.country_code,

-                    t.c.linegeo.ST_Centroid().label('centroid'),
-                    _select_column_geometry(t.c.linegeo, details.geometry_output))
+                    t.c.linegeo.ST_Centroid().label('centroid'))

 
     if isinstance(place, ntyp.PlaceID):
         sql = sql.where(t.c.place_id == place.place_id)
 
     if isinstance(place, ntyp.PlaceID):
         sql = sql.where(t.c.place_id == place.place_id)


@@ -89,19 +77,22 @@ async def find_in_osmline(conn: SearchConnection, place: ntyp.PlaceRef,
         sql = sql.where(t.c.osm_id == place.osm_id).limit(1)
         if place.osm_class and place.osm_class.isdigit():
             sql = sql.order_by(sa.func.greatest(0,
         sql = sql.where(t.c.osm_id == place.osm_id).limit(1)
         if place.osm_class and place.osm_class.isdigit():
             sql = sql.order_by(sa.func.greatest(0,

-                                    sa.func.least(int(place.osm_class) - t.c.endnumber),
-                                           t.c.startnumber - int(place.osm_class)))
+                                                int(place.osm_class) - t.c.endnumber,
+                                                t.c.startnumber - int(place.osm_class)))

     else:
         return None
 
     else:
         return None
 

-    return (await conn.execute(sql)).one_or_none()
+    return (await conn.execute(add_geometries(sql, t.c.linegeo))).one_or_none()

 
 
 async def find_in_tiger(conn: SearchConnection, place: ntyp.PlaceRef,
 
 
 async def find_in_tiger(conn: SearchConnection, place: ntyp.PlaceRef,

-                        details: ntyp.LookupDetails) -> Optional[SaRow]:
+                        add_geometries: GeomFunc) -> Optional[SaRow]:

     """ Search for the given place in the table of Tiger addresses and return
         the base information. Only lookup by place ID is supported.
     """
     """ Search for the given place in the table of Tiger addresses and return
         the base information. Only lookup by place ID is supported.
     """

+    if not isinstance(place, ntyp.PlaceID):
+        return None
+

     log().section("Find in TIGER table")
     t = conn.t.tiger
     parent = conn.t.placex
     log().section("Find in TIGER table")
     t = conn.t.tiger
     parent = conn.t.placex


@@ -109,61 +100,54 @@ async def find_in_tiger(conn: SearchConnection, place: ntyp.PlaceRef,
                     parent.c.osm_type, parent.c.osm_id,
                     t.c.startnumber, t.c.endnumber, t.c.step,
                     t.c.postcode,
                     parent.c.osm_type, parent.c.osm_id,
                     t.c.startnumber, t.c.endnumber, t.c.step,
                     t.c.postcode,

-                    t.c.linegeo.ST_Centroid().label('centroid'),
-                    _select_column_geometry(t.c.linegeo, details.geometry_output))
-
-    if isinstance(place, ntyp.PlaceID):
-        sql = sql.where(t.c.place_id == place.place_id)\
-                 .join(parent, t.c.parent_place_id == parent.c.place_id, isouter=True)
-    else:
-        return None
+                    t.c.linegeo.ST_Centroid().label('centroid'))\
+            .where(t.c.place_id == place.place_id)\
+            .join(parent, t.c.parent_place_id == parent.c.place_id, isouter=True)

 
 

-    return (await conn.execute(sql)).one_or_none()
+    return (await conn.execute(add_geometries(sql, t.c.linegeo))).one_or_none()

 
 
 async def find_in_postcode(conn: SearchConnection, place: ntyp.PlaceRef,
 
 
 async def find_in_postcode(conn: SearchConnection, place: ntyp.PlaceRef,

-                           details: ntyp.LookupDetails) -> Optional[SaRow]:
+                           add_geometries: GeomFunc) -> Optional[SaRow]:

     """ Search for the given place in the postcode table and return the
         base information. Only lookup by place ID is supported.
     """
     """ Search for the given place in the postcode table and return the
         base information. Only lookup by place ID is supported.
     """

+    if not isinstance(place, ntyp.PlaceID):
+        return None
+

     log().section("Find in postcode table")
     t = conn.t.postcode
     sql = sa.select(t.c.place_id, t.c.parent_place_id,
                     t.c.rank_search, t.c.rank_address,
                     t.c.indexed_date, t.c.postcode, t.c.country_code,
     log().section("Find in postcode table")
     t = conn.t.postcode
     sql = sa.select(t.c.place_id, t.c.parent_place_id,
                     t.c.rank_search, t.c.rank_address,
                     t.c.indexed_date, t.c.postcode, t.c.country_code,

-                    t.c.geometry.label('centroid'),
-                    _select_column_geometry(t.c.geometry, details.geometry_output))
-
-    if isinstance(place, ntyp.PlaceID):
-        sql = sql.where(t.c.place_id == place.place_id)
-    else:
-        return None
+                    t.c.geometry.label('centroid')) \
+            .where(t.c.place_id == place.place_id)

 
 

-    return (await conn.execute(sql)).one_or_none()
+    return (await conn.execute(add_geometries(sql, t.c.geometry))).one_or_none()

 
 
 async def find_in_all_tables(conn: SearchConnection, place: ntyp.PlaceRef,
 
 
 async def find_in_all_tables(conn: SearchConnection, place: ntyp.PlaceRef,

-                             details: ntyp.LookupDetails
+                             add_geometries: GeomFunc

                             ) -> Tuple[Optional[SaRow], RowFunc[nres.BaseResultT]]:
     """ Search for the given place in all data tables
         and return the base information.
     """
                             ) -> Tuple[Optional[SaRow], RowFunc[nres.BaseResultT]]:
     """ Search for the given place in all data tables
         and return the base information.
     """

-    row = await find_in_placex(conn, place, details)
+    row = await find_in_placex(conn, place, add_geometries)

     log().var_dump('Result (placex)', row)
     if row is not None:
         return row, nres.create_from_placex_row
 
     log().var_dump('Result (placex)', row)
     if row is not None:
         return row, nres.create_from_placex_row
 

-    row = await find_in_osmline(conn, place, details)
+    row = await find_in_osmline(conn, place, add_geometries)

     log().var_dump('Result (osmline)', row)
     if row is not None:
         return row, nres.create_from_osmline_row
 
     log().var_dump('Result (osmline)', row)
     if row is not None:
         return row, nres.create_from_osmline_row
 

-    row = await find_in_postcode(conn, place, details)
+    row = await find_in_postcode(conn, place, add_geometries)

     log().var_dump('Result (postcode)', row)
     if row is not None:
         return row, nres.create_from_postcode_row
 
     log().var_dump('Result (postcode)', row)
     if row is not None:
         return row, nres.create_from_postcode_row
 

-    row = await find_in_tiger(conn, place, details)
+    row = await find_in_tiger(conn, place, add_geometries)

     log().var_dump('Result (tiger)', row)
     return row, nres.create_from_tiger_row
 
     log().var_dump('Result (tiger)', row)
     return row, nres.create_from_tiger_row
 


@@ -172,13 +156,23 @@ async def get_detailed_place(conn: SearchConnection, place: ntyp.PlaceRef,
                              details: ntyp.LookupDetails) -> Optional[nres.DetailedResult]:
     """ Retrieve a place with additional details from the database.
     """
                              details: ntyp.LookupDetails) -> Optional[nres.DetailedResult]:
     """ Retrieve a place with additional details from the database.
     """

-    log().function('get_place_by_id', place=place, details=details)
+    log().function('get_detailed_place', place=place, details=details)

 
     if details.geometry_output and details.geometry_output != ntyp.GeometryFormat.GEOJSON:
         raise ValueError("lookup only supports geojosn polygon output.")
 
 
     if details.geometry_output and details.geometry_output != ntyp.GeometryFormat.GEOJSON:
         raise ValueError("lookup only supports geojosn polygon output.")
 

+    if details.geometry_output & ntyp.GeometryFormat.GEOJSON:
+        def _add_geometry(sql: SaSelect, column: SaColumn) -> SaSelect:
+            return sql.add_columns(sa.func.ST_AsGeoJSON(
+                                    sa.case((sa.func.ST_NPoints(column) > 5000,
+                                             sa.func.ST_SimplifyPreserveTopology(column, 0.0001)),
+                                            else_=column), 7).label('geometry_geojson'))
+    else:
+        def _add_geometry(sql: SaSelect, column: SaColumn) -> SaSelect:
+            return sql.add_columns(sa.func.ST_GeometryType(column).label('geometry_type'))
+

     row_func: RowFunc[nres.DetailedResult]
     row_func: RowFunc[nres.DetailedResult]

-    row, row_func = await find_in_all_tables(conn, place, details)
+    row, row_func = await find_in_all_tables(conn, place, _add_geometry)

 
     if row is None:
         return None
 
     if row is None:
         return None


@@ -188,13 +182,70 @@ async def get_detailed_place(conn: SearchConnection, place: ntyp.PlaceRef,
 
     # add missing details
     assert result is not None
 
     # add missing details
     assert result is not None

-    result.parent_place_id = row.parent_place_id
-    result.linked_place_id = getattr(row, 'linked_place_id', None)
-    result.admin_level = getattr(row, 'admin_level', 15)
+    if 'type' in result.geometry:
+        result.geometry['type'] = GEOMETRY_TYPE_MAP.get(result.geometry['type'],
+                                                        result.geometry['type'])

     indexed_date = getattr(row, 'indexed_date', None)
     if indexed_date is not None:
         result.indexed_date = indexed_date.replace(tzinfo=dt.timezone.utc)
 
     indexed_date = getattr(row, 'indexed_date', None)
     if indexed_date is not None:
         result.indexed_date = indexed_date.replace(tzinfo=dt.timezone.utc)
 

-    await nres.add_result_details(conn, result, details)
+    await nres.add_result_details(conn, [result], details)

 
     return result
 
     return result

+
+
+async def get_simple_place(conn: SearchConnection, place: ntyp.PlaceRef,
+                           details: ntyp.LookupDetails) -> Optional[nres.SearchResult]:
+    """ Retrieve a place as a simple search result from the database.
+    """
+    log().function('get_simple_place', place=place, details=details)
+
+    def _add_geometry(sql: SaSelect, col: SaColumn) -> SaSelect:
+        if not details.geometry_output:
+            return sql
+
+        out = []
+
+        if details.geometry_simplification > 0.0:
+            col = sa.func.ST_SimplifyPreserveTopology(col, details.geometry_simplification)
+
+        if details.geometry_output & ntyp.GeometryFormat.GEOJSON:
+            out.append(sa.func.ST_AsGeoJSON(col, 7).label('geometry_geojson'))
+        if details.geometry_output & ntyp.GeometryFormat.TEXT:
+            out.append(sa.func.ST_AsText(col).label('geometry_text'))
+        if details.geometry_output & ntyp.GeometryFormat.KML:
+            out.append(sa.func.ST_AsKML(col, 7).label('geometry_kml'))
+        if details.geometry_output & ntyp.GeometryFormat.SVG:
+            out.append(sa.func.ST_AsSVG(col, 0, 7).label('geometry_svg'))
+
+        return sql.add_columns(*out)
+
+
+    row_func: RowFunc[nres.SearchResult]
+    row, row_func = await find_in_all_tables(conn, place, _add_geometry)
+
+    if row is None:
+        return None
+
+    result = row_func(row, nres.SearchResult)
+    assert result is not None
+
+    # add missing details
+    assert result is not None
+    if hasattr(row, 'bbox'):
+        result.bbox = ntyp.Bbox.from_wkb(row.bbox)
+
+    await nres.add_result_details(conn, [result], details)
+
+    return result
+
+
+GEOMETRY_TYPE_MAP = {
+    'POINT': 'ST_Point',
+    'MULTIPOINT': 'ST_MultiPoint',
+    'LINESTRING': 'ST_LineString',
+    'MULTILINESTRING': 'ST_MultiLineString',
+    'POLYGON': 'ST_Polygon',
+    'MULTIPOLYGON': 'ST_MultiPolygon',
+    'GEOMETRYCOLLECTION': 'ST_GeometryCollection'
+}