feat(streamsorter): flush on timeout, adjust logging (#279)

Signed-off-by: Martin Thøgersen <[email protected]>

Author: th0ger

packages/pynumaflow/examples/accumulator/streamsorter/example.py

@@ -38,10 +38,14 @@ async def handler(
             # If watermark has moved forward
             if datum.watermark and datum.watermark > self.latest_wm:
                 self.latest_wm = datum.watermark
+                _LOGGER.info(f"Watermark updated: {self.latest_wm}")
                 await self.flush_buffer(output)
 
             self.insert_sorted(datum)
 
+        _LOGGER.info("Timeout reached")
+        await self.flush_buffer(output, flush_all=True)
+
     def insert_sorted(self, datum: Datum):
         # Binary insert to keep sorted buffer in order
         left, right = 0, len(self.sorted_buffer)
@@ -53,11 +57,14 @@ def insert_sorted(self, datum: Datum):
                 left = mid + 1
         self.sorted_buffer.insert(left, datum)
 
-    async def flush_buffer(self, output: NonBlockingIterator):
-        _LOGGER.info(f"Watermark updated, flushing sortedBuffer: {self.latest_wm}")
+    async def flush_buffer(self, output: NonBlockingIterator, flush_all: bool = False):
+        if flush_all:
+            _LOGGER.info("Flushing entire sortedBuffer")
+        else:
+            _LOGGER.info(f"Flushing sortedBuffer above watermark: {self.latest_wm}")
         i = 0
         for datum in self.sorted_buffer:
-            if datum.event_time > self.latest_wm:
+            if datum.event_time > self.latest_wm and not flush_all:
                 break
             await output.put(Message.from_datum(datum))
             _LOGGER.info(f"Sent datum with event time: {datum.event_time}")