Updates to chunk out analysis, reducing peak memory use.

Author: njhenry

.claude/TODO.md

@@ -6,6 +6,8 @@ Short running list of in-progress / upcoming work. Edit freely; trim older compl
 
 ## Upcoming
 
+- [ ] Trim peak memory in `scripts/conflation/conflate.py` so the OSM(8.7M) × Overture(13M) run fits inside 16GB+4GB swap. Widened Overture allowlist (6.29M → 13.05M) overflowed WSL2 and rebooted the VM on 2026-04-17. Tactics: drop `osm_gdf`/`overture_gdf` down to the minimal column set `build_merge_parts` needs before the scoring pass, free normalized name arrays right after `compute_name_scores`, and narrow `osm_idx`/`overture_idx` to int32. See [src/openpois/conflation/match.py](../src/openpois/conflation/match.py) and [scripts/conflation/conflate.py](../scripts/conflation/conflate.py). Added 2026-04-17.
+- [ ] Instrument setup / matching-reload phase of `scripts/conflation/conflate.py` to pin down the ~17 GB VmHWM spike observed on the 2026-04-17 chunked run (checkpoints reloaded, merge phase bounded — spike is upstream of both). Likely culprits: `pd.concat` of 128 chunk parquets, name/brand array construction holding dual refs, or taxonomy crosswalk transient. Add `psutil` RSS logging at each phase boundary so we can see exactly which step jumps. Added 2026-04-17.
 - [ ] Watch for a DuckDB release that fixes the WSL2 httpfs "Information loss on integer cast" crash (issue #21669, fix PR #21395). Once a tagged release ships with the fix and a full `scripts/overture/download.py` run on WSL2 completes, we can unpin from `duckdb==1.4.1` and revert the per-part download to a single-query DuckDB scan. Added 2026-04-17.
 - [ ] Auto-check taxonomy changes whenever we switch to a new Overture Maps version (detect new/removed L0/L1/L2 categories vs. `taxonomy_crosswalk_overture_maps.csv` and flag gaps). Added 2026-04-16.
 - [ ] Watch for Overture L0/L1 → flat `basic_category` migration (~June 2026). Crosswalk CSV + `assign_overture_shared_label` will need updating. See [docs/taxonomy-setup.md](docs/taxonomy-setup.md).

config.yaml

@@ -40,9 +40,9 @@ download:
     history_pbf_url: "https://osm-internal.download.geofabrik.de/north-america/us-internal.osh.pbf"
     pr_history_pbf_url: "https://osm-internal.download.geofabrik.de/north-america/us/puerto-rico-internal.osh.pbf"
     history_cookie_file: "~/data/openpois/.creds/geofabrik_cookies.txt"
-    overwrite_download: false
-    overwrite_filter: false
-    overwrite_parse: false
+    overwrite_download: true
+    overwrite_filter: true
+    overwrite_parse: true
     source_label: "osm"
     keep_all_keys: false
     chunk_size: 100_000
@@ -223,6 +223,7 @@ conflation:
   type_weight: 0.30
   identifier_weight: 0.20
   chunk_size: 500_000
+  chunk_target_pois: 200_000
   test_bbox:
     xmin: -122.45
     ymin: 47.50

scripts/conflation/conflate.py

@@ -37,7 +37,9 @@
 
 import argparse
 import gc
+import shutil
 import time
+from pathlib import Path
 
 import geopandas as gpd
 import numpy as np
@@ -47,11 +49,15 @@
 
 from openpois.conflation.match import (
     compute_match_scores,
+    find_and_score_matches_chunked,
     find_spatial_candidates,
     select_best_matches,
 )
+
+CHECKPOINT_SUBDIR = "chunk_matches"
 from openpois.conflation.merge import (
     build_merge_parts,
+    build_merge_parts_chunked,
     save_conflated_from_parts,
 )
 from openpois.conflation.taxonomy import (
@@ -89,6 +95,9 @@
 TYPE_WEIGHT = config.get("conflation", "type_weight")
 IDENTIFIER_WEIGHT = config.get("conflation", "identifier_weight")
 CHUNK_SIZE = config.get("conflation", "chunk_size")
+CHUNK_TARGET_POIS = config.get(
+    "conflation", "chunk_target_pois",
+)
 TEST_BBOX = config.get("conflation", "test_bbox")
 
 # Columns needed for matching (memory optimization)
@@ -146,6 +155,15 @@ def _load_gdf(
             "(Seattle area) for testing."
         ),
     )
+    parser.add_argument(
+        "--no-chunk",
+        action = "store_true",
+        help = (
+            "Disable spatial chunking and run matching over the "
+            "full dataset in one pass. Uses more memory; kept as "
+            "a debug/baseline option."
+        ),
+    )
     args = parser.parse_args()
     t0 = time.time()
 
@@ -209,75 +227,135 @@ def _load_gdf(
             overture_gdf.drop(columns = col, inplace = True)
     gc.collect()
 
-    # -- Spatial candidates ----------------------------------------
-    print(f"\nFinding spatial candidates (max {MAX_RADIUS_M}m) ...")
-    candidates = find_spatial_candidates(
-        osm_geom = osm_gdf.geometry.values,
-        overture_geom = overture_gdf.geometry.values,
-        osm_radii_m = osm_radii,
-        max_radius_m = MAX_RADIUS_M,
-        chunk_size = CHUNK_SIZE,
+    # -- Matching --------------------------------------------------
+    # Prepare name/brand arrays once (used by both code paths).
+    osm_names = osm_gdf["name"].to_numpy()
+    osm_brands = (
+        osm_gdf["brand"].to_numpy()
+        if "brand" in osm_gdf.columns
+        else np.full(len(osm_gdf), None, dtype = object)
+    )
+    overture_names = overture_gdf["overture_name"].to_numpy()
+    overture_brands = (
+        overture_gdf["brand_name"].to_numpy()
+        if "brand_name" in overture_gdf.columns
+        else np.full(
+            len(overture_gdf), None, dtype = object
+        )
     )
-    print(f"  Found {len(candidates):,} candidate pairs")
-    gc.collect()
 
-    if candidates.empty:
-        print("No spatial candidates found. Merging all as unmatched.")
-        matches = candidates
-    else:
-        # -- Scoring -----------------------------------------------
-        print("\nScoring candidates ...")
-        osm_names = osm_gdf["name"].to_numpy()
-        osm_brands = (
-            osm_gdf["brand"].to_numpy()
-            if "brand" in osm_gdf.columns
-            else np.full(len(osm_gdf), None, dtype = object)
+    chunk_summary: dict | None = None
+    checkpoint_dir: Path | None = None
+
+    if args.no_chunk:
+        # -- Non-chunked baseline (full-dataset pipeline) ----------
+        print(
+            f"\nFinding spatial candidates (max {MAX_RADIUS_M}m) ..."
         )
-        overture_names = overture_gdf["overture_name"].to_numpy()
-        overture_brands = (
-            overture_gdf["brand_name"].to_numpy()
-            if "brand_name" in overture_gdf.columns
-            else np.full(
-                len(overture_gdf), None, dtype = object
-            )
+        candidates = find_spatial_candidates(
+            osm_geom = osm_gdf.geometry.values,
+            overture_geom = overture_gdf.geometry.values,
+            osm_radii_m = osm_radii,
+            max_radius_m = MAX_RADIUS_M,
+            chunk_size = CHUNK_SIZE,
         )
+        print(f"  Found {len(candidates):,} candidate pairs")
+        gc.collect()
 
-        scored = compute_match_scores(
-            candidates = candidates,
-            osm_names = osm_names,
-            osm_brands = osm_brands,
-            overture_names = overture_names,
-            overture_brands = overture_brands,
+        if candidates.empty:
+            print(
+                "No spatial candidates found. "
+                "Merging all as unmatched."
+            )
+            matches = candidates
+        else:
+            print("\nScoring candidates ...")
+            scored = compute_match_scores(
+                candidates = candidates,
+                osm_names = osm_names,
+                osm_brands = osm_brands,
+                overture_names = overture_names,
+                overture_brands = overture_brands,
+                osm_shared_labels = osm_shared_labels,
+                overture_shared_labels = overture_shared_labels,
+                osm_radii_m = osm_radii,
+                osm_l0_bits = osm_l0_bits,
+                overture_l0_bits = overture_l0_bits,
+                distance_weight = DISTANCE_WEIGHT,
+                name_weight = NAME_WEIGHT,
+                type_weight = TYPE_WEIGHT,
+                identifier_weight = IDENTIFIER_WEIGHT,
+            )
+            print(
+                f"  Mean composite score: "
+                f"{scored['composite_score'].mean():.3f}"
+            )
+
+            print(
+                f"\nSelecting best matches "
+                f"(min_score={MIN_MATCH_SCORE}) ..."
+            )
+            matches = select_best_matches(
+                scored, min_score = MIN_MATCH_SCORE,
+            )
+            print(
+                f"  Selected {len(matches):,} one-to-one matches"
+            )
+
+            del scored, candidates
+            gc.collect()
+    else:
+        # -- Chunked driver (default) ------------------------------
+        conflation_dir = Path(OUTPUT_PATH).parent
+        checkpoint_dir = conflation_dir / CHECKPOINT_SUBDIR
+        print(
+            f"\nRunning chunked matching "
+            f"(target ~{CHUNK_TARGET_POIS:,} POIs/chunk, "
+            f"max {MAX_RADIUS_M}m) ..."
+        )
+        matches, chunk_summary = find_and_score_matches_chunked(
+            osm_geom = osm_gdf.geometry.values,
+            overture_geom = overture_gdf.geometry.values,
+            osm_radii_m = osm_radii,
             osm_shared_labels = osm_shared_labels,
             overture_shared_labels = overture_shared_labels,
-            osm_radii_m = osm_radii,
             osm_l0_bits = osm_l0_bits,
             overture_l0_bits = overture_l0_bits,
+            osm_names = osm_names,
+            osm_brands = osm_brands,
+            overture_names = overture_names,
+            overture_brands = overture_brands,
             distance_weight = DISTANCE_WEIGHT,
             name_weight = NAME_WEIGHT,
             type_weight = TYPE_WEIGHT,
             identifier_weight = IDENTIFIER_WEIGHT,
+            min_match_score = MIN_MATCH_SCORE,
+            max_radius_m = MAX_RADIUS_M,
+            chunk_target_pois = CHUNK_TARGET_POIS,
+            chunk_size = CHUNK_SIZE,
+            checkpoint_dir = checkpoint_dir,
         )
         print(
-            f"  Mean composite score: "
-            f"{scored['composite_score'].mean():.3f}"
+            f"  Selected {len(matches):,} one-to-one matches "
+            f"across {chunk_summary['n_chunks']} chunks "
+            f"(Overture dedup drops: "
+            f"{chunk_summary['n_overture_dedup_drops']:,})"
         )
 
-        # -- Best matches ------------------------------------------
-        print(
-            f"\nSelecting best matches "
-            f"(min_score={MIN_MATCH_SCORE}) ..."
-        )
-        matches = select_best_matches(
-            scored, min_score = MIN_MATCH_SCORE,
-        )
-        print(f"  Selected {len(matches):,} one-to-one matches")
+    del osm_names, osm_brands
+    del overture_names, overture_brands
+    gc.collect()
 
-        # Free scoring intermediates
-        del scored, candidates
-        del osm_names, osm_brands
-        del overture_names, overture_brands
-        gc.collect()
+    # Drop scoring intermediates the merge doesn't need. Keeps only
+    # the four columns read by ``_build_matched_gdf``.
+    keep_cols = [
+        "osm_idx", "overture_idx",
+        "composite_score", "distance_m",
+    ]
+    matches = matches[
+        [c for c in keep_cols if c in matches.columns]
+    ].reset_index(drop = True)
+    gc.collect()
 
     # -- Merge (disk-backed to limit memory) -----------------------
     print("\nMerging into unified dataset ...")
@@ -291,14 +369,27 @@ def _load_gdf(
     )
     n_matches = len(matches)
 
-    part_paths = build_merge_parts(
-        osm_gdf = osm_gdf,
-        overture_gdf = overture_gdf,
-        matches = matches,
-        osm_shared_labels = osm_shared_labels,
-        overture_shared_labels = overture_shared_labels,
-        overture_confidence_weight = OVERTURE_CONF_WEIGHT,
-    )
+    if chunk_summary is not None:
+        part_paths = build_merge_parts_chunked(
+            osm_gdf = osm_gdf,
+            overture_gdf = overture_gdf,
+            matches = matches,
+            osm_shared_labels = osm_shared_labels,
+            overture_shared_labels = overture_shared_labels,
+            osm_primary = chunk_summary["osm_primary"],
+            overture_primary = chunk_summary["overture_primary"],
+            n_chunks = chunk_summary["n_chunks"],
+            overture_confidence_weight = OVERTURE_CONF_WEIGHT,
+        )
+    else:
+        part_paths = build_merge_parts(
+            osm_gdf = osm_gdf,
+            overture_gdf = overture_gdf,
+            matches = matches,
+            osm_shared_labels = osm_shared_labels,
+            overture_shared_labels = overture_shared_labels,
+            overture_confidence_weight = OVERTURE_CONF_WEIGHT,
+        )
 
     # Free ALL source data before concat+save
     del osm_gdf, overture_gdf, matches
@@ -311,6 +402,10 @@ def _load_gdf(
     n_total = save_conflated_from_parts(part_paths, OUTPUT_PATH)
     config.write_self("conflation")
 
+    # Clear chunk checkpoints after a successful save.
+    if checkpoint_dir is not None and checkpoint_dir.exists():
+        shutil.rmtree(checkpoint_dir)
+
     # -- Summary ---------------------------------------------------
     elapsed = time.time() - t0
     print(f"\n{'=' * 60}")
@@ -324,4 +419,15 @@ def _load_gdf(
         print(
             f"  Mean match distance: {match_dist_mean:.1f}m"
         )
+    if chunk_summary is not None:
+        print(
+            f"  Chunks:         {chunk_summary['n_chunks']} "
+            f"(OSM/chunk: "
+            f"{chunk_summary['min_chunk_pois']:,}"
+            f"–{chunk_summary['max_chunk_pois']:,})"
+        )
+        print(
+            f"  Overture dedup drops: "
+            f"{chunk_summary['n_overture_dedup_drops']:,}"
+        )
     print(f"  Output: {OUTPUT_PATH}")