Add minute bar feature pipeline

2026-06-16 13:57:17 +08:00
parent 17fa75495d
commit 83a006bbe4
19 changed files with 1289 additions and 11 deletions
@@ -3,7 +3,7 @@
 import click
 from datetime import date

-from pipeline.data.downloader import download_universe
+from pipeline.data.downloader import download_minute_universe, download_universe


@click.group(name="data")
@@ -42,3 +42,37 @@ def download(universe, start_date, end_date, output_dir, symbols, chunk_size, ad
        f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}"
    )
    click.echo(f"Dataset: {stats['dataset_path']}")
+
+
+@data.command("download-minute")
+@click.option(
+    "--universe", default="csi500",
+    help="Which universe: hs300, csi500, all (~5000 A-shares), file path, or comma-separated symbols",
+)
+@click.option("--start-date", default="2017-01-01", help="Start date YYYY-MM-DD")
+@click.option("--end-date", default=str(date.today()), help="End date YYYY-MM-DD")
+@click.option("--output-dir", default="data/minute_bars", help="Root for the partitioned dataset")
+@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
+@click.option("--chunk-size", default=100, type=int, help="Symbols per durability flush")
+@click.option("--frequency", default="5", help="Minute frequency: 5, 15, 30, or 60")
+def download_minute(universe, start_date, end_date, output_dir, symbols, chunk_size, frequency):
+    """Download raw Baostock minute bars into a partitioned parquet dataset.
+
+    Writes ``{output_dir}/{universe}/frequency=5m/month=YYYY-MM/*.pq`` for the
+    default 5-minute frequency.
+    """
+    stats = download_minute_universe(
+        universe=universe,
+        start_date=start_date,
+        end_date=end_date,
+        output_dir=output_dir,
+        max_symbols=symbols,
+        chunk_size=chunk_size,
+        frequency=frequency,
+    )
+    click.echo(
+        f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
+        f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}, "
+        f"frequency={stats['frequency']}"
+    )
+    click.echo(f"Dataset: {stats['dataset_path']}")
@@ -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()),
+    }