Feature: TimestampedStore

quixstreams/state/base/partition.py

@@ -25,6 +25,8 @@ class StorePartition(ABC):
     the persistent storage).
     """
 
+    partition_transaction_class = PartitionTransaction
+
     def __init__(
         self,
         dumps: DumpsFunc,
@@ -112,7 +114,7 @@ def begin(self) -> PartitionTransaction:
 
         Using `PartitionTransaction` is a recommended way for accessing the data.
         """
-        return PartitionTransaction(
+        return self.partition_transaction_class(
             partition=self,
             dumps=self._dumps,
             loads=self._loads,

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/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,8 +1,8 @@
 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.exceptions import ColumnFamilyDoesNotExist
@@ -42,6 +42,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 +62,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 +143,53 @@ 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 exists(self, key: bytes, cf_name: str = "default") -> bool:
         """
         Check if a key is present in the DB.
@@ -328,3 +379,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
         )