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

[tchaton]

Summary

Small, safe read-path speedup found by profiling the streaming iteration loop. Two per-item costs turned out to be invariant per chunk and are now memoized. No behavioural change.

How it was found

Profiled a pure local read (no network, num_workers=0) with cProfile over StreamingDataset + StreamingDataLoader iteration. The hottest per-item entries that weren't intrinsic work (deserialize / file IO / pytree) were:

• posixpath.join — 40k calls (one per item) from ChunksConfig.__getitem__, even though the chunk path depends only on chunk_index.
• serializer dict lookups per leaf and a config["data_spec"] lookup per item in PyTreeLoader.deserialize.

Changes

• ChunksConfig.__getitem__ memoizes its (local_path, begin, filesize) tuple per chunk_index. These are all fixed once the config is loaded (_cache_dir, _chunks, _intervals, _compressor never change), so this is a pure lookup cache — no invalidation needed.
• BaseItemLoader.setup precomputes the per-leaf serializer list and the data spec; deserialize uses them instead of re-looking-up a dict per leaf and the config per item.

Benchmark

Local dataset, 40k items ({"x": float32[128], "y": int}), chunk_size=2000, num_workers=0, median of 5 runs:

	items/s
before	~104,800
after	~110,600
	~+6%

The gain is per-item decode CPU, so it helps most when decode-bound (num_workers=0, or many workers saturating cores); it won't move an I/O- or GPU-bound run. Reproduced repeatedly; the posixpath.join per-item calls disappear from the profile after the change.

Follow-ups (not in this PR)

The remaining per-item hot spots are the vendored pytree unflatten (~16%, deliberately left alone — it mirrors torch's pytree) and the 2 seek+2 read per item in _load_data. A future change could cache each chunk's offset table on open to halve the per-item file syscalls; it touches the core read path so it deserves its own reviewed PR.

Verification

No regressions: test_reader.py, test_config.py, test_item_loader.py, test_serializer.py (48 passed). ruff + mypy clean.

🤖 Generated with Claude Code

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 87.50000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81%. Comparing base (ac1b064) to head (c9a0f6e).
⚠️ Report is 2 commits behind head on main.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@         Coverage Diff         @@
##           main   #848   +/-   ##
===================================
  Coverage    81%    81%           
===================================
  Files        54     54           
  Lines      7683   7736   +53     
===================================
+ Hits       6210   6262   +52     
- Misses     1473   1474    +1     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.