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.

651 lines
23 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
from contextlib import contextmanager
import os
import threading
from pathlib import Path
from typing import Any
from uuid import uuid4
import pyarrow as pa
import pyarrow.parquet as pq
try:
import fcntl
except ImportError: # pragma: no cover - Windows fallback
fcntl = None
_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"}
_PATH_WRITE_LOCKS_GUARD = threading.Lock()
_PATH_WRITE_LOCKS: dict[str, threading.Lock] = {}
def _write_lock_for(path: Path) -> threading.Lock:
key = str(path)
with _PATH_WRITE_LOCKS_GUARD:
lock = _PATH_WRITE_LOCKS.get(key)
if lock is None:
lock = threading.Lock()
_PATH_WRITE_LOCKS[key] = lock
return lock
def _write_table_atomic(table: pa.Table, path: Path) -> None:
tmp = path.with_suffix(f".{uuid4().hex}.tmp")
try:
pq.write_table(table, str(tmp), compression="snappy")
os.replace(str(tmp), str(path))
except Exception:
if tmp.exists():
tmp.unlink(missing_ok=True)
raise
def _lock_path_for(path: Path) -> Path:
return path.with_suffix(f"{path.suffix}.lock")
@contextmanager
def _exclusive_write_guard(path: Path):
"""Serialize writes for one cache file across threads and processes."""
with _write_lock_for(path):
if fcntl is None:
yield
return
lock_path = _lock_path_for(path)
lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+b") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
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
# Newly listed tickers can legitimately have no rows near coverage_start.
# Treat the metadata as corrupt only when the payload is also very sparse;
# this catches intraday-fallback fragments that were written with the
# caller's wide requested range as their cache coverage.
if (first_row - coverage_start_dt).days > 10 and len(rows) < 20:
return False
if (coverage_end_dt - last_row).days > 10:
return False
return True
class IntradayCache:
"""Disk-based cache for intraday bars using Parquet.
Thread-safe for reads; serializes writes per path within and across processes.
Supports arbitrary bar resolution via interval_minutes (default 5 for backward compat).
"""
def __init__(self, cache_dir: str = "data/cache/intraday",
interval_minutes: int = 5) -> None:
self._root = Path(cache_dir)
interval_label = f"{interval_minutes}min"
self._cache_metadata = {
**_CACHE_METADATA,
b"intraday_cache_interval": interval_label.encode(),
}
self._schema_with_metadata = _SCHEMA.with_metadata({
**self._cache_metadata,
_INTRADAY_KIND_KEY: _INTRADAY_KIND_POSITIVE,
})
def _path(self, ticker: str, date: str) -> Path:
return self._root / ticker.upper() / f"{date}.parquet"
def _metadata_status(self, 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 self._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=self._schema_with_metadata)
with _exclusive_write_guard(p):
_write_table_atomic(table, p)
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({
**self._cache_metadata,
_INTRADAY_KIND_KEY: _INTRADAY_KIND_NEGATIVE,
_INTRADAY_NEGATIVE_REASON_KEY: reason.encode(),
})
table = pa.Table.from_pylist([], schema=schema)
with _exclusive_write_guard(p):
_write_table_atomic(table, p)
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 end_date > coverage_end:
return None
effective_start = max(start_date, coverage_start)
try:
table = pq.read_table(
str(p),
filters=[
("date", ">=", effective_start),
("date", "<=", end_date),
],
)
rows = self._normalize_rows(table.to_pylist())
if not _daily_rows_match_coverage_sanity(rows, coverage_start, coverage_end):
coverage_start = rows[0]["date"] if rows else coverage_start
coverage_end = rows[-1]["date"] if rows else coverage_end
if end_date > coverage_end:
return None
# Late-start coverage is valid for IPOs/recent listings when enough
# history exists for rolling features; sparse fragments remain misses.
if start_date < coverage_start and len(rows) < 20:
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 end_date < coverage_start or start_date > coverage_end:
return None, None
effective_start = max(start_date, coverage_start)
effective_end = coverage_end if end_date > coverage_end else end_date
try:
table = pq.read_table(
str(p),
filters=[("date", ">=", effective_start), ("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):
if not rows:
return None, None
coverage_start = rows[0]["date"]
coverage_end = rows[-1]["date"]
if end_date < coverage_start or start_date > coverage_end:
return None, None
effective_start = max(start_date, coverage_start)
effective_end = coverage_end if end_date > coverage_end else end_date
rows = [
row
for row in rows
if effective_start <= str(row.get("date", ""))[:10] <= effective_end
]
if not rows:
return None, None
# Late-start coverage is valid for IPOs/recent listings when enough
# history exists for rolling features; sparse fragments remain misses.
if start_date < coverage_start and len(rows) < 20:
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)
with _exclusive_write_guard(p):
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)]
if rows and not _daily_rows_match_coverage_sanity(rows, coverage_start, coverage_end):
coverage_start = rows[0]["date"]
coverage_end = rows[-1]["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)
_write_table_atomic(table, p)
class LayeredDailyBarCache:
"""Read through a snapshot cache, but write only to a mutable overlay."""
def __init__(self, snapshot_cache: DailyBarCache, overlay_cache: DailyBarCache) -> None:
self._snapshot_cache = snapshot_cache
self._overlay_cache = overlay_cache
def get(
self,
ticker: str,
start_date: str,
end_date: str,
) -> list[dict[str, Any]] | None:
rows = self._overlay_cache.get(ticker, start_date, end_date)
if rows is not None:
return rows
return self._snapshot_cache.get(ticker, start_date, end_date)
def get_with_tail(
self,
ticker: str,
start_date: str,
end_date: str,
) -> tuple[list[dict[str, Any]] | None, str | None]:
rows, tail_start = self._overlay_cache.get_with_tail(ticker, start_date, end_date)
if rows is not None:
return rows, tail_start
return self._snapshot_cache.get_with_tail(ticker, start_date, end_date)
def put(
self,
ticker: str,
start_date: str,
end_date: str,
bars: list[dict[str, Any]],
) -> None:
self._overlay_cache.put(ticker, start_date, end_date, bars)
class ReadOnlyDailyBarCache:
"""Read-only daily cache for immutable snapshot backtests.
The screener treats strict misses as unavailable data instead of repairing
them from mutable Oracle/cache sources.
"""
strict_misses = True
def __init__(self, cache: DailyBarCache) -> None:
self._cache = cache
def get(
self,
ticker: str,
start_date: str,
end_date: str,
) -> list[dict[str, Any]] | None:
return self._cache.get(ticker, start_date, end_date)
def get_with_tail(
self,
ticker: str,
start_date: str,
end_date: str,
) -> tuple[list[dict[str, Any]] | None, str | None]:
return self._cache.get_with_tail(ticker, start_date, end_date)
def put(
self,
ticker: str,
start_date: str,
end_date: str,
bars: list[dict[str, Any]],
) -> None:
return None