SPB Git

spb/findata

Public
3.8 KB · 63 lines python
1"""Compaction: merge session_date delta files older than `min_age_days` into base files per (symbol, year).2Asserts the (instrument_id, timestamp_utc) uniqueness invariant before publishing."""3from __future__ import annotations45from datetime import UTC, datetime, timedelta6from typing import Any78import duckdb9import pyarrow as pa10from findata_core.db import connection11from findata_schema import BARS_SCHEMA1213from findata_ingest.lake import Lake141516def run(job: dict[str, Any], lake: Lake, log) -> dict[str, Any]:17    p = job.get("params") or {}18    dataset_id = job["dataset_id"]19    min_age = int(p.get("min_age_days", 35))20    max_files = int(p.get("max_delta_files", 40))21    cutoff = (datetime.now(UTC) - timedelta(days=min_age)).date()22    with connection() as conn:23        busy = conn.execute("SELECT count(*) AS n FROM ingest_jobs WHERE dataset_id=%s AND kind IN ('backfill','derive') AND status IN ('queued','running') AND id<>%s",24                            (dataset_id, job["id"])).fetchone()["n"]25        if busy:26            log.info("compaction.deferred", reason=f"{busy} backfill/derive job(s) pending on {dataset_id}")27            return {"compacted_files": 0, "deferred": True}28        deltas = conn.execute("SELECT path, date_min FROM lake_files WHERE dataset_id=%s AND layout='delta' AND version_removed IS NULL AND date_max <= %s ORDER BY date_min LIMIT %s",29                              (dataset_id, cutoff, max_files)).fetchall()30    if not deltas:31        return {"compacted_files": 0}32    paths = [str(lake.abs(d["path"])) for d in deltas]33    con = duckdb.connect(config={"threads": 8, "memory_limit": "16GB"})34    # read every delta once (sorted), then slice in memory per (symbol, year) — avoids thousands of file scans35    all_rows = con.sql("SELECT * FROM read_parquet($p, hive_partitioning=false) ORDER BY symbol_as_reported, session_date, timestamp_utc",36                       params={"p": paths}).to_arrow_table().cast(BARS_SCHEMA)37    import polars as pl38    df = pl.from_arrow(all_rows).with_columns(pl.col("session_date").dt.year().alias("_y"))39    written, retire = [], [d["path"] for d in deltas]40    rows = 041    groups = df.partition_by(["symbol_as_reported", "_y"], as_dict=True, maintain_order=True)42    log.info("compaction.loaded", delta_files=len(deltas), rows=df.height, groups=len(groups))43    for (s, y), new_df in groups.items():44        new = new_df.drop("_y").to_arrow().cast(BARS_SCHEMA)45        with connection() as conn:46            old = conn.execute("SELECT path FROM lake_files WHERE dataset_id=%s AND layout='base' AND version_removed IS NULL AND symbol_min=%s AND symbol_max=%s AND date_min>=%s AND date_max<=%s",47                               (dataset_id, s, s, f"{y}-01-01", f"{y}-12-31")).fetchall()48        parts = [new]49        for o in old:50            parts.append(con.sql("SELECT * FROM read_parquet($p, hive_partitioning=false)", params={"p": [str(lake.abs(o["path"]))]}).to_arrow_table().cast(BARS_SCHEMA))51            retire.append(o["path"])52        merged = pa.concat_tables(parts)53        con.register("m", merged)54        dup = con.sql("SELECT count(*) AS n FROM (SELECT instrument_id, timestamp_utc FROM m GROUP BY 1,2 HAVING count(*)>1)").fetchall()[0][0]55        con.unregister("m")56        if dup:57            raise RuntimeError(f"compaction invariant violated: {dup} duplicate (instrument_id, timestamp_utc) for {s} {y}")58        written.append(lake.write_table(dataset_id, merged, "base", {"symbol": s, "year": str(y)}))59        rows += merged.num_rows60    pub = lake.publish(dataset_id, written, sorted(set(retire)), job.get("run_id"), note=f"compaction {len(deltas)} delta files")61    log.info("compaction.done", deltas=len(deltas), base_files=len(written), rows=rows, **pub)62    return {"compacted_files": len(deltas), "base_files_written": len(written), "rows": rows, **pub}63
Files · main
Menu

Theme