Join two "ordered" streams (streamsorter)

[th0ger]

In continuation of the "ordered streams" workaround discussed in #2395 (by setting readBatchSize:1 and scale.max: 1).

I have two ordered source streams that I want to combine into a single ordered stream using a join vertex. Which vertex type could support me writing a UDF that sorts the output events?

[vigith]

is there an ID to join by? if yes, try session-window with a timeout?

[th0ger]

@vigith
Well, I wrote the following StreamSortMapper that I thought would work. Unit test is passing.
Unfortunately it doesn't work with a real pipeline, because I misunderstood how handlers are threaded by Numaflow. (The first handler call turns out to block the following.) So my unit test setup is not representable of the actual usage.

I tried to see if it could be fixed with you async servers, but did not manage.

mapper.py

import heapq
import threading
from datetime import datetime, timedelta
from pynumaflow.mapper import Datum, Message, Messages, Mapper
from .dtypes import DatumHashable


class StreamSortMapper(Mapper):
    """Sort a Numaflow event stream based on event_time."""

    def __init__(self, idle_threshold: float, idle_increment_by: float) -> None:
        self.idle_threshold = idle_threshold
        self.idle_increment_by = idle_increment_by
        super().__init__()
        self.queue = []
        self.watermark = datetime(1970, 1, 1)
        self.lock = threading.Lock()
        self.condition = threading.Condition(self.lock)

    def handler(self, keys: list[str], datum: Datum) -> Messages:
        event_time = datum.event_time
        watermark = datum.watermark
        datum_hash = hash(DatumHashable.from_datum(datum))

        with self.condition:
            if event_time <= self.watermark:
                return Messages(Message.to_drop())

            heapq.heappush(self.queue, (event_time, datum_hash))
            if watermark > self.watermark:
                self.watermark = watermark
            self.condition.notify_all()

            while True:
                if self.queue and self.queue[0][1] == datum_hash and self.queue[0][0] < self.watermark:
                    heapq.heappop(self.queue)
                    self.condition.notify_all()
                    break
                else:
                    idling = not self.condition.wait(timeout=self.idle_threshold)
                    if idling:
                        _increment_idle_watermark = self.watermark + timedelta(seconds=self.idle_increment_by)
                        self.watermark = _increment_idle_watermark

        return Messages(Message(value=datum.value, keys=keys))

__main__.py

import click
from pynumaflow.mapper import MapServer
from .mapper import StreamSortMapper


@click.option(
    "--idle-threshold",
    required=True,
    type=float,
    default=10.0,
    help="""Threshold for idle source detection, in seconds.
    Must be positive. Default 10.0.
    """,
)
@click.option(
    "--idle-increment-by",
    required=True,
    type=float,
    default=1.0,
    help="""Watermark increment interval for idle detection, in seconds.
    Must be positive. Default 1.0.
    """,
)
@click.command
def cli(idle_threshold: float, idle_increment_by: float):
    grpc_server = MapServer(
        mapper_instance=StreamSortMapper(idle_threshold=idle_threshold, idle_increment_by=idle_increment_by)
    )
    grpc_server.start()


if __name__ == "__main__":
    cli()

dtypes.py

from dataclasses import dataclass
from datetime import datetime
from pynumaflow.mapper import Datum


@dataclass(frozen=True)
class DatumHashable:
    """
    Hashable version of the Datum class.
    - keys: stored as tuple (instead of list)
    - headers: stored as frozenset of items (instead of dict)
    """

    keys: tuple[str, ...]
    value: bytes
    event_time: datetime
    watermark: datetime
    headers: frozenset[tuple[str, str]]

    @classmethod
    def from_datum(cls, datum: Datum) -> "DatumHashable":
        """Create a DatumHashable instance from a Datum instance."""
        return cls(
            keys=tuple(datum.keys),
            value=datum.value,
            event_time=datum.event_time,
            watermark=datum.watermark,
            headers=frozenset(datum.headers.items()),
        )

test_mapper.py

from datetime import datetime, timedelta
import threading
import time
from pynumaflow.mapper import Datum, Messages, DROP
from .mapper import StreamSortMapper
from pydantic import BaseModel


class DatumValue(BaseModel):
    id: int
    timestamp: datetime


def test_mapper():
    """This test is best inspected with `pytest -vs`."""
    speed_factor = 0.01  # Virtual test speed for faster test execution.
    process_time_step = 1

    mapper = StreamSortMapper(idle_threshold=10 * speed_factor, idle_increment_by=1 * speed_factor)

    # Create datums with event times close to now, slightly unordered
    now = datetime.now()
    event_times = [
        now + timedelta(seconds=delta * speed_factor)
        for delta in [0, 2, 1, 3, 4, -1, 5, 5]  # Slightly unordered event times
    ]
    watermarks = [
        now + timedelta(seconds=delta * speed_factor)
        for delta in [-3, -1, -1, 0, 1, 1, 2, 2]  # Monotonic watermark 3s delay
    ]
    datums = [
        Datum(
            keys=[],
            value=DatumValue(id=id, timestamp=event_time).model_dump_json().encode(),
            event_time=event_time,
            watermark=watermark,
            headers=None,
        )
        for id, (event_time, watermark) in enumerate(zip(event_times, watermarks))
    ]

    results = []

    def handler_thread(datum):
        res: Messages = mapper.handler(keys=[], datum=datum)
        if res[0].tags != [DROP]:
            results.append(res)

    # Create and start threads
    threads = []
    for datum in datums:
        t = threading.Thread(target=handler_thread, args=(datum,))
        threads.append(t)
        t.start()
        time.sleep(process_time_step * speed_factor)

    # Wait for all threads to complete
    for t in threads:
        t.join()
    print("All handler threads have completed.")

    # Assert output
    assert len(results) == len(event_times) - 1

    result_vals = [DatumValue.model_validate_json(r[0].value.decode()) for r in results]

    result_ids = [v.id for v in result_vals]
    expected_ids = [[0, 2, 1, 3, 4, 6, 7], [0, 2, 1, 3, 4, 7, 6]]
    assert result_ids in expected_ids

    result_timestamps = [v.timestamp for v in result_vals]
    assert result_timestamps == sorted(result_timestamps)

test output

$pytest -vvs
Init StreamSortMapper
[thread=0] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":0,"timestamp":"2025-02-28T07:53:30.333571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 333571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 303571), _headers={})
[thread=0] Heapqueue push, heap_size=1, event_time=2025-02-28 07:53:30.333571
[thread=0] Watermark incremented 1740729210.303571s to 2025-02-28 07:53:30.303571
[thread=1] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":1,"timestamp":"2025-02-28T07:53:30.353571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 353571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 323571), _headers={})
[thread=1] Heapqueue push, heap_size=2, event_time=2025-02-28 07:53:30.353571
[thread=1] Watermark incremented 0.02s to 2025-02-28 07:53:30.323571
[thread=2] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":2,"timestamp":"2025-02-28T07:53:30.343571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 343571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 323571), _headers={})
[thread=2] Heapqueue push, heap_size=3, event_time=2025-02-28 07:53:30.343571
[thread=2] Watermark not incremented
[thread=3] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":3,"timestamp":"2025-02-28T07:53:30.363571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 363571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 333571), _headers={})
[thread=3] Heapqueue push, heap_size=4, event_time=2025-02-28 07:53:30.363571
[thread=3] Watermark incremented 0.01s to 2025-02-28 07:53:30.333571
[thread=4] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":4,"timestamp":"2025-02-28T07:53:30.373571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 373571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 343571), _headers={})
[thread=4] Heapqueue push, heap_size=5, event_time=2025-02-28 07:53:30.373571
[thread=4] Watermark incremented 0.01s to 2025-02-28 07:53:30.343571
[thread=0] Heapqueue pop, heap_size=4, event_time=2025-02-28 07:53:30.333571
[thread=5] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":5,"timestamp":"2025-02-28T07:53:30.323571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 323571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 343571), _headers={})
[thread=5] WARNING: Dropping late event (2025-02-28 07:53:30.323571)
[thread=6] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":6,"timestamp":"2025-02-28T07:53:30.383571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 383571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 353571), _headers={})
[thread=6] Heapqueue push, heap_size=5, event_time=2025-02-28 07:53:30.383571
[thread=6] Watermark incremented 0.01s to 2025-02-28 07:53:30.353571
[thread=2] Heapqueue pop, heap_size=4, event_time=2025-02-28 07:53:30.343571
[thread=7] Started thread, keys=[], datum=Datum(_keys=[], _value=b'{"id":7,"timestamp":"2025-02-28T07:53:30.383571"}', _event_time=datetime.datetime(2025, 2, 28, 7, 53, 30, 383571), _watermark=datetime.datetime(2025, 2, 28, 7, 53, 30, 353571), _headers={})
[thread=7] Heapqueue push, heap_size=5, event_time=2025-02-28 07:53:30.383571
[thread=7] Watermark not incremented
[thread=3] Idling detected, incrementing watermark by 0.01s to 2025-02-28 07:53:30.363571
[thread=7] Idling detected, incrementing watermark by 0.01s to 2025-02-28 07:53:30.373571
[thread=4] Idling detected, incrementing watermark by 0.01s to 2025-02-28 07:53:30.383571
[thread=6] Idling detected, incrementing watermark by 0.01s to 2025-02-28 07:53:30.393571
[thread=1] Idling detected, incrementing watermark by 0.01s to 2025-02-28 07:53:30.403571
[thread=1] Heapqueue pop, heap_size=4, event_time=2025-02-28 07:53:30.353571
[thread=3] Heapqueue pop, heap_size=3, event_time=2025-02-28 07:53:30.363571
[thread=4] Heapqueue pop, heap_size=2, event_time=2025-02-28 07:53:30.373571
[thread=7] Heapqueue pop, heap_size=1, event_time=2025-02-28 07:53:30.383571
[thread=6] Heapqueue pop, heap_size=0, event_time=2025-02-28 07:53:30.383571
All handler threads have completed.

[th0ger]

1. Can you provide a functional udf for stream sort? (I can't.)
2. Otherwise I suggest that you commit to prioritize Ordered Processing #1677.
3. Maybe an intermediate solution exists where you slightly open up the numflow API/SDK, such that a udf implementation 1) becomes possible.

[vigith]

thank you @th0ger for sharing the code. I understand the problem now. Let me convert this to an issue and make sure we give a way for users to solve this problem.

Let me convert this to an issue.

[th0ger]

@vigith sounds great!
Maybe one approach is to allow udfs that return the full Datum instead of just Message(s). Without me knowing what complications that might bring.

[vigith]

#2440