Optimize state operations (#891)

- WindowedRocksDBPartitionTransaction._get_timestamp(): return the cached value if it has been set previously
- set_bytes(): remove extra @validate_transaction_state decorator
- set_bytes(): check if the value is bytes only at the top level

Author: daniil-quix

quixstreams/state/base/transaction.py

@@ -388,7 +388,6 @@ def _get_bytes(
 
         return cached
 
-    @validate_transaction_status(PartitionTransactionStatus.STARTED)
     def set(self, key: K, value: V, prefix: bytes, cf_name: str = "default") -> None:
         """
         Set value for the key.
@@ -404,9 +403,8 @@ def set(self, key: K, value: V, prefix: bytes, cf_name: str = "default") -> None
             self._status = PartitionTransactionStatus.FAILED
             raise
 
-        self.set_bytes(key, value_serialized, prefix, cf_name=cf_name)
+        self._set_bytes(key, value_serialized, prefix, cf_name=cf_name)
 
-    @validate_transaction_status(PartitionTransactionStatus.STARTED)
     def set_bytes(
         self, key: K, value: bytes, prefix: bytes, cf_name: str = "default"
     ) -> None:
@@ -417,10 +415,17 @@ def set_bytes(
         :param value: value
         :param cf_name: column family name
         """
-        try:
-            if not isinstance(value, bytes):
-                raise StateSerializationError("Value must be bytes")
+        if not isinstance(value, bytes):
+            self._status = PartitionTransactionStatus.FAILED
+            raise StateSerializationError("Value must be bytes")
+
+        self._set_bytes(key=key, value=value, prefix=prefix, cf_name=cf_name)
 
+    @validate_transaction_status(PartitionTransactionStatus.STARTED)
+    def _set_bytes(
+        self, key: K, value: bytes, prefix: bytes, cf_name: str = "default"
+    ) -> None:
+        try:
             key_serialized = self._serialize_key(key, prefix=prefix)
             self._update_cache.set(
                 key=key_serialized,

quixstreams/state/rocksdb/windowed/transaction.py

@@ -38,7 +38,7 @@
     from .partition import WindowedRocksDBStorePartition
 
 
-class WindowedRocksDBPartitionTransaction(PartitionTransaction[bytes, dict]):
+class WindowedRocksDBPartitionTransaction(PartitionTransaction[bytes, Any]):
     def __init__(
         self,
         partition: "WindowedRocksDBStorePartition",
@@ -108,13 +108,11 @@ def keys(self, cf_name: str = "default") -> Iterable[Any]:
             yield key
 
     def get_latest_timestamp(self, prefix: bytes) -> int:
-        return self._get_timestamp(
-            prefix=prefix, cache=self._latest_timestamps, default=0
-        )
+        return self._get_timestamp(prefix=prefix, cache=self._latest_timestamps) or 0
 
     def get_latest_expired(self, prefix: bytes) -> int:
-        return self._get_timestamp(
-            prefix=prefix, cache=self._last_expired_timestamps, default=0
+        return (
+            self._get_timestamp(prefix=prefix, cache=self._last_expired_timestamps) or 0
         )
 
     def get_window(
@@ -315,9 +313,7 @@ def expire_all_windows(
         :param delete: If True, expired windows will be deleted.
         :param collect: If True, values will be collected into windows.
         """
-        last_expired = self._get_timestamp(
-            prefix=b"", cache=self._last_expired_timestamps, default=0
-        )
+        last_expired = self.get_latest_expired(prefix=b"")
 
         to_delete: set[tuple[bytes, int, int]] = set()
         collected = []
@@ -526,18 +522,15 @@ def _get_items(
         # Sort and deserialize items merged from the cache and store
         return sorted(merged_items.items(), key=lambda kv: kv[0], reverse=backwards)
 
-    def _get_timestamp(
-        self, cache: TimestampsCache, prefix: bytes, default: Any = None
-    ) -> int:
-        cached_ts = cache.timestamps.get(prefix)
-        if cached_ts is not None:
-            return cached_ts
+    def _get_timestamp(self, cache: TimestampsCache, prefix: bytes) -> Optional[int]:
+        if prefix in cache.timestamps:
+            # Return the cached value if it has been set at least once
+            return cache.timestamps[prefix]
 
         stored_ts = self.get(
             key=cache.key,
             prefix=prefix,
             cf_name=cache.cf_name,
-            default=default,
         )
         if stored_ts is not None and not isinstance(stored_ts, int):
             raise ValueError(f"invalid timestamp {stored_ts}")