fix(streaming): do not auto-window JPEG and other per-item media

src/litdata/CHANGELOG.md

@@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [unreleased] - YYYY-MM-DD
 
+### Fixed
+
+- Auto ``batch_decode`` is 1 for JPEG / image / audio (and other per-item media). A 16-row window plus ``item_shuffle_window=256`` was re-decoding each ImageNet JPEG ~16×. Cheap leaves (text, ints, small tensors) still batch. ([#897](https://github.com/Lightning-AI/litData/pull/897) follow-up)
+
 ## [0.2.74] - 2026-09-01
 
 ### Added

src/litdata/streaming/dataset.py

@@ -117,11 +117,11 @@ def __init__(
                 IDs from that first-run shuffler assignment, not a different bucket layout.
             batch_decode: How many items to deserialize together after a chunk is local.
                 ``"auto"`` (default) picks from the data format and mean sample size
-                (256 for text/nested, down to 1 for multi-MB images/video).
-                ``0`` is per item, ``N`` is an aligned window, ``"all"`` is the whole chunk.
-                Shuffle permutes items inside the same window so training still hits the
-                cache. ``LITDATA_BATCH_DECODE`` / ``LITDATA_BATCH_ROWS`` apply only when
-                this is ``"auto"``.
+                (256 for text/nested; 1 for JPEG/image/audio so each row is decoded
+                once). ``0`` is per item, ``N`` is an aligned window, ``"all"`` is the
+                whole chunk. Shuffle permutes items inside the same window so cheap
+                leaves still hit the cache. ``LITDATA_BATCH_DECODE`` /
+                ``LITDATA_BATCH_ROWS`` apply only when this is ``"auto"``.
             item_shuffle_window: In-chunk shuffle block size (pairs with ``batch_decode``).
                 ``None`` / ``"auto"`` (default) is 256, or ``LITDATA_ITEM_SHUFFLE_WINDOW``.
                 ``0`` / ``"full"`` is a full in-chunk permutation. Blocks are shuffled,

src/litdata/streaming/item_loader.py

@@ -63,7 +63,9 @@
 # Cap for cheap leaves (text / nested JSON). Images/video scale down from avg bytes.
 _DEFAULT_BATCH_ROWS = 256
 _AUTO_WINDOW_BYTES = 16 << 20  # ~16MB of on-disk samples per decode window
-_HEAVY_LEAF = frozenset(
+# C++ / per-row deserialize is not amortizable. A window smaller than
+# ``item_shuffle_window`` (default 256) re-decodes the same JPEG ~16×.
+_PER_ITEM_LEAF = frozenset(
     {
         "jpeg",
         "pil",
@@ -76,6 +78,11 @@
         "nifti",
         "tiff",
         "jpeg_array",
+    }
+)
+_HEAVY_LEAF = frozenset(
+    {
+        *_PER_ITEM_LEAF,
         "tensor",
         "no_header_tensor",
         "graph",
@@ -129,14 +136,22 @@ def _leaf_key(name: str) -> str:
     return name.split(":", 1)[0].lower()
 
 
+def _item_only_batch(batch_rows: int | None) -> bool:
+    """``0`` / ``1`` decode the requested row; ``-1`` / ``N>1`` keep a window."""
+    return batch_rows is not None and 0 <= batch_rows <= 1
+
+
 def _auto_batch_rows(data_format: list[str] | None, chunks: list | None = None) -> int:
     """Pick a window from leaf types and mean on-disk sample size.
 
-    Text/nested stay at 256 (measured winner). A 2MB JPEG would turn 256 into
-    a multi-GB Python spike, so heavy leaves scale toward 1.
+    Text/nested stay at 256 (measured winner). JPEG / image / audio always use
+    1 (decode only the requested row) so shuffle cannot re-decode a dropped
+    window. Large tensors still scale toward 1 from mean on-disk bytes.
     """
     avg = _avg_sample_bytes(chunks)
     keys = [_leaf_key(name) for name in (data_format or [])]
+    if any(key in _PER_ITEM_LEAF for key in keys):
+        return 1
     heavy = any(key in _HEAVY_LEAF for key in keys)
     if avg >= 1 << 20:
         return 1
@@ -695,6 +710,7 @@ def __init__(self, batch_decode: Any = "auto") -> None:
         self._framed_meta: dict[int, FramedHeader] = {}
         self._framed_decompressor: Any | None = None
         self._framed_inflate_buf: bytes | memoryview | None = None
+        self._framed_inflate_key: tuple[int, int] | None = None
         self._compression_level = "chunk"
         self._sample_compression = False
 
@@ -891,10 +907,19 @@ def _fill_framed_window(
             and len(rows) == n_items
         ):
             return rows[table_idx - first]
-        raw = inflate_frame(view, header, frame_i, self._framed_compressor())
-        self._framed_inflate_buf = raw
+        inflate_key = (chunk_index, frame_i)
+        if self._framed_inflate_key == inflate_key and self._framed_inflate_buf is not None:
+            raw = self._framed_inflate_buf
+        else:
+            raw = inflate_frame(view, header, frame_i, self._framed_compressor())
+            self._framed_inflate_buf = raw
+            self._framed_inflate_key = inflate_key
         base = header.offsets[first]
         local_offsets = [int(off) - base for off in header.offsets[first : first + n_items + 1]]
+        if _item_only_batch(self._batch_rows):
+            local_i = table_idx - first
+            decoded = self._batch_deserialize_payload(raw, local_offsets, chunk_index, local_i, local_i + 1)
+            return decoded[0]
         decoded = self._batch_deserialize_payload(raw, local_offsets, chunk_index, 0, n_items)
         return self._store_decode_window(chunk_index, first, decoded, table_idx)
 
@@ -927,10 +952,10 @@ def _load_batched_item(self, chunk_index: int, chunk_filepath: str, table_idx: i
             arrow = self._try_arrow_footer_rows(view, chunk_index, table_idx)
             if arrow is not _BATCH_SKIP:
                 return arrow
-            if not batch_rows:
+            if _item_only_batch(batch_rows):
                 return _BATCH_SKIP
             return self._fill_decode_window_mmap(chunk_index, table_idx, batch_rows)
-        if not batch_rows:
+        if _item_only_batch(batch_rows):
             with open(chunk_filepath, "rb") as handle:
                 blob = handle.read()
             header = self._resolve_framed_header(chunk_index, blob)
@@ -1418,6 +1443,7 @@ def __getstate__(self) -> dict[str, Any]:
         state["_framed_meta"] = {}
         state["_framed_decompressor"] = None
         state["_framed_inflate_buf"] = None
+        state["_framed_inflate_key"] = None
         # Compiled unflatten closures aren't picklable; rebuild after unpickle.
         state["_unflatten"] = None
         state["_sizes_struct"] = None
@@ -1453,6 +1479,8 @@ def __setstate__(self, state: dict[str, Any]) -> None:
             self._framed_meta = {}
             self._framed_decompressor = None
             self._framed_inflate_buf = None
+        if not hasattr(self, "_framed_inflate_key"):
+            self._framed_inflate_key = None
         data_spec = getattr(self, "_data_spec", None)
         if isinstance(data_spec, TreeSpec):
             self._unflatten = _compile_treespec_unflatten(data_spec)

tests/streaming/test_item_loader.py

@@ -23,6 +23,7 @@
     TokensLoader,
     _auto_batch_rows,
     _batch_rows_for_format,
+    _item_only_batch,
     _parse_batch_decode,
 )
 from litdata.streaming.sampler import ChunkedIndex
@@ -39,8 +40,20 @@ def test_batch_rows_for_format(monkeypatch):
     assert _batch_rows_for_format(["str", "int"]) == 256
     assert _batch_rows_for_format(["str", "str", "json", "json"]) == 256
     assert _batch_rows_for_format(["pickle"]) == 256
+    imagenet_chunks = [{"chunk_bytes": 63968405, "chunk_size": 556}]
+    jpeg_auto = _auto_batch_rows(["jpeg", "int"], imagenet_chunks)
+    assert jpeg_auto == 1
+    assert _item_only_batch(jpeg_auto)
+    jpeg_115k = _auto_batch_rows(["jpeg"], [{"chunk_bytes": 512 << 10, "chunk_size": 8}])
+    assert jpeg_115k == 1
+    assert _item_only_batch(jpeg_115k)
     assert _auto_batch_rows(["jpeg", "int"], [{"chunk_bytes": 8 << 20, "chunk_size": 4}]) == 1
-    assert _auto_batch_rows(["jpeg"], [{"chunk_bytes": 512 << 10, "chunk_size": 8}]) == 16
+    assert _auto_batch_rows(["image"], None) == 1
+    assert _auto_batch_rows(["audio"], None) == 1
+    tensor_auto = _auto_batch_rows(["tensor"], [{"chunk_bytes": 512 << 10, "chunk_size": 8}])
+    assert tensor_auto == 16
+    assert not _item_only_batch(tensor_auto)
+    assert not _item_only_batch(_auto_batch_rows(["str", "int"]))
     assert _batch_rows_for_format(["jpeg"], [{"chunk_bytes": 8 << 20, "chunk_size": 4}]) == 1
     assert _batch_rows_for_format(["str"], batch_decode=32) == 32
     assert _batch_rows_for_format(["jpeg"], [{"chunk_bytes": 8 << 20, "chunk_size": 4}], batch_decode=8) == 8
@@ -66,6 +79,49 @@ def test_decode_window_is_aligned():
     assert loader._window_bounds(3, 10, 1) == (3, 4)
 
 
+def test_decode_window_is_single_slot():
+    """JPEG must not keep extra decoded windows around for shuffle."""
+    loader = PyTreeLoader()
+    assert loader._store_decode_window(0, 0, list(range(16)), 3) == 3
+    assert loader._store_decode_window(0, 16, list(range(16, 32)), 20) == 20
+    assert loader._chunk_rows == list(range(16, 32))
+    assert loader._win_start == 16
+    assert loader._chunk_rows_index == 0
+
+
+def test_shuffled_jpeg_auto_batch_decodes_once(tmp_path, monkeypatch):
+    """Shuffle + auto JPEG must decode_jpeg once per item, not a 16-row window (~15×)."""
+    pytest.importorskip("PIL")
+
+    from litdata import optimize
+    from litdata.streaming import serializers
+
+    n = 64
+    calls = {"n": 0}
+    orig = serializers.JPEGSerializer.deserialize
+
+    def counting(self, data: bytes) -> torch.Tensor:
+        calls["n"] += 1
+        return orig(self, data)
+
+    monkeypatch.setattr(serializers.JPEGSerializer, "deserialize", counting)
+
+    optimize(
+        fn=_jpeg_label_sample,
+        inputs=list(range(n)),
+        output_dir=str(tmp_path / "jpeg-ds"),
+        chunk_size=n,
+        num_workers=1,
+    )
+    ds = StreamingDataset(str(tmp_path / "jpeg-ds"), shuffle=True, seed=42, item_shuffle_window=256)
+    items = list(ds)
+    assert len(items) == n
+    assert calls["n"] == n
+    loader = ds.cache._reader._item_loader
+    assert loader._batch_rows == 1
+    assert _item_only_batch(loader._batch_rows)
+
+
 def test_streaming_dataset_exposes_batch_decode(tmp_path):
     from litdata import optimize
 
@@ -580,6 +636,19 @@ def _nested_arrow_sample(i: int):
     }
 
 
+def _jpeg_label_sample(i: int):
+    import io
+
+    from PIL import Image as PILImage
+
+    from litdata.types import Jpeg
+
+    buf = io.BytesIO()
+    array = np.full((16, 16, 3), i % 256, dtype=np.uint8)
+    PILImage.fromarray(array).save(buf, format="JPEG", quality=80)
+    return Jpeg(bytes=buf.getvalue()), int(i)
+
+
 def _flat_arrow_sample(i: int):
     return {"text": f"row {i}", "label": i % 2}