From 83a006bbe42528fa81609ffc206148248755a19d Mon Sep 17 00:00:00 2001 From: Yuxuan Yan Date: Tue, 16 Jun 2026 13:57:17 +0800 Subject: [PATCH] Add minute bar feature pipeline --- cli.py | 3 + data/downloader.py | 158 +++++++++++++ docs/minute_bar_data.md | 67 ++++++ pipeline/alpha/base.py | 24 +- pipeline/alpha/cli.py | 7 +- pipeline/alpha/compute.py | 53 ++++- pipeline/common/schema.py | 18 ++ pipeline/data/cli.py | 36 ++- pipeline/data/downloader.py | 128 ++++++++++- pipeline/features/__init__.py | 1 + pipeline/features/base.py | 34 +++ pipeline/features/cli.py | 108 +++++++++ pipeline/features/compute.py | 75 ++++++ pipeline/features/library/__init__.py | 3 + .../features/library/minute_daily_summary.py | 84 +++++++ pipeline/features/registry.py | 79 +++++++ tests/test_alpha.py | 66 ++++++ tests/test_features.py | 141 ++++++++++++ tests/test_minute_downloader.py | 215 ++++++++++++++++++ 19 files changed, 1289 insertions(+), 11 deletions(-) create mode 100644 docs/minute_bar_data.md create mode 100644 pipeline/features/__init__.py create mode 100644 pipeline/features/base.py create mode 100644 pipeline/features/cli.py create mode 100644 pipeline/features/compute.py create mode 100644 pipeline/features/library/__init__.py create mode 100644 pipeline/features/library/minute_daily_summary.py create mode 100644 pipeline/features/registry.py create mode 100644 tests/test_features.py create mode 100644 tests/test_minute_downloader.py diff --git a/cli.py b/cli.py index 678a7e8..0078e34 100644 --- a/cli.py +++ b/cli.py @@ -4,6 +4,7 @@ Phases: data — Download daily bars to parquet alpha — Compute alpha weights from data + feature — Compute daily features from minute bars combo — Combine alphas into a single weight portfolio — Build tradable positions and simulate execution """ @@ -14,6 +15,7 @@ import click from pipeline.data.cli import data from pipeline.alpha.cli import alpha +from pipeline.features.cli import feature from pipeline.combo.cli import combo from pipeline.portfolio.cli import portfolio from tools.pqcat import pqcat @@ -40,6 +42,7 @@ def cli(log_level): cli.add_command(data) cli.add_command(alpha) +cli.add_command(feature) cli.add_command(combo) cli.add_command(portfolio) cli.add_command(pqcat) diff --git a/data/downloader.py b/data/downloader.py index 29d2786..5a0cea8 100644 --- a/data/downloader.py +++ b/data/downloader.py @@ -31,11 +31,68 @@ _BATCH_COLUMNS = [ "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM", ] +# Raw Baostock minute bars. The ``time`` field is usually compact +# YYYYMMDDHHMMSSmmm; parsing below also tolerates HH:MM:SS strings in tests. +_MINUTE_FIELDS = "date,time,code,open,high,low,close,volume,amount,adjustflag" +_MINUTE_NUMERIC = ["open", "high", "low", "close", "volume", "amount"] +_MINUTE_COLUMNS = [ + "symbol", "datetime", "date", "time", "frequency", + "open", "high", "low", "close", "volume", "amount", "vwap", "adjustflag", +] +_MINUTE_FREQUENCIES = {"5", "15", "30", "60"} + class _SessionLost(Exception): """baostock reported the session was dropped (``用户未登录``).""" +def _normalize_minute_frequency(frequency: str | int) -> tuple[str, str]: + """Return Baostock frequency and partition label for a minute interval.""" + raw = str(frequency).strip().lower() + if raw.endswith("m"): + raw = raw[:-1] + if raw not in _MINUTE_FREQUENCIES: + raise ValueError( + f"Unsupported minute frequency '{frequency}'. " + f"Expected one of {sorted(_MINUTE_FREQUENCIES)} minutes." + ) + return raw, f"{raw}m" + + +def _parse_minute_datetime(date: pd.Series, time: pd.Series) -> pd.Series: + """Parse Baostock minute timestamps into pandas datetimes.""" + date_dt = pd.to_datetime(date, errors="coerce") + date_compact = date_dt.dt.strftime("%Y%m%d") + time_text = time.astype(str).str.strip() + time_digits = time_text.str.replace(r"\D", "", regex=True) + + full_digits = time_digits.str.slice(0, 14) + from_full = pd.to_datetime(full_digits, format="%Y%m%d%H%M%S", errors="coerce") + + from_short = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]") + short_time = time_digits.str.len().between(1, 6) + if short_time.any(): + short_digits = ( + time_digits.loc[short_time] + .str.pad(6, side="right", fillchar="0") + .str.slice(0, 6) + ) + from_short.loc[short_time] = pd.to_datetime( + date_compact.loc[short_time] + short_digits, + format="%Y%m%d%H%M%S", + errors="coerce", + ) + + from_text = pd.Series(pd.NaT, index=date.index, dtype="datetime64[ns]") + text_time = time_text.str.contains(":", regex=False) + if text_time.any(): + from_text.loc[text_time] = pd.to_datetime( + date.astype(str).loc[text_time] + " " + time_text.loc[text_time], + errors="coerce", + ) + return from_full.fillna(from_short).fillna(from_text) + + def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]: """Download daily bars from akshare. Returns DataFrame with OHLCV columns.""" try: @@ -239,3 +296,104 @@ def download_daily_batch( except Exception: pass + +def download_minute_batch( + symbols: Iterable[str], + start: str, + end: str, + frequency: str | int = 5, + relogin_every: int = 200, +) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]: + """Download raw Baostock minute bars for many symbols. + + Minute bars are intentionally unadjusted (`adjustflag='3'`) because the + output is raw intraday market data for downstream feature aggregation, not a + tradable daily price series. + + Args: + symbols: Internal-form symbols (``sh600000`` / ``sz000001``). + start, end: ``YYYY-MM-DD`` bounds. + frequency: Baostock minute frequency. ``5``/``"5"``/``"5m"`` all mean + 5-minute bars. + relogin_every: Proactively refresh the baostock session every N symbols. + + Yields: + ``(symbol, df)`` where ``df`` has raw minute bars or ``None`` when no + data is available. + """ + query_frequency, frequency_label = _normalize_minute_frequency(frequency) + adjustflag = _BAOSTOCK_ADJUST["none"] + + def _relogin() -> None: + try: + bs.logout() + except Exception: + pass + bs.login() + + def _fetch(symbol: str) -> Optional[pd.DataFrame]: + """One Baostock minute query; returns df, None, or raises _SessionLost.""" + code = f"{symbol[:2]}.{symbol[2:]}" + rs = bs.query_history_k_data_plus( + code=code, + fields=_MINUTE_FIELDS, + start_date=start, + end_date=end, + frequency=query_frequency, + adjustflag=adjustflag, + ) + if rs.error_code != "0": + if "未登录" in (rs.error_msg or ""): + raise _SessionLost(rs.error_msg) + logger.warning("baostock minute error for %s: %s", symbol, rs.error_msg) + return None + rows = [] + while rs.next(): + rows.append(rs.get_row_data()) + if not rows: + return None + + df = pd.DataFrame(rows, columns=_MINUTE_FIELDS.split(",")) + df[_MINUTE_NUMERIC] = df[_MINUTE_NUMERIC].apply(pd.to_numeric, errors="coerce") + df["datetime"] = _parse_minute_datetime(df["date"], df["time"]) + bad_timestamps = df["datetime"].isna() + if bad_timestamps.any(): + raise ValueError( + f"Could not parse {int(bad_timestamps.sum())} minute timestamp(s)" + ) + df["date"] = pd.to_datetime(df["date"], errors="coerce") + df["time"] = df["datetime"].dt.strftime("%H:%M:%S") + df["frequency"] = frequency_label + df["vwap"] = (df["amount"] / df["volume"]).where(df["volume"] > 0) + df["symbol"] = symbol + return df[_MINUTE_COLUMNS].sort_values("datetime").reset_index(drop=True) + + bs.login() + try: + for i, symbol in enumerate(symbols): + if i and relogin_every and i % relogin_every == 0: + _relogin() + + df: Optional[pd.DataFrame] = None + for attempt in (1, 2): + try: + df = _fetch(symbol) + break + except _SessionLost: + if attempt == 1: + _relogin() + continue + logger.warning("baostock minute session lost for %s after relogin", symbol) + except Exception as e: + logger.warning("baostock minute download failed for %s: %s", symbol, e) + break + + if df is not None and not df.empty: + yield symbol, df + else: + yield symbol, None + finally: + try: + bs.logout() + except Exception: + pass diff --git a/docs/minute_bar_data.md b/docs/minute_bar_data.md new file mode 100644 index 0000000..ba1a139 --- /dev/null +++ b/docs/minute_bar_data.md @@ -0,0 +1,67 @@ +# Minute Bar Data Notes + +The minute-bar path downloads raw Baostock intraday bars and stores them as a +Hive-partitioned dataset: + +```bash +uv run python cli.py data download-minute \ + --universe sh600000 \ + --start-date 2024-01-02 --end-date 2024-01-05 \ + --frequency 5 +``` + +The default layout is: + +```text +data/minute_bars/{universe}/frequency=5m/month=YYYY-MM/*.pq +``` + +Feature plugins can aggregate those bars to daily `symbol_id,date` feature +files, for example: + +```bash +uv run python cli.py feature compute \ + --minute-path data/minute_bars/sh600000 \ + --daily-path data/daily_bars/sh600000 \ + --feature-type minute_daily_summary \ + --feature-name minute_summary +``` + +## Daily vs Minute Reconciliation + +Baostock's daily raw bars and 5-minute raw bars are close, but they should not +be treated as perfectly reconstructible from each other. + +When checking consistency, compare daily raw bars (`data download --adjust none`) +against minute bars on the same raw price scale. The minute aggregation should +use: + +- `open`: first minute open +- `high`: max minute high +- `low`: min minute low +- `close`: last minute close +- `volume`: sum minute volume +- `amount`: sum minute amount +- `vwap`: `sum(amount) / sum(volume)` + +In a sanity check for `sh600000` from `2024-01-02` through `2024-01-05`, Baostock +returned 4 daily rows and 192 5-minute bars, exactly 48 bars per day. Open, low, +and close matched daily exactly on all 4 days. High matched on 3 of 4 days; on +`2024-01-04`, the daily high was `6.67` while the max 5-minute high was `6.66`. +Minute-summed volume and amount were higher than daily by roughly `0.16%` to +`1.23%`. VWAP remained very close, with max relative difference around +`0.0043%`. + +This appears to be a source-level Baostock reconciliation caveat, not a parser +or ordering issue: the minute bars covered the regular `09:35:00` through +`15:00:00` range and sorted correctly by timestamp. + +Practical guidance: + +- Use tolerance-based daily-vs-minute checks; do not require exact equality for + high, volume, or amount. +- Expect open/close alignment to be a stronger sanity check than exact volume + reconstruction. +- Use minute-derived values as separate daily features, not as replacements for + the canonical daily bar dataset unless a strategy explicitly wants that + source convention. diff --git a/pipeline/alpha/base.py b/pipeline/alpha/base.py index 9115d40..2d140eb 100644 --- a/pipeline/alpha/base.py +++ b/pipeline/alpha/base.py @@ -5,7 +5,7 @@ position weights. Subclasses implement :meth:`signal` — the raw, unnormalized score. The base class turns a signal into cross-sectionally z-scored weights via :meth:`to_weights` (override it for a different normalization). """ -from abc import ABC, abstractmethod +from abc import ABC import numpy as np import pandas as pd @@ -15,15 +15,14 @@ class BaseAlpha(ABC): """A position-weight alpha over a cross-section of stocks. Concrete subclasses must set a unique class-level :attr:`name` (the registry - key) and implement :meth:`signal`. Construct subclasses with their own typed - parameters (e.g. ``lookback``); the factory passes only the parameters a - given ``__init__`` accepts. + key) and implement either :meth:`signal` or :meth:`signal_from_data`. + Construct subclasses with their own typed parameters (e.g. ``lookback``); + the factory passes only the parameters a given ``__init__`` accepts. """ #: Unique registry key. Every concrete alpha must set this to a non-empty str. name: str = "" - @abstractmethod def signal(self, close: pd.DataFrame) -> pd.DataFrame: """Compute the raw signal. @@ -34,6 +33,21 @@ class BaseAlpha(ABC): A wide DataFrame aligned to ``close`` where higher values indicate a stronger long. Use NaN where the signal is undefined. """ + raise NotImplementedError( + f"{type(self).__name__} must implement signal() or signal_from_data()" + ) + + def signal_from_data( + self, + data: pd.DataFrame, + close: pd.DataFrame, + ) -> pd.DataFrame: + """Compute the raw signal from long daily data plus wide closes. + + Feature-aware alphas can override this to pivot joined feature columns + from ``data``. The default preserves the existing close-only alpha API. + """ + return self.signal(close) def to_weights(self, signal: pd.DataFrame) -> pd.DataFrame: """Cross-sectionally z-score a signal into signed position weights. diff --git a/pipeline/alpha/cli.py b/pipeline/alpha/cli.py index 8702d96..c01b54f 100644 --- a/pipeline/alpha/cli.py +++ b/pipeline/alpha/cli.py @@ -56,6 +56,10 @@ def list_(alpha_modules): @click.option("--output-dir", default="alphas", help="Directory to save alpha parquet") @click.option("--lookback", default=5, type=int, help="Lookback days") @click.option("--vol-window", default=20, type=int, help="Volatility window (reversal_vol only)") +@click.option( + "--feature-path", "feature_paths", multiple=True, + help="Daily feature parquet file/dataset to left-join on symbol_id,date (repeatable)", +) @click.option( "--alpha-module", "alpha_modules", multiple=True, help="External module(s) to import so their alphas register (dotted path or .py file)", @@ -74,7 +78,7 @@ def list_(alpha_modules): help="Most-liquid names kept per date when --liquid-universe is set", ) def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window, - alpha_modules, extra_params, liquid_universe, universe_top_n): + feature_paths, alpha_modules, extra_params, liquid_universe, universe_top_n): """Compute one alpha from raw data and save as parquet.""" for spec in alpha_modules: load_alpha_module(spec) @@ -100,6 +104,7 @@ def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window, alpha_name=alpha_name, alpha_type=alpha_type, universe=universe, + feature_paths=feature_paths, **params, ) diff --git a/pipeline/alpha/compute.py b/pipeline/alpha/compute.py index 9f238cc..857a9cd 100644 --- a/pipeline/alpha/compute.py +++ b/pipeline/alpha/compute.py @@ -7,12 +7,15 @@ through :mod:`pipeline.alpha.registry`. """ import logging +from pathlib import Path +from typing import Iterable import numpy as np import pandas as pd from pipeline.alpha.registry import get_alpha from pipeline.common.schema import ALPHA_COLUMNS +from pipeline.features.compute import FEATURE_KEY_COLUMNS, read_feature_frames, validate_feature_frame logger = logging.getLogger(__name__) @@ -33,6 +36,38 @@ def _pivot_open(df: pd.DataFrame) -> pd.DataFrame: return pivot.sort_index() +def join_feature_frames( + data: pd.DataFrame, + feature_frames: Iterable[pd.DataFrame], +) -> pd.DataFrame: + """Left-join validated daily feature frames onto long daily data.""" + out = data.copy() + out["date"] = pd.to_datetime(out["date"]) + existing = set(out.columns) + joined_cols: list[str] = [] + + for frame in feature_frames: + features = validate_feature_frame(frame) + feature_cols = [col for col in features.columns if col not in FEATURE_KEY_COLUMNS] + overlap = sorted(existing.intersection(feature_cols)) + if overlap: + raise ValueError( + f"Feature columns conflict with existing daily data columns: {overlap}" + ) + out = out.merge( + features, + on=FEATURE_KEY_COLUMNS, + how="left", + validate="many_to_one", + ) + existing.update(feature_cols) + joined_cols.extend(feature_cols) + + if joined_cols: + logger.info("Joined feature columns into daily data: %s", joined_cols) + return out + + def _forward_open_to_open_returns(open_: pd.DataFrame) -> pd.DataFrame: """Return earned by a close-formed signal after next-open execution. @@ -105,6 +140,8 @@ def compute_alpha( alpha_name: str, alpha_type: str, universe: dict | None = None, + feature_paths: Iterable[str | Path] | None = None, + feature_frames: Iterable[pd.DataFrame] | None = None, **params, ) -> pd.DataFrame: """Compute alpha weights from raw data. @@ -118,6 +155,10 @@ def compute_alpha( :func:`investable_universe_mask`) *before* it is turned into weights, so unheld names get weight 0. Keys are forwarded as keyword arguments to :func:`investable_universe_mask`. + feature_paths: Optional parquet files/datasets keyed by ``symbol_id`` + and ``date``. Their numeric feature columns are left-joined onto + ``data`` before alpha logic runs. + feature_frames: Optional in-memory feature frames with the same schema. **params: Constructor parameters for the alpha (e.g. ``lookback``, ``vol_window``). Only the params the alpha's ``__init__`` accepts are used; extras are ignored. @@ -128,12 +169,20 @@ def compute_alpha( Raises: KeyError: If ``alpha_type`` is not registered. """ + feature_inputs: list[pd.DataFrame] = [] + if feature_paths: + feature_inputs.extend(read_feature_frames(feature_paths)) + if feature_frames: + feature_inputs.extend(feature_frames) + if feature_inputs: + data = join_feature_frames(data, feature_inputs) + alpha = get_alpha(alpha_type, **params) close = _pivot_close(data) + signal = alpha.signal_from_data(data, close) if universe is None: - weights = alpha.weights(close) + weights = alpha.to_weights(signal) else: - signal = alpha.signal(close) mask = investable_universe_mask(data, signal, **universe) weights = alpha.to_weights(signal.where(mask)) diff --git a/pipeline/common/schema.py b/pipeline/common/schema.py index da27ce8..4857179 100644 --- a/pipeline/common/schema.py +++ b/pipeline/common/schema.py @@ -26,6 +26,24 @@ DATA_COLUMNS: Final[list[str]] = [ "pcfNcfTTM", # float64: P/CF (net cash flow, TTM) ] +# Required columns for raw intraday minute bar parquet files. +MINUTE_BAR_COLUMNS: Final[list[str]] = [ + "symbol_id", # str: internal code like 'sh600000' + "symbol_name", # str: stock name like '浦发银行' + "datetime", # datetime64: intraday bar timestamp + "date", # date component, aligned with daily DATA_COLUMNS date + "time", # str: HH:MM:SS bar time + "frequency", # str: e.g. '5m' + "open", # float64 + "high", # float64 + "low", # float64 + "close", # float64 + "volume", # float64 (shares) + "amount", # float64 (turnover in yuan, raw/unadjusted) + "vwap", # float64: amount / volume + "adjustflag", # str: baostock adjustment flag; '3' for raw/unadjusted +] + # Required columns for alpha parquet files. # Alphas are position WEIGHTS: positive=long, negative=short. ALPHA_COLUMNS: Final[list[str]] = [ diff --git a/pipeline/data/cli.py b/pipeline/data/cli.py index 4d9e41f..c71825e 100644 --- a/pipeline/data/cli.py +++ b/pipeline/data/cli.py @@ -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']}") diff --git a/pipeline/data/downloader.py b/pipeline/data/downloader.py index 9eff224..b7e996f 100644 --- a/pipeline/data/downloader.py +++ b/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()), + } diff --git a/pipeline/features/__init__.py b/pipeline/features/__init__.py new file mode 100644 index 0000000..2e5861a --- /dev/null +++ b/pipeline/features/__init__.py @@ -0,0 +1 @@ +"""Daily feature plugin package.""" diff --git a/pipeline/features/base.py b/pipeline/features/base.py new file mode 100644 index 0000000..55a8697 --- /dev/null +++ b/pipeline/features/base.py @@ -0,0 +1,34 @@ +"""Base class for daily feature plugins.""" + +from abc import ABC, abstractmethod + +import pandas as pd + + +class BaseFeature(ABC): + """Aggregate raw minute bars into daily, symbol-keyed feature columns.""" + + #: Unique registry key. Every concrete feature must set this to a non-empty str. + name: str = "" + + @abstractmethod + def compute( + self, + minute: pd.DataFrame, + daily: pd.DataFrame | None = None, + ) -> pd.DataFrame: + """Compute daily features. + + Args: + minute: Raw minute bars with ``symbol_id`` and ``date`` keys. + daily: Optional daily data frame for calendar alignment or + reference daily columns. + + Returns: + DataFrame with ``symbol_id``, ``date``, and one or more numeric + feature columns. + """ + + def __repr__(self) -> str: + params = ", ".join(f"{k}={v!r}" for k, v in vars(self).items()) + return f"{type(self).__name__}({params})" diff --git a/pipeline/features/cli.py b/pipeline/features/cli.py new file mode 100644 index 0000000..aba382e --- /dev/null +++ b/pipeline/features/cli.py @@ -0,0 +1,108 @@ +"""CLI for daily feature computation.""" + +import os + +import click +import pandas as pd + +from pipeline.features.compute import compute_feature +from pipeline.features.registry import available_features, load_feature_module + + +@click.group(name="feature") +def feature(): + """Compute daily feature parquet files from minute bars.""" + + +def _coerce(value: str): + """Best-effort coercion of a CLI string to int, then float, else str.""" + for cast in (int, float): + try: + return cast(value) + except ValueError: + continue + return value + + +def _parse_params(pairs: tuple[str, ...]) -> dict: + """Parse repeated ``name=value`` options into a params dict.""" + params: dict = {} + for pair in pairs: + if "=" not in pair: + raise click.BadParameter(f"--param must be name=value, got '{pair}'") + key, value = pair.split("=", 1) + params[key.strip()] = _coerce(value.strip()) + return params + + +@feature.command("list") +@click.option( + "--feature-module", "feature_modules", multiple=True, + help="External module(s) to import first (dotted path or .py file)", +) +def list_(feature_modules): + """List the registered feature types.""" + for spec in feature_modules: + load_feature_module(spec) + for name in available_features(): + click.echo(name) + + +@feature.command("compute") +@click.option("--minute-path", required=True, help="Path to minute parquet dataset/file") +@click.option("--daily-path", default=None, help="Optional daily data parquet for alignment") +@click.option("--feature-type", required=True, help="Registry key of the feature class") +@click.option("--feature-name", required=True, help="Name for this feature run/output file") +@click.option("--output-dir", default="features", help="Directory to save feature parquet") +@click.option( + "--feature-module", "feature_modules", multiple=True, + help="External module(s) to import so their features register (dotted path or .py file)", +) +@click.option( + "--param", "extra_params", multiple=True, + help="Extra feature constructor param as name=value (repeatable)", +) +def compute( + minute_path, + daily_path, + feature_type, + feature_name, + output_dir, + feature_modules, + extra_params, +): + """Compute one daily feature file from raw minute bars.""" + for spec in feature_modules: + load_feature_module(spec) + + options = available_features() + if feature_type not in options: + raise click.BadParameter( + f"Unknown feature-type '{feature_type}'. Available: {options}. " + f"Use --feature-module to register an external feature.", + param_hint="--feature-type", + ) + + minute = pd.read_parquet(minute_path) + click.echo(f"Loaded minute bars: {len(minute):,} rows from {minute_path}") + + daily = None + if daily_path: + daily = pd.read_parquet(daily_path) + click.echo(f"Loaded daily data: {len(daily):,} rows from {daily_path}") + + result = compute_feature( + minute=minute, + daily=daily, + feature_type=feature_type, + **_parse_params(extra_params), + ) + + os.makedirs(output_dir, exist_ok=True) + out_path = f"{output_dir}/{feature_name}.pq" + result.to_parquet(out_path, index=False) + feature_cols = [col for col in result.columns if col not in ("symbol_id", "date")] + click.echo( + f"Saved feature: {out_path} ({len(result):,} rows, " + f"{len(feature_cols)} columns)" + ) diff --git a/pipeline/features/compute.py b/pipeline/features/compute.py new file mode 100644 index 0000000..2394745 --- /dev/null +++ b/pipeline/features/compute.py @@ -0,0 +1,75 @@ +"""Feature computation and validation.""" + +import logging +from pathlib import Path +from typing import Iterable + +import pandas as pd +from pandas.api.types import is_numeric_dtype + +from pipeline.features.registry import get_feature + +logger = logging.getLogger(__name__) + +FEATURE_KEY_COLUMNS = ["symbol_id", "date"] + + +def validate_feature_frame(features: pd.DataFrame) -> pd.DataFrame: + """Validate and normalize a daily feature frame. + + A valid feature frame is keyed by unique ``symbol_id,date`` rows and has at + least one numeric feature column beyond those keys. + """ + duplicated = features.columns[features.columns.duplicated()].tolist() + if duplicated: + raise ValueError(f"Feature output has duplicate columns: {duplicated}") + + missing = [col for col in FEATURE_KEY_COLUMNS if col not in features.columns] + if missing: + raise ValueError(f"Feature output missing required columns: {missing}") + + out = features.copy() + out["date"] = pd.to_datetime(out["date"]) + + if out.duplicated(FEATURE_KEY_COLUMNS).any(): + raise ValueError("Feature output has duplicate symbol_id,date rows") + + feature_cols = [col for col in out.columns if col not in FEATURE_KEY_COLUMNS] + if not feature_cols: + raise ValueError("Feature output must include at least one feature column") + + non_numeric = [col for col in feature_cols if not is_numeric_dtype(out[col])] + if non_numeric: + raise ValueError(f"Feature columns must be numeric: {non_numeric}") + + out = out[FEATURE_KEY_COLUMNS + feature_cols].copy() + return out.sort_values(FEATURE_KEY_COLUMNS).reset_index(drop=True) + + +def compute_feature( + minute: pd.DataFrame, + feature_type: str, + daily: pd.DataFrame | None = None, + **params, +) -> pd.DataFrame: + """Compute one registered feature from raw minute bars.""" + feature = get_feature(feature_type, **params) + result = validate_feature_frame(feature.compute(minute=minute, daily=daily)) + feature_cols = [col for col in result.columns if col not in FEATURE_KEY_COLUMNS] + logger.info( + "Feature '%s' (%r): %d symbols × %d dates, columns=%s", + feature_type, + feature, + result["symbol_id"].nunique(), + result["date"].nunique(), + feature_cols, + ) + return result + + +def read_feature_frames(feature_paths: Iterable[str | Path]) -> list[pd.DataFrame]: + """Read and validate feature parquet files.""" + return [ + validate_feature_frame(pd.read_parquet(path)) + for path in feature_paths + ] diff --git a/pipeline/features/library/__init__.py b/pipeline/features/library/__init__.py new file mode 100644 index 0000000..1d1b1a7 --- /dev/null +++ b/pipeline/features/library/__init__.py @@ -0,0 +1,3 @@ +"""Built-in feature library.""" + +from pipeline.features.library import minute_daily_summary # noqa: F401 diff --git a/pipeline/features/library/minute_daily_summary.py b/pipeline/features/library/minute_daily_summary.py new file mode 100644 index 0000000..75b3b7e --- /dev/null +++ b/pipeline/features/library/minute_daily_summary.py @@ -0,0 +1,84 @@ +"""Daily summary features derived from raw minute bars.""" + +import numpy as np +import pandas as pd + +from pipeline.features.base import BaseFeature +from pipeline.features.registry import register_feature + + +@register_feature +class MinuteDailySummaryFeature(BaseFeature): + """Aggregate intraday bars into daily summary columns.""" + + name = "minute_daily_summary" + + def compute( + self, + minute: pd.DataFrame, + daily: pd.DataFrame | None = None, + ) -> pd.DataFrame: + minute = minute.copy() + minute["date"] = pd.to_datetime(minute["date"]) + sort_cols = ["symbol_id", "date"] + if "datetime" in minute.columns: + minute["datetime"] = pd.to_datetime(minute["datetime"]) + sort_cols.append("datetime") + elif "time" in minute.columns: + sort_cols.append("time") + minute = minute.sort_values(sort_cols) + + grouped = minute.groupby(["symbol_id", "date"], sort=True) + summary = grouped.agg( + minute_bar_count=("close", "count"), + first_open=("open", "first"), + last_close=("close", "last"), + high=("high", "max"), + low=("low", "min"), + volume_sum=("volume", "sum"), + amount_sum=("amount", "sum"), + ) + summary["minute_intraday_return"] = ( + summary["last_close"] / summary["first_open"] - 1.0 + ) + summary["minute_intraday_range"] = summary["high"] / summary["low"] - 1.0 + summary["minute_vwap"] = ( + summary["amount_sum"] / summary["volume_sum"].where(summary["volume_sum"] > 0) + ) + summary = summary.reset_index() + + if daily is not None: + daily_keys = daily[["symbol_id", "date"]].copy() + daily_keys["date"] = pd.to_datetime(daily_keys["date"]) + daily_keys = daily_keys.drop_duplicates(["symbol_id", "date"]) + result = daily_keys.merge(summary, on=["symbol_id", "date"], how="left") + if "close" in daily.columns: + daily_close = daily[["symbol_id", "date", "close"]].copy() + daily_close["date"] = pd.to_datetime(daily_close["date"]) + daily_close = daily_close.drop_duplicates(["symbol_id", "date"]) + result = result.merge( + daily_close.rename(columns={"close": "daily_close"}), + on=["symbol_id", "date"], + how="left", + ) + reference_close = result["daily_close"].fillna(result["last_close"]) + else: + reference_close = result["last_close"] + else: + result = summary + reference_close = result["last_close"] + + result["minute_vwap_deviation"] = ( + result["minute_vwap"] / reference_close.replace(0.0, np.nan) - 1.0 + ) + return result[ + [ + "symbol_id", + "date", + "minute_bar_count", + "minute_intraday_return", + "minute_intraday_range", + "minute_vwap", + "minute_vwap_deviation", + ] + ] diff --git a/pipeline/features/registry.py b/pipeline/features/registry.py new file mode 100644 index 0000000..b026dfc --- /dev/null +++ b/pipeline/features/registry.py @@ -0,0 +1,79 @@ +"""Registry and factory for daily feature plugins.""" + +import importlib +import importlib.util +import inspect +from pathlib import Path +from typing import Optional, Type + +from pipeline.features.base import BaseFeature + +_REGISTRY: dict[str, Type[BaseFeature]] = {} +_builtins_loaded = False + + +def register_feature(cls: Type[BaseFeature]) -> Type[BaseFeature]: + """Class decorator that registers a feature under ``BaseFeature.name``.""" + if not (isinstance(cls, type) and issubclass(cls, BaseFeature)): + raise TypeError(f"{cls!r} is not a BaseFeature subclass") + key = getattr(cls, "name", "") + if not key: + raise ValueError(f"{cls.__name__} must set a non-empty class attribute `name`") + existing = _REGISTRY.get(key) + if existing is not None and existing is not cls: + raise ValueError( + f"Feature name '{key}' already registered by {existing.__name__}" + ) + _REGISTRY[key] = cls + return cls + + +def available_features() -> list[str]: + """Sorted names of all registered features (built-ins are loaded lazily).""" + _ensure_builtins() + return sorted(_REGISTRY) + + +def get_feature(name: str, **params) -> BaseFeature: + """Instantiate a registered feature by name. + + Only parameters accepted by the feature class's ``__init__`` are forwarded. + """ + _ensure_builtins() + if name not in _REGISTRY: + raise KeyError(f"Unknown feature '{name}'. Available: {sorted(_REGISTRY)}") + cls = _REGISTRY[name] + accepted = _accepted_params(cls) + kwargs = params if accepted is None else {k: v for k, v in params.items() if k in accepted} + return cls(**kwargs) + + +def load_feature_module(spec: str) -> None: + """Import an external module so its ``@register_feature`` classes register.""" + looks_like_file = spec.endswith(".py") or Path(spec).expanduser().exists() + if looks_like_file: + path = Path(spec).expanduser().resolve() + if not path.exists(): + raise FileNotFoundError(f"Feature module not found: {path}") + module_spec = importlib.util.spec_from_file_location(path.stem, path) + if module_spec is None or module_spec.loader is None: + raise ImportError(f"Cannot load feature module from {path}") + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + else: + importlib.import_module(spec) + + +def _accepted_params(cls: Type[BaseFeature]) -> Optional[set[str]]: + """Param names ``cls.__init__`` accepts, or None if it takes ``**kwargs``.""" + sig = inspect.signature(cls.__init__) + if any(p.kind is p.VAR_KEYWORD for p in sig.parameters.values()): + return None + return {name for name in sig.parameters if name != "self"} + + +def _ensure_builtins() -> None: + global _builtins_loaded + if not _builtins_loaded: + import pipeline.features.library # noqa: F401 + _builtins_loaded = True diff --git a/tests/test_alpha.py b/tests/test_alpha.py index 17ed739..08eb155 100644 --- a/tests/test_alpha.py +++ b/tests/test_alpha.py @@ -347,3 +347,69 @@ def test_universe_filter_does_not_corrupt_signal_history(): held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique()) # The two most liquid names (highest amount) are sh600519, sz300750. assert held == {"sh600519", "sz300750"} + + +# --- feature-aware alpha integration ---------------------------------------- + +def test_compute_alpha_without_feature_path_matches_empty_feature_paths(): + data = _make_data() + + base = compute_alpha(data, "rev5", "reversal", lookback=5) + with_empty_features = compute_alpha( + data, + "rev5", + "reversal", + lookback=5, + feature_paths=[], + ) + + pd.testing.assert_frame_equal(base, with_empty_features) + + +def test_feature_aware_alpha_reads_joined_feature_column(tmp_path): + module_path = tmp_path / "feature_aware_alpha.py" + module_path.write_text(textwrap.dedent(''' + import pandas as pd + from pipeline.alpha.base import BaseAlpha + from pipeline.alpha.registry import register_alpha + + @register_alpha + class FeatureAwareAlpha(BaseAlpha): + name = "feature_aware_test_alpha" + + def signal_from_data( + self, + data: pd.DataFrame, + close: pd.DataFrame, + ) -> pd.DataFrame: + signal = data.pivot_table( + index="date", + columns="symbol_id", + values="toy_feature", + aggfunc="first", + ) + return signal.reindex(index=close.index, columns=close.columns) + ''')) + + data = _make_data() + feature = data[["symbol_id", "date"]].copy() + feature["toy_feature"] = feature["symbol_id"].map({ + "sh600000": 1.0, + "sz000001": 2.0, + "sh600519": 3.0, + }) + feature_path = tmp_path / "toy_feature.pq" + feature.to_parquet(feature_path, index=False) + + load_alpha_module(str(module_path)) + result = compute_alpha( + data, + "feature_run", + "feature_aware_test_alpha", + feature_paths=[str(feature_path)], + ) + + assert list(result.columns) == ALPHA_COLUMNS + assert (result["alpha_name"] == "feature_run").all() + last = result[result["date"] == result["date"].max()] + assert last.set_index("symbol_id")["weight"].idxmax() == "sh600519" diff --git a/tests/test_features.py b/tests/test_features.py new file mode 100644 index 0000000..dc0ad55 --- /dev/null +++ b/tests/test_features.py @@ -0,0 +1,141 @@ +"""Tests for minute-derived daily feature plugins.""" + +import textwrap + +import numpy as np +import pandas as pd +import pytest + +from pipeline.features.compute import compute_feature, validate_feature_frame +from pipeline.features.registry import ( + available_features, + get_feature, + load_feature_module, +) + + +def _minute_bars() -> pd.DataFrame: + return pd.DataFrame({ + "symbol_id": ["sh600000", "sh600000", "sz000001"], + "symbol_name": ["PF Bank", "PF Bank", "Ping An"], + "datetime": pd.to_datetime([ + "2024-01-02 09:35:00", + "2024-01-02 09:40:00", + "2024-01-02 09:35:00", + ]), + "date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-02"]), + "time": ["09:35:00", "09:40:00", "09:35:00"], + "frequency": ["5m", "5m", "5m"], + "open": [10.0, 10.5, 20.0], + "high": [11.0, 12.0, 21.0], + "low": [9.0, 10.0, 19.0], + "close": [10.5, 11.0, 20.5], + "volume": [100.0, 300.0, 200.0], + "amount": [1000.0, 3300.0, 4100.0], + "vwap": [10.0, 11.0, 20.5], + "adjustflag": ["3", "3", "3"], + }) + + +def test_built_in_minute_daily_summary(): + daily = pd.DataFrame({ + "symbol_id": ["sh600000", "sz000001", "sh600000"], + "date": pd.to_datetime(["2024-01-02", "2024-01-02", "2024-01-03"]), + "close": [11.0, 20.5, 12.0], + }) + + result = compute_feature( + minute=_minute_bars(), + daily=daily, + feature_type="minute_daily_summary", + ) + + assert "minute_daily_summary" in available_features() + row = result[ + (result["symbol_id"] == "sh600000") + & (result["date"] == pd.Timestamp("2024-01-02")) + ].iloc[0] + assert row["minute_bar_count"] == 2 + assert np.isclose(row["minute_intraday_return"], 11.0 / 10.0 - 1.0) + assert np.isclose(row["minute_intraday_range"], 12.0 / 9.0 - 1.0) + assert np.isclose(row["minute_vwap"], 4300.0 / 400.0) + assert np.isclose(row["minute_vwap_deviation"], (4300.0 / 400.0) / 11.0 - 1.0) + + missing = result[ + (result["symbol_id"] == "sh600000") + & (result["date"] == pd.Timestamp("2024-01-03")) + ].iloc[0] + assert pd.isna(missing["minute_vwap"]) + + +def test_load_external_feature_module_and_filter_params(tmp_path): + module_path = tmp_path / "external_feature.py" + module_path.write_text(textwrap.dedent(''' + import pandas as pd + from pipeline.features.base import BaseFeature + from pipeline.features.registry import register_feature + + @register_feature + class ExternalVolumeFeature(BaseFeature): + name = "external_volume_feature" + + def __init__(self, scale: float = 1.0): + self.scale = scale + + def compute(self, minute: pd.DataFrame, daily=None) -> pd.DataFrame: + out = ( + minute.groupby(["symbol_id", "date"], as_index=False)["volume"] + .sum() + .rename(columns={"volume": "scaled_volume"}) + ) + out["scaled_volume"] *= self.scale + return out + ''')) + + load_feature_module(str(module_path)) + assert "external_volume_feature" in available_features() + + instance = get_feature("external_volume_feature", scale=2.0, ignored=99) + assert instance.scale == 2.0 + assert not hasattr(instance, "ignored") + + result = compute_feature( + minute=_minute_bars(), + feature_type="external_volume_feature", + scale=2.0, + ignored=99, + ) + row = result[result["symbol_id"] == "sh600000"].iloc[0] + assert np.isclose(row["scaled_volume"], 800.0) + + +def test_validate_feature_frame_rejects_missing_keys(): + with pytest.raises(ValueError, match="missing required"): + validate_feature_frame(pd.DataFrame({"symbol_id": ["sh600000"], "x": [1.0]})) + + +def test_validate_feature_frame_rejects_duplicate_keys_after_date_normalization(): + with pytest.raises(ValueError, match="duplicate symbol_id,date"): + validate_feature_frame(pd.DataFrame({ + "symbol_id": ["sh600000", "sh600000"], + "date": ["2024-01-02", pd.Timestamp("2024-01-02")], + "x": [1.0, 2.0], + })) + + +def test_validate_feature_frame_rejects_duplicate_columns(): + bad = pd.DataFrame( + [["sh600000", pd.Timestamp("2024-01-02"), 1.0, 2.0]], + columns=["symbol_id", "date", "dup", "dup"], + ) + with pytest.raises(ValueError, match="duplicate columns"): + validate_feature_frame(bad) + + +def test_validate_feature_frame_rejects_non_numeric_feature_columns(): + with pytest.raises(ValueError, match="numeric"): + validate_feature_frame(pd.DataFrame({ + "symbol_id": ["sh600000"], + "date": [pd.Timestamp("2024-01-02")], + "bad": ["not numeric"], + })) diff --git a/tests/test_minute_downloader.py b/tests/test_minute_downloader.py new file mode 100644 index 0000000..127e320 --- /dev/null +++ b/tests/test_minute_downloader.py @@ -0,0 +1,215 @@ +"""Tests for raw Baostock minute bar download plumbing.""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +import data.downloader as low_level_downloader +import pipeline.data.downloader as pipeline_downloader +from data.downloader import download_minute_batch +from pipeline.common.schema import MINUTE_BAR_COLUMNS +from pipeline.data.downloader import download_minute_universe + + +class _FakeResult: + def __init__(self, rows, error_code="0", error_msg=""): + self.rows = rows + self.error_code = error_code + self.error_msg = error_msg + self._idx = -1 + + def next(self): + self._idx += 1 + return self._idx < len(self.rows) + + def get_row_data(self): + return self.rows[self._idx] + + +def test_download_minute_batch_maps_and_parses_baostock_rows(monkeypatch): + rows = [ + [ + "2024-01-02", + "20240102093500000", + "sh.600000", + "10", + "11", + "9", + "10.5", + "1000", + "10500", + "3", + ], + [ + "2024-01-02", + "20240102094000000", + "sh.600000", + "10.5", + "12", + "10", + "11", + "2000", + "22000", + "3", + ], + ] + calls = [] + + def fake_query(**kwargs): + calls.append(kwargs) + return _FakeResult(rows) + + monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None) + monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None) + monkeypatch.setattr( + low_level_downloader.bs, + "query_history_k_data_plus", + fake_query, + ) + + [(symbol, df)] = list( + download_minute_batch( + ["sh600000"], + "2024-01-02", + "2024-01-02", + frequency=5, + ) + ) + + assert symbol == "sh600000" + assert calls[0]["code"] == "sh.600000" + assert calls[0]["frequency"] == "5" + assert calls[0]["adjustflag"] == "3" + assert df is not None + assert df["datetime"].iloc[0] == pd.Timestamp("2024-01-02 09:35:00") + assert df["time"].tolist() == ["09:35:00", "09:40:00"] + assert (df["frequency"] == "5m").all() + assert np.isclose(df["open"].iloc[0], 10.0) + assert np.isclose(df["vwap"].iloc[0], 10.5) + assert pd.api.types.is_numeric_dtype(df["volume"]) + + +def test_download_minute_batch_empty_result_yields_none(monkeypatch): + monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None) + monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None) + monkeypatch.setattr( + low_level_downloader.bs, + "query_history_k_data_plus", + lambda **kwargs: _FakeResult([]), + ) + + assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [ + ("sh600000", None) + ] + + +def test_download_minute_batch_rejects_unparsed_timestamps(monkeypatch): + bad_rows = [[ + "2024-01-02", + "not-a-time", + "sh.600000", + "10", + "11", + "9", + "10.5", + "1000", + "10500", + "3", + ]] + monkeypatch.setattr(low_level_downloader.bs, "login", lambda: None) + monkeypatch.setattr(low_level_downloader.bs, "logout", lambda: None) + monkeypatch.setattr( + low_level_downloader.bs, + "query_history_k_data_plus", + lambda **kwargs: _FakeResult(bad_rows), + ) + + assert list(download_minute_batch(["sh600000"], "2024-01-02", "2024-01-02")) == [ + ("sh600000", None) + ] + + +def test_download_minute_universe_writes_frequency_month_partitions(tmp_path, monkeypatch): + minute = pd.DataFrame({ + "symbol": ["sh600000", "sh600000"], + "datetime": pd.to_datetime(["2024-01-02 09:35:00", "2024-01-02 09:40:00"]), + "date": pd.to_datetime(["2024-01-02", "2024-01-02"]), + "time": ["09:35:00", "09:40:00"], + "frequency": ["5m", "5m"], + "open": [10.0, 10.5], + "high": [11.0, 12.0], + "low": [9.0, 10.0], + "close": [10.5, 11.0], + "volume": [1000.0, 2000.0], + "amount": [10500.0, 22000.0], + "vwap": [10.5, 11.0], + "adjustflag": ["3", "3"], + }) + + monkeypatch.setattr( + pipeline_downloader, + "_resolve_universe", + lambda universe, max_symbols=0: pd.DataFrame({ + "symbol_id": ["sh600000"], + "symbol_name": ["PF Bank"], + }), + ) + + def fake_batch(symbols, start, end, frequency=5): + assert symbols == ["sh600000"] + assert frequency == "5" + yield "sh600000", minute + + monkeypatch.setattr(pipeline_downloader, "download_minute_batch", fake_batch) + + preserved = tmp_path / "toy" / "frequency=15m" / "month=2024-01" / "old.pq" + preserved.parent.mkdir(parents=True) + preserved_minute = minute.copy() + preserved_minute["frequency"] = "15m" + preserved_minute["symbol_id"] = "sh600000" + preserved_minute["symbol_name"] = "PF Bank" + preserved_minute[MINUTE_BAR_COLUMNS].to_parquet(preserved, index=False) + + stats = download_minute_universe( + universe="toy", + start_date="2024-01-02", + end_date="2024-01-02", + output_dir=str(tmp_path), + chunk_size=1, + frequency="5", + ) + + dataset_path = Path(stats["dataset_path"]) + assert (dataset_path / "frequency=5m" / "month=2024-01").is_dir() + assert preserved.exists() + out = pd.read_parquet(dataset_path / "frequency=5m") + assert (set(MINUTE_BAR_COLUMNS) - {"frequency"}) <= set(out.columns) + assert set(out["symbol_id"]) == {"sh600000"} + assert set(out["symbol_name"]) == {"PF Bank"} + assert stats["n_rows"] == 2 + + +def test_download_minute_universe_raises_when_all_symbols_empty(tmp_path, monkeypatch): + monkeypatch.setattr( + pipeline_downloader, + "_resolve_universe", + lambda universe, max_symbols=0: pd.DataFrame({ + "symbol_id": ["sh600000"], + "symbol_name": ["PF Bank"], + }), + ) + monkeypatch.setattr( + pipeline_downloader, + "download_minute_batch", + lambda symbols, start, end, frequency=5: iter([("sh600000", None)]), + ) + + with pytest.raises(RuntimeError, match="No minute data"): + download_minute_universe( + universe="toy", + start_date="2024-01-02", + end_date="2024-01-02", + output_dir=str(tmp_path), + )