perf: cache _replica_dict on Tablet for O(1) host/shard lookup

Build a {host_id: shard_id} dict once at Tablet construction time so
that policies.py and pool.py can replace set(map(lambda ...)) and
linear scans with O(1) dict operations.

- Add _replica_dict to __slots__
- Build dict from the materialized tuple (not the raw replicas arg)
  to avoid double-consuming a one-shot iterator
- Update DCAwareRoundRobinPolicy to use tablet._replica_dict keys
- Update HostConnection to use tablet._replica_dict.get() for shard
- Rewrite replica_contains_host_id() to use dict membership
- Add 7 unit tests covering dict construction, lookup, host membership,
  tuple storage, and the iterator edge case

Author: mykaul

cassandra/policies.py

@@ -507,10 +507,10 @@ def make_query_plan(self, working_keyspace=None, query=None):
             keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key))
 
         if tablet is not None:
-            replicas_mapped = set(map(lambda r: r[0], tablet.replicas))
+            replica_dict = tablet._replica_dict
             child_plan = child.make_query_plan(keyspace, query)
 
-            replicas = [host for host in child_plan if host.host_id in replicas_mapped]
+            replicas = [host for host in child_plan if host.host_id in replica_dict]
         else:
             replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
 

cassandra/pool.py

@@ -462,10 +462,7 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table
                 tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t)
 
                 if tablet is not None:
-                    for replica in tablet.replicas:
-                        if replica[0] == self.host.host_id:
-                            shard_id = replica[1]
-                            break
+                    shard_id = tablet._replica_dict.get(self.host.host_id)
 
             if shard_id is None:
                 shard_id = self.host.sharding_info.shard_id_from_token(t.value)

cassandra/tablets.py

@@ -15,12 +15,18 @@ class Tablet(object):
     It stores information about each replica, its host and shard,
     and the token interval in the format (first_token, last_token].
     """
-    __slots__ = ('first_token', 'last_token', 'replicas')
+    __slots__ = ('first_token', 'last_token', 'replicas', '_replica_dict')
 
     def __init__(self, first_token=0, last_token=0, replicas=None):
         self.first_token = first_token
         self.last_token = last_token
-        self.replicas = tuple(replicas) if replicas is not None else None
+        if replicas is not None:
+            replicas_tuple = tuple(replicas)
+            self.replicas = replicas_tuple
+            self._replica_dict = {r[0]: r[1] for r in replicas_tuple}
+        else:
+            self.replicas = None
+            self._replica_dict = {}
 
     def __str__(self):
         return "<Tablet: first_token=%s last_token=%s replicas=%s>" \
@@ -39,10 +45,7 @@ def from_row(first_token, last_token, replicas):
         return None
 
     def replica_contains_host_id(self, uuid: UUID) -> bool:
-        for replica in self.replicas:
-            if replica[0] == uuid:
-                return True
-        return False
+        return uuid in self._replica_dict
 
 
 class Tablets(object):

tests/unit/test_tablets.py

@@ -1,4 +1,5 @@
 import unittest
+from uuid import UUID
 
 from cassandra.tablets import Tablets, Tablet
 
@@ -124,3 +125,62 @@ def __init__(self, v):
         # Token value 50 is not > first_token (100) of the tablet whose
         # last_token (200) is >= 50, so no match.
         self.assertIsNone(tablets.get_tablet_for_key("ks", "tb", Token(50)))
+
+
+class TabletReplicaDictTest(unittest.TestCase):
+    """Tests for Tablet._replica_dict cached lookup."""
+
+    def test_replica_dict_built_from_replicas(self):
+        u1 = UUID('12345678-1234-5678-1234-567812345678')
+        u2 = UUID('87654321-4321-8765-4321-876543218765')
+        t = Tablet(0, 100, [(u1, 3), (u2, 7)])
+        self.assertEqual(t._replica_dict, {u1: 3, u2: 7})
+
+    def test_replica_dict_empty_when_no_replicas(self):
+        t = Tablet(0, 100, None)
+        self.assertEqual(t._replica_dict, {})
+
+    def test_replica_dict_contains_host(self):
+        u1 = UUID('12345678-1234-5678-1234-567812345678')
+        u2 = UUID('87654321-4321-8765-4321-876543218765')
+        u3 = UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee')
+        t = Tablet(0, 100, [(u1, 3), (u2, 7)])
+        self.assertIn(u1, t._replica_dict)
+        self.assertIn(u2, t._replica_dict)
+        self.assertNotIn(u3, t._replica_dict)
+
+    def test_replica_dict_shard_lookup(self):
+        u1 = UUID('12345678-1234-5678-1234-567812345678')
+        u2 = UUID('87654321-4321-8765-4321-876543218765')
+        t = Tablet(0, 100, [(u1, 3), (u2, 7)])
+        self.assertEqual(t._replica_dict.get(u1), 3)
+        self.assertEqual(t._replica_dict.get(u2), 7)
+        self.assertIsNone(t._replica_dict.get(UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee')))
+
+    def test_replica_contains_host_id_uses_dict(self):
+        u1 = UUID('12345678-1234-5678-1234-567812345678')
+        u2 = UUID('87654321-4321-8765-4321-876543218765')
+        t = Tablet(0, 100, [(u1, 3), (u2, 7)])
+        self.assertTrue(t.replica_contains_host_id(u1))
+        self.assertTrue(t.replica_contains_host_id(u2))
+        self.assertFalse(t.replica_contains_host_id(UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee')))
+
+    def test_replicas_stored_as_tuple(self):
+        t = Tablet(0, 100, [("host1", 0), ("host2", 1)])
+        self.assertIsInstance(t.replicas, tuple)
+
+    def test_replica_dict_from_iterator(self):
+        """Ensure _replica_dict is correctly built even when replicas is a
+        one-shot iterator (generator), not a reusable list."""
+        u1 = UUID('12345678-1234-5678-1234-567812345678')
+        u2 = UUID('87654321-4321-8765-4321-876543218765')
+
+        def gen():
+            yield (u1, 3)
+            yield (u2, 7)
+
+        t = Tablet(0, 100, gen())
+        self.assertEqual(t.replicas, ((u1, 3), (u2, 7)))
+        self.assertEqual(t._replica_dict, {u1: 3, u2: 7})
+        self.assertTrue(t.replica_contains_host_id(u1))
+        self.assertTrue(t.replica_contains_host_id(u2))