chinese-equity-quant/pipeline/data/downloader.py

"""Download daily bar data for a universe and save as a partitioned parquet dataset."""

import logging
import shutil
import sys
from pathlib import Path

import pandas as pd
import pyarrow as pa
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.universe import get_all_stocks, get_hs300_stocks, get_zz500_stocks
from pipeline.common.schema import DATA_COLUMNS

logger = logging.getLogger(__name__)


def _fix_baostock_columns(df: pd.DataFrame) -> pd.DataFrame:
    """baostock constituent queries return (update_date, code, name) — 
    detect columns by value patterns rather than assuming column order."""
    cols = df.columns.tolist()
    result = {}
    for col in cols:
        vals = df[col].astype(str)
        # Stock code: matches sh.NNNNNN or sz.NNNNNN (possibly with dot)
        if vals.str.match(r"^(sh|sz)\.?\d{6}$").all():
            result["symbol_id"] = df[col].str.replace(".", "", regex=False)
        # Stock name: Chinese characters (detected by byte length > str length)
        elif vals.apply(lambda x: len(x.encode("utf-8")) > len(x)).any() and vals.str.len().max() < 10:
            result["symbol_name"] = df[col]
        # Skip date column
    return pd.DataFrame(result)


def _resolve_universe(universe: str, max_symbols: int = 0) -> pd.DataFrame:
    """Resolve a universe name or file path to symbol list with names.

    Returns DataFrame with columns: code (symbol_id), name (symbol_name).
    """
    name = universe.lower()
    if name == "hs300":
        df = get_hs300_stocks()
        # baostock returns (date, code, name) — detect columns by value patterns
        df = _fix_baostock_columns(df)
    elif name == "csi500":
        df = get_zz500_stocks()
        df = _fix_baostock_columns(df)
    elif name in ("all", "full"):
        # Every listed A-share (~5000); already (code, name) with prefixed codes.
        all_df = get_all_stocks()
        df = all_df.rename(columns={"code": "symbol_id", "name": "symbol_name"})
    elif Path(universe).exists():
        # File with one symbol_id per line
        with open(universe) as f:
            symbols = [line.strip() for line in f if line.strip()]
        df = pd.DataFrame({"symbol_id": symbols, "symbol_name": symbols})
    else:
        # Assume comma-separated list
        symbols = [s.strip() for s in universe.split(",") if s.strip()]
        df = pd.DataFrame({"symbol_id": symbols, "symbol_name": symbols})

    if max_symbols and max_symbols > 0 and len(df) > max_symbols:
        df = df.head(max_symbols).copy()

    return df


def _write_month_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: str) -> None:
    """Append rows to a Hive-partitioned (month=YYYY-MM) parquet dataset.

    ``existing_data_behavior='overwrite_or_ignore'`` plus a per-chunk
    ``basename_prefix`` means each flush adds new ``.pq`` files into the month
    directories without deleting earlier chunks' files.
    """
    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=["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",
    end_date: str = "2026-12-31",
    output_dir: str = "data/daily_bars",
    max_symbols: int = 0,
    chunk_size: int = 300,
    adjust: str = "qfq",
) -> dict:
    """Download a universe's daily bars into a month-partitioned parquet dataset.

    Streams downloads under a single baostock session and flushes every
    ``chunk_size`` symbols, so memory stays bounded and a crash keeps the
    partitions already written. The dataset is rebuilt from scratch: any
    existing ``output_dir/{universe}`` directory is removed first.

    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}/month=YYYY-MM/*.pq`` is written.
        max_symbols: Cap on symbols (0 = all).
        chunk_size: Symbols per durability flush.
        adjust: ``qfq`` / ``hfq`` / ``''``.

    Returns:
        Stats dict: ``dataset_path``, ``n_symbols`` (succeeded), ``n_requested``,
        ``n_rows``, ``date_min``, ``date_max``.
    """
    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("Universe %s: %d symbols, %s → %s", universe, n_requested, start_date, end_date)

    base_dir = Path(output_dir) / universe
    if base_dir.exists():
        shutil.rmtree(base_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_month_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 chunk %d: %d rows (%d symbols done)", chunk_idx, len(chunk), succeeded)
        buffer = []
        chunk_idx += 1

    for i, (symbol, df) in enumerate(
        download_daily_batch(symbols, start_date, end_date, adjust=adjust), start=1
    ):
        if df is None:
            logger.warning("  %s: no data", symbol)
        else:
            df["symbol_id"] = symbol
            df["symbol_name"] = names.get(symbol, symbol)
            buffer.append(df[DATA_COLUMNS])
            succeeded += 1
        if len(buffer) >= chunk_size:
            flush()
        if i % 100 == 0:
            logger.info("Progress: %d/%d symbols", i, n_requested)
    flush()

    if succeeded == 0:
        raise RuntimeError("No data downloaded for any symbol")

    return {
        "dataset_path": str(base_dir),
        "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()),
    }