You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

498 lines
17 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

"""Disk-based Parquet cache for intraday bar data.
Layout: {cache_dir}/{TICKER}/{YYYY-MM-DD}.parquet
Each file stores one trading day of 5-minute bars for one ticker.
~78 bars × 7 columns ≈ 2-3 KB per file (snappy compressed).
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
_CACHE_METADATA = {
b"intraday_cache_version": b"3", # v3: SIP data (v2 was IEX ~2.5% tape)
b"intraday_cache_source": b"api_v1_alpaca_intraday",
b"intraday_cache_interval": b"5min",
}
_INTRADAY_KIND_KEY = b"intraday_cache_kind"
_INTRADAY_KIND_POSITIVE = b"bars"
_INTRADAY_KIND_NEGATIVE = b"negative"
_INTRADAY_NEGATIVE_REASON_KEY = b"intraday_negative_reason"
_SCHEMA = pa.schema([
pa.field("timestamp", pa.string()),
pa.field("open", pa.float64()),
pa.field("high", pa.float64()),
pa.field("low", pa.float64()),
pa.field("close", pa.float64()),
pa.field("volume", pa.float64()),
pa.field("vwap", pa.float64()),
])
_REQUIRED_COLUMNS = {"timestamp", "open", "high", "low", "close", "volume"}
_SCHEMA_WITH_METADATA = _SCHEMA.with_metadata({
**_CACHE_METADATA,
_INTRADAY_KIND_KEY: _INTRADAY_KIND_POSITIVE,
})
_MIN_VALID_INTRADAY_ROWS = 10
_DAILY_CACHE_STATIC_METADATA = {
b"daily_cache_version": b"1",
b"daily_cache_source": b"api_v1_price_data",
b"daily_cache_interval": b"1d",
}
_DAILY_COVERAGE_START_KEY = b"daily_cache_coverage_start"
_DAILY_COVERAGE_END_KEY = b"daily_cache_coverage_end"
_DAILY_SCHEMA = pa.schema([
pa.field("date", pa.string()),
pa.field("open", pa.float64()),
pa.field("high", pa.float64()),
pa.field("low", pa.float64()),
pa.field("close", pa.float64()),
pa.field("volume", pa.float64()),
])
_DAILY_REQUIRED_COLUMNS = {"date", "open", "high", "low", "close", "volume"}
def _daily_rows_match_coverage_sanity(
rows: list[dict[str, Any]],
coverage_start: str | None,
coverage_end: str | None,
) -> bool:
"""Reject egregiously partial daily caches that claim much wider coverage.
Small mismatches are expected around weekends / holidays because the cache
stores calendar coverage while the rows only contain trading sessions. Large
gaps, however, usually indicate a truncated bulk Oracle response that should
not be reused for feature warmup.
"""
if not rows or coverage_start is None or coverage_end is None:
return True
try:
from datetime import date as _date
first_row = _date.fromisoformat(str(rows[0]["date"])[:10])
last_row = _date.fromisoformat(str(rows[-1]["date"])[:10])
coverage_start_dt = _date.fromisoformat(coverage_start)
coverage_end_dt = _date.fromisoformat(coverage_end)
except Exception:
return True
if (first_row - coverage_start_dt).days > 10:
return False
if (coverage_end_dt - last_row).days > 10:
return False
return True
class IntradayCache:
"""Disk-based cache for 5-minute intraday bars using Parquet.
Thread-safe for reads; uses atomic write (tmp → rename) for writes.
"""
def __init__(self, cache_dir: str = "data/cache/intraday") -> None:
self._root = Path(cache_dir)
def _path(self, ticker: str, date: str) -> Path:
return self._root / ticker.upper() / f"{date}.parquet"
@staticmethod
def _metadata_status(path: Path) -> str | None:
try:
meta = pq.read_metadata(str(path))
arrow_schema = meta.schema.to_arrow_schema()
if not _REQUIRED_COLUMNS.issubset(set(arrow_schema.names)):
return None
schema_meta = arrow_schema.metadata or {}
if not all(schema_meta.get(k) == v for k, v in _CACHE_METADATA.items()):
return None
kind = schema_meta.get(_INTRADAY_KIND_KEY, _INTRADAY_KIND_POSITIVE)
if kind == _INTRADAY_KIND_NEGATIVE:
return "negative" if meta.num_rows == 0 else None
if kind not in (_INTRADAY_KIND_POSITIVE, None):
return None
return "positive" if meta.num_rows >= _MIN_VALID_INTRADAY_ROWS else None
except Exception:
return None
@staticmethod
def is_complete_enough(bars: list[dict[str, Any]]) -> bool:
"""Heuristic: keep only intraday responses with enough bars to be useful."""
return len(bars) >= _MIN_VALID_INTRADAY_ROWS
def has(self, ticker: str, date: str) -> bool:
"""Return True if cached bars exist for ticker on date."""
p = self._path(ticker, date)
if not p.exists():
return False
if not self._metadata_status(p):
p.unlink(missing_ok=True)
return False
return True
def get(self, ticker: str, date: str) -> list[dict[str, Any]] | None:
"""Read cached bars. Returns None on cache miss or read error."""
p = self._path(ticker, date)
if not p.exists():
return None
status = self._metadata_status(p)
if not status:
p.unlink(missing_ok=True)
return None
if status == "negative":
return []
try:
table = pq.read_table(str(p))
return table.to_pylist()
except Exception:
return None
def put(self, ticker: str, date: str, bars: list[dict[str, Any]]) -> None:
"""Write bars to cache using atomic write (tmp → rename).
Silently skips if bars is empty.
"""
if not bars:
return
p = self._path(ticker, date)
p.parent.mkdir(parents=True, exist_ok=True)
# Normalize bars to match schema
rows = []
for b in bars:
rows.append({
"timestamp": str(b.get("timestamp", "")),
"open": float(b.get("open", 0.0)),
"high": float(b.get("high", 0.0)),
"low": float(b.get("low", 0.0)),
"close": float(b.get("close", 0.0)),
"volume": float(b.get("volume", 0.0)),
"vwap": float(b.get("vwap", 0.0)) if b.get("vwap") is not None else 0.0,
})
table = pa.Table.from_pylist(rows, schema=_SCHEMA_WITH_METADATA)
tmp = p.with_suffix(".tmp")
try:
pq.write_table(table, str(tmp), compression="snappy")
os.replace(str(tmp), str(p)) # atomic on POSIX
except Exception:
if tmp.exists():
tmp.unlink(missing_ok=True)
raise
def put_negative(self, ticker: str, date: str, reason: str = "empty_or_sparse") -> None:
"""Cache a stable empty/sparse response to avoid repeating futile API calls."""
p = self._path(ticker, date)
p.parent.mkdir(parents=True, exist_ok=True)
schema = _SCHEMA.with_metadata({
**_CACHE_METADATA,
_INTRADAY_KIND_KEY: _INTRADAY_KIND_NEGATIVE,
_INTRADAY_NEGATIVE_REASON_KEY: reason.encode(),
})
table = pa.Table.from_pylist([], schema=schema)
tmp = p.with_suffix(".tmp")
try:
pq.write_table(table, str(tmp), compression="snappy")
os.replace(str(tmp), str(p))
except Exception:
if tmp.exists():
tmp.unlink(missing_ok=True)
raise
def evict(
self,
ticker: str | None = None,
before_date: str | None = None,
) -> int:
"""Remove cached files. Returns count of files removed.
Args:
ticker: If set, only evict this ticker's files.
before_date: If set, only evict files for dates < this (YYYY-MM-DD).
"""
removed = 0
if ticker:
ticker_dir = self._root / ticker.upper()
if not ticker_dir.exists():
return 0
dirs_to_scan = [ticker_dir]
else:
if not self._root.exists():
return 0
dirs_to_scan = [d for d in self._root.iterdir() if d.is_dir()]
for d in dirs_to_scan:
for f in d.glob("*.parquet"):
if before_date and f.stem >= before_date:
continue
f.unlink(missing_ok=True)
removed += 1
return removed
def available_dates(
self,
ticker: str,
start_date: str | None = None,
end_date: str | None = None,
) -> list[str]:
"""List cache dates for one ticker within an optional date range.
Invalid cache files are discarded as they are encountered. Negative cache
entries are included so callers can preserve the original trading-day
shape while deciding how to handle empty responses.
"""
ticker_dir = self._root / ticker.upper()
if not ticker_dir.exists():
return []
dates: list[str] = []
for path in ticker_dir.glob("*.parquet"):
date_str = path.stem
if start_date and date_str < start_date:
continue
if end_date and date_str > end_date:
continue
if not self._metadata_status(path):
path.unlink(missing_ok=True)
continue
dates.append(date_str)
return sorted(dates)
def stats(self) -> dict[str, Any]:
"""Return cache statistics."""
if not self._root.exists():
return {"total_files": 0, "total_bytes": 0, "tickers": 0}
total_files = 0
total_bytes = 0
tickers = set()
dates: list[str] = []
for ticker_dir in self._root.iterdir():
if not ticker_dir.is_dir():
continue
tickers.add(ticker_dir.name)
for f in ticker_dir.glob("*.parquet"):
total_files += 1
total_bytes += f.stat().st_size
dates.append(f.stem)
return {
"total_files": total_files,
"total_bytes": total_bytes,
"total_mb": round(total_bytes / 1_048_576, 2),
"tickers": len(tickers),
"date_min": min(dates) if dates else None,
"date_max": max(dates) if dates else None,
}
class DailyBarCache:
"""Disk-based cache for daily OHLCV bars using one Parquet file per ticker."""
def __init__(self, cache_dir: str = "data/cache/daily") -> None:
self._root = Path(cache_dir)
def _path(self, ticker: str) -> Path:
return self._root / f"{ticker.upper()}.parquet"
@staticmethod
def _read_schema_metadata(path: Path) -> dict[bytes, bytes] | None:
try:
meta = pq.read_metadata(str(path))
if meta.num_rows < 0:
return None
return meta.schema.to_arrow_schema().metadata or {}
except Exception:
return None
@classmethod
def _metadata_valid(cls, path: Path) -> bool:
try:
meta = pq.read_metadata(str(path))
if meta.num_rows < 0:
return False
arrow_schema = meta.schema.to_arrow_schema()
if not _DAILY_REQUIRED_COLUMNS.issubset(set(arrow_schema.names)):
return False
schema_meta = arrow_schema.metadata or {}
if not all(schema_meta.get(k) == v for k, v in _DAILY_CACHE_STATIC_METADATA.items()):
return False
return (
schema_meta.get(_DAILY_COVERAGE_START_KEY) is not None
and schema_meta.get(_DAILY_COVERAGE_END_KEY) is not None
)
except Exception:
return False
@classmethod
def _coverage_from_metadata(cls, path: Path) -> tuple[str | None, str | None]:
schema_meta = cls._read_schema_metadata(path) or {}
start = schema_meta.get(_DAILY_COVERAGE_START_KEY)
end = schema_meta.get(_DAILY_COVERAGE_END_KEY)
return (
start.decode() if start else None,
end.decode() if end else None,
)
@staticmethod
def _normalize_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
normalized: list[dict[str, Any]] = []
for b in rows:
normalized.append(
{
"date": str(b.get("date", ""))[:10],
"open": float(b.get("open", 0.0) or 0.0),
"high": float(b.get("high", 0.0) or 0.0),
"low": float(b.get("low", 0.0) or 0.0),
"close": float(b.get("close", 0.0) or 0.0),
"volume": float(b.get("volume", 0.0) or 0.0),
}
)
return normalized
def get(
self,
ticker: str,
start_date: str,
end_date: str,
) -> list[dict[str, Any]] | None:
"""Read cached daily bars for an exact date range if covered."""
p = self._path(ticker)
if not p.exists():
return None
if not self._metadata_valid(p):
p.unlink(missing_ok=True)
return None
coverage_start, coverage_end = self._coverage_from_metadata(p)
if coverage_start is None or coverage_end is None:
p.unlink(missing_ok=True)
return None
if start_date < coverage_start or end_date > coverage_end:
return None
try:
table = pq.read_table(
str(p),
filters=[
("date", ">=", start_date),
("date", "<=", end_date),
],
)
rows = self._normalize_rows(table.to_pylist())
if not _daily_rows_match_coverage_sanity(rows, coverage_start, coverage_end):
p.unlink(missing_ok=True)
return None
return rows
except Exception:
return None
def get_with_tail(
self,
ticker: str,
start_date: str,
end_date: str,
) -> tuple[list[dict[str, Any]] | None, str | None]:
"""Like get(), but supports partial hits when end_date > coverage_end.
Returns:
(bars, None) — full hit (end_date <= coverage_end)
(bars, tail_start) — partial hit; bars cover start_date..coverage_end,
tail_start is the first calendar day to fetch from Oracle
(None, None) — true miss (no file, or start_date not covered)
"""
from datetime import date as _date, timedelta as _td
p = self._path(ticker)
if not p.exists():
return None, None
if not self._metadata_valid(p):
p.unlink(missing_ok=True)
return None, None
coverage_start, coverage_end = self._coverage_from_metadata(p)
if coverage_start is None or coverage_end is None:
p.unlink(missing_ok=True)
return None, None
if start_date < coverage_start:
return None, None # true miss: need earlier data than cached
effective_end = coverage_end if end_date > coverage_end else end_date
try:
table = pq.read_table(
str(p),
filters=[("date", ">=", start_date), ("date", "<=", effective_end)],
)
rows = self._normalize_rows(table.to_pylist())
except Exception:
return None, None
if not _daily_rows_match_coverage_sanity(rows, coverage_start, coverage_end):
p.unlink(missing_ok=True)
return None, None
if end_date <= coverage_end:
return rows, None # full hit
# Partial hit — caller must fetch from tail_start to end_date
tail_start = (_date.fromisoformat(coverage_end) + _td(days=1)).isoformat()
return rows, tail_start
def put(
self,
ticker: str,
start_date: str,
end_date: str,
bars: list[dict[str, Any]],
) -> None:
"""Merge a fetched daily-bar range into the ticker cache."""
p = self._path(ticker)
p.parent.mkdir(parents=True, exist_ok=True)
normalized_new = self._normalize_rows(bars)
merged_by_date: dict[str, dict[str, Any]] = {}
coverage_start = start_date
coverage_end = end_date
if p.exists() and self._metadata_valid(p):
try:
existing = pq.read_table(str(p)).to_pylist()
for row in self._normalize_rows(existing):
merged_by_date[row["date"]] = row
existing_start, existing_end = self._coverage_from_metadata(p)
if existing_start:
coverage_start = min(coverage_start, existing_start)
if existing_end:
coverage_end = max(coverage_end, existing_end)
except Exception:
p.unlink(missing_ok=True)
merged_by_date = {}
elif p.exists():
p.unlink(missing_ok=True)
for row in normalized_new:
merged_by_date[row["date"]] = row
rows = [merged_by_date[d] for d in sorted(merged_by_date)]
metadata = dict(_DAILY_CACHE_STATIC_METADATA)
metadata[_DAILY_COVERAGE_START_KEY] = coverage_start.encode()
metadata[_DAILY_COVERAGE_END_KEY] = coverage_end.encode()
schema = _DAILY_SCHEMA.with_metadata(metadata)
table = pa.Table.from_pylist(rows, schema=schema)
tmp = p.with_suffix(".tmp")
try:
pq.write_table(table, str(tmp), compression="snappy")
os.replace(str(tmp), str(p))
except Exception:
if tmp.exists():
tmp.unlink(missing_ok=True)
raise