wehub-resource-sync
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""ETF universe loader."""
|
|
|
|
import polars as pl
|
|
|
|
from data.exceptions import DataNotFoundError
|
|
from utils import ML4T_DATA_PATH
|
|
from utils.data_quality import apply_max_symbols
|
|
|
|
|
|
def list_etfs() -> list[str]:
|
|
"""List ETF symbols available in the local data store.
|
|
|
|
Returns:
|
|
Sorted list of ETF tickers (e.g., ``["ACWI", "AGG", ..., "XLK"]``).
|
|
|
|
Raises:
|
|
DataNotFoundError: If ``etfs/etf_universe.parquet`` is missing.
|
|
|
|
Example:
|
|
>>> list_etfs()[:3]
|
|
['ACWI', 'ACWX', 'AGG']
|
|
"""
|
|
path = ML4T_DATA_PATH / "etfs" / "market" / "etf_universe.parquet"
|
|
if not path.exists():
|
|
raise DataNotFoundError(
|
|
dataset_name="ETF Universe",
|
|
path=path,
|
|
download_script="data/etfs/market/download.py",
|
|
readme="data/etfs/README.md",
|
|
)
|
|
return pl.scan_parquet(path).select("symbol").unique().collect().to_series().sort().to_list()
|
|
|
|
|
|
def load_etfs(
|
|
symbols: list[str] | None = None,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
max_symbols: int = 0,
|
|
) -> pl.DataFrame:
|
|
"""Load ETF universe for momentum case study.
|
|
|
|
Args:
|
|
symbols: Optional list of symbols to filter (e.g., ["SPY", "QQQ"])
|
|
start_date: Optional start date (YYYY-MM-DD format)
|
|
end_date: Optional end date (YYYY-MM-DD format)
|
|
max_symbols: Limit to N random symbols (0 = all). Seed-deterministic.
|
|
|
|
Returns:
|
|
DataFrame with columns: timestamp, symbol, open, high, low, close, volume
|
|
"""
|
|
path = ML4T_DATA_PATH / "etfs" / "market" / "etf_universe.parquet"
|
|
if not path.exists():
|
|
raise DataNotFoundError(
|
|
dataset_name="ETF Universe",
|
|
path=path,
|
|
download_script="data/etfs/market/download.py",
|
|
readme="data/etfs/README.md",
|
|
)
|
|
|
|
lf = pl.scan_parquet(path)
|
|
ts_type = lf.collect_schema()["timestamp"]
|
|
|
|
# Apply filters lazily (parquet pushdown / row-group pruning)
|
|
if symbols:
|
|
lf = lf.filter(pl.col("symbol").is_in(symbols))
|
|
if start_date:
|
|
lit = (
|
|
pl.lit(start_date).str.to_date()
|
|
if ts_type == pl.Date
|
|
else pl.lit(start_date).str.to_datetime()
|
|
)
|
|
lf = lf.filter(pl.col("timestamp") >= lit)
|
|
if end_date:
|
|
lit = (
|
|
pl.lit(end_date).str.to_date()
|
|
if ts_type == pl.Date
|
|
else pl.lit(end_date).str.to_datetime()
|
|
)
|
|
lf = lf.filter(pl.col("timestamp") <= lit)
|
|
|
|
# Normalize daily data to Date type (post-filter so pushdown works on raw type)
|
|
if ts_type != pl.Date:
|
|
lf = lf.with_columns(pl.col("timestamp").cast(pl.Date))
|
|
|
|
return apply_max_symbols(lf.collect(), max_symbols)
|