Calculate window watermarks by individual keys instead of partitions (#591)

daniil-quix · web-flow · commit 16d8637a6833 · 2024-11-04T10:04:47.000+01:00
diff --git a/examples/bank_example/quix_platform_version/producer.py b/examples/bank_example/quix_platform_version/producer.py
@@ -6,9 +6,7 @@
 from dotenv import load_dotenv
 
 from quixstreams import Application
-from quixstreams.models.serializers import (
-    QuixTimeseriesSerializer,
-)
+from quixstreams.models.serializers import QuixTimeseriesSerializer
 
 load_dotenv("./bank_example/quix_platform_version/quix_vars.env")
 
diff --git a/quixstreams/state/rocksdb/windowed/metadata.py b/quixstreams/state/rocksdb/windowed/metadata.py
@@ -1,5 +1,5 @@
 LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY = b"__expired_start_gt__"
 
 LATEST_EXPIRED_WINDOW_CF_NAME = "__expiration-index__"
-
-LATEST_TIMESTAMP_KEY = b"__topic_latest_timestamp__"
+LATEST_TIMESTAMPS_CF_NAME = "__latest-timestamps__"
+LATEST_TIMESTAMP_KEY = b"__latest_timestamp__"
diff --git a/quixstreams/state/rocksdb/windowed/partition.py b/quixstreams/state/rocksdb/windowed/partition.py
@@ -1,17 +1,17 @@
 import logging
 from typing import Optional
 
-from rocksdict import RdictItems, ReadOptions, WriteBatch  # type: ignore
+from rocksdict import RdictItems, ReadOptions  # type: ignore
 
-from quixstreams.state.base import PartitionTransactionCache
 from quixstreams.state.recovery import ChangelogProducer
-from quixstreams.state.serialization import int_from_int64_bytes, int_to_int64_bytes
 
 from ..exceptions import ColumnFamilyDoesNotExist
-from ..metadata import METADATA_CF_NAME
 from ..partition import RocksDBStorePartition
 from ..types import RocksDBOptionsType
-from .metadata import LATEST_EXPIRED_WINDOW_CF_NAME, LATEST_TIMESTAMP_KEY
+from .metadata import (
+    LATEST_EXPIRED_WINDOW_CF_NAME,
+    LATEST_TIMESTAMPS_CF_NAME,
+)
 from .transaction import WindowedRocksDBPartitionTransaction
 
 logger = logging.getLogger(__name__)
@@ -38,8 +38,8 @@ def __init__(
         super().__init__(
             path=path, options=options, changelog_producer=changelog_producer
         )
-        self._latest_timestamp_ms = self._get_latest_timestamp_from_db()
         self._ensure_column_family(LATEST_EXPIRED_WINDOW_CF_NAME)
+        self._ensure_column_family(LATEST_TIMESTAMPS_CF_NAME)
 
     def iter_items(
         self, from_key: bytes, read_opt: ReadOptions, cf_name: str = "default"
@@ -52,48 +52,11 @@ def begin(self) -> "WindowedRocksDBPartitionTransaction":
             partition=self,
             dumps=self._dumps,
             loads=self._loads,
-            latest_timestamp_ms=self._latest_timestamp_ms,
             changelog_producer=self._changelog_producer,
         )
 
-    def set_latest_timestamp(self, timestamp_ms: int):
-        self._latest_timestamp_ms = timestamp_ms
-
-    def _get_latest_timestamp_from_db(self) -> int:
-        value = self.get(LATEST_TIMESTAMP_KEY, cf_name=METADATA_CF_NAME)
-        if value is None:
-            return 0
-        return int_from_int64_bytes(value)
-
     def _ensure_column_family(self, cf_name: str):
         try:
             self.get_column_family(cf_name)
         except ColumnFamilyDoesNotExist:
             self.create_column_family(cf_name)
-
-    def write(
-        self,
-        cache: PartitionTransactionCache,
-        processed_offset: Optional[int],
-        changelog_offset: Optional[int],
-        batch: Optional[WriteBatch] = None,
-        latest_timestamp_ms: Optional[int] = None,
-    ):
-        batch = WriteBatch(raw_mode=True)
-
-        if latest_timestamp_ms is not None:
-            cf_handle = self.get_column_family_handle(METADATA_CF_NAME)
-            batch.put(
-                LATEST_TIMESTAMP_KEY,
-                int_to_int64_bytes(latest_timestamp_ms),
-                cf_handle,
-            )
-        super().write(
-            cache=cache,
-            processed_offset=processed_offset,
-            changelog_offset=changelog_offset,
-            batch=batch,
-        )
-
-        if latest_timestamp_ms is not None:
-            self.set_latest_timestamp(latest_timestamp_ms)
diff --git a/quixstreams/state/rocksdb/windowed/state.py b/quixstreams/state/rocksdb/windowed/state.py
@@ -59,15 +59,15 @@ def update_window(self, start_ms: int, end_ms: int, value: Any, timestamp_ms: in
 
     def get_latest_timestamp(self) -> int:
         """
-        Get the latest observed timestamp for the current state partition.
+        Get the latest observed timestamp for the current message key.
 
         Use this timestamp to determine if the arriving event is late and should be
         discarded from the processing.
 
         :return: latest observed event timestamp in milliseconds
         """
 
-        return self._transaction.get_latest_timestamp()
+        return self._transaction.get_latest_timestamp(prefix=self._prefix)
 
     def expire_windows(
         self, duration_ms: int, grace_ms: int = 0
diff --git a/quixstreams/state/rocksdb/windowed/transaction.py b/quixstreams/state/rocksdb/windowed/transaction.py
@@ -4,16 +4,16 @@
 from rocksdict import ReadOptions
 
 from quixstreams.state.base.transaction import PartitionTransaction
-from quixstreams.state.exceptions import InvalidChangelogOffset
 from quixstreams.state.metadata import DEFAULT_PREFIX, PREFIX_SEPARATOR
 from quixstreams.state.recovery import ChangelogProducer
-from quixstreams.state.serialization import (
-    DumpsFunc,
-    LoadsFunc,
-    serialize,
-)
+from quixstreams.state.serialization import DumpsFunc, LoadsFunc, serialize
 
-from .metadata import LATEST_EXPIRED_WINDOW_CF_NAME, LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY
+from .metadata import (
+    LATEST_EXPIRED_WINDOW_CF_NAME,
+    LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY,
+    LATEST_TIMESTAMP_KEY,
+    LATEST_TIMESTAMPS_CF_NAME,
+)
 from .serialization import encode_window_key, encode_window_prefix, parse_window_key
 from .state import WindowedTransactionState
 
@@ -22,14 +22,11 @@
 
 
 class WindowedRocksDBPartitionTransaction(PartitionTransaction):
-    __slots__ = ("_latest_timestamp_ms",)
-
     def __init__(
         self,
         partition: "WindowedRocksDBStorePartition",
         dumps: DumpsFunc,
         loads: LoadsFunc,
-        latest_timestamp_ms: int,
         changelog_producer: Optional[ChangelogProducer] = None,
     ):
         super().__init__(
@@ -39,7 +36,11 @@ def __init__(
             changelog_producer=changelog_producer,
         )
         self._partition = cast("WindowedRocksDBStorePartition", self._partition)
-        self._latest_timestamp_ms = latest_timestamp_ms
+        # Cache the metadata separately to avoid serdes on each access
+        # (we are 100% sure that the underlying types are immutable, while windows'
+        # values are not)
+        self._latest_timestamps: dict[bytes, int] = {}
+        self._last_expired_timestamps: dict[bytes, int] = {}
 
     def as_state(self, prefix: Any = DEFAULT_PREFIX) -> WindowedTransactionState:
         return WindowedTransactionState(
@@ -51,15 +52,19 @@ def as_state(self, prefix: Any = DEFAULT_PREFIX) -> WindowedTransactionState:
             ),
         )
 
-    def get_latest_timestamp(self) -> int:
-        return self._latest_timestamp_ms
+    def get_latest_timestamp(self, prefix: bytes) -> int:
+        cached_ts = self._latest_timestamps.get(prefix)
+        if cached_ts is not None:
+            return cached_ts
 
-    def _validate_duration(self, start_ms: int, end_ms: int):
-        if end_ms <= start_ms:
-            raise ValueError(
-                f"Invalid window duration: window end {end_ms} is smaller or equal "
-                f"than window start {start_ms}"
-            )
+        stored_ts = self.get(
+            key=LATEST_TIMESTAMP_KEY,
+            prefix=prefix,
+            cf_name=LATEST_TIMESTAMPS_CF_NAME,
+            default=0,
+        )
+        self._latest_timestamps[prefix] = stored_ts
+        return stored_ts
 
     def get_window(
         self,
@@ -81,34 +86,16 @@ def update_window(
 
         key = encode_window_key(start_ms, end_ms)
         self.set(key=key, value=value, prefix=prefix)
-        self._latest_timestamp_ms = max(self._latest_timestamp_ms, timestamp_ms)
+        latest_timestamp_ms = self.get_latest_timestamp(prefix=prefix)
+        self._set_latest_timestamp(
+            prefix=prefix, timestamp_ms=max(latest_timestamp_ms, timestamp_ms)
+        )
 
     def delete_window(self, start_ms: int, end_ms: int, prefix: bytes):
         self._validate_duration(start_ms=start_ms, end_ms=end_ms)
         key = encode_window_key(start_ms, end_ms)
         self.delete(key=key, prefix=prefix)
 
-    def _flush(self, processed_offset: Optional[int], changelog_offset: Optional[int]):
-        if self._update_cache.is_empty():
-            return
-
-        if changelog_offset is not None:
-            current_changelog_offset = self._partition.get_changelog_offset()
-            if (
-                current_changelog_offset is not None
-                and changelog_offset < current_changelog_offset
-            ):
-                raise InvalidChangelogOffset(
-                    "Cannot set changelog offset lower than already saved one"
-                )
-
-        self._partition.write(
-            cache=self._update_cache,
-            processed_offset=processed_offset,
-            changelog_offset=changelog_offset,
-            latest_timestamp_ms=self._latest_timestamp_ms,
-        )
-
     def expire_windows(
         self, duration_ms: int, prefix: bytes, grace_ms: int = 0
     ) -> list[tuple[tuple[int, int], Any]]:
@@ -134,16 +121,12 @@ def expire_windows(
             Defaults to 0, meaning no grace period is applied.
         :return: A generator that yields sorted tuples in the format `((start, end), value)`.
         """
-        latest_timestamp = self._latest_timestamp_ms
+        latest_timestamp = self.get_latest_timestamp(prefix=prefix)
         start_to = latest_timestamp - duration_ms - grace_ms
         start_from = -1
 
         # Find the latest start timestamp of the expired windows for the given key
-        last_expired = self.get(
-            key=LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY,
-            prefix=prefix,
-            cf_name=LATEST_EXPIRED_WINDOW_CF_NAME,
-        )
+        last_expired = self._get_last_expired_timestamp(prefix=prefix)
         if last_expired is not None:
             start_from = max(start_from, last_expired)
 
@@ -160,22 +143,15 @@ def expire_windows(
             # Save the start of the latest expired window to the expiration index
             latest_window = expired_windows[-1]
             last_expired__gt = latest_window[0][0]
-            self.set(
-                key=LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY,
-                value=last_expired__gt,
-                prefix=prefix,
-                cf_name=LATEST_EXPIRED_WINDOW_CF_NAME,
+
+            self._set_last_expired_timestamp(
+                prefix=prefix, timestamp_ms=last_expired__gt
             )
             # Delete expired windows from the state
             for (start, end), _ in expired_windows:
                 self.delete_window(start, end, prefix=prefix)
         return expired_windows
 
-    def _serialize_key(self, key: Any, prefix: bytes) -> bytes:
-        # Allow bytes keys in WindowedStore
-        key_bytes = key if isinstance(key, bytes) else serialize(key, dumps=self._dumps)
-        return prefix + PREFIX_SEPARATOR + key_bytes
-
     def get_windows(
         self,
         start_from_ms: int,
@@ -240,3 +216,46 @@ def get_windows(
                 result.append(((start, end), value))
 
         return result
+
+    def _set_latest_timestamp(self, prefix: bytes, timestamp_ms: int):
+        self._latest_timestamps[prefix] = timestamp_ms
+        self.set(
+            key=LATEST_TIMESTAMP_KEY,
+            value=timestamp_ms,
+            prefix=prefix,
+            cf_name=LATEST_TIMESTAMPS_CF_NAME,
+        )
+
+    def _get_last_expired_timestamp(self, prefix: bytes) -> Optional[int]:
+        cached_ts = self._last_expired_timestamps.get(prefix)
+        if cached_ts is not None:
+            return cached_ts
+
+        stored_ts = self.get(
+            key=LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY,
+            prefix=prefix,
+            cf_name=LATEST_EXPIRED_WINDOW_CF_NAME,
+        )
+        self._last_expired_timestamps[prefix] = stored_ts
+        return stored_ts
+
+    def _set_last_expired_timestamp(self, prefix: bytes, timestamp_ms: int):
+        self._last_expired_timestamps[prefix] = timestamp_ms
+        self.set(
+            key=LATEST_EXPIRED_WINDOW_TIMESTAMP_KEY,
+            value=timestamp_ms,
+            prefix=prefix,
+            cf_name=LATEST_EXPIRED_WINDOW_CF_NAME,
+        )
+
+    def _validate_duration(self, start_ms: int, end_ms: int):
+        if end_ms <= start_ms:
+            raise ValueError(
+                f"Invalid window duration: window end {end_ms} is smaller or equal "
+                f"than window start {start_ms}"
+            )
+
+    def _serialize_key(self, key: Any, prefix: bytes) -> bytes:
+        # Allow bytes keys in WindowedStore
+        key_bytes = key if isinstance(key, bytes) else serialize(key, dumps=self._dumps)
+        return prefix + PREFIX_SEPARATOR + key_bytes
diff --git a/quixstreams/state/types.py b/quixstreams/state/types.py
@@ -163,9 +163,10 @@ def update_window(
         """
         ...
 
-    def get_latest_timestamp(self) -> int:
+    def get_latest_timestamp(self, prefix: bytes) -> int:
         """
-        Get the latest observed timestamp for the current state partition.
+        Get the latest observed timestamp for the current state prefix
+        (same as message key).
 
         Use this timestamp to determine if the arriving event is late and should be
         discarded from the processing.
diff --git a/tests/test_quixstreams/test_app.py b/tests/test_quixstreams/test_app.py
diff --git a/tests/test_quixstreams/test_state/test_rocksdb/test_windowed/test_state.py b/tests/test_quixstreams/test_state/test_rocksdb/test_windowed/test_state.py
diff --git a/tests/test_quixstreams/test_state/test_rocksdb/test_windowed/test_transaction.py b/tests/test_quixstreams/test_state/test_rocksdb/test_windowed/test_transaction.py
