Replace cache-level iter-items with get_updates_for_prefix

Author: gwaramadze

quixstreams/state/base/transaction.py

@@ -6,7 +6,6 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
     Generic,
     Optional,
     Set,
@@ -97,21 +96,6 @@ def get(
         # UNDEFINED to signify that
         return self._updated[cf_name][prefix].get(key, Marker.UNDEFINED)
 
-    def iter_items(
-        self,
-        prefix: bytes,
-        backwards: bool = False,
-        cf_name: str = "default",
-    ) -> list[tuple[bytes, bytes]]:
-        """
-        Iterate over sorted items in the cache
-        for the given prefix and column family.
-        """
-        return sorted(
-            self._updated[cf_name][prefix].items(),
-            reverse=backwards,
-        )
-
     def set(self, key: bytes, value: bytes, prefix: bytes, cf_name: str = "default"):
         """
         Set a value for the key.
@@ -150,15 +134,29 @@ def get_column_families(self) -> Set[str]:
         """
         return set(self._updated.keys()) | set(self._deleted.keys())
 
-    def get_updates(self, cf_name: str = "default") -> Dict[bytes, Dict[bytes, bytes]]:
+    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>}}".
+        in the format "{<prefix>: {<key>: <value>, ...}, ...}".
 
         :param: cf_name: column family name
         """
         return self._updated.get(cf_name, {})
 
+    def get_updates_for_prefix(
+        self,
+        prefix: bytes,
+        cf_name: str = "default",
+    ) -> dict[bytes, bytes]:
+        """
+        Get all updated keys (excluding deleted)
+        in the format "{<key>: <value>, ...}".
+
+        :param: prefix: key prefix
+        :param: cf_name: column family name
+        """
+        return self._updated.get(cf_name, {}).get(prefix, {})
+
     def get_deletes(self, cf_name: str = "default") -> Set[bytes]:
         """
         Get all deleted keys (excluding updated) as a set.

quixstreams/state/rocksdb/timestamped.py

@@ -61,12 +61,11 @@ def get_last(
         value: Optional[bytes] = None
 
         deletes = self._update_cache.get_deletes(cf_name=cf_name)
-
-        cached = self._update_cache.iter_items(
-            prefix=prefix,
-            backwards=True,
-            cf_name=cf_name,
+        updates = self._update_cache.get_updates_for_prefix(
+            cf_name=cf_name, prefix=prefix
         )
+
+        cached = sorted(updates.items(), reverse=True)
         for cached_key, cached_value in cached:
             if prefix < cached_key < key and cached_key not in deletes:
                 value = cached_value
@@ -110,8 +109,12 @@ def set(self, timestamp: int, value: Any, prefix: Any, cf_name: str = "default")
     def expire(self, timestamp: int, prefix: bytes, cf_name: str = "default"):
         key = self._serialize_key(timestamp + 1, prefix)
 
-        cached = self._update_cache.iter_items(prefix=prefix, cf_name=cf_name)
-        for cached_key, _ in cached:
+        cached = self._update_cache.get_updates_for_prefix(
+            prefix=prefix,
+            cf_name=cf_name,
+        )
+        # Cast to list to avoid RuntimeError: dictionary changed size during iteration
+        for cached_key in list(cached):
             if cached_key < key:
                 self._update_cache.delete(cached_key, prefix, cf_name=cf_name)
 

tests/test_quixstreams/test_state/test_transaction.py

@@ -534,6 +534,29 @@ def test_get_updates_after_delete(self, cache: PartitionTransactionCache):
         cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
         assert cache.get_updates(cf_name="cf_name") == {b"prefix": {}}
 
+    def test_get_updates_for_prefix_empty(self, cache: PartitionTransactionCache):
+        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
+
+        # Delete an item and make sure it's not in "updates"
+        cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
+        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
+
+    def test_get_updates_for_prefix_present(self, cache: PartitionTransactionCache):
+        cache.set(key=b"key", value=b"value", prefix=b"prefix", cf_name="cf_name")
+        cache.set(key=b"key", value=b"value", prefix=b"other_prefix", cf_name="cf_name")
+        cache.set(key=b"key", value=b"value", prefix=b"other", cf_name="other_cf_name")
+
+        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {
+            b"key": b"value"
+        }
+
+    def test_get_updates_for_prefix_after_delete(
+        self, cache: PartitionTransactionCache
+    ):
+        cache.set(key=b"key", value=b"value", prefix=b"prefix", cf_name="cf_name")
+        cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
+        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
+
     def test_get_deletes_empty(self, cache: PartitionTransactionCache):
         assert cache.get_deletes(cf_name="cf_name") == set()
 
@@ -567,25 +590,3 @@ def test_get_deletes_after_set(self, cache: PartitionTransactionCache):
     def test_empty(self, action, expected, cache):
         action(cache)
         assert cache.is_empty() == expected
-
-    def test_iter_items(self, cache: PartitionTransactionCache):
-        cache.set(key=b"key1", value=b"value1", prefix=b"prefix")
-        cache.set(key=b"key2", value=b"value2", prefix=b"prefix")
-        cache.set(key=b"key3", value=b"value3", prefix=b"prefix")
-        cache.set(key=b"key4", value=b"value4", prefix=b"prefix", cf_name="other")
-
-        assert cache.iter_items(prefix=b"prefix") == [
-            (b"key1", b"value1"),
-            (b"key2", b"value2"),
-            (b"key3", b"value3"),
-        ]
-
-        assert cache.iter_items(prefix=b"prefix", backwards=True) == [
-            (b"key3", b"value3"),
-            (b"key2", b"value2"),
-            (b"key1", b"value1"),
-        ]
-
-        assert cache.iter_items(prefix=b"prefix", cf_name="other") == [
-            (b"key4", b"value4"),
-        ]