lgtm

pipeline/transform/price_estimation/estimate.py

@@ -1,4 +1,4 @@
-"""Augment properties.parquet with estimated current prices.
+"""Estimate current prices for the merged properties, as a standalone artifact.
 
 For properties with a known prior sale, applies the repeat-sales price index
 to adjust the last known price to the current date, then blends with kNN
@@ -6,8 +6,13 @@ estimates from nearby recently-sold properties. Includes:
 - Capping extreme index adjustments
 - kNN spatial blending
 
-Modifies properties.parquet in-place. Temporarily joins postcode.parquet
-for lat/lon needed by kNN, then drops those columns before writing.
+Reads the slim price_inputs.parquet (built by property_base, independently of
+merge's area features) plus postcode.parquet for the lat/lon kNN needs, and
+writes a slim price_estimates.parquet of just the natural key (Postcode +
+coalesced address) and "Estimated current price" / "Est. price per sqm".
+join_price_estimates.py joins those two columns back onto properties.parquet.
+Because the inputs do not depend on merge's area columns, adding such a column
+does not invalidate this step.
 """
 
 import argparse
@@ -26,12 +31,27 @@ from pipeline.transform.price_estimation.knn import (
 from pipeline.transform.price_estimation.utils import (
     CURRENT_FRAC_YEAR,
     CURRENT_YEAR,
+    ESTIMATE_COLUMNS,
+    JOIN_KEYS,
     MAX_LOG_ADJUSTMENT,
     interpolate_log_index,
+    join_address_expr,
     sector_expr,
     type_group_expr,
 )
 
+# Columns estimate reads from price_inputs.parquet. The two address columns are
+# only carried to build the natural join key (Postcode + coalesced address).
+INPUT_COLUMNS = [
+    "Postcode",
+    "Property type",
+    "Total floor area (sqm)",
+    "Last known price",
+    "Date of last transaction",
+    "Address per Property Register",
+    "Address per EPC",
+]
+
 MAX_KNN_TO_INDEX_RATIO = 2.0
 MIN_KNN_TO_INDEX_RATIO = 0.5
 # Cap the final estimate at this multiple of the last known price as a guard
@@ -161,13 +181,14 @@ def guarded_blend_estimates(
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Augment properties.parquet with estimated current prices"
+        description="Estimate current prices for the merged properties"
     )
     parser.add_argument(
-        "--properties",
+        "--input",
         type=Path,
         required=True,
-        help="Path to properties.parquet (modified in-place)",
+        help="Path to price_inputs.parquet (slim per-dwelling inputs from "
+        "property_base)",
     )
     parser.add_argument(
         "--postcodes",
@@ -178,22 +199,23 @@ def main():
     parser.add_argument(
         "--index", type=Path, required=True, help="Path to price_index.parquet"
     )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        required=True,
+        help="Output price_estimates.parquet (natural key + estimate columns)",
+    )
     args = parser.parse_args()
 
-    print("Loading properties.parquet...")
-    df = pl.read_parquet(args.properties)
-    print(f"  {len(df):,} rows, {len(df.columns)} columns")
+    print("Loading price inputs (projection)...")
+    df = pl.read_parquet(args.input, columns=INPUT_COLUMNS)
+    print(f"  {len(df):,} rows, {len(INPUT_COLUMNS)} input columns")
 
     # Join lat/lon from postcode.parquet for kNN spatial queries
     postcodes = pl.read_parquet(args.postcodes).select("Postcode", "lat", "lon")
     df = df.join(postcodes, on="Postcode", how="left")
     print(f"  Joined lat/lon from {len(postcodes):,} postcodes")
 
-    # Drop existing estimated columns if re-running
-    for col in ["Estimated current price", "Est. price per sqm"]:
-        if col in df.columns:
-            df = df.drop(col)
-
     # Derive helper columns
     df = df.with_columns(
         sector_expr().alias("_sector"),
@@ -355,16 +377,15 @@ def main():
         .alias("Est. price per sqm"),
     )
 
-    # Drop all temporary columns and joined lat/lon (those belong in postcode.parquet)
-    temp_cols = [c for c in df.columns if c.startswith("_") or c.startswith("log_idx_")]
-    df = df.drop(temp_cols).drop("lat", "lon")
+    # Emit only the natural join key and the two estimate columns.
+    # join_price_estimates.py joins these back onto properties.parquet.
+    result = df.with_columns(join_address_expr()).select(*JOIN_KEYS, *ESTIMATE_COLUMNS)
 
-    df.write_parquet(args.properties)
-    size_mb = args.properties.stat().st_size / (1024 * 1024)
-    print(f"\nWrote {args.properties} ({size_mb:.1f} MB)")
-    print(
-        f"  {len(df):,} rows, {len(df.columns)} columns (including 'Estimated current price')"
-    )
+    result.write_parquet(args.output)
+    size_mb = args.output.stat().st_size / (1024 * 1024)
+    print(f"\nWrote {args.output} ({size_mb:.1f} MB)")
+    n_priced = result.filter(pl.col("Estimated current price").is_not_null()).height
+    print(f"  {len(result):,} rows, {n_priced:,} with an estimated current price")
 
 
 if __name__ == "__main__":

pipeline/transform/price_estimation/utils.py

@@ -16,6 +16,26 @@ LATEST_COMPLETE_YEAR = CURRENT_YEAR - 1
 _today = date.today()
 CURRENT_FRAC_YEAR = _today.year + (_today.month - 1) / 12
 
+# The two columns price estimation contributes to properties.parquet, kept here
+# so both the producer (estimate) and the joiner (join_price_estimates) agree.
+ESTIMATE_COLUMNS = ["Estimated current price", "Est. price per sqm"]
+
+# Natural join key from estimates back onto properties: postcode plus the
+# coalesced register/EPC address. This is unique and non-null on the deduped
+# dwelling universe (see property_base._dedupe_collapsed_properties), so it maps
+# estimates 1:1 onto properties regardless of row order — estimates are computed
+# from a separate price_inputs.parquet, so a positional key would not line up.
+JOIN_ADDRESS = "_join_address"
+JOIN_KEYS = ["Postcode", JOIN_ADDRESS]
+
+
+def join_address_expr() -> pl.Expr:
+    """The coalesced address half of the natural key, aliased to JOIN_ADDRESS."""
+    return pl.coalesce("Address per Property Register", "Address per EPC").alias(
+        JOIN_ADDRESS
+    )
+
+
 # Cap on log(index_ratio) to prevent wild estimates from thin sectors
 MAX_LOG_ADJUSTMENT = 3.0  # ~20x max price change
 TERRACE_TYPES = [