hamuna-quant-cli 0.1.0.dev93__py3-none-any.whl

hamuna_quant_cli/prebuilt_resolver.py

@@ -0,0 +1,470 @@
+"""ADR-0040 Phase C — 自适应 prebuilt dataset resolver (v2 自有版本, B5).
+
+来源: 原仓根 `hamuna_quant_cli/references/prebuilt_resolver.py`.
+本文件是 wholesale copy, 算法/filters/prefix 不变. 迁移原因: v2 skill 独立分发.
+
+背景:
+  旧 _try_prebuilt_parquet 硬编码 `__bundle__all_a_D.parquet`, 用户若只查单股
+  (如 ['600000.SH']) → bundle 不存在 / 不含该标 → 全 fallback 到慢速
+  market_mod 路径 (~92s). 同时不论 universe 多小, 都把整张 402MB parquet 全部
+  读入内存再 pandas filter — 浪费 IO + 内存.
+
+设计:
+  1) **registry**: 描述本地已知 prebuilt 的覆盖范围 (kind=bundle/single,
+     universe / symbol, 路径模板). 用 sklearn 风格 scan → 取第一个匹配的
+     prebuilt, 没匹配 → None.
+  2) **filter pushdown**: 用 pyarrow `filters=[(stockCode, 'in', codes),
+     (time, '>=', s), (time, '<=', e)]`, 只把匹配行解压出内存 (400MB
+     → 1MB 量级). 比 `read_parquet(columns=...)` + pandas isin 省 10×
+     IO (实测).
+  3) **优先级**: bundle 优先 (一次 IO 覆盖多标); 单股 fallback.
+  4) **增量扩展**: 新加 prebuilt (etf_b, convertible_b 等) 只需往
+     KNOWN_PREBUILT 加一行.
+
+隔离保证:
+  - v2 自有: 不依赖 v1 strategy_cli (避免 cycle). 用与 market._CLOUD_UNIVERSE_PREFIX
+    同等的 prefix 规则本地复刻.
+  - 单文件, 无外部 import (除 stdlib + pyarrow).
+
+API:
+  resolve(universe, start, end, data_cache_dir=None, period='1d')
+    -> (DataFrame | None, source_path: str | None)
+  discover(data_cache_dir=None)
+    -> list[dict]: 扫盘已知 prebuilt, 报每个的 metadata (含覆盖 symbol 数)
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import pandas as pd
+
+# ---- prebuilt registry ----
+
+@dataclass(frozen=True)
+class PrebuiltSpec:
+    """单个 prebuilt 文件的元信息.
+
+    path_template: 相对 data_cache_dir 的文件名 (或绝对路径)
+    kind: 'bundle' = 含多标的 stockCode 列, 'single' = 仅一支股票 (无 stockCode 列)
+    """
+    path_template: str            # e.g. '__bundle__all_a_D.parquet'
+    kind: str                     # 'bundle' | 'single'
+    universe_label: str           # e.g. 'all_a' / '600000.SH' — 仅供日志/UI
+    # bundle 用: 每个 stockCode 是否在覆盖范围 (按前缀匹配)
+    bundle_prefix_predicate: tuple[str, ...] | None = None
+    # single 用: 该 parquet 对应的完整带后缀 symbol
+    single_symbol: str | None = None
+
+    @property
+    def path(self) -> Path:
+        # 默认 ~/.hamuna/data_cache/datasets/; 调用方可覆盖 data_cache_dir
+        # 这里延迟到 resolve() 时再拼 (避免模块级读 home dir 副作用)
+        return Path()
+
+
+# 已知 prebuilt (按优先级: bundle 优先 → single).
+# bundle_prefix_predicate 与 strategy_cli.fundamental.data.market._CLOUD_UNIVERSE_PREFIX
+# 等价 — 本地复刻避免跨模块 import.
+#   all_a:      ('60', '68', '00', '30', '40', '43', '83', '87', '88')
+#   etf:        ('51', '56', '15', '16')
+#   convertible_bond: ('11', '12')
+KNOWN_PREBUILT: tuple[PrebuiltSpec, ...] = (
+    PrebuiltSpec(
+        path_template='__bundle__all_a_D.parquet',
+        kind='bundle',
+        universe_label='all_a',
+        bundle_prefix_predicate=('60', '68', '00', '30', '40', '43', '83', '87', '88'),
+    ),
+    PrebuiltSpec(
+        path_template='__bundle__etf_D.parquet',
+        kind='bundle',
+        universe_label='etf',
+        bundle_prefix_predicate=('51', '56', '15', '16'),
+    ),
+    PrebuiltSpec(
+        path_template='__bundle__cb_D.parquet',
+        kind='bundle',
+        universe_label='convertible_bond',
+        bundle_prefix_predicate=('11', '12'),
+    ),
+    PrebuiltSpec(
+        path_template='000300_D_fq1.parquet',
+        kind='single',
+        universe_label='000300.SH (HS300 基准)',
+        single_symbol='000300.SH',
+    ),
+    PrebuiltSpec(
+        path_template='600000_D_fq1.parquet',
+        kind='single',
+        universe_label='600000.SH',
+        single_symbol='600000.SH',
+    ),
+)
+
+
+def _default_cache_dir() -> Path:
+    return Path.home() / '.hamuna' / 'data_cache' / 'datasets'
+
+
+def _resolve_path(spec: PrebuiltSpec, cache_dir: Path) -> Path:
+    return cache_dir / spec.path_template
+
+
+def _bundle_covers(spec: PrebuiltSpec, universe: list[str]) -> set[str] | None:
+    """bundle: 计算 universe ∩ 覆盖集 (按 prefix) → 返命中集合.
+
+    返回 None = spec 不是 bundle; 空集 = 完全不命中; 非空集 = 命中子集.
+    """
+    if spec.kind != 'bundle' or spec.bundle_prefix_predicate is None:
+        return None
+    prefixes = spec.bundle_prefix_predicate
+    hit = set()
+    for sym in universe:
+        bare = sym.split('.')[0]
+        if any(bare.startswith(p) for p in prefixes):
+            hit.add(sym)
+    return hit
+
+
+def _build_arrow_filters(codes: list[str], start: str, end: str) -> list:
+    """pyarrow Dataset.filter 接受的谓词列表.
+
+    codes: 期望含后缀 (e.g. '600000.SH'). bundle parquet 的 stockCode 列就
+    是这个格式. start/end: 'YYYYMMDD'. time 列是 YYYYMMDDhhmmss int64 →
+    起止用 int 比较.
+    """
+    s_int = int(start) * 1_000_000  # YYYYMMDD → YYYYMMDD000000
+    e_int = int(end) * 1_000_000 + 99_999_999  # 含 end 当日最后 bar
+    return [
+        ('stockCode', 'in', codes),
+        ('time', '>=', s_int),
+        ('time', '<=', e_int),
+    ]
+
+
+def _read_with_pushdown(path: Path, columns: list[str],
+                        filters: list | None = None) -> pd.DataFrame:
+    """pyarrow filter pushdown 读子集.
+
+    走 pyarrow.parquet.ParquetDataset + read_table(filters=...) — 只把
+    满足谓词的 row group 解压到内存. 与 `pd.read_parquet(columns=...)`
+    比, 后者读所有行, 只是稀疏列; 前者真"只读子集".
+
+    大文件子集提取推荐路径 (400MB → 1MB). filters=None 走 plain read.
+    """
+    import pyarrow.parquet as pq
+    table = pq.read_table(path, columns=columns, filters=filters)
+    return table.to_pandas()
+
+
+def _read_single_with_pushdown(path: Path, start: str, end: str) -> pd.DataFrame | None:
+    """single parquet (无 stockCode 列) → 读 [start, end] 窗口.
+
+    schema 容错:
+      - `time` int64 = YYYYMMDDhhmmss (14 位) → pyarrow filter int 比较
+      - `date` int / str = YYYYMMDD (8 位)
+          * int64 → pyarrow filter int 比较
+          * string → pyarrow 不支持 str/int 跨类型 filter, 退回 pandas str slice + 比较
+      - 都没 → 返 None (schema 不兼容)
+    """
+    import pyarrow.parquet as pq
+    try:
+        schema = pq.read_schema(path)
+        cols_all = [c for c in ('time', 'date', 'open', 'high', 'low', 'close', 'volume', 'stockCode')
+                    if c in schema.names]
+        time_type = schema.field('time').type if 'time' in schema.names else None
+        date_type = schema.field('date').type if 'date' in schema.names else None
+    except Exception:
+        return None
+
+    if time_type is not None:
+        s_int = int(start) * 1_000_000
+        e_int = int(end) * 1_000_000 + 99_999_999
+        try:
+            table = pq.read_table(
+                path, columns=cols_all,
+                filters=[('time', '>=', s_int), ('time', '<=', e_int)],
+            )
+        except Exception:
+            return None
+        return table.to_pandas()
+
+    if date_type is not None:
+        # date 列: pyarrow 可对 int 比较; string 列需退回 pandas.
+        from pyarrow import types as pat
+        if pat.is_integer(date_type):
+            try:
+                table = pq.read_table(
+                    path, columns=cols_all,
+                    filters=[('date', '>=', int(start)), ('date', '<=', int(end))],
+                )
+            except Exception:
+                return None
+            return table.to_pandas()
+        # string: read full, pandas slice 头 8 位 + int 比较
+        try:
+            table = pq.read_table(path, columns=cols_all)
+        except Exception:
+            return None
+        df = table.to_pandas()
+        if df.empty:
+            return None
+        d_str = df['date'].astype(str).str[:8]
+        mask = (d_str >= start) & (d_str <= end)
+        return df[mask]
+
+    return None
+
+
+def _normalize_bundle_df(df: pd.DataFrame) -> pd.DataFrame:
+    """bundle raw df → akquant-ready schema.
+
+    time → date (datetime64); stockCode → symbol; 选列.
+    """
+    df['date'] = pd.to_datetime(df['time'].astype(str).str[:8], format='%Y%m%d')
+    df = df.rename(columns={'stockCode': 'symbol'})
+    return df[['date', 'open', 'high', 'low', 'close', 'volume', 'symbol']].reset_index(drop=True)
+
+
+def _normalize_single_df(df: pd.DataFrame, symbol: str) -> pd.DataFrame:
+    """single raw df → akquant-ready schema (注入 symbol 列).
+
+    容错: df 既有 `time` 又有 `date` 优先用 `time` (更精确).
+    date 列必须 datetime64 — akquant utils df_to_arrays 要求 (会找
+    ["日期","date","datetime","time","timestamp"] 候选列, 但 string 类的
+    date 不被认作 timestamp, 需要 pandas 端 to_datetime 转).
+    """
+    if 'time' in df.columns:
+        df['date'] = pd.to_datetime(df['time'].astype(str).str[:8], format='%Y%m%d')
+    elif 'date' in df.columns:
+        # date 可能是 YYYYMMDD 字符串 / int, 统一 to_datetime
+        df['date'] = pd.to_datetime(df['date'].astype(str).str[:8], format='%Y%m%d')
+    else:
+        return df  # 无 date/time 列, schema 不兼容
+    df['symbol'] = symbol
+    keep = [c for c in ['date', 'open', 'high', 'low', 'close', 'volume', 'symbol']
+            if c in df.columns]
+    return df[keep].reset_index(drop=True)
+
+
+def resolve(universe: list[str], start: str, end: str,
+            data_cache_dir: Path | None = None,
+            period: str = '1d') -> tuple[pd.DataFrame | None, str | None]:
+    """自适应选 prebuilt + pyarrow filter pushdown 读子集.
+
+    返回 (df, source_path):
+      - df 为 None: 所有 prebuilt 都不命中 (调用方回退 market_mod)
+      - source_path 描述命中哪个文件 + kind ('bundle:__bundle__all_a_D' 等)
+
+    选择策略:
+      1) bundle 整包命中 → 一次 IO, filter pushdown 取 universe ∩ 覆盖集
+      2) 否则按 universe 中每标的找 single parquet → 每个标一次 IO
+      3) 都不命中 → 返 None
+
+    注: '1d' 之外的 period 当前不支持 (与旧 _try_prebuilt_parquet 一致).
+    """
+    if period != '1d':
+        return None, None
+    cache_dir = data_cache_dir or _default_cache_dir()
+    if not universe:
+        return None, None
+
+    # 裸码 → 带后缀 (config 标准写法是裸码; bundle stockCode 是 '600000.SH' 形态,
+    # filter 必须带后缀才命中 — 2026-08-19 e2e 修复: run/manifest/parity 全走这里).
+    from .akquant_schema_adapter import normalize_symbol
+    universe = [normalize_symbol(s) for s in universe]
+
+    # ---- 1) bundle 整包 (累积式 — 跨 universe 时各 bundle 各取子集, 最后 concat) ----
+    bundle_pieces: list[pd.DataFrame] = []
+    bundle_sources: list[str] = []
+    remaining = set(universe)
+    for spec in KNOWN_PREBUILT:
+        if spec.kind != 'bundle':
+            continue
+        path = _resolve_path(spec, cache_dir)
+        if not path.exists():
+            continue
+        hit = _bundle_covers(spec, remaining)
+        if not hit:
+            continue
+        try:
+            df = _read_with_pushdown(
+                path,
+                columns=['time', 'stockCode', 'open', 'high', 'low', 'close', 'volume'],
+                filters=_build_arrow_filters(sorted(hit), start, end),
+            )
+        except Exception:  # noqa: BLE001
+            return None, None
+        if df.empty:
+            continue
+        bundle_pieces.append(_normalize_bundle_df(df))
+        bundle_sources.append(f'{spec.path_template}:{len(hit)}')
+        remaining -= hit
+        if not remaining:
+            break
+    if bundle_pieces:
+        merged = pd.concat(bundle_pieces, ignore_index=True)
+        return merged, 'bundle:' + '+'.join(bundle_sources)
+
+    # ---- 2) single parquet 拼 ----
+    by_symbol = {s: s for s in universe}
+    pieces: list[pd.DataFrame] = []
+    used: list[str] = []
+    for sym in universe:
+        for spec in KNOWN_PREBUILT:
+            if spec.kind != 'single' or spec.single_symbol != sym:
+                continue
+            path = _resolve_path(spec, cache_dir)
+            if not path.exists():
+                continue
+            raw = _read_single_with_pushdown(path, start, end)
+            if raw is None or raw.empty:
+                break
+            # single parquet 可能没 stockCode 列 (旧 schema) — 显式注入
+            if 'stockCode' not in raw.columns:
+                raw = _normalize_single_df(raw, sym)
+            else:
+                raw = _normalize_bundle_df(raw)
+            pieces.append(raw)
+            used.append(sym)
+            break
+        else:
+            continue
+
+    if not pieces:
+        return None, None
+    df = pd.concat(pieces, ignore_index=True)
+    return df, f'single:{",".join(used)}'
+
+
+def discover(data_cache_dir: Path | None = None) -> list[dict]:
+    """扫盘已知 prebuilt → 报每个 metadata (path/exists/size).
+
+    供 CLI / smoke 打印 prebuilt 状态. 不读 parquet 内容, 只 stat.
+    """
+    cache_dir = data_cache_dir or _default_cache_dir()
+    out: list[dict] = []
+    for spec in KNOWN_PREBUILT:
+        path = _resolve_path(spec, cache_dir)
+        out.append({
+            'path': str(path),
+            'exists': path.exists(),
+            'size_mb': round(path.stat().st_size / 1e6, 1) if path.exists() else 0.0,
+            'kind': spec.kind,
+            'universe_label': spec.universe_label,
+        })
+    return out
+
+
+def _selfcheck() -> None:
+    """冒烟:
+      1) discover() 不抛
+      2) resolve(全 A 子集, 1y) → 命中 __bundle__all_a_D, df 非空 + pyarrow filter 推下行
+      3) resolve(单股 600000.SH) → 命中 600000_D_fq1.parquet (新增能力)
+      4) resolve([未知标]) → None
+    """
+    cache_dir = _default_cache_dir()
+    discovered = discover(cache_dir)
+    assert isinstance(discovered, list) and len(discovered) >= 3, \
+        f'KNOWN_PREBUILT 至少 3 个 (bundle+2 single), 实际 {len(discovered)}'
+    print(f'OK: discover() → {len(discovered)} known prebuilt')
+    for d in discovered:
+        flag = '✓' if d['exists'] else '✗'
+        print(f'  {flag} {d["path"]} ({d["size_mb"]} MB, {d["kind"]})')
+
+    # ---- 2) bundle 命中 ----
+    bundle_path = cache_dir / '__bundle__all_a_D.parquet'
+    if bundle_path.exists():
+        # 拿全 A 池 stockCode list (复用 _try_prebuilt_parquet 风格, 但要全量)
+        all_codes_df = pd.read_parquet(bundle_path, columns=['stockCode'])
+        codes = sorted(all_codes_df['stockCode'].unique().tolist())
+        sub = codes[:10]  # 前 10 标的, 验证 resolver 真过滤
+        df, src = resolve(sub, '20240701', '20250731', data_cache_dir=cache_dir)
+        assert df is not None and src is not None and src.startswith('bundle:'), \
+            f'bundle 路径应命中, src={src}'
+        assert df['symbol'].nunique() <= len(sub), \
+            f'filter pushdown 失效? unique_syms={df["symbol"].nunique()} > sub={len(sub)}'
+        # 验证 df 真的只有 sub 里那些标的
+        assert set(df['symbol'].unique()) <= set(sub), \
+            f'filter 漏标: {set(df["symbol"].unique()) - set(sub)}'
+        print(f'OK: bundle 命中 + filter pushdown (子集 {len(sub)} 标 → {df["symbol"].nunique()} 唯一, src={src})')
+    else:
+        print('SKIP: __bundle__all_a_D 不存在, bundle 路径未测')
+
+    # ---- 3) single 命中 — bundle 优先策略下 600000.SH 会被 bundle 抢先命中;
+    #       这里直接测 _read_single_with_pushdown (走 single IO + filter pushdown).
+    single_path = cache_dir / '600000_D_fq1.parquet'
+    if single_path.exists():
+        raw = _read_single_with_pushdown(single_path, '20240701', '20250731')
+        assert raw is not None and not raw.empty, 'single 路径应非空'
+        # 测 full normalize (注入 symbol + date)
+        norm = _normalize_single_df(raw, '600000.SH')
+        assert norm['symbol'].unique().tolist() == ['600000.SH']
+        assert list(norm.columns) == ['date', 'open', 'high', 'low', 'close', 'volume', 'symbol']
+        print(f'OK: _read_single_with_pushdown + normalize (600000.SH → {len(norm)} rows)')
+        # 单股场景通过 resolve() 也会被 bundle 抢命中, 这是预期 (bundle 优先).
+        df, src = resolve(['600000.SH'], '20240701', '20250731', data_cache_dir=cache_dir)
+        assert df is not None and src is not None
+        print(f'OK: resolve([600000.SH]) → bundle 优先 (src={src})')
+    else:
+        print('SKIP: 600000_D_fq1 不存在, single 路径未测')
+
+    # ---- 4) 完全不命中 → None ----
+    df, src = resolve(['999999.SH'], '20240101', '20241231', data_cache_dir=cache_dir)
+    assert df is None, f'未知标应 None, 实际 df 非空 ({len(df)} rows)'
+    print(f'OK: 不命中标 → (None, None)')
+
+    # ---- 5) prefix 隔离 — 直接验 KNOWN_PREBUILT prefix 谓词, 不读 parquet (无需数据) ----
+    from ..prebuilt_resolver import _bundle_covers, KNOWN_PREBUILT as _KP
+    by_label = {s.universe_label: s for s in _KP}
+    if 'all_a' in by_label and 'etf' in by_label and 'convertible_bond' in by_label:
+        all_a_spec = by_label['all_a']
+        etf_spec = by_label['etf']
+        cb_spec = by_label['convertible_bond']
+        assert _bundle_covers(all_a_spec, ['113021.SH']) == set(), \
+            'all_a bundle 不应覆盖 CB 11x'
+        assert _bundle_covers(all_a_spec, ['510500.SH']) == set(), \
+            'all_a bundle 不应覆盖 ETF 51x'
+        assert _bundle_covers(etf_spec, ['510500.SH']) == {'510500.SH'}
+        assert _bundle_covers(etf_spec, ['600000.SH']) == set(), \
+            'etf bundle 不应覆盖 A 股 60x'
+        assert _bundle_covers(cb_spec, ['113021.SH']) == {'113021.SH'}
+        assert _bundle_covers(cb_spec, ['600000.SH']) == set()
+        print('OK: prefix 隔离 (all_a/etf/cb 三域互不交叉) — KNOWN_PREBUILT 注册项行为正确')
+    else:
+        print('SKIP: prefix 隔离 (KNOWN_PREBUILT 缺 etf/cb/all_a 注册)')
+
+    # ---- 6) 跨 bundle 合并 — 仅当本地同时有 ≥2 个 bundle 时测 (否则 SKIP) ----
+    bundle_specs_real = [s for s in _KP if s.kind == 'bundle' and (cache_dir / s.path_template).exists()]
+    if len(bundle_specs_real) >= 2:
+        # 拿每个 bundle 的真实 stockCode, 跨 bundle 拼 universe
+        all_codes: list[str] = []
+        for s in bundle_specs_real[:2]:
+            df_b = pd.read_parquet(cache_dir / s.path_template, columns=['stockCode'])
+            all_codes.append(df_b['stockCode'].iloc[0])  # 每 bundle 取第一个
+        df, src = resolve(all_codes, '20240101', '20241231', data_cache_dir=cache_dir)
+        if df is not None:
+            print(f'OK: 跨 bundle 合并 (本地 {len(bundle_specs_real)} bundle → {df["symbol"].nunique()} 唯一标, src={src})')
+        else:
+            print(f'SKIP: 跨 bundle 合并 (resolve 返 None, src={src})')
+    else:
+        print(f'SKIP: 跨 bundle 合并 (本地仅 {len(bundle_specs_real)} 个 bundle 真存在, 不足 2 个无法测跨 bundle)')
+
+    # ---- 7) 真实本地无 etf/cb bundle → None (验证 KNOWN_PREBUILT 注册但不误命中) ----
+    df, src = resolve(['510500.SH'], '20240101', '20241231', data_cache_dir=cache_dir)
+    if (cache_dir / '__bundle__etf_D.parquet').exists():
+        print(f'SKIP: 本地已有 etf bundle, 此 case 应改走 ETF 命中 ({src})')
+    else:
+        assert df is None, f'本地无 etf bundle 应 None, 实际 src={src}'
+        print('OK: 本地无 etf bundle → (None, None) — fallback market_mod')
+    df, src = resolve(['113021.SH'], '20240101', '20241231', data_cache_dir=cache_dir)
+    if (cache_dir / '__bundle__cb_D.parquet').exists():
+        print(f'SKIP: 本地已有 cb bundle, 此 case 应改走 CB 命中 ({src})')
+    else:
+        assert df is None, f'本地无 cb bundle 应 None, 实际 src={src}'
+        print('OK: 本地无 cb bundle → (None, None) — fallback market_mod')
+
+
+if __name__ == '__main__':
+    _selfcheck()