Add keyed sample lookup and in-place dataset_update

[tchaton]

Summary

• Add optimize(..., key_fn=...) support that writes keys/shard-*.parquet and records layout in index.json["keys"].
• Support entity-key reads via StreamingDataset["entity-id"] (int indexes stay sample indexes).
• Add Session-style dataset_update: stage replaces, then commit() to rewrite only touched chunks; exit without commit discards.
• Include build_keys_index to backfill keys for datasets optimized without key_fn.

This is aimed at surgical value patches under a stable schema (fix/relabel specific entity IDs), not schema-wide table redesigns that touch every joined row.

API design

Layout

dataset/
  index.json          # includes "keys": { version, num_shards, sharding }
  keys/
    shard-00000.parquet
    shard-XXXXX.parquet   # optional; hash(key) % num_shards
  chunk-*.bin

Shard parquet columns (sorted by key): key, index, chunk_index, chunk_offset.

Key-store metadata lives in index.json["keys"] — no separate manifest file.
Default is one shard (num_shards=1). Lookups use lazy Polars scan_parquet + predicate pushdown (remote-friendly).

Write keys at optimize time

from litdata import optimize, StreamingDataset, dataset_update, build_keys_index

optimize(
    fn=...,
    inputs=...,
    output_dir=out,
    key_fn=lambda s: s["id"],  # str (preferred) or int entity id
)

Read by entity key

ds = StreamingDataset(out)
sample = ds["play-123"]   # entity key
# ds[3] is still a global sample index, not an entity key
# integer entity ids: ds.get_by_key(entity_id)

In-place update (Session-style)

with dataset_update(out) as update:
    update["play-123"] = new_sample
    # edit any other samples ...
    update.commit()               # required; exiting without commit discards
    # update[...] raises after commit

Backfill keys for an existing dataset

build_keys_index(out, key_fn=lambda s: s["id"])

Test plan

• pytest tests/streaming/test_dataset_update.py (18 passed)
• End-to-end on /teamspace/lightning_storage/testing/keys_demo: optimize + keyed read + commit + discard-without-commit

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 73.17961% with 221 lines in your changes missing coverage. Please review.
✅ Project coverage is 81%. Comparing base (cf72e73) to head (64d4d08).
⚠️ Report is 1 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    #868    +/-   ##
=====================================
- Coverage    81%     81%    -1%     
=====================================
  Files        56      58     +2     
  Lines      9303   10124   +821     
=====================================
+ Hits       7563    8166   +603     
- Misses     1740    1958   +218     

🚀 New features to boost your workflow:

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