faster lookup

lookup.py

@@ -2,9 +2,10 @@ import json
 import click
 import pygeohash as pgh
 import pyarrow.parquet as pq
+import geopandas as gpd
 from shapely.geometry import Point, LineString, MultiLineString
-from shapely import wkb
 from pyproj import Geod
+import numpy as np
 
 @click.command()
 @click.argument("latitude", type=float)
@@ -15,15 +16,13 @@ def geohash_lookup(latitude, longitude, distance, parquet_dir):
     """
     Finds polygons containing the given point and lines within a threshold (in meters).
     """
-
     # Compute geohash (precision can be adjusted)
     geohash = pgh.encode(latitude, longitude, precision=4)  # Adjust precision as needed
     parquet_file = f"{parquet_dir}/{geohash}.parquet"
 
     try:
-        # Load the Parquet file
-        table = pq.read_table(parquet_file)
-        df = table.to_pandas()
+        # Load the Parquet file directly into a GeoDataFrame
+        df = gpd.read_parquet(parquet_file)
     except FileNotFoundError:
         click.echo(json.dumps({"error": f"No Parquet file found for geohash {geohash}."}))
         return
@@ -33,46 +32,27 @@ def geohash_lookup(latitude, longitude, distance, parquet_dir):
 
     found_geometries = []
 
-    for _, row in df.iterrows():
-        # Decode geometry from WKB
-        try:
-            geom = wkb.loads(bytes(row["geometry"]))  # Ensure WKB is properly read
-        except Exception as e:
-            click.echo(json.dumps({"error": f"Failed to decode geometry: {str(e)}"}))
-            continue
+    # Vectorized distance calculation for LineStrings and MultiLineStrings
+    def calculate_min_distance(geom):
+        if isinstance(geom, LineString):
+            return min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in geom.coords)
+        elif isinstance(geom, MultiLineString):
+            return min(min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in line.coords) for line in geom.geoms)
+        return np.inf
+
+    for row in df.itertuples(index=False):
+        geom = row.geometry  # Directly use the Shapely geometry object
 
-        # ✅ Ensure `tags` is always a dictionary
         metadata = {}
-        if isinstance(row["tags"], dict):
-            metadata = {k: v for k, v in row["tags"].items() if v is not None}  # Remove empty values
+        if isinstance(row.tags, dict):
+            metadata = {k: v for k, v in row.tags.items() if v is not None}  # Remove empty values
 
-        # ✅ Check if the point is inside a polygon
         if geom.contains(point):
             found_geometries.append({"type": "polygon", **metadata})
-
-        # ✅ Handle LineString and MultiLineString
-        elif isinstance(geom, LineString):
-            try:
-                min_distance = min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in geom.coords)
-                if min_distance <= distance:
-                    found_geometries.append({"type": "line", **metadata})
-            except Exception as e:
-                click.echo(json.dumps({"error": f"Failed distance calculation (LineString): {str(e)}"}))
-                continue
-
-        elif isinstance(geom, MultiLineString):
-            try:
-                # ✅ Iterate through all LineStrings inside the MultiLineString
-                min_distance = min(
-                    min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in line.coords)
-                    for line in geom.geoms
-                )
-
-                if min_distance <= distance:
-                    found_geometries.append({"type": "multi_line", **metadata})
-            except Exception as e:
-                click.echo(json.dumps({"error": f"Failed distance calculation (MultiLineString): {str(e)}"}))
-                continue
+        elif isinstance(geom, (LineString, MultiLineString)):
+            min_distance = calculate_min_distance(geom)
+            if min_distance <= distance:
+                found_geometries.append({"type": "line" if isinstance(geom, LineString) else "multi_line", **metadata})
 
     # Output JSON result
     click.echo(json.dumps(found_geometries, indent=2, ensure_ascii=False))  # Ensure proper Unicode display