perf(streaming): memoize per-chunk metadata on the read hot path

src/litdata/streaming/config.py

@@ -108,6 +108,9 @@ def __init__(
 
         self._skip_chunk_indexes_deletion: list[int] | None = None
         self.zero_based_roi: list[tuple[int, int]] | None = None
+        # Memoizes ``__getitem__`` results per chunk_index (invariant once the config is loaded);
+        # avoids rebuilding the chunk path on every item read.
+        self._chunk_meta_cache: dict[int, tuple[str, int, int]] = {}
         self.filename_to_size_map: dict[str, int] = {}
         for cnk in _original_chunks:
             # since files downloaded while reading will be decompressed, we need to store the name without compression
@@ -287,7 +290,17 @@ def _get_chunk_index_from_index(self, index: int) -> tuple[int, int]:
         )
 
     def __getitem__(self, index: ChunkedIndex) -> tuple[str, int, int]:
-        """Find the associated chunk metadata."""
+        """Find the associated chunk metadata.
+
+        This is called once per item on the read hot path, but its result depends only on
+        ``index.chunk_index`` (the local path, the chunk's begin offset and its byte size are all
+        fixed once the config is loaded). The per-chunk tuple is therefore memoized to avoid
+        rebuilding the path (``os.path.join`` + decompression-suffix stripping) on every item.
+        """
+        cached = self._chunk_meta_cache.get(index.chunk_index)
+        if cached is not None:
+            return cached
+
         assert self._chunks is not None
         chunk = self._chunks[index.chunk_index]
 
@@ -300,7 +313,9 @@ def __getitem__(self, index: ChunkedIndex) -> tuple[str, int, int]:
 
         filesize_bytes = chunk["chunk_bytes"]
 
-        return local_chunkpath, begin, filesize_bytes
+        meta = (local_chunkpath, begin, filesize_bytes)
+        self._chunk_meta_cache[index.chunk_index] = meta
+        return meta
 
     def download_filepath(self, chunk_index: int) -> str:
         """The raw on-disk path that the chunk is downloaded to before any decompression."""

src/litdata/streaming/item_loader.py

@@ -69,6 +69,11 @@ def setup(
             serializer.setup(data_format)
             self._serializers[data_format] = serializer
 
+        # Precompute the per-leaf serializer list and the data spec so the per-item `deserialize`
+        # hot path avoids a dict lookup per leaf and a config lookup per item.
+        self._serializers_list = [self._serializers[data_format] for data_format in self._data_format]
+        self._data_spec = self._config["data_spec"]
+
     def force_download(self, chunk_index: int) -> None:
         if self._force_download_queue:
             self._force_download_queue.put(chunk_index)
@@ -316,12 +321,11 @@ def deserialize(self, raw_item_data: bytes) -> "PyTree":
         idx = self._shift_idx
         sizes = np.frombuffer(raw_item_data[:idx], np.uint32)
         data = []
-        for size, data_format in zip(sizes, self._data_format):
-            serializer = self._serializers[data_format]
+        for size, serializer in zip(sizes, self._serializers_list):
             data_bytes = raw_item_data[idx : idx + size]
             data.append(serializer.deserialize(data_bytes))
             idx += size
-        return tree_unflatten(data, self._config["data_spec"])
+        return tree_unflatten(data, self._data_spec)
 
     def close(self, chunk_index: int) -> None:
         """Close the open file handle."""