Feature: TimestampedStore (#841)

Co-authored-by: Daniil Gusev <daniil@quix.io>
Co-authored-by: Quentin Dawans <quentin@quix.io>

Author: 3 people

quixstreams/sources/base/source.py

@@ -528,7 +528,7 @@ def flush(self, timeout: Optional[float] = None) -> None:
         :raises CheckpointProducerTimeout: if any message fails to produce before the timeout
         """
         if self._store_transaction:
-            self._store_transaction.prepare(None)
+            self._store_transaction.prepare()
             self._store_transaction.flush()
             self._store_transaction = None
             self._store_state = None

quixstreams/state/base/partition.py

@@ -106,18 +106,14 @@ def recover_from_changelog_message(
         :param offset: changelog message offset
         """
 
+    @abstractmethod
     def begin(self) -> PartitionTransaction:
         """
         Start a new `PartitionTransaction`
 
         Using `PartitionTransaction` is a recommended way for accessing the data.
         """
-        return PartitionTransaction(
-            partition=self,
-            dumps=self._dumps,
-            loads=self._loads,
-            changelog_producer=self._changelog_producer,
-        )
+        ...
 
     def __enter__(self):
         return self

quixstreams/state/base/transaction.py

@@ -44,6 +44,7 @@
     "PartitionTransactionStatus",
     "PartitionTransaction",
     "PartitionTransactionCache",
+    "validate_transaction_status",
 )
 
 logger = logging.getLogger(__name__)
@@ -471,7 +472,7 @@ def exists(self, key: K, prefix: bytes, cf_name: str = "default") -> bool:
             return self._partition.exists(key_serialized, cf_name=cf_name)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def prepare(self, processed_offsets: Optional[dict[str, int]]):
+    def prepare(self, processed_offsets: Optional[dict[str, int]] = None) -> None:
         """
         Produce changelog messages to the changelog topic for all changes accumulated
         in this transaction and prepare transaction to flush its state to the state
@@ -581,4 +582,5 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_val is None and not self.failed:
+            self.prepare()
             self.flush()

quixstreams/state/manager.py

@@ -15,6 +15,7 @@
 from .memory import MemoryStore
 from .recovery import ChangelogProducerFactory, RecoveryManager
 from .rocksdb import RocksDBOptionsType, RocksDBStore
+from .rocksdb.timestamped import TimestampedStore
 from .rocksdb.windowed.store import WindowedRocksDBStore
 
 __all__ = ("StateStoreManager", "DEFAULT_STATE_STORE_NAME", "StoreTypes")
@@ -24,7 +25,7 @@
 
 DEFAULT_STATE_STORE_NAME = "default"
 
-StoreTypes = Union[Type[RocksDBStore], Type[MemoryStore]]
+StoreTypes = Union[Type[RocksDBStore], Type[MemoryStore], Type[TimestampedStore]]
 SUPPORTED_STORES = [RocksDBStore, MemoryStore]
 
 
@@ -189,23 +190,31 @@ def register_store(
 
             store_type = store_type or self.default_store_type
             if store_type == RocksDBStore:
-                factory: Store = RocksDBStore(
+                store: Store = RocksDBStore(
+                    name=store_name,
+                    stream_id=stream_id,
+                    base_dir=str(self._state_dir),
+                    changelog_producer_factory=changelog_producer_factory,
+                    options=self._rocksdb_options,
+                )
+            elif store_type == TimestampedStore:
+                store = TimestampedStore(
                     name=store_name,
                     stream_id=stream_id,
                     base_dir=str(self._state_dir),
                     changelog_producer_factory=changelog_producer_factory,
                     options=self._rocksdb_options,
                 )
             elif store_type == MemoryStore:
-                factory = MemoryStore(
+                store = MemoryStore(
                     name=store_name,
                     stream_id=stream_id,
                     changelog_producer_factory=changelog_producer_factory,
                 )
             else:
                 raise ValueError(f"invalid store type: {store_type}")
 
-            self._stores.setdefault(stream_id, {})[store_name] = factory
+            self._stores.setdefault(stream_id, {})[store_name] = store
 
     def register_windowed_store(
         self,

quixstreams/state/memory/partition.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Any, Dict, Literal, Optional, Union
 
+from quixstreams.state import PartitionTransaction
 from quixstreams.state.base import PartitionTransactionCache, StorePartition
 from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.metadata import METADATA_CF_NAME, Marker
@@ -62,6 +63,14 @@ def closed(self) -> bool:
     def close(self) -> None:
         self._closed = True
 
+    def begin(self) -> PartitionTransaction:
+        return PartitionTransaction(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )
+
     @_validate_partition_state()
     def write(
         self,

quixstreams/state/rocksdb/cache.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class TimestampsCache:
+    key: bytes
+    cf_name: str
+    timestamps: dict[bytes, Optional[int]] = field(default_factory=dict)
+
+
+@dataclass
+class CounterCache:
+    key: bytes
+    cf_name: str
+    counter: Optional[int] = None

quixstreams/state/rocksdb/partition.py

@@ -1,17 +1,26 @@
 import logging
 import time
-from typing import Dict, List, Literal, Optional, Union, cast
+from typing import (
+    Dict,
+    Iterator,
+    List,
+    Literal,
+    Optional,
+    Union,
+    cast,
+)
 
-from rocksdict import AccessType, ColumnFamily, Rdict, WriteBatch
+from rocksdict import AccessType, ColumnFamily, Rdict, ReadOptions, WriteBatch
 
-from quixstreams.state.base import PartitionTransactionCache, StorePartition
+from quixstreams.state.base import (
+    PartitionTransaction,
+    PartitionTransactionCache,
+    StorePartition,
+)
 from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.metadata import METADATA_CF_NAME, Marker
 from quixstreams.state.recovery import ChangelogProducer
-from quixstreams.state.serialization import (
-    int_from_int64_bytes,
-    int_to_int64_bytes,
-)
+from quixstreams.state.serialization import int_from_int64_bytes, int_to_int64_bytes
 
 from .exceptions import ColumnFamilyAlreadyExists
 from .metadata import (
@@ -22,6 +31,7 @@
 
 __all__ = ("RocksDBStorePartition",)
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -42,6 +52,8 @@ class RocksDBStorePartition(StorePartition):
     :param options: RocksDB options. If `None`, the default options will be used.
     """
 
+    additional_column_families: tuple[str, ...] = ()
+
     def __init__(
         self,
         path: str,
@@ -60,6 +72,8 @@ def __init__(
         self._db = self._init_rocksdb()
         self._cf_cache: Dict[str, Rdict] = {}
         self._cf_handle_cache: Dict[str, ColumnFamily] = {}
+        for cf_name in self.additional_column_families:
+            self._ensure_column_family(cf_name)
 
     def recover_from_changelog_message(
         self, key: bytes, value: Optional[bytes], cf_name: str, offset: int
@@ -139,6 +153,61 @@ def get(
         # RDict accept Any type as value but we only write bytes so we should only get bytes back.
         return cast(Union[bytes, Literal[Marker.UNDEFINED]], result)
 
+    def iter_items(
+        self,
+        lower_bound: bytes,  # inclusive
+        upper_bound: bytes,  # exclusive
+        backwards: bool = False,
+        cf_name: str = "default",
+    ) -> Iterator[tuple[bytes, bytes]]:
+        """
+        Iterate over key-value pairs within a specified range in a column family.
+
+        :param lower_bound: The lower bound key (inclusive) for the iteration range.
+        :param upper_bound: The upper bound key (exclusive) for the iteration range.
+        :param backwards: If `True`, iterate in reverse order (descending).
+            Default is `False` (ascending).
+        :param cf_name: The name of the column family to iterate over.
+            Default is "default".
+        :return: An iterator yielding (key, value) tuples.
+        """
+        cf = self.get_column_family(cf_name=cf_name)
+
+        # Set iterator bounds to reduce IO by limiting the range of keys fetched
+        read_opt = ReadOptions()
+        read_opt.set_iterate_lower_bound(lower_bound)
+        read_opt.set_iterate_upper_bound(upper_bound)
+
+        from_key = upper_bound if backwards else lower_bound
+
+        # RDict accepts Any type as value but we only write bytes so we should only get bytes back.
+        items = cast(
+            Iterator[tuple[bytes, bytes]],
+            cf.items(from_key=from_key, read_opt=read_opt, backwards=backwards),
+        )
+
+        if not backwards:
+            # NOTE: Forward iteration respects bounds correctly.
+            # Also, we need to use yield from notation to replace RdictItems
+            # with Python-native generator or else garbage collection
+            # will make the result unpredictable.
+            yield from items
+        else:
+            # NOTE: When iterating backwards, the `read_opt` lower bound
+            # is not respected by Rdict for some reason. We need to manually
+            # filter it here.
+            for key, value in items:
+                if lower_bound <= key:
+                    yield key, value
+
+    def begin(self) -> PartitionTransaction:
+        return PartitionTransaction(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )
+
     def exists(self, key: bytes, cf_name: str = "default") -> bool:
         """
         Check if a key is present in the DB.
@@ -328,3 +397,9 @@ def _update_changelog_offset(self, batch: WriteBatch, offset: int):
             int_to_int64_bytes(offset),
             self.get_column_family_handle(METADATA_CF_NAME),
         )
+
+    def _ensure_column_family(self, cf_name: str) -> None:
+        try:
+            self.get_column_family(cf_name)
+        except ColumnFamilyDoesNotExist:
+            self.create_column_family(cf_name)

quixstreams/state/rocksdb/store.py

@@ -23,6 +23,8 @@ class RocksDBStore(Store):
     partitions' transactions.
     """
 
+    store_partition_class = RocksDBStorePartition
+
     def __init__(
         self,
         name: str,
@@ -61,6 +63,6 @@ def create_new_partition(
                 self._changelog_producer_factory.get_partition_producer(partition)
             )
 
-        return RocksDBStorePartition(
+        return self.store_partition_class(
             path=path, options=self._options, changelog_producer=changelog_producer
         )