ftdt-quant-lab/data/store.py

"""
Parquet-based raw message storage with background writer.

Messages are partitioned by channel/date/ and stored as Parquet files.
Thread-safe: collectors push dicts to a queue, a background thread flushes
to disk periodically.

Schema per row:
  exchange_ts  int64    — exchange timestamp (ms since epoch)
  local_ts     float64  — wall clock at message receipt (seconds since epoch)
  channel      str      — e.g. 'l2book', 'trades', 'funding', 'mark', 'oi'
  coin         str      — e.g. 'BTC', 'ETH'
  payload      bytes    — gzipped JSON blob of the raw message

Usage:
    store = RawMessageStore(data_dir="/data/ftdt-raw")
    store.start()
    store.push(channel="l2book", coin="BTC", exchange_ts=..., payload={...})
    ...
    store.stop()
"""

from __future__ import annotations

import gzip
import json
import logging
import os
import queue
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import pyarrow as pa
import pyarrow.parquet as pq

logger = logging.getLogger(__name__)

SCHEMA = pa.schema([
    pa.field("exchange_ts", pa.int64()),
    pa.field("local_ts", pa.float64()),
    pa.field("channel", pa.string()),
    pa.field("coin", pa.string()),
    pa.field("payload", pa.binary()),
])


class RawMessageStore:
    """Thread-safe Parquet store for raw market data messages."""

    def __init__(
        self,
        data_dir: str = "data/raw",
        flush_interval_sec: float = 5.0,
        max_queue_size: int = 500_000,
        compression: str = "zstd",
        cleanup_days: int = 30,
    ):
        self._data_dir = Path(data_dir)
        self._flush_interval = flush_interval_sec
        self._cleanup_days = cleanup_days
        self._compression = compression
        self._queue: queue.Queue = queue.Queue(maxsize=max_queue_size)
        self._writer_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._buffer: dict[str, list[dict]] = {}
        self._total_written = 0
        self._lock = threading.Lock()

    @property
    def total_written(self) -> int:
        return self._total_written

    def start(self):
        if self._writer_thread and self._writer_thread.is_alive():
            return
        self._stop_event.clear()
        self._data_dir.mkdir(parents=True, exist_ok=True)
        self._writer_thread = threading.Thread(target=self._flush_loop, daemon=True, name="raw-store-writer")
        self._writer_thread.start()
        logger.info("RawMessageStore started (%s)", self._data_dir)

    def stop(self):
        self._stop_event.set()
        if self._writer_thread:
            self._writer_thread.join(timeout=10)
        self._flush_all()
        logger.info("RawMessageStore stopped (%d total written)", self._total_written)

    def push(
        self,
        channel: str,
        coin: str,
        exchange_ts: int,
        payload: dict,
    ):
        """Enqueue a raw message. Non-blocking — drops if queue full."""
        local_ts = time.time()
        try:
            self._queue.put_nowait({
                "exchange_ts": exchange_ts,
                "local_ts": local_ts,
                "channel": channel,
                "coin": coin,
                "payload": payload,
            })
        except queue.Full:
            logger.warning("Store queue full — dropping message (channel=%s coin=%s)", channel, coin)

    # ── internals ───────────────────────────────────────────────

    def _flush_loop(self):
        while not self._stop_event.is_set():
            self._drain_queue()
            self._stop_event.wait(self._flush_interval)
        self._drain_queue()

    def _drain_queue(self):
        drained = 0
        while True:
            try:
                msg = self._queue.get_nowait()
                key = self._partition_key(msg["channel"], msg["coin"])
                with self._lock:
                    self._buffer.setdefault(key, []).append(msg)
                drained += 1
            except queue.Empty:
                break
        if drained:
            self._flush_all()

    def _flush_all(self):
        with self._lock:
            if not self._buffer:
                return
            for key, rows in list(self._buffer.items()):
                if not rows:
                    continue
                self._write_partition(key, rows)
                self._total_written += len(rows)
                self._buffer[key] = []

    def _partition_key(self, channel: str, coin: str) -> str:
        now = datetime.now(timezone.utc)
        return f"{channel}/{coin.upper()}/{now.strftime('%Y-%m-%d')}"

    def _write_partition(self, key: str, rows: list[dict]):
        out_path = self._data_dir / f"{key}.parquet"
        out_path.parent.mkdir(parents=True, exist_ok=True)

        columns = {
            "exchange_ts": [r["exchange_ts"] for r in rows],
            "local_ts": [r["local_ts"] for r in rows],
            "channel": [r["channel"] for r in rows],
            "coin": [r["coin"] for r in rows],
            "payload": [gzip.compress(json.dumps(r["payload"], default=str).encode()) for r in rows],
        }
        table = pa.table(columns, schema=SCHEMA)

        if out_path.exists():
            existing = pq.read_table(out_path)
            table = pa.concat_tables([existing, table])

        pq.write_table(
            table,
            out_path,
            compression=self._compression,
        )


# ── Read helpers ───────────────────────────────────────────────

def read_range(
    data_dir: str,
    channel: str,
    coin: str,
    start_date: str,
    end_date: str,
) -> list[dict]:
    """Read stored messages for a channel/coin/date range. Returns decoded dicts."""
    root = Path(data_dir)
    results = []
    from datetime import date, timedelta

    s = date.fromisoformat(start_date)
    e = date.fromisoformat(end_date)
    current = s
    while current <= e:
        date_str = current.isoformat()
        fpath = root / channel / coin / f"{date_str}.parquet"
        if fpath.exists():
            table = pq.read_table(fpath)
            for i in range(table.num_rows):
                payload_bytes = table["payload"][i].as_py()
                payload = json.loads(gzip.decompress(payload_bytes))
                row = {
                    "exchange_ts": table["exchange_ts"][i].as_py(),
                    "local_ts": table["local_ts"][i].as_py(),
                    "channel": table["channel"][i].as_py(),
                    "coin": table["coin"][i].as_py(),
                    "payload": payload,
                }
                results.append(row)
        current += timedelta(days=1)
    return sorted(results, key=lambda r: r["exchange_ts"] or 0)