112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from pipeline.local_temp import local_tmp_dir
|
|
from pipeline.utils.postcode_mapping import build_postcode_mapping
|
|
|
|
from .memory import release_memory
|
|
|
|
|
|
def _canonical_postcode_expr(name: str) -> pl.Expr:
|
|
return pl.col(name).str.strip_chars().str.to_uppercase()
|
|
|
|
|
|
def load_uprns(
|
|
uprn_path: Path, arcgis_path: Path | None = None
|
|
) -> tuple[pl.DataFrame, dict[str, tuple[int, int]]]:
|
|
"""Load UPRNs as a sorted polars DataFrame with OA offset lookup.
|
|
|
|
Returns (df, offsets) where offsets[oa_code] = (start_row, end_row).
|
|
Peak ~5GB during sort, steady state ~1.5GB (Arrow columnar with compact strings).
|
|
"""
|
|
import tempfile
|
|
|
|
print("Loading UPRN lookup...")
|
|
mapping = None
|
|
if arcgis_path is not None:
|
|
mapping = (
|
|
build_postcode_mapping(arcgis_path)
|
|
.with_columns(
|
|
_canonical_postcode_expr("old_postcode").alias("old_postcode"),
|
|
_canonical_postcode_expr("new_postcode").alias("new_postcode"),
|
|
)
|
|
.unique("old_postcode")
|
|
)
|
|
|
|
# Sort via streaming sink to avoid polars doubling memory during in-memory sort
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=".parquet", delete=False, dir=local_tmp_dir()
|
|
) as tmp:
|
|
tmp_path = Path(tmp.name)
|
|
uprns = (
|
|
pl.scan_parquet(uprn_path)
|
|
.select("GRIDGB1E", "GRIDGB1N", "PCDS", "OA21CD")
|
|
.filter(pl.col("OA21CD").str.starts_with("E"))
|
|
.filter(pl.col("GRIDGB1E").is_not_null() & pl.col("GRIDGB1N").is_not_null())
|
|
.with_columns(_canonical_postcode_expr("PCDS").alias("PCDS"))
|
|
.filter(pl.col("PCDS").is_not_null() & (pl.col("PCDS") != ""))
|
|
)
|
|
|
|
if mapping is not None and mapping.height > 0:
|
|
uprns = (
|
|
uprns.join(mapping.lazy(), left_on="PCDS", right_on="old_postcode", how="left")
|
|
.with_columns(pl.coalesce("new_postcode", "PCDS").alias("PCDS"))
|
|
.select("GRIDGB1E", "GRIDGB1N", "PCDS", "OA21CD")
|
|
)
|
|
|
|
uprns.sort("OA21CD").sink_parquet(tmp_path)
|
|
release_memory()
|
|
|
|
# Read the sorted data — only one copy in memory (~2GB)
|
|
df = pl.read_parquet(tmp_path)
|
|
tmp_path.unlink()
|
|
n = len(df)
|
|
print(f" Loaded {n:,} UPRNs (England)")
|
|
|
|
# Compute OA group offsets using polars (avoids 37M Python string creation)
|
|
boundary_df = (
|
|
df.lazy()
|
|
.with_row_index("_i")
|
|
.filter(
|
|
pl.col("OA21CD").shift(1).is_null()
|
|
| (pl.col("OA21CD") != pl.col("OA21CD").shift(1))
|
|
)
|
|
.select("_i", "OA21CD")
|
|
.collect()
|
|
)
|
|
starts_list = boundary_df["_i"].to_list()
|
|
oa_list = boundary_df["OA21CD"].to_list()
|
|
del boundary_df
|
|
offsets: dict[str, tuple[int, int]] = {}
|
|
for j in range(len(starts_list)):
|
|
end = starts_list[j + 1] if j + 1 < len(starts_list) else n
|
|
offsets[oa_list[j]] = (starts_list[j], end)
|
|
del starts_list, oa_list
|
|
|
|
# Drop OA column (no longer needed) to save ~400MB
|
|
df = df.select("GRIDGB1E", "GRIDGB1N", "PCDS")
|
|
release_memory()
|
|
|
|
print(f" Grouped into {len(offsets)} OAs")
|
|
return df, offsets
|
|
|
|
|
|
def get_oa_uprns(
|
|
df: pl.DataFrame, offsets: dict[str, tuple[int, int]], oa_code: str
|
|
) -> tuple[np.ndarray, list[str]]:
|
|
"""Get UPRN coordinates and postcodes for a single OA.
|
|
|
|
Returns (points_nx2, postcodes_list).
|
|
"""
|
|
s, e = offsets[oa_code]
|
|
sub = df[s:e]
|
|
points = np.column_stack(
|
|
[
|
|
sub["GRIDGB1E"].to_numpy(),
|
|
sub["GRIDGB1N"].to_numpy(),
|
|
]
|
|
)
|
|
postcodes = sub["PCDS"].to_list()
|
|
return points, postcodes
|