Add minute bar feature pipeline

pipeline/data/downloader.py

@@ -11,9 +11,13 @@ import pyarrow.dataset as pads
 
 # Reuse existing downloader and universe modules
 sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
-from data.downloader import download_daily_batch
+from data.downloader import (
+    _normalize_minute_frequency,
+    download_daily_batch,
+    download_minute_batch,
+)
 from data.universe import get_all_stocks, get_hs300_stocks, get_zz500_stocks
-from pipeline.common.schema import DATA_COLUMNS
+from pipeline.common.schema import DATA_COLUMNS, MINUTE_BAR_COLUMNS
 
 logger = logging.getLogger(__name__)
 
@@ -89,6 +93,25 @@ def _write_month_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: s
     )
 
 
+def _write_minute_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: str) -> None:
+    """Append rows to a Hive-partitioned minute dataset.
+
+    Layout: ``frequency=5m/month=YYYY-MM/*.pq``.
+    """
+    out = df.copy()
+    out["month"] = pd.to_datetime(out["date"]).dt.strftime("%Y-%m")
+    table = pa.Table.from_pandas(out, preserve_index=False)
+    pads.write_dataset(
+        table,
+        str(base_dir),
+        format="parquet",
+        partitioning=["frequency", "month"],
+        partitioning_flavor="hive",
+        basename_template=f"{basename_prefix}-{{i}}.pq",
+        existing_data_behavior="overwrite_or_ignore",
+    )
+
+
 def download_universe(
     universe: str = "csi500",
     start_date: str = "2017-01-01",
@@ -177,3 +200,104 @@ def download_universe(
         "date_min": None if date_min is None else str(pd.Timestamp(date_min).date()),
         "date_max": None if date_max is None else str(pd.Timestamp(date_max).date()),
     }
+
+
+def download_minute_universe(
+    universe: str = "csi500",
+    start_date: str = "2017-01-01",
+    end_date: str = "2026-12-31",
+    output_dir: str = "data/minute_bars",
+    max_symbols: int = 0,
+    chunk_size: int = 100,
+    frequency: str | int = 5,
+) -> dict:
+    """Download raw minute bars into a frequency/month-partitioned dataset.
+
+    Args:
+        universe: ``hs300``, ``csi500``, ``all``/``full``, a file path, or a
+            comma-separated symbol list.
+        start_date, end_date: ``YYYY-MM-DD`` bounds.
+        output_dir: Root under which ``{universe}/frequency=5m/month=YYYY-MM``
+            is written.
+        max_symbols: Cap on symbols (0 = all).
+        chunk_size: Symbols per durability flush.
+        frequency: Minute interval. ``5``/``"5"``/``"5m"`` are 5-minute bars.
+
+    Returns:
+        Stats dict with dataset path, row count, symbol count, date range, and
+        frequency label.
+    """
+    _, frequency_label = _normalize_minute_frequency(frequency)
+    constituents = _resolve_universe(universe, max_symbols)
+    symbols = constituents["symbol_id"].tolist()
+    names = dict(zip(constituents["symbol_id"], constituents["symbol_name"]))
+    n_requested = len(symbols)
+    logger.info(
+        "Minute universe %s: %d symbols, %s → %s, frequency=%s",
+        universe,
+        n_requested,
+        start_date,
+        end_date,
+        frequency,
+    )
+
+    base_dir = Path(output_dir) / universe
+    target_frequency_dir = base_dir / f"frequency={frequency_label}"
+    if target_frequency_dir.exists():
+        shutil.rmtree(target_frequency_dir)
+    base_dir.mkdir(parents=True, exist_ok=True)
+
+    buffer: list[pd.DataFrame] = []
+    chunk_idx = 0
+    succeeded = 0
+    n_rows = 0
+    date_min = None
+    date_max = None
+
+    def flush() -> None:
+        nonlocal buffer, chunk_idx, n_rows, date_min, date_max
+        if not buffer:
+            return
+        chunk = pd.concat(buffer, ignore_index=True)
+        _write_minute_partitions(chunk, base_dir, basename_prefix=f"chunk{chunk_idx:04d}")
+        n_rows += len(chunk)
+        cmin, cmax = chunk["date"].min(), chunk["date"].max()
+        date_min = cmin if date_min is None else min(date_min, cmin)
+        date_max = cmax if date_max is None else max(date_max, cmax)
+        logger.info(
+            "Flushed minute chunk %d: %d rows (%d symbols done)",
+            chunk_idx,
+            len(chunk),
+            succeeded,
+        )
+        buffer = []
+        chunk_idx += 1
+
+    for i, (symbol, df) in enumerate(
+        download_minute_batch(symbols, start_date, end_date, frequency=frequency), start=1
+    ):
+        if df is None:
+            logger.warning("  %s: no minute data", symbol)
+        else:
+            df["symbol_id"] = symbol
+            df["symbol_name"] = names.get(symbol, symbol)
+            buffer.append(df[MINUTE_BAR_COLUMNS])
+            succeeded += 1
+        if len(buffer) >= chunk_size:
+            flush()
+        if i % 100 == 0:
+            logger.info("Minute progress: %d/%d symbols", i, n_requested)
+    flush()
+
+    if succeeded == 0:
+        raise RuntimeError("No minute data downloaded for any symbol")
+
+    return {
+        "dataset_path": str(base_dir),
+        "frequency": frequency_label,
+        "n_symbols": succeeded,
+        "n_requested": n_requested,
+        "n_rows": n_rows,
+        "date_min": None if date_min is None else str(pd.Timestamp(date_min).date()),
+        "date_max": None if date_max is None else str(pd.Timestamp(date_max).date()),
+    }