perf(cosmos): optimize partition key range cache memory usage

Three optimizations to reduce CollectionRoutingMap memory footprint:

1. Convert to compact PKRange namedtuples: Strip 9 unused fields from
   service response during _discard_parent_ranges (single pass, no extra
   iteration). PKRange supports dict-style access for backward compat.

2. Add __slots__ to Range class: Eliminates per-instance __dict__ overhead
   (~100 bytes per Range object, measured at 64 bytes with slots).

3. Skip redundant .upper() calls: Check if string is already uppercase
   before creating a copy. Service returns uppercase hex.

PPCB overhead reduction (tracemalloc, 150 clients):

| Clients | Original | Patched | Reduction |
|---------|----------|---------|-----------|
|      25 |  5.1 MB  | 2.1 MB  |     -58%  |
|      50 | 10.3 MB  | 4.1 MB  |     -60%  |
|     100 | 15.4 MB  | 7.2 MB  |     -53%  |
|     150 | 27.4 MB  | 6.9 MB  |     -74%  |

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: tvaron3

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py

@@ -28,6 +28,7 @@
 from ... import _base
 from ..collection_routing_map import CollectionRoutingMap
 from .. import routing_range
+from ..routing_range import PKRange
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -124,10 +125,18 @@ async def get_range_by_partition_key_range_id(
     def _discard_parent_ranges(partitionKeyRanges):
         parentIds = set()
         for r in partitionKeyRanges:
-            if isinstance(r, dict) and routing_range.PartitionKeyRange.Parents in r:
-                for parentId in r[routing_range.PartitionKeyRange.Parents]:
+            parents = r.get(routing_range.PartitionKeyRange.Parents)
+            if parents:
+                for parentId in parents:
                     parentIds.add(parentId)
-        return (r for r in partitionKeyRanges if r[routing_range.PartitionKeyRange.Id] not in parentIds)
+        return (
+            PKRange(id=r[routing_range.PartitionKeyRange.Id],
+                    minInclusive=r[routing_range.PartitionKeyRange.MinInclusive],
+                    maxExclusive=r[routing_range.PartitionKeyRange.MaxExclusive],
+                    parents=r.get(routing_range.PartitionKeyRange.Parents))
+            for r in partitionKeyRanges
+            if r[routing_range.PartitionKeyRange.Id] not in parentIds
+        )
 
 
 def _second_range_is_after_first_range(range1, range2):

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py

@@ -28,6 +28,7 @@
 from .. import _base
 from .collection_routing_map import CollectionRoutingMap
 from . import routing_range
+from .routing_range import PKRange
 from .routing_range import PartitionKeyRange
 
 _LOGGER = logging.getLogger(__name__)
@@ -125,10 +126,18 @@ def get_range_by_partition_key_range_id(
     def _discard_parent_ranges(partitionKeyRanges):
         parentIds = set()
         for r in partitionKeyRanges:
-            if isinstance(r, dict) and PartitionKeyRange.Parents in r:
-                for parentId in r[PartitionKeyRange.Parents]:
+            parents = r.get(PartitionKeyRange.Parents)
+            if parents:
+                for parentId in parents:
                     parentIds.add(parentId)
-        return (r for r in partitionKeyRanges if r[PartitionKeyRange.Id] not in parentIds)
+        return (
+            PKRange(id=r[PartitionKeyRange.Id],
+                    minInclusive=r[PartitionKeyRange.MinInclusive],
+                    maxExclusive=r[PartitionKeyRange.MaxExclusive],
+                    parents=r.get(PartitionKeyRange.Parents))
+            for r in partitionKeyRanges
+            if r[PartitionKeyRange.Id] not in parentIds
+        )
 
 
 def _second_range_is_after_first_range(range1, range2):

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py

@@ -27,6 +27,31 @@
 import json
 
 
+from collections import namedtuple
+
+_PKRangeBase = namedtuple('PKRange', ['id', 'minInclusive', 'maxExclusive', 'parents'])
+
+
+class PKRange(_PKRangeBase):
+    """Compact partition key range with dict-compatible access."""
+    __slots__ = ()
+
+    def __getitem__(self, key):
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            raise KeyError(key)
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+    def __contains__(self, key):
+        return key in self._fields
+
+    def items(self):
+        return zip(self._fields, self)
+
+
 class PartitionKeyRange(object):
     """Partition Key Range Constants"""
 
@@ -39,6 +64,8 @@ class PartitionKeyRange(object):
 class Range(object):
     """description of class"""
 
+    __slots__ = ('min', 'max', 'isMinInclusive', 'isMaxInclusive')
+
     MinPath = "min"
     MaxPath = "max"
     IsMinInclusivePath = "isMinInclusive"
@@ -50,8 +77,8 @@ def __init__(self, range_min, range_max, isMinInclusive, isMaxInclusive):
         if range_max is None:
             raise ValueError("max is missing")
 
-        self.min = range_min.upper()
-        self.max = range_max.upper()
+        self.min = range_min if range_min == range_min.upper() else range_min.upper()
+        self.max = range_max if range_max == range_max.upper() else range_max.upper()
         self.isMinInclusive = isMinInclusive
         self.isMaxInclusive = isMaxInclusive
 

sdk/cosmos/azure-cosmos/tests/routing/test_pk_range_optimizations.py

@@ -0,0 +1,117 @@
+# The MIT License (MIT)
+# Copyright (c) Microsoft Corporation. All rights reserved.
+
+import sys
+import unittest
+
+import pytest
+
+from azure.cosmos._routing.routing_range import Range, PKRange, PartitionKeyRange
+from azure.cosmos._routing.collection_routing_map import CollectionRoutingMap
+from azure.cosmos._routing.routing_map_provider import PartitionKeyRangeCache
+
+
+@pytest.mark.cosmosEmulator
+class TestPKRangeOptimizations(unittest.TestCase):
+    """Tests for partition key range memory optimizations."""
+
+    def test_pkrange_dict_access(self):
+        """PKRange supports dict-style [key] access for backward compatibility."""
+        pkr = PKRange(id="1", minInclusive="00", maxExclusive="FF", parents=["0"])
+        self.assertEqual(pkr["id"], "1")
+        self.assertEqual(pkr["minInclusive"], "00")
+        self.assertEqual(pkr["maxExclusive"], "FF")
+        self.assertEqual(pkr["parents"], ["0"])
+
+    def test_pkrange_get_method(self):
+        """PKRange.get() works like dict.get() with defaults."""
+        pkr = PKRange(id="1", minInclusive="00", maxExclusive="FF", parents=None)
+        self.assertEqual(pkr.get("id"), "1")
+        self.assertIsNone(pkr.get("parents"))
+        self.assertEqual(pkr.get("nonexistent", "default"), "default")
+
+    def test_pkrange_contains(self):
+        """PKRange supports 'in' operator for field checking."""
+        pkr = PKRange(id="1", minInclusive="00", maxExclusive="FF", parents=["0"])
+        self.assertIn("id", pkr)
+        self.assertIn("parents", pkr)
+        self.assertNotIn("_rid", pkr)
+
+    def test_pkrange_getitem_keyerror(self):
+        """PKRange raises KeyError for invalid keys, like a dict."""
+        pkr = PKRange(id="1", minInclusive="00", maxExclusive="FF", parents=None)
+        with self.assertRaises(KeyError):
+            _ = pkr["_rid"]
+
+    def test_pkrange_in_collection_routing_map(self):
+        """CollectionRoutingMap works with PKRange namedtuples instead of dicts."""
+        pk_ranges = [
+            PKRange(id="0", minInclusive="", maxExclusive="05C1C9CD673398", parents=None),
+            PKRange(id="1", minInclusive="05C1C9CD673398", maxExclusive="05C1D9CD673398", parents=None),
+            PKRange(id="2", minInclusive="05C1D9CD673398", maxExclusive="FF", parents=None),
+        ]
+        crm = CollectionRoutingMap.CompleteRoutingMap(
+            [(r, True) for r in pk_ranges], "test-collection"
+        )
+        self.assertIsNotNone(crm)
+
+        # Test get_overlapping_ranges
+        full_range = Range("", "FF", True, False)
+        overlapping = crm.get_overlapping_ranges(full_range)
+        self.assertEqual(len(overlapping), 3)
+
+        # Test get_range_by_effective_partition_key
+        result = crm.get_range_by_effective_partition_key("")
+        self.assertEqual(result["id"], "0")
+
+        # Test get_range_by_partition_key_range_id
+        result = crm.get_range_by_partition_key_range_id("1")
+        self.assertEqual(result["id"], "1")
+
+    def test_discard_parent_ranges_returns_pkrange(self):
+        """_discard_parent_ranges returns PKRange namedtuples, filtering parents."""
+        raw_ranges = [
+            {"id": "0", "minInclusive": "", "maxExclusive": "FF", "parents": []},
+            {"id": "1", "minInclusive": "", "maxExclusive": "80", "parents": ["0"]},
+            {"id": "2", "minInclusive": "80", "maxExclusive": "FF", "parents": ["0"]},
+        ]
+        result = list(PartitionKeyRangeCache._discard_parent_ranges(raw_ranges))
+        # Parent "0" should be discarded, only children "1" and "2" remain
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0].id, "1")
+        self.assertEqual(result[1].id, "2")
+        # Verify they are PKRange instances
+        self.assertIsInstance(result[0], PKRange)
+        self.assertIsInstance(result[1], PKRange)
+
+    def test_range_has_slots(self):
+        """Range class uses __slots__ to minimize per-instance memory."""
+        r = Range("00", "FF", True, False)
+        self.assertFalse(hasattr(r, "__dict__"))
+        self.assertTrue(hasattr(Range, "__slots__"))
+        self.assertEqual(set(Range.__slots__), {"min", "max", "isMinInclusive", "isMaxInclusive"})
+
+    def test_range_memory_smaller_with_slots(self):
+        """Range with __slots__ should be smaller than a regular object."""
+        r = Range("00", "FF", True, False)
+        # __slots__ objects are typically 48-72 bytes, without __slots__ ~250+ bytes
+        self.assertLess(sys.getsizeof(r), 100)
+
+    def test_range_upper_skips_when_already_uppercase(self):
+        """Range.__init__ should not create new string objects for already-uppercase input."""
+        original_min = "05C1C9CD673398"
+        original_max = "05C1D9CD673398"
+        r = Range(original_min, original_max, True, False)
+        # If .upper() is skipped, the same string object is reused (is, not just ==)
+        self.assertIs(r.min, original_min)
+        self.assertIs(r.max, original_max)
+
+    def test_range_upper_applied_when_lowercase(self):
+        """Range.__init__ should uppercase lowercase input."""
+        r = Range("05c1c9cd", "05c1d9cd", True, False)
+        self.assertEqual(r.min, "05C1C9CD")
+        self.assertEqual(r.max, "05C1D9CD")
+
+
+if __name__ == "__main__":
+    unittest.main()