perf(data): node cache for component readers

docs/data/storage_and_access.rst

@@ -361,6 +361,14 @@ Performance Recommendations
   This is especially important for remote stores where each metadata
   lookup would otherwise be a separate round-trip.
 
+* **Keep the node cache enabled** (the default). Component readers cache the
+  groups / arrays they open (including their decoded metadata, but not their
+  data), so repeated accesses (e.g., of frames) skip re-opening them. The
+  ``node_cache_size`` parameter on
+  :class:`~ncore.data.v4.SequenceComponentGroupsReader` bounds the number of
+  cached nodes per component reader for very long sequences (``None`` caches
+  all accessed nodes, ``0`` disables caching).
+
 * **Tune the block size and cache type** for your workload.  The
   ``default_block_size`` and ``default_cache_type`` parameters on ``UPath``
   control how much data is fetched per HTTP request and how it is cached

ncore/impl/data/nodes.py

@@ -22,6 +22,9 @@
 
 from __future__ import annotations
 
+import collections
+import threading
+
 from typing import TYPE_CHECKING, Any, Iterator, List, Literal, Optional, Tuple, Type, TypeVar, Union, cast
 
 import numpy as np
@@ -48,6 +51,53 @@
 T = TypeVar("T")
 
 
+class NodeCache:
+    """Cache of the opened nodes of a hierarchy opened for reading, by node path.
+
+    Opening a node constructs its zarr-python object (decoding its metadata), which dominates repeated reads of small
+    arrays (e.g., of frames). Cached nodes reference their metadata only, never array data. ``max_entries=None`` caches
+    all opened nodes, ``max_entries=N`` the N most recently used nodes (evicted nodes are re-opened on access).
+
+    All methods are thread-safe.
+    """
+
+    def __init__(self, max_entries: Optional[int] = None) -> None:
+        if max_entries is not None and max_entries < 1:
+            raise ValueError(f"max_entries needs to be positive or None (unbounded), got {max_entries}")
+        self._max_entries = max_entries
+        self._nodes: collections.OrderedDict[str, ZarrNode] = collections.OrderedDict()
+        self._lock = threading.Lock()
+
+    @property
+    def max_entries(self) -> Optional[int]:
+        return self._max_entries
+
+    def __len__(self) -> int:
+        return len(self._nodes)
+
+    def keys(self) -> List[str]:
+        """The paths of the cached nodes, from least to most recently used"""
+        with self._lock:
+            return list(self._nodes)
+
+    def _open(self, group: ZarrGroup, path: str) -> ZarrNode:
+        key = _Node._join(group.path, path)
+        with self._lock:
+            if (node := self._nodes.get(key)) is not None:
+                self._nodes.move_to_end(key)
+                return node
+
+        # (concurrent first accesses open nodes redundantly, which is cheaper than serializing all accesses)
+        node = _backend.open_member(group, path)
+
+        with self._lock:
+            self._nodes[key] = node
+            # (at most one entry is added per call, so at most the least recently used node is evicted)
+            if self._max_entries is not None and len(self._nodes) > self._max_entries:
+                self._nodes.popitem(last=False)
+        return node
+
+
 class _Node:
     """Common functionality of groups and arrays"""
 
@@ -100,9 +150,10 @@ def _typed_attr(self, name: str, expected: Type[T]) -> T:
             raise TypeError(f"Attribute {name!r} of {self.path!r} is not of type {expected.__name__}: {value!r}")
         return cast(T, value)
 
-    def _member_path(self, path: str) -> str:
-        """The path of a (nested) member of the node within its store"""
-        return f"{self.path}/{path}" if self.path else path
+    @staticmethod
+    def _join(parent: str, path: str) -> str:
+        """The path of a (nested) member of a node within its store"""
+        return f"{parent}/{path}" if parent else path
 
     def __repr__(self) -> str:
         return f"{type(self).__name__}({self.path!r}, zarr_format={self.zarr_format})"
@@ -159,11 +210,15 @@ def read_bytes(self) -> bytes:
 
 
 class Group(_Node):
-    """A group of an NCore data store"""
+    """A group of an NCore data store.
 
-    def __init__(self, group: ZarrGroup, store: Optional[Store] = None) -> None:
+    Members of groups using a :class:`NodeCache` are opened via the cache (also members of member groups).
+    """
+
+    def __init__(self, group: ZarrGroup, cache: Optional[NodeCache] = None, store: Optional[Store] = None) -> None:
         super().__init__(group)
         self._group = group
+        self._cache = cache
         self._store = store
 
     @property
@@ -176,23 +231,37 @@ def store(self) -> Optional[Store]:
         """The store of the group (None for groups not opened / created via a :class:`~ncore.impl.data.stores.Store`)"""
         return self._store
 
+    @property
+    def cache(self) -> Optional[NodeCache]:
+        """The node cache used to open members (None if not cached)"""
+        return self._cache
+
+    def with_cache(self, cache: Optional[NodeCache]) -> Group:
+        """Returns the group, using the given node cache to open its members"""
+        return Group(self._group, cache, self._store)
+
     # -- reading ---------------------------------------------------------------
 
+    def _open(self, path: str) -> ZarrNode:
+        if self._cache is not None:
+            return self._cache._open(self._group, path)
+        return _backend.open_member(self._group, path)
+
     def node(self, path: str) -> Union[Group, Array]:
         """Returns a (nested) member group or array (raises KeyError if it doesn't exist)"""
-        node = _backend.open_member(self._group, path)
-        return Group(node, self._store) if isinstance(node, ZarrGroup) else Array(node)
+        node = self._open(path)
+        return Group(node, self._cache, self._store) if isinstance(node, ZarrGroup) else Array(node)
 
     def group(self, path: str) -> Group:
         """Returns a (nested) member group (raises KeyError if it doesn't exist or isn't a group)"""
-        if not isinstance(node := _backend.open_member(self._group, path), ZarrGroup):
-            raise KeyError(f"{self._member_path(path)} is not a group")
-        return Group(node, self._store)
+        if not isinstance(node := self._open(path), ZarrGroup):
+            raise KeyError(f"{_Node._join(self.path, path)} is not a group")
+        return Group(node, self._cache, self._store)
 
     def array(self, path: str) -> Array:
         """Returns a (nested) member array (raises KeyError if it doesn't exist or isn't an array)"""
-        if not isinstance(node := _backend.open_member(self._group, path), ZarrArray):
-            raise KeyError(f"{self._member_path(path)} is not an array")
+        if not isinstance(node := self._open(path), ZarrArray):
+            raise KeyError(f"{_Node._join(self.path, path)} is not an array")
         return Array(node)
 
     def members(self) -> List[str]:
@@ -206,7 +275,7 @@ def __contains__(self, name: str) -> bool:
     def groups(self) -> Iterator[Tuple[str, Group]]:
         """Iterates over the direct member groups (ordered by name)"""
         for name, group in sorted(_backend.member_groups(self._group), key=lambda member: member[0]):
-            yield name, Group(group, self._store)
+            yield name, Group(group, self._cache, self._store)
 
     # (zarr-python style read access, without the equality semantics of mappings: groups compare by identity)
 
@@ -242,11 +311,11 @@ def __len__(self) -> int:
 
     def create_group(self, name: str, *, attributes: Optional[Attributes] = None) -> Group:
         """Creates a member group (raises if a member with the name exists already)"""
-        return Group(_backend.create_group(self._group, name, attributes), self._store)
+        return Group(_backend.create_group(self._group, name, attributes), store=self._store)
 
     def require_group(self, name: str) -> Group:
         """Returns a member group, creating it if it doesn't exist yet"""
-        return Group(_backend.require_group(self._group, name), self._store)
+        return Group(_backend.require_group(self._group, name), store=self._store)
 
     def create_array(
         self,

ncore/impl/data/stores.py

@@ -54,12 +54,12 @@ def zarr(self) -> ZarrStore:
 
     def create_root_group(self, attributes: Attributes) -> Group:
         """Creates the root group of the hierarchy of the store"""
-        return Group(_backend.create_root_group(self._zarr_store, 2, attributes), self)
+        return Group(_backend.create_root_group(self._zarr_store, 2, attributes), store=self)
 
     def open_root_group(self, consolidated: bool = True) -> Group:
         """Opens the root group of the hierarchy of the store for reading, via its compressed consolidated metadata if
         consolidated (raises KeyError if the store has no compressed consolidated metadata)"""
-        return Group(_backend.open_root_group(self._zarr_store, consolidated), self)
+        return Group(_backend.open_root_group(self._zarr_store, consolidated), store=self)
 
     def consolidate(self) -> None:
         """Stores the compressed consolidated metadata of the hierarchy of the store (after writing all nodes)"""

ncore/impl/data/stores_test.py

@@ -546,6 +546,78 @@ def test_names(self):
         self.assertEqual(self.root["a/b/c"].name, "/a/b/c")
 
 
+class TestNodeCache(unittest.TestCase):
+    """Node caches of groups (disabled / unbounded / least recently used)"""
+
+    def setUp(self):
+        self._tmp = tempfile.TemporaryDirectory()
+        path = os.path.join(self._tmp.name, "store.itar")
+        with IndexedTarZarrStore(path, mode="w") as store:
+            sub = store.create_root_group({}).create_group("g")
+            for i in range(8):
+                sub.create_array(f"a{i}", np.full(3, i))
+            store.consolidate()
+        self.root = IndexedTarZarrStore(path).open_root_group()
+
+    def tearDown(self):
+        self._tmp.cleanup()
+
+    def test_uncached(self):
+        self.assertIsNone(self.root.cache)
+        self.assertIsNot(self.root.array("g/a1").zarr, self.root.array("g/a1").zarr)
+
+    def test_unbounded(self):
+        root = self.root.with_cache(cache := nodes.NodeCache())
+        for i in range(8):
+            np.testing.assert_array_equal(root.array(f"g/a{i}").read(), [i] * 3)
+        self.assertEqual(len(cache), 8)
+        self.assertIs(root.array("g/a3").zarr, root.array("g/a3").zarr)
+
+    def test_member_groups_share_the_cache(self):
+        root = self.root.with_cache(cache := nodes.NodeCache())
+        group = root.group("g")
+        self.assertIs(group.cache, cache)
+        node = root.node("g")
+        assert isinstance(node, nodes.Group)
+        self.assertIs(node.cache, cache)
+        self.assertIs(cast(nodes.Array, root.node("g/a2")).zarr, root.array("g/a2").zarr)
+        self.assertIs(group.array("a3").zarr, root.array("g/a3").zarr)
+        self.assertIs(next(root.groups())[1].cache, cache)
+        self.assertEqual(cache.keys(), ["g", "g/a2", "g/a3"])
+
+    def test_zarr_python_style_access_uses_the_cache(self):
+        root = self.root.with_cache(cache := nodes.NodeCache())
+        assert isinstance(group := root["g"], nodes.Group)
+        self.assertIs(group.cache, cache)
+        self.assertIs(cast(nodes.Array, root["g/a1"]).zarr, cast(nodes.Array, group.get("a1")).zarr)
+        self.assertIsNone(group.get("missing"))
+        members = dict(group.items())
+        self.assertIs(cast(nodes.Array, members["a2"]).zarr, root.array("g/a2").zarr)
+        self.assertEqual(sorted(cache.keys()), ["g"] + [f"g/a{i}" for i in range(8)])
+
+    def test_least_recently_used(self):
+        root = self.root.with_cache(cache := nodes.NodeCache(max_entries=3))
+        a0 = root.array("g/a0").zarr
+        root.array("g/a1")
+        root.array("g/a2")
+        self.assertIs(root.array("g/a0").zarr, a0)  # (a1 is the least recently used node now)
+        root.array("g/a3")
+        self.assertEqual(cache.keys(), ["g/a2", "g/a0", "g/a3"])
+        for i in range(8):  # evicted nodes are re-opened
+            np.testing.assert_array_equal(root.array(f"g/a{i}").read(), [i] * 3)
+        self.assertEqual(len(cache), 3)
+
+    def test_missing_nodes_are_not_cached(self):
+        root = self.root.with_cache(cache := nodes.NodeCache())
+        with self.assertRaises(KeyError):
+            root.array("g/missing")
+        self.assertEqual(len(cache), 0)
+
+    def test_invalid_sizes(self):
+        with self.assertRaises(ValueError):
+            nodes.NodeCache(max_entries=0)
+
+
 class TestZeroCopyReads(unittest.TestCase):
     """Read results are read-only, arrays stored as a single uncompressed chunk are read as views of their records"""
 

ncore/impl/data/v4/components.py

@@ -346,6 +346,7 @@ def __init__(
         open_consolidated: bool = True,
         itar_index_tail_read_size: int = 1 << 20,  # 1 MiB default
         max_threads: int | None = None,
+        node_cache_size: int | None = None,
     ):
         """Initialize a SequenceComponentReader for a virtual sequence represented by a list of components.
 
@@ -361,21 +362,32 @@ def __init__(
                                single I/O call to initialize the tar file index to minimize reads.
                                Default 1 MiB size fits typical use cases, but may be increased if the compressed
                                index size is expected to be larger or decreased if remote access chunk sizes are smaller.
-            max_threads:       The maximum number of threads used to load the different components (if None,
-                               use interpreter-default number of threads for a ThreadPoolExecutor)
+            max_threads:       The maximum number of threads used to open the component stores concurrently (if
+                               None, use interpreter-default number of threads for a ThreadPoolExecutor)
+            node_cache_size:   Number of opened groups / arrays (with their decoded meta-data, but not their data)
+                               each component reader caches for repeated accesses (e.g., of frames). None caches
+                               all accessed nodes for the lifetime of the reader, a positive number the most recently
+                               used nodes only, and 0 disables caching.
         """
 
+        if node_cache_size is not None and node_cache_size < 0:
+            raise ValueError(f"node_cache_size needs to be None or non-negative, got {node_cache_size}")
+        self._node_cache_size = node_cache_size
+
         component_group_upaths: List[UPath] = self.expand_component_group_paths(component_group_paths)
 
         assert len(component_group_upaths), "No component inputs provided"
 
         # Load component stores concurrently (to hide latency) and check for sequence consistency
         self._component_stores: Dict[str, Tuple[nodes.Group, UPath]] = {}  # use str as the generic path / URL type
 
+        # Component stores are opened concurrently (hiding the latency of remote archive index reads), their
+        # hierarchies are opened from their consolidated metadata sequentially on the calling thread (pure-Python CPU
+        # work that threads only slow down), each as soon as its store is opened
         with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
 
-            def thread_load_component_store(component_store_upath: UPath) -> Tuple[nodes.Group, UPath]:
-                """Thread-executed shard opening"""
+            def thread_open_component_store(component_store_upath: UPath) -> Tuple[stores.Store, UPath]:
+                """Thread-executed component store opening"""
 
                 # Make sure paths are absolute at this point
                 component_store_upath = component_store_upath.absolute()
@@ -396,16 +408,14 @@ def thread_load_component_store(component_store_upath: UPath) -> Tuple[nodes.Gro
                 else:
                     component_store = stores.ZarrStoreAdapter.directory(component_store_upath)
 
-                return component_store.open_root_group(open_consolidated), component_store_upath
+                return component_store, component_store_upath
 
-            for future in concurrent.futures.as_completed(
-                [
-                    executor.submit(thread_load_component_store, component_store_upath)
-                    for component_store_upath in component_group_upaths
-                ]
+            # (hierarchies are opened as soon as their stores are opened, in order, while the remaining stores
+            # are still being opened)
+            for component_store, component_store_path in executor.map(
+                thread_open_component_store, component_group_upaths
             ):
-                # Note: thread completion order is not relevant here
-                component_root, component_store_path = future.result()
+                component_root = component_store.open_root_group(open_consolidated)
 
                 # Check sequence compatibility
                 component_root_version = component_root.attr_str("version")
@@ -491,7 +501,11 @@ def open_component_readers(
                 if not component_reader_type.supports_component_version(component_group.attr_str("component_version")):
                     continue
 
-                ret[component_instance_name] = component_reader_type(component_instance_name, component_group)
+                # each reader caches the nodes of its component
+                cache = None if self._node_cache_size == 0 else nodes.NodeCache(self._node_cache_size)
+                ret[component_instance_name] = component_reader_type(
+                    component_instance_name, component_group.with_cache(cache)
+                )
 
         return ret