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__":