SPB Git

spb/findata

Public
4.2 KB · 88 lines python
1"""Derive 5-minute (or other) bars from validated 1-minute bars, per session_date (delta layout, supersede)."""2from __future__ import annotations34from datetime import date5from typing import Any67import pyarrow as pa8from findata_core.db import connection, transaction9from findata_quality import aggregate_bars10from findata_schema import Interval1112from findata_ingest.lake import Lake, live_files13from findata_ingest.reader import read_bars14from findata_ingest.tasks.common import live_delta_dates, live_delta_paths_for_dates, upsert_coverage151617def run(job: dict[str, Any], lake: Lake, log) -> dict[str, Any]:18    p = job.get("params") or {}19    src = p["source_dataset"]20    dst = job.get("dataset_id") or src.replace(".1m.", ".5m.")21    target = Interval(p.get("target_interval", "5m"))22    dates = [date.fromisoformat(d) for d in p.get("dates", [])]23    symbols = p.get("symbols")24    run_id = job.get("run_id")25    with connection() as conn:26        if not dates:27            # all live source dates for requested symbols28            fl = live_files(conn, src, symbols)29            ds = set()30            for f in fl:31                d = f["date_min"]32                while d <= f["date_max"]:33                    ds.add(d)34                    d = d.fromordinal(d.toordinal() + 1)35            dates = sorted(ds)36        all_syms = symbols or [r["symbol"] for r in conn.execute("SELECT DISTINCT symbol FROM coverage WHERE dataset_id=%s AND session_date = ANY(%s)", (src, dates)).fetchall()]37    written, rows = [], 038    tables = {}39    if symbols:40        # backfill path: aggregate per symbol over all its sessions and write base files per (symbol, year)41        import pyarrow.compute as pc42        retire: list[str] = []43        for s in sorted(set(x.upper() for x in symbols)):44            t, _ = read_bars(src, [s], None, None, lake=lake)45            if t.num_rows == 0:46                continue47            agg = aggregate_bars(t, target, Interval.m1)48            years = pc.year(agg.column("session_date"))49            with connection() as conn:50                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",51                                   (dst, s, s)).fetchall()52                delta_dates = live_delta_dates(conn, dst)53            if delta_dates:54                agg = agg.filter(pc.invert(pc.is_in(agg.column("session_date"), value_set=pa.array(sorted(delta_dates), pa.date32()))))55                years = pc.year(agg.column("session_date"))56            for y in sorted(set(years.to_pylist())):57                ty = agg.filter(pc.equal(years, y))58                written.append(lake.write_table(dst, ty, "base", {"symbol": s, "year": str(y)}))59            retire.extend(o["path"] for o in old)60            rows += agg.num_rows61            tables[s] = agg62        if not written:63            return {"rows": 0, "symbols": len(symbols)}64        pub = lake.publish(dst, written, sorted(set(retire)), run_id, note=f"derive {target.value} (base) from {src} for {len(symbols)} symbols")65        with transaction() as conn:66            for t in tables.values():67                upsert_coverage(conn, dst, t, target)68        log.info("derive.done", rows=rows, files=len(written), **pub)69        return {"rows": rows, "files": len(written), "symbols": len(tables), **pub}70    for d in dates:71        t, _ = read_bars(src, all_syms, d, d, lake=lake)72        if t.num_rows == 0:73            continue74        agg = aggregate_bars(t, target, Interval.m1)75        tables[d] = agg76        rows += agg.num_rows77        written.append(lake.write_table(dst, agg, "delta", {"session_date": d.isoformat()}))78    if not written:79        return {"rows": 0, "dates": len(dates)}80    with connection() as conn:81        retire = live_delta_paths_for_dates(conn, dst, list(tables)) if not symbols else []82    pub = lake.publish(dst, written, retire, run_id, note=f"derive {target.value} from {src}")83    with transaction() as conn:84        for t in tables.values():85            upsert_coverage(conn, dst, t, target)86    log.info("derive.done", rows=rows, files=len(written), **pub)87    return {"rows": rows, "files": len(written), "dates": len(tables), **pub}88
Files · ee1dcbcefcb46b74e2ae222c3c00dc318d19817f
Menu

Theme