Abstract away the state update cache (#576)

Author: daniil-quix

quixstreams/state/base/partition.py

@@ -12,7 +12,7 @@
 from quixstreams.state.serialization import DumpsFunc, LoadsFunc
 from quixstreams.utils.json import loads as json_loads
 
-from .transaction import PartitionTransaction, CACHE_TYPE
+from .transaction import PartitionTransaction, PartitionTransactionCache
 
 if TYPE_CHECKING:
     from quixstreams.state.recovery import ChangelogProducer
@@ -72,14 +72,14 @@ def get_changelog_offset(self) -> Optional[int]:
     @abstractmethod
     def write(
         self,
-        data: CACHE_TYPE,
+        cache: PartitionTransactionCache,
         processed_offset: Optional[int],
         changelog_offset: Optional[int],
     ):
         """
-        Update the state with data
+        Update the state with data from the update cache
 
-        :param data: The modified data
+        :param cache: The modified data
         :param processed_offset: The offset processed to generate the data.
         :param changelog_offset: The changelog message offset of the data.
         """

quixstreams/state/base/transaction.py

@@ -1,21 +1,11 @@
 import enum
-import logging
 import functools
-from typing import (
-    Any,
-    Optional,
-    Dict,
-    Tuple,
-    Union,
-    TYPE_CHECKING,
-)
-
+import logging
 from abc import ABC
+from collections import defaultdict
+from typing import Any, Optional, Dict, Tuple, Union, TYPE_CHECKING, Set
 
-from quixstreams.state.exceptions import (
-    StateTransactionError,
-    InvalidChangelogOffset,
-)
+from quixstreams.state.exceptions import StateTransactionError, InvalidChangelogOffset
 from quixstreams.state.metadata import (
     CHANGELOG_CF_MESSAGE_HEADER,
     CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER,
@@ -25,24 +15,129 @@
     Undefined,
     DEFAULT_PREFIX,
 )
-from quixstreams.state.serialization import (
-    serialize,
-    deserialize,
-    LoadsFunc,
-    DumpsFunc,
-)
+from quixstreams.state.serialization import serialize, deserialize, LoadsFunc, DumpsFunc
 from quixstreams.utils.json import dumps as json_dumps
-
 from .state import State, TransactionState
 
 if TYPE_CHECKING:
     from quixstreams.state.recovery import ChangelogProducer
     from .partition import StorePartition
 
-__all__ = ("PartitionTransactionStatus", "PartitionTransaction", "CACHE_TYPE")
+__all__ = (
+    "PartitionTransactionStatus",
+    "PartitionTransaction",
+    "PartitionTransactionCache",
+)
 
 logger = logging.getLogger(__name__)
-CACHE_TYPE = Dict[str, Dict[bytes, Dict[bytes, Union[bytes, Undefined]]]]
+
+
+class PartitionTransactionCache:
+    """
+    A cache with the data updated in the current PartitionTransaction.
+    It is used to read-your-own-writes before the transaction is committed to the Store.
+
+    Internally, updates and deletes are separated into two separate structures
+    to simplify the querying over them.
+    """
+
+    def __init__(self):
+        # A map with updated keys in format {<cf>: {<prefix>: {<key>: <value>}}}
+        # Note: "updates" are bucketed per prefix to speed up iterating over the
+        # specific set of keys when we merge updates with data from the stores.
+        # Using a prefix like that allows us to perform fewer iterations.
+        self._updated: dict[str, dict[bytes, dict[bytes, bytes]]] = defaultdict(
+            lambda: defaultdict(dict)
+        )
+        # Dict of sets with deleted keys in format {<cf>: set[<key1>, <key2>]}
+        # Deletes are stored without prefixes because we don't need to iterate over
+        # them.
+        self._deleted: dict[str, set[bytes]] = defaultdict(set)
+        self._empty = True
+
+    def get(
+        self,
+        key: bytes,
+        prefix: bytes,
+        cf_name: str = "default",
+    ) -> Union[bytes, Undefined]:
+        """
+        Get a value for the key.
+
+        Returns the key value if it has been updated during the transaction.
+
+        If the key has already been deleted, returns "DELETED" sentinel
+        (we don't need to check the actual store).
+        If the key is not present in the cache, returns "UNDEFINED sentinel
+        (we need to check the store).
+
+        :param: key: key as bytes
+        :param: prefix: key prefix as bytes
+        :param: cf_name: column family name
+        """
+        # Check if the key has been deleted
+        if key in self._deleted[cf_name]:
+            # The key is deleted and the store doesn't need to be checked
+            return DELETED
+
+        # Check if the key has been updated
+        # If the key is not present in the cache, we need to check the store and return
+        # UNDEFINED to signify that
+        return self._updated[cf_name][prefix].get(key, UNDEFINED)
+
+    def set(self, key: bytes, value: bytes, prefix: bytes, cf_name: str = "default"):
+        """
+        Set a value for the key.
+
+        :param: key: key as bytes
+        :param: value: value as bytes
+        :param: prefix: key prefix as bytes
+        :param: cf_name: column family name
+        """
+        self._updated[cf_name][prefix][key] = value
+        self._deleted[cf_name].discard(key)
+        self._empty = False
+
+    def delete(self, key: Any, prefix: bytes, cf_name: str = "default"):
+        """
+        Delete a key.
+
+        :param: key: key as bytes
+        :param: value: value as bytes
+        :param: prefix: key prefix as bytes
+        :param: cf_name: column family name
+        """
+        self._updated[cf_name][prefix].pop(key, None)
+        self._deleted[cf_name].add(key)
+        self._empty = False
+
+    def is_empty(self) -> bool:
+        """
+        Return True if any changes have been made (updates or deletes), otherwise
+        return False.
+        """
+        return self._empty
+
+    def get_column_families(self) -> Set[str]:
+        """
+        Get all update column families.
+        """
+        return set(self._updated.keys()) | set(self._deleted.keys())
+
+    def get_updates(self, cf_name: str = "default") -> Dict[bytes, Dict[bytes, bytes]]:
+        """
+        Get all updated keys (excluding deleted)
+        in the format "{<prefix>: {<key>: <value>}}".
+
+        :param: cf_name: column family name
+        """
+        return self._updated.get(cf_name, {})
+
+    def get_deletes(self, cf_name: str = "default") -> Set[bytes]:
+        """
+        Get all deleted keys (excluding updated) as a set.
+        """
+        return self._deleted[cf_name]
 
 
 class PartitionTransactionStatus(enum.Enum):
@@ -97,7 +192,7 @@ def __init__(
         self._loads = loads
         self._partition = partition
 
-        self._update_cache: CACHE_TYPE = {}
+        self._update_cache = PartitionTransactionCache()
 
     @property
     def changelog_producer(self) -> Optional["ChangelogProducer"]:
@@ -197,14 +292,13 @@ def get(
         :param key: key
         :param prefix: a key prefix
         :param default: default value to return if the key is not found
+        :param cf_name: column family name
         :return: value or None if the key is not found and `default` is not provided
         """
         key_serialized = self._serialize_key(key, prefix=prefix)
 
-        cached = (
-            self._update_cache.get(cf_name, {})
-            .get(prefix, {})
-            .get(key_serialized, UNDEFINED)
+        cached = self._update_cache.get(
+            key=key_serialized, prefix=prefix, cf_name=cf_name
         )
         if cached is DELETED:
             return default
@@ -225,14 +319,18 @@ def set(self, key: Any, value: Any, prefix: bytes, cf_name: str = "default"):
         :param key: key
         :param prefix: a key prefix
         :param value: value
+        :param cf_name: column family name
         """
 
         try:
             key_serialized = self._serialize_key(key, prefix=prefix)
             value_serialized = self._serialize_value(value)
-            self._update_cache.setdefault(cf_name, {}).setdefault(prefix, {})[
-                key_serialized
-            ] = value_serialized
+            self._update_cache.set(
+                key=key_serialized,
+                value=value_serialized,
+                prefix=prefix,
+                cf_name=cf_name,
+            )
         except Exception:
             self._status = PartitionTransactionStatus.FAILED
             raise
@@ -245,12 +343,13 @@ def delete(self, key: Any, prefix: bytes, cf_name: str = "default"):
         This function always returns `None`, even if value is not found.
         :param key: key
         :param prefix: a key prefix
+        :param cf_name: column family name
         """
         try:
             key_serialized = self._serialize_key(key, prefix=prefix)
-            self._update_cache.setdefault(cf_name, {}).setdefault(prefix, {})[
-                key_serialized
-            ] = DELETED
+            self._update_cache.delete(
+                key=key_serialized, prefix=prefix, cf_name=cf_name
+            )
         except Exception:
             self._status = PartitionTransactionStatus.FAILED
             raise
@@ -261,21 +360,19 @@ def exists(self, key: Any, prefix: bytes, cf_name: str = "default") -> bool:
         Check if the key exists in state.
         :param key: key
         :param prefix: a key prefix
+        :param cf_name: column family name
         :return: True if key exists, False otherwise
         """
         key_serialized = self._serialize_key(key, prefix=prefix)
-        cached = (
-            self._update_cache.get(cf_name, {})
-            .get(prefix, {})
-            .get(key_serialized, UNDEFINED)
+        cached = self._update_cache.get(
+            key=key_serialized, prefix=prefix, cf_name=cf_name
         )
         if cached is DELETED:
             return False
-
-        if cached is not UNDEFINED:
+        elif cached is not UNDEFINED:
             return True
-
-        return self._partition.exists(key_serialized, cf_name=cf_name)
+        else:
+            return self._partition.exists(key_serialized, cf_name=cf_name)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
     def prepare(self, processed_offset: int):
@@ -310,21 +407,32 @@ def _prepare(self, processed_offset: int):
             f"partition={self._changelog_producer.partition} "
             f"processed_offset={processed_offset}"
         )
-        for cf_name, cf_update_cache in self._update_cache.items():
-            source_tp_offset_header = json_dumps(processed_offset)
+        source_tp_offset_header = json_dumps(processed_offset)
+        column_families = self._update_cache.get_column_families()
+
+        for cf_name in column_families:
             headers = {
                 CHANGELOG_CF_MESSAGE_HEADER: cf_name,
                 CHANGELOG_PROCESSED_OFFSET_MESSAGE_HEADER: source_tp_offset_header,
             }
-            for _, prefix_update_cache in cf_update_cache.items():
+
+            updates = self._update_cache.get_updates(cf_name=cf_name)
+            for prefix_update_cache in updates.values():
                 for key, value in prefix_update_cache.items():
-                    # Produce changes to the changelog topic
                     self._changelog_producer.produce(
                         key=key,
-                        value=value if value is not DELETED else None,
+                        value=value,
                         headers=headers,
                     )
 
+            deletes = self._update_cache.get_deletes(cf_name=cf_name)
+            for key in deletes:
+                self._changelog_producer.produce(
+                    key=key,
+                    value=None,
+                    headers=headers,
+                )
+
     @validate_transaction_status(
         PartitionTransactionStatus.STARTED, PartitionTransactionStatus.PREPARED
     )
@@ -357,7 +465,7 @@ def flush(
             raise
 
     def _flush(self, processed_offset: Optional[int], changelog_offset: Optional[int]):
-        if not self._update_cache:
+        if self._update_cache.is_empty():
             return
 
         if changelog_offset is not None:
@@ -371,7 +479,7 @@ def _flush(self, processed_offset: Optional[int], changelog_offset: Optional[int
                 )
 
         self._partition.write(
-            data=self._update_cache,
+            cache=self._update_cache,
             processed_offset=processed_offset,
             changelog_offset=changelog_offset,
         )

quixstreams/state/rocksdb/partition.py

@@ -5,9 +5,8 @@
 from rocksdict import WriteBatch, Rdict, ColumnFamily, AccessType
 
 from quixstreams.models import ConfluentKafkaMessageProto
-from quixstreams.state.metadata import DELETED
 from quixstreams.state.recovery import ChangelogProducer
-from quixstreams.state.base import StorePartition, CACHE_TYPE
+from quixstreams.state.base import StorePartition, PartitionTransactionCache
 from quixstreams.state.serialization import (
     int_from_int64_bytes,
     int_to_int64_bytes,
@@ -113,32 +112,37 @@ def _recover_from_changelog_message(
 
     def write(
         self,
-        data: CACHE_TYPE,
+        cache: PartitionTransactionCache,
         processed_offset: Optional[int],
         changelog_offset: Optional[int],
         batch: Optional[WriteBatch] = None,
     ):
         """
         Write data to RocksDB
 
-        :param data: The modified data
+        :param cache: The modified data
         :param processed_offset: The offset processed to generate the data.
         :param changelog_offset: The changelog message offset of the data.
+        :param batch: prefilled `rocksdict.WriteBatch`, optional.
         """
         if batch is None:
             batch = WriteBatch(raw_mode=True)
 
         meta_cf_handle = self.get_column_family_handle(METADATA_CF_NAME)
+
         # Iterate over the transaction update cache
-        for cf_name, cf_update_cache in data.items():
+        column_families = cache.get_column_families()
+        for cf_name in column_families:
             cf_handle = self.get_column_family_handle(cf_name)
-            for _prefix, prefix_update_cache in cf_update_cache.items():
+
+            updates = cache.get_updates(cf_name=cf_name)
+            for prefix_update_cache in updates.values():
                 for key, value in prefix_update_cache.items():
-                    # Apply changes to the Writebatch
-                    if value is DELETED:
-                        batch.delete(key, cf_handle)
-                    else:
-                        batch.put(key, value, cf_handle)
+                    batch.put(key, value, cf_handle)
+
+            deletes = cache.get_deletes(cf_name=cf_name)
+            for key in deletes:
+                batch.delete(key, cf_handle)
 
         # Save the latest processed input topic offset
         if processed_offset is not None: