Move transform logic around

pipeline/download/pois.py

@@ -0,0 +1,171 @@
+import argparse
+import tempfile
+import urllib.request
+from pathlib import Path
+from tempfile import mkdtemp
+
+import osmium
+import polars as pl
+from tqdm import tqdm
+
+from pathlib import Path
+
+
+BATCH_SIZE = 50_000
+
+MIN_OCCURENCE_COUNT = 20
+
+GEOFABRIK_GB_URL = (
+    "https://download.geofabrik.de/europe/great-britain-latest.osm.pbf"
+)
+
+UK_BBOX_WEST = -7.57
+UK_BBOX_SOUTH = 49.96
+UK_BBOX_EAST = 1.68
+UK_BBOX_NORTH = 58.64
+
+POI_TAG_KEYS: list[str] = [
+    "amenity",
+    "building",
+    "craft",
+    "emergency",
+    "healthcare",
+    "leisure",
+    "office",
+    "shop",
+    "tourism",
+    "public_transport",
+]
+
+
+
+def download_pbf(pbf_file: Path) -> None:
+    pbf_file.parent.mkdir(parents=True, exist_ok=True)
+    tmp = pbf_file.with_suffix(".pbf.tmp")
+    print(f"Downloading {GEOFABRIK_GB_URL}")
+
+    with (
+        tqdm(unit="B", unit_scale=True, desc="Downloading") as bar,
+        urllib.request.urlopen(GEOFABRIK_GB_URL) as resp,
+        open(tmp, "wb") as f,
+    ):
+        length = resp.headers.get("Content-Length")
+        if length:
+            bar.total = int(length)
+        while chunk := resp.read(1 << 20):
+            f.write(chunk)
+            bar.update(len(chunk))
+
+    tmp.rename(pbf_file)
+    print(f"Saved to {pbf_file}")
+
+
+class POIHandler(osmium.SimpleHandler):
+    def __init__(self, progress: tqdm, tmp_dir: Path) -> None:
+        super().__init__()
+        self._batch: list[dict] = []
+        self._tmp_dir = tmp_dir
+        self._batch_num = 0
+        self.poi_count = 0
+        self._progress = progress
+
+    def _in_uk(self, lat: float, lon: float) -> bool:
+        return (
+            UK_BBOX_SOUTH <= lat <= UK_BBOX_NORTH
+            and UK_BBOX_WEST <= lon <= UK_BBOX_EAST
+        )
+
+    def _match_tags(self, tags: osmium.osm.TagList) -> list[str]:
+        return [f"{key}/{tags[key]}" for key in POI_TAG_KEYS if key in tags]
+
+    def _get_name(self, tags: osmium.osm.TagList) -> str:
+        return tags.get("name:en", tags.get("name", ""))
+
+    def _flush_batch(self) -> None:
+        if not self._batch:
+            return
+        df = pl.DataFrame(self._batch)
+        out = self._tmp_dir / f"batch_{self._batch_num:05d}.parquet"
+        df.write_parquet(out)
+        self._batch_num += 1
+        self._batch.clear()
+
+    def _add_poi(
+        self, osm_id: str, tags: osmium.osm.TagList, category: str, lat: float, lng: float
+    ) -> None:
+        self._batch.append(
+            {
+                "id": osm_id,
+                "name": self._get_name(tags),
+                "category": category,
+                "lat": lat,
+                "lng": lng,
+            }
+        )
+        self.poi_count += 1
+        self._progress.set_postfix(pois=f"{self.poi_count:,}", refresh=False)
+        if len(self._batch) >= BATCH_SIZE:
+            self._flush_batch()
+
+    def _tick(self) -> None:
+        self._progress.update(1)
+
+    def node(self, n: osmium.osm.Node) -> None:
+        self._tick()
+        if not n.location.valid:
+            return
+        lat, lon = n.location.lat, n.location.lon
+        if not self._in_uk(lat, lon):
+            return
+        categories = self._match_tags(n.tags)
+        for category in categories:
+            self._add_poi(f"n{n.id}", n.tags, category, lat, lon)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Download and extract POIs from OpenStreetMap")
+    parser.add_argument("--output", type=Path, required=True, help="Output parquet file path")
+    args = parser.parse_args()
+
+    with tempfile.TemporaryDirectory() as cache_dir:
+        pbf_file = Path(cache_dir) / "great-britain-latest.osm.pbf"
+
+        if not pbf_file.exists():
+            download_pbf(pbf_file)
+        else:
+            print(f"Using cached PBF file at {pbf_file}")
+
+        print(f"Tag keys: {POI_TAG_KEYS}")
+
+        tmp_dir = Path(mkdtemp(prefix="pois_"))
+        with tqdm(
+            unit=" elements",
+            unit_scale=True,
+            desc="Streaming",
+            smoothing=0.05,
+            mininterval=1.0,
+        ) as progress:
+            handler = POIHandler(progress, tmp_dir)
+            handler.apply_file(str(pbf_file), locations=True)
+            handler._flush_batch()  # write any remaining POIs
+
+        print(f"Extracted {handler.poi_count:,} POIs")
+
+        batch_files = sorted(tmp_dir.glob("batch_*.parquet"))
+        df = pl.concat([pl.scan_parquet(f) for f in batch_files])
+
+        # Only keep categories with enough occurrences
+        valid_categories = (
+            df.group_by("category")
+            .agg(pl.len().alias("count"))
+            .filter(pl.col("count") >= MIN_OCCURENCE_COUNT)
+        )
+        df = df.join(valid_categories.select("category"), on="category", how="semi")
+
+        print(f"Total POIs: {handler.poi_count:,}")
+        df.sink_parquet(args.output)
+        print(f"Saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()