refactor: change operations to R-tree from shapely

kraina-ai · May 15, 2024 · 6123e27 · 6123e27
1 parent 25b1ca1
commit 6123e27
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/quackosm/pbf_file_reader.py b/quackosm/pbf_file_reader.py
@@ -31,8 +31,8 @@
 from geoarrow.pyarrow import io
 from pandas.util._decorators import deprecate, deprecate_kwarg
 from pyarrow_ops import drop_duplicates
-from shapely import contains_xy
-from shapely.geometry import LinearRing, Polygon
+from shapely import STRtree
+from shapely.geometry import LinearRing, Point, Polygon
 from shapely.geometry.base import BaseGeometry, BaseMultipartGeometry
 
 from quackosm._constants import FEATURES_INDEX, GEOMETRY_COLUMN, WGS84_CRS
@@ -1163,8 +1163,20 @@ def _prefilter_elements_ids(
                     #     (lon.as_py(), lat.as_py()) for lon, lat in zip(batch["lon"], batch["lat"])
                     # ]
 
+                    tree = STRtree(
+                        [
+                            Point(lon.as_py(), lat.as_py())
+                            for lon, lat in zip(batch["lon"], batch["lat"])
+                        ]
+                    )
+
+                    # tree = STRtree([GeocodeGeometryParser().convert("Monaco-Ville, Monaco")])
 
-                    mask = contains_xy(self.geometry_filter, x=batch["lon"], y=batch["lat"])
+                    intersecting_ids_array = ids.take(
+                        tree.query(self.geometry_filter, predicate="intersects")
+                    )
+
+                    # mask = contains_xy(self.geometry_filter, x=batch["lon"], y=batch["lat"])
 
                     # pool.imap(f_wrapped, zip(da, repeat(db))), total=len(da)
                     # with Pool() as pool:
@@ -1178,7 +1190,8 @@ def _prefilter_elements_ids(
                     # ]
 
                     # total_mask = reduce(np.logical_or, masks)
-                    intersecting_ids_array = ids.filter(pa.array(mask))
+                    # intersecting_ids_array = ids.filter(pa.array(mask))
+                    # intersecting_ids_array = ids.filter(pa.array(mask))
                     intersecting_ids_batch = pa.RecordBatch.from_arrays(
                         [intersecting_ids_array], names=["id"]
                     )
@@ -2667,6 +2680,7 @@ def _group_ways_with_polars(current_ways_group_path: Path, current_destination_p
         current_destination_path
     )
 
+
 # def _check_points_in_polygon(polygon, points)-> np.ndarray[bool]:
 #     mask = point_in_polygon(points, polygon) == 1
 #     return mask