changes

pipeline/transform/price_estimation/knn.py

@@ -29,7 +29,7 @@ def _scale_coords(lat: np.ndarray, lon: np.ndarray) -> np.ndarray:
 
 
 def build_knn_pool(
-    input_path: Path,
+    source: Path | pl.LazyFrame,
     index: pl.DataFrame,
     ref_frac_year: float,
     max_sale_year: int | None = None,
@@ -42,8 +42,9 @@ def build_knn_pool(
     Returns dict mapping type_group -> (KDTree over scaled lat/lon, adjusted_psm array).
     """
     print("Building kNN pool...")
+    lf = pl.scan_parquet(source) if isinstance(source, Path) else source
     query = (
-        pl.scan_parquet(input_path)
+        lf
         .select(
             "Postcode",
             "Property type",