Rewrites to reduce peak memory use.

Author: njhenry

scripts/osm_snapshot/apply_model.py

@@ -21,17 +21,16 @@
   model_version                      — which model version was used
   model_group                        — which group was matched, or "constant"
 
-Saves as a spatially-sorted GeoParquet (Hilbert curve order, zstd compression,
-50k-row row groups) for efficient cloud-native range reads from S3.
+Streams input row-groups via pyarrow, computes predictions per batch in numpy,
+and appends the new columns before writing each batch to the output parquet.
+Input row order is preserved — the downstream `format_for_upload.py` step
+re-sorts by geohash, so no intermediate spatial ordering is needed here.
 """
 from __future__ import annotations
 
 import argparse
-import gc
-import io
 from pathlib import Path
 
-import geopandas as gpd
 import numpy as np
 import pandas as pd
 import pyarrow as pa
@@ -60,6 +59,77 @@
 # Base directory containing all versioned model subdirectories
 MODEL_BASE = Path(config.get_dir_path("model_output")).parent
 
+BATCH_ROWS = 500_000
+ROW_GROUP_SIZE = 50_000
+
+
+# -----------------------------------------------------------------------------
+# Per-batch prediction logic (numpy only)
+# -----------------------------------------------------------------------------
+
+def _compute_batch_predictions(
+    df_lookup: pd.DataFrame,
+    const_arr: np.ndarray,
+    by_key_lookups: dict[str, tuple[list[str], np.ndarray]],
+    constant_version: str,
+) -> tuple[dict[str, np.ndarray], np.ndarray]:
+    """
+    Given only the `last_edited` + `FILTER_KEYS` columns of a batch, compute
+    the 6 prediction columns plus a boolean `matched` mask (True where a
+    per-key random-effects model was used, False where the constant fallback
+    applied). Mirrors the logic of the original in-memory implementation.
+    """
+    n = len(df_lookup)
+    today = pd.Timestamp.now(tz = "UTC")
+    last_edited = df_lookup["last_edited"]
+    if last_edited.dt.tz is None:
+        last_edited = last_edited.dt.tz_localize("UTC")
+    n_null = last_edited.isna().sum()
+    if n_null:
+        raise ValueError(
+            f"{n_null} rows have a null last_edited timestamp. "
+            "Remove or impute these rows before applying the model."
+        )
+    elapsed_secs = (today - last_edited).dt.total_seconds().to_numpy()
+    elapsed_years = elapsed_secs / (365.25 * 86_400)
+    t2_years = np.clip(np.round(elapsed_years * 10) / 10, 0.0, 10.0)
+    t2_int_arr = np.round(t2_years * 10).astype(int)
+
+    p_arr = const_arr[t2_int_arr].copy()
+    model_version_arr = np.full(n, constant_version, dtype = object)
+    model_group_arr = np.full(n, "constant", dtype = object)
+    matched = np.zeros(n, dtype = bool)
+
+    for key in FILTER_KEYS:
+        if key not in by_key_lookups:
+            continue
+        groups, group_arr = by_key_lookups[key]
+        group_to_idx = {g: i for i, g in enumerate(groups)}
+        group_ids = (
+            df_lookup[key].map(group_to_idx).fillna(-1).astype(int).to_numpy()
+        )
+        eligible = ~matched & (group_ids >= 0)
+        eli_pos = np.where(eligible)[0]
+        if len(eli_pos) == 0:
+            continue
+        p_arr[eli_pos] = group_arr[group_ids[eli_pos], t2_int_arr[eli_pos]]
+        model_version_arr[eli_pos] = f"{MODEL_STUB}_by_{key}"
+        model_group_arr[eli_pos] = df_lookup[key].to_numpy()[eli_pos]
+        matched[eli_pos] = True
+
+    return (
+        {
+            "t2_years": t2_years,
+            "conf_mean": 1.0 - p_arr[:, 0],
+            "conf_lower": 1.0 - p_arr[:, 2],   # 1 - p_upper
+            "conf_upper": 1.0 - p_arr[:, 1],   # 1 - p_lower
+            "model_version": model_version_arr,
+            "model_group": model_group_arr,
+        },
+        matched,
+    )
+
+
 # -----------------------------------------------------------------------------
 # Main
 # -----------------------------------------------------------------------------
@@ -71,7 +141,7 @@
     parser.add_argument(
         "--test",
         action = "store_true",
-        help = "Load only the first 10,000 rows of the snapshot for testing.",
+        help = "Process only the first 10,000 rows of the snapshot for testing.",
     )
     args = parser.parse_args()
 
@@ -99,120 +169,74 @@
         else:
             print(f"  No predictions found for {version}; will skip")
 
-    # -- Read snapshot ----------------------------------------------------------
+    # -- Open input, build output schema ----------------------------------------
     print(f"\nReading OSM snapshot from {SNAPSHOT_PATH} ...")
-    if args.test:
-        # Read only the first row group from disk rather than all 7.8M rows,
-        # then round-trip through BytesIO to preserve GeoParquet metadata.
-        print("  (--test mode: loading first 10,000 rows only)")
-        pf = pq.ParquetFile(SNAPSHOT_PATH)
-        buf = io.BytesIO()
-        pq.write_table(pf.read_row_group(0).slice(0, 10_000), buf)
-        buf.seek(0)
-        gdf = gpd.read_parquet(buf)
-    else:
-        gdf = gpd.read_parquet(SNAPSHOT_PATH)
-    print(f"  {len(gdf):,} POIs loaded")
-
-    n = len(gdf)
-
-    # -- Compute years since last edit ------------------------------------------
-    today = pd.Timestamp.now(tz = "UTC")
-    last_edited = gdf["last_edited"]
-    if last_edited.dt.tz is None:
-        last_edited = last_edited.dt.tz_localize("UTC")
-    n_null = last_edited.isna().sum()
-    if n_null:
-        raise ValueError(
-            f"{n_null} rows have a null last_edited timestamp. "
-            "Remove or impute these rows before applying the model."
-        )
-    elapsed_secs = (today - last_edited).dt.total_seconds().to_numpy()
-    elapsed_years = elapsed_secs / (365.25 * 86_400)
-    t2_years = np.clip(np.round(elapsed_years * 10) / 10, 0.0, 10.0)
-    # t2_int_arr stays in numpy only; never written to gdf
-    t2_int_arr = np.round(t2_years * 10).astype(int)
-    gdf["t2_years"] = t2_years
-
-    # -- Assign predictions via numpy arrays ------------------------------------
-    # All matching and lookup work is done in numpy and written to gdf once at
-    # the end, avoiding repeated pandas indexing overhead across 7.8M rows.
-
-    # Initialize from constant model: single vectorized index → shape (n, 3)
-    p_arr = const_arr[t2_int_arr].copy()   # columns: p_mean, p_lower, p_upper
-    model_version_arr = np.full(n, constant_version, dtype = object)
-    model_group_arr = np.full(n, "constant", dtype = object)
-    matched = np.zeros(n, dtype = bool)
-
-    for key in FILTER_KEYS:
-        if key not in by_key_lookups:
-            continue
-        groups, group_arr = by_key_lookups[key]   # group_arr: (n_groups, 101, 3)
-
-        # Map tag values to group indices; NaN and unknown values become -1.
-        group_to_idx = {g: i for i, g in enumerate(groups)}
-        group_ids = (
-            gdf[key].map(group_to_idx).fillna(-1).astype(int).to_numpy()
-        )
-
-        eligible = ~matched & (group_ids >= 0)
-        eli_pos = np.where(eligible)[0]
-        if len(eli_pos) == 0:
-            continue
+    pf = pq.ParquetFile(SNAPSHOT_PATH)
+    n_total = pf.metadata.num_rows
+    print(f"  {n_total:,} POIs across {pf.num_row_groups} row groups")
+
+    # Preserve the input GeoParquet file-level metadata (contains the `geo`
+    # block that marks `geometry` as the primary geometry column + its CRS).
+    # We only append new columns — existing schema + metadata carry through.
+    input_schema = pf.schema_arrow
+    new_fields = [
+        pa.field("t2_years", pa.float64()),
+        pa.field("conf_mean", pa.float64()),
+        pa.field("conf_lower", pa.float64()),
+        pa.field("conf_upper", pa.float64()),
+        pa.field("model_version", pa.string()),
+        pa.field("model_group", pa.string()),
+    ]
+    output_schema = pa.schema(
+        list(input_schema) + new_fields,
+        metadata = input_schema.metadata,
+    )
 
-        # Vectorized 2D fancy indexing: group_arr[m_gids, m_t2s] → (m, 3)
-        p_arr[eli_pos] = group_arr[group_ids[eli_pos], t2_int_arr[eli_pos]]
-        model_version_arr[eli_pos] = f"{MODEL_STUB}_by_{key}"
-        model_group_arr[eli_pos] = gdf[key].to_numpy()[eli_pos]
-        matched[eli_pos] = True
+    # -- Stream: read batch → append prediction columns → write -----------------
+    lookup_cols = ["last_edited"] + [k for k in FILTER_KEYS if k in by_key_lookups]
+    version_counts: dict[str, int] = {
+        f"{MODEL_STUB}_by_{k}": 0 for k in by_key_lookups
+    }
+    version_counts[constant_version] = 0
+    n_written = 0
 
-        print(f"  {MODEL_STUB}_by_{key}: matched {len(eli_pos):,} POIs")
-
-    n_constant = int((~matched).sum())
-    print(f"  {constant_version}: {n_constant:,} POIs (fallback)")
-
-    # -- Assign back to GeoDataFrame --------------------------------------------
-    # Convert change probability to confidence (1 - p).
-    # Note: conf_lower = 1 - p_upper and conf_upper = 1 - p_lower.
-    gdf["conf_mean"] = 1.0 - p_arr[:, 0]
-    gdf["conf_lower"] = 1.0 - p_arr[:, 2]   # 1 - p_upper
-    gdf["conf_upper"] = 1.0 - p_arr[:, 1]   # 1 - p_lower
-    # Categorical dtype stores integer codes + a small lookup, saving ~90%
-    # memory vs. object strings for these low-cardinality columns.
-    gdf["model_version"] = pd.Categorical(model_version_arr)
-    gdf["model_group"] = pd.Categorical(model_group_arr)
-
-    # -- Spatial sort for cloud-native S3 reads ---------------------------------
-    # Two-pass approach avoids holding two GDF copies in memory at once:
-    #   1. compute sorted indices, report model breakdown, write unsorted
-    #      GeoParquet to a temp file, drop the GDF
-    #   2. read the temp back as a pa.Table, reorder via .take(), write final
-    print("\nComputing Hilbert curve order ...")
-    sorted_indices = gdf.hilbert_distance().to_numpy().argsort()
+    OUTPUT_PATH.parent.mkdir(parents = True, exist_ok = True)
+    print(f"\nWriting to {OUTPUT_PATH} ...", flush = True)
+
+    with pq.ParquetWriter(
+        OUTPUT_PATH, output_schema, compression = "zstd"
+    ) as writer:
+        if args.test:
+            print("  (--test mode: first 10,000 rows only)")
+            batches = [next(pf.iter_batches(batch_size = 10_000))]
+        else:
+            batches = pf.iter_batches(batch_size = BATCH_ROWS)
+
+        for batch in batches:
+            tbl = pa.Table.from_batches([batch])
+            df_lookup = tbl.select(lookup_cols).to_pandas()
+            preds, matched = _compute_batch_predictions(
+                df_lookup, const_arr, by_key_lookups, constant_version,
+            )
+
+            for key in by_key_lookups:
+                version = f"{MODEL_STUB}_by_{key}"
+                version_counts[version] += int(
+                    (preds["model_version"] == version).sum()
+                )
+            version_counts[constant_version] += int(len(df_lookup) - matched.sum())
+
+            for field in new_fields:
+                tbl = tbl.append_column(
+                    field.name,
+                    pa.array(preds[field.name], type = field.type),
+                )
+
+            writer.write_table(tbl, row_group_size = ROW_GROUP_SIZE)
+            n_written += batch.num_rows
+            print(f"  {n_written:,}/{n_total:,} rows written", flush = True)
 
     print("\nModel version breakdown:")
-    print(gdf["model_version"].value_counts().to_string())
-
-    OUTPUT_PATH.parent.mkdir(parents = True, exist_ok = True)
-    tmp_path = OUTPUT_PATH.with_suffix(OUTPUT_PATH.suffix + ".unsorted")
-    print(f"\nWriting unsorted snapshot to {tmp_path} ...")
-    gdf.to_parquet(tmp_path, compression = "zstd")
-    n_total = len(gdf)
-    del gdf
-    gc.collect()
-
-    print(f"Reordering via PyArrow and writing final output to {OUTPUT_PATH} ...")
-    table = pq.read_table(tmp_path)
-    sorted_table = table.take(pa.array(sorted_indices))
-    del table
-    gc.collect()
-    pq.write_table(
-        sorted_table,
-        OUTPUT_PATH,
-        compression = "zstd",
-        row_group_size = 50_000,
-    )
-    del sorted_table
-    gc.collect()
-    tmp_path.unlink()
-    print(f"Done. Saved {n_total:,} POIs.")
+    for version, count in sorted(version_counts.items(), key = lambda kv: -kv[1]):
+        print(f"  {version}: {count:,}")
+    print(f"\nDone. Saved {n_written:,} POIs.", flush = True)