Add keyed sample lookup and in-place dataset_update

CONTRIBUTING.md

@@ -2,7 +2,7 @@
 
 Welcome to the PyTorch Lightning community! We're building the most advanced research platform on the planet to implement the latest, best practices and integrations that the amazing PyTorch team and other research organization rolls out!
 
-If you are new to open source, check out [this blog to get started with your first Open Source contribution](https://devblog.pytorchlightning.ai/quick-contribution-guide-86d977171b3a).
+If you are new to open source, check out [GitHub's guide to making your first contribution](https://docs.github.com/en/get-started/quickstart/contributing-to-projects).
 
 ## Main Core Value: One less thing to remember
 

src/litdata/__init__.py

@@ -19,11 +19,13 @@
 from litdata.streaming.combined import CombinedStreamingDataset
 from litdata.streaming.dataloader import StreamingDataLoader
 from litdata.streaming.dataset import StreamingDataset
+from litdata.streaming.dataset_update import dataset_update
 from litdata.streaming.item_loader import TokensLoader
 from litdata.streaming.parallel import ParallelStreamingDataset
 from litdata.streaming.writer import index_parquet_dataset
 from litdata.utilities.breakpoint import breakpoint
 from litdata.utilities.hf_dataset import index_hf_dataset
+from litdata.utilities.keys_index import build_keys_index
 from litdata.utilities.train_test_split import train_test_split
 
 warnings.filterwarnings(
@@ -41,6 +43,8 @@
     "ParallelStreamingDataset",
     "map",
     "optimize",
+    "dataset_update",
+    "build_keys_index",
     "walk",
     "train_test_split",
     "merge_datasets",

src/litdata/constants.py

@@ -20,6 +20,12 @@
 from lightning_utilities.core.imports import RequirementCache
 
 _INDEX_FILENAME = "index.json"
+_KEYS_DIRNAME = "keys"
+_KEYS_SHARD_TEMPLATE = "shard-{:05d}.parquet"
+# Legacy single-file sidecar (still read if present).
+_KEYS_FILENAME = "keys.parquet"
+_RANK_KEYS_SUFFIX = ".keys.parquet"
+_DEFAULT_KEYS_NUM_SHARDS = 1
 _DEFAULT_CHUNK_BYTES = 1 << 26  # 64M B
 _DEFAULT_FAST_DEV_RUN_ITEMS = 10
 _DEFAULT_CACHE_DIR = os.path.join(Path.home(), ".lightning", "chunks")

src/litdata/processing/data_processor.py

@@ -25,6 +25,7 @@
 import traceback
 import warnings
 from abc import abstractmethod
+from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass
 from multiprocessing import Process, Queue
@@ -566,6 +567,8 @@ def __init__(
         self.checkpoint_next_index: int | None = checkpoint_next_index
         self.storage_options = storage_options
         self.using_queue_optimize = using_queue_optimize
+        # Explicit (writer_sample_index, key) pairs — index is the same value passed to Cache._add_item.
+        self._key_pairs: list[tuple[int, Any]] = []
 
     def run(self) -> None:
         try:
@@ -868,14 +871,25 @@ def _handle_data_chunk_recipe(self, index: int, item: Any) -> None:
                 return
 
             item_data_or_generator = self.data_recipe.prepare_item(current_item)
+            key_fn = getattr(self.data_recipe, "key_fn", None)
             if self.data_recipe.is_generator:
                 for item_data in item_data_or_generator:
                     if item_data is not None:
-                        chunk_filepath = self.cache._add_item(self._index_counter, item_data)
+                        sample_index = self._index_counter
+                        if key_fn is not None:
+                            from litdata.utilities.keys_index import normalize_key
+
+                            self._key_pairs.append((sample_index, normalize_key(key_fn(item_data))))
+                        chunk_filepath = self.cache._add_item(sample_index, item_data)
                         self._try_upload(chunk_filepath)
                         self._index_counter += 1
             elif item_data_or_generator is not None:
-                chunk_filepath = self.cache._add_item(self._index_counter, item_data_or_generator)
+                sample_index = self._index_counter
+                if key_fn is not None:
+                    from litdata.utilities.keys_index import normalize_key
+
+                    self._key_pairs.append((sample_index, normalize_key(key_fn(item_data_or_generator))))
+                chunk_filepath = self.cache._add_item(sample_index, item_data_or_generator)
                 self._try_upload(chunk_filepath)
                 self._index_counter += 1
                 if self.use_checkpoint:
@@ -896,6 +910,15 @@ def _handle_data_chunk_recipe_end(self) -> None:
                 if isinstance(chunk_filepath, str) and os.path.exists(chunk_filepath):
                     self.to_upload_queues[i % self.num_uploaders].put(chunk_filepath)
 
+        if getattr(self.data_recipe, "key_fn", None) is not None:
+            from litdata.utilities.keys_index import save_rank_keys
+
+            # Keep rank key files in the cache until `_merge_and_upload_keys` runs.
+            # Do not `_try_upload` them — upload removes the local file and merge would miss them.
+            rank = _get_node_rank() * self.num_workers + self.worker_index
+            keys_filepath = os.path.join(self.cache_chunks_dir, f"{rank}.keys.parquet")
+            save_rank_keys(keys_filepath, self._key_pairs)
+
         if self.use_checkpoint and not self.data_recipe.is_generator:
             checkpoint_filepath = self.cache.save_checkpoint()
             self._try_upload(checkpoint_filepath)
@@ -986,6 +1009,7 @@ def __init__(
         compression: str | None = None,
         encryption: Encryption | None = None,
         storage_options: dict[str, Any] = {},
+        key_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(storage_options)
         if chunk_size is not None and chunk_bytes is not None:
@@ -995,6 +1019,7 @@ def __init__(
         self.chunk_bytes = 1 << 26 if chunk_size is None and chunk_bytes is None else chunk_bytes  # 1<<26 = 64 MB
         self.compression = compression
         self.encryption = encryption
+        self.key_fn = key_fn
 
     @abstractmethod
     def prepare_structure(self, input_dir: str | None) -> list[T]:
@@ -1020,6 +1045,7 @@ def _done(self, size: int | None, delete_cached_files: bool, output_dir: Dir) ->
         node_rank = _get_node_rank()
         merge_cache._merge_no_wait(node_rank if num_nodes > 1 else None, getattr(self, "existing_index", None))
 
+        self._merge_and_upload_keys(output_dir, cache_dir, num_nodes, node_rank)
         self._upload_index(output_dir, cache_dir, num_nodes, node_rank)
 
         if num_nodes == node_rank + 1:
@@ -1050,6 +1076,124 @@ def _done(self, size: int | None, delete_cached_files: bool, output_dir: Dir) ->
             size=size,
         )
 
+    def _merge_and_upload_keys(self, output_dir: Dir, cache_dir: str, num_nodes: int, node_rank: int | None) -> None:
+        """Merge per-rank key sidecars and publish ``keys/shard-*.parquet`` next to ``index.json``."""
+        if getattr(self, "key_fn", None) is None:
+            return
+
+        from litdata.constants import _INDEX_FILENAME
+        from litdata.utilities.keys_index import (
+            concatenate_key_files,
+            enrich_keys_with_chunks,
+            has_keys_index,
+            list_key_parquet_files,
+            merge_rank_key_files,
+        )
+
+        def _enrich_if_index(dataset_dir: str) -> None:
+            index_path = os.path.join(cache_dir, _INDEX_FILENAME)
+            if os.path.isfile(index_path):
+                # Close before enrich: enrich rewrites index.json, and Windows cannot
+                # replace a path that still has an open handle.
+                with open(index_path, encoding="utf-8") as f:
+                    index_json = json.load(f)
+                enrich_keys_with_chunks(dataset_dir, index_json)
+
+        # Single-node (or per-node partial): merge rank files present in this cache.
+        if num_nodes <= 1:
+            merged = merge_rank_key_files(cache_dir)
+            if merged is None:
+                return
+            existing = getattr(self, "existing_index", None)
+            # Append mode: prepend keys from the previous dataset if present on output_dir.
+            if existing is not None and output_dir.path and has_keys_index(output_dir.path):
+                concatenate_key_files(
+                    list_key_parquet_files(output_dir.path) + list_key_parquet_files(cache_dir),
+                    cache_dir,
+                )
+            _enrich_if_index(cache_dir)
+            self._upload_keys_store(output_dir, cache_dir)
+            return
+
+        # Multi-node: each node merges local rank keys into ``{node_rank}-keys.parquet``,
+        # then the last node concatenates into the final ``keys/`` store.
+        assert node_rank is not None
+        node_keys_name = f"{node_rank}-keys.parquet"
+        merged = merge_rank_key_files(cache_dir, output_filename=node_keys_name)
+        if merged is None:
+            return
+        self._upload_file(output_dir, cache_dir, node_keys_name)
+
+        if num_nodes != node_rank + 1:
+            return
+
+        obj = parse.urlparse(output_dir.url if output_dir.url else output_dir.path)
+        local_paths: list[str] = []
+        for nr in range(num_nodes):
+            name = f"{nr}-keys.parquet"
+            local_path = os.path.join(cache_dir, name)
+            if nr != node_rank:
+                remote_base = output_dir.url if output_dir.url else output_dir.path
+                assert remote_base
+                remote_filepath = os.path.join(remote_base, name)
+                if obj.scheme in _SUPPORTED_PROVIDERS:
+                    merged_storage_options = construct_storage_options(self.storage_options, output_dir)
+                    _wait_for_file_to_exist(remote_filepath, storage_options=merged_storage_options)
+                    fs_provider = _get_fs_provider(remote_filepath, merged_storage_options)
+                    fs_provider.download_file(remote_filepath, local_path)
+                elif output_dir.path and os.path.isdir(output_dir.path):
+                    shutil.copyfile(remote_filepath, local_path)
+            local_paths.append(local_path)
+
+        concatenate_key_files(local_paths, cache_dir)
+        _enrich_if_index(cache_dir)
+        self._upload_keys_store(output_dir, cache_dir)
+
+    def _upload_keys_store(self, output_dir: Dir, cache_dir: str) -> None:
+        """Upload ``keys/shard-*.parquet`` next to the dataset (layout lives in ``index.json``)."""
+        from litdata.constants import _KEYS_DIRNAME
+        from litdata.utilities.keys_index import keys_dir, list_shard_files
+
+        if output_dir.path is None and output_dir.url is None:
+            return
+        local_keys = keys_dir(cache_dir)
+        if not os.path.isdir(local_keys):
+            return
+
+        rel_files = [os.path.basename(p) for p in list_shard_files(local_keys)]
+        obj = parse.urlparse(output_dir.url if output_dir.url else output_dir.path)
+        for name in rel_files:
+            local_filepath = os.path.join(local_keys, name)
+            if not os.path.isfile(local_filepath):
+                continue
+            remote_rel = os.path.join(_KEYS_DIRNAME, name)
+            if obj.scheme in _SUPPORTED_PROVIDERS:
+                assert output_dir.url
+                merged_storage_options = construct_storage_options(self.storage_options, output_dir)
+                fs_provider = _get_fs_provider(output_dir.url, merged_storage_options)
+                fs_provider.upload_file(local_filepath, os.path.join(output_dir.url, remote_rel))
+            elif output_dir.path and os.path.isdir(output_dir.path):
+                dest_dir = os.path.join(output_dir.path, _KEYS_DIRNAME)
+                os.makedirs(dest_dir, exist_ok=True)
+                dest = os.path.join(dest_dir, name)
+                if os.path.abspath(local_filepath) != os.path.abspath(dest):
+                    shutil.copyfile(local_filepath, dest)
+
+    def _upload_file(self, output_dir: Dir, cache_dir: str, filename: str) -> None:
+        if output_dir.path is None and output_dir.url is None:
+            return
+        local_filepath = os.path.join(cache_dir, filename)
+        if not os.path.isfile(local_filepath):
+            return
+        obj = parse.urlparse(output_dir.url if output_dir.url else output_dir.path)
+        if obj.scheme in _SUPPORTED_PROVIDERS:
+            assert output_dir.url
+            merged_storage_options = construct_storage_options(self.storage_options, output_dir)
+            fs_provider = _get_fs_provider(output_dir.url, merged_storage_options)
+            fs_provider.upload_file(local_filepath, os.path.join(output_dir.url, filename))
+        elif output_dir.path and os.path.isdir(output_dir.path):
+            shutil.copyfile(local_filepath, os.path.join(output_dir.path, filename))
+
     def _upload_index(self, output_dir: Dir, cache_dir: str, num_nodes: int, node_rank: int | None) -> None:
         """Upload the index file to the remote cloud directory."""
         if output_dir.path is None and output_dir.url is None:

src/litdata/processing/functions.py

@@ -169,13 +169,15 @@ def __init__(
         encryption: Encryption | None = None,
         existing_index: dict[str, Any] | None = None,
         storage_options: dict[str, Any] = {},
+        key_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(
             chunk_size=chunk_size,
             chunk_bytes=chunk_bytes,
             compression=compression,
             encryption=encryption,
             storage_options=storage_options,
+            key_fn=key_fn,
         )
         self._fn = fn
         self._inputs = inputs
@@ -220,13 +222,15 @@ def __init__(
         encryption: Encryption | None = None,
         existing_index: dict[str, Any] | None = None,
         storage_options: dict[str, Any] = {},
+        key_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(
             chunk_size=chunk_size,
             chunk_bytes=chunk_bytes,
             compression=compression,
             encryption=encryption,
             storage_options=storage_options,
+            key_fn=key_fn,
         )
         self._fn = fn
         self._queue = queue
@@ -422,6 +426,7 @@ def optimize(
     keep_data_ordered: bool = True,
     verbose: bool = True,
     broadcast_paths: bool = False,
+    key_fn: Callable[[Any], Any] | None = None,
 ) -> None:
     """This function converts a dataset into chunks, possibly in a distributed way.
 
@@ -473,6 +478,9 @@ def optimize(
         broadcast_paths: Broadcast resolved input/output dirs across multi-node ranks. Defaults to ``False``.
             Auto-enabled when ``input_dir`` or ``output_dir`` contains a ``{%strftime}`` time template so ranks
             share one expanded path. When off, each rank uses its locally resolved path.
+        key_fn: Optional callable ``sample -> key`` (str or int). When set, writes a ``keys/``
+            sidecar mapping each key to its sample ``index`` (and chunk location) for
+            ``dataset[key]`` / ``dataset_update``.
     """
     _check_version_and_prompt_upgrade(__version__)
 
@@ -608,6 +616,7 @@ def optimize(
                     encryption=encryption,
                     existing_index=existing_index_file_content,
                     storage_options=storage_options,
+                    key_fn=key_fn,
                 )
             else:
                 assert queue is not None
@@ -620,6 +629,7 @@ def optimize(
                     encryption=encryption,
                     existing_index=existing_index_file_content,
                     storage_options=storage_options,
+                    key_fn=key_fn,
                 )
             assert recipe is not None, "Recipe should be defined at this point."
             data_processor.run(recipe)

src/litdata/streaming/dataset.py

@@ -183,6 +183,8 @@ def __init__(
         # Upcoming in-chunk sample indexes for this worker. ``deque`` keeps ``popleft`` O(1);
         # a list ``pop(0)`` would be O(n) per sample for large chunks.
         self.upcoming_indexes: deque[int] = deque()
+        # Lazy parquet key index (``keys/shard-*.parquet``) for ``dataset[key]`` lookups.
+        self._key_index: Any | None = None
 
         # which index of the array `self.worker_chunks` will we work on after this chunk is completely consumed
         self.worker_next_chunk_index = 0
@@ -484,11 +486,15 @@ def _resume(self, workers_chunks: list[list[int]], workers_intervals: list[Any])
         # bump the chunk_index
         self.worker_next_chunk_index += 1
 
-    def __getitem__(self, index: ChunkedIndex | int | slice) -> Any:
+    def __getitem__(self, index: ChunkedIndex | int | slice | str) -> Any:
         if self.cache is None:
             self.worker_env = _WorkerEnv.detect()
             self.cache = self._create_cache(worker_env=self.worker_env)
             self.shuffler = self._create_shuffler(self.cache)
+        # String keys go through the keys/ store. Int remains a global sample index;
+        # use ``get_by_key`` for int entity keys.
+        if isinstance(index, str):
+            return self.get_by_key(index)
         if isinstance(index, int):
             index = ChunkedIndex(*self.cache._get_chunk_index_from_index(index))
         elif isinstance(index, slice):
@@ -507,6 +513,39 @@ def __getitem__(self, index: ChunkedIndex | int | slice) -> Any:
 
         return item
 
+    def get_by_key(self, key: Any) -> Any:
+        """Load a sample by entity key from the ``keys/`` store (str or int keys).
+
+        The key store is read from the dataset root (local cache if present, else
+        the remote URL). Remote lookups use Polars ``scan_parquet`` with predicate
+        pushdown so shards stay in object storage.
+        """
+        from litdata.utilities.keys_index import KeyIndex, has_keys_index
+
+        if self.cache is None:
+            self.worker_env = _WorkerEnv.detect()
+            self.cache = self._create_cache(worker_env=self.worker_env)
+            self.shuffler = self._create_shuffler(self.cache)
+
+        if self._key_index is None:
+            local_root = self.input_dir.path
+            remote_root = self.input_dir.url
+            root: str | None = None
+            if local_root and has_keys_index(local_root):
+                root = local_root
+            elif remote_root and has_keys_index(remote_root, self.storage_options):
+                root = remote_root
+            if root is None:
+                raise KeyError(f"Keyed access requires a keys/ index next to the dataset. Missing for key={key!r}.")
+            self._key_index = KeyIndex(root, storage_options=self.storage_options)
+
+        global_index, chunk_index, chunk_offset = self._key_index.resolve(key)
+        if chunk_index >= 0:
+            chunked = ChunkedIndex(index=chunk_offset, chunk_index=chunk_index)
+        else:
+            chunked = ChunkedIndex(*self.cache._get_chunk_index_from_index(global_index))
+        return self[chunked]
+
     def __next__(self) -> Any:
         # check if we have reached the end of the dataset (i.e., all the chunks have been processed)
         if self.global_index >= self.stop_length: