Windowing Based Aggregations?

[primalleous]

I've built my aggregations to have some functionality that I'll likely need that I'm not seeing here. For instance, in a session for NYSE I may have odd intervals I want to aggregate over so I define all the intervals in advance and bucket explicitly. Is even something like using a tumbling windower with an explicit length and align to possible in nautilus currently? That's the base version of mine.

Currently I'm using some custom bytewax code for aggregations and I do have a dependency on accessing the market hours for that date but I don't believe nautilus uses trading calendars at all.

Here's my test code and some of the custom functions. I'd love to have something like this. How big of an overhaul would it be to do this?

import logging

import bytewax.operators as op
import numpy as np
import pandas as pd
import pytest
from bytewax.dataflow import Dataflow
from bytewax.testing import TestingSink, TestingSource, run_main

from src.contexts.data_transformation.aggregation.session.window_event_processor import (
    ProcessEventResult,
    event_processor,
)
from src.contexts.data_transformation.aggregation.session.windowing import (
    ContinuousWindower,
    builder,
    folder,
    merger,
    session_fold_window,
)
from src.contexts.data_transformation.tests.aggregation.session.fixtures.data_loaders import (
    base_agg_test_cases_w_historical,
    base_agg_test_cases_wo_historical,
)

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


@pytest.mark.parametrize("test_case", base_agg_test_cases_wo_historical())
def test_single_session_wo_historical(
    event_clock, base_lower_timeframe_data, test_case
):
    logger.info(
        f"{test_case['agg_dict']['timeframe']}, {test_case['agg_dict']['agg_key']}"
    )

    # Arrange
    flow = Dataflow(test_case["description"])

    custom_windower = ContinuousWindower(intervals=test_case["agg_tasks"])
    logger.debug(f"agg_tasks: {test_case['agg_tasks'][:5]}")  # Log first few agg_tasks

    events = op.input(
        "base_lower_timeframe_data_input",
        flow,
        TestingSource(base_lower_timeframe_data),
    )
    keyed_events = op.key_on("key_on_id", events, lambda x: "id")
    fold_out = session_fold_window(
        step_id=f"aggregation_{test_case['agg_dict']['timeframe']}_{test_case['agg_dict']['agg_key']}",
        up=keyed_events,
        clock=event_clock,
        windower=custom_windower,
        builder=builder,
        folder=folder,
        merger=merger,
        ordered=True,
    )
    unkeyed = op.key_rm("key_rm", fold_out.down)

    processed = op.map("process_event", unkeyed, event_processor)

    def is_close_event(event: ProcessEventResult) -> bool:
        return event.update_type == "close"

    closes = op.filter("filter_closes", processed, is_close_event)

    sink_data = []
    sink = TestingSink(sink_data)
    op.output("out", closes, sink)
    logger.debug(f"Sink data (first 10 items): {sink_data[:10]}")
    logger.debug(f"Sink data length: {len(sink_data)}")

    # Act
    run_main(flow)

    # Assert
    # Retrieve the output from the sink
    actual_aggregated = sink_data
    # logger.debug(f"Sink data collected: {actual_aggregated}")

    # Extract the aggregated state dictionaries directly from ProcessEventResult instances
    actual_state_dicts = [
        result.state for result in actual_aggregated if result is not None
    ]

    # logger.debug(f"Extracted aggregated states: {actual_state_dicts}")

    # Prepare the expected aggregation
    expected_aggregated = test_case["agg_single_session"]

    # Convert to DataFrame for easier comparison
    actual_df = pd.DataFrame(actual_state_dicts)
    actual_df["timestamp"] = pd.to_datetime(actual_df["timestamp"])
    actual_df.set_index("timestamp", inplace=True)

    expected_df = pd.DataFrame(expected_aggregated)
    expected_df["timestamp"] = pd.to_datetime(expected_df["timestamp"])
    expected_df.set_index("timestamp", inplace=True)

    # Done to remove fuzzy generation of last columns
    expected_df.dropna(subset=["open"], inplace=True)

    # Ensure timestamps are in datetime format and sorted
    actual_df_sorted = actual_df.sort_index()
    expected_df_sorted = expected_df.sort_index()

    # Select only OHLCV columns for comparison
    ohlcv_columns = ["open", "high", "low", "close", "volume"]

    # Verify that both DataFrames contain the required columns
    for column in ohlcv_columns:
        assert (
            column in actual_df_sorted.columns
        ), f"Column '{column}' missing in actual data"
        assert (
            column in expected_df_sorted.columns
        ), f"Column '{column}' missing in expected data"

    actual_ohlcv = actual_df_sorted[ohlcv_columns]
    expected_ohlcv = expected_df_sorted[ohlcv_columns]

    # Log the DataFrames for debugging
    logger.debug(f"Actual Aggregated Data:\n{actual_ohlcv}")
    logger.debug(f"Expected Aggregated Data:\n{expected_ohlcv}")

    # Use pandas testing utility for DataFrame comparison
    try:
        pd.testing.assert_frame_equal(
            actual_ohlcv,
            expected_ohlcv,
            check_dtype=False,  # Add this if there are dtype mismatches
        )
    except AssertionError as e:
        logger.error(f"Aggregation mismatch for {test_case['description']}: {e}")
        raise


@pytest.mark.parametrize("test_case", base_agg_test_cases_w_historical())
def test_single_session_w_historical(event_clock, base_lower_timeframe_data, test_case):
    """Test session aggregation with historical data merged in."""
    # Arrange
    flow = Dataflow(test_case["description"])

    custom_windower = ContinuousWindower(intervals=test_case["agg_tasks"])
    logger.debug(f"agg_tasks: {test_case['agg_tasks'][:5]}")  # Log first few agg_tasks

    # Get the first element from historical_overlap
    overlap_element = test_case["historical_overlap"][0]

    if not np.isnan(overlap_element["open"]):

        # Get the first element from base_lower to copy missing fields
        base_element = base_lower_timeframe_data[0]

        # Create new dict with all fields from overlap_element
        new_element = overlap_element.copy()

        # Add missing fields from base_element
        missing_fields = [
            "created_at",
            "created_by",
            "data_provider",
            "updated_at",
            "updated_by",
            "updated_with",
        ]

        for field in missing_fields:
            new_element[field] = base_element[field]
        new_element["id"] = -1

        # Create new list with the new element at the front
        base_lower_timeframe_data_w_historical = [
            new_element
        ] + base_lower_timeframe_data[:]
    else:
        base_lower_timeframe_data_w_historical = base_lower_timeframe_data

    # Input streams for real-time and historical data
    events = op.input(
        "base_lower_timeframe_data_input",
        flow,
        TestingSource(base_lower_timeframe_data_w_historical),
    )
    keyed_events = op.key_on("key_on_id", events, lambda x: "id")
    fold_out = session_fold_window(
        step_id=f"aggregation_{test_case['agg_dict']['timeframe']}_{test_case['agg_dict']['agg_key']}",
        up=keyed_events,
        clock=event_clock,
        windower=custom_windower,
        builder=builder,
        folder=folder,
        merger=merger,
        ordered=True,
    )
    unkeyed = op.key_rm("key_rm", fold_out.down)

    processed = op.map("process_event", unkeyed, event_processor)

    def is_close_event(event: ProcessEventResult) -> bool:
        return event.update_type == "close"

    closes = op.filter("filter_closes", processed, is_close_event)

    sink_data = []
    sink = TestingSink(sink_data)
    op.output("out", closes, sink)
    logger.debug(f"Sink data (first 10 items): {sink_data[:10]}")
    logger.debug(f"Sink data length: {len(sink_data)}")

    # Act
    run_main(flow)

    # Assert
    # Retrieve the output from the sink
    actual_aggregated = sink_data
    # logger.debug(f"Sink data collected: {actual_aggregated}")

    # Extract the aggregated state dictionaries directly from ProcessEventResult instances
    actual_state_dicts = [
        result.state for result in actual_aggregated if result is not None
    ]

    # logger.debug(f"Extracted aggregated states: {actual_state_dicts}")

    # Prepare the expected aggregation
    expected_aggregated = test_case["agg_single_session"]

    # Convert to DataFrame for easier comparison
    actual_df = pd.DataFrame(actual_state_dicts)
    actual_df["timestamp"] = pd.to_datetime(actual_df["timestamp"])
    actual_df.set_index("timestamp", inplace=True)

    expected_df = pd.DataFrame(expected_aggregated)
    expected_df["timestamp"] = pd.to_datetime(expected_df["timestamp"])
    expected_df.set_index("timestamp", inplace=True)

    # Done to remove fuzzy generation of last columns
    expected_df.dropna(subset=["open"], inplace=True)

    # Ensure timestamps are in datetime format and sorted
    actual_df_sorted = actual_df.sort_index()
    expected_df_sorted = expected_df.sort_index()

    # Select only OHLCV columns for comparison
    ohlcv_columns = ["open", "high", "low", "close", "volume"]

    # Verify that both DataFrames contain the required columns
    for column in ohlcv_columns:
        assert (
            column in actual_df_sorted.columns
        ), f"Column '{column}' missing in actual data"
        assert (
            column in expected_df_sorted.columns
        ), f"Column '{column}' missing in expected data"

    actual_ohlcv = actual_df_sorted[ohlcv_columns]
    expected_ohlcv = expected_df_sorted[ohlcv_columns]

    # Log the DataFrames for debugging
    logger.debug(f"Actual Aggregated Data:\n{actual_ohlcv}")
    logger.debug(f"Expected Aggregated Data:\n{expected_ohlcv}")

    # Use pandas testing utility for DataFrame comparison
    try:
        pd.testing.assert_frame_equal(
            actual_ohlcv,
            expected_ohlcv,
            check_dtype=False,  # Add this if there are dtype mismatches
        )
    except AssertionError as e:
        logger.error(f"Aggregation mismatch for {test_case['description']}: {e}")
        raise

Custom windower operator

import copy
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    Iterable,
    List,
    Optional,
    Self,
    Tuple,
    override,
)

import bytewax.operators as op
from bytewax._utils import partition
from bytewax.operators import KeyedStream, S, V, W, operator
from bytewax.operators.windowing import (
    SC,
    SW,
    UTC_MIN,
    Clock,
    ClockLogic,
    StatefulBatchLogic,
    WindowLogic,
    WindowMetadata,
    WindowOut,
    _unwrap_emit,
    _unwrap_late,
    _unwrap_meta,
    _WindowEvent,
    _WindowQueueEntry,
    _WindowSnapshot,
)

from src.contexts.data_transformation.common.enums.aggregation import WindowingType

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class CustomWindowerLogic(ABC, Generic[S]):
    """Abstract class which defines a type of window.

    If you're just getting started with windowing operators, see the
    concrete subclasses of {py:obj}`Windower` for built-in options.

    This type is instantiated within Bytewax's windowing operators and
    has methods called by the runtime. A new subclass of this will
    need to be paired with a new subclass of {py:obj}`Windower`, which
    holds the configuration data for this type.

    This is instantiated for each key which is encountered.

    """

    @abstractmethod
    def open_for(
        self,
        timestamp: datetime,
    ) -> Iterable[int]:
        """Find which windows an item is in and mark them as open.

        Called when a new, non-late item arrives.

        You'll need to do whatever bookkeeping is necessary internally
        to be able to satisfy the other abstract methods, like being
        able to return the metadata for a window.

        :arg timestamp: Of the incoming item.

        :returns: A list of window IDs that this item is in.

        """
        ...

    @abstractmethod
    def late_for(self, timestamp: datetime) -> Iterable[int]:
        """Find which windows an item would have been in, if on-time.

        Called when a late item arrives.

        If there isn't a way to know for sure which specific windows
        this item would have been in, you can return a sentinel value.

        This should not persist any state internally for the returned
        window IDs; the watermark has already passed these windows and
        so they will not be processed by {py:obj}`close_for`.

        :arg timestamp: Of the incoming item.

        :returns: A list of window IDs that this item is late for.

        """
        ...

    @abstractmethod
    def merged(self) -> Iterable[Tuple[int, int]]:
        """Report any windows that have been merged.

        A merge is reported by saying an **original window** should be
        merged into a **target window**.

        This will be called after a batch of items has been processed.
        The system assumes merges can only occur because of incoming
        items.

        You only need to report a merge once. After a window ID has
        been returned as an original window, that ID does not need to
        be reported as closed.

        :returns: An iterable of 2-tuples of `(original_window_id,
            target_window_id)`.

        """
        ...

    @abstractmethod
    def close_for(self, watermark: datetime) -> Iterable[Tuple[int, WindowMetadata]]:
        """Report what windows are now closed.

        Called periodically once the watermark has advanced.

        Once a window ID is returned from this method, it should never
        be opened again (although items may be reported as late for
        it). Internally, you can discard whatever state was being
        maintained about this window.

        :arg watermark: Current watermark. Guaranteed to never
            regress across calls.


        :returns: A list of now closed window IDs and the their final
            metadata.

        """
        ...

    @abstractmethod
    def notify_at(self) -> Optional[datetime]:
        """Next system time timestamp for which a window might close.

        This will be called periodically to know when to wake up this
        operator to close windows, even if no items were received.

        :returns: System time of when to next wake up.

        """
        ...

    @abstractmethod
    def is_empty(self) -> bool:
        """If no state needs to be maintained anymore.

        This, in general, is if there are no windows currently open.
        But you might have other state that needs to persist across
        resumes, and so might return `True` here to maintain that.

        If `False` is returned here, the entire windower will be
        discarded as to not "leak" memory. A new windower will be
        built if this state key is encountered again.

        :returns: If no state needs to be maintained.

        """
        ...

    @abstractmethod
    def snapshot(self) -> S:
        """Return a immutable copy of the state for recovery.

        This will be called periodically by the runtime.

        The value returned here will be passed back to
        {py:obj}`Windower.build` when resuming.

        The state must be {py:obj}`pickle`-able.

        :::{danger}

        **The state must be effectively immutable!** If any of the
        other functions in this class might be able to mutate the
        state, you must {py:obj}`copy.deepcopy` or something
        equivalent before returning it here.

        :::

        :returns: The immutable state to be {py:obj}`pickle`d.

        """
        ...

    @abstractmethod
    def get_window_metadata(self, window_id: int) -> WindowMetadata | None:
        pass


class CustomWindower(ABC, Generic[S]):
    """A type of window.

    If you're just getting started with windowing operators, see the
    concrete subclasses of {py:obj}`Windower` for built-in options.

    Every subclass of this class must also have a matching
    implementation of {py:obj}`WindowerLogic` that actually implements
    the window.

    """

    @abstractmethod
    def build(self, resume_state: Optional[S]) -> CustomWindowerLogic[S]:
        """Construct a new windower logic instance.

        This should close over any non-state configuration and combine
        it with the resume state to return a prepared
        {py:obj}`WindowerLogic` instance.

        :arg resume_state: Snapshot from the previous call to
            {py:obj}`WindowerLogic.snapshot`, if any.

        :returns: A prepared windower logic.

        """
        ...


@dataclass
class _CustomWindowLogic(
    StatefulBatchLogic[
        V,
        _WindowEvent[V, W],
        _WindowSnapshot[V, SC, SW, S],
    ],
):
    clock: ClockLogic[V, SC]
    windower: CustomWindowerLogic[SW]
    builder: Callable[[Optional[S], Optional[WindowMetadata]], WindowLogic[V, W, S]]
    ordered: bool

    logics: Dict[int, WindowLogic[V, W, S]] = field(default_factory=dict)
    queue: List[_WindowQueueEntry[V]] = field(default_factory=list)

    _last_watermark: datetime = UTC_MIN

    def _handle_inserts(
        self, due_entries: List[_WindowQueueEntry[V]]
    ) -> Iterable[_WindowEvent[V, W]]:
        for entry in due_entries:
            value, timestamp = entry
            for window_id in self.windower.open_for(timestamp):
                if window_id in self.logics:
                    logic = self.logics[window_id]
                else:
                    metadata = self.windower.get_window_metadata(window_id)
                    logic = self.builder(None, metadata)
                    self.logics[window_id] = logic

                ws = logic.on_value(value)
                yield from ((window_id, "E", w) for w in ws)  # type: ignore

    def _handle_merged(self) -> Iterable[_WindowEvent[V, W]]:
        for orig_window_id, targ_window_id in self.windower.merged():
            if targ_window_id != orig_window_id:
                orig_logic = self.logics.pop(orig_window_id)
                into_logic = self.logics[targ_window_id]
                ws = into_logic.on_merge(orig_logic)
                yield from ((targ_window_id, "E", w) for w in ws)  # type: ignore

    def _handle_closed(self, watermark: datetime) -> Iterable[_WindowEvent[V, W]]:
        for window_id, meta in self.windower.close_for(watermark):
            logic = self.logics.pop(window_id)
            ws = logic.on_close()
            yield from ((window_id, "E", w) for w in ws)  # type: ignore

            yield (window_id, "M", meta)

    def _flush_queue(self, watermark: datetime) -> Iterable[_WindowEvent[V, W]]:
        # logger.debug(f"Flushing queue with watermark: {watermark}")
        if self.ordered:
            # Remove due entries from the queue.
            due_entries, self.queue = partition(
                self.queue, lambda entry: entry[1] <= watermark
            )
            due_entries.sort(key=lambda entry: entry[1])
        else:
            due_entries = self.queue
            self.queue = []

        yield from self._handle_inserts(due_entries)
        yield from self._handle_merged()
        yield from self._handle_closed(watermark)

    def _is_empty(self) -> bool:
        return (
            len(self.logics) <= 0 and len(self.queue) <= 0 and self.windower.is_empty()
        )

    @override
    def on_batch(self, values: List[V]) -> Tuple[Iterable[_WindowEvent[V, W]], bool]:
        self.clock.before_batch()

        events: List[_WindowEvent[V, W]] = []
        for value in values:
            value_timestamp, watermark = self.clock.on_item(value)
            # logger.debug(
            #     f"on_batch - Value timestamp: {value_timestamp}, Watermark: {watermark}"
            # )
            assert watermark >= self._last_watermark
            self._last_watermark = watermark

            if value_timestamp < watermark:
                late_window_ids = self.windower.late_for(value_timestamp)
                events.extend((window_id, "L", value) for window_id in late_window_ids)
            else:
                entry = (value, value_timestamp)
                self.queue.append(entry)

        events.extend(self._flush_queue(watermark))  # type: ignore
        return (events, self._is_empty())

    @override
    def on_notify(self) -> Tuple[Iterable[_WindowEvent[V, W]], bool]:
        watermark = self.clock.on_notify()
        assert watermark >= self._last_watermark
        self._last_watermark = watermark

        events = list(self._flush_queue(watermark))
        return (events, self._is_empty())

    @override
    def on_eof(self) -> Tuple[Iterable[_WindowEvent[V, W]], bool]:
        watermark = self.clock.on_eof()
        assert watermark >= self._last_watermark
        self._last_watermark = watermark

        events = list(self._flush_queue(watermark))
        return (events, self._is_empty())

    @override
    def notify_at(self) -> Optional[datetime]:
        notify_at = self.windower.notify_at()
        if notify_at is not None:
            notify_at = self.clock.to_system_utc(notify_at)
        return notify_at

    @override
    def snapshot(self) -> _WindowSnapshot[V, SC, SW, S]:
        return _WindowSnapshot(
            self.clock.snapshot(),
            self.windower.snapshot(),
            {window_id: logic.snapshot() for window_id, logic in self.logics.items()},
            # No need to deepcopy since entries are frozen.
            list(self.queue),
        )


@operator
def custom_window(
    step_id: str,
    up: KeyedStream[V],
    clock: Clock[V, Any],
    windower: CustomWindower[Any],
    builder: Callable[[Optional[S], Optional[WindowMetadata]], WindowLogic[V, W, S]],
    ordered: bool = True,
) -> WindowOut[V, W]:
    """Advanced generic windowing operator.

    This is a lower-level operator Bytewax and gives you control over
    when a windowing operator emits items. {py:obj}`WindowLogic` works
    in tandem with {py:obj}`ClockLogic` and {py:obj}`WindowerLogic` to
    implement its behavior. See documentation of those interfaces.

    :arg step_id: Unique ID.

    :arg up: Keyed upstream.

    :arg clock: Time definition.

    :arg windower: Window definition.

    :arg builder: Called whenever a new window is opened with the
        resume state returned from {py:obj}`WindowLogic.snapshot` for
        that window, if any. This should close over any non-state
        configuration and combine it with the resume state to return a
        prepared {py:obj}`WindowLogic` for this window.

    :arg ordered: Whether to apply values to the logic in timestamp
        order. If not, they'll be in upstream order. There is a
        performance and latency penalty to ordering by timestamp.
        Defaults to `True`.

    :returns: Window result streams.

    """

    def shim_builder(
        resume_state: Optional[_WindowSnapshot],
    ) -> _CustomWindowLogic[V, W, SC, SW, S]:
        if resume_state is not None:
            clock_logic = clock.build(resume_state.clock_state)
            windower_logic = windower.build(resume_state.windower_state)
            logics = {
                window_id: builder(
                    logic_state, windower_logic.get_window_metadata(window_id)
                )
                for window_id, logic_state in resume_state.logic_states.items()
            }
            return _CustomWindowLogic(
                clock_logic,
                windower_logic,
                builder,
                ordered,
                logics,
                resume_state.queue,
            )
        else:
            clock_logic = clock.build(None)
            windower_logic = windower.build(None)
            return _CustomWindowLogic(
                clock_logic,
                windower_logic,
                builder,
                ordered,
            )

    events = op.stateful_batch("stateful_batch", up, shim_builder)

    downs: KeyedStream[Tuple[int, W]] = op.filter_map_value(
        "unwrap_down",
        events,
        _unwrap_emit,
    )
    lates: KeyedStream[Tuple[int, V]] = op.filter_map_value(
        "unwrap_late",
        events,
        _unwrap_late,
    )
    metas: KeyedStream[Tuple[int, WindowMetadata]] = op.filter_map_value(
        "unwrap_meta",
        events,
        _unwrap_meta,
    )
    return WindowOut(downs, lates, metas)


@dataclass
class _PorousFoldWindowLogic(WindowLogic[V, S, S]):
    """Implements porous fold window logic for custom windowing operations."""

    # Accept builder function that takes resume_state
    builder: Callable[[Optional[S], Optional[WindowMetadata]], S]
    folder: Callable[[S, V], S]
    merger: Callable[[S, S], S]
    state: S
    metadata: Optional[WindowMetadata] = None

    def __init__(
        self,
        builder: Callable[[Optional[S], Optional[WindowMetadata]], S],
        folder: Callable[[S, V], S],
        merger: Callable[[S, S], S],
        resume_state: Optional[S] = None,
        metadata: Optional[WindowMetadata] = None,
    ):
        """Initialize with optional resume state."""
        self.builder = builder
        self.folder = folder
        self.merger = merger
        self.state = (
            builder(resume_state, metadata) if resume_state is None else resume_state
        )
        self.metadata = metadata

    @override
    def on_value(self, value: V) -> Iterable[S]:
        """Process a new value and update the state."""
        self.state = self.folder(self.state, value)
        return (("value", self.state),)  # type: ignore

    @override
    def on_merge(self, original: Self) -> Iterable[S]:
        """Merge this window's state with another window's state."""
        self.state = self.merger(self.state, original.state)
        return (("merge", self.state),)  # type: ignore

    @override
    def on_close(self) -> Iterable[S]:
        """Return the final state when the window closes."""
        return (("close", self.state),)  # type: ignore

    @override
    def snapshot(self) -> S:
        """Create a deep copy of the current state."""
        return copy.deepcopy(self.state)


@operator
def session_fold_window(
    step_id: str,
    up: KeyedStream[V],
    clock: Clock[V, Any],
    windower: CustomWindower[Any],
    builder: Callable[[Optional[S], Optional[WindowMetadata]], S],
    folder: Callable[[S, V], S],
    merger: Callable[[S, S], S],
    ordered: bool = True,
) -> WindowOut[V, S]:
    """
    Implement a session-based folding window operation.
    Args:
        builder: Function that takes optional resume state and returns initial state
        folder: Function to fold new values into state
        merger: Function to merge states when windows merge
    """

    def shim_builder(
        resume_state: Optional[S], metadata: Optional[WindowMetadata]
    ) -> _PorousFoldWindowLogic[V, S]:
        return _PorousFoldWindowLogic(
            builder=builder,
            folder=folder,
            merger=merger,
            resume_state=resume_state,
            metadata=metadata,
        )

    return custom_window(
        "session_fold_window", up, clock, windower, shim_builder, ordered
    )


def builder(
    state: Optional[Dict[str, Any]], metadata: Optional[WindowMetadata]
) -> Dict[str, Any]:
    try:
        state = state or {}
        return {
            "metadata": metadata,
            "timestamp": metadata.open_time if metadata else None,
            "open": state.get("open", None),
            "high": state.get("high", float("-inf")),
            "low": state.get("low", float("inf")),
            "close": state.get("close", None),
            "volume": state.get("volume", 0),
            **state,
        }
    except Exception as e:
        logger.error(f"Error in builder: {e}")
        raise


def folder(state, event):
    logger.debug(f"state {state}")
    logger.debug(f"event {event}")
    try:
        # if state["timestamp"] is None:
        #     state["timestamp"] = event["timestamp"]
        if state["open"] is None:
            state["open"] = event["open"]
        state["high"] = max(state["high"], event["high"])
        state["low"] = min(state["low"], event["low"])
        state["close"] = event["close"]
        state["volume"] += event["volume"]
        return state
    except Exception as e:
        logger.error(f"Error in folder: {e}")
        raise


def merger(state1, state2):
    try:
        merged_state = {
            "timestamp": state1["timestamp"],
            "open": state1["open"] if state1["open"] is not None else state2["open"],
            "high": max(state1["high"], state2["high"]),
            "low": min(state1["low"], state2["low"]),
            "close": (
                state2["close"] if state2["close"] is not None else state1["close"]
            ),
            "volume": state1["volume"] + state2["volume"],
        }
        return merged_state
    except Exception as e:
        logger.error(f"Error in merger: {e}")
        raise


IntervalInfo = Tuple[WindowingType, Optional[timedelta], Tuple[datetime, datetime]]


@dataclass
class _ContinuousWindowState:
    """State for continuous windowing with no gaps."""

    opened: Dict[int, WindowMetadata] = field(default_factory=dict)


@dataclass
class _ContinuousWindowLogic(CustomWindowerLogic[_ContinuousWindowState]):
    """Logic for handling both tumbling and session windows."""

    intervals: List[IntervalInfo]
    state: _ContinuousWindowState

    def get_window_metadata(self, window_id: int) -> WindowMetadata | None:
        return self.state.opened.get(window_id)

    def intersects(
        self, timestamp: datetime, window_length: timedelta, align_to: datetime
    ) -> List[int]:
        since_origin = timestamp - align_to
        window_ids = list(
            range(
                ((since_origin - window_length) // window_length) + 1,
                (since_origin // window_length) + 1,
            )
        )
        # logger.debug(
        #     f"For timestamp {timestamp}: since_origin={since_origin}, window_ids={window_ids}"
        # )
        return window_ids

    def _metadata_for_tumbling(
        self, window_id: int, window_length: timedelta, align_to: datetime
    ) -> WindowMetadata:
        """Calculate metadata for tumbling windows using SlidingWindowerLogic approach."""
        open_time = align_to + window_length * window_id
        close_time = open_time + window_length
        return WindowMetadata(open_time, close_time)

    def _find_interval(self, timestamp: datetime) -> Optional[Tuple[int, IntervalInfo]]:
        """Find the interval containing the timestamp."""
        for i, interval_info in enumerate(self.intervals):
            _, _, (start, end) = interval_info
            if start <= timestamp < end:
                return i, interval_info
        return None

    def open_for(self, timestamp: datetime) -> Iterable[int]:
        """Find which windows a timestamp belongs to."""
        interval_info = self._find_interval(timestamp)
        if interval_info is None:
            return iter(())

        interval_idx, (window_type, offset, (start, end)) = interval_info

        window_id = interval_idx
        meta = WindowMetadata(open_time=start, close_time=end)
        self.state.opened.setdefault(window_id, meta)
        yield window_id

        # TODO: Just going to force everything to be WindowType.SESSION due to bugs for now
        # bug is related to switching between windowing types and off-by-one erros in window boundaries
        #   Eventual consistency
        # https://chatgpt.com/c/678c096e-a658-8005-b393-06bcfa5d0287
        # https://claude.ai/chat/f53c9e97-dd50-493a-843f-3c6c972dc40e
        # https://claude.ai/chat/f5dbc885-f61f-412b-a67d-0d842c354958
        # https://claude.ai/chat/0c4b3bb3-9cd3-4ff8-a4dd-937eb7c84367
        # https://claude.ai/chat/f970be2b-9ecc-4b62-acf1-e68532494c4f
        # https://claude.ai/chat/10fceb80-50b3-4be9-8f55-0d85a116c5a7
        # https://claude.ai/chat/6f3fedaa-e754-46b5-834c-3f1bef8f2e61
        # if window_type == WindowingType.TUMBLING and offset is not None:
        #     # For tumbling windows, use SlidingWindowerLogic approach
        #     for window_id in self.intersects(timestamp, offset, start):
        #         meta = self.state.opened.get(window_id) or self._metadata_for_tumbling(
        #             window_id, offset, start
        #         )
        #         self.state.opened.setdefault(window_id, meta)
        #         yield window_id
        # else:
        #     # For session windows, use original approach
        #     window_id = interval_idx
        #     meta = WindowMetadata(open_time=start, close_time=end)
        #     self.state.opened.setdefault(window_id, meta)
        #     yield window_id

    def late_for(self, timestamp: datetime) -> Iterable[int]:
        """Find windows that a late timestamp would have belonged to."""
        interval_info = self._find_interval(timestamp)
        if interval_info is None:
            return iter(())

        interval_idx, (window_type, offset, (start, end)) = interval_info
        if window_type == WindowingType.TUMBLING and offset is not None:
            yield from self.intersects(timestamp, offset, start)
        else:
            # For session windows
            for i, (_, _, (open_time, _)) in enumerate(self.intervals):
                if timestamp < open_time:
                    yield i

    def merged(self) -> Iterable[Tuple[int, int]]:
        """No merging needed."""
        return iter(())

    def close_for(self, watermark: datetime) -> Iterable[Tuple[int, WindowMetadata]]:
        # logger.debug(f"Checking window closure with watermark: {watermark}")
        """Close windows based on watermark."""
        closed_windows = []
        for window_id, meta in list(self.state.opened.items()):
            # logger.debug(f"Window {window_id}: close_time={meta.close_time}")
            if meta.close_time <= watermark:
                # logger.debug(f"Closing window {window_id}")
                closed_windows.append((window_id, meta))
                del self.state.opened[window_id]
        return closed_windows

    def notify_at(self) -> Optional[datetime]:
        """Get next notification time."""
        return min(
            (meta.close_time for meta in self.state.opened.values()), default=None
        )

    def is_empty(self) -> bool:
        """Check if there are no active windows."""
        return len(self.state.opened) == 0

    def snapshot(self) -> _ContinuousWindowState:
        """Return the current state of the windower."""
        return self.state


@dataclass
class ContinuousWindower(CustomWindower[_ContinuousWindowState]):
    """Windower that handles both tumbling and session windows."""

    intervals: List[IntervalInfo]

    def build(
        self, resume_state: Optional[_ContinuousWindowState]
    ) -> _ContinuousWindowLogic:
        """Build windower logic."""
        state = resume_state if resume_state is not None else _ContinuousWindowState()
        return _ContinuousWindowLogic(intervals=self.intervals, state=state)

[faysou]

That's a lot of code. If it's important for you, then best to find how to integrate it yourself in nautilus then do a PR. That's the path I followed for computing option greeks with nautilus.

[primalleous]

I'll try to look into it. I doubt I'll be able to do it fully since never having coded in Rust/Cython. I think what I want may be doable with just adding an Actor too although I'm sure not nearly performant then.

I know bytewax uses timely dataflow in Rust. Not even sure if that programming paradigm is compatible with nautilus trader.