Yuxuan Yan
116 lines
4.0 KiB
Python
116 lines
4.0 KiB
Python
"""Derived-data computation and validation."""
|
||
|
||
import csv
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
import pandas as pd
|
||
from pandas.api.types import is_bool_dtype, is_numeric_dtype
|
||
|
||
from pipeline.common.schema import DERIVED_KEY_COLUMNS
|
||
from pipeline.derived.registry import get_derived
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def validate_derived_frame(derived: pd.DataFrame) -> pd.DataFrame:
|
||
"""Validate and normalize a daily derived-data frame.
|
||
|
||
A valid derived frame is keyed by unique ``symbol_id,date`` rows and has at
|
||
least one numeric value column beyond those keys. Dates are normalized to
|
||
daily timestamps before duplicate-key checks.
|
||
"""
|
||
duplicated = derived.columns[derived.columns.duplicated()].tolist()
|
||
if duplicated:
|
||
raise ValueError(f"Derived data has duplicate columns: {duplicated}")
|
||
|
||
missing = [col for col in DERIVED_KEY_COLUMNS if col not in derived.columns]
|
||
if missing:
|
||
raise ValueError(f"Derived data missing required columns: {missing}")
|
||
|
||
out = derived.copy()
|
||
out["date"] = pd.to_datetime(out["date"]).dt.normalize()
|
||
|
||
if out.duplicated(DERIVED_KEY_COLUMNS).any():
|
||
raise ValueError("Derived data has duplicate symbol_id,date rows")
|
||
|
||
value_cols = [col for col in out.columns if col not in DERIVED_KEY_COLUMNS]
|
||
if not value_cols:
|
||
raise ValueError("Derived data must include at least one value column")
|
||
|
||
non_numeric = [
|
||
col
|
||
for col in value_cols
|
||
if is_bool_dtype(out[col]) or not is_numeric_dtype(out[col])
|
||
]
|
||
if non_numeric:
|
||
raise ValueError(f"Derived data value columns must be numeric: {non_numeric}")
|
||
|
||
out = out[DERIVED_KEY_COLUMNS + value_cols].copy()
|
||
return out.sort_values(DERIVED_KEY_COLUMNS).reset_index(drop=True)
|
||
|
||
|
||
def compute_derived(
|
||
derived_type: str,
|
||
daily: pd.DataFrame | None = None,
|
||
minute: pd.DataFrame | None = None,
|
||
**params,
|
||
) -> pd.DataFrame:
|
||
"""Compute one registered derived-data plugin."""
|
||
if daily is None and minute is None:
|
||
raise ValueError("Derived data computation requires --daily-path or --minute-path")
|
||
|
||
derived = get_derived(derived_type, **params)
|
||
result = validate_derived_frame(derived.compute(daily=daily, minute=minute))
|
||
value_cols = [col for col in result.columns if col not in DERIVED_KEY_COLUMNS]
|
||
logger.info(
|
||
"Derived data '%s' (%r): %d symbols × %d dates, columns=%s",
|
||
derived_type,
|
||
derived,
|
||
result["symbol_id"].nunique(),
|
||
result["date"].nunique(),
|
||
value_cols,
|
||
)
|
||
return result
|
||
|
||
|
||
def read_derived_frame(path: str | Path) -> pd.DataFrame:
|
||
"""Read and validate one derived CSV/parquet file or parquet dataset."""
|
||
path = Path(path)
|
||
if path.suffix.lower() == ".csv":
|
||
return validate_derived_frame(_read_csv_with_duplicate_header_check(path))
|
||
return validate_derived_frame(pd.read_parquet(path))
|
||
|
||
|
||
def read_derived_frames(derived_paths: Iterable[str | Path]) -> list[pd.DataFrame]:
|
||
"""Read and validate derived-data files."""
|
||
return [read_derived_frame(path) for path in derived_paths]
|
||
|
||
|
||
def write_derived_frame(
|
||
derived: pd.DataFrame,
|
||
derived_name: str,
|
||
output_dir: str | Path = "derived",
|
||
) -> Path:
|
||
"""Validate and write derived data to ``{output_dir}/{derived_name}.pq``."""
|
||
result = validate_derived_frame(derived)
|
||
output_dir = Path(output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
out_path = output_dir / f"{derived_name}.pq"
|
||
result.to_parquet(out_path, index=False)
|
||
return out_path
|
||
|
||
|
||
def _read_csv_with_duplicate_header_check(path: Path) -> pd.DataFrame:
|
||
with path.open(newline="") as fh:
|
||
reader = csv.reader(fh)
|
||
try:
|
||
header = next(reader)
|
||
except StopIteration as exc:
|
||
raise ValueError("CSV input is empty") from exc
|
||
duplicated = sorted({col for col in header if header.count(col) > 1})
|
||
if duplicated:
|
||
raise ValueError(f"Derived data has duplicate columns: {duplicated}")
|
||
return pd.read_csv(path)
|