42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
from pathlib import Path
|
|
import polars as pl
|
|
|
|
from pipeline.config import AGGREGATES_DIR, H3_RESOLUTIONS
|
|
|
|
|
|
def aggregate(df: pl.LazyFrame, resolution: int) -> pl.LazyFrame:
|
|
"""Aggregate property data by H3 cell and year."""
|
|
h3_col = f"h3_res{resolution}"
|
|
|
|
return (
|
|
df.group_by(h3_col, "year")
|
|
.agg(
|
|
pl.len().alias("count"),
|
|
pl.col("price").mean().alias("avg_price"),
|
|
pl.col("price").median().alias("median_price"),
|
|
pl.col("price").min().alias("min_price"),
|
|
pl.col("price").max().alias("max_price"),
|
|
)
|
|
.rename({h3_col: "h3"})
|
|
)
|
|
|
|
|
|
def aggregate_all(df: pl.LazyFrame) -> dict[int, pl.LazyFrame]:
|
|
"""Aggregate at all H3 resolutions."""
|
|
return {res: aggregate(df, res) for res in H3_RESOLUTIONS}
|
|
|
|
|
|
def save_aggregates(df: pl.LazyFrame, output_dir: Path | None = None) -> list[Path]:
|
|
"""Aggregate and save at all H3 resolutions."""
|
|
output_dir = output_dir or AGGREGATES_DIR
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
saved_paths = []
|
|
aggregates = aggregate_all(df)
|
|
|
|
for res, agg_df in aggregates.items():
|
|
output_path = output_dir / f"res{res}.parquet"
|
|
agg_df.collect().write_parquet(output_path)
|
|
saved_paths.append(output_path)
|
|
|
|
return saved_paths
|