Merge pull request #3859 from alphagov/use-rtreelib

Use pure Python Rtree library

app/broadcast_areas/__init__.py

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
 from notifications_utils.formatters import formatted_list
 from notifications_utils.polygons import Polygons
 from notifications_utils.serialised_model import SerialisedModelCollection
+from rtreelib import Rect
 from werkzeug.utils import cached_property
 
 from .populations import CITY_OF_LONDON
@@ -157,9 +158,11 @@ class CustomBroadcastArea(BaseBroadcastArea):
     def overlapping_areas(self):
         if not self.polygons:
             return []
-        return broadcast_area_libraries.get_areas(
-            *rtree_index.intersection(self.polygons.bounds, objects='raw')
-        )
+        return broadcast_area_libraries.get_areas([
+            overlap.data for overlap in rtree_index.query(
+                Rect(*self.polygons.bounds)
+            )
+        ])
 
     @cached_property
     def count_of_phones(self):

app/broadcast_areas/create-broadcast-areas-db.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import csv
+import pickle
 import sys
 from pathlib import Path
 
@@ -15,10 +16,12 @@ from populations import (
     SMARTPHONE_OWNERSHIP_BY_AGE_RANGE,
     estimate_number_of_smartphones_for_population,
 )
-from repo import BroadcastAreasRepository, rtree_index
+from repo import BroadcastAreasRepository, rtree_index_path
+from rtreelib import Rect, RTree
 
 source_files_path = Path(__file__).resolve().parent / 'source_files'
 point_counts = []
+rtree_index = RTree()
 
 
 def simplify_geometry(feature):
@@ -227,7 +230,7 @@ def add_wards_local_authorities_and_counties():
 def _add_electoral_wards(dataset_id):
     areas_to_add = []
 
-    for index, feature in enumerate(geojson.loads(wd20_filepath.read_text())["features"]):
+    for feature in geojson.loads(wd20_filepath.read_text())["features"]:
         ward_code = feature["properties"]["wd20cd"]
         ward_name = feature["properties"]["wd20nm"]
         ward_id = "wd20-" + ward_code
@@ -243,7 +246,7 @@ def _add_electoral_wards(dataset_id):
             )
 
             if feature:
-                rtree_index.insert(index, Polygons(feature).bounds, obj=ward_id)
+                rtree_index.insert(ward_id, Rect(*Polygons(feature).bounds))
 
             areas_to_add.append([
                 ward_id, ward_name,
@@ -255,6 +258,7 @@ def _add_electoral_wards(dataset_id):
         except KeyError:
             print("Skipping", ward_code, ward_name)  # noqa: T001
 
+    rtree_index_path.open('wb').write(pickle.dumps(rtree_index))
     repo.insert_broadcast_areas(areas_to_add, keep_old_polygons)
 
 
@@ -340,6 +344,4 @@ print(  # noqa: T001
     'DONE\n'
     f'    Processed {len(point_counts):,} polygons.\n'
     f'    Highest point counts once simplifed: {most_detailed_polygons}\n'
-    f'    RTree bounds: {rtree_index.bounds}\n'
-    f'    Number of objects in Rtree: {rtree_index.get_size():,}\n'
 )

app/broadcast_areas/repo.py

@@ -1,14 +1,11 @@
 import json
 import os
+import pickle
 import sqlite3
 from pathlib import Path
 
-from rtree import index
-
-rtree_index = index.Rtree(
-    str((Path(__file__).parent / 'rtree').absolute()),
-    interleaved=True,
-)
+rtree_index_path = Path(__file__).parent / 'rtree.pickle'
+rtree_index = pickle.loads(rtree_index_path.read_bytes())
 
 
 class BroadcastAreasRepository(object):

requirements.in

@@ -19,7 +19,7 @@ gunicorn==20.1.0
 eventlet==0.30.2
 notifications-python-client==6.0.2
 Shapely==1.7.1
-Rtree==0.9.7
+rtreelib==0.2.0
 
 # PaaS
 awscli-cwlogs>=1.4,<1.5

requirements.txt

@@ -171,7 +171,7 @@ requests==2.25.1
     #   notifications-utils
 rsa==4.7.2
     # via awscli
-rtree==0.9.7
+rtreelib==0.2.0
     # via -r requirements.in
 s3transfer==0.3.6
     # via