stream orderbook row-groups via pyarrow (fix OOM, peak 5MB/slug)

data_loader.py

@@ -199,10 +199,9 @@ def iter_orderbook_batches(
     slugs: Iterable[str],
     batch_size: int = 500,
 ):
-    """Yield polars DataFrames, each containing orderbook rows for up to
-    `batch_size` slugs. Relies on parquet row-group pushdown of the slug
-    filter — keeps peak memory at O(batch_size * per_slug_bytes) instead of
-    loading the full asset parquet (~37 GB for BTC) into RAM.
+    """DEPRECATED: polars scan-filter-collect reads the full 37 GB parquet even
+    when filtering to a small slug list (is_in doesn't do row-group pushdown).
+    Kept for backwards-compat callers; use `iter_orderbook_slug_pairs` instead.
     """
     asset = asset.lower()
     local = _orderbook_local_path(asset, hf_token, cache_dir)
@@ -216,6 +215,99 @@ def iter_orderbook_batches(
         yield df, batch
 
 
+def _arrow_rg_to_polars(tbl) -> "pl.DataFrame":
+    """Convert an arrow row-group Table to a polars DataFrame with the right
+    dtypes: prices → Float32, sizes → Float64 (strings in storage)."""
+    df = pl.from_arrow(tbl)
+    casts = []
+    for c in _OB_PX_COLS:
+        if c in df.columns:
+            casts.append(pl.col(c).cast(pl.Float32, strict=False).alias(c))
+    for c in _OB_SZ_COLS:
+        if c in df.columns:
+            casts.append(pl.col(c).cast(pl.Float64, strict=False).alias(c))
+    if casts:
+        df = df.with_columns(casts)
+    return df
+
+
+def iter_orderbook_slug_pairs(
+    asset: str,
+    hf_token: str,
+    cache_dir: Path,
+    wanted_slugs: Iterable[str],
+):
+    """Stream (slug, ob_up, ob_dn) tuples directly from parquet row groups.
+
+    The seeder wrote each (slug, outcome) intermediate via a single
+    `ParquetWriter.write_table()` call → each row group in the final parquet
+    contains exactly one (slug, outcome) pair. We iterate row groups in file
+    order, grouping Down+Up pairs per slug, and yield only slugs in
+    `wanted_slugs`.
+
+    Peak memory: ~2 row groups (~5 MB for BTC) regardless of asset size.
+    Works for the BTC 37 GB parquet on a 32 GB Space.
+    """
+    import pyarrow.parquet as pq
+
+    asset = asset.lower()
+    local = _orderbook_local_path(asset, hf_token, cache_dir)
+    wanted = set(wanted_slugs)
+    if not wanted:
+        return
+
+    pf = pq.ParquetFile(str(local))
+    avail_cols = pf.schema.names
+    cols = [c for c in _OB_BASE_COLS + _OB_PX_COLS + _OB_SZ_COLS if c in avail_cols]
+
+    current_slug: Optional[str] = None
+    ob_up_tbls: list = []
+    ob_dn_tbls: list = []
+
+    def _emit(slug, up_tbls, dn_tbls):
+        if slug not in wanted:
+            return None
+        if up_tbls:
+            up_tbl = up_tbls[0] if len(up_tbls) == 1 else __import__("pyarrow").concat_tables(up_tbls)
+            ob_up = _arrow_rg_to_polars(up_tbl).sort("timestamp_us")
+        else:
+            ob_up = pl.DataFrame()
+        if dn_tbls:
+            dn_tbl = dn_tbls[0] if len(dn_tbls) == 1 else __import__("pyarrow").concat_tables(dn_tbls)
+            ob_dn = _arrow_rg_to_polars(dn_tbl).sort("timestamp_us")
+        else:
+            ob_dn = pl.DataFrame()
+        return slug, ob_up, ob_dn
+
+    for rg_idx in range(pf.num_row_groups):
+        rg_tbl = pf.read_row_group(rg_idx, columns=cols)
+        if rg_tbl.num_rows == 0:
+            continue
+        slug_val = rg_tbl.column("slug")[0].as_py()
+        outcome_val = rg_tbl.column("outcome")[0].as_py()
+
+        if current_slug is None:
+            current_slug = slug_val
+
+        if slug_val != current_slug:
+            res = _emit(current_slug, ob_up_tbls, ob_dn_tbls)
+            if res is not None:
+                yield res
+            ob_up_tbls = []
+            ob_dn_tbls = []
+            current_slug = slug_val
+
+        if outcome_val == "Up":
+            ob_up_tbls.append(rg_tbl)
+        elif outcome_val == "Down":
+            ob_dn_tbls.append(rg_tbl)
+
+    if current_slug is not None:
+        res = _emit(current_slug, ob_up_tbls, ob_dn_tbls)
+        if res is not None:
+            yield res
+
+
 def load_orderbook_filtered(
     asset: str,
     hf_token: str,

train.py

@@ -123,7 +123,7 @@ def _build_training_dataset(
     import gc
     import polars as pl  # local import to keep module import-light
 
-    from data_loader import iter_orderbook_batches
+    from data_loader import iter_orderbook_slug_pairs
 
     log(f"[data] loading markets_index for {asset}")
     markets = load_markets_index(asset, hf_token, cache_dir)
@@ -135,64 +135,59 @@ def _build_training_dataset(
     slugs = markets["slug"].to_list()
     slug_ts_list = markets["slug_ts"].to_list()
     slug_ts_map = dict(zip(slugs, [int(t) for t in slug_ts_list]))
+    wanted = set(slugs)
 
-    log(f"[data] streaming book_snapshot_5 in batches of {ob_batch_size} "
-        f"(~{(len(slugs) + ob_batch_size - 1) // ob_batch_size} batches)")
+    log(f"[data] streaming book_snapshot_5 row-groups (~{len(slugs)} slugs, "
+        f"peak ~5 MB per slug)")
 
     rows: List[Dict] = []
     built = 0
     skipped = 0
+    processed = 0
 
-    for batch_idx, (ob_batch, batch_slugs) in enumerate(
-        iter_orderbook_batches(asset, hf_token, cache_dir, slugs, batch_size=ob_batch_size)
+    for slug, ob_up, ob_dn in iter_orderbook_slug_pairs(
+        asset, hf_token, cache_dir, wanted
     ):
-        ob_by_slug_up: Dict[str, pl.DataFrame] = {}
-        ob_by_slug_dn: Dict[str, pl.DataFrame] = {}
-        if len(ob_batch) > 0:
-            for (slug_val, outcome_val), sub in ob_batch.group_by(["slug", "outcome"]):
-                sub_sorted = sub.sort("timestamp_us")
-                if outcome_val == "Up":
-                    ob_by_slug_up[slug_val] = sub_sorted
-                elif outcome_val == "Down":
-                    ob_by_slug_dn[slug_val] = sub_sorted
-
-        for slug in batch_slugs:
-            slug_ts = slug_ts_map[slug]
-            try:
-                spot = get_window_label(slug_ts, ohlcv)
-                if spot is None:
-                    skipped += 1
-                    continue
-                ob_up = ob_by_slug_up.get(slug, pl.DataFrame())
-                ob_dn = ob_by_slug_dn.get(slug, pl.DataFrame())
-                wf = build_window_frame(slug, slug_ts, ob_up, ob_dn, ohlcv)
-                feats = _extract_all(wf, at_tick=120)
-                arb = compute_optimal_arb(wf, spot_label=spot)
-                row = {
-                    "slug": slug,
-                    "slug_ts": slug_ts,
-                    "spot_label": spot,
-                    "optimal_pnl": arb["optimal_pnl"],
-                    "binary_label": arb["binary_label"],
-                }
-                for name, v in zip(ALL_FEATURES, feats):
-                    row[name] = float(v)
-                rows.append(row)
-                built += 1
-            except Exception as e:  # noqa: BLE001
+        processed += 1
+        slug_ts = slug_ts_map.get(slug)
+        if slug_ts is None:
+            continue
+        try:
+            spot = get_window_label(slug_ts, ohlcv)
+            if spot is None:
                 skipped += 1
-                if skipped <= 3:
-                    import traceback
-                    log(f"[data] window error slug={slug}: {e!r}\n"
-                        f"{traceback.format_exc()}")
-                else:
-                    log(f"[data] window error slug={slug}: {e!r}")
                 continue
-
-        # free polars memory before advancing to next batch
-        del ob_batch, ob_by_slug_up, ob_by_slug_dn
-        gc.collect()
-        log(f"[data] batch {batch_idx + 1} done; built={built} skipped={skipped}")
+            wf = build_window_frame(slug, slug_ts, ob_up, ob_dn, ohlcv)
+            feats = _extract_all(wf, at_tick=120)
+            arb = compute_optimal_arb(wf, spot_label=spot)
+            row = {
+                "slug": slug,
+                "slug_ts": slug_ts,
+                "spot_label": spot,
+                "optimal_pnl": arb["optimal_pnl"],
+                "binary_label": arb["binary_label"],
+            }
+            for name, v in zip(ALL_FEATURES, feats):
+                row[name] = float(v)
+            rows.append(row)
+            built += 1
+        except Exception as e:  # noqa: BLE001
+            skipped += 1
+            if skipped <= 3:
+                import traceback
+                log(f"[data] window error slug={slug}: {e!r}\n"
+                    f"{traceback.format_exc()}")
+            else:
+                log(f"[data] window error slug={slug}: {e!r}")
+
+        # free the row-group tables before the next one
+        del ob_up, ob_dn
+        if processed % 1000 == 0:
+            gc.collect()
+            log(f"[data] processed={processed} built={built} skipped={skipped}")
+
+    gc.collect()
+    log(f"[data] done streaming. processed={processed} built={built} skipped={skipped}")
 
     df = pd.DataFrame(rows)
     if len(df) == 0: