pfeed

41def get_config() -> PFeedConfig:
42    """Lazy singleton - only creates config when first called.
43    Also loads the .env file.
44    """
45    global _config
46    if _config is None:
47        _config = PFeedConfig()
48    return _config

55def configure(
56    data_path: str | Path | None = None,
57    log_path: str | Path | None = None,
58    cache_path: str | Path | None = None,
59    show_progress_bar: bool | None = None,
60    allow_timestamp_precision_loss: bool | None = None,
61    stream_latency_log_interval: float | None = None,
62    stream_optimize_interval: float | None = None,
63    persist: bool = False,
64) -> PFeedConfig:
65    """
66    Configures the global config object.
67    Args:
68        data_path: Path to the data directory.
69        log_path: Path to the log directory.
70        cache_path: Path to the cache directory.
71        show_progress_bar: Whether pfeed progress bars are displayed.
72        allow_timestamp_precision_loss: Whether IOs that store timestamps in a coarser unit
73            (e.g. DeltaLakeIO, microseconds) may truncate finer ones (e.g. nanoseconds), with a warning.
74            If False, such writes raise. Defaults to False.
75        stream_latency_log_interval: Seconds between a stream's latency logs (at DEBUG level), e.g. 60.
76            Defaults to None, i.e. latency is not tracked.
77        stream_optimize_interval: Seconds between optimize() calls on the table IOs (e.g. DuckLakeIO) a feed streams to,
78            e.g. to move DuckLake's inlined rows out of its catalog into parquet files. 0 runs it only once, when the streams end.
79            Defaults to 1800.
80        persist: If True, the config will be saved to the config file.
81    """
82    config = get_config()
83    config_dict = config.to_dict()
84
85    # Apply updates for non-None values
86    for k in config_dict:
87        v = locals().get(k)
88        if v is not None:
89            if "_path" in k:
90                v = Path(v)
91            setattr(config, k, v)
92
93    config.ensure_dirs()
94
95    if persist:
96        config.save()
97
98    return config

101def configure_logging(
102    logging_config: dict[str, Any] | None = None, debug: bool = False
103) -> dict[str, Any]:
104    return kit_logging.configure_logging(
105        get_config(), overrides=logging_config, debug=debug
106    )

 33class DataEngine:
 34    """Runs the streams or replays of several feeds together and sends out their output, e.g. for pfund.
 35
 36    The output is each streamed message after the feeds' transformations (e.g. a resampled bar), with or without Ray, sent to:
 37    - the callback, if given
 38    - the iteration, i.e. `async for data in engine`
 39    - the subscribers of the engine's XPUB, if setup_publishing() was called
 40    """
 41
 42    QUEUE_MAXSIZE: ClassVar[int] = 1000
 43
 44    def __init__(
 45        self,
 46        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 47        isolate_failures: bool = True,
 48    ):
 49        """
 50        Args:
 51            callback: Sync or async callable called with each output, i.e. a message after the feeds' transformations.
 52                It is awaited before the next output, so a slow callback delays every feed, offload heavy work, e.g. to a task or a thread.
 53            isolate_failures: If True (default), a failed feed ends alone, the others keep running.
 54                If False, the first failed feed ends all the others.
 55        """
 56        self._logger = logging.getLogger("pfeed")
 57        self._callback = callback
 58        self._isolate_failures = isolate_failures
 59        self._feeds: list[StreamingFeed] = []
 60        self._is_running = False
 61        # XPUB publishing the output to other processes, only set up by setup_publishing()
 62        self._msg_queue: ZeroMQ | None = None
 63        # PULL receiving the Ray workers' output, set up by run_async() when a feed uses Ray
 64        self._worker_msg_queue: ZeroMQ | None = None
 65        self._worker_thread: threading.Thread | None = None
 66        self._stop_worker_thread = threading.Event()
 67        # NOTE: only set while replaying, merges the replays of all feeds by time
 68        self._replay_clock: ReplayClock | None = None
 69        # NOTE: only set while iterating the engine with __aiter__()
 70        self._iter_queue: asyncio.Queue[Message] | None = None
 71        self._zmq_routes: dict[
 72            DataFlow[Any, Any], tuple[ZeroMQChannel, ZeroMQTopic]
 73        ] = {}
 74        # NOTE: only set while running, so that end() can cancel the feeds from another thread
 75        self._loop: asyncio.AbstractEventLoop | None = None
 76        self._feed_tasks: list[asyncio.Task[RunResult]] = []
 77
 78    def is_running(self) -> bool:
 79        return self._is_running
 80
 81    @property
 82    def feeds(self) -> list[BaseFeed[Any, Any, Any, Any]]:
 83        return cast("list[BaseFeed[Any, Any, Any, Any]]", self._feeds)
 84
 85    def add_feed(
 86        self,
 87        data_source: str,
 88        data_category: DataCategory | str = DataCategory.MARKET_DATA,
 89        num_workers: int | None = None,
 90    ) -> BaseFeed[Any, Any, Any, Any]:
 91        """Creates a feed in pipeline mode, to call stream() or replay() on, run by the engine.
 92
 93        Args:
 94            num_workers: Number of Ray workers for the feed's streams, None to run them without Ray.
 95                Feeds with and without Ray can be mixed.
 96        """
 97        from pfeed.base import get_feed
 98        from pfeed.streaming.mixin import StreamingMixin
 99
100        if self.is_running():
101            raise RuntimeError("cannot add a feed while the data engine is running")
102        feed: BaseFeed[Any, Any, Any, Any] = get_feed(
103            data_source, data_category, pipeline_mode=True, num_workers=num_workers
104        )
105        if not isinstance(feed, StreamingMixin) or not (
106            feed._supports("stream") or feed._supports("replay")
107        ):
108            raise ValueError(
109                f"{feed.name} supports neither stream() nor replay(), the data engine only runs streams and replays"
110            )
111        feed._set_engine(self)
112        self._feeds.append(feed)
113        return feed
114
115    def setup_publishing(
116        self,
117        zmq_url: str | None = None,
118        zmq_sender_port: int | None = None,
119        io_threads: int = 1,
120    ) -> None:
121        """Sets up the XPUB publishing the output to other processes, e.g. pfund's.
122
123        Each output is sent on its dataflow's channel, see BaseDataModel.to_zmq_route().
124        The XPUB is closed when the run ends.
125
126        Args:
127            zmq_url: e.g. "tcp://localhost" (default).
128            zmq_sender_port: The XPUB's port, a random free one if None, see `engine._msg_queue.get_ports_in_use()`.
129        """
130        import zmq
131
132        from pfeed.streaming.zeromq import ZeroMQ
133
134        if self._msg_queue is not None:
135            raise RuntimeError("Publishing is already set up, cannot set it up again")
136        self._msg_queue = ZeroMQ(
137            name="data_engine",
138            logger=self._logger,
139            io_threads=io_threads,
140            sender_type=zmq.XPUB,
141        )
142        self._msg_queue.bind(
143            self._msg_queue.sender,
144            port=zmq_sender_port,
145            url=zmq_url or ZeroMQ.DEFAULT_URL,
146        )
147
148    def _get_worker_ports(self) -> list[int]:
149        """The ports of the PULL the Ray workers push their output to."""
150        assert self._worker_msg_queue is not None, (
151            "the data engine's PULL for the Ray workers is not set up"
152        )
153        return self._worker_msg_queue.get_ports_in_use(self._worker_msg_queue.receiver)
154
155    def _validate_feeds(self) -> None:
156        if not self._feeds:
157            raise RuntimeError("No feeds to run, add one with add_feed()")
158        if idle := [feed.name for feed in self._feeds if not feed.is_streaming()]:
159            raise RuntimeError(
160                f"{idle} have nothing to run, call stream() or replay() on them first"
161            )
162        envs = {
163            feed: {request.env for request in feed._requests} for feed in self._feeds
164        }
165        if len(set().union(*envs.values())) > 1:
166            details = ", ".join(
167                f"{feed.name}={sorted(feed_envs)}" for feed, feed_envs in envs.items()
168            )
169            raise ValueError(
170                f"cannot run feeds of different envs in one data engine: {details}"
171            )
172        if self._is_replaying():
173            replay_paces = {
174                request.replay_pace
175                for feed in self._feeds
176                for request in feed._requests
177            }
178            if len(replay_paces) > 1:
179                raise ValueError(
180                    f"the replays of a data engine share one pace, got different replay_pace: {replay_paces}"
181                )
182
183    def _register_codecs(self, msg_queue: ZeroMQ) -> None:
184        """Encodes and decodes each dataflow's channel with its data model's wire format, e.g. Bars for market data."""
185        for dataflow, (channel, _topic) in self._zmq_routes.items():
186            data_model = dataflow.data_model
187            msg_queue.register_codec(
188                channel,
189                data_model.StreamMessage,
190                data_model.enc_hook,
191                data_model.dec_hook,
192            )
193
194    def _is_replaying(self) -> bool:
195        # NOTE: no mixed envs, see _validate_feeds()
196        return self._feeds[0]._requests[0].env == Environment.BACKTEST
197
198    # --- output ---
199
200    async def _forward_dataflow_message(
201        self, dataflow: DataFlow[Any, Any], data: Message
202    ) -> None:
203        """Called by a feed's dataflow (without Ray) with each message's output."""
204        channel, topic = self._zmq_routes[dataflow]
205        await self._emit(channel, topic, data)
206
207    async def _emit(
208        self, channel: ZeroMQChannel, topic: ZeroMQTopic, data: Message
209    ) -> None:
210        """Sends an output to the callback, the iteration and the XPUB's subscribers."""
211        if self._callback:
212            # NOTE: a failed callback skips only this output, with or without Ray
213            try:
214                result = self._callback(data)
215                if inspect.isawaitable(result):
216                    await result
217            except Exception:
218                self._logger.exception(
219                    f"Data Engine callback failed on {channel}, skipping the output:"
220                )
221        if (iter_queue := self._iter_queue) is not None:
222            # NOTE: drop the oldest output instead of blocking, a slow consumer would otherwise stall every feed
223            if iter_queue.full():
224                self._logger.warning(
225                    "data engine iteration queue is full, dropping the oldest output - consider increasing QUEUE_MAXSIZE "
226                    + f"(current: {self.QUEUE_MAXSIZE}) or improving the consumer's speed"
227                )
228                _ = iter_queue.get_nowait()
229            iter_queue.put_nowait(data)
230        if self._msg_queue is not None:
231            # NOTE: a failed send skips only this output, it must not stop the dataflow that produced it
232            try:
233                self._msg_queue.send(channel=channel, topic=topic, data=data)
234            except Exception:
235                self._logger.exception(
236                    f"Data Engine failed to send the output on {channel}, skipping it:"
237                )
238
239    def _setup_worker_messaging(self, inbox: asyncio.Queue[Output]) -> None:
240        """Binds the PULL the Ray workers push their output to, a thread receives from it and puts the output in the inbox."""
241        import zmq
242
243        from pfeed.streaming.zeromq import ZeroMQ, ZeroMQDataChannel
244
245        loop = asyncio.get_running_loop()
246        msg_queue = ZeroMQ(
247            name="data_engine.workers",
248            logger=self._logger,
249            receiver_type=zmq.PULL,
250            receiver_method="bind",
251        )
252        # NOTE: decodes the outputs back into their types, e.g. Bar/Tick for market data, like the outputs without Ray
253        self._register_codecs(msg_queue)
254        msg_queue.bind(msg_queue.receiver)
255        self._worker_msg_queue = msg_queue
256        self._stop_worker_thread.clear()
257
258        def _receive():
259            # NOTE: only this thread uses the PULL after this, a zmq socket isn't thread-safe
260            try:
261                while True:
262                    if (msg := msg_queue.recv()) is None:
263                        # NOTE: stops only once nothing is left to receive, the workers' last outputs can arrive after they have stopped
264                        if self._stop_worker_thread.is_set():
265                            break
266                        continue
267                    channel, topic, data, _msg_ts = msg
268                    # NOTE: each worker sends STOP when it terminates, it is not an output
269                    if channel == ZeroMQDataChannel.signal:
270                        continue
271                    _ = loop.call_soon_threadsafe(
272                        inbox.put_nowait, (channel, topic, data)
273                    )
274            finally:
275                msg_queue.terminate()
276
277        self._worker_thread = threading.Thread(
278            target=_receive, name="data_engine.workers", daemon=True
279        )
280        self._worker_thread.start()
281
282    def _teardown_worker_messaging(self) -> None:
283        self._stop_worker_thread.set()
284        if self._worker_thread is not None:
285            self._worker_thread.join(timeout=10)
286            if self._worker_thread.is_alive():
287                self._logger.warning(
288                    "data engine's thread receiving from the Ray workers did not stop in time"
289                )
290        self._worker_thread = None
291        self._worker_msg_queue = None
292
293    async def _forward_worker_output(self, inbox: asyncio.Queue[Output]) -> None:
294        while True:
295            channel, topic, data = await inbox.get()
296            try:
297                await self._emit(channel, topic, data)
298            # NOTE: keeps forwarding, run_async() waits for the inbox to be emptied when the feeds have ended
299            except Exception:
300                self._logger.exception(
301                    f"Data Engine failed to send the output on {channel}:"
302                )
303            finally:
304                inbox.task_done()
305
306    # --- running ---
307
308    def run(self) -> dict[BaseFeed[Any, Any, Any, Any], RunResult]:
309        """Runs every feed's streams or replays until they end or end() is called, see run_async()."""
310        try:
311            _ = asyncio.get_running_loop()
312        except RuntimeError:
313            # no running event loop, safe to use asyncio.run()
314            pass
315        else:
316            raise RuntimeError(
317                "Cannot call engine.run() from within a running event loop.\n"
318                + "Did you mean to call engine.run_async()?"
319            )
320        return asyncio.run(self.run_async())
321
322    async def run_async(self) -> dict[BaseFeed[Any, Any, Any, Any], RunResult]:
323        """Runs every feed's streams or replays until they end or end() is called.
324
325        Returns:
326            Each feed's result. A feed that raised gets its dataflows failed with the error.
327        """
328        from pfeed.dataflow.result import DataFlowResult, RunResult
329        from pfeed.streaming.replay_clock import ReplayClock
330
331        if self.is_running():
332            raise RuntimeError("Data Engine is already running, cannot run again")
333        self._validate_feeds()
334        self._is_running = True
335        self._loop = asyncio.get_running_loop()
336        inbox: asyncio.Queue[Output] = asyncio.Queue()
337        forwarding: asyncio.Task[None] | None = None
338        # NOTE: kept since a feed clears its dataflows when its run ends
339        dataflows = {feed: feed.dataflows for feed in self._feeds}
340        try:
341            # NOTE: each dataflow's channel, its outputs are sent on it, e.g. to the XPUB's subscribers
342            self._zmq_routes = {
343                dataflow: dataflow.data_model.to_zmq_route()
344                for feed_dataflows in dataflows.values()
345                for dataflow in feed_dataflows
346            }
347            if self._msg_queue is not None:
348                self._register_codecs(self._msg_queue)
349            if any(feed._is_using_ray() for feed in self._feeds):
350                self._setup_worker_messaging(inbox)
351                forwarding = asyncio.create_task(self._forward_worker_output(inbox))
352            if self._is_replaying():
353                self._replay_clock = ReplayClock(
354                    self._feeds[0]._requests[0].replay_pace
355                )
356                # NOTE: every feed joins before any starts, otherwise a feed could replay ahead of one that hasn't joined yet
357                for feed in self._feeds:
358                    self._replay_clock.join(feed)
359            self._feed_tasks = [
360                asyncio.create_task(feed.run_async()) for feed in self._feeds
361            ]
362            try:
363                await self._wait_for_feeds()
364            except asyncio.CancelledError:
365                self._logger.warning("Data Engine was cancelled, ending the feeds...")
366            finally:
367                # NOTE: e.g. cancelled, or a feed failed with isolate_failures=False
368                if pending := [task for task in self._feed_tasks if not task.done()]:
369                    for task in pending:
370                        _ = task.cancel()
371                    _ = await asyncio.wait(pending)
372            results: dict[BaseFeed[Any, Any, Any, Any], RunResult] = {}
373            for feed, task in zip(self._feeds, self._feed_tasks, strict=True):
374                # NOTE: a cancelled feed's dataflows keep their results, a cancelled stream is not a failure
375                if task.cancelled():
376                    results[feed] = RunResult(data=None, dataflows=dataflows[feed])
377                elif (err := task.exception()) is not None:
378                    self._logger.error(f"{feed.name} failed: {err!r}")
379                    for dataflow in dataflows[feed]:
380                        dataflow._result = DataFlowResult.failed(error=err)
381                    results[feed] = RunResult(data=None, dataflows=dataflows[feed])
382                else:
383                    results[feed] = task.result()
384            if failed := [
385                feed.name for feed, result in results.items() if not result.success
386            ]:
387                self._logger.warning(
388                    f"Data Engine has {len(failed)} failed feeds: {failed}"
389                )
390            return results
391        finally:
392            # NOTE: every feed has ended, and with them the Ray workers
393            self._teardown_worker_messaging()
394            if forwarding is not None:
395                # NOTE: lets the thread's last call_soon_threadsafe() puts run, then sends what is left in the inbox
396                await asyncio.sleep(0)
397                if not forwarding.done():
398                    await inbox.join()
399                _ = forwarding.cancel()
400                _ = await asyncio.wait([forwarding])
401            if self._msg_queue is not None:
402                self._msg_queue.terminate()
403                self._msg_queue = None
404            self._replay_clock = None
405            self._zmq_routes.clear()
406            self._feed_tasks = []
407            self._loop = None
408            self._is_running = False
409
410    async def _wait_for_feeds(self) -> None:
411        """Waits until every feed has ended, or until the first failed one if isolate_failures is False."""
412        pending: set[asyncio.Task[RunResult]] = set(self._feed_tasks)
413        while pending:
414            done, pending = await asyncio.wait(
415                pending, return_when=asyncio.FIRST_COMPLETED
416            )
417            if self._isolate_failures:
418                continue
419            for task in done:
420                if not task.cancelled() and (
421                    task.exception() is not None or not task.result().success
422                ):
423                    self._logger.warning(
424                        "a feed has failed, ending the other feeds since isolate_failures=False"
425                    )
426                    return
427
428    def end(self) -> None:
429        """Ends the feeds, run() / run_async() then return their results. Can be called from any thread."""
430        if (loop := self._loop) is None:
431            return
432        self._logger.debug("Data Engine is ending")
433
434        def _cancel_feeds():
435            for task in self._feed_tasks:
436                _ = task.cancel()
437
438        # NOTE: raised if the loop has just closed, i.e. the run has ended
439        with contextlib.suppress(RuntimeError):
440            _ = loop.call_soon_threadsafe(_cancel_feeds)
441
442    async def __aiter__(self) -> AsyncGenerator[Message]:
443        """Runs the feeds and yields each output, like the callback gets it.
444
445        Breaking out of the loop ends the feeds only when the generator is closed, which Python does lazily,
446        wrap it in `contextlib.aclosing()` to end them right at the break.
447        """
448        queue: asyncio.Queue[Message] = asyncio.Queue(maxsize=self.QUEUE_MAXSIZE)
449        self._iter_queue = queue
450        producer = asyncio.create_task(self.run_async())
451        # NOTE: a non-immediate shutdown lets the consumer drain the queued outputs before get() raises QueueShutDown
452        producer.add_done_callback(lambda _: queue.shutdown())
453        try:
454            while True:
455                try:
456                    data = await queue.get()
457                except asyncio.QueueShutDown:
458                    break
459                yield data
460            # re-raises the producer's exception, if any
461            _ = await producer
462        finally:
463            # NOTE: the consumer stopped early, e.g. break or cancelled, end the feeds
464            if not producer.done():
465                _ = producer.cancel()
466                # NOTE: wait() instead of awaiting the producer, so the producer's cancellation isn't raised here, but this task's own is
467                _ = await asyncio.wait([producer])
468            self._iter_queue = None

11def get_feed(
12    data_source: str,
13    data_category: DataCategory | str,
14    pipeline_mode: bool = False,
15    num_workers: int | None = None,
16) -> BaseFeed:
17    from pfeed import registry
18
19    Feed = registry.get_feed(data_source, data_category)
20    # feeds are only created by their source, so build the source and return its feed
21    source = registry.get_source(data_source)(
22        pipeline_mode=pipeline_mode, num_workers=num_workers
23    )
24    return getattr(source, Feed.data_domain.feed_name)

  42class MarketFeed[
  43    SourceT: BaseSource[bool],
  44    RequestT: MarketFeedBaseRequest = MarketFeedBaseRequest,
  45    DataModelT: MarketDataModel = MarketDataModel,
  46    PipelineT: bool = Literal[False],
  47](StreamingMixin, TimeBasedFeed[SourceT, RequestT, DataModelT, PipelineT], ABC):
  48    class Requests:  # ty: ignore[incompatible-attribute-override]
  49        """Request class per verb, a subclass sets a verb to None if it doesn't support it."""
  50
  51        download: ClassVar[type[MarketFeedDownloadRequest] | None] = (
  52            MarketFeedDownloadRequest
  53        )
  54        retrieve: ClassVar[type[MarketFeedRetrieveRequest] | None] = (
  55            MarketFeedRetrieveRequest
  56        )
  57        stream: ClassVar[type[MarketFeedStreamRequest] | None] = MarketFeedStreamRequest
  58        replay: ClassVar[type[MarketFeedReplayRequest] | None] = MarketFeedReplayRequest
  59
  60    DataModel: ClassVar[type[MarketDataModel]] = MarketDataModel  # ty: ignore[incompatible-mutable-override]
  61    data_domain: ClassVar[DataCategory] = DataCategory.MARKET_DATA
  62    _required_methods: ClassVar[dict[str, tuple[str, ...]]] = {
  63        MarketFeedDownloadRequest.extract_type: (
  64            "_download_impl",
  65            "_normalize_downloaded_data",
  66        ),
  67        MarketFeedStreamRequest.extract_type: ("_parse_streaming_message",),
  68    }
  69
  70    @staticmethod
  71    def _standardize_columns(
  72        df: pl.LazyFrame, product: BaseProduct, resolution: Resolution
  73    ) -> pl.LazyFrame:
  74        """Standardizes the columns of a DataFrame.
  75        Adds columns 'product', 'resolution'
  76        """
  77        from pfeed.schemas.market_data_schema import PRODUCT, RESOLUTION
  78
  79        return df.with_columns(
  80            pl.lit(product.name).alias(PRODUCT),
  81            # NOTE: str(resolution) is used, NOT repr(resolution)
  82            pl.lit(str(resolution)).alias(RESOLUTION),
  83        )
  84
  85    @staticmethod
  86    def _filter_columns(df: pl.LazyFrame, product: BaseProduct) -> pl.LazyFrame:
  87        """Keeps only the standard columns of cleaned data, plus the optional ones present."""
  88        from pfeed.schemas.bar_data_schema import CLOSE, HIGH, LOW, OPEN
  89        from pfeed.schemas.market_data_schema import KEY_COLS, VOLUME
  90        from pfeed.schemas.tick_data_schema import PRICE, SIDE
  91
  92        cols = df.collect_schema().names()
  93        is_tick_data = PRICE in cols
  94        if is_tick_data:
  95            value_cols = [PRICE, VOLUME, SIDE]
  96            # side is a standard column even if the source has no trade direction, null means unknown
  97            if SIDE not in cols:
  98                df = df.with_columns(pl.lit(None, dtype=pl.Int8).alias(SIDE))
  99        else:
 100            value_cols = [OPEN, HIGH, LOW, CLOSE, VOLUME]
 101        standard_cols = [*KEY_COLS, *value_cols]
 102
 103        product_specific_cols = []
 104        if product.is_stock() or product.is_etf():
 105            from pfeed.schemas.stock_data_schema import DIVIDENDS, SPLITS
 106
 107            product_specific_cols.extend([DIVIDENDS, SPLITS])
 108
 109        for col in product_specific_cols:
 110            if col in cols:
 111                standard_cols.append(col)
 112        return df.select(standard_cols)
 113
 114    @staticmethod
 115    def _organize_columns(df: pl.LazyFrame) -> pl.LazyFrame:
 116        """Moves the key columns ('date', 'product', 'resolution') to the leftmost side."""
 117        from pfeed.schemas.market_data_schema import KEY_COLS
 118
 119        left_cols = list(KEY_COLS)
 120        current_cols = df.collect_schema().names()
 121        target_cols = left_cols + [c for c in current_cols if c not in left_cols]
 122        if current_cols == target_cols:
 123            return df
 124        return df.select(target_cols)
 125
 126    @staticmethod
 127    def _resample_data(
 128        df: pl.LazyFrame,
 129        product: BaseProduct,
 130        resolution: str | Resolution,
 131        offset: str | None = None,
 132    ) -> pl.LazyFrame:
 133        """Resamples the input lazyframe based on the target resolution.
 134
 135        Args:
 136            df: The input lazyframe to be resampled.
 137            product: Product instance of the data, e.g. created from `'BTC_USDT_PERP'`.
 138                Decides the product-specific columns to keep and aggregate,
 139                e.g. `dividends` and `splits` for stocks and ETFs.
 140            resolution: The target resolution to resample the data to.
 141            offset: Optional polars duration string (e.g., '30m') to shift the resampling window.
 142                For example, with '1h' resolution and '30m' offset, timestamps will be XX:30:00.
 143
 144        Returns:
 145            The resampled lazyframe.
 146        """
 147        if isinstance(resolution, str):
 148            resolution = Resolution(resolution)
 149
 150        from pfeed.schemas.bar_data_schema import CLOSE, HIGH, LOW, OPEN
 151        from pfeed.schemas.market_data_schema import PRODUCT, RESOLUTION, VOLUME
 152        from pfeed.schemas.tick_data_schema import PRICE
 153        from pfeed.schemas.time_based_data_schema import DATE
 154
 155        df_resolution_str = df.select(pl.col(RESOLUTION).first()).collect().item()
 156        df_resolution = Resolution(df_resolution_str)
 157        is_resample_required = resolution < df_resolution
 158        if not is_resample_required:
 159            return df
 160
 161        df = MarketFeed._filter_columns(df, product)
 162
 163        # pfund's timeframe units (s, m, h, d, w, mo, y) are the same as polars' duration units
 164        eresolution = f"{resolution.period}{resolution.timeframe.canonical}"
 165
 166        cols = df.collect_schema().names()
 167        is_tick_data = PRICE in cols
 168
 169        aggs: list[pl.Expr] = []
 170        if is_tick_data:
 171            aggs.extend(
 172                [
 173                    pl.col(PRICE).first().alias(OPEN),
 174                    pl.col(PRICE).max().alias(HIGH),
 175                    pl.col(PRICE).min().alias(LOW),
 176                    pl.col(PRICE).last().alias(CLOSE),
 177                    pl.col(VOLUME).sum().alias(VOLUME),
 178                ]
 179            )
 180        else:
 181            aggs.extend(
 182                [
 183                    pl.col(OPEN).first().alias(OPEN),
 184                    pl.col(HIGH).max().alias(HIGH),
 185                    pl.col(LOW).min().alias(LOW),
 186                    pl.col(CLOSE).last().alias(CLOSE),
 187                    pl.col(VOLUME).sum().alias(VOLUME),
 188                ]
 189            )
 190        aggs.append(pl.col(PRODUCT).first().alias(PRODUCT))
 191
 192        if product.is_stock() or product.is_etf():
 193            from pfeed.schemas.stock_data_schema import DIVIDENDS, SPLITS
 194
 195            if DIVIDENDS in cols:
 196                aggs.append(pl.col(DIVIDENDS).sum().alias(DIVIDENDS))
 197            if SPLITS in cols:
 198                aggs.append(pl.col(SPLITS).product().alias(SPLITS))
 199
 200        aggs.append(pl.len().alias("n_data_points"))
 201
 202        resampled = (
 203            df.sort(DATE)
 204            .group_by_dynamic(
 205                DATE,
 206                every=eresolution,
 207                offset=offset,
 208                label="left",
 209                closed="left",
 210            )
 211            .agg(aggs)
 212            .with_columns(pl.lit(str(resolution)).alias(RESOLUTION))
 213        )
 214        return resampled
 215
 216    def get_supported_resolutions(
 217        self, include_resampled: bool = False
 218    ) -> list[Resolution]:
 219        """Get all supported resolutions for batch processing for the data source.
 220
 221        Args:
 222            include_resampled: If False (default), return only the resolutions the
 223                data source literally provides. If True, also include all coarser
 224                resolutions derivable by resampling from the finest native one —
 225                e.g. a source providing '1t' also "supports" '1s', '1m', '1h', '1d', etc.
 226                Quote resolutions pass through but are never a resampling base, so a
 227                quote-only source returns only its native resolutions.
 228        """
 229        native = [
 230            Resolution(dtype_or_resol)
 231            for dtype_or_resol in self.data_source.data_categories[self.data_domain]
 232        ]
 233        if not include_resampled or not native:
 234            return native
 235        resample_bases = [r for r in native if not r.is_quote()]
 236        if not resample_bases:
 237            return native
 238        finest = max(resample_bases)
 239        return sorted(
 240            set(native) | set(finest.get_coarser_resolutions(exclude_quote=True)),
 241            reverse=True,
 242        )
 243
 244    @classmethod
 245    def create_data_model(
 246        cls,
 247        product: BaseProduct | str,
 248        resolution: Resolution | MarketDataType | str,
 249        start_date: datetime.date | str,
 250        end_date: datetime.date | str | None = None,
 251        env: Environment | str = Environment.BACKTEST,
 252        data_origin: str = "",
 253        **product_specs: Any,
 254    ) -> MarketDataModel:
 255        """Create a MarketDataModel instance.
 256
 257        Args:
 258            product: Product basis, e.g. `'BTC_USDT_PERP'`, or a product instance.
 259            resolution: Resolution, e.g. `'1t'`, `'1m'`, a `MarketDataType` e.g. "minute", or a `Resolution` instance.
 260            start_date: First date (inclusive), e.g. `'2025-01-01'`, or a `datetime.date`.
 261            end_date: Last date (inclusive). If None, defaults to `start_date` (a single-day model).
 262            env: Trading environment.
 263            data_origin: Sub-label for data from the same source but different origins.
 264                Defaults to the source name.
 265            product_specs: Extra product attributes, e.g. `expiration='2025-12-26'` for
 266                futures. Leave them out to get an error listing the required ones.
 267        """
 268        if isinstance(product, str):
 269            product = cls.DataSource.create_product(product, **product_specs)
 270        return cls.DataModel(
 271            env=env,
 272            data_source=cls.DataSource.metadata.name,
 273            data_origin=data_origin,
 274            product=product,
 275            resolution=resolution,
 276            start_date=start_date,
 277            end_date=end_date or start_date,
 278        )
 279
 280    # NOTE: the return type follows the pipeline mode, see BaseFeed's PipelineT
 281    @overload
 282    def download[F: MarketFeed[Any, Any, Any, Literal[True]]](
 283        self: F,
 284        product: str,
 285        resolution: Resolution | MarketDataType | str,
 286        *,
 287        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 288        start_date: datetime.date | str | None = None,
 289        end_date: datetime.date | str | None = None,
 290        data_origin: str = "",
 291        data_layer: DataLayer | str = DataLayer.CLEANED,
 292        io: BaseIO | None = None,
 293        **product_specs: Any,
 294    ) -> F: ...
 295    @overload
 296    def download(
 297        self: MarketFeed[Any, Any, Any, Literal[False]],
 298        product: str,
 299        resolution: Resolution | MarketDataType | str,
 300        *,
 301        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 302        start_date: datetime.date | str | None = None,
 303        end_date: datetime.date | str | None = None,
 304        data_origin: str = "",
 305        data_layer: DataLayer | str = DataLayer.CLEANED,
 306        io: BaseIO | None = None,
 307        **product_specs: Any,
 308    ) -> RunResult: ...
 309    @overload
 310    def download[F: MarketFeed[Any, Any, Any, bool]](
 311        self: F,
 312        product: str,
 313        resolution: Resolution | MarketDataType | str,
 314        *,
 315        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 316        start_date: datetime.date | str | None = None,
 317        end_date: datetime.date | str | None = None,
 318        data_origin: str = "",
 319        data_layer: DataLayer | str = DataLayer.CLEANED,
 320        io: BaseIO | None = None,
 321        **product_specs: Any,
 322    ) -> F | RunResult: ...
 323    @requires("download")
 324    def download(
 325        self,
 326        product: str,
 327        resolution: Resolution | MarketDataType | str,
 328        *,
 329        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 330        start_date: datetime.date | str | None = None,
 331        end_date: datetime.date | str | None = None,
 332        data_origin: str = "",
 333        data_layer: DataLayer | str = DataLayer.CLEANED,
 334        io: BaseIO | None = None,
 335        **product_specs: Any,
 336    ) -> Self | RunResult:
 337        """Download historical data from the data source.
 338
 339        Args:
 340            product: Product basis, e.g. `'BTC_USDT_PERP'`, `'AAPL_USD_STK'`.
 341            resolution: Target resolution, e.g. `'1t'` (tick), `'1m'`, `'1d'`.
 342                If the source doesn't provide it, the closest finer resolution is
 343                downloaded and downsampled.
 344            rollback_period: How far to roll back from today (UTC) when `start_date`
 345                is None, e.g. `'7d'` = the last 7 days up to yesterday.
 346                Also accepts `'ytd'` (year to date) and `'max'` (all available data).
 347            start_date: First date (inclusive), e.g. `'2025-01-01'`.
 348            end_date: Last date (inclusive). If None, defaults to yesterday (UTC).
 349                Requires `start_date`.
 350            data_origin: Sub-label for data from the same source but different origins.
 351                Defaults to the source name.
 352            data_layer: Data layer of the downloaded data, `'raw'` or `'cleaned'` (default).
 353                `'cleaned'` runs the default transformations (normalize, standardize columns,
 354                downsample); `'raw'` returns the data as-is. The data is stored in this layer if `io` is given.
 355            io: Where and how to store the data. If None, data is not stored.
 356                e.g. `ParquetIO(base_path='./data')`, `DuckLakeIO()`
 357            product_specs: Extra product attributes, e.g. `expiration='2025-12-26'` for
 358                futures. Leave them out to get an error listing the required ones.
 359
 360        Returns:
 361            `RunResult` with the downloaded data, or `self` in pipeline mode.
 362        """
 363        env = Environment.BACKTEST
 364        setup_logging(env=env)
 365        product: BaseProduct = self.data_source.create_product(product, **product_specs)
 366        resolution = Resolution(resolution)
 367        start_date, end_date = self._standardize_dates(
 368            resolution, start_date, end_date, rollback_period
 369        )
 370        if not (
 371            candidates := [
 372                r for r in self.get_supported_resolutions() if r >= resolution
 373            ]
 374        ):
 375            raise ValueError(f"{resolution} is not supported by {self.name}")
 376        # find the first resolution that is >= (finer/higher) the target resolution
 377        data_resolution = min(candidates)
 378        request = cast(
 379            "RequestT",
 380            cast("type[MarketFeedDownloadRequest]", self.Requests.download)(
 381                data_source=self.data_source,
 382                data_origin=data_origin,
 383                io=io,
 384                data_layer=data_layer,
 385                env=env,
 386                product=product,
 387                target_resolution=resolution,
 388                data_resolution=data_resolution,
 389                start_date=start_date,
 390                end_date=end_date,
 391            ),
 392        )
 393        self._append_request(request)
 394        data_model = cast("DataModelT", request.to_data_model())
 395        self._create_batch_dataflows(
 396            extract_func=self._download_impl,
 397            data_models=data_model.split_by_date()
 398            if self.download_dataflow_per_date
 399            else [data_model],
 400        )
 401        return self.run() if not self.is_pipeline() else self
 402
 403    @staticmethod
 404    def _normalize_downloaded_data(df: pl.LazyFrame) -> pl.LazyFrame:
 405        raise NotImplementedError
 406
 407    @staticmethod
 408    def _download_impl(
 409        request: RequestT, data_model: DataModelT
 410    ) -> pl.LazyFrame | None:
 411        raise NotImplementedError
 412
 413    def _get_default_transformations(
 414        self, request: RequestT, /
 415    ) -> list[Callable[..., Any]]:
 416        match request.extract_type:
 417            case "download":
 418                return self._get_default_transformations_for_download(request)
 419            case "retrieve":
 420                return self._get_default_transformations_for_retrieve(request)
 421            case "stream":
 422                return self._get_default_transformations_for_stream(request)
 423            case "replay":
 424                return self._get_default_transformations_for_replay(request)
 425            case _:
 426                raise ValueError(f"Unknown extract type: {request.extract_type}")
 427
 428    def _get_default_transformations_for_download(
 429        self, request: RequestT, /
 430    ) -> list[Callable[..., Any]]:
 431        from pfeed.utils import lambda_with_name
 432
 433        # bind the class and plain values so the lambdas don't capture self (the whole feed) when pickled by Ray
 434        cls = type(self)
 435        is_raw_data = not request.should_clean_data()
 436        product, data_resolution, target_resolution = (
 437            request.product,
 438            request.data_resolution,
 439            request.target_resolution,
 440        )
 441        default_transformations = [
 442            lambda_with_name(
 443                "standardize_date_column",
 444                lambda df: cls._standardize_date_column(df, is_raw_data=is_raw_data),
 445            ),
 446        ]
 447        if is_raw_data:
 448            return default_transformations
 449        default_transformations.extend(
 450            [
 451                cls._normalize_downloaded_data,
 452                lambda_with_name(
 453                    "standardize_columns",
 454                    lambda df: cls._standardize_columns(df, product, data_resolution),
 455                ),
 456                lambda_with_name(
 457                    "resample_data_if_necessary",
 458                    lambda df: cls._resample_data(df, product, target_resolution),
 459                ),
 460                cls._organize_columns,
 461            ]
 462        )
 463        return default_transformations
 464
 465    # NOTE: the return type follows the pipeline mode, see BaseFeed's PipelineT
 466    @overload
 467    def retrieve[F: MarketFeed[Any, Any, Any, Literal[True]]](
 468        self: F,
 469        product: str,
 470        resolution: Resolution | MarketDataType | str,
 471        *,
 472        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 473        start_date: datetime.date | str | None = None,
 474        end_date: datetime.date | str | None = None,
 475        data_origin: str = "",
 476        data_layer: DataLayer | str = DataLayer.CLEANED,
 477        env: Environment | str = Environment.BACKTEST,
 478        dataflow_per_date: bool = False,
 479        io: BaseIO | None = None,
 480        **product_specs: Any,
 481    ) -> F: ...
 482    @overload
 483    def retrieve(
 484        self: MarketFeed[Any, Any, Any, Literal[False]],
 485        product: str,
 486        resolution: Resolution | MarketDataType | str,
 487        *,
 488        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 489        start_date: datetime.date | str | None = None,
 490        end_date: datetime.date | str | None = None,
 491        data_origin: str = "",
 492        data_layer: DataLayer | str = DataLayer.CLEANED,
 493        env: Environment | str = Environment.BACKTEST,
 494        dataflow_per_date: bool = False,
 495        io: BaseIO | None = None,
 496        **product_specs: Any,
 497    ) -> RunResult: ...
 498    @overload
 499    def retrieve[F: MarketFeed[Any, Any, Any, bool]](
 500        self: F,
 501        product: str,
 502        resolution: Resolution | MarketDataType | str,
 503        *,
 504        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 505        start_date: datetime.date | str | None = None,
 506        end_date: datetime.date | str | None = None,
 507        data_origin: str = "",
 508        data_layer: DataLayer | str = DataLayer.CLEANED,
 509        env: Environment | str = Environment.BACKTEST,
 510        dataflow_per_date: bool = False,
 511        io: BaseIO | None = None,
 512        **product_specs: Any,
 513    ) -> F | RunResult: ...
 514    @requires("retrieve")
 515    def retrieve(
 516        self,
 517        product: str,
 518        resolution: Resolution | MarketDataType | str,
 519        *,
 520        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 521        start_date: datetime.date | str | None = None,
 522        end_date: datetime.date | str | None = None,
 523        data_origin: str = "",
 524        data_layer: DataLayer | str = DataLayer.CLEANED,
 525        env: Environment | str = Environment.BACKTEST,
 526        dataflow_per_date: bool = False,
 527        io: BaseIO | None = None,
 528        **product_specs: Any,
 529    ) -> Self | RunResult:
 530        """Retrieve data from storage.
 531
 532        Args:
 533            product: Product basis, e.g. `'BTC_USDT_PERP'`, `'AAPL_USD_STK'`.
 534            resolution: Target resolution, e.g. `'1t'` (tick), `'1m'`, `'1d'`.
 535                If not stored, finer stored data is retrieved and downsampled.
 536            rollback_period: How far to roll back from today (UTC) when `start_date`
 537                is None, e.g. `'7d'` = the last 7 days up to yesterday.
 538                Also accepts `'ytd'` (year to date) and `'max'` (all available data).
 539            start_date: First date (inclusive), e.g. `'2025-01-01'`.
 540            end_date: Last date (inclusive). If None, defaults to yesterday (UTC).
 541                Requires `start_date`.
 542            data_origin: Sub-label for data from the same source but different origins.
 543                Defaults to the source name.
 544            data_layer: Data layer to retrieve the data from:
 545                `'raw'`, `'cleaned'` (default) or `'curated'`. Data is returned as stored, raw data is not cleaned.
 546            env: Trading environment the data was stored in: `'BACKTEST'` (default) for
 547                downloaded data, `'PAPER'` or `'LIVE'` for streamed data.
 548            dataflow_per_date: Whether to create one dataflow per date.
 549                If False (default), all dates are retrieved in one dataflow
 550                (one multi-file scan), fastest for typical queries.
 551                Set True when:
 552
 553                - all dates don't fit in memory at once, e.g. when downsampling
 554                - using Ray, to parallelize per-date tasks across workers
 555            io: Where and how to retrieve the data from. If None, uses `ParquetIO()`.
 556                e.g. `ParquetIO(base_path='./data')`, `DuckLakeIO()`
 557            product_specs: Extra product attributes, e.g. `expiration='2025-12-26'` for
 558                futures. Leave them out to get an error listing the required ones.
 559
 560        Returns:
 561            `RunResult` with the retrieved data, or `self` in pipeline mode.
 562        """
 563        from functools import partial
 564
 565        from pfeed.io.parquet_io import ParquetIO
 566
 567        env = Environment(env.upper())
 568        setup_logging(env=env)
 569        product: BaseProduct = self.data_source.create_product(product, **product_specs)
 570        resolution = Resolution(resolution)
 571        start_date, end_date = self._standardize_dates(
 572            resolution, start_date, end_date, rollback_period
 573        )
 574        data_layer = DataLayer.parse(data_layer)
 575        # search for higher resolutions (closest first), e.g. if resolution is '1m', search '1m' -> '1s' -> '1t'
 576        # only cleaned data has a known schema to resample, raw and curated data only use the requested resolution
 577        search_resolutions = (
 578            [resolution]
 579            if data_layer != DataLayer.CLEANED
 580            else [
 581                resolution,
 582                *sorted(
 583                    [
 584                        _resolution
 585                        for _resolution in self.get_supported_resolutions(
 586                            include_resampled=True
 587                        )
 588                        # quote data has no price or OHLC columns to resample into bars
 589                        if _resolution > resolution and not _resolution.is_quote()
 590                    ]
 591                ),
 592            ]
 593        )
 594
 595        # find the data resolution: the first search resolution stored for every date,
 596        # falls back to the target resolution if none is fully stored
 597        partial_request = partial(
 598            cast("type[MarketFeedRetrieveRequest]", self.Requests.retrieve),
 599            data_source=self.data_source,
 600            data_origin=data_origin,
 601            io_for_retrieval=io or ParquetIO(),
 602            data_layer=data_layer,
 603            env=env,
 604            product=product,
 605            target_resolution=resolution,
 606            start_date=start_date,
 607            end_date=end_date,
 608        )
 609        # data probing, checks storage before any dataflow is created, so the request carries the exact data_resolution to read
 610        missing_dates: list[datetime.date] = []
 611        for search_resolution in search_resolutions:
 612            request = partial_request(data_resolution=search_resolution)
 613            if not (
 614                _missing_dates
 615                := request.create_data_resolution_handler().find_missing_dates_in_storage()
 616            ):
 617                break
 618            # the target resolution is the one fallen back to, report its missing dates
 619            if search_resolution == resolution:
 620                missing_dates = _missing_dates
 621        else:
 622            # list at most 10 missing dates, a long range can miss thousands
 623            shown_dates = ", ".join(str(date) for date in missing_dates[:10]) + (
 624                ", ..." if len(missing_dates) > 10 else ""
 625            )
 626            self.logger.warning(
 627                f"{product} data from {start_date} to {end_date} is not fully stored at any of {search_resolutions} in {request.io_for_retrieval!r}, "
 628                + f"falling back to {resolution}, missing {len(missing_dates)} dates: {shown_dates}"
 629            )
 630            request = partial_request(data_resolution=resolution)
 631        self._append_request(request)
 632        data_model = cast("DataModelT", request.to_data_model())
 633        self._create_batch_dataflows(
 634            extract_func=self._retrieve_impl,
 635            data_models=data_model.split_by_date()
 636            if dataflow_per_date
 637            else [data_model],
 638        )
 639        return self.run() if not self.is_pipeline() else self
 640
 641    @staticmethod
 642    def _retrieve_impl(
 643        request: RequestT, data_model: DataModelT
 644    ) -> pl.LazyFrame | None:
 645        # data_model is of target resolution and has the dataflow's date range (one date when dataflow_per_date=True),
 646        # the handler reads that range at data_resolution
 647        assert isinstance(request, MarketFeedRetrieveRequest)
 648        lf, _ = request.create_data_resolution_handler(data_model).read()
 649        return lf
 650
 651    def _get_default_transformations_for_retrieve(
 652        self, request: RequestT, /
 653    ) -> list[Callable[..., Any]]:
 654        from pfeed.utils import lambda_with_name
 655
 656        default_transformations = []
 657
 658        if cast("MarketFeedRetrieveRequest", request).is_streamed_data():
 659            import polars.selectors as cs
 660
 661            from pfeed.schemas.market_data_schema import RESOLUTION
 662            from pfeed.schemas.streaming.bar_message_schema import START_TS
 663            from pfeed.schemas.streaming.market_message_schema import RECEIVED_AT
 664            from pfeed.schemas.streaming.tick_message_schema import TS
 665            from pfeed.schemas.time_based_data_schema import DATE
 666
 667            # a bar's date is its start, like in batch data, see MarketDataHandler._get_date_col()
 668            date_col = START_TS if request.data_resolution.is_bar() else TS
 669            data_resolution = request.data_resolution
 670            default_transformations.append(
 671                lambda_with_name(
 672                    "streaming_to_batch_schema",
 673                    # stored as UTC, returned tz-naive UTC like batch data, e.g. start_ts, received_at;
 674                    # streamed data has no resolution column;
 675                    # sorted like batch data, the table IO doesn't keep the stored order, e.g. DuckLake's inlined rows and files,
 676                    # received_at orders the updates of the same bar, e.g. incremental bars
 677                    lambda df: (
 678                        df.with_columns(
 679                            cs.datetime(time_zone="*").dt.replace_time_zone(None),
 680                        )
 681                        .with_columns(
 682                            pl.col(date_col).alias(DATE),
 683                            pl.lit(str(data_resolution)).alias(RESOLUTION),
 684                        )
 685                        .sort(date_col, RECEIVED_AT)
 686                    ),
 687                ),
 688            )
 689        # only cleaned data has a known schema to resample and organize, raw and curated data are returned as stored
 690        if request.data_layer == DataLayer.CLEANED:
 691            # bind the class and plain values so the lambdas don't capture self (the whole feed) when pickled by Ray
 692            cls = type(self)
 693            product, target_resolution = request.product, request.target_resolution
 694            default_transformations.extend(
 695                [
 696                    lambda_with_name(
 697                        "resample_data_if_necessary",
 698                        lambda df: cls._resample_data(df, product, target_resolution),
 699                    ),
 700                    cls._organize_columns,
 701                ]
 702            )
 703        return default_transformations
 704
 705    # NOTE: the return type follows the pipeline mode, see BaseFeed's PipelineT
 706    @overload
 707    def stream[F: MarketFeed[Any, Any, Any, Literal[True]]](
 708        self: F,
 709        product: str,
 710        resolution: Resolution | MarketDataType | str,
 711        *,
 712        data_origin: str = "",
 713        io: BaseIO | None = None,
 714        env: Literal[
 715            Environment.PAPER, Environment.LIVE, "PAPER", "LIVE"
 716        ] = Environment.LIVE,
 717        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 718        store_incremental_bars: bool = False,
 719        flush_interval: float = 100,
 720        **product_specs: Any,
 721    ) -> F: ...
 722    @overload
 723    def stream(
 724        self: MarketFeed[Any, Any, Any, Literal[False]],
 725        product: str,
 726        resolution: Resolution | MarketDataType | str,
 727        *,
 728        data_origin: str = "",
 729        io: BaseIO | None = None,
 730        env: Literal[
 731            Environment.PAPER, Environment.LIVE, "PAPER", "LIVE"
 732        ] = Environment.LIVE,
 733        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 734        store_incremental_bars: bool = False,
 735        flush_interval: float = 100,
 736        **product_specs: Any,
 737    ) -> RunResult: ...
 738    @overload
 739    def stream[F: MarketFeed[Any, Any, Any, bool]](
 740        self: F,
 741        product: str,
 742        resolution: Resolution | MarketDataType | str,
 743        *,
 744        data_origin: str = "",
 745        io: BaseIO | None = None,
 746        env: Literal[
 747            Environment.PAPER, Environment.LIVE, "PAPER", "LIVE"
 748        ] = Environment.LIVE,
 749        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 750        store_incremental_bars: bool = False,
 751        flush_interval: float = 100,
 752        **product_specs: Any,
 753    ) -> F | RunResult: ...
 754    @requires("stream")
 755    def stream(
 756        self,
 757        product: str,
 758        resolution: Resolution | MarketDataType | str,
 759        *,
 760        data_origin: str = "",
 761        io: BaseIO | None = None,
 762        env: Literal[
 763            Environment.PAPER, Environment.LIVE, "PAPER", "LIVE"
 764        ] = Environment.LIVE,
 765        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 766        store_incremental_bars: bool = False,
 767        flush_interval: float = 100,
 768        **product_specs: Any,
 769    ) -> Self | RunResult:
 770        """Stream data live from the data source, to replay stored data use replay().
 771
 772        Args:
 773            product: Product basis, e.g. `'BTC_USDT_PERP'`, `'AAPL_USD_STK'`.
 774            resolution: Target resolution, e.g. `'1t'` (tick), `'1m'`, `'1d'`.
 775                If the source doesn't stream it, a finer streamed resolution is used and resampled.
 776            data_origin: Sub-label for data from the same source but different origins.
 777                Defaults to the source name.
 778            io: Where to store the streamed data. If None, data is not stored.
 779                e.g. `ParquetIO(base_path='./data')`, `DuckLakeIO()`
 780            env: Trading environment, `'PAPER'` or `'LIVE'` (default).
 781            callback: Sync or async callable called with each streamed `Message`, before it is transformed or stored.
 782                `Message['data']` holds the parsed data as a list, e.g. bars or ticks,
 783                except when using Ray (see `set_num_workers()`), where the workers parse it and `Message['data']` is None.
 784                It is awaited before the next message is read, so a slow callback delays every stream of the feed,
 785                offload heavy work, e.g. to a task or a thread.
 786            store_incremental_bars: Whether to also store the updates of a bar before it closes,
 787                e.g. for a venue that only streams such updates. If False (default), only closed bars are stored.
 788                Only meaningful for bar resolutions when `io` is given.
 789            flush_interval: Seconds between writes of the buffered streamed data to `io`.
 790                Frequent writes create many small files, which slow down reads until the table is compacted.
 791                Only meaningful when `io` is given.
 792            product_specs: Extra product attributes, e.g. `expiration='2025-12-26'` for
 793                futures. Leave them out to get an error listing the required ones.
 794
 795        Returns:
 796            `RunResult` once the streams end, or `self` in pipeline mode.
 797        """
 798        from pfund_kit.utils.temporal import get_utc_now
 799
 800        setup_logging(env=env)
 801        product: BaseProduct = self.data_source.create_product(product, **product_specs)
 802        resolution = Resolution(resolution)
 803        start_date = end_date = get_utc_now().date()
 804
 805        data_resolution = resolution
 806        if resolution.is_bar():
 807            supported = self.data_source.get_stream_api(env).get_supported_resolutions(
 808                product
 809            )
 810            if resolution not in supported:
 811                from pfund_blocks.datas.bar_data import BarData
 812
 813                target = BarData(product, resolution)
 814                candidates = [r for r in supported if target.can_resample_from(r)]
 815                if not candidates:
 816                    raise ValueError(
 817                        f"{resolution!r} is not supported in streaming by {self.name}"
 818                    )
 819                # the coarsest one that fits, e.g. 6m with [1m, 3m, 5m, 1t] -> 3m
 820                data_resolution = min(candidates)
 821                self.logger.warning(
 822                    f"{product} {resolution} is not supported in streaming, using {data_resolution} instead to resample data",
 823                )
 824
 825        request = cast(
 826            "RequestT",
 827            cast("type[MarketFeedStreamRequest]", self.Requests.stream)(
 828                data_source=self.data_source,
 829                data_origin=data_origin,
 830                io=io,
 831                data_layer=DataLayer.CLEANED,
 832                env=env,
 833                product=product,
 834                target_resolution=resolution,
 835                data_resolution=data_resolution,
 836                start_date=start_date,
 837                end_date=end_date,
 838                store_incremental_bars=store_incremental_bars,
 839                flush_interval=flush_interval,
 840            ),
 841        )
 842        self._append_request(request)
 843        data_model = cast("DataModelT", request.to_data_model())
 844        self._create_stream_dataflow(
 845            extract_func=self._stream_impl,
 846            data_model=data_model,
 847            callback=callback,
 848        )
 849        return self.run() if not self.is_pipeline() else self
 850
 851    @staticmethod
 852    def _create_queue_key_from_request(request: RequestT) -> tuple[ProductName, str]:
 853        return (request.product.name, repr(request.data_resolution))
 854
 855    @staticmethod
 856    def _create_queue_key_from_message(message: Message) -> Hashable:
 857        return (message["product"], message["resolution"])
 858
 859    def _get_handler_kwargs(self, request: RequestT) -> dict[str, Any]:
 860        if request.extract_type == "stream":
 861            return {
 862                "store_incremental_bars": cast(
 863                    "MarketFeedStreamRequest", request
 864                ).store_incremental_bars
 865            }
 866        return {}
 867
 868    @staticmethod
 869    def _parse_streaming_message(msg: Message) -> Message:
 870        raise NotImplementedError
 871
 872    async def _stream_impl(self, request: RequestT):
 873        """Streams from the data source's stream API (env=PAPER/LIVE), replaying (env=BACKTEST) is _replay_impl()."""
 874        stream_api = self.data_source.get_stream_api(
 875            cast("Literal[Environment.PAPER, Environment.LIVE]", request.env)
 876        )
 877        try:
 878            await stream_api.connect()
 879            # NOTE: connect() returns right away, this is what notices that the stream has stopped
 880            await stream_api.wait_until_disconnected()
 881        finally:
 882            try:
 883                # NOTE: cancelling wait_until_disconnected() doesn't disconnect the websockets
 884                await stream_api.disconnect()
 885            finally:
 886                for queue in self._queues.values():
 887                    queue.shutdown()
 888
 889    # NOTE: ALL transformation functions MUST be static methods so that they can be serialized by Ray
 890    def _get_default_transformations_for_stream(
 891        self, request: RequestT
 892    ) -> list[Callable[..., Any]]:
 893        from pfund_blocks.datas.bar_data import BarData
 894        from pfund_blocks.datas.tick_data import TickData
 895
 896        from pfeed.utils import lambda_with_name
 897
 898        default_transformations: list[Callable[..., Any]] = [
 899            self._parse_streaming_message
 900        ]
 901
 902        is_resampling = bool(request.target_resolution < request.data_resolution)
 903        if is_resampling:
 904            self.logger.warning(
 905                f"{request.product.name} {request.target_resolution!r} bars will be created by resampling {request.data_resolution!r} data, "
 906                + "incremental bars (not closed yet) will also be sent, check bar.is_closed if you only want closed bars"
 907            )
 908            product = request.product
 909            target_resolution = request.target_resolution
 910            data_resolution = request.data_resolution
 911            # use data_bar to resample data, e.g. bybit doesn't support '1s' data (target resolution), use '1t' (data resolution) instead
 912            data_bar = BarData(product=product, resolution=target_resolution)
 913            # the resampler takes every update in data resolution, and only passes what can be merged to data_bar (its resamplee),
 914            # i.e. every tick, but only closed bars, an unclosed bar's volume would otherwise be counted at each of its updates
 915            if data_resolution.is_tick():
 916                resampler = TickData(product=product, resolution=data_resolution)
 917            elif data_resolution.is_bar():
 918                resampler = BarData(product=product, resolution=data_resolution)
 919            else:
 920                raise TypeError(f"Invalid {data_resolution=} for resampling")
 921            # raises if the resampler's bars don't fit in data_bar's, e.g. 3m bars can't be resampled to 5m bars
 922            data_bar.bind_resampler(resampler)
 923            # bind the class and plain values so the lambda doesn't capture self (the whole feed) when pickled by Ray
 924            cls = type(self)
 925            default_transformations.append(
 926                lambda_with_name(
 927                    "resample_streaming_message",
 928                    lambda msg: cls._resample_streaming_message(
 929                        message=msg,
 930                        resampler=resampler,
 931                        data_bar=data_bar,
 932                        target_resolution=target_resolution,
 933                    ),
 934                )
 935            )
 936        return default_transformations
 937
 938    @staticmethod
 939    def _resample_streaming_message(
 940        message: Message,
 941        resampler: MarketData,
 942        data_bar: BarData,
 943        target_resolution: Resolution,
 944    ) -> Message:
 945        """Updates the resampler with the message's data points, returns the message with data_bar's bars as its data.
 946
 947        The data is every bar closed by the update in time order, or data_bar's open bar if it changed but nothing closed.
 948        It is None if the update didn't change data_bar, e.g. an unclosed finer bar, which the resampler keeps until it closes.
 949        """
 950        current_bar = data_bar.bar
 951        # NOTE: a parsed message can hold several data points (e.g. Bybit sends a list), each is updated in order
 952        data_points = cast("list[Bar | Tick]", message["data"])
 953        # finest first, so the resampler's own closed bars come before data_bar's
 954        closed_bars = [
 955            bar
 956            for data_point in data_points
 957            for bar in resampler.update(data_point)
 958            if bar.resolution == target_resolution
 959        ]
 960        # case 1: the update closed bars, e.g. the 1m bar ending at 9:05 closes the 9:00 5m bar,
 961        # they are sent with their final states, the bar opened by the same update (if any) is sent with the next update
 962        # NOTE: a message with several data points can close several bars, e.g. a batch of trades spanning seconds
 963        if closed_bars:
 964            resampled_bars = closed_bars
 965        # case 2: the update changed the open bar, e.g. the 9:01 1m bar closes and is merged into the open 9:00 5m bar
 966        # Bar is immutable, so a changed data_bar has a new bar object
 967        elif data_bar.bar is not current_bar:
 968            resampled_bars = [data_bar.bar]
 969        # case 3: nothing changed, e.g. an unclosed 1m bar, or no 1m bar has closed since the stream started
 970        else:
 971            resampled_bars = None
 972        return {
 973            **message,
 974            "resolution": repr(target_resolution),
 975            "data": resampled_bars,
 976        }
 977
 978    # NOTE: the return type follows the pipeline mode, see BaseFeed's PipelineT
 979    @overload
 980    def replay[F: MarketFeed[Any, Any, Any, Literal[True]]](
 981        self: F,
 982        product: str,
 983        resolution: Resolution | MarketDataType | str,
 984        *,
 985        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
 986        start_date: datetime.date | str = "",
 987        end_date: datetime.date | str = "",
 988        data_origin: str = "",
 989        io: BaseIO | None = None,
 990        callback: Callable[[Message], Awaitable[None] | None] | None = None,
 991        replay_pace: float | None = 0,
 992        **product_specs: Any,
 993    ) -> F: ...
 994    @overload
 995    def replay(
 996        self: MarketFeed[Any, Any, Any, Literal[False]],
 997        product: str,
 998        resolution: Resolution | MarketDataType | str,
 999        *,
1000        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
1001        start_date: datetime.date | str = "",
1002        end_date: datetime.date | str = "",
1003        data_origin: str = "",
1004        io: BaseIO | None = None,
1005        callback: Callable[[Message], Awaitable[None] | None] | None = None,
1006        replay_pace: float | None = 0,
1007        **product_specs: Any,
1008    ) -> RunResult: ...
1009    @overload
1010    def replay[F: MarketFeed[Any, Any, Any, bool]](
1011        self: F,
1012        product: str,
1013        resolution: Resolution | MarketDataType | str,
1014        *,
1015        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
1016        start_date: datetime.date | str = "",
1017        end_date: datetime.date | str = "",
1018        data_origin: str = "",
1019        io: BaseIO | None = None,
1020        callback: Callable[[Message], Awaitable[None] | None] | None = None,
1021        replay_pace: float | None = 0,
1022        **product_specs: Any,
1023    ) -> F | RunResult: ...
1024    @requires("replay")
1025    def replay(
1026        self,
1027        product: str,
1028        resolution: Resolution | MarketDataType | str,
1029        *,
1030        rollback_period: Resolution | str | Literal["ytd", "max"] = "1d",
1031        start_date: datetime.date | str = "",
1032        end_date: datetime.date | str = "",
1033        data_origin: str = "",
1034        io: BaseIO | None = None,
1035        callback: Callable[[Message], Awaitable[None] | None] | None = None,
1036        replay_pace: float | None = 0,
1037        **product_specs: Any,
1038    ) -> Self | RunResult:
1039        """Replay stored data at the cleaned data layer as a stream, e.g. for backtesting.
1040
1041        Every replay of the feed is merged into one stream, ordered by when each message would have been received live,
1042        i.e. a bar at its end, a tick at its timestamp.
1043
1044        Args:
1045            product: Product basis, e.g. `'BTC_USDT_PERP'`, `'AAPL_USD_STK'`.
1046            resolution: Resolution of the stored data to replay, e.g. `'1t'` (tick), `'1m'`, `'1d'`.
1047            rollback_period: How far to roll back from today (UTC) when `start_date`
1048                is empty, e.g. `'7d'` = the last 7 days up to yesterday.
1049                Also accepts `'ytd'` (year to date) and `'max'` (all available data).
1050            start_date: First date (inclusive), e.g. `'2025-01-01'`.
1051            end_date: Last date (inclusive). If empty, defaults to yesterday (UTC).
1052                Requires `start_date`.
1053            data_origin: Sub-label for data from the same source but different origins.
1054                Defaults to the source name.
1055            io: Where to replay the data from. If None, uses `ParquetIO()`.
1056                e.g. `ParquetIO(base_path='./data')`, `DuckLakeIO()`
1057            callback: Sync or async callable called with each replayed `Message`, before it is transformed.
1058                `Message['data']` holds a list with one closed bar or one tick, like a streamed message holds a list of them.
1059                It is awaited before the next message is replayed, so a slow callback slows down the replay.
1060            replay_pace: Pacing between messages.
1061                - 0 (default): ASAP — no sleep between messages. Backtests process the whole
1062                  range as fast as possible regardless of resolution or row count.
1063                - >0: fixed cadence in seconds (e.g. 1.0 → one message per wall-second).
1064                  Useful for watching a replay at a steady, human-readable rate.
1065                - None: realistic — sleeps the time between consecutive messages, at most one bar period for bars,
1066                  so data gaps (e.g. weekends) aren't slept through.
1067                  Opt-in only: for fine resolutions or tick data this can take hours, so it is not the default.
1068            product_specs: Extra product attributes, e.g. `expiration='2025-12-26'` for
1069                futures. Leave them out to get an error listing the required ones.
1070
1071        Returns:
1072            `RunResult` once the replay ends, or `self` in pipeline mode.
1073        """
1074        from pfeed.io.parquet_io import ParquetIO
1075
1076        env = Environment.BACKTEST
1077        setup_logging(env=env)
1078        product: BaseProduct = self.data_source.create_product(product, **product_specs)
1079        resolution = Resolution(resolution)
1080        start_date, end_date = self._standardize_dates(
1081            resolution, start_date, end_date, rollback_period
1082        )
1083        request = cast(
1084            "RequestT",
1085            cast("type[MarketFeedReplayRequest]", self.Requests.replay)(
1086                data_source=self.data_source,
1087                data_origin=data_origin,
1088                io_for_replay=io or ParquetIO(),
1089                data_layer=DataLayer.CLEANED,
1090                env=env,
1091                product=product,
1092                target_resolution=resolution,
1093                # NOTE: replays the stored data as is, no resampling
1094                data_resolution=resolution,
1095                start_date=start_date,
1096                end_date=end_date,
1097                replay_pace=replay_pace,
1098            ),
1099        )
1100        self._append_request(request)
1101        data_model = cast("DataModelT", request.to_data_model())
1102        self._create_stream_dataflow(
1103            extract_func=self._replay_impl,
1104            data_model=data_model,
1105            callback=callback,
1106        )
1107        return self.run() if not self.is_pipeline() else self
1108
1109    async def _replay_impl(self, request: RequestT):
1110        """Replays the stored data of every stream request of the feed (env=BACKTEST), merged by the time each message would have been received live.
1111
1112        The extract function of replay(), like _stream_impl() is for stream().
1113        """
1114        import asyncio
1115        from collections import defaultdict
1116        from operator import itemgetter
1117
1118        from pfund_blocks.datas.bar_data import Bar
1119        from pfund_blocks.datas.tick_data import Tick
1120        from pfund_blocks.enums.side import Side
1121        from pfund_blocks.utils.temporal import NS_PER_SECOND
1122
1123        from pfeed.schemas.bar_data_schema import CLOSE, HIGH, LOW, OPEN
1124        from pfeed.schemas.market_data_schema import VOLUME
1125        from pfeed.schemas.tick_data_schema import PRICE, SIDE
1126        from pfeed.schemas.time_based_data_schema import DATE
1127        from pfeed.streaming.replay_clock import ReplayClock
1128
1129        # NOTE: all streams of a feed share one faucet, which runs this once with the first request;
1130        # like a stream API holding every channel, replay reads every request and merges their data by time
1131        requests = [cast("MarketFeedReplayRequest", req) for req in self._requests]
1132        # NOTE: with the data engine, its clock merges the replays of all its feeds by time, it has joined this feed already
1133        clock = self._engine._replay_clock if self._engine else None
1134        if clock is None:
1135            clock = ReplayClock(cast("MarketFeedReplayRequest", request).replay_pace)
1136            clock.join(self)
1137
1138        def _read_day(
1139            handlers: list[tuple[MarketFeedReplayRequest, MarketDataHandler]],
1140        ) -> list[tuple[int, Message]]:
1141            """Returns the day's messages of every request, each with the time it would have been received live, ordered by it."""
1142            messages: list[tuple[int, Message]] = []
1143            for req, handler in handlers:
1144                lf, _ = handler.read()
1145                if lf is None:
1146                    continue
1147                # NOTE: as int ns, iter_rows() would return a python datetime, which only has microseconds
1148                df = lf.with_columns(pl.col(DATE).cast(pl.Int64)).collect()
1149                resolution = req.data_resolution
1150                for row in df.iter_rows(named=True):
1151                    data: Bar | Tick
1152                    if resolution.is_bar():
1153                        # NOTE: a bar's date is its start, but it is only known once closed, so it is received at its end.
1154                        # columns a Bar doesn't have are not replayed, e.g. a stock's dividends and splits
1155                        start_ts: int = row[DATE]
1156                        end_ts = start_ts + resolution.to_seconds() * NS_PER_SECOND
1157                        data = Bar(
1158                            resolution=resolution,
1159                            start_ts=start_ts,
1160                            end_ts=end_ts,
1161                            updated_ts=end_ts,
1162                            open=row[OPEN],
1163                            high=row[HIGH],
1164                            low=row[LOW],
1165                            close=row[CLOSE],
1166                            volume=row[VOLUME],
1167                            is_closed=True,
1168                        )
1169                        received_at = end_ts
1170                    else:
1171                        side = row.get(SIDE)
1172                        data = Tick(
1173                            ts=row[DATE],
1174                            price=row[PRICE],
1175                            volume=row[VOLUME],
1176                            side=Side(side) if side is not None else None,
1177                        )
1178                        received_at = data.ts
1179                    message: Message = {
1180                        "source": req.data_source.name,
1181                        "account": None,
1182                        "channel": None,
1183                        "product": req.product.name,
1184                        # NOTE: must match the queue key, see _create_queue_key_from_request()
1185                        "resolution": repr(resolution),
1186                        "received_at": received_at,
1187                        "message": None,  # raw message as None
1188                        # NOTE: a list, like a streamed message, which can hold several data points
1189                        "data": [data],
1190                        "error": None,
1191                        "extra": {},
1192                    }
1193                    messages.append((received_at, message))
1194            # NOTE: a stable sort, messages received at the same time stay in request order, then in stored order
1195            messages.sort(key=itemgetter(0))
1196            return messages
1197
1198        try:
1199            handlers_per_date: dict[
1200                datetime.date, list[tuple[MarketFeedReplayRequest, MarketDataHandler]]
1201            ] = defaultdict(list)
1202            for req in requests:
1203                if not (req.data_resolution.is_bar() or req.data_resolution.is_tick()):
1204                    raise NotImplementedError(
1205                        f"replaying {req.data_resolution!r} data is not supported yet"
1206                    )
1207                replay_data_model = req.to_data_model()
1208                if missing_dates := req.create_data_resolution_handler(
1209                    replay_data_model
1210                ).find_missing_dates_in_storage():
1211                    self.logger.warning(
1212                        f"{req.product} {req.data_resolution!r} data from {replay_data_model.start_date} to {replay_data_model.end_date} "
1213                        + f"is missing {len(missing_dates)} dates in {req.io_for_replay!r}, they are skipped in the replay"
1214                    )
1215                for daily_data_model in replay_data_model.split_by_date():
1216                    handlers_per_date[daily_data_model.date].append(
1217                        (req, req.create_data_resolution_handler(daily_data_model))
1218                    )
1219            # NOTE: one day at a time to bound memory
1220            for date in sorted(handlers_per_date):
1221                # NOTE: reading blocks, so it runs in a thread, the dataflows keep consuming their queues meanwhile
1222                messages = await asyncio.to_thread(_read_day, handlers_per_date[date])
1223                if not messages:
1224                    self.logger.debug(f"No data to replay on {date}")
1225                    continue
1226                for received_at, message in messages:
1227                    data = cast("list[Bar | Tick]", message["data"])[0]
1228                    # NOTE: at most one bar period for the realistic pace, so data gaps (e.g. weekends) aren't slept through
1229                    max_delay = (
1230                        data.resolution.to_seconds() if isinstance(data, Bar) else None
1231                    )
1232                    await clock.wait_until(self, received_at, max_delay=max_delay)
1233                    await self._callback(message)
1234        finally:
1235            clock.leave(self)
1236            # NOTE: the replay has ended (or failed / was cancelled), ends the dataflows reading the queues
1237            for queue in self._queues.values():
1238                queue.shutdown()
1239
1240    def _get_default_transformations_for_replay(
1241        self, request: RequestT, /
1242    ) -> list[Callable[..., Any]]:
1243        # NOTE: the replayed data is already cleaned
1244        return []