diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1c123da7b..fbdae4c81 100644 --- a/CONTRIBUTING.md +++ b/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 diff --git a/src/litdata/__init__.py b/src/litdata/__init__.py index e46acecda..eb89b39b1 100644 --- a/src/litdata/__init__.py +++ b/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", diff --git a/src/litdata/constants.py b/src/litdata/constants.py index bb9f64a24..ae18b9068 100644 --- a/src/litdata/constants.py +++ b/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") diff --git a/src/litdata/processing/data_processor.py b/src/litdata/processing/data_processor.py index a6c35f4f9..7ab9b44cc 100644 --- a/src/litdata/processing/data_processor.py +++ b/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: diff --git a/src/litdata/processing/functions.py b/src/litdata/processing/functions.py index d0f8e88b8..4c73d144d 100644 --- a/src/litdata/processing/functions.py +++ b/src/litdata/processing/functions.py @@ -169,6 +169,7 @@ 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, @@ -176,6 +177,7 @@ def __init__( compression=compression, encryption=encryption, storage_options=storage_options, + key_fn=key_fn, ) self._fn = fn self._inputs = inputs @@ -220,6 +222,7 @@ 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, @@ -227,6 +230,7 @@ def __init__( 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) diff --git a/src/litdata/streaming/dataset.py b/src/litdata/streaming/dataset.py index 1535b1b57..9d7094431 100644 --- a/src/litdata/streaming/dataset.py +++ b/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: diff --git a/src/litdata/streaming/dataset_update.py b/src/litdata/streaming/dataset_update.py new file mode 100644 index 000000000..1e7ff1d84 --- /dev/null +++ b/src/litdata/streaming/dataset_update.py @@ -0,0 +1,299 @@ +# Copyright The Lightning AI team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""In-place keyed updates for optimized LitData datasets. + +Requires a ``keys/`` sidecar (``keys/shard-*.parquet``) produced by +``optimize(..., key_fn=...)`` or ``build_keys_index``. +Only local dataset directories are supported in v1. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import re +import shutil +import tempfile +from collections import defaultdict +from time import time +from typing import Any + +from litdata.constants import _INDEX_FILENAME, _KEYS_DIRNAME +from litdata.streaming.cache import Cache +from litdata.streaming.resolver import _resolve_dir +from litdata.streaming.writer import BinaryWriter +from litdata.utilities.keys_index import ( + KeyIndex, + _atomic_replace, + has_keys_index, + normalize_key, +) + +_CHUNK_NAME_RE = re.compile( + r"^chunk-(?P\d+)-(?P\d+)" + r"(?:-u(?P\d+))?" + r"(?:\.(?P[^.]+))?\.bin$" +) + + +def dataset_update(input_dir: str) -> DatasetUpdate: + """Open a keyed update session for an optimized dataset. + + Example:: + + with dataset_update("optimized_data") as update: + update["sample-id"] = {"x": 1, "y": 2} + update.commit() + """ + return DatasetUpdate(input_dir) + + +class DatasetUpdate: + """Session-style context manager for keyed sample replaces. + + Stage changes with ``update[key] = sample``, then call :meth:`commit`. + Exiting the context without ``commit()`` discards pending changes. + After a successful commit, the session is closed for further writes. + """ + + def __init__(self, input_dir: str) -> None: + resolved = _resolve_dir(input_dir) + # Studio ``lightning_storage`` paths often resolve with both a local FUSE + # ``path`` and a remote ``url``. Prefer the local path so commits write + # through the mount. Pure remote URLs (no local dir) are not supported yet. + if resolved.path is None or not os.path.isdir(resolved.path): + if resolved.url is not None: + raise NotImplementedError( + "dataset_update() requires a local dataset directory " + "(or a lightning_storage FUSE path). " + f"Got remote-only url={resolved.url!r}." + ) + raise FileNotFoundError(f"Dataset directory not found: {input_dir}") + + self._dir = resolved.path + index_path = os.path.join(self._dir, _INDEX_FILENAME) + if not os.path.isfile(index_path): + raise FileNotFoundError(f"Missing {_INDEX_FILENAME} in {self._dir}. Did you run optimize()?") + if not has_keys_index(self._dir): + raise FileNotFoundError( + f"Missing {_KEYS_DIRNAME}/ key index in {self._dir}. " + "Re-run optimize(..., key_fn=...) or call build_keys_index(dir, key_fn)." + ) + + with open(index_path, encoding="utf-8") as f: + self._index: dict[str, Any] = json.load(f) + self._key_index = KeyIndex(self._dir) + self._pending: dict[str | int, Any] = {} + self._entered = False + self._committed = False + + # Build global_index → (chunk_list_index, local_offset) and chunk intervals + self._chunk_starts: list[int] = [] + start = 0 + for chunk in self._index["chunks"]: + self._chunk_starts.append(start) + start += int(chunk["chunk_size"]) + self._length = start + + def __enter__(self) -> DatasetUpdate: + self._entered = True + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + # Like SQLAlchemy Session: no implicit commit — uncommitted work is dropped. + try: + self._pending.clear() + finally: + self._entered = False + self._key_index.close() + + def __setitem__(self, key: Any, sample: Any) -> None: + self.replace(key, sample) + + def replace(self, key: Any, sample: Any) -> None: + if self._committed: + raise RuntimeError("Cannot modify dataset_update after commit().") + # Existence is validated on commit via a single batched parquet scan. + self._pending[normalize_key(key)] = sample + + def commit(self) -> None: + """Persist pending keyed replaces. Further modifications are rejected.""" + if self._committed: + raise RuntimeError("dataset_update session already committed.") + if not self._pending: + self._committed = True + return + + resolved = self._key_index.resolve_many(list(self._pending.keys())) + missing = [k for k in self._pending if k not in resolved] + if missing: + raise KeyError(f"Unknown dataset key: {missing[0]!r}") + + # Group pending updates by chunk index in index["chunks"] + updates_by_chunk: dict[int, dict[int, Any]] = defaultdict(dict) + for key, sample in self._pending.items(): + global_index, chunk_i, _chunk_off = resolved[key] + if chunk_i < 0: + chunk_i, local_i = self._locate(global_index) + else: + local_i = global_index - self._chunk_starts[chunk_i] + updates_by_chunk[chunk_i][local_i] = sample + + config = self._index["config"] + compression = config.get("compression") + + for chunk_i, local_updates in updates_by_chunk.items(): + self._rewrite_chunk(chunk_i, local_updates, compression, config) + + self._index["updated_at"] = str(time()) + index_path = os.path.join(self._dir, _INDEX_FILENAME) + tmp_index = f"{index_path}.tmp" + with open(tmp_index, "w", encoding="utf-8") as f: + json.dump(self._index, f, sort_keys=True) + _atomic_replace(tmp_index, index_path) + + # keys/ store unchanged for whole-sample replace (same key set / indices) + self._pending.clear() + self._committed = True + self._sync_studio_read_cache() + + def _sync_studio_read_cache(self) -> None: + """Copy the updated dataset into the Studio chunk cache. + + Paths under ``/teamspace/lightning_storage/`` (and similar) are rewritten by + ``StreamingDataset`` to ``/cache/chunks/...``, which is filled from the remote + URL. FUSE writes are not always visible on R2 immediately, so without this + sync a post-commit read can still see stale chunks. + """ + from litdata.utilities.dataset_utilities import _should_replace_path, _try_create_cache_dir + + if not _should_replace_path(self._dir): + return + cache_path = _try_create_cache_dir(self._dir) + if cache_path is None: + return + + for name in os.listdir(self._dir): + src = os.path.join(self._dir, name) + dst = os.path.join(cache_path, name) + if name == "keys" and os.path.isdir(src): + if os.path.isdir(dst): + shutil.rmtree(dst) + shutil.copytree(src, dst) + elif os.path.isfile(src): + shutil.copy2(src, dst) + + def _locate(self, global_index: int) -> tuple[int, int]: + if global_index < 0 or global_index >= self._length: + raise IndexError(global_index) + # linear scan is fine for typical chunk counts; chunks are contiguous + for chunk_i, start in enumerate(self._chunk_starts): + size = int(self._index["chunks"][chunk_i]["chunk_size"]) + if start <= global_index < start + size: + return chunk_i, global_index - start + raise IndexError(global_index) + + def _rewrite_chunk( + self, + chunk_i: int, + local_updates: dict[int, Any], + compression: str | None, + config: dict[str, Any], + ) -> None: + chunk_info = self._index["chunks"][chunk_i] + filename = chunk_info["filename"] + match = _CHUNK_NAME_RE.match(filename) + if not match: + raise ValueError(f"Unrecognized chunk filename: {filename}") + rank = int(match.group("rank")) + chunk_index = int(match.group("chunk_index")) + chunk_size = int(chunk_info["chunk_size"]) + global_start = self._chunk_starts[chunk_i] + + cache = Cache(self._dir, chunk_bytes=1) + samples: list[Any] + try: + samples = [cache[global_start + i] for i in range(chunk_size)] + finally: + # Windows cannot replace a chunk that still has an open mmap/handle. + item_loader = getattr(cache._reader, "_item_loader", None) + close_open = getattr(item_loader, "_close_open_chunk", None) + if callable(close_open): + close_open() + elif item_loader is not None and hasattr(item_loader, "close"): + item_loader.close(chunk_i) + del cache + + for local_i, sample in local_updates.items(): + samples[local_i] = sample + + tmp_dir = tempfile.mkdtemp(prefix="litdata-update-") + try: + writer = BinaryWriter( + tmp_dir, + chunk_size=len(samples), + compression=compression, + chunk_index=chunk_index, + ) + writer._rank = rank + + for i, sample in enumerate(samples): + writer.add_item(i, sample) + writer.done() + + # Ensure rewritten samples stay schema-compatible with the dataset. + new_config = writer.get_config() + if new_config.get("data_format") != config.get("data_format"): + raise ValueError( + "Updated sample is incompatible with the dataset data_format. " + f"Expected {config.get('data_format')}, got {new_config.get('data_format')}." + ) + + new_chunk_path = os.path.join(tmp_dir, filename) + if not os.path.isfile(new_chunk_path): + # BinaryWriter may have used a slightly different name if compression unset + produced = [f for f in os.listdir(tmp_dir) if f.startswith("chunk-") and f.endswith(".bin")] + if len(produced) != 1: + raise RuntimeError(f"Expected one rewritten chunk, found {produced} in {tmp_dir}") + new_chunk_path = os.path.join(tmp_dir, produced[0]) + filename = produced[0] + + dest = os.path.join(self._dir, filename) + tmp_dest = dest + ".tmp" + shutil.copyfile(new_chunk_path, tmp_dest) + try: + _atomic_replace(tmp_dest, dest) + except PermissionError: + # Destination still locked (e.g. another StreamingDataset mmap on Windows). + # Publish under a new filename and retarget the index entry. + compression_part = match.group("compression") + update_token = int(time() * 1000) % 1_000_000_000 + if compression_part: + filename = f"chunk-{rank}-{chunk_index}-u{update_token}.{compression_part}.bin" + else: + filename = f"chunk-{rank}-{chunk_index}-u{update_token}.bin" + dest = os.path.join(self._dir, filename) + _atomic_replace(tmp_dest, dest) + old_path = os.path.join(self._dir, chunk_info["filename"]) + with contextlib.suppress(OSError, PermissionError): + if old_path != dest and os.path.isfile(old_path): + os.remove(old_path) + + new_size = os.path.getsize(dest) + chunk_info["filename"] = filename + chunk_info["chunk_bytes"] = new_size + chunk_info["chunk_size"] = len(samples) + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) diff --git a/src/litdata/utilities/format.py b/src/litdata/utilities/format.py index 7a02955ae..16eb09e0f 100644 --- a/src/litdata/utilities/format.py +++ b/src/litdata/utilities/format.py @@ -51,7 +51,7 @@ def _get_tqdm_iterator_if_available() -> Any: return _tqdm - def _pass_through(iterator: Any) -> Any: + def _pass_through(iterator: Any, *args: Any, **kwargs: Any) -> Any: yield from iterator return _pass_through diff --git a/src/litdata/utilities/keys_index.py b/src/litdata/utilities/keys_index.py new file mode 100644 index 000000000..377ffa5bb --- /dev/null +++ b/src/litdata/utilities/keys_index.py @@ -0,0 +1,868 @@ +# Copyright The Lightning AI team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parquet key / metadata sidecar for optimized datasets (Polars-backed). + +Default on-disk layout:: + + index.json # includes a ``keys`` section (num_shards / sharding) + keys/ + shard-00000.parquet + shard-00001.parquet # when num_shards > 1 + +Each shard parquet is sorted by ``key`` and contains: + +* ``key`` — str or int64 (opaque entity id) +* ``index`` — int64 global sample index +* ``chunk_index`` — int32 chunk list index in ``index.json`` (optional until enriched) +* ``chunk_offset`` — int32 value passed as ``ChunkedIndex.index`` (optional until enriched) + +With ``num_shards > 1``, rows are assigned with a stable hash of ``key``. +Legacy single-file ``keys.parquet`` is still readable. + +Written by ``optimize(..., key_fn=...)`` / ``build_keys_index``. Used by +``dataset_update`` and ``StreamingDataset.get_by_key`` / ``dataset[str_key]``. + +Remote datasets keep shards in object storage. Lookups use +``scan_parquet(s3://.../keys/shard-*.parquet)`` with predicate pushdown so only +parquet footer + matching row groups are fetched — shards are not downloaded +wholesale into the chunk cache. Shard layout is described in ``index.json``. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import os +import re +import shutil +import tempfile +from collections import defaultdict +from collections.abc import Callable, Iterator, Sequence +from time import sleep +from typing import Any +from urllib import parse + +from litdata.constants import ( + _DEFAULT_KEYS_NUM_SHARDS, + _INDEX_FILENAME, + _KEYS_DIRNAME, + _KEYS_FILENAME, + _KEYS_SHARD_TEMPLATE, + _POLARS_AVAILABLE, + _RANK_KEYS_SUFFIX, + _SUPPORTED_PROVIDERS, +) + + +def _require_polars() -> Any: + if not _POLARS_AVAILABLE: + raise ModuleNotFoundError( + "Polars is required for optimize(key_fn=...) / dataset_update. Install with `pip install 'polars>1.0.0'`." + ) + import polars as pl + + return pl + + +def _atomic_replace(tmp_path: str, dest_path: str) -> None: + """Replace ``dest_path`` with ``tmp_path``, retrying Windows file-lock races. + + On Windows, ``os.replace`` fails with ``PermissionError`` if another handle + still has ``dest_path`` open (or antivirus briefly locks it). Retry with a + short backoff instead of failing the write. + """ + last_err: PermissionError | None = None + for _ in range(20): + try: + os.replace(tmp_path, dest_path) + return + except PermissionError as e: + last_err = e + sleep(0.05) + assert last_err is not None + raise last_err + + +def normalize_key(key: Any) -> str | int: + """Normalize a user key to str or int.""" + if isinstance(key, bool) or key is None: + raise TypeError(f"Unsupported key type: {type(key)!r} ({key!r})") + if isinstance(key, int): + return key + if isinstance(key, str): + return key + if isinstance(key, bytes): + return key.decode("utf-8") + try: + import numpy as np + + if isinstance(key, np.integer): + return int(key) + except ImportError: + pass + raise TypeError(f"Unsupported key type: {type(key)!r}. Use str or int.") + + +def _natural_key(s: str) -> list: + return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", s)] + + +def key_shard(key: Any, num_shards: int) -> int: + """Stable shard id for ``key`` in ``[0, num_shards)``.""" + if num_shards <= 1: + return 0 + nkey = normalize_key(key) + digest = hashlib.blake2b(str(nkey).encode("utf-8"), digest_size=8).digest() + return int.from_bytes(digest, "little") % num_shards + + +def _is_remote_uri(path: str) -> bool: + return parse.urlparse(path).scheme in _SUPPORTED_PROVIDERS + + +def _join_uri(base: str, *parts: str) -> str: + if _is_remote_uri(base): + return "/".join([base.rstrip("/")] + [p.strip("/") for p in parts]) + return os.path.join(base, *parts) + + +def keys_dir(dataset_dir: str) -> str: + return _join_uri(dataset_dir, _KEYS_DIRNAME) + + +def shard_filename(shard: int) -> str: + return _KEYS_SHARD_TEMPLATE.format(shard) + + +def shard_path(dataset_dir: str, shard: int = 0) -> str: + return _join_uri(keys_dir(dataset_dir), shard_filename(shard)) + + +def keys_path(dataset_dir: str) -> str: + """Path/URI to the default shard (or legacy file if that is all that exists).""" + default_shard = shard_path(dataset_dir, 0) + if _is_remote_uri(dataset_dir): + return default_shard + if os.path.isfile(default_shard) or os.path.isdir(keys_dir(dataset_dir)): + return default_shard + legacy = os.path.join(dataset_dir, _KEYS_FILENAME) + if os.path.isfile(legacy): + return legacy + return default_shard + + +def has_keys_index(dataset_dir: str, storage_options: dict[str, Any] | None = None) -> bool: + if _is_remote_uri(dataset_dir): + from litdata.streaming.fs_provider import _get_fs_provider + + fs = _get_fs_provider(dataset_dir, storage_options) + return fs.exists(shard_path(dataset_dir, 0)) or fs.exists(_join_uri(dataset_dir, _KEYS_FILENAME)) + return os.path.isfile(shard_path(dataset_dir, 0)) or os.path.isfile(os.path.join(dataset_dir, _KEYS_FILENAME)) + + +def keys_config(num_shards: int, sharding: str = "hash") -> dict[str, Any]: + """Build the ``keys`` section stored inside ``index.json``.""" + return { + "version": 1, + "num_shards": int(num_shards), + "sharding": sharding if num_shards > 1 else "none", + } + + +def set_keys_config_in_index( + dataset_dir: str, + num_shards: int, + sharding: str = "hash", +) -> None: + """Write / update the ``keys`` section on a local ``index.json``.""" + index_path = os.path.join(dataset_dir, _INDEX_FILENAME) + if not os.path.isfile(index_path): + return + with open(index_path, encoding="utf-8") as f: + data = json.load(f) + data["keys"] = keys_config(num_shards, sharding=sharding) + tmp = f"{index_path}.tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(data, f, sort_keys=True) + _atomic_replace(tmp, index_path) + + +def read_keys_config(dataset_dir: str, storage_options: dict[str, Any] | None = None) -> dict[str, Any] | None: + """Read the ``keys`` section from ``index.json`` (local or remote).""" + index_path = _join_uri(dataset_dir, _INDEX_FILENAME) + if _is_remote_uri(dataset_dir): + from litdata.streaming.fs_provider import _get_fs_provider + + fs = _get_fs_provider(dataset_dir, storage_options) + if not fs.exists(index_path): + return None + data = _read_remote_json(index_path, storage_options) + cfg = data.get("keys") + return cfg if isinstance(cfg, dict) else None + + if not os.path.isfile(index_path): + return None + with open(index_path, encoding="utf-8") as f: + data = json.load(f) + cfg = data.get("keys") + return cfg if isinstance(cfg, dict) else None + + +def _infer_keys_config_from_shards(dataset_dir: str) -> dict[str, Any]: + shards = list_shard_files(keys_dir(dataset_dir)) + n = max(len(shards), 1) + return keys_config(n, sharding="hash" if n > 1 else "none") + + +def list_shard_files(keys_directory: str) -> list[str]: + if not os.path.isdir(keys_directory): + return [] + files = [ + os.path.join(keys_directory, name) + for name in os.listdir(keys_directory) + if name.startswith("shard-") and name.endswith(".parquet") + ] + return sorted(files, key=lambda p: _natural_key(os.path.basename(p))) + + +def list_key_parquet_files(dataset_dir: str) -> list[str]: + """Parquet files that make up the key index (sharded store or legacy file).""" + shards = list_shard_files(keys_dir(dataset_dir)) + if shards: + return shards + legacy = os.path.join(dataset_dir, _KEYS_FILENAME) + if os.path.isfile(legacy): + return [legacy] + return [] + + +def _read_remote_json(url: str, storage_options: dict[str, Any] | None) -> dict[str, Any]: + from litdata.streaming.fs_provider import _get_fs_provider + + fs = _get_fs_provider(url, storage_options) + with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as tmp: + local_path = tmp.name + try: + fs.download_file(url, local_path) + with open(local_path, encoding="utf-8") as f: + return json.load(f) + finally: + with contextlib.suppress(OSError): + os.remove(local_path) + + +def _resolve_remote_key_paths(dataset_url: str, storage_options: dict[str, Any] | None) -> tuple[list[str], int, str]: + """Resolve remote shard URIs using ``index.json``'s ``keys`` section.""" + from litdata.streaming.fs_provider import _get_fs_provider + + fs = _get_fs_provider(dataset_url, storage_options) + shard0_url = shard_path(dataset_url, 0) + legacy_url = _join_uri(dataset_url, _KEYS_FILENAME) + + if fs.exists(shard0_url): + cfg = read_keys_config(dataset_url, storage_options) or {"num_shards": 1, "sharding": "none"} + num_shards = int(cfg.get("num_shards", 1)) + sharding = str(cfg.get("sharding", "hash" if num_shards > 1 else "none")) + paths = [shard_path(dataset_url, i) for i in range(num_shards)] + return paths, num_shards, sharding + + if fs.exists(legacy_url): + return [legacy_url], 1, "none" + + raise FileNotFoundError(f"No key index under {dataset_url!r} (expected {_KEYS_DIRNAME}/ or {_KEYS_FILENAME})") + + +def _resolve_key_paths(path: str, storage_options: dict[str, Any] | None = None) -> tuple[list[str], int, str]: + """Resolve a dataset dir/URL, ``keys/`` dir, or parquet file into shard paths/URIs.""" + if _is_remote_uri(path): + # Remote parquet file directly, or dataset / keys prefix. + if path.rstrip("/").endswith(".parquet"): + return [path], 1, "none" + if path.rstrip("/").endswith(f"/{_KEYS_DIRNAME}") or path.rstrip("/").endswith(_KEYS_DIRNAME): + # keys/ URL without dataset root — read manifest beside shards + dataset_url = path.rstrip("/").rsplit("/", 1)[0] + return _resolve_remote_key_paths(dataset_url, storage_options) + return _resolve_remote_key_paths(path, storage_options) + + if os.path.isfile(path): + return [path], 1, "none" + + if not os.path.isdir(path): + raise FileNotFoundError(path) + + # Dataset directory or a keys/ directory with shards. + direct_shards = list_shard_files(path) + if direct_shards: + keys_directory = path + dataset_root = os.path.dirname(path) if os.path.basename(path.rstrip(os.sep)) == _KEYS_DIRNAME else path + else: + nested = keys_dir(path) + nested_shards = list_shard_files(nested) + if nested_shards: + keys_directory = nested + dataset_root = path + else: + legacy = os.path.join(path, _KEYS_FILENAME) + if os.path.isfile(legacy): + return [legacy], 1, "none" + raise FileNotFoundError(f"No key index under {path!r} (expected {_KEYS_DIRNAME}/ or {_KEYS_FILENAME})") + + shards = list_shard_files(keys_directory) + if not shards: + raise FileNotFoundError(f"No shard-*.parquet files in {keys_directory}") + cfg = read_keys_config(dataset_root, storage_options) or _infer_keys_config_from_shards(dataset_root) + num_shards = int(cfg.get("num_shards", len(shards))) + sharding = str(cfg.get("sharding", "hash" if num_shards > 1 else "none")) + return shards, num_shards, sharding + + +def _polars_cloud_paths_and_options( + paths: Sequence[str], storage_options: dict[str, Any] | None +) -> tuple[list[str], dict[str, Any] | None]: + """Adapt URIs/credentials for Polars ``scan_parquet`` (object_store).""" + if not paths or not _is_remote_uri(paths[0]): + return list(paths), None + + opts = dict(storage_options or {}) + scheme = parse.urlparse(paths[0]).scheme + + # R2: Polars talks S3 protocol; resolve Lightning data-connection creds if needed. + if scheme == "r2": + if opts.get("data_connection_id") and not opts.get("aws_access_key_id"): + from litdata.streaming.client import R2Client + + client = R2Client(storage_options=opts) + creds = client.get_r2_bucket_credentials(opts["data_connection_id"]) + opts = {k: v for k, v in opts.items() if k != "data_connection_id"} + opts.update(creds) + paths = ["s3://" + p[len("r2://") :] for p in paths] + + # Normalize endpoint key for object_store / Polars. + if "endpoint_url" in opts and "aws_endpoint_url" not in opts: + opts["aws_endpoint_url"] = opts.pop("endpoint_url") + opts.pop("data_connection_id", None) + + return list(paths), (opts or None) + + +def _to_dataframe( + keys: Sequence[str | int], + indices: Sequence[int], + chunk_indices: Sequence[int] | None = None, + chunk_offsets: Sequence[int] | None = None, +) -> Any: + pl = _require_polars() + if len(keys) != len(indices): + raise ValueError("keys and indices must have the same length") + + if keys and all(isinstance(k, int) and not isinstance(k, bool) for k in keys): + key_series = pl.Series("key", list(keys), dtype=pl.Int64) + else: + key_series = pl.Series("key", [str(k) for k in keys], dtype=pl.Utf8) + + data: dict[str, Any] = { + "key": key_series, + "index": pl.Series("index", list(indices), dtype=pl.Int64), + } + if chunk_indices is not None: + data["chunk_index"] = pl.Series("chunk_index", list(chunk_indices), dtype=pl.Int32) + if chunk_offsets is not None: + data["chunk_offset"] = pl.Series("chunk_offset", list(chunk_offsets), dtype=pl.Int32) + return pl.DataFrame(data) + + +def save_keys( + path: str, + keys: Sequence[str | int], + indices: Sequence[int] | None = None, + chunk_indices: Sequence[int] | None = None, + chunk_offsets: Sequence[int] | None = None, + *, + sort_by_key: bool = True, +) -> None: + """Write a single keys parquet file via Polars.""" + if indices is None: + indices = list(range(len(keys))) + df = _to_dataframe(keys, indices, chunk_indices, chunk_offsets) + + n_unique = df["key"].n_unique() + if n_unique != df.height: + raise ValueError("Duplicate keys are not allowed in the keys sidecar.") + + if sort_by_key and df.height > 0: + df = df.sort("key") + + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + tmp = f"{path}.tmp" + df.write_parquet(tmp, compression="zstd", row_group_size=min(1_000_000, max(df.height, 1))) + _atomic_replace(tmp, path) + + +def write_keys_store( + dataset_dir: str, + keys: Sequence[str | int], + indices: Sequence[int] | None = None, + chunk_indices: Sequence[int] | None = None, + chunk_offsets: Sequence[int] | None = None, + *, + num_shards: int = _DEFAULT_KEYS_NUM_SHARDS, +) -> str: + """Write ``keys/shard-*.parquet`` and record layout in ``index.json`` ``keys`` section.""" + if num_shards < 1: + raise ValueError("num_shards must be >= 1") + if indices is None: + indices = list(range(len(keys))) + if len(keys) != len(indices): + raise ValueError("keys and indices must have the same length") + + df = _to_dataframe(keys, indices, chunk_indices, chunk_offsets) + n_unique = df["key"].n_unique() + if n_unique != df.height: + raise ValueError("Duplicate keys are not allowed in the keys sidecar.") + + out_dir = keys_dir(dataset_dir) + if os.path.isdir(out_dir): + shutil.rmtree(out_dir) + os.makedirs(out_dir, exist_ok=True) + + if num_shards == 1: + path = shard_path(dataset_dir, 0) + if df.height > 0: + df = df.sort("key") + tmp = f"{path}.tmp" + df.write_parquet(tmp, compression="zstd", row_group_size=min(1_000_000, max(df.height, 1))) + _atomic_replace(tmp, path) + else: + shard_ids = [key_shard(k, num_shards) for k in df["key"].to_list()] + pl = _require_polars() + df = df.with_columns(pl.Series("_shard", shard_ids, dtype=pl.Int32)) + for shard in range(num_shards): + part = df.filter(pl.col("_shard") == shard).drop("_shard") + if part.height > 0: + part = part.sort("key") + path = shard_path(dataset_dir, shard) + tmp = f"{path}.tmp" + part.write_parquet(tmp, compression="zstd", row_group_size=min(1_000_000, max(part.height, 1))) + _atomic_replace(tmp, path) + + set_keys_config_in_index(dataset_dir, num_shards=num_shards) + # Drop legacy single-file sidecar if present so readers prefer the store. + legacy = os.path.join(dataset_dir, _KEYS_FILENAME) + if os.path.isfile(legacy): + os.remove(legacy) + return out_dir + + +def save_rank_keys(path: str, index_key_pairs: Sequence[tuple[int, str | int]]) -> None: + """Write a per-rank keys file from explicit ``(sample_index, key)`` pairs.""" + if not index_key_pairs: + save_keys(path, [], indices=[]) + return + ordered = sorted(index_key_pairs, key=lambda x: x[0]) + indices = [i for i, _ in ordered] + keys = [k for _, k in ordered] + if indices != list(range(len(indices))): + raise ValueError(f"Rank key indexes must be contiguous from 0. Found {indices[:5]}... (len={len(indices)})") + save_keys(path, keys, indices=indices, sort_by_key=False) + + +def _chunk_columns_for_indices( + global_indices: Sequence[int], index_json: dict[str, Any] +) -> tuple[list[int], list[int]]: + chunk_starts: list[int] = [] + start = 0 + for chunk in index_json["chunks"]: + chunk_starts.append(start) + start += int(chunk["chunk_size"]) + + chunk_indices: list[int] = [] + chunk_offsets: list[int] = [] + for gidx in global_indices: + for ci, cstart in enumerate(chunk_starts): + csize = int(index_json["chunks"][ci]["chunk_size"]) + if cstart <= gidx < cstart + csize: + chunk_indices.append(ci) + # For default (non-subsampled) layouts this matches + # ChunksConfig._get_chunk_index_from_index's first return value. + chunk_offsets.append(int(gidx)) + break + else: + raise IndexError(f"Sample index {gidx} out of range for chunk layout") + return chunk_indices, chunk_offsets + + +def enrich_keys_with_chunks(path: str, index_json: dict[str, Any]) -> None: + """Add ``chunk_index`` / ``chunk_offset`` columns to a parquet file or keys store.""" + pl = _require_polars() + + # Keys store directory (dataset dir or keys/) + if os.path.isdir(path): + try: + paths, num_shards, _sharding = _resolve_key_paths(path) + except FileNotFoundError: + paths, num_shards = [], _DEFAULT_KEYS_NUM_SHARDS + if not paths: + return + # Re-load rows, enrich, rewrite store under the dataset dir. + dataset_dir = path if os.path.basename(path.rstrip(os.sep)) != _KEYS_DIRNAME else os.path.dirname(path) + df = pl.concat([pl.read_parquet(p) for p in paths]) + chunk_indices, chunk_offsets = _chunk_columns_for_indices(df["index"].to_list(), index_json) + write_keys_store( + dataset_dir, + df["key"].to_list(), + indices=df["index"].to_list(), + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + num_shards=num_shards, + ) + return + + df = pl.read_parquet(path) + chunk_indices, chunk_offsets = _chunk_columns_for_indices(df["index"].to_list(), index_json) + save_keys( + path, + df["key"].to_list(), + indices=df["index"].to_list(), + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + ) + + +class KeyIndex: + """Lazy Parquet key lookup — no full in-memory key map. + + ``path`` may be a local dataset directory, a ``keys/`` directory, a parquet + file, or a remote dataset URL (``s3://``, ``gs://``, ``r2://``). + + Lookups use ``scan_parquet`` + predicate pushdown (and hash routing when + sharded). For remote URLs this fetches only parquet metadata / matching row + groups — not the whole shard into the local chunk cache. + """ + + def __init__(self, path: str, storage_options: dict[str, Any] | None = None) -> None: + _require_polars() + self.path = path + self._storage_options = storage_options + self._paths, self._num_shards, self._sharding = _resolve_key_paths(path, storage_options) + schema_names = list(self._scan(self._paths[0]).collect_schema().names()) + self._has_chunks = "chunk_index" in schema_names and "chunk_offset" in schema_names + self._count: int | None = None + + def close(self) -> None: + return None + + def __enter__(self) -> KeyIndex: + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + def _scan(self, paths: str | Sequence[str] | None = None) -> Any: + pl = _require_polars() + target = self._paths if paths is None else paths + target_list = [target] if isinstance(target, str) else list(target) + cloud_paths, cloud_opts = _polars_cloud_paths_and_options(target_list, self._storage_options) + if cloud_opts is not None: + return pl.scan_parquet(cloud_paths if len(cloud_paths) > 1 else cloud_paths[0], storage_options=cloud_opts) + return pl.scan_parquet(cloud_paths if len(cloud_paths) > 1 else cloud_paths[0]) + + def _paths_for_keys(self, keys: Sequence[Any]) -> list[str]: + if self._num_shards <= 1 or self._sharding == "none": + return self._paths + by_shard: dict[int, None] = {} + for key in keys: + by_shard[key_shard(key, self._num_shards)] = None + return [self._paths[s] for s in sorted(by_shard) if s < len(self._paths)] + + def __len__(self) -> int: + if self._count is None: + pl = _require_polars() + self._count = int(self._scan().select(pl.len()).collect().item()) + return self._count + + def __contains__(self, key: Any) -> bool: + return self.get(key) is not None + + def _lookup_rows(self, keys: Sequence[Any]) -> Any: + """Fetch rows for the given keys.""" + pl = _require_polars() + nkeys = [normalize_key(k) for k in keys] + cols = ["key", "index"] + if self._has_chunks: + cols.extend(["chunk_index", "chunk_offset"]) + paths = self._paths_for_keys(nkeys) + return self._scan(paths).filter(pl.col("key").is_in(nkeys)).select(cols).collect() + + def get(self, key: Any, default: int | None = None) -> int | None: + df = self._lookup_rows([key]) + if df.height == 0: + return default + return int(df["index"][0]) + + def __getitem__(self, key: Any) -> int: + idx = self.get(key) + if idx is None: + raise KeyError(key) + return idx + + def resolve(self, key: Any) -> tuple[int, int, int]: + """Return ``(global_index, chunk_index, chunk_offset)`` for ``key``.""" + df = self._lookup_rows([key]) + if df.height == 0: + raise KeyError(normalize_key(key)) + gidx = int(df["index"][0]) + if self._has_chunks: + return gidx, int(df["chunk_index"][0]), int(df["chunk_offset"][0]) + return gidx, -1, -1 + + def resolve_many(self, keys: Sequence[Any]) -> dict[Any, tuple[int, int, int]]: + """Batch resolve keys (routed per shard when sharded).""" + if not keys: + return {} + nkeys = [normalize_key(k) for k in keys] + if self._num_shards <= 1 or self._sharding == "none": + df = self._lookup_rows(nkeys) + else: + pl = _require_polars() + grouped: dict[int, list[Any]] = defaultdict(list) + for key in nkeys: + grouped[key_shard(key, self._num_shards)].append(key) + cols = ["key", "index"] + if self._has_chunks: + cols.extend(["chunk_index", "chunk_offset"]) + frames = [] + for shard, shard_keys in grouped.items(): + if shard >= len(self._paths): + continue + frames.append( + self._scan(self._paths[shard]).filter(pl.col("key").is_in(shard_keys)).select(cols).collect() + ) + df = pl.concat(frames) if frames else pl.DataFrame({c: [] for c in cols}) + + out: dict[Any, tuple[int, int, int]] = {} + for row in df.iter_rows(named=True): + key = row["key"] + if self._has_chunks: + out[key] = (int(row["index"]), int(row["chunk_index"]), int(row["chunk_offset"])) + else: + out[key] = (int(row["index"]), -1, -1) + return out + + def key_at(self, index: int) -> str | int: + """Return the key for a global sample index (filter — avoid on hot paths).""" + pl = _require_polars() + matched = self._scan().filter(pl.col("index") == index).select("key").collect() + if matched.height == 0: + raise IndexError(index) + return matched["key"][0] + + def keys(self) -> Iterator[str | int]: + """Stream keys from parquet. Avoid on huge sidecars — materializes column.""" + yield from self._scan().select("key").collect().to_series().to_list() + + def indices(self) -> list[int]: + """Load all global indexes. Avoid on huge sidecars.""" + return self._scan().select("index").collect().to_series().to_list() + + +def load_keys(path: str) -> list[str | int]: + """Load all keys into a list (tests / small datasets only — not for millions of keys).""" + with KeyIndex(path) as idx: + return list(idx.keys()) + + +def load_key_index(path: str) -> KeyIndex: + return KeyIndex(path) + + +def merge_rank_key_files( + cache_dir: str, + output_filename: str | None = None, + *, + num_shards: int = _DEFAULT_KEYS_NUM_SHARDS, +) -> str | None: + """Merge ``{rank}.keys.parquet`` in rank order. + + By default writes the sharded store under ``cache_dir/keys/``. + When ``output_filename`` is set (multi-node partials), writes a single flat + parquet at ``cache_dir/output_filename`` instead. + """ + pl = _require_polars() + files = [f for f in os.listdir(cache_dir) if f.endswith(_RANK_KEYS_SUFFIX)] + if not files: + return None + + frames = [] + global_base = 0 + for filename in sorted(files, key=_natural_key): + filepath = os.path.join(cache_dir, filename) + df = pl.read_parquet(filepath) + df = df.with_columns((pl.col("index") + global_base).alias("index")) + global_base += df.height + frames.append(df.select(["key", "index"])) + os.remove(filepath) + + merged = pl.concat(frames) if frames else pl.DataFrame({"key": [], "index": []}) + keys = merged["key"].to_list() if merged.height else [] + indices = merged["index"].to_list() if merged.height else [] + + if output_filename is not None: + out_path = os.path.join(cache_dir, output_filename) + save_keys(out_path, keys, indices=indices, sort_by_key=True) + return out_path + + return write_keys_store(cache_dir, keys, indices=indices, num_shards=num_shards) + + +def concatenate_key_files( + paths: Sequence[str], + dataset_dir: str, + *, + num_shards: int = _DEFAULT_KEYS_NUM_SHARDS, +) -> str: + """Concatenate key parquet files into ``dataset_dir/keys/``.""" + pl = _require_polars() + frames = [] + base = 0 + for filepath in paths: + df = pl.read_parquet(filepath).select(["key", "index"]) + if df.height == 0: + continue + local_min = int(df["index"].min()) + df = df.with_columns((pl.col("index") - local_min + base).alias("index")) + base += df.height + frames.append(df) + + merged = pl.concat(frames) if frames else pl.DataFrame({"key": [], "index": []}) + return write_keys_store( + dataset_dir, + merged["key"].to_list() if merged.height else [], + indices=merged["index"].to_list() if merged.height else [], + num_shards=num_shards, + ) + + +def iter_key_indexes( + input_dir: str, + key_fn: Callable[[Any], Any], + *, + verbose: bool = True, +) -> Iterator[tuple[str | int, int]]: + """Yield ``(key, global_index)`` by scanning an already-optimized dataset. + + Walks chunks sequentially through :class:`~litdata.streaming.cache.Cache` + (keeps each chunk mmap warm) instead of per-sample ``dataset[i]`` lookups. + + Does not materialize a full key→index map. Prefer :func:`build_keys_index` + when you need a durable sidecar for ``dataset_update`` / keyed reads. + """ + from litdata.constants import _TQDM_AVAILABLE + from litdata.streaming.cache import Cache + from litdata.streaming.resolver import _resolve_dir + from litdata.streaming.sampler import ChunkedIndex + + resolved = _resolve_dir(input_dir) + if resolved.path is None or not os.path.isdir(resolved.path): + raise FileNotFoundError(f"Dataset directory not found: {input_dir}") + + index_path = os.path.join(resolved.path, _INDEX_FILENAME) + if not os.path.isfile(index_path): + raise FileNotFoundError(f"Missing {_INDEX_FILENAME} in {resolved.path}. Did you run optimize()?") + + cache = Cache(resolved.path, chunk_bytes=1) + intervals = cache.get_chunk_intervals() + total = len(cache) + + pbar: Any = None + if verbose and _TQDM_AVAILABLE: + from tqdm.auto import tqdm as _tqdm + + pbar = _tqdm(total=total, desc="Building keys index", unit="sample") + + global_index = 0 + try: + for chunk_index, interval in enumerate(intervals): + begin = int(interval.roi_start_idx) + end = int(interval.roi_end_idx) + chunk_size = end - begin + for offset in range(chunk_size): + sample = cache[ + ChunkedIndex( + index=begin + offset, + chunk_index=chunk_index, + chunk_size=chunk_size, + ) + ] + yield normalize_key(key_fn(sample)), global_index + global_index += 1 + if pbar is not None: + pbar.update(1) + finally: + if pbar is not None: + pbar.close() + # Release Windows file locks from mmap before callers rewrite chunks. + item_loader = getattr(cache._reader, "_item_loader", None) + close_open = getattr(item_loader, "_close_open_chunk", None) + if callable(close_open): + close_open() + elif item_loader is not None and hasattr(item_loader, "close") and intervals: + item_loader.close(len(intervals) - 1) + + if global_index != total: + raise RuntimeError(f"Key scan produced {global_index} samples but index reports {total}.") + + +def build_keys_index( + input_dir: str, + key_fn: Callable[[Any], Any], + *, + output_dir: str | None = None, + overwrite: bool = False, + verbose: bool = True, + num_shards: int = _DEFAULT_KEYS_NUM_SHARDS, +) -> str: + """Scan an optimized dataset with ``key_fn`` and write ``keys/shard-*.parquet``. + + Use this to backfill a key sidecar for datasets produced without + ``optimize(..., key_fn=...)``. Returns the path to the ``keys/`` directory. + """ + from litdata.streaming.resolver import _resolve_dir + + resolved = _resolve_dir(input_dir) + if resolved.path is None or not os.path.isdir(resolved.path): + raise FileNotFoundError(f"Dataset directory not found: {input_dir}") + + dataset_dir = output_dir or resolved.path + out = keys_dir(dataset_dir) + if has_keys_index(dataset_dir) and not overwrite: + raise FileExistsError(f"Key index already exists under {dataset_dir}. Pass overwrite=True to replace it.") + + keys: list[str | int] = [] + indices: list[int] = [] + for key, index in iter_key_indexes(resolved.path, key_fn, verbose=verbose): + keys.append(key) + indices.append(index) + + write_keys_store(dataset_dir, keys, indices=indices, num_shards=num_shards) + + # Load+close before enrich: enrich rewrites index.json via write_keys_store, and + # Windows cannot os.replace a file that still has an open handle. + index_path = os.path.join(resolved.path, _INDEX_FILENAME) + with open(index_path, encoding="utf-8") as f: + index_json = json.load(f) + enrich_keys_with_chunks(dataset_dir, index_json) + return out diff --git a/tests/streaming/test_dataset_update.py b/tests/streaming/test_dataset_update.py new file mode 100644 index 000000000..d559d6768 --- /dev/null +++ b/tests/streaming/test_dataset_update.py @@ -0,0 +1,405 @@ +# Copyright The Lightning AI team. +# Licensed under the Apache License, Version 2.0. + +import json +import os + +import polars as pl +import pytest + +from litdata import StreamingDataset, build_keys_index, dataset_update, optimize +from litdata.constants import _INDEX_FILENAME, _KEYS_DIRNAME +from litdata.utilities.keys_index import ( + KeyIndex, + enrich_keys_with_chunks, + iter_key_indexes, + keys_dir, + save_keys, + save_rank_keys, + shard_path, + write_keys_store, +) + + +@pytest.fixture(autouse=True) +def _isolate_optimizer_cache(tmpdir, monkeypatch): + """Keep optimize() scratch dirs unique per test (xdist-safe). + + Without this, workers share ``/tmp/chunks`` and leftover ``*.bin`` / + ``*-index.json`` files from other tests fail ``_done`` cleanup checks and + can merge inconsistent configs into the dataset under test. + """ + monkeypatch.setenv("DATA_OPTIMIZER_CACHE_FOLDER", os.path.join(str(tmpdir), "opt_chunks")) + monkeypatch.setenv("DATA_OPTIMIZER_DATA_CACHE_FOLDER", os.path.join(str(tmpdir), "opt_data")) + + +def _fn(i: int) -> dict: + return {"id": f"item-{i}", "value": i} + + +def _key_fn(sample: dict) -> str: + return sample["id"] + + +def _int_fn(i: int) -> dict: + return {"id": i, "value": i * 10} + + +def _int_key_fn(sample: dict) -> int: + return sample["id"] + + +def test_keys_parquet_int64_roundtrip(tmpdir): + path = os.path.join(tmpdir, "keys.parquet") + keys = list(range(1000)) + save_keys(path, keys) + with KeyIndex(path) as idx: + assert len(idx) == 1000 + assert idx[0] == 0 + assert idx[999] == 999 + assert 42 in idx + assert idx.key_at(42) == 42 + assert idx.get(-1) is None + gidx, chunk_i, chunk_off = idx.resolve(42) + assert gidx == 42 + assert chunk_i == -1 # not enriched yet + + +def test_keys_parquet_utf8_roundtrip(tmpdir): + path = os.path.join(tmpdir, "keys.parquet") + keys = [f"id-{i}" for i in range(500)] + save_keys(path, keys) + df = pl.read_parquet(path) + assert df.columns == ["key", "index"] + with KeyIndex(path) as idx: + assert len(idx) == 500 + assert idx["id-7"] == 7 + assert idx.key_at(7) == "id-7" + + +def test_write_keys_store_default_shard_layout(tmpdir): + # Minimal index.json so keys config is recorded there (not a second manifest). + with open(os.path.join(tmpdir, _INDEX_FILENAME), "w", encoding="utf-8") as f: + json.dump({"chunks": [], "config": {}}, f) + + out = write_keys_store(str(tmpdir), ["a", "b", "c"], indices=[0, 1, 2]) + assert out == keys_dir(str(tmpdir)) + assert os.path.isfile(shard_path(str(tmpdir), 0)) + assert not os.path.isfile(os.path.join(out, "manifest.json")) + with open(os.path.join(tmpdir, _INDEX_FILENAME), encoding="utf-8") as f: + index = json.load(f) + assert index["keys"]["num_shards"] == 1 + assert index["keys"]["sharding"] == "none" + with KeyIndex(str(tmpdir)) as idx: + assert idx["b"] == 1 + + +def test_write_keys_store_multi_shard(tmpdir): + keys = [f"k-{i}" for i in range(100)] + write_keys_store(str(tmpdir), keys, num_shards=4) + for shard in range(4): + assert os.path.isfile(shard_path(str(tmpdir), shard)) + with KeyIndex(str(tmpdir)) as idx: + assert len(idx) == 100 + assert idx["k-7"] == 7 + assert idx.resolve_many(["k-1", "k-50", "k-99"])["k-99"][0] == 99 + + +def test_save_rank_keys_pairs_and_merge_remaps_indexes(tmpdir): + from litdata.utilities.keys_index import merge_rank_key_files + + # Rank 0: local indexes 0..2 + save_rank_keys( + os.path.join(tmpdir, "0.keys.parquet"), + [(0, "a"), (1, "b"), (2, "c")], + ) + # Rank 1: local indexes 0..1 + save_rank_keys( + os.path.join(tmpdir, "1.keys.parquet"), + [(0, "d"), (1, "e")], + ) + merged = merge_rank_key_files(str(tmpdir)) + assert merged is not None + assert merged == keys_dir(str(tmpdir)) + with KeyIndex(merged) as idx: + assert len(idx) == 5 + assert idx["a"] == 0 + assert idx["c"] == 2 + assert idx["d"] == 3 # remapped after rank 0 + assert idx["e"] == 4 + + +def test_enrich_keys_with_chunks(tmpdir): + write_keys_store(str(tmpdir), ["x", "y", "z"], indices=[0, 1, 2]) + index_json = { + "chunks": [ + {"filename": "chunk-0-0.bin", "chunk_bytes": 10, "chunk_size": 2}, + {"filename": "chunk-0-1.bin", "chunk_bytes": 10, "chunk_size": 1}, + ] + } + enrich_keys_with_chunks(str(tmpdir), index_json) + df = pl.read_parquet(shard_path(str(tmpdir), 0)) + assert "chunk_index" in df.columns + assert "chunk_offset" in df.columns + with KeyIndex(str(tmpdir)) as idx: + g0, c0, o0 = idx.resolve("x") + g2, c2, o2 = idx.resolve("z") + assert (g0, c0) == (0, 0) + assert (g2, c2) == (2, 1) + + +def test_optimize_writes_keys_store(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(10)), + output_dir=out, + chunk_size=4, + num_workers=2, + reorder_files=False, + key_fn=_key_fn, + ) + + keys_file = shard_path(out, 0) + assert os.path.isfile(keys_file) + assert os.path.isdir(os.path.join(out, _KEYS_DIRNAME)) + with open(os.path.join(out, _INDEX_FILENAME), encoding="utf-8") as f: + index = json.load(f) + assert index["keys"]["num_shards"] == 1 + df = pl.read_parquet(keys_file) + assert "key" in df.columns + assert "index" in df.columns + assert "chunk_index" in df.columns + assert df.height == 10 + + ds = StreamingDataset(out, shuffle=False) + assert len(ds) == 10 + sample = ds["item-3"] + assert sample["id"] == "item-3" + assert sample["value"] == 3 + + +def test_streaming_dataset_key_access_matches_int_index(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(8)), + output_dir=out, + chunk_size=3, + num_workers=2, + reorder_files=False, + key_fn=_key_fn, + ) + ds = StreamingDataset(out, shuffle=False) + with KeyIndex(out) as index: + for key in ("item-0", "item-4", "item-7"): + assert ds[key] == ds[index[key]] + + +def test_dataset_update_replaces_sample_by_key(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(12)), + output_dir=out, + chunk_size=5, + num_workers=2, + reorder_files=False, + key_fn=_key_fn, + ) + + with dataset_update(out) as update: + update["item-3"] = {"id": "item-3", "value": 999} + update["item-11"] = {"id": "item-11", "value": -1} + update.commit() + with pytest.raises(RuntimeError, match="after commit"): + update["item-0"] = {"id": "item-0", "value": 0} + + ds = StreamingDataset(out, shuffle=False) + assert ds["item-3"]["value"] == 999 + assert ds["item-11"]["value"] == -1 + assert ds["item-0"]["value"] == 0 + + +def test_dataset_update_without_commit_discards_changes(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(4)), + output_dir=out, + chunk_size=2, + num_workers=1, + key_fn=_key_fn, + ) + with dataset_update(out) as update: + update["item-1"] = {"id": "item-1", "value": 999} + # no commit() + assert StreamingDataset(out, shuffle=False)["item-1"]["value"] == 1 + + +def test_dataset_update_int_keys(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_int_fn, + inputs=list(range(8)), + output_dir=out, + chunk_size=3, + num_workers=2, + reorder_files=False, + key_fn=_int_key_fn, + ) + with dataset_update(out) as update: + update[3] = {"id": 3, "value": 123} + update.commit() + ds = StreamingDataset(out, shuffle=False) + # Int entity keys use get_by_key (ds[3] remains a global sample index). + assert ds.get_by_key(3)["value"] == 123 + + +def test_dataset_update_unknown_key_raises(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(4)), + output_dir=out, + chunk_size=2, + num_workers=1, + key_fn=_key_fn, + ) + with dataset_update(out) as update: + update["missing"] = {"id": "missing", "value": 0} + with pytest.raises(KeyError, match="Unknown dataset key"): + update.commit() + + +def test_dataset_update_requires_keys_file(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(4)), + output_dir=out, + chunk_size=2, + num_workers=1, + ) + with pytest.raises(FileNotFoundError, match="keys/"): + dataset_update(out) + + +def test_build_keys_index_backfills_sidecar(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(9)), + output_dir=out, + chunk_size=4, + num_workers=2, + reorder_files=False, + ) + assert not os.path.isfile(shard_path(out, 0)) + + pairs = list(iter_key_indexes(out, _key_fn, verbose=False)) + assert pairs[0] == ("item-0", 0) + assert pairs[-1] == ("item-8", 8) + + path = build_keys_index(out, _key_fn, verbose=False) + assert path == keys_dir(out) + df = pl.read_parquet(shard_path(out, 0)) + assert df.height == 9 + assert "chunk_index" in df.columns + + ds = StreamingDataset(out, shuffle=False) + assert ds["item-5"]["value"] == 5 + # Release chunk mmaps before in-place rewrite (required on Windows). + item_loader = getattr(getattr(ds, "cache", None), "_reader", None) + item_loader = getattr(item_loader, "_item_loader", None) if item_loader is not None else None + close_open = getattr(item_loader, "_close_open_chunk", None) + if callable(close_open): + close_open() + del ds + + with dataset_update(out) as update: + update["item-5"] = {"id": "item-5", "value": 50} + update.commit() + assert StreamingDataset(out, shuffle=False)["item-5"]["value"] == 50 + + with pytest.raises(FileExistsError, match="already exists"): + build_keys_index(out, _key_fn, verbose=False) + + build_keys_index(out, _key_fn, overwrite=True, verbose=False) + + +def test_duplicate_keys_rejected(tmpdir): + path = os.path.join(tmpdir, "keys.parquet") + with pytest.raises(ValueError, match="Duplicate key"): + save_keys(path, ["a", "b", "a"]) + + +def test_legacy_keys_parquet_still_readable(tmpdir): + out = os.path.join(tmpdir, "data") + os.makedirs(out) + save_keys(os.path.join(out, "keys.parquet"), ["legacy-a", "legacy-b"], indices=[0, 1]) + with KeyIndex(out) as idx: + assert idx["legacy-a"] == 0 + assert idx["legacy-b"] == 1 + + +def test_remote_key_index_scans_without_full_download(monkeypatch, tmpdir): + """Remote KeyIndex should scan cloud URIs with storage_options, not cache shards locally.""" + import polars as pl + + import litdata.streaming.fs_provider as fs_mod + + dataset_url = "s3://bucket/dataset" + with open(os.path.join(tmpdir, _INDEX_FILENAME), "w", encoding="utf-8") as f: + json.dump({"chunks": [], "config": {}, "keys": {"version": 1, "num_shards": 1, "sharding": "none"}}, f) + write_keys_store(str(tmpdir), ["a", "b"], indices=[0, 1], num_shards=1) + local_shard = shard_path(str(tmpdir), 0) + local_index = os.path.join(tmpdir, _INDEX_FILENAME) + + class _FakeFS: + def exists(self, path: str) -> bool: + return path.endswith("index.json") or path.endswith("shard-00000.parquet") + + def download_file(self, remote_path: str, local_path: str) -> None: + assert remote_path.endswith("index.json") + with open(local_index, encoding="utf-8") as src, open(local_path, "w", encoding="utf-8") as dst: + dst.write(src.read()) + + scanned: dict[str, object] = {} + real_scan = pl.scan_parquet + + def _tracking_scan(source, *args, **kwargs): + scanned["source"] = source + scanned["storage_options"] = kwargs.get("storage_options") + if isinstance(source, str) and source.startswith("s3://"): + source = local_shard + elif isinstance(source, list): + source = [local_shard if str(s).startswith("s3://") else s for s in source] + return real_scan(source, *args, **{k: v for k, v in kwargs.items() if k != "storage_options"}) + + monkeypatch.setattr(fs_mod, "_get_fs_provider", lambda url, opts=None: _FakeFS()) + monkeypatch.setattr(pl, "scan_parquet", _tracking_scan) + + storage_options = {"aws_region": "us-east-1"} + with KeyIndex(dataset_url, storage_options=storage_options) as idx: + assert idx["a"] == 0 + + source = scanned["source"] + assert (isinstance(source, str) and source.startswith("s3://")) or ( + isinstance(source, list) and str(source[0]).startswith("s3://") + ) + assert scanned["storage_options"] == storage_options + + +def test_streaming_dataset_missing_keys_raises_on_str_key(tmpdir): + out = os.path.join(tmpdir, "data") + optimize( + fn=_fn, + inputs=list(range(4)), + output_dir=out, + chunk_size=2, + num_workers=1, + ) + ds = StreamingDataset(out, shuffle=False) + with pytest.raises(KeyError, match="keys/"): + _ = ds["item-0"]