Add minute bar feature pipeline

This commit is contained in:
Yuxuan Yan
2026-06-16 13:57:17 +08:00
parent 17fa75495d
commit 83a006bbe4
19 changed files with 1289 additions and 11 deletions
+3
View File
@@ -4,6 +4,7 @@
Phases: Phases:
data — Download daily bars to parquet data — Download daily bars to parquet
alpha — Compute alpha weights from data alpha — Compute alpha weights from data
feature — Compute daily features from minute bars
combo — Combine alphas into a single weight combo — Combine alphas into a single weight
portfolio — Build tradable positions and simulate execution portfolio — Build tradable positions and simulate execution
""" """
@@ -14,6 +15,7 @@ import click
from pipeline.data.cli import data from pipeline.data.cli import data
from pipeline.alpha.cli import alpha from pipeline.alpha.cli import alpha
from pipeline.features.cli import feature
from pipeline.combo.cli import combo from pipeline.combo.cli import combo
from pipeline.portfolio.cli import portfolio from pipeline.portfolio.cli import portfolio
from tools.pqcat import pqcat from tools.pqcat import pqcat
@@ -40,6 +42,7 @@ def cli(log_level):
cli.add_command(data) cli.add_command(data)
cli.add_command(alpha) cli.add_command(alpha)
cli.add_command(feature)
cli.add_command(combo) cli.add_command(combo)
cli.add_command(portfolio) cli.add_command(portfolio)
cli.add_command(pqcat) cli.add_command(pqcat)
+158
View File
@@ -31,11 +31,68 @@ _BATCH_COLUMNS = [
"peTTM", "pbMRQ", "psTTM", "pcfNcfTTM", "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): class _SessionLost(Exception):
"""baostock reported the session was dropped (``用户未登录``).""" """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]: 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.""" """Download daily bars from akshare. Returns DataFrame with OHLCV columns."""
try: try:
@@ -239,3 +296,104 @@ def download_daily_batch(
except Exception: except Exception:
pass 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
+67
View File
@@ -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.
+19 -5
View File
@@ -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 score. The base class turns a signal into cross-sectionally z-scored weights
via :meth:`to_weights` (override it for a different normalization). via :meth:`to_weights` (override it for a different normalization).
""" """
from abc import ABC, abstractmethod from abc import ABC
import numpy as np import numpy as np
import pandas as pd import pandas as pd
@@ -15,15 +15,14 @@ class BaseAlpha(ABC):
"""A position-weight alpha over a cross-section of stocks. """A position-weight alpha over a cross-section of stocks.
Concrete subclasses must set a unique class-level :attr:`name` (the registry Concrete subclasses must set a unique class-level :attr:`name` (the registry
key) and implement :meth:`signal`. Construct subclasses with their own typed key) and implement either :meth:`signal` or :meth:`signal_from_data`.
parameters (e.g. ``lookback``); the factory passes only the parameters a Construct subclasses with their own typed parameters (e.g. ``lookback``);
given ``__init__`` accepts. the factory passes only the parameters a given ``__init__`` accepts.
""" """
#: Unique registry key. Every concrete alpha must set this to a non-empty str. #: Unique registry key. Every concrete alpha must set this to a non-empty str.
name: str = "" name: str = ""
@abstractmethod
def signal(self, close: pd.DataFrame) -> pd.DataFrame: def signal(self, close: pd.DataFrame) -> pd.DataFrame:
"""Compute the raw signal. """Compute the raw signal.
@@ -34,6 +33,21 @@ class BaseAlpha(ABC):
A wide DataFrame aligned to ``close`` where higher values indicate a A wide DataFrame aligned to ``close`` where higher values indicate a
stronger long. Use NaN where the signal is undefined. 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: def to_weights(self, signal: pd.DataFrame) -> pd.DataFrame:
"""Cross-sectionally z-score a signal into signed position weights. """Cross-sectionally z-score a signal into signed position weights.
+6 -1
View File
@@ -56,6 +56,10 @@ def list_(alpha_modules):
@click.option("--output-dir", default="alphas", help="Directory to save alpha parquet") @click.option("--output-dir", default="alphas", help="Directory to save alpha parquet")
@click.option("--lookback", default=5, type=int, help="Lookback days") @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("--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( @click.option(
"--alpha-module", "alpha_modules", multiple=True, "--alpha-module", "alpha_modules", multiple=True,
help="External module(s) to import so their alphas register (dotted path or .py file)", 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", 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, 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.""" """Compute one alpha from raw data and save as parquet."""
for spec in alpha_modules: for spec in alpha_modules:
load_alpha_module(spec) 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_name=alpha_name,
alpha_type=alpha_type, alpha_type=alpha_type,
universe=universe, universe=universe,
feature_paths=feature_paths,
**params, **params,
) )
+51 -2
View File
@@ -7,12 +7,15 @@ through :mod:`pipeline.alpha.registry`.
""" """
import logging import logging
from pathlib import Path
from typing import Iterable
import numpy as np import numpy as np
import pandas as pd import pandas as pd
from pipeline.alpha.registry import get_alpha from pipeline.alpha.registry import get_alpha
from pipeline.common.schema import ALPHA_COLUMNS 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__) logger = logging.getLogger(__name__)
@@ -33,6 +36,38 @@ def _pivot_open(df: pd.DataFrame) -> pd.DataFrame:
return pivot.sort_index() 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: def _forward_open_to_open_returns(open_: pd.DataFrame) -> pd.DataFrame:
"""Return earned by a close-formed signal after next-open execution. """Return earned by a close-formed signal after next-open execution.
@@ -105,6 +140,8 @@ def compute_alpha(
alpha_name: str, alpha_name: str,
alpha_type: str, alpha_type: str,
universe: dict | None = None, universe: dict | None = None,
feature_paths: Iterable[str | Path] | None = None,
feature_frames: Iterable[pd.DataFrame] | None = None,
**params, **params,
) -> pd.DataFrame: ) -> pd.DataFrame:
"""Compute alpha weights from raw data. """Compute alpha weights from raw data.
@@ -118,6 +155,10 @@ def compute_alpha(
:func:`investable_universe_mask`) *before* it is turned into :func:`investable_universe_mask`) *before* it is turned into
weights, so unheld names get weight 0. Keys are forwarded as keyword weights, so unheld names get weight 0. Keys are forwarded as keyword
arguments to :func:`investable_universe_mask`. 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``, **params: Constructor parameters for the alpha (e.g. ``lookback``,
``vol_window``). Only the params the alpha's ``__init__`` accepts are ``vol_window``). Only the params the alpha's ``__init__`` accepts are
used; extras are ignored. used; extras are ignored.
@@ -128,12 +169,20 @@ def compute_alpha(
Raises: Raises:
KeyError: If ``alpha_type`` is not registered. 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) alpha = get_alpha(alpha_type, **params)
close = _pivot_close(data) close = _pivot_close(data)
signal = alpha.signal_from_data(data, close)
if universe is None: if universe is None:
weights = alpha.weights(close) weights = alpha.to_weights(signal)
else: else:
signal = alpha.signal(close)
mask = investable_universe_mask(data, signal, **universe) mask = investable_universe_mask(data, signal, **universe)
weights = alpha.to_weights(signal.where(mask)) weights = alpha.to_weights(signal.where(mask))
+18
View File
@@ -26,6 +26,24 @@ DATA_COLUMNS: Final[list[str]] = [
"pcfNcfTTM", # float64: P/CF (net cash flow, TTM) "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. # Required columns for alpha parquet files.
# Alphas are position WEIGHTS: positive=long, negative=short. # Alphas are position WEIGHTS: positive=long, negative=short.
ALPHA_COLUMNS: Final[list[str]] = [ ALPHA_COLUMNS: Final[list[str]] = [
+35 -1
View File
@@ -3,7 +3,7 @@
import click import click
from datetime import date 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") @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']}" f"{stats['n_rows']:,} bars, {stats['date_min']}{stats['date_max']}"
) )
click.echo(f"Dataset: {stats['dataset_path']}") 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']}")
+126 -2
View File
@@ -11,9 +11,13 @@ import pyarrow.dataset as pads
# Reuse existing downloader and universe modules # Reuse existing downloader and universe modules
sys.path.insert(0, str(Path(__file__).resolve().parents[2])) 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 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__) 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( def download_universe(
universe: str = "csi500", universe: str = "csi500",
start_date: str = "2017-01-01", 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_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()), "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()),
}
+1
View File
@@ -0,0 +1 @@
"""Daily feature plugin package."""
+34
View File
@@ -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})"
+108
View File
@@ -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)"
)
+75
View File
@@ -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
]
+3
View File
@@ -0,0 +1,3 @@
"""Built-in feature library."""
from pipeline.features.library import minute_daily_summary # noqa: F401
@@ -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",
]
]
+79
View File
@@ -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
+66
View File
@@ -347,3 +347,69 @@ def test_universe_filter_does_not_corrupt_signal_history():
held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique()) held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique())
# The two most liquid names (highest amount) are sh600519, sz300750. # The two most liquid names (highest amount) are sh600519, sz300750.
assert held == {"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"
+141
View File
@@ -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"],
}))
+215
View File
@@ -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),
)