From cb577eb9b86d796e1f6b800735ab1b66e3d06976 Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 12:39:18 +0000 Subject: [PATCH 1/9] Add MultiJoinStreamingDataset design plan (V1/V2). Customer-facing proposal for co-partitioned multi-table streaming with an atomic version manifest, covering phase V1 aligned LitData chunks and phase V2 independently compacted column families. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 1531 +++++++++++++++++ 1 file changed, 1531 insertions(+) create mode 100644 .claude/plans/multi_join_streaming_dataset.md diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md new file mode 100644 index 00000000..44bf1cf2 --- /dev/null +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -0,0 +1,1531 @@ +# MultiJoinStreamingDataset + +## Customer-facing product and implementation plan + +Status: proposed design +Audience: customers, ML infrastructure engineers, LitData maintainers +Scope: independently versioned tables that describe the same training entity and must be joined while streaming + +## Executive summary + +`MultiJoinStreamingDataset` is intended for datasets in which: + +- Several tables or modalities describe the same logical entity. +- One or more large tables are stable and expensive to rebuild. +- Smaller tables change more often, including schema-level changes. +- Training still requires the throughput, distributed sampling, shuffle, prefetch, cache locality, and checkpoint-resume behavior of an optimized LitData dataset. + +The core design is **co-partitioned column families with an atomic version manifest**. + +The central constraint cannot be removed: if tables are not joined before training, they must remain aligned by key in the storage layout. Otherwise, every sampled key requires unrelated random reads from every table, turning prefetch into a distributed query-planning problem. + +The work is split into two deliberately different phases: + +- **Phase V1 — strict aligned LitData chunks.** Every table is a normal optimized LitData dataset with exactly the same logical chunk boundaries, item counts, and key order. `MultiJoinStreamingDataset` validates the layout and reuses `ParallelStreamingDataset`. This minimizes new read-path code and preserves existing LitData behavior. +- **Phase V2 — logical partitions with independently compacted column families.** Logical sampling buckets remain aligned, but physical files no longer need a one-to-one correspondence. Small tables can pack many logical partitions into larger Parquet or LitData objects, while large feature tables keep appropriately sized binary chunks. V2 introduces one shared sampler and a coordinated multi-table reader. + +The public read API is designed once and remains stable across both phases. V2 is primarily an internal storage and reader improvement. + +```mermaid +flowchart LR + canonicalKeys[Canonical ordered entity keys] + canonicalKeys --> logicalBuckets[Shared logical sampling buckets] + logicalBuckets --> features[Features table version] + logicalBuckets --> labels[Labels table version] + logicalBuckets --> tracking[Tracking table version] + activeManifest[Atomic join manifest] --> features + activeManifest --> labels + activeManifest --> tracking + features --> joinedReader[MultiJoinStreamingDataset] + labels --> joinedReader + tracking --> joinedReader + joinedReader --> training[Training batches] +``` + +## 1. Customer problem + +The current high-throughput LitData workflow materializes the complete training sample before or during `optimize()`. This is efficient at training time because one sample is stored in one streaming layout, but it couples the lifecycle of every source table: + +1. A small table changes. +2. The joined sample schema changes. +3. The complete optimized dataset is rebuilt. +4. Large unchanged feature tables are read, serialized, and uploaded again. + +`dataset_update` is complementary but does not solve the schema-change case. It is useful when a bounded set of existing samples can be replaced under a compatible optimized schema. If every entity in a table gains or loses columns, that table must be re-optimized. + +The requested behavior is: + +1. Optimize each logical table independently. +2. Re-optimize only the table whose values or schema changed. +3. Atomically publish a new combination of table versions. +4. Stream the selected versions as one training sample. +5. Preserve deterministic shuffle, DDP sharding, DataLoader workers, prefetch, caching, and exact resume. + +## 2. Why alignment is required + +LitData chunks are both: + +- Physical I/O units downloaded and cached by `BinaryReader`. +- Sampling buckets assigned to ranks and workers. + +The current sampling path has two stages: + +1. Chunks are assigned and shuffled across ranks and DataLoader workers. +2. Item positions are shuffled within each selected chunk. + +If separately optimized tables have different chunk boundaries or key order, the same seed does not make them align. + +Without alignment, one sampled key needs a location in every table: + +```text +entity-123: + features -> chunk 42, offset 731 + labels -> chunk 3, offset 18 + tracking -> chunk 91, offset 204 +``` + +Successive shuffled keys are likely to reference unrelated object combinations. The system then needs to: + +1. Resolve every key in every table. +2. Plan downloads across unrelated chunks. +3. Deduplicate and prioritize those downloads. +4. Coordinate cache admission and eviction. +5. Preserve distributed sampling and exact resume. +6. Join the decoded values. + +That is closer to implementing a distributed feature store or query engine than extending a streaming dataset. + +With alignment, one location is valid for every table: + +```text +entity-123 -> logical bucket 42, position 731 +``` + +The reader can fetch bucket group `42`, apply one item permutation, and read position `731` from every table. + +## 3. Terminology + +- **Entity key:** Stable identifier used to correlate tables, such as `sumer_play_id`. +- **Canonical key order:** Immutable ordered sequence of entity keys for one layout. +- **Logical item:** One training entity in the canonical order. +- **Table sample:** The value contributed by one table for one logical item. +- **Logical bucket:** Contiguous range of canonical item positions used as one sampling bucket. +- **Physical chunk or object:** File downloaded from local or object storage. +- **Table version:** Immutable physical representation of one table under one canonical layout. +- **Snapshot:** Immutable manifest selecting exactly one version of every active table. +- **Active manifest:** Small `join.json` document pointing to the active snapshot. +- **Layout ID:** Content-derived identifier for the canonical key order and logical bucket boundaries. + +## 4. Cardinality model + +Phase V1 uses a strict **one logical table sample per entity key** contract. + +This does not require every source table to contain one physical row per key. A table sample may contain: + +- One scalar or dictionary. +- A tensor. +- A list of rows. +- A variable-length NumPy array. +- An Arrow-like or serialized tabular bundle. +- An explicit empty collection when that table has no source rows for the entity. + +For example, a player-frame table may contribute a variable-length bundle containing every player-frame row for one play. The logical join remains one-to-one even though the source table is one-to-many. + +V1 rejects: + +- A missing logical table sample. +- Duplicate logical entity keys. +- Filtering that removes an entity. +- A generator that emits zero or multiple logical samples for one input key. + +An empty table contribution must be represented as an explicit empty value. This keeps every table positionally aligned. + +V2 can add a native offset-based representation for zero-to-many rows per key, but the canonical entity key and logical bucket alignment remain mandatory. + +## 5. Design goals + +### Correctness + +- Never silently combine values from different entity keys. +- Validate alignment before a table version can become active. +- Pin every training run to an immutable snapshot. +- Resume only against the same snapshot and sampling configuration. +- Fail closed when metadata is missing, incompatible, or ambiguous. + +### Performance + +- Keep the iterative training path sequential and prefetchable. +- Avoid per-sample key lookups during normal training. +- Reuse existing LitData downloader, cache, serializers, item loaders, and shuffle behavior in V1. +- Avoid reading or rewriting unchanged table versions. +- Preserve direct object-store access for S3, GCS, R2, and Lightning Storage paths. + +### Operability + +- Publish a table version immutably. +- Publish a new snapshot atomically only after validation. +- Allow instant rollback to a previous snapshot. +- Keep failed or incomplete versions unreachable. +- Expose clear diagnostics for alignment and manifest failures. + +### API quality + +- A simple root-path read API. +- Named table outputs rather than positional tuples. +- One shared set of sampling options. +- Explicit commit semantics for write operations. +- No need for users to configure `align_chunking`, `reorder_files`, child seeds, or child cache directories. +- The same user-facing read API in V1 and V2. + +## 6. Non-goals + +The following are not goals of V1: + +- Arbitrary SQL joins at training time. +- Joining independently optimized legacy datasets without a one-time aligned rebuild. +- Different key populations per required table. +- Independent per-table shuffle. +- Byte-based chunking for aligned table versions. +- Dynamic filtering in one table. +- Many-to-many joins. +- Replacing `dataset_update`. +- Replacing a feature store or query engine. +- Automatically changing the canonical key population without rebuilding all table versions. + +The full baked dataset remains the recommended default when all tables normally change together or when maximum simplicity and minimum object count are more important than independent table versioning. + +## 7. Stable public API + +The API below is the target public contract. Some advanced storage arguments become effective only in V2, but normal training code does not change. + +### 7.1 Create a canonical layout and initial snapshot + +```python +from litdata import MultiJoinWriter + +play_ids = load_canonical_play_ids() + +with MultiJoinWriter.create( + "s3://bucket/football-training", + keys=play_ids, + key_name="sumer_play_id", + chunk_size=2048, +) as writer: + writer.optimize_table( + "features", + fn=build_features, + version="features-2026-08-10", + num_workers=32, + num_nodes=8, + compression="zstd", + ) + writer.optimize_table( + "labels", + fn=build_labels, + version="labels-2026-08-10", + num_workers=8, + ) + snapshot = writer.commit() + +print(snapshot.id) +``` + +Semantics: + +- `keys` defines the canonical order once. +- Users should intentionally randomize or otherwise curate this order before layout creation when source order has structure. +- When `inputs` is omitted, `fn` receives each canonical key. +- `chunk_size` is the logical bucket size and is shared by every V1 table. +- Version paths are immutable. +- `commit()` validates all table versions and publishes the snapshot. +- Exiting the context without `commit()` does not publish changes. +- Calling `optimize_table()`, `remove_table()`, or `commit()` after a successful commit raises an error. + +### 7.2 Optimize from pre-grouped inputs + +Large pipelines may already produce one grouped input object per entity: + +```python +with MultiJoinWriter.create( + output_dir, + keys=play_ids, + key_name="sumer_play_id", + chunk_size=2048, +) as writer: + writer.optimize_table( + "tracking", + fn=encode_tracking_group, + inputs=tracking_groups, + input_key=lambda group: group.sumer_play_id, + ) + writer.commit() +``` + +Contract: + +- `inputs` must already follow the canonical order. +- `input_key` is checked against the expected canonical key at every position. +- The writer does not build an in-memory key-to-input map or silently reorder billions of records. +- A mismatch reports the table, global index, expected key, and actual key. +- The preparation pipeline may align data by joining against the canonical `(key, global_index)` mapping and ordering by `global_index`. + +### 7.3 Re-optimize one table after a schema change + +```python +from litdata import MultiJoinWriter + +with MultiJoinWriter.open( + "s3://bucket/football-training", + expected_snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", +) as writer: + writer.optimize_table( + "labels", + fn=build_labels_v2, + version="labels-2026-08-17", + num_workers=8, + ) + new_snapshot = writer.commit() +``` + +Only the new labels prefix and new snapshot metadata are written. The active features version is referenced unchanged. + +`expected_snapshot` provides optimistic concurrency protection. If another publisher changes the active snapshot first, commit fails rather than overwriting that change. + +### 7.4 Stream the active snapshot + +```python +from litdata import MultiJoinStreamingDataset, StreamingDataLoader + +dataset = MultiJoinStreamingDataset( + "s3://bucket/football-training", + tables=("features", "labels"), + shuffle=True, + seed=42, + drop_last=True, + transform=lambda parts: { + **parts["features"], + **parts["labels"], + }, + max_cache_size="200GB", + max_pre_download=4, +) + +loader = StreamingDataLoader( + dataset, + batch_size=64, + num_workers=8, +) +``` + +Default output without `transform`: + +```python +{ + "features": , + "labels": , +} +``` + +The table namespace is preserved by default. LitData does not implicitly merge dictionaries because duplicate field names would otherwise be ambiguous. + +### 7.5 Pin an exact snapshot + +```python +dataset = MultiJoinStreamingDataset( + "s3://bucket/football-training", + snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", + tables=("features", "labels"), + shuffle=True, + seed=42, +) +``` + +Useful properties: + +```python +dataset.snapshot_id +dataset.layout_id +dataset.table_versions +dataset.tables +``` + +### 7.6 Select only required tables + +```python +dataset = MultiJoinStreamingDataset( + root, + tables=("features", "labels"), +) +``` + +Selection changes which column families are downloaded, but it does not alter canonical sampling order. + +Unknown, duplicated, or inactive table names raise a clear error during construction. + +### 7.7 Per-table decoding options + +Sampling options cannot vary by table. Decoding-specific options may: + +```python +dataset = MultiJoinStreamingDataset( + root, + tables=("features", "labels"), + table_options={ + "features": {"encryption": feature_key}, + "labels": {"serializers": custom_serializers}, + }, +) +``` + +The following are always shared and cannot appear in `table_options`: + +- `shuffle` +- `seed` +- `drop_last` +- `subsample` +- epoch +- number of workers +- batch size +- distributed rank and world size + +### 7.8 Validate without training + +```python +from litdata import validate_multi_join + +report = validate_multi_join( + "s3://bucket/football-training", + snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", + deep=True, +) + +report.raise_for_errors() +``` + +Validation levels: + +- Construction always performs mandatory constant-size metadata validation. +- `deep=False` verifies manifests, table indexes, counts, layout IDs, and alignment roots. +- `deep=True` streams partition alignment metadata and verifies every ordered-key digest. +- Writer commit always performs the validation required to prove a newly published table matches the canonical layout. + +### 7.9 Keyed debugging access + +```python +sample = dataset.get_by_key("play-123") +``` + +Keyed lookup is intended for inspection, debugging, and bounded retrieval. It is not used by the iterative training path. + +For integer entity keys, `get_by_key()` remains explicit so integer positional indexing is unambiguous. + +## 8. Common storage and snapshot model + +The root is a versioned store: + +```text +football-training/ + join.json + snapshots/ + 01J4Y7M4VPH6V9B8P3N6K5W2HQ.json + 01J5A1QGMGT8P4J8Q1Z0AXEF3R.json + layouts/ + 6ce6f3.../ + index.json + keys/ + shard-00000.parquet + shard-00001.parquet + partitions.parquet + tables/ + features/ + features-2026-08-10/ + index.json + alignment.parquet + chunk-0-0.bin + ... + labels/ + labels-2026-08-10/ + index.json + alignment.parquet + chunk-0-0.bin + ... + labels-2026-08-17/ + index.json + alignment.parquet + chunk-0-0.bin + ... +``` + +### 8.1 Active manifest + +`join.json` is small and published last: + +```json +{ + "format": "litdata-multi-join", + "format_version": 1, + "active_snapshot": "01J5A1QGMGT8P4J8Q1Z0AXEF3R", + "updated_at": "2026-08-17T10:42:11Z" +} +``` + +### 8.2 Immutable snapshot + +```json +{ + "format": "litdata-multi-join-snapshot", + "format_version": 1, + "snapshot_id": "01J5A1QGMGT8P4J8Q1Z0AXEF3R", + "parent_snapshot_id": "01J4Y7M4VPH6V9B8P3N6K5W2HQ", + "created_at": "2026-08-17T10:42:10Z", + "layout": { + "id": "6ce6f3...", + "path": "layouts/6ce6f3...", + "key_name": "sumer_play_id", + "key_type": "string", + "length": 1000000000, + "chunk_size": 2048, + "num_chunks": 488282, + "alignment_root": "blake2b-256:..." + }, + "tables": { + "features": { + "version": "features-2026-08-10", + "path": "tables/features/features-2026-08-10", + "format": "litdata", + "layout_id": "6ce6f3...", + "alignment_root": "blake2b-256:...", + "schema_fingerprint": "sha256:..." + }, + "labels": { + "version": "labels-2026-08-17", + "path": "tables/labels/labels-2026-08-17", + "format": "litdata", + "layout_id": "6ce6f3...", + "alignment_root": "blake2b-256:...", + "schema_fingerprint": "sha256:..." + } + } +} +``` + +All paths are relative to the root in V1. Manifest parsing rejects: + +- Absolute paths. +- Parent traversal. +- A different URI scheme or bucket. +- Duplicate normalized table names. +- Unknown format versions. + +This keeps one snapshot within one trust and credential boundary. + +### 8.3 Canonical layout + +The layout is immutable and contains: + +- Canonical entity key type and order. +- Global item count. +- Logical bucket size. +- Tail bucket size. +- Per-bucket global start and stop positions. +- Per-bucket ordered-key digest. +- A digest root covering the complete ordered layout. +- A key index for debugging lookup and table re-optimization. + +The key store must be sharded and streamed. Creating or validating a billion-row layout must not require a Python dictionary containing every key. + +The canonical key order, rather than lexical key order, defines training positions. Key-index shards may be physically sorted or hash-partitioned for lookup as long as the stored `global_index`, `chunk_index`, and `chunk_offset` preserve the canonical order. + +### 8.4 Digest encoding + +Digest computation must be deterministic across Python versions and machines: + +1. Normalize the key to the supported integer or UTF-8 string representation. +2. Prefix each key with a type tag. +3. Prefix variable-length bytes with an explicit fixed-width length. +4. Hash keys in canonical order. +5. Include bucket index, start position, and item count in the bucket digest. +6. Hash the ordered bucket metadata into the layout alignment root. + +The initial algorithm is `BLAKE2b-256`. The manifest stores the algorithm and canonical encoding version so a future algorithm can coexist without ambiguity. + +Using `str(key)` concatenation without type and length framing is not acceptable because it can create ambiguous encodings. + +### 8.5 Table alignment metadata + +Every table version contains an `alignment.parquet` sidecar with: + +- `chunk_index` +- `global_start` +- `num_items` +- `ordered_key_digest` + +The table `index.json` gains a backward-compatible `multi_join` section: + +```json +{ + "multi_join": { + "format_version": 1, + "table": "labels", + "table_version": "labels-2026-08-17", + "layout_id": "6ce6f3...", + "length": 1000000000, + "num_chunks": 488282, + "alignment_root": "blake2b-256:..." + } +} +``` + +Per-chunk digests live in compact Parquet rather than expanding an already large `index.json`. Normal dataset construction compares constant-size roots. Deep validation reads the Parquet sidecars. + +## 9. Atomic publication and reader isolation + +### 9.1 Publication protocol + +`MultiJoinWriter.commit()` follows this order: + +1. Read and retain the expected active snapshot. +2. Write new table data to a unique immutable version path. +3. Upload all chunk objects. +4. Upload table alignment metadata. +5. Upload the table `index.json` and completion marker last. +6. Validate the table against the canonical layout. +7. Write a new immutable snapshot document. +8. Recheck the active snapshot or storage generation. +9. Atomically replace `join.json` with the new active snapshot. +10. Mark the writer committed and reject further mutation. + +If any operation fails before step 9, the active snapshot remains unchanged. Unreferenced objects are safe to garbage-collect later. + +### 9.2 Backend behavior + +- Local files use a temporary file, `fsync` where appropriate, and `os.replace`. +- S3 and R2 use immutable version objects and a conditional active-manifest write where supported. +- GCS uses generation-match preconditions. +- A backend without safe compare-and-swap must use a single-publisher lease or fail closed for concurrent publication. + +The storage abstraction must expose the precondition required for `expected_snapshot`; a read-then-unconditional-write sequence is not sufficient to prevent lost updates. + +### 9.3 Reader behavior + +`MultiJoinStreamingDataset` resolves the active snapshot exactly once during construction. It never polls `join.json` during iteration. + +Therefore: + +- Existing training jobs continue reading the old immutable table paths. +- New training jobs see the newly active snapshot. +- No job sees a mixture of old and new table versions. +- Checkpoint state records the snapshot ID. +- Resume against a different snapshot fails unless the caller explicitly opts to start a new data epoch. + +## 10. Phase V1 — strict aligned LitData chunks + +### 10.1 V1 objective + +Deliver a safe, production-testable implementation with the smallest possible change to LitData’s proven training read path. + +In V1: + +- Every table version is a standard optimized LitData dataset. +- One physical LitData chunk is one logical sampling bucket. +- Corresponding table chunks contain the same keys in the same order. +- Chunk byte sizes, schemas, compression, serializers, and payload types may differ. +- `MultiJoinStreamingDataset` is a validated named wrapper around `ParallelStreamingDataset`. + +### 10.2 V1 hard invariants + +For every active table: + +1. `layout_id` matches the snapshot layout. +2. Total logical length matches. +3. Number of chunks matches. +4. Chunk `i` has the canonical item count. +5. Chunk `i` has the canonical ordered-key digest. +6. The final partial chunk appears in the same position. +7. Every canonical key appears exactly once. +8. No additional key appears. +9. One input produces one logical table sample. +10. Table paths and versions are immutable. + +For the read configuration: + +1. All children receive the same `shuffle`. +2. All children receive the same `seed`. +3. All children receive the same epoch. +4. All children receive the same `drop_last`. +5. All children receive the same `subsample`. +6. All children see the same distributed environment. +7. All children receive the same DataLoader worker count and batch size. +8. A loaded state dict references the same snapshot and ordered table list. + +### 10.3 V1 write path + +`MultiJoinWriter.optimize_table()` internally calls the existing optimize pipeline with required safe settings: + +- `chunk_size` comes from the canonical layout. +- `align_chunking=True`. +- `reorder_files=False`. +- `keep_data_ordered=True`. +- Static ordered inputs. +- One output sample per input. +- Key capture enabled for alignment validation. + +The following are hidden or rejected: + +- `chunk_bytes` +- `weights` +- shared dynamic work queue +- `keep_data_ordered=False` +- `reorder_files=True` +- filtering with `None` +- variable-yield generators +- append into an existing table version +- overwrite of a published table version + +`align_chunking=True` is important because it makes logical chunk boundaries independent of the number of optimize workers or nodes. Workers receive complete item-count chunks, and rank-index merge order reconstructs the canonical sequence. + +A labels table may therefore be rebuilt with 8 workers while a feature table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. + +### 10.4 V1 read path + +Construction: + +1. Resolve the root through LitData’s normal path resolver. +2. Load and pin a snapshot. +3. Validate the snapshot and selected table metadata. +4. Create one `StreamingDataset` per selected table version. +5. Give every child identical sampling options. +6. Allocate a unique cache namespace per snapshot, table, and version. +7. Pass the children to `ParallelStreamingDataset`. +8. Adapt positional tuples to a named dictionary. + +Iteration: + +1. Existing `FullShuffle` computes the same chunk ordering for every child. +2. Existing worker assignment computes the same chunk intervals because chunk counts and item counts match. +3. Existing in-chunk shuffle computes the same item permutation. +4. Each child `BinaryReader` prefetches its corresponding chunks. +5. `ParallelStreamingDataset` pulls one aligned value from each child. +6. `MultiJoinStreamingDataset` creates the named table mapping. +7. The optional transform creates the final training sample. + +No key lookup occurs on this path. + +### 10.5 Why existing shuffle remains aligned in V1 + +Corresponding children have identical: + +- Chunk interval arrays. +- Number of chunks. +- Seed. +- Epoch. +- Chunk index. +- Worker and distributed topology. +- Batch and `drop_last` configuration. + +As a result, both chunk-to-worker assignment and within-chunk permutation are deterministic and equal across children. + +This must be proven with integration tests rather than assumed from unit-level seed equality. + +### 10.6 V1 resume behavior + +The state dict includes: + +- Snapshot ID. +- Layout ID. +- Ordered table names and versions. +- Existing child `StreamingDataset` states. +- Current epoch. +- Per-worker yielded counts. +- Transform RNG state inherited from `ParallelStreamingDataset`. + +On load: + +- A different snapshot is rejected. +- A different selected table set or order is rejected. +- Existing LitData checks still reject incompatible seed, shuffle, worker count, batch size, item loader, or distributed world size. +- `force_override_state_dict` remains an advanced escape hatch and emits a strong warning because repeated or skipped samples are possible. + +### 10.7 V1 cache and prefetch behavior + +Each child keeps its existing: + +- `BinaryReader` +- `PrepareChunksThread` +- Downloader +- Asynchronous remote prefetch +- Refcount and deletion logic + +The wrapper does not introduce another download engine. + +An explicit join cache root is namespaced: + +```text +cache/ + / + features/ + / + labels/ + / +``` + +This prevents same-named LitData chunk files from different tables from colliding. + +`max_cache_size` is defined as the aggregate user budget. V1 approximates this by assigning per-table budgets in proportion to average chunk bytes, with a documented minimum and an optional per-table override. + +Peak in-flight storage is approximately: + +```text +num_workers × max_pre_download × sum(corresponding table chunk bytes) +``` + +The slowest or largest member of a chunk group determines when that group is fully available. + +### 10.8 V1 limitations + +V1 intentionally accepts the following trade-offs: + +- Every table has the same number of physical chunk objects. +- Tiny tables may produce many small objects. +- Object GET count is approximately multiplied by the number of selected tables. +- A single `chunk_size` must balance feature chunk bytes, sampling diversity, and small-table object count. +- The wrapper relies on several independently operating child prefetchers. +- Aggregate cache enforcement is approximate because children retain independent cache managers. +- Required tables must have the same logical key population. +- Changing the canonical key population requires a new layout and rebuilding all tables. +- Native per-key zero-to-many row ranges are not yet represented; V1 bundles them into one table sample. + +These limitations are the primary motivation for V2. + +### 10.9 V1 error model + +Introduce specific public exceptions: + +- `MultiJoinError` +- `MultiJoinManifestError` +- `MultiJoinAlignmentError` +- `MultiJoinSnapshotMismatchError` +- `MultiJoinCommitConflictError` + +An alignment error includes: + +- Table name and version. +- Layout ID. +- Chunk or global index. +- Expected item count or key digest. +- Actual item count or key digest. +- Recommended remediation. + +No warning-only mode is provided for alignment failures. + +### 10.10 V1 implementation work + +#### Manifest and layout utilities + +Add `src/litdata/utilities/multi_join.py`: + +- Typed dataclasses for active manifest, snapshot, layout, table version, and validation report. +- Strict JSON parsing and format-version checks. +- Relative-path validation. +- Canonical key encoding and digest computation. +- Streaming layout creation. +- Snapshot loading and publication helpers. +- Sharded key iteration without a full Python dictionary. + +Do not add a heavy schema dependency solely for these small metadata models. + +#### Write API + +Add `src/litdata/processing/multi_join.py`: + +- `MultiJoinWriter.create` +- `MultiJoinWriter.open` +- `optimize_table` +- `remove_table` +- `validate` +- `commit` +- abort and cleanup bookkeeping + +Extend the optimize internals where necessary to: + +- Stream ordered `(global_index, key)` metadata. +- Compute per-chunk key digests. +- Avoid materializing all keys during alignment generation. +- Attach the backward-compatible `multi_join` section to `index.json`. +- Return a typed table-version result to the writer. + +Likely integration points: + +- `src/litdata/processing/functions.py` +- `src/litdata/processing/data_processor.py` +- `src/litdata/streaming/writer.py` +- `src/litdata/utilities/keys_index.py` + +#### Read API + +Add `src/litdata/streaming/multi_join.py`: + +- Manifest resolution. +- Mandatory alignment validation. +- Named child construction. +- Cache namespacing. +- Tuple-to-mapping transform adapter. +- Snapshot-aware state dict. +- Optional keyed debugging access. + +Reuse: + +- `src/litdata/streaming/parallel.py` +- `src/litdata/utilities/base.py` +- `src/litdata/streaming/dataset.py` +- `src/litdata/streaming/dataloader.py` +- `src/litdata/streaming/shuffle.py` +- `src/litdata/streaming/reader.py` + +#### Storage publication + +Extend `src/litdata/streaming/fs_provider.py` or a narrowly scoped manifest-storage abstraction with: + +- Read object metadata or generation. +- Conditional active-manifest write. +- Immutable JSON upload. +- Local atomic replacement. + +Training downloads continue to use `Downloader`; `FsProvider` must not be introduced into `PrepareChunksThread`. + +#### Public exports + +Update: + +- `src/litdata/__init__.py` +- Public API documentation. +- The LitData skill and reference docs. + +### 10.11 V1 correctness test matrix + +Add focused tests under: + +- `tests/processing/test_multi_join.py` +- `tests/streaming/test_multi_join.py` +- `tests/utilities/test_multi_join.py` + +Cover: + +- One, two, and many tables. +- Empty dataset rejection. +- Full and partial final chunks. +- Different table payload byte sizes. +- Different schemas and serializers. +- Different compression settings. +- Different optimize worker counts. +- Different optimize node counts. +- Stable boundaries with `align_chunking=True`. +- Key order mismatch. +- Missing key. +- Duplicate key. +- Extra key. +- `fn` returning `None`. +- Variable-yield generator rejection. +- Corrupt layout ID. +- Corrupt per-chunk digest. +- Missing table `index.json`. +- Incomplete table version. +- Manifest path traversal. +- Unsupported manifest format version. + +Read-path matrix: + +- `shuffle=False` and `shuffle=True`. +- Multiple epochs. +- Seeds. +- DataLoader workers `0`, `1`, `2`, and a higher stress count. +- DDP world sizes and ranks. +- Chunk intervals split across workers. +- `drop_last=True` and `False`. +- Different batch sizes. +- Common subsampling. +- Persistent workers. +- Early break. +- Complete and partial checkpoint resume. +- Transform RNG resume. +- Cache pressure and eviction. +- Asynchronous remote prefetch on and off. + +Snapshot tests: + +- A reader opened before commit continues on the old snapshot. +- A reader opened after commit sees the new snapshot. +- Features objects are not rewritten by a labels-only update. +- Rollback activates the previous immutable snapshot. +- Concurrent stale writer commit fails. +- Failure before active-manifest publication leaves the old snapshot active. + +Remote tests: + +- S3-compatible mocked backend. +- R2 path resolution. +- Lightning Storage connection path resolution. +- GCS manifest generation behavior where available. +- Direct object access rather than hand-reading a FUSE mount. + +### 10.12 V1 performance validation + +Benchmark against a baked single LitData dataset using the same: + +- Entity population and canonical order. +- Batch size. +- DataLoader workers. +- DDP topology. +- Cache disk. +- `max_pre_download`. +- Compression. +- Cold-cache and warm-cache conditions. + +Measure: + +- Time to first batch. +- Median and tail batch wait. +- Samples per second. +- GPU data-wait percentage. +- Object GET count. +- Bytes downloaded. +- Peak cache usage. +- CPU decode time. +- Cost and elapsed time to rebuild only the changed table. + +No production throughput claim should be made before this benchmark runs on the customer-shaped workload. Correctness and independent table replacement are hard release gates; throughput is a measured release gate agreed from the baseline. + +### 10.13 V1 completion criteria + +V1 is complete when: + +1. A two-table and a many-table dataset can be created through the public writer API. +2. Re-optimizing one table does not write under unchanged table-version prefixes. +3. A new table snapshot is atomically activated. +4. Old readers continue without observing mixed versions. +5. Every tested shuffle, worker, DDP, and resume configuration preserves key alignment. +6. Intentional misalignment fails before iteration. +7. The Lightning Storage end-to-end test is green. +8. Customer-shaped cold-cache benchmarks are documented. +9. API documentation includes creation, update, rollback, training, and troubleshooting. + +## 11. Phase V2 — logical partitions and independent physical compaction + +### 11.1 V2 objective + +Remove V1’s requirement that every logical sampling bucket be a separate physical object in every table, while retaining canonical key alignment and deterministic bucket sampling. + +V2 does **not** remove alignment. It separates: + +- Logical sampling layout. +- Table-specific physical storage layout. + +This permits small or schema-volatile tables to use an efficient physical representation without changing the training entity order. + +### 11.2 Why V2 is needed + +Assume the canonical bucket contains 2,048 entities: + +- A feature bucket may be 64–256 MB. +- A label bucket may be only tens or hundreds of KB. + +V1 writes one object for each in both tables. At very large scale, the labels table may have hundreds of thousands of tiny objects. This increases: + +- Object listing and metadata cost. +- GET request count. +- Time to first batch. +- Index size. +- Cache bookkeeping. +- Publication and garbage-collection overhead. + +V2 allows many label buckets to be packed into one appropriately sized object while preserving the original logical bucket boundaries for sampling. + +### 11.3 V2 storage model + +```text +football-training/ + join.json + snapshots/ + layouts/ + / + index.json + keys/ + partitions.parquet + tables/ + features/ + / + table.json + mapping.parquet + objects/ + pack-00000.bin + pack-00001.bin + labels/ + / + table.json + mapping.parquet + objects/ + pack-00000.parquet + pack-00001.parquet +``` + +Each table’s `mapping.parquet` maps canonical logical buckets to physical storage: + +- Logical bucket ID. +- Canonical global start and item count. +- Physical object path. +- Encoding or item-loader type. +- Byte, row-group, or row interval. +- Ordered-key digest. +- Optional per-key offset-vector location. + +One physical object may contain several consecutive logical buckets. + +V2 initially avoids splitting one logical bucket across many physical objects unless a single bucket exceeds the configured maximum. Supporting one object span per table per logical bucket keeps prefetch planning bounded. + +### 11.4 Heterogeneous column families + +V2 can support different physical formats behind one logical contract: + +- LitData binary for large tensors, images, audio, or nested Python structures. +- Parquet for schema-volatile tabular data. +- A compact offset-based binary representation for variable-length rows. + +Every format adapter must provide: + +- Metadata loading. +- Physical dependency resolution for a logical bucket. +- Prefetch request generation. +- Loading a logical item by position. +- Resource release. +- Format-specific schema fingerprint. + +The join sampler must not contain format-specific decode logic. + +### 11.5 V2 read engine + +V2 replaces the V1 `ParallelStreamingDataset` composition internally with a shared engine: + +1. Load the canonical logical partition intervals once. +2. Run chunk-to-rank and chunk-to-worker assignment once. +3. Generate one within-bucket item permutation. +4. Resolve each selected table’s physical dependencies for upcoming buckets. +5. Deduplicate dependencies when several logical buckets share a physical object. +6. Prefetch the complete dependency group. +7. Apply the same item position to every table loader. +8. Return the named table mapping. + +Potential internal components: + +- `MultiJoinChunksConfig` +- `MultiJoinShuffle` +- `MultiJoinReader` +- `PrepareJoinPartitionsThread` +- `ColumnFamilyLoader` +- `LitDataColumnFamilyLoader` +- `ParquetColumnFamilyLoader` + +Existing `Downloader` implementations, async remote fetch, serializers, and cache locking should be reused wherever possible. + +V2 must not use `FsProvider` on the training read path. + +### 11.6 V2 cache behavior + +The shared reader owns one aggregate cache budget: + +- Cache keys include snapshot, table, version, and physical object. +- One downloaded packed object can satisfy several upcoming logical buckets. +- Admission and eviction use the shared future-use schedule. +- Refcounts cover all workers and all logical buckets referencing an object. +- Prefetch depth is measured in logical bucket groups, not independently per table. + +This eliminates V1’s approximate per-child budget split. + +### 11.7 V2 writer API extension + +The public writer gains optional per-table storage configuration: + +```python +from litdata import MultiJoinWriter, TableStorage + +with MultiJoinWriter.open(root) as writer: + writer.optimize_table( + "labels", + fn=build_labels_v3, + storage=TableStorage( + format="parquet", + target_chunk_bytes="128MB", + pack_logical_chunks=True, + ), + ) + writer.commit() +``` + +Large feature table: + +```python +writer.optimize_table( + "features", + fn=build_features, + storage=TableStorage( + format="litdata", + target_chunk_bytes="256MB", + pack_logical_chunks=False, + ), +) +``` + +The normal `MultiJoinStreamingDataset` construction remains unchanged. + +### 11.8 V2 variable-cardinality column families + +After the shared reader and packed-object model are stable, V2 can represent source-table cardinality natively: + +- Each logical key position maps to a start and stop row offset. +- An empty range represents no rows for that table and key. +- A non-empty range returns one or many rows as the table contribution. +- Duplicate source rows are valid within that range. +- Duplicate logical entity keys remain invalid. + +This is similar to a compressed sparse row layout: + +```text +logical keys: [k0, k1, k2, k3] +row offsets: [0, 3, 3, 8, 9] +rows for k0: [0:3] +rows for k1: [3:3] -> empty +rows for k2: [3:8] +rows for k3: [8:9] +``` + +The entity sampler still chooses `k0..k3`; the table loader resolves the associated row slice without a hash join. + +This feature should be staged after V2’s one-envelope-per-key implementation because it changes collation, memory bounds, and schema semantics. + +### 11.9 V2 state and resume + +V2 stores one canonical sampling state rather than one child state per table: + +- Snapshot and layout ID. +- Epoch. +- Canonical bucket permutation. +- Worker assignment inputs. +- Current logical bucket. +- Consumed positions in the bucket. +- Shared transform RNG state. + +Table versions are data dependencies of the snapshot, not independent samplers. + +This reduces the possibility of child state divergence and makes the alignment contract explicit in code. + +### 11.10 V2 failure behavior + +- A missing physical object reports the table, version, logical bucket, and object. +- A corrupt mapping entry fails before returning a sample. +- A loader returning a different item count from the mapping fails the bucket. +- Repeated download failure follows existing LitData retry and timeout behavior. +- Partial table publication remains unreachable because snapshots reference only complete versions. +- Unsupported table formats fail during construction. + +### 11.11 V2 implementation work + +Add or evolve: + +- `src/litdata/streaming/multi_join.py` +- `src/litdata/streaming/multi_join_reader.py` +- `src/litdata/streaming/multi_join_config.py` +- `src/litdata/streaming/multi_join_shuffle.py` +- `src/litdata/processing/multi_join.py` +- Format-specific column-family loaders. + +Refactor the V1 implementation behind internal protocols so the public class and state schema can evolve without breaking customer code. + +### 11.12 V2 test matrix + +In addition to all V1 correctness tests: + +- Different physical object counts per table. +- Several logical buckets in one small-table object. +- Large object spanning upcoming worker buckets. +- Dependency deduplication. +- LitData binary plus Parquet in one snapshot. +- Row-group and byte-interval mapping. +- Shared aggregate cache eviction. +- Multiple workers referencing the same packed object. +- Resume in the middle of a packed object. +- Missing and corrupt mapping entries. +- Optional empty per-key row ranges. +- Variable-cardinality collation. +- Very large partition manifests. +- Bounded metadata memory usage. + +### 11.13 V2 performance validation + +Compare: + +1. Baked single dataset. +2. V1 strict aligned physical chunks. +3. V2 independently compacted column families. + +The expected V2 improvement is primarily: + +- Fewer object GETs for small tables. +- Smaller index and file-management overhead. +- Better cache reuse across adjacent logical buckets. +- Better control over each table’s physical object size. + +V2 must preserve V1’s correctness and independent-version benefits. + +### 11.14 V2 completion criteria + +V2 is complete when: + +1. V1 read code runs unchanged against the public API. +2. Different table object counts are supported under one canonical layout. +3. One shared sampler controls all tables. +4. Packed small-table objects reduce object GET count on the representative workload. +5. Aggregate cache usage is bounded by the configured budget. +6. Resume reproduces exactly the same logical entity sequence. +7. Mixed LitData and Parquet column families pass local, cloud, worker, and DDP tests. +8. V2 meets the customer-agreed throughput target relative to the baked baseline. + +## 12. V1 and V2 boundary + +The boundary is intentionally simple: + +### V1 aligns physical chunks + +```text +logical bucket 0 -> features chunk 0 -> labels chunk 0 +logical bucket 1 -> features chunk 1 -> labels chunk 1 +logical bucket 2 -> features chunk 2 -> labels chunk 2 +``` + +The existing LitData sampler and independent readers can be reused. + +### V2 aligns logical buckets + +```text +logical buckets 0..9: + features -> 10 feature objects + labels -> 1 packed label object +``` + +A new shared sampler and mapping-aware reader are required. + +### What V2 still does not do + +V2 does not accept arbitrary unrelated table sharding and repair it with per-sample random joins. Logical key order and bucket membership remain shared. This is the requirement that preserves predictable training I/O. + +## 13. Customer migration workflow + +### One-time migration + +1. Choose the entity key. +2. Produce a unique canonical key list. +3. Intentionally choose its training order. +4. Choose the logical bucket item count based primarily on feature bytes and desired bucket sampling. +5. Create the immutable layout. +6. Group each source table into one logical contribution per key. +7. Align each table input to canonical `global_index`. +8. Optimize stable large tables once. +9. Optimize schema-volatile tables. +10. Validate and publish the first snapshot. +11. Benchmark against the current baked dataset. + +### Small-table schema update + +1. Open the current store with its expected snapshot. +2. Rebuild only the changed table under a new version. +3. Validate keys, bucket counts, and digests. +4. Publish a new snapshot. +5. Start new training jobs on the new snapshot. +6. Leave existing jobs pinned to their original snapshot. + +### Rollback + +1. Select a known-good immutable snapshot. +2. Atomically update `join.json`. +3. Start new jobs. + +No table data is copied during rollback. + +### Canonical key population change + +Adding or removing entity keys changes: + +- Global indexes. +- Logical bucket membership. +- Chunk counts or tail. +- Ordered-key digests. + +V1 therefore creates a new layout and rebuilds every table. This is a deliberate correctness boundary. + +A future append-only layout extension could preserve complete existing buckets and add new buckets, but it requires separate sampling, snapshot, and resume semantics and is not included in the initial plan. + +## 14. Operational guidance + +### Version naming + +- User-provided readable versions are accepted. +- Internally generated immutable IDs prevent collisions. +- Reusing an existing table version path is rejected. + +### Retention and garbage collection + +Initial behavior: + +- Commits never delete table versions. +- Rollback remains possible while snapshots and versions are retained. +- Failed staging outputs are recorded as unreachable. + +Later garbage collection: + +1. Enumerate retained snapshots. +2. Mark referenced layouts and table versions. +3. Apply a minimum age safety window. +4. Delete only unreferenced immutable prefixes. +5. Support dry-run output before deletion. + +Garbage collection is never part of `commit()`. + +### Observability + +Expose structured diagnostics for: + +- Snapshot resolution. +- Selected table versions. +- Alignment validation duration. +- Per-table download bytes and latency. +- Prefetch queue depth. +- Cache hit and eviction counts. +- Bucket wait on the slowest table. +- Keyed debug lookup latency. + +Tracing should identify table and version without including customer entity keys by default. + +## 15. Risks and mitigations + +### Silent positional mis-join + +Risk: the most serious failure mode. + +Mitigation: + +- Canonical key registry. +- Per-bucket ordered-key digests. +- Alignment-root comparison. +- Validation before publication. +- Mandatory metadata validation on open. +- End-to-end tests that include the key in test payloads. + +### Tiny-object explosion in V1 + +Risk: schema-small tables produce many tiny chunks. + +Mitigation: + +- Measure V1 object count and GET overhead. +- Keep V1 as the minimal safe implementation. +- Implement V2 physical packing if the measured cost is material. + +### Cache multiplication + +Risk: independent V1 child readers each reserve and prefetch data. + +Mitigation: + +- Namespace caches. +- Interpret the public cache budget as aggregate. +- Allocate proportional child budgets. +- Move to one shared cache coordinator in V2. + +### Concurrent publisher race + +Risk: two updates overwrite each other’s active snapshot. + +Mitigation: + +- `expected_snapshot`. +- Conditional object-store publication. +- Immutable snapshots. +- Explicit commit-conflict error. + +### Metadata size at billion-row scale + +Risk: keys and per-chunk digests cannot be held in Python memory or expanded into large JSON documents. + +Mitigation: + +- Sharded Parquet key storage. +- Compact Parquet partition metadata. +- Streaming digest computation. +- Constant-size roots in normal open validation. +- Deep validation as a sequential metadata scan. + +### Optimize worker-count dependence + +Risk: independent table runs use different worker or node counts. + +Mitigation: + +- Force item-count chunking. +- Force ordered static inputs. +- Use `align_chunking=True`. +- Verify merged chunk counts and digests. +- Test all worker-count combinations. + +### Manifest or path injection + +Risk: a modified manifest points outside the dataset root. + +Mitigation: + +- Strict relative paths in V1. +- Reject traversal and cross-storage references. +- Validate format and field types. +- Store no credentials in manifests. + +## 16. Decision summary + +1. Alignment is required; arbitrary random keyed joins are not the product direction. +2. The entity key and canonical order are immutable within one layout. +3. One logical table contribution per entity is the V1 contract. +4. V1 uses standard LitData datasets with exactly aligned physical chunks. +5. V1 reuses `ParallelStreamingDataset` and existing read/prefetch infrastructure. +6. V2 retains logical alignment but permits independent physical compaction and formats. +7. V2 introduces one shared sampler and coordinated reader. +8. Table versions and snapshots are immutable. +9. `join.json` is published last and readers pin one snapshot. +10. Writer publication requires explicit `commit()`. +11. Training iteration performs no per-sample key lookup. +12. The public `MultiJoinStreamingDataset` API remains stable across V1 and V2. + +## 17. Recommended delivery order + +### Phase V1 + +1. Freeze manifest, layout, digest, and API specifications. +2. Implement metadata models and validation. +3. Implement scalable canonical layout creation. +4. Implement `MultiJoinWriter` with strict aligned optimize settings. +5. Implement atomic immutable snapshots. +6. Implement the validated `ParallelStreamingDataset` wrapper. +7. Add snapshot-aware state and resume. +8. Add local, remote, DDP, worker, cache, and failure tests. +9. Validate with Lightning Storage. +10. Run the representative customer benchmark. +11. Publish documentation and migration tooling. + +### Phase V2 + +1. Freeze logical-to-physical mapping format. +2. Define the column-family loader protocol. +3. Implement one canonical sampler. +4. Implement coordinated dependency prefetch. +5. Implement shared aggregate cache ownership. +6. Add LitData packed-object support. +7. Add Parquet column-family support. +8. Add optional offset-based variable-cardinality tables. +9. Run comparative V1, V2, and baked benchmarks. +10. Migrate internals while keeping the public read API unchanged. + From 148ecf17acb75710b252e5dcae4a38a76f9ff6db Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:39:31 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .claude/plans/multi_join_streaming_dataset.md | 115 +++++++++--------- 1 file changed, 57 insertions(+), 58 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 44bf1cf2..2992fba9 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -2,9 +2,9 @@ ## Customer-facing product and implementation plan -Status: proposed design -Audience: customers, ML infrastructure engineers, LitData maintainers -Scope: independently versioned tables that describe the same training entity and must be joined while streaming +Status: proposed design +Audience: customers, ML infrastructure engineers, LitData maintainers +Scope: independently versioned tables that describe the same training entity and must be joined while streaming ## Executive summary @@ -584,15 +584,15 @@ Per-chunk digests live in compact Parquet rather than expanding an already large `MultiJoinWriter.commit()` follows this order: -1. Read and retain the expected active snapshot. -2. Write new table data to a unique immutable version path. -3. Upload all chunk objects. -4. Upload table alignment metadata. -5. Upload the table `index.json` and completion marker last. -6. Validate the table against the canonical layout. -7. Write a new immutable snapshot document. -8. Recheck the active snapshot or storage generation. -9. Atomically replace `join.json` with the new active snapshot. +01. Read and retain the expected active snapshot. +02. Write new table data to a unique immutable version path. +03. Upload all chunk objects. +04. Upload table alignment metadata. +05. Upload the table `index.json` and completion marker last. +06. Validate the table against the canonical layout. +07. Write a new immutable snapshot document. +08. Recheck the active snapshot or storage generation. +09. Atomically replace `join.json` with the new active snapshot. 10. Mark the writer committed and reject further mutation. If any operation fails before step 9, the active snapshot remains unchanged. Unreferenced objects are safe to garbage-collect later. @@ -636,15 +636,15 @@ In V1: For every active table: -1. `layout_id` matches the snapshot layout. -2. Total logical length matches. -3. Number of chunks matches. -4. Chunk `i` has the canonical item count. -5. Chunk `i` has the canonical ordered-key digest. -6. The final partial chunk appears in the same position. -7. Every canonical key appears exactly once. -8. No additional key appears. -9. One input produces one logical table sample. +01. `layout_id` matches the snapshot layout. +02. Total logical length matches. +03. Number of chunks matches. +04. Chunk `i` has the canonical item count. +05. Chunk `i` has the canonical ordered-key digest. +06. The final partial chunk appears in the same position. +07. Every canonical key appears exactly once. +08. No additional key appears. +09. One input produces one logical table sample. 10. Table paths and versions are immutable. For the read configuration: @@ -1320,15 +1320,15 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam ### One-time migration -1. Choose the entity key. -2. Produce a unique canonical key list. -3. Intentionally choose its training order. -4. Choose the logical bucket item count based primarily on feature bytes and desired bucket sampling. -5. Create the immutable layout. -6. Group each source table into one logical contribution per key. -7. Align each table input to canonical `global_index`. -8. Optimize stable large tables once. -9. Optimize schema-volatile tables. +01. Choose the entity key. +02. Produce a unique canonical key list. +03. Intentionally choose its training order. +04. Choose the logical bucket item count based primarily on feature bytes and desired bucket sampling. +05. Create the immutable layout. +06. Group each source table into one logical contribution per key. +07. Align each table input to canonical `global_index`. +08. Optimize stable large tables once. +09. Optimize schema-volatile tables. 10. Validate and publish the first snapshot. 11. Benchmark against the current baked dataset. @@ -1487,15 +1487,15 @@ Mitigation: ## 16. Decision summary -1. Alignment is required; arbitrary random keyed joins are not the product direction. -2. The entity key and canonical order are immutable within one layout. -3. One logical table contribution per entity is the V1 contract. -4. V1 uses standard LitData datasets with exactly aligned physical chunks. -5. V1 reuses `ParallelStreamingDataset` and existing read/prefetch infrastructure. -6. V2 retains logical alignment but permits independent physical compaction and formats. -7. V2 introduces one shared sampler and coordinated reader. -8. Table versions and snapshots are immutable. -9. `join.json` is published last and readers pin one snapshot. +01. Alignment is required; arbitrary random keyed joins are not the product direction. +02. The entity key and canonical order are immutable within one layout. +03. One logical table contribution per entity is the V1 contract. +04. V1 uses standard LitData datasets with exactly aligned physical chunks. +05. V1 reuses `ParallelStreamingDataset` and existing read/prefetch infrastructure. +06. V2 retains logical alignment but permits independent physical compaction and formats. +07. V2 introduces one shared sampler and coordinated reader. +08. Table versions and snapshots are immutable. +09. `join.json` is published last and readers pin one snapshot. 10. Writer publication requires explicit `commit()`. 11. Training iteration performs no per-sample key lookup. 12. The public `MultiJoinStreamingDataset` API remains stable across V1 and V2. @@ -1504,28 +1504,27 @@ Mitigation: ### Phase V1 -1. Freeze manifest, layout, digest, and API specifications. -2. Implement metadata models and validation. -3. Implement scalable canonical layout creation. -4. Implement `MultiJoinWriter` with strict aligned optimize settings. -5. Implement atomic immutable snapshots. -6. Implement the validated `ParallelStreamingDataset` wrapper. -7. Add snapshot-aware state and resume. -8. Add local, remote, DDP, worker, cache, and failure tests. -9. Validate with Lightning Storage. +01. Freeze manifest, layout, digest, and API specifications. +02. Implement metadata models and validation. +03. Implement scalable canonical layout creation. +04. Implement `MultiJoinWriter` with strict aligned optimize settings. +05. Implement atomic immutable snapshots. +06. Implement the validated `ParallelStreamingDataset` wrapper. +07. Add snapshot-aware state and resume. +08. Add local, remote, DDP, worker, cache, and failure tests. +09. Validate with Lightning Storage. 10. Run the representative customer benchmark. 11. Publish documentation and migration tooling. ### Phase V2 -1. Freeze logical-to-physical mapping format. -2. Define the column-family loader protocol. -3. Implement one canonical sampler. -4. Implement coordinated dependency prefetch. -5. Implement shared aggregate cache ownership. -6. Add LitData packed-object support. -7. Add Parquet column-family support. -8. Add optional offset-based variable-cardinality tables. -9. Run comparative V1, V2, and baked benchmarks. +01. Freeze logical-to-physical mapping format. +02. Define the column-family loader protocol. +03. Implement one canonical sampler. +04. Implement coordinated dependency prefetch. +05. Implement shared aggregate cache ownership. +06. Add LitData packed-object support. +07. Add Parquet column-family support. +08. Add optional offset-based variable-cardinality tables. +09. Run comparative V1, V2, and baked benchmarks. 10. Migrate internals while keeping the public read API unchanged. - From af9a613ffe42b772f24c3363707e298c5a74ca6e Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 12:41:16 +0000 Subject: [PATCH 3/9] Make MultiJoin plan domain-agnostic. Replace customer-specific entity names, paths, and examples with generic multi-table terminology so the design can be shared publicly. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 92 +++++++++---------- 1 file changed, 46 insertions(+), 46 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 2992fba9..1d0aaac0 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -1,9 +1,9 @@ # MultiJoinStreamingDataset -## Customer-facing product and implementation plan +## Product and implementation plan Status: proposed design -Audience: customers, ML infrastructure engineers, LitData maintainers +Audience: LitData users, ML infrastructure engineers, LitData maintainers Scope: independently versioned tables that describe the same training entity and must be joined while streaming ## Executive summary @@ -32,17 +32,17 @@ flowchart LR canonicalKeys --> logicalBuckets[Shared logical sampling buckets] logicalBuckets --> features[Features table version] logicalBuckets --> labels[Labels table version] - logicalBuckets --> tracking[Tracking table version] + logicalBuckets --> metadata[Metadata table version] activeManifest[Atomic join manifest] --> features activeManifest --> labels - activeManifest --> tracking + activeManifest --> metadata features --> joinedReader[MultiJoinStreamingDataset] labels --> joinedReader - tracking --> joinedReader + metadata --> joinedReader joinedReader --> training[Training batches] ``` -## 1. Customer problem +## 1. Problem The current high-throughput LitData workflow materializes the complete training sample before or during `optimize()`. This is efficient at training time because one sample is stored in one streaming layout, but it couples the lifecycle of every source table: @@ -53,7 +53,7 @@ The current high-throughput LitData workflow materializes the complete training `dataset_update` is complementary but does not solve the schema-change case. It is useful when a bounded set of existing samples can be replaced under a compatible optimized schema. If every entity in a table gains or loses columns, that table must be re-optimized. -The requested behavior is: +The target behavior is: 1. Optimize each logical table independently. 2. Re-optimize only the table whose values or schema changed. @@ -81,7 +81,7 @@ Without alignment, one sampled key needs a location in every table: entity-123: features -> chunk 42, offset 731 labels -> chunk 3, offset 18 - tracking -> chunk 91, offset 204 + metadata -> chunk 91, offset 204 ``` Successive shuffled keys are likely to reference unrelated object combinations. The system then needs to: @@ -105,7 +105,7 @@ The reader can fetch bucket group `42`, apply one item permutation, and read pos ## 3. Terminology -- **Entity key:** Stable identifier used to correlate tables, such as `sumer_play_id`. +- **Entity key:** Stable identifier used to correlate tables, such as `entity_id`. - **Canonical key order:** Immutable ordered sequence of entity keys for one layout. - **Logical item:** One training entity in the canonical order. - **Table sample:** The value contributed by one table for one logical item. @@ -129,7 +129,7 @@ This does not require every source table to contain one physical row per key. A - An Arrow-like or serialized tabular bundle. - An explicit empty collection when that table has no source rows for the entity. -For example, a player-frame table may contribute a variable-length bundle containing every player-frame row for one play. The logical join remains one-to-one even though the source table is one-to-many. +For example, a nested events table may contribute a variable-length bundle containing every event row for one parent entity. The logical join remains one-to-one even though the source table is one-to-many. V1 rejects: @@ -203,18 +203,18 @@ The API below is the target public contract. Some advanced storage arguments bec ```python from litdata import MultiJoinWriter -play_ids = load_canonical_play_ids() +entity_ids = load_canonical_entity_ids() with MultiJoinWriter.create( - "s3://bucket/football-training", - keys=play_ids, - key_name="sumer_play_id", + "s3://bucket/multi-join-dataset", + keys=entity_ids, + key_name="entity_id", chunk_size=2048, ) as writer: writer.optimize_table( "features", fn=build_features, - version="features-2026-08-10", + version="features-v1", num_workers=32, num_nodes=8, compression="zstd", @@ -222,7 +222,7 @@ with MultiJoinWriter.create( writer.optimize_table( "labels", fn=build_labels, - version="labels-2026-08-10", + version="labels-v1", num_workers=8, ) snapshot = writer.commit() @@ -248,15 +248,15 @@ Large pipelines may already produce one grouped input object per entity: ```python with MultiJoinWriter.create( output_dir, - keys=play_ids, - key_name="sumer_play_id", + keys=entity_ids, + key_name="entity_id", chunk_size=2048, ) as writer: writer.optimize_table( - "tracking", - fn=encode_tracking_group, - inputs=tracking_groups, - input_key=lambda group: group.sumer_play_id, + "metadata", + fn=encode_metadata_group, + inputs=metadata_groups, + input_key=lambda group: group.entity_id, ) writer.commit() ``` @@ -275,13 +275,13 @@ Contract: from litdata import MultiJoinWriter with MultiJoinWriter.open( - "s3://bucket/football-training", + "s3://bucket/multi-join-dataset", expected_snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", ) as writer: writer.optimize_table( "labels", fn=build_labels_v2, - version="labels-2026-08-17", + version="labels-v2", num_workers=8, ) new_snapshot = writer.commit() @@ -297,7 +297,7 @@ Only the new labels prefix and new snapshot metadata are written. The active fea from litdata import MultiJoinStreamingDataset, StreamingDataLoader dataset = MultiJoinStreamingDataset( - "s3://bucket/football-training", + "s3://bucket/multi-join-dataset", tables=("features", "labels"), shuffle=True, seed=42, @@ -332,7 +332,7 @@ The table namespace is preserved by default. LitData does not implicitly merge d ```python dataset = MultiJoinStreamingDataset( - "s3://bucket/football-training", + "s3://bucket/multi-join-dataset", snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", tables=("features", "labels"), shuffle=True, @@ -394,7 +394,7 @@ The following are always shared and cannot appear in `table_options`: from litdata import validate_multi_join report = validate_multi_join( - "s3://bucket/football-training", + "s3://bucket/multi-join-dataset", snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", deep=True, ) @@ -412,7 +412,7 @@ Validation levels: ### 7.9 Keyed debugging access ```python -sample = dataset.get_by_key("play-123") +sample = dataset.get_by_key("entity-123") ``` Keyed lookup is intended for inspection, debugging, and bounded retrieval. It is not used by the iterative training path. @@ -424,7 +424,7 @@ For integer entity keys, `get_by_key()` remains explicit so integer positional i The root is a versioned store: ```text -football-training/ +multi-join-dataset/ join.json snapshots/ 01J4Y7M4VPH6V9B8P3N6K5W2HQ.json @@ -438,18 +438,18 @@ football-training/ partitions.parquet tables/ features/ - features-2026-08-10/ + features-v1/ index.json alignment.parquet chunk-0-0.bin ... labels/ - labels-2026-08-10/ + labels-v1/ index.json alignment.parquet chunk-0-0.bin ... - labels-2026-08-17/ + labels-v2/ index.json alignment.parquet chunk-0-0.bin @@ -481,7 +481,7 @@ football-training/ "layout": { "id": "6ce6f3...", "path": "layouts/6ce6f3...", - "key_name": "sumer_play_id", + "key_name": "entity_id", "key_type": "string", "length": 1000000000, "chunk_size": 2048, @@ -490,16 +490,16 @@ football-training/ }, "tables": { "features": { - "version": "features-2026-08-10", - "path": "tables/features/features-2026-08-10", + "version": "features-v1", + "path": "tables/features/features-v1", "format": "litdata", "layout_id": "6ce6f3...", "alignment_root": "blake2b-256:...", "schema_fingerprint": "sha256:..." }, "labels": { - "version": "labels-2026-08-17", - "path": "tables/labels/labels-2026-08-17", + "version": "labels-v2", + "path": "tables/labels/labels-v2", "format": "litdata", "layout_id": "6ce6f3...", "alignment_root": "blake2b-256:...", @@ -567,7 +567,7 @@ The table `index.json` gains a backward-compatible `multi_join` section: "multi_join": { "format_version": 1, "table": "labels", - "table_version": "labels-2026-08-17", + "table_version": "labels-v2", "layout_id": "6ce6f3...", "length": 1000000000, "num_chunks": 488282, @@ -993,7 +993,7 @@ Measure: - CPU decode time. - Cost and elapsed time to rebuild only the changed table. -No production throughput claim should be made before this benchmark runs on the customer-shaped workload. Correctness and independent table replacement are hard release gates; throughput is a measured release gate agreed from the baseline. +No production throughput claim should be made before this benchmark runs on a representative multi-table workload. Correctness and independent table replacement are hard release gates; throughput is a measured release gate agreed against the baseline. ### 10.13 V1 completion criteria @@ -1006,7 +1006,7 @@ V1 is complete when: 5. Every tested shuffle, worker, DDP, and resume configuration preserves key alignment. 6. Intentional misalignment fails before iteration. 7. The Lightning Storage end-to-end test is green. -8. Customer-shaped cold-cache benchmarks are documented. +8. Representative cold-cache benchmarks are documented. 9. API documentation includes creation, update, rollback, training, and troubleshooting. ## 11. Phase V2 — logical partitions and independent physical compaction @@ -1043,7 +1043,7 @@ V2 allows many label buckets to be packed into one appropriately sized object wh ### 11.3 V2 storage model ```text -football-training/ +multi-join-dataset/ join.json snapshots/ layouts/ @@ -1237,7 +1237,7 @@ Add or evolve: - `src/litdata/processing/multi_join.py` - Format-specific column-family loaders. -Refactor the V1 implementation behind internal protocols so the public class and state schema can evolve without breaking customer code. +Refactor the V1 implementation behind internal protocols so the public class and state schema can evolve without breaking user code. ### 11.12 V2 test matrix @@ -1286,7 +1286,7 @@ V2 is complete when: 5. Aggregate cache usage is bounded by the configured budget. 6. Resume reproduces exactly the same logical entity sequence. 7. Mixed LitData and Parquet column families pass local, cloud, worker, and DDP tests. -8. V2 meets the customer-agreed throughput target relative to the baked baseline. +8. V2 meets the agreed throughput target relative to the baked baseline. ## 12. V1 and V2 boundary @@ -1316,7 +1316,7 @@ A new shared sampler and mapping-aware reader are required. V2 does not accept arbitrary unrelated table sharding and repair it with per-sample random joins. Logical key order and bucket membership remain shared. This is the requirement that preserves predictable training I/O. -## 13. Customer migration workflow +## 13. Migration workflow ### One-time migration @@ -1401,7 +1401,7 @@ Expose structured diagnostics for: - Bucket wait on the slowest table. - Keyed debug lookup latency. -Tracing should identify table and version without including customer entity keys by default. +Tracing should identify table and version without including entity keys by default. ## 15. Risks and mitigations @@ -1513,7 +1513,7 @@ Mitigation: 07. Add snapshot-aware state and resume. 08. Add local, remote, DDP, worker, cache, and failure tests. 09. Validate with Lightning Storage. -10. Run the representative customer benchmark. +10. Run the representative multi-table benchmark. 11. Publish documentation and migration tooling. ### Phase V2 From b712d9e096a73846e22711f5fe769bbbca868393 Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 12:45:30 +0000 Subject: [PATCH 4/9] Clarify snapshot IDs and use generic table_N examples. Snapshot IDs are human-readable publish names, not chunk content hashes. Examples now use table_0 / table_1 / table_2 throughout. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 184 +++++++++--------- 1 file changed, 96 insertions(+), 88 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 1d0aaac0..74c8931e 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -22,7 +22,7 @@ The central constraint cannot be removed: if tables are not joined before traini The work is split into two deliberately different phases: - **Phase V1 — strict aligned LitData chunks.** Every table is a normal optimized LitData dataset with exactly the same logical chunk boundaries, item counts, and key order. `MultiJoinStreamingDataset` validates the layout and reuses `ParallelStreamingDataset`. This minimizes new read-path code and preserves existing LitData behavior. -- **Phase V2 — logical partitions with independently compacted column families.** Logical sampling buckets remain aligned, but physical files no longer need a one-to-one correspondence. Small tables can pack many logical partitions into larger Parquet or LitData objects, while large feature tables keep appropriately sized binary chunks. V2 introduces one shared sampler and a coordinated multi-table reader. +- **Phase V2 — logical partitions with independently compacted column families.** Logical sampling buckets remain aligned, but physical files no longer need a one-to-one correspondence. Small tables can pack many logical partitions into larger Parquet or LitData objects, while large tables keep appropriately sized binary chunks. V2 introduces one shared sampler and a coordinated multi-table reader. The public read API is designed once and remains stable across both phases. V2 is primarily an internal storage and reader improvement. @@ -30,15 +30,15 @@ The public read API is designed once and remains stable across both phases. V2 i flowchart LR canonicalKeys[Canonical ordered entity keys] canonicalKeys --> logicalBuckets[Shared logical sampling buckets] - logicalBuckets --> features[Features table version] - logicalBuckets --> labels[Labels table version] - logicalBuckets --> metadata[Metadata table version] - activeManifest[Atomic join manifest] --> features - activeManifest --> labels - activeManifest --> metadata - features --> joinedReader[MultiJoinStreamingDataset] - labels --> joinedReader - metadata --> joinedReader + logicalBuckets --> t0[table_0 version] + logicalBuckets --> t1[table_1 version] + logicalBuckets --> t2[table_2 version] + activeManifest[Atomic join manifest] --> t0 + activeManifest --> t1 + activeManifest --> t2 + t0 --> joinedReader[MultiJoinStreamingDataset] + t1 --> joinedReader + t2 --> joinedReader joinedReader --> training[Training batches] ``` @@ -49,7 +49,7 @@ The current high-throughput LitData workflow materializes the complete training 1. A small table changes. 2. The joined sample schema changes. 3. The complete optimized dataset is rebuilt. -4. Large unchanged feature tables are read, serialized, and uploaded again. +4. Large unchanged tables are read, serialized, and uploaded again. `dataset_update` is complementary but does not solve the schema-change case. It is useful when a bounded set of existing samples can be replaced under a compatible optimized schema. If every entity in a table gains or loses columns, that table must be re-optimized. @@ -79,9 +79,9 @@ Without alignment, one sampled key needs a location in every table: ```text entity-123: - features -> chunk 42, offset 731 - labels -> chunk 3, offset 18 - metadata -> chunk 91, offset 204 + table_0 -> chunk 42, offset 731 + table_1 -> chunk 3, offset 18 + table_2 -> chunk 91, offset 204 ``` Successive shuffled keys are likely to reference unrelated object combinations. The system then needs to: @@ -212,17 +212,17 @@ with MultiJoinWriter.create( chunk_size=2048, ) as writer: writer.optimize_table( - "features", - fn=build_features, - version="features-v1", + "table_0", + fn=build_table_0, + version="v1", num_workers=32, num_nodes=8, compression="zstd", ) writer.optimize_table( - "labels", - fn=build_labels, - version="labels-v1", + "table_1", + fn=build_table_1, + version="v1", num_workers=8, ) snapshot = writer.commit() @@ -253,9 +253,9 @@ with MultiJoinWriter.create( chunk_size=2048, ) as writer: writer.optimize_table( - "metadata", - fn=encode_metadata_group, - inputs=metadata_groups, + "table_2", + fn=encode_table_2_group, + inputs=table_2_groups, input_key=lambda group: group.entity_id, ) writer.commit() @@ -276,20 +276,20 @@ from litdata import MultiJoinWriter with MultiJoinWriter.open( "s3://bucket/multi-join-dataset", - expected_snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", + expected_snapshot="snap-baseline", ) as writer: writer.optimize_table( - "labels", - fn=build_labels_v2, - version="labels-v2", + "table_1", + fn=build_table_1_v2, + version="v2", num_workers=8, ) new_snapshot = writer.commit() ``` -Only the new labels prefix and new snapshot metadata are written. The active features version is referenced unchanged. +Only the new `table_1` version prefix and a new snapshot document are written. The active `table_0` version is referenced unchanged. -`expected_snapshot` provides optimistic concurrency protection. If another publisher changes the active snapshot first, commit fails rather than overwriting that change. +`expected_snapshot` is the **name** of the currently active snapshot (for example `"snap-baseline"`), not a hash of chunk bytes. Writers pass it as an optimistic concurrency guard: if another publisher advances the active snapshot first, commit fails rather than overwriting that change. Snapshot IDs are opaque human-readable identifiers assigned at publish time; LitData does not compute content hashes of chunk payloads to form them. ### 7.4 Stream the active snapshot @@ -298,13 +298,13 @@ from litdata import MultiJoinStreamingDataset, StreamingDataLoader dataset = MultiJoinStreamingDataset( "s3://bucket/multi-join-dataset", - tables=("features", "labels"), + tables=("table_0", "table_1"), shuffle=True, seed=42, drop_last=True, transform=lambda parts: { - **parts["features"], - **parts["labels"], + **parts["table_0"], + **parts["table_1"], }, max_cache_size="200GB", max_pre_download=4, @@ -321,8 +321,8 @@ Default output without `transform`: ```python { - "features": , - "labels": , + "table_0": , + "table_1": , } ``` @@ -333,8 +333,8 @@ The table namespace is preserved by default. LitData does not implicitly merge d ```python dataset = MultiJoinStreamingDataset( "s3://bucket/multi-join-dataset", - snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", - tables=("features", "labels"), + snapshot="snap-baseline", + tables=("table_0", "table_1"), shuffle=True, seed=42, ) @@ -354,7 +354,7 @@ dataset.tables ```python dataset = MultiJoinStreamingDataset( root, - tables=("features", "labels"), + tables=("table_0", "table_1"), ) ``` @@ -369,10 +369,10 @@ Sampling options cannot vary by table. Decoding-specific options may: ```python dataset = MultiJoinStreamingDataset( root, - tables=("features", "labels"), + tables=("table_0", "table_1"), table_options={ - "features": {"encryption": feature_key}, - "labels": {"serializers": custom_serializers}, + "table_0": {"encryption": table_0_key}, + "table_1": {"serializers": custom_serializers}, }, ) ``` @@ -395,7 +395,7 @@ from litdata import validate_multi_join report = validate_multi_join( "s3://bucket/multi-join-dataset", - snapshot="01J4Y7M4VPH6V9B8P3N6K5W2HQ", + snapshot="snap-baseline", deep=True, ) @@ -427,29 +427,29 @@ The root is a versioned store: multi-join-dataset/ join.json snapshots/ - 01J4Y7M4VPH6V9B8P3N6K5W2HQ.json - 01J5A1QGMGT8P4J8Q1Z0AXEF3R.json + snap-baseline.json + snap-table1-v2.json layouts/ - 6ce6f3.../ + layout-entity-v1/ index.json keys/ shard-00000.parquet shard-00001.parquet partitions.parquet tables/ - features/ - features-v1/ + table_0/ + v1/ index.json alignment.parquet chunk-0-0.bin ... - labels/ - labels-v1/ + table_1/ + v1/ index.json alignment.parquet chunk-0-0.bin ... - labels-v2/ + v2/ index.json alignment.parquet chunk-0-0.bin @@ -464,7 +464,7 @@ multi-join-dataset/ { "format": "litdata-multi-join", "format_version": 1, - "active_snapshot": "01J5A1QGMGT8P4J8Q1Z0AXEF3R", + "active_snapshot": "snap-table1-v2", "updated_at": "2026-08-17T10:42:11Z" } ``` @@ -475,12 +475,12 @@ multi-join-dataset/ { "format": "litdata-multi-join-snapshot", "format_version": 1, - "snapshot_id": "01J5A1QGMGT8P4J8Q1Z0AXEF3R", - "parent_snapshot_id": "01J4Y7M4VPH6V9B8P3N6K5W2HQ", + "snapshot_id": "snap-table1-v2", + "parent_snapshot_id": "snap-baseline", "created_at": "2026-08-17T10:42:10Z", "layout": { - "id": "6ce6f3...", - "path": "layouts/6ce6f3...", + "id": "layout-entity-v1", + "path": "layouts/layout-entity-v1", "key_name": "entity_id", "key_type": "string", "length": 1000000000, @@ -489,19 +489,19 @@ multi-join-dataset/ "alignment_root": "blake2b-256:..." }, "tables": { - "features": { - "version": "features-v1", - "path": "tables/features/features-v1", + "table_0": { + "version": "v1", + "path": "tables/table_0/v1", "format": "litdata", - "layout_id": "6ce6f3...", + "layout_id": "layout-entity-v1", "alignment_root": "blake2b-256:...", "schema_fingerprint": "sha256:..." }, - "labels": { - "version": "labels-v2", - "path": "tables/labels/labels-v2", + "table_1": { + "version": "v2", + "path": "tables/table_1/v2", "format": "litdata", - "layout_id": "6ce6f3...", + "layout_id": "layout-entity-v1", "alignment_root": "blake2b-256:...", "schema_fingerprint": "sha256:..." } @@ -536,9 +536,17 @@ The key store must be sharded and streamed. Creating or validating a billion-row The canonical key order, rather than lexical key order, defines training positions. Key-index shards may be physically sorted or hash-partitioned for lookup as long as the stored `global_index`, `chunk_index`, and `chunk_offset` preserve the canonical order. -### 8.4 Digest encoding +### 8.4 Identifiers vs digests -Digest computation must be deterministic across Python versions and machines: +Three different identifiers appear in the format. They must not be confused: + +- **Snapshot ID** (for example `snap-baseline`): a human-readable name for one published combination of table versions. Assigned at commit time. It is **not** a hash of chunk bytes. +- **Table version** (for example `v1`, `v2`): a human-readable name for one immutable build of a single table. +- **Ordered-key digest / alignment root**: a compact fingerprint of the **canonical entity-key order** inside each logical bucket. Used only to prove tables are aligned. LitData does **not** compute content hashes of chunk payloads for this purpose. + +### 8.5 Digest encoding + +Digest computation (for ordered keys only) must be deterministic across Python versions and machines: 1. Normalize the key to the supported integer or UTF-8 string representation. 2. Prefix each key with a type tag. @@ -551,7 +559,7 @@ The initial algorithm is `BLAKE2b-256`. The manifest stores the algorithm and ca Using `str(key)` concatenation without type and length framing is not acceptable because it can create ambiguous encodings. -### 8.5 Table alignment metadata +### 8.6 Table alignment metadata Every table version contains an `alignment.parquet` sidecar with: @@ -566,9 +574,9 @@ The table `index.json` gains a backward-compatible `multi_join` section: { "multi_join": { "format_version": 1, - "table": "labels", - "table_version": "labels-v2", - "layout_id": "6ce6f3...", + "table": "table_1", + "table_version": "v2", + "layout_id": "layout-entity-v1", "length": 1000000000, "num_chunks": 488282, "alignment_root": "blake2b-256:..." @@ -684,7 +692,7 @@ The following are hidden or rejected: `align_chunking=True` is important because it makes logical chunk boundaries independent of the number of optimize workers or nodes. Workers receive complete item-count chunks, and rank-index merge order reconstructs the canonical sequence. -A labels table may therefore be rebuilt with 8 workers while a feature table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. +A small table may therefore be rebuilt with 8 workers while a large table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. ### 10.4 V1 read path @@ -763,9 +771,9 @@ An explicit join cache root is namespaced: ```text cache/ / - features/ + table_0/ / - labels/ + table_1/ / ``` @@ -955,7 +963,7 @@ Snapshot tests: - A reader opened before commit continues on the old snapshot. - A reader opened after commit sees the new snapshot. -- Features objects are not rewritten by a labels-only update. +- `table_0` objects are not rewritten by a `table_1`-only update. - Rollback activates the previous immutable snapshot. - Concurrent stale writer commit fails. - Failure before active-manifest publication leaves the old snapshot active. @@ -1026,10 +1034,10 @@ This permits small or schema-volatile tables to use an efficient physical repres Assume the canonical bucket contains 2,048 entities: -- A feature bucket may be 64–256 MB. -- A label bucket may be only tens or hundreds of KB. +- A large `table_0` bucket may be 64–256 MB. +- A small `table_1` bucket may be only tens or hundreds of KB. -V1 writes one object for each in both tables. At very large scale, the labels table may have hundreds of thousands of tiny objects. This increases: +V1 writes one object for each in both tables. At very large scale, a small table may have hundreds of thousands of tiny objects. This increases: - Object listing and metadata cost. - GET request count. @@ -1038,7 +1046,7 @@ V1 writes one object for each in both tables. At very large scale, the labels ta - Cache bookkeeping. - Publication and garbage-collection overhead. -V2 allows many label buckets to be packed into one appropriately sized object while preserving the original logical bucket boundaries for sampling. +V2 allows many small-table buckets to be packed into one appropriately sized object while preserving the original logical bucket boundaries for sampling. ### 11.3 V2 storage model @@ -1052,14 +1060,14 @@ multi-join-dataset/ keys/ partitions.parquet tables/ - features/ + table_0/ / table.json mapping.parquet objects/ pack-00000.bin pack-00001.bin - labels/ + table_1/ / table.json mapping.parquet @@ -1149,8 +1157,8 @@ from litdata import MultiJoinWriter, TableStorage with MultiJoinWriter.open(root) as writer: writer.optimize_table( - "labels", - fn=build_labels_v3, + "table_1", + fn=build_table_1_v3, storage=TableStorage( format="parquet", target_chunk_bytes="128MB", @@ -1160,12 +1168,12 @@ with MultiJoinWriter.open(root) as writer: writer.commit() ``` -Large feature table: +Large table (`table_0`): ```python writer.optimize_table( - "features", - fn=build_features, + "table_0", + fn=build_table_0, storage=TableStorage( format="litdata", target_chunk_bytes="256MB", @@ -1306,8 +1314,8 @@ The existing LitData sampler and independent readers can be reused. ```text logical buckets 0..9: - features -> 10 feature objects - labels -> 1 packed label object + table_0 -> 10 objects + table_1 -> 1 packed object ``` A new shared sampler and mapping-aware reader are required. @@ -1323,12 +1331,12 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam 01. Choose the entity key. 02. Produce a unique canonical key list. 03. Intentionally choose its training order. -04. Choose the logical bucket item count based primarily on feature bytes and desired bucket sampling. +04. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. 05. Create the immutable layout. 06. Group each source table into one logical contribution per key. 07. Align each table input to canonical `global_index`. -08. Optimize stable large tables once. -09. Optimize schema-volatile tables. +08. Optimize large stable tables once. +09. Optimize schema-volatile tables (for example `table_1`). 10. Validate and publish the first snapshot. 11. Benchmark against the current baked dataset. From c41c395e4e502523e7de865c5912599432cd9785 Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 12:45:47 +0000 Subject: [PATCH 5/9] Finish renaming V1 alignment diagram to table_N. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 74c8931e..189be3da 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -1303,9 +1303,9 @@ The boundary is intentionally simple: ### V1 aligns physical chunks ```text -logical bucket 0 -> features chunk 0 -> labels chunk 0 -logical bucket 1 -> features chunk 1 -> labels chunk 1 -logical bucket 2 -> features chunk 2 -> labels chunk 2 +logical bucket 0 -> table_0 chunk 0 -> table_1 chunk 0 +logical bucket 1 -> table_0 chunk 1 -> table_1 chunk 1 +logical bucket 2 -> table_0 chunk 2 -> table_1 chunk 2 ``` The existing LitData sampler and independent readers can be reused. From f82dfb14cfa8df7ce6fc2cbe8ee30ab9c62257e7 Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 13:08:22 +0000 Subject: [PATCH 6/9] Rewrite write API as independent joint_optimize jobs. Drop MultiJoinWriter / key_name in favor of optimize-like joint_optimize calls that can run as distributed jobs and publish snapshots on success. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 272 +++++++++--------- 1 file changed, 139 insertions(+), 133 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 189be3da..11503c15 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -170,10 +170,12 @@ V2 can add a native offset-based representation for zero-to-many rows per key, b ### API quality +- Write path feels like `optimize`: independent `joint_optimize(...)` calls, including multi-node jobs. +- No long-lived writer context manager that must wrap every table build. +- No required `key_name` / schema-field argument; keys come from `key_fn` like existing `optimize`. - A simple root-path read API. - Named table outputs rather than positional tuples. - One shared set of sampling options. -- Explicit commit semantics for write operations. - No need for users to configure `align_chunking`, `reorder_files`, child seeds, or child cache directories. - The same user-facing read API in V1 and V2. @@ -198,98 +200,98 @@ The full baked dataset remains the recommended default when all tables normally The API below is the target public contract. Some advanced storage arguments become effective only in V2, but normal training code does not change. -### 7.1 Create a canonical layout and initial snapshot +The write path mirrors existing `optimize`: each table is an independent function call that can run as a local or multi-node job. There is no `MultiJoinWriter` context manager and no `key_name` argument. + +### 7.1 Build tables with `joint_optimize` ```python -from litdata import MultiJoinWriter +from litdata import joint_optimize -entity_ids = load_canonical_entity_ids() +# Independent jobs — each call can use num_workers / num_nodes like optimize(). +# The first successful call into an empty join root creates the shared layout +# from the ordered inputs + key_fn. Later calls must align to that layout. -with MultiJoinWriter.create( - "s3://bucket/multi-join-dataset", - keys=entity_ids, - key_name="entity_id", +joint_optimize( + fn=build_table_0, + inputs=inputs_0, + output_dir="s3://bucket/multi-join-dataset", + table="table_0", + chunk_size=2048, + key_fn=lambda sample: sample["id"], + num_workers=32, + num_nodes=8, + compression="zstd", +) + +joint_optimize( + fn=build_table_1, + inputs=inputs_1, + output_dir="s3://bucket/multi-join-dataset", + table="table_1", chunk_size=2048, -) as writer: - writer.optimize_table( - "table_0", - fn=build_table_0, - version="v1", - num_workers=32, - num_nodes=8, - compression="zstd", - ) - writer.optimize_table( - "table_1", - fn=build_table_1, - version="v1", - num_workers=8, - ) - snapshot = writer.commit() - -print(snapshot.id) + key_fn=lambda sample: sample["id"], + num_workers=8, +) ``` +`joint_optimize` accepts the same core knobs as `optimize` (`fn`, `inputs`, `output_dir`, `chunk_size`, `key_fn`, `num_workers`, `num_nodes`, `compression`, `encryption`, `storage_options`, …) plus join-specific arguments: + +- `table`: logical table name under the join root (for example `"table_0"`). +- `version`: optional immutable table-version name; defaults to an auto-generated unique version. +- `snapshot`: optional human-readable snapshot name to publish on success. +- `expected_snapshot`: optional concurrency guard against the currently active snapshot. + Semantics: -- `keys` defines the canonical order once. -- Users should intentionally randomize or otherwise curate this order before layout creation when source order has structure. -- When `inputs` is omitted, `fn` receives each canonical key. -- `chunk_size` is the logical bucket size and is shared by every V1 table. -- Version paths are immutable. -- `commit()` validates all table versions and publishes the snapshot. -- Exiting the context without `commit()` does not publish changes. -- Calling `optimize_table()`, `remove_table()`, or `commit()` after a successful commit raises an error. +- Each call is a complete, distributed-capable job. Jobs do not share a process-local writer session. +- `inputs` order is the stream / layout order. Shuffle or otherwise curate inputs before calling when source order has structure. +- `key_fn` works like existing `optimize(..., key_fn=...)`: it extracts an opaque key from each sample for the alignment sidecar. There is no separate `key_name` schema field. +- `chunk_size` is required (item count). `chunk_bytes` is rejected. +- The first table that successfully initializes an empty join root writes the shared layout (ordered keys, bucket boundaries, digests). +- Later `joint_optimize` calls for other tables must reproduce the same length, bucket sizes, and ordered-key digests. +- On success, the job validates alignment and atomically publishes a new snapshot that includes this table version, keeping previously published sibling table versions from the prior active snapshot. +- Failed jobs leave the previous active snapshot unchanged. +- Rebuilding one table later is just another `joint_optimize` into a new `version=` under the same `table=`. -### 7.2 Optimize from pre-grouped inputs +### 7.2 Same ordered inputs across tables -Large pipelines may already produce one grouped input object per entity: +Every table must be optimized from inputs that follow the same entity order. Typical patterns: ```python -with MultiJoinWriter.create( - output_dir, - keys=entity_ids, - key_name="entity_id", - chunk_size=2048, -) as writer: - writer.optimize_table( - "table_2", - fn=encode_table_2_group, - inputs=table_2_groups, - input_key=lambda group: group.entity_id, - ) - writer.commit() +# inputs_0 / inputs_1 are already aligned to the same entity order +joint_optimize(fn=build_table_0, inputs=inputs_0, output_dir=root, table="table_0", chunk_size=2048, key_fn=get_id) +joint_optimize(fn=build_table_1, inputs=inputs_1, output_dir=root, table="table_1", chunk_size=2048, key_fn=get_id) ``` Contract: -- `inputs` must already follow the canonical order. -- `input_key` is checked against the expected canonical key at every position. -- The writer does not build an in-memory key-to-input map or silently reorder billions of records. -- A mismatch reports the table, global index, expected key, and actual key. -- The preparation pipeline may align data by joining against the canonical `(key, global_index)` mapping and ordering by `global_index`. +- Do not silently reorder inputs inside `joint_optimize`. +- A key mismatch reports table, global index, expected key, and actual key. +- Users may prepare inputs by joining against the published layout key store and sorting by `global_index`. +- Parallel first-time builds of multiple tables into an empty root are unsafe without a shared layout. Create the layout with one `joint_optimize` (or a small driver table) first, then launch sibling table jobs. ### 7.3 Re-optimize one table after a schema change ```python -from litdata import MultiJoinWriter - -with MultiJoinWriter.open( - "s3://bucket/multi-join-dataset", +from litdata import joint_optimize + +joint_optimize( + fn=build_table_1_v2, + inputs=inputs_1, + output_dir="s3://bucket/multi-join-dataset", + table="table_1", + version="v2", + chunk_size=2048, + key_fn=lambda sample: sample["id"], + num_workers=8, expected_snapshot="snap-baseline", -) as writer: - writer.optimize_table( - "table_1", - fn=build_table_1_v2, - version="v2", - num_workers=8, - ) - new_snapshot = writer.commit() + snapshot="snap-table1-v2", +) ``` -Only the new `table_1` version prefix and a new snapshot document are written. The active `table_0` version is referenced unchanged. +Only the new `table_1` version prefix and a new snapshot document are written. Other tables keep their previously published versions. -`expected_snapshot` is the **name** of the currently active snapshot (for example `"snap-baseline"`), not a hash of chunk bytes. Writers pass it as an optimistic concurrency guard: if another publisher advances the active snapshot first, commit fails rather than overwriting that change. Snapshot IDs are opaque human-readable identifiers assigned at publish time; LitData does not compute content hashes of chunk payloads to form them. +`expected_snapshot` / `snapshot` are **human-readable names** (for example `"snap-baseline"`), not hashes of chunk bytes. `expected_snapshot` is an optimistic concurrency guard: if another publisher advances the active snapshot first, this job fails rather than overwriting that change. LitData does not compute content hashes of chunk payloads to form snapshot IDs. ### 7.4 Stream the active snapshot @@ -407,7 +409,7 @@ Validation levels: - Construction always performs mandatory constant-size metadata validation. - `deep=False` verifies manifests, table indexes, counts, layout IDs, and alignment roots. - `deep=True` streams partition alignment metadata and verifies every ordered-key digest. -- Writer commit always performs the validation required to prove a newly published table matches the canonical layout. +- Writer / `joint_optimize` success always performs the validation required to prove a newly published table matches the canonical layout. ### 7.9 Keyed debugging access @@ -481,7 +483,6 @@ multi-join-dataset/ "layout": { "id": "layout-entity-v1", "path": "layouts/layout-entity-v1", - "key_name": "entity_id", "key_type": "string", "length": 1000000000, "chunk_size": 2048, @@ -540,7 +541,7 @@ The canonical key order, rather than lexical key order, defines training positio Three different identifiers appear in the format. They must not be confused: -- **Snapshot ID** (for example `snap-baseline`): a human-readable name for one published combination of table versions. Assigned at commit time. It is **not** a hash of chunk bytes. +- **Snapshot ID** (for example `snap-baseline`): a human-readable name for one published combination of table versions. Assigned when a `joint_optimize` job publishes successfully. It is **not** a hash of chunk bytes. - **Table version** (for example `v1`, `v2`): a human-readable name for one immutable build of a single table. - **Ordered-key digest / alignment root**: a compact fingerprint of the **canonical entity-key order** inside each logical bucket. Used only to prove tables are aligned. LitData does **not** compute content hashes of chunk payloads for this purpose. @@ -590,20 +591,22 @@ Per-chunk digests live in compact Parquet rather than expanding an already large ### 9.1 Publication protocol -`MultiJoinWriter.commit()` follows this order: +Each successful `joint_optimize(...)` publishes as part of the same job (no separate writer `commit()`): -01. Read and retain the expected active snapshot. +01. Read and retain the expected active snapshot when `expected_snapshot` is set. 02. Write new table data to a unique immutable version path. -03. Upload all chunk objects. +03. Upload all chunk objects (existing multi-worker / multi-node optimize upload path). 04. Upload table alignment metadata. 05. Upload the table `index.json` and completion marker last. -06. Validate the table against the canonical layout. -07. Write a new immutable snapshot document. -08. Recheck the active snapshot or storage generation. -09. Atomically replace `join.json` with the new active snapshot. -10. Mark the writer committed and reject further mutation. +06. If this is the first table in an empty join root, publish the shared layout from the ordered `key_fn` stream. +07. Validate the table against the canonical layout. +08. Write a new immutable snapshot document that merges this table version with sibling versions from the previous active snapshot. +09. Recheck the active snapshot or storage generation. +10. Atomically replace `join.json` with the new active snapshot. + +If any operation fails before step 10, the active snapshot remains unchanged. Unreferenced objects are safe to garbage-collect later. -If any operation fails before step 9, the active snapshot remains unchanged. Unreferenced objects are safe to garbage-collect later. +Because each `joint_optimize` is an independent distributed job, publication must be safe across machines and processes. There is no process-local `MultiJoinWriter` session holding uncommitted state. ### 9.2 Backend behavior @@ -668,15 +671,16 @@ For the read configuration: ### 10.3 V1 write path -`MultiJoinWriter.optimize_table()` internally calls the existing optimize pipeline with required safe settings: +`joint_optimize(...)` wraps the existing optimize pipeline with required safe settings: -- `chunk_size` comes from the canonical layout. +- `chunk_size` is required and becomes the shared logical bucket size. - `align_chunking=True`. - `reorder_files=False`. - `keep_data_ordered=True`. - Static ordered inputs. - One output sample per input. -- Key capture enabled for alignment validation. +- `key_fn` required so alignment digests can be computed (same mechanism as `optimize(..., key_fn=...)`). +- Output lands under `tables/{table}/{version}/` inside the join root rather than a standalone dataset root. The following are hidden or rejected: @@ -689,10 +693,11 @@ The following are hidden or rejected: - variable-yield generators - append into an existing table version - overwrite of a published table version +- a `key_name` / schema-field argument `align_chunking=True` is important because it makes logical chunk boundaries independent of the number of optimize workers or nodes. Workers receive complete item-count chunks, and rank-index merge order reconstructs the canonical sequence. -A small table may therefore be rebuilt with 8 workers while a large table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. +A small table may therefore be rebuilt with 8 workers while a large table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. Because each call is a normal optimize-style job, `num_nodes` / Studio multi-node execution works the same way as `optimize`. ### 10.4 V1 read path @@ -844,23 +849,21 @@ Do not add a heavy schema dependency solely for these small metadata models. #### Write API -Add `src/litdata/processing/multi_join.py`: +Add `joint_optimize` in `src/litdata/processing/functions.py` (or a thin wrapper module re-exported from `__init__.py`): -- `MultiJoinWriter.create` -- `MultiJoinWriter.open` -- `optimize_table` -- `remove_table` -- `validate` -- `commit` -- abort and cleanup bookkeeping +- Same distributed execution path as `optimize` (`DataProcessor`, `num_nodes`, uploaders, index merge). +- Extra args: `table`, optional `version`, `snapshot`, `expected_snapshot`. +- Force aligned write settings listed above. +- After merge, validate against the join layout and publish/update the snapshot atomically. +- No `MultiJoinWriter` context manager and no process-local uncommitted session. Extend the optimize internals where necessary to: -- Stream ordered `(global_index, key)` metadata. +- Stream ordered `(global_index, key)` metadata from `key_fn`. - Compute per-chunk key digests. - Avoid materializing all keys during alignment generation. - Attach the backward-compatible `multi_join` section to `index.json`. -- Return a typed table-version result to the writer. +- Create the shared layout on first publish into an empty join root. Likely integration points: @@ -868,6 +871,7 @@ Likely integration points: - `src/litdata/processing/data_processor.py` - `src/litdata/streaming/writer.py` - `src/litdata/utilities/keys_index.py` +- `src/litdata/utilities/multi_join.py` #### Read API @@ -965,7 +969,7 @@ Snapshot tests: - A reader opened after commit sees the new snapshot. - `table_0` objects are not rewritten by a `table_1`-only update. - Rollback activates the previous immutable snapshot. -- Concurrent stale writer commit fails. +- Concurrent stale `joint_optimize` publish fails. - Failure before active-manifest publication leaves the old snapshot active. Remote tests: @@ -1007,7 +1011,7 @@ No production throughput claim should be made before this benchmark runs on a re V1 is complete when: -1. A two-table and a many-table dataset can be created through the public writer API. +1. A two-table and a many-table dataset can be created through independent `joint_optimize` jobs. 2. Re-optimizing one table does not write under unchanged table-version prefixes. 3. A new table snapshot is atomically activated. 4. Old readers continue without observing mixed versions. @@ -1148,32 +1152,38 @@ The shared reader owns one aggregate cache budget: This eliminates V1’s approximate per-child budget split. -### 11.7 V2 writer API extension +### 11.7 V2 `joint_optimize` storage extension -The public writer gains optional per-table storage configuration: +`joint_optimize` gains optional per-table storage configuration: ```python -from litdata import MultiJoinWriter, TableStorage - -with MultiJoinWriter.open(root) as writer: - writer.optimize_table( - "table_1", - fn=build_table_1_v3, - storage=TableStorage( - format="parquet", - target_chunk_bytes="128MB", - pack_logical_chunks=True, - ), - ) - writer.commit() +from litdata import joint_optimize, TableStorage + +joint_optimize( + fn=build_table_1_v3, + inputs=inputs_1, + output_dir=root, + table="table_1", + chunk_size=2048, + key_fn=get_id, + storage=TableStorage( + format="parquet", + target_chunk_bytes="128MB", + pack_logical_chunks=True, + ), +) ``` Large table (`table_0`): ```python -writer.optimize_table( - "table_0", +joint_optimize( fn=build_table_0, + inputs=inputs_0, + output_dir=root, + table="table_0", + chunk_size=2048, + key_fn=get_id, storage=TableStorage( format="litdata", target_chunk_bytes="256MB", @@ -1328,26 +1338,21 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam ### One-time migration -01. Choose the entity key. -02. Produce a unique canonical key list. -03. Intentionally choose its training order. +01. Choose the entity key extraction (`key_fn`). +02. Produce unique ordered inputs for every table (same entity order). +03. Intentionally choose that training order before the first `joint_optimize`. 04. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. -05. Create the immutable layout. -06. Group each source table into one logical contribution per key. -07. Align each table input to canonical `global_index`. -08. Optimize large stable tables once. -09. Optimize schema-volatile tables (for example `table_1`). -10. Validate and publish the first snapshot. -11. Benchmark against the current baked dataset. +05. Run `joint_optimize` for the first / driver table to create the join root and layout. +06. Run independent `joint_optimize` jobs for the remaining tables (local or multi-node). +07. Validate the published snapshot. +08. Benchmark against the current baked dataset. ### Small-table schema update -1. Open the current store with its expected snapshot. -2. Rebuild only the changed table under a new version. -3. Validate keys, bucket counts, and digests. -4. Publish a new snapshot. -5. Start new training jobs on the new snapshot. -6. Leave existing jobs pinned to their original snapshot. +01. Keep existing sibling table versions as-is. +02. Run `joint_optimize(..., table=..., version=..., expected_snapshot=...)` for the changed table only. +03. Start new training jobs on the newly published snapshot. +04. Leave existing jobs pinned to their original snapshot. ### Rollback @@ -1382,7 +1387,7 @@ A future append-only layout extension could preserve complete existing buckets a Initial behavior: -- Commits never delete table versions. +- Successful `joint_optimize` jobs never delete previous table versions. - Rollback remains possible while snapshots and versions are retained. - Failed staging outputs are recorded as unreachable. @@ -1394,7 +1399,7 @@ Later garbage collection: 4. Delete only unreferenced immutable prefixes. 5. Support dry-run output before deletion. -Garbage collection is never part of `commit()`. +Garbage collection is never part of `joint_optimize` publication. ### Observability @@ -1504,9 +1509,10 @@ Mitigation: 07. V2 introduces one shared sampler and coordinated reader. 08. Table versions and snapshots are immutable. 09. `join.json` is published last and readers pin one snapshot. -10. Writer publication requires explicit `commit()`. +10. Each `joint_optimize` job validates and publishes atomically; no writer context / deferred `commit()`. 11. Training iteration performs no per-sample key lookup. 12. The public `MultiJoinStreamingDataset` API remains stable across V1 and V2. +13. Keys come from `key_fn`, not from a `key_name` schema-field argument. ## 17. Recommended delivery order @@ -1515,7 +1521,7 @@ Mitigation: 01. Freeze manifest, layout, digest, and API specifications. 02. Implement metadata models and validation. 03. Implement scalable canonical layout creation. -04. Implement `MultiJoinWriter` with strict aligned optimize settings. +04. Implement `joint_optimize` with strict aligned optimize settings and multi-node support. 05. Implement atomic immutable snapshots. 06. Implement the validated `ParallelStreamingDataset` wrapper. 07. Add snapshot-aware state and resume. From 2744a38946cf7d99866d98003bdbc04e7a63d6c0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:08:36 +0000 Subject: [PATCH 7/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .claude/plans/multi_join_streaming_dataset.md | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 11503c15..55f46378 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -1338,21 +1338,21 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam ### One-time migration -01. Choose the entity key extraction (`key_fn`). -02. Produce unique ordered inputs for every table (same entity order). -03. Intentionally choose that training order before the first `joint_optimize`. -04. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. -05. Run `joint_optimize` for the first / driver table to create the join root and layout. -06. Run independent `joint_optimize` jobs for the remaining tables (local or multi-node). -07. Validate the published snapshot. -08. Benchmark against the current baked dataset. +1. Choose the entity key extraction (`key_fn`). +2. Produce unique ordered inputs for every table (same entity order). +3. Intentionally choose that training order before the first `joint_optimize`. +4. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. +5. Run `joint_optimize` for the first / driver table to create the join root and layout. +6. Run independent `joint_optimize` jobs for the remaining tables (local or multi-node). +7. Validate the published snapshot. +8. Benchmark against the current baked dataset. ### Small-table schema update -01. Keep existing sibling table versions as-is. -02. Run `joint_optimize(..., table=..., version=..., expected_snapshot=...)` for the changed table only. -03. Start new training jobs on the newly published snapshot. -04. Leave existing jobs pinned to their original snapshot. +1. Keep existing sibling table versions as-is. +2. Run `joint_optimize(..., table=..., version=..., expected_snapshot=...)` for the changed table only. +3. Start new training jobs on the newly published snapshot. +4. Leave existing jobs pinned to their original snapshot. ### Rollback From b9596c94b82f7b3b251ed6ebdc684f1dcb95195f Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Mon, 10 Aug 2026 13:25:28 +0000 Subject: [PATCH 8/9] Freeze TablesStreamingDataset write/read contracts for V1. Lock join_optimize + join_update_optimize, input planning, sequential key validation, uniqueness, and snapshot concurrency rules into the design plan. Co-authored-by: Cursor --- .claude/plans/multi_join_streaming_dataset.md | 718 ++++++++++++------ 1 file changed, 485 insertions(+), 233 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index 11503c15..ccae6490 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -1,4 +1,4 @@ -# MultiJoinStreamingDataset +# TablesStreamingDataset ## Product and implementation plan @@ -8,7 +8,7 @@ Scope: independently versioned tables that describe the same training entity and ## Executive summary -`MultiJoinStreamingDataset` is intended for datasets in which: +`TablesStreamingDataset` is intended for datasets in which: - Several tables or modalities describe the same logical entity. - One or more large tables are stable and expensive to rebuild. @@ -21,7 +21,7 @@ The central constraint cannot be removed: if tables are not joined before traini The work is split into two deliberately different phases: -- **Phase V1 — strict aligned LitData chunks.** Every table is a normal optimized LitData dataset with exactly the same logical chunk boundaries, item counts, and key order. `MultiJoinStreamingDataset` validates the layout and reuses `ParallelStreamingDataset`. This minimizes new read-path code and preserves existing LitData behavior. +- **Phase V1 — strict aligned LitData chunks.** Every table is a normal optimized LitData dataset with exactly the same logical chunk boundaries, item counts, and key order. `TablesStreamingDataset` validates the layout and reuses `ParallelStreamingDataset`. This minimizes new read-path code and preserves existing LitData behavior. - **Phase V2 — logical partitions with independently compacted column families.** Logical sampling buckets remain aligned, but physical files no longer need a one-to-one correspondence. Small tables can pack many logical partitions into larger Parquet or LitData objects, while large tables keep appropriately sized binary chunks. V2 introduces one shared sampler and a coordinated multi-table reader. The public read API is designed once and remains stable across both phases. V2 is primarily an internal storage and reader improvement. @@ -36,7 +36,7 @@ flowchart LR activeManifest[Atomic join manifest] --> t0 activeManifest --> t1 activeManifest --> t2 - t0 --> joinedReader[MultiJoinStreamingDataset] + t0 --> joinedReader[TablesStreamingDataset] t1 --> joinedReader t2 --> joinedReader joinedReader --> training[Training batches] @@ -55,8 +55,8 @@ The current high-throughput LitData workflow materializes the complete training The target behavior is: -1. Optimize each logical table independently. -2. Re-optimize only the table whose values or schema changed. +1. Optimize all required tables together from one ordered input stream so chunk alignment is guaranteed. +2. Later, update (or add) one table with `join_update_optimize`, validating against the existing chunk alignment. 3. Atomically publish a new combination of table versions. 4. Stream the selected versions as one training sample. 5. Preserve deterministic shuffle, DDP sharding, DataLoader workers, prefetch, caching, and exact resume. @@ -140,6 +140,10 @@ V1 rejects: An empty table contribution must be represented as an explicit empty value. This keeps every table positionally aligned. +**Runtime rule for builders:** In V1, a table builder must return exactly one serializable Python value for each input item. Returning `None` is always rejected and is not a representation of an empty contribution. Nullable payloads use an explicit envelope such as `{"value": None}`. Zero-row payloads use `[]`, `{}`, or a typed empty array such as `np.empty((0,), dtype=np.float32)`. + +A builder exception aborts the whole job and prevents snapshot publication, even if other table writers already produced partial local output. + V2 can add a native offset-based representation for zero-to-many rows per key, but the canonical entity key and logical bucket alignment remain mandatory. ## 5. Design goals @@ -170,11 +174,13 @@ V2 can add a native offset-based representation for zero-to-many rows per key, b ### API quality -- Write path feels like `optimize`: independent `joint_optimize(...)` calls, including multi-node jobs. -- No long-lived writer context manager that must wrap every table build. -- No required `key_name` / schema-field argument; keys come from `key_fn` like existing `optimize`. -- A simple root-path read API. -- Named table outputs rather than positional tuples. +- Write path feels like `optimize`: one distributed `join_optimize(...)` job writes all selected tables from a single ordered input stream. +- Single-table updates (replace or add) use `join_update_optimize(...)`, which validates against the persisted layout. +- Public vocabulary is consistent: `TablesStreamingDataset`, `join_optimize`, `join_update_optimize`, `validate_tables_dataset`, `TableBuild`, `Tables*Error`. +- `key_fn` always receives the original input item, never a table output. +- No long-lived writer context manager. +- No required `key_name` / schema-field argument. +- A simple root-path read API with named table outputs. - One shared set of sampling options. - No need for users to configure `align_chunking`, `reorder_files`, child seeds, or child cache directories. - The same user-facing read API in V1 and V2. @@ -200,113 +206,203 @@ The full baked dataset remains the recommended default when all tables normally The API below is the target public contract. Some advanced storage arguments become effective only in V2, but normal training code does not change. -The write path mirrors existing `optimize`: each table is an independent function call that can run as a local or multi-node job. There is no `MultiJoinWriter` context manager and no `key_name` argument. +There is no `TablesWriter` context manager and no `key_name` argument. -### 7.1 Build tables with `joint_optimize` +Alignment is enforced at write time by construction for the initial build: **one** `join_optimize` job walks a single ordered `inputs` stream and writes every selected table under its own subpath. Independent per-table optimize jobs with separate input lists are rejected as the primary API because they cannot guarantee the same entity order or chunk boundaries. -```python -from litdata import joint_optimize +### 7.1 Initial build with `join_optimize` + +Canonical public form: a `tables` mapping of builders (or `TableBuild` configs). There is no dual `fn=` / `fns=` create API in V1. -# Independent jobs — each call can use num_workers / num_nodes like optimize(). -# The first successful call into an empty join root creates the shared layout -# from the ordered inputs + key_fn. Later calls must align to that layout. +```python +from litdata import TableBuild, join_optimize -joint_optimize( - fn=build_table_0, - inputs=inputs_0, - output_dir="s3://bucket/multi-join-dataset", - table="table_0", +join_optimize( + inputs=canonical_inputs, + output_dir="s3://bucket/tables-dataset", + key_fn=lambda item: item["entity_id"], chunk_size=2048, - key_fn=lambda sample: sample["id"], + tables={ + "images": TableBuild(fn=build_images, compression="zstd"), + "labels": TableBuild(fn=build_labels, compression="zstd"), + "metadata": TableBuild(fn=build_metadata), + }, num_workers=32, num_nodes=8, - compression="zstd", + snapshot="baseline", ) +``` -joint_optimize( - fn=build_table_1, - inputs=inputs_1, - output_dir="s3://bucket/multi-join-dataset", - table="table_1", - chunk_size=2048, - key_fn=lambda sample: sample["id"], - num_workers=8, -) +Callables are accepted as shorthand and normalized to `TableBuild` internally: + +```python +tables={ + "images": build_images, + "labels": build_labels, +} ``` -`joint_optimize` accepts the same core knobs as `optimize` (`fn`, `inputs`, `output_dir`, `chunk_size`, `key_fn`, `num_workers`, `num_nodes`, `compression`, `encryption`, `storage_options`, …) plus join-specific arguments: +`join_optimize` accepts the same core knobs as `optimize` (`inputs`, `output_dir`, `chunk_size`, `key_fn`, `num_workers`, `num_nodes`, `storage_options`, …) plus: -- `table`: logical table name under the join root (for example `"table_0"`). -- `version`: optional immutable table-version name; defaults to an auto-generated unique version. +- `tables`: ordered mapping `{name: TableBuild | Callable}`. Insertion order becomes the snapshot `table_order`. +- `version` / per-table versions: optional; default is an auto-generated unique version shared by all tables in the job, or a mapping `{table: version}`. - `snapshot`: optional human-readable snapshot name to publish on success. -- `expected_snapshot`: optional concurrency guard against the currently active snapshot. +- `expected_snapshot`: optional concurrency guard when updating an existing root. -Semantics: +#### `key_fn` contract + +`key_fn(input_item)` extracts the canonical entity key from the **source input item before any table builder runs**. It is invoked exactly once per logical input entity. Table outputs are **not** required to contain the key. A later schema change may remove or rename an ID field inside a table without changing alignment behavior. -- Each call is a complete, distributed-capable job. Jobs do not share a process-local writer session. -- `inputs` order is the stream / layout order. Shuffle or otherwise curate inputs before calling when source order has structure. -- `key_fn` works like existing `optimize(..., key_fn=...)`: it extracts an opaque key from each sample for the alignment sidecar. There is no separate `key_name` schema field. -- `chunk_size` is required (item count). `chunk_bytes` is rejected. -- The first table that successfully initializes an empty join root writes the shared layout (ordered keys, bucket boundaries, digests). -- Later `joint_optimize` calls for other tables must reproduce the same length, bucket sizes, and ordered-key digests. -- On success, the job validates alignment and atomically publishes a new snapshot that includes this table version, keeping previously published sibling table versions from the prior active snapshot. -- Failed jobs leave the previous active snapshot unchanged. -- Rebuilding one table later is just another `joint_optimize` into a new `version=` under the same `table=`. +#### Semantics -### 7.2 Same ordered inputs across tables +- One distributed-capable job, same execution model as `optimize` (workers, nodes, upload, index merge). +- For each input item, every table builder in `tables` is invoked with that same item and must return exactly one serializable non-`None` value. +- `input_item` is shared **read-only** input. Mutating it inside a builder is unsupported behavior (later builders must not observe different data). +- The job uses one canonical input order. All tables share entity order, global indexes, `chunk_size` boundaries, and within-chunk key order **by construction**. +- Chunks land under `tables/{table}/{version}/` with separate `index.json` files, identical item counts per chunk index. +- `chunk_size` is required (item count). `chunk_bytes` is rejected for join writes. +- Root-level `compression` / `encryption` apply as defaults; per-table `TableBuild` fields override them. +- On success the job writes the shared layout, **proves normalized entity-key uniqueness**, validates alignment metadata, and atomically publishes `join.json`. +- Failed jobs leave the previous active snapshot unchanged. They may leave unreferenced staging objects under unique version prefixes; those are never valid without a completion marker and snapshot reference. GC removes them only after the configured safety window. +- Initial layout creation is **conditional**. A `join_optimize` job may initialize a root only if no active manifest and no completed layout exist. The backend must use an exclusive create, generation precondition, or publisher lease. If another initializer wins, the losing job fails with `TablesCommitConflictError`. +- Adding a table to an existing layout uses `join_update_optimize` (see 7.3). -Every table must be optimized from inputs that follow the same entity order. Typical patterns: +### 7.2 Why not one optimize call per table ```python -# inputs_0 / inputs_1 are already aligned to the same entity order -joint_optimize(fn=build_table_0, inputs=inputs_0, output_dir=root, table="table_0", chunk_size=2048, key_fn=get_id) -joint_optimize(fn=build_table_1, inputs=inputs_1, output_dir=root, table="table_1", chunk_size=2048, key_fn=get_id) +# Rejected as the primary write API +optimize(..., inputs=inputs_0) # into tables/table_0/... +optimize(..., inputs=inputs_1) # into tables/table_1/... ``` -Contract: +Separate input lists or separate distributed schedules can diverge in order, filtering, retries, or worker assignment. Digests can detect mismatch after the fact, but they cannot invent a shared chunk layout. V1 therefore builds alignment by writing all tables together in one `join_optimize` job. -- Do not silently reorder inputs inside `joint_optimize`. -- A key mismatch reports table, global index, expected key, and actual key. -- Users may prepare inputs by joining against the published layout key store and sorting by `global_index`. -- Parallel first-time builds of multiple tables into an empty root are unsafe without a shared layout. Create the layout with one `joint_optimize` (or a small driver table) first, then launch sibling table jobs. +### 7.3 Update or add one table with `join_update_optimize` -### 7.3 Re-optimize one table after a schema change +| Operation | Alignment property | +|---|---| +| `join_optimize` initial build | Guaranteed by one canonical traversal and synchronized writers | +| `join_update_optimize` | Validated against the persisted canonical layout | +| Future layout-driven update (post-V1) | Guaranteed by streaming persisted layout keys through `lookup_fn` | -```python -from litdata import joint_optimize +`join_update_optimize` does **not** infer or repair entity order. It validates that the provided input stream reproduces the persisted layout exactly. A key mismatch, missing entity, duplicate displacement, extra entity, or different final length fails the job before snapshot publication. -joint_optimize( - fn=build_table_1_v2, - inputs=inputs_1, - output_dir="s3://bucket/multi-join-dataset", - table="table_1", +```python +from litdata import join_update_optimize + +join_update_optimize( + inputs=canonical_inputs, + output_dir="s3://bucket/tables-dataset", + table="metadata", + fn=build_metadata_v2, + key_fn=lambda item: item["entity_id"], version="v2", - chunk_size=2048, - key_fn=lambda sample: sample["id"], + base_snapshot="baseline", + expected_snapshot="baseline", + snapshot="metadata-v2", num_workers=8, - expected_snapshot="snap-baseline", - snapshot="snap-table1-v2", ) ``` -Only the new `table_1` version prefix and a new snapshot document are written. Other tables keep their previously published versions. +Semantics: + +- Reads `base_snapshot` (default: the active snapshot observed at job start) and its canonical layout. +- In V1, `base_snapshot` identifies both the input dependency set and the snapshot from which sibling table versions are inherited. When `expected_snapshot` is provided, it **must equal** `base_snapshot`. Publication succeeds only if the active snapshot still equals that value at compare-and-swap time. Differing values are rejected because it would be unclear which snapshot supplies inherited siblings, and a merge from a stale base could silently discard newer sibling updates. A stale base always raises `TablesCommitConflictError`; LitData does not perform automatic snapshot merges. +- `key_fn` always receives the original input item (read-only). +- Forces the layout `chunk_size` and aligned write settings. +- Writes only `tables/{table}/{new_version}/`. +- During the write, every work unit validates by **streaming** expected keys for its assigned contiguous global-index range (not random per-item metadata lookup): + +```python +for global_index, expected_key, input_item in zip( + assigned_positions, + layout.iter_keys(start=range_start, stop=range_stop), + assigned_inputs, +): + actual_key = key_fn(input_item) + if actual_key != expected_key: + raise TablesAlignmentError( + table=table, + global_index=global_index, + expected_key=expected_key, + actual_key=actual_key, + ) +``` + +The conceptual `layout.key_at(global_index)` idea is implemented as sequential `layout.iter_keys(start, stop)` over each assigned range. Implementations must not perform one remote random key-index lookup per input item. + +- At stream end for each range / overall: fail if inputs ended early (missing expected position) or continued past layout length (extra input). +- After merge, also verify length, chunk counts, per-chunk item counts, and ordered-key digests (compact completion / deep-validation metadata). +- Uniqueness is **not** recomputed on update/add if the base layout was already validated as unique; sequential reproduction is sufficient. +- On success, publishes a new snapshot that inherits sibling versions from `base_snapshot` and points `table` at the new version. +- Reject overwrite of an existing published version path before writing. +- Reject `fn` returning `None` immediately. +- On any failure, leave no completion marker and do not advance `join.json`. Failed jobs may leave unreferenced staging objects under the unique version prefix; they are unreachable without a completion marker and snapshot reference. + +#### Adding a new table (supported in V1) + +```python +join_update_optimize( + inputs=canonical_inputs, + output_dir=root, + table="new_table", + fn=build_new_table, + key_fn=lambda item: item["entity_id"], + base_snapshot="baseline", + expected_snapshot="baseline", + snapshot="snap-add-new-table", +) +``` + +Adding a table creates a new snapshot containing inherited sibling versions plus the new table. Existing snapshots remain valid and do **not** retroactively gain the table. A reader requesting a table absent from its pinned snapshot raises during construction (`TablesManifestError`). + +`table_order` for the new snapshot is the previous order with the new table appended. -`expected_snapshot` / `snapshot` are **human-readable names** (for example `"snap-baseline"`), not hashes of chunk bytes. `expected_snapshot` is an optimistic concurrency guard: if another publisher advances the active snapshot first, this job fails rather than overwriting that change. LitData does not compute content hashes of chunk payloads to form snapshot IDs. +#### Future layout-driven update (out of scope for V1) + +```python +# Post-V1 sketch — alignment by construction for updates +join_update_optimize( + output_dir=root, + base_snapshot="baseline", + table="metadata", + lookup_fn=lookup_entity, # key -> source payload + fn=build_metadata_v2, + version="v2", +) + +# layout position i -> key ki -> lookup_fn(ki) -> table builder +``` + +V1 still requires users to supply `inputs` in layout order. Documenting this future mode clarifies the long-term path without expanding V1 scope. + +`expected_snapshot` / `snapshot` / `base_snapshot` are **human-readable names**, not hashes of chunk bytes. + +#### Concurrency example + +```python +# Valid: base and expected are the same parent +join_update_optimize(..., base_snapshot="baseline", expected_snapshot="baseline", ...) + +# Rejected in V1: unclear inheritance / possible silent sibling loss +join_update_optimize(..., base_snapshot="baseline", expected_snapshot="metadata-v2", ...) +``` ### 7.4 Stream the active snapshot ```python -from litdata import MultiJoinStreamingDataset, StreamingDataLoader +from litdata import TablesStreamingDataset, StreamingDataLoader -dataset = MultiJoinStreamingDataset( - "s3://bucket/multi-join-dataset", - tables=("table_0", "table_1"), +dataset = TablesStreamingDataset( + "s3://bucket/tables-dataset", + tables=("images", "labels", "metadata"), shuffle=True, seed=42, drop_last=True, transform=lambda parts: { - **parts["table_0"], - **parts["table_1"], + **parts["images"], + **parts["labels"], + **parts["metadata"], }, max_cache_size="200GB", max_pre_download=4, @@ -323,8 +419,9 @@ Default output without `transform`: ```python { - "table_0": , - "table_1": , + "images": , + "labels": , + "metadata": , } ``` @@ -333,10 +430,10 @@ The table namespace is preserved by default. LitData does not implicitly merge d ### 7.5 Pin an exact snapshot ```python -dataset = MultiJoinStreamingDataset( - "s3://bucket/multi-join-dataset", - snapshot="snap-baseline", - tables=("table_0", "table_1"), +dataset = TablesStreamingDataset( + "s3://bucket/tables-dataset", + snapshot="baseline", + tables=("images", "labels", "metadata"), shuffle=True, seed=42, ) @@ -354,9 +451,9 @@ dataset.tables ### 7.6 Select only required tables ```python -dataset = MultiJoinStreamingDataset( +dataset = TablesStreamingDataset( root, - tables=("table_0", "table_1"), + tables=("images", "metadata"), ) ``` @@ -369,12 +466,12 @@ Unknown, duplicated, or inactive table names raise a clear error during construc Sampling options cannot vary by table. Decoding-specific options may: ```python -dataset = MultiJoinStreamingDataset( +dataset = TablesStreamingDataset( root, - tables=("table_0", "table_1"), + tables=("images", "metadata"), table_options={ - "table_0": {"encryption": table_0_key}, - "table_1": {"serializers": custom_serializers}, + "images": {"encryption": images_key}, + "metadata": {"serializers": custom_serializers}, }, ) ``` @@ -386,18 +483,19 @@ The following are always shared and cannot appear in `table_options`: - `drop_last` - `subsample` - epoch -- number of workers -- batch size +- number of DataLoader workers - distributed rank and world size +Batch size is owned by `StreamingDataLoader`, not by each child `StreamingDataset`. The invariant is that all children observe the same DataLoader / worker topology and therefore the same batch semantics. + ### 7.8 Validate without training ```python -from litdata import validate_multi_join +from litdata import validate_tables_dataset -report = validate_multi_join( - "s3://bucket/multi-join-dataset", - snapshot="snap-baseline", +report = validate_tables_dataset( + "s3://bucket/tables-dataset", + snapshot="baseline", deep=True, ) @@ -407,9 +505,9 @@ report.raise_for_errors() Validation levels: - Construction always performs mandatory constant-size metadata validation. -- `deep=False` verifies manifests, table indexes, counts, layout IDs, and alignment roots. +- `deep=False` verifies manifests, table indexes, counts, layout IDs, `table_order`, and alignment roots. - `deep=True` streams partition alignment metadata and verifies every ordered-key digest. -- Writer / `joint_optimize` success always performs the validation required to prove a newly published table matches the canonical layout. +- `join_optimize` / `join_update_optimize` success always performs the validation required to prove published table(s) match the canonical layout. ### 7.9 Keyed debugging access @@ -426,11 +524,11 @@ For integer entity keys, `get_by_key()` remains explicit so integer positional i The root is a versioned store: ```text -multi-join-dataset/ +tables-dataset/ join.json snapshots/ - snap-baseline.json - snap-table1-v2.json + baseline.json + metadata-v2.json layouts/ layout-entity-v1/ index.json @@ -439,13 +537,19 @@ multi-join-dataset/ shard-00001.parquet partitions.parquet tables/ - table_0/ + images/ v1/ index.json alignment.parquet chunk-0-0.bin ... - table_1/ + labels/ + v1/ + index.json + alignment.parquet + chunk-0-0.bin + ... + metadata/ v1/ index.json alignment.parquet @@ -464,9 +568,9 @@ multi-join-dataset/ ```json { - "format": "litdata-multi-join", + "format": "litdata-tables-dataset", "format_version": 1, - "active_snapshot": "snap-table1-v2", + "active_snapshot": "metadata-v2", "updated_at": "2026-08-17T10:42:11Z" } ``` @@ -475,10 +579,10 @@ multi-join-dataset/ ```json { - "format": "litdata-multi-join-snapshot", + "format": "litdata-tables-snapshot", "format_version": 1, - "snapshot_id": "snap-table1-v2", - "parent_snapshot_id": "snap-baseline", + "snapshot_id": "metadata-v2", + "parent_snapshot_id": "baseline", "created_at": "2026-08-17T10:42:10Z", "layout": { "id": "layout-entity-v1", @@ -489,18 +593,27 @@ multi-join-dataset/ "num_chunks": 488282, "alignment_root": "blake2b-256:..." }, + "table_order": ["images", "labels", "metadata"], "tables": { - "table_0": { + "images": { "version": "v1", - "path": "tables/table_0/v1", + "path": "tables/images/v1", "format": "litdata", "layout_id": "layout-entity-v1", "alignment_root": "blake2b-256:...", "schema_fingerprint": "sha256:..." }, - "table_1": { + "labels": { + "version": "v1", + "path": "tables/labels/v1", + "format": "litdata", + "layout_id": "layout-entity-v1", + "alignment_root": "blake2b-256:...", + "schema_fingerprint": "sha256:..." + }, + "metadata": { "version": "v2", - "path": "tables/table_1/v2", + "path": "tables/metadata/v2", "format": "litdata", "layout_id": "layout-entity-v1", "alignment_root": "blake2b-256:...", @@ -510,12 +623,28 @@ multi-join-dataset/ } ``` +`table_order` is explicit and immutable within a snapshot. JSON object key order is not relied upon for: + +- `ParallelStreamingDataset` child ordering; +- state-dict compatibility; +- transform / diagnostics determinism; +- manifest serialization stability. + +Rules: + +- Snapshot `table_order` lists every active table exactly once. +- Dataset `tables=(...)` is an ordered selection from that set. +- State dict records the requested table order; reordering selected tables is state-incompatible unless the caller explicitly starts a new data epoch. +- Default named output preserves the requested selection order. +- Adding a table appends it to `table_order` in the new snapshot. + All paths are relative to the root in V1. Manifest parsing rejects: - Absolute paths. - Parent traversal. - A different URI scheme or bucket. - Duplicate normalized table names. +- Missing / inconsistent `table_order`. - Unknown format versions. This keeps one snapshot within one trust and credential boundary. @@ -537,11 +666,41 @@ The key store must be sharded and streamed. Creating or validating a billion-row The canonical key order, rather than lexical key order, defines training positions. Key-index shards may be physically sorted or hash-partitioned for lookup as long as the stored `global_index`, `chunk_index`, and `chunk_offset` preserve the canonical order. +The layout key store must expose efficient sequential range iteration for update validation: + +```python +layout.iter_keys(start: int, stop: int) # contiguous global indexes +``` + +`get_by_key()` may use a separately optimized key-sorted or hash-partitioned index. That debugging path is distinct from sequential update validation. + +### 8.3.1 Initial key uniqueness (required before first publish) + +Encounter-order recording alone does not prove uniqueness. An accidental duplicate in the initial layout would still be reproducible by later updates and would match digests. + +Before an initial snapshot is published, the canonical layout builder must prove that every normalized entity key occurs **exactly once**. At billion-row scale this cannot use an in-memory Python set. Suitable approaches: + +1. **External sort validation** — emit `(normalized_key, global_index)` records, externally sort by key, detect adjacent equals, report both positions. +2. **Hash partition plus local sort** — hash-partition keys into bounded shards, sort each shard, detect duplicates within each shard. +3. **Existing scalable key-index infrastructure** — if LitData already has a distributed sorted key index, use it as the uniqueness authority. + +Canonical training order remains encounter order. Sorting is only for uniqueness verification and keyed lookup index construction. + +```text +TablesAlignmentError: + duplicate entity key: "entity-123" + first_global_index: 1024 + duplicate_global_index: 8439921 + layout_id: layout-entity-v1 +``` + +Update and add-table jobs do not recompute uniqueness when the base layout was already validated. + ### 8.4 Identifiers vs digests Three different identifiers appear in the format. They must not be confused: -- **Snapshot ID** (for example `snap-baseline`): a human-readable name for one published combination of table versions. Assigned when a `joint_optimize` job publishes successfully. It is **not** a hash of chunk bytes. +- **Snapshot ID** (for example `snap-baseline`): a human-readable name for one published combination of table versions. Assigned when a `join_optimize` or `join_update_optimize` job publishes successfully (create, replace, or add-table). It is **not** a hash of chunk bytes. - **Table version** (for example `v1`, `v2`): a human-readable name for one immutable build of a single table. - **Ordered-key digest / alignment root**: a compact fingerprint of the **canonical entity-key order** inside each logical bucket. Used only to prove tables are aligned. LitData does **not** compute content hashes of chunk payloads for this purpose. @@ -569,13 +728,13 @@ Every table version contains an `alignment.parquet` sidecar with: - `num_items` - `ordered_key_digest` -The table `index.json` gains a backward-compatible `multi_join` section: +The table `index.json` gains a backward-compatible `tables_dataset` section: ```json { - "multi_join": { + "tables_dataset": { "format_version": 1, - "table": "table_1", + "table": "metadata", "table_version": "v2", "layout_id": "layout-entity-v1", "length": 1000000000, @@ -591,22 +750,22 @@ Per-chunk digests live in compact Parquet rather than expanding an already large ### 9.1 Publication protocol -Each successful `joint_optimize(...)` publishes as part of the same job (no separate writer `commit()`): +Each successful `join_optimize` / `join_update_optimize` publishes as part of the same job (no separate writer `commit()`): 01. Read and retain the expected active snapshot when `expected_snapshot` is set. -02. Write new table data to a unique immutable version path. +02. Write new table data to unique immutable version path(s) under `tables/{table}/{version}/`. 03. Upload all chunk objects (existing multi-worker / multi-node optimize upload path). -04. Upload table alignment metadata. -05. Upload the table `index.json` and completion marker last. -06. If this is the first table in an empty join root, publish the shared layout from the ordered `key_fn` stream. -07. Validate the table against the canonical layout. -08. Write a new immutable snapshot document that merges this table version with sibling versions from the previous active snapshot. +04. Upload table alignment metadata for every rewritten table. +05. Upload each rewritten table `index.json` and completion marker last. +06. For an initial `join_optimize`, publish the shared layout from the single ordered `key_fn` stream over canonical inputs. +07. Validate rewritten table(s) against the layout (position checks during update; counts + per-chunk ordered-key digests at completion). +08. Write a new immutable snapshot document: the base snapshot's table set plus newly written or replaced table versions (`table_order` preserved; new tables appended). 09. Recheck the active snapshot or storage generation. 10. Atomically replace `join.json` with the new active snapshot. -If any operation fails before step 10, the active snapshot remains unchanged. Unreferenced objects are safe to garbage-collect later. +If any operation fails before step 10, the active snapshot remains unchanged. Failed jobs may leave unreferenced staging objects under unique version prefixes; they are never considered valid without a completion marker and snapshot reference. Garbage collection removes them only after the configured safety window. -Because each `joint_optimize` is an independent distributed job, publication must be safe across machines and processes. There is no process-local `MultiJoinWriter` session holding uncommitted state. +Jobs remain distributed-capable like `optimize`. There is no process-local writer session holding uncommitted state across tables outside the single job. ### 9.2 Backend behavior @@ -619,7 +778,7 @@ The storage abstraction must expose the precondition required for `expected_snap ### 9.3 Reader behavior -`MultiJoinStreamingDataset` resolves the active snapshot exactly once during construction. It never polls `join.json` during iteration. +`TablesStreamingDataset` resolves the active snapshot exactly once during construction. It never polls `join.json` during iteration. Therefore: @@ -641,7 +800,7 @@ In V1: - One physical LitData chunk is one logical sampling bucket. - Corresponding table chunks contain the same keys in the same order. - Chunk byte sizes, schemas, compression, serializers, and payload types may differ. -- `MultiJoinStreamingDataset` is a validated named wrapper around `ParallelStreamingDataset`. +- `TablesStreamingDataset` is a validated named wrapper around `ParallelStreamingDataset`. ### 10.2 V1 hard invariants @@ -666,38 +825,95 @@ For the read configuration: 4. All children receive the same `drop_last`. 5. All children receive the same `subsample`. 6. All children see the same distributed environment. -7. All children receive the same DataLoader worker count and batch size. +7. All children observe the same DataLoader worker topology and batch semantics (batch size is owned by the DataLoader). 8. A loaded state dict references the same snapshot and ordered table list. ### 10.3 V1 write path -`joint_optimize(...)` wraps the existing optimize pipeline with required safe settings: +#### Canonical input planning (required capability) -- `chunk_size` is required and becomes the shared logical bucket size. -- `align_chunking=True`. -- `reorder_files=False`. -- `keep_data_ordered=True`. -- Static ordered inputs. -- One output sample per input. -- `key_fn` required so alignment digests can be computed (same mechanism as `optimize(..., key_fn=...)`). -- Output lands under `tables/{table}/{version}/` inside the join root rather than a standalone dataset root. +The indexing invariant below is not implementable for an arbitrary Python generator without materializing or indexing the input stream. V1 therefore requires inputs that support deterministic planning: -The following are hidden or rejected: +1. A deterministic **indexable sequence** (list-like: `__len__` + positional access); or +2. A deterministic **shard manifest** whose shard order and per-shard item order are fixed, from which LitData can derive canonical global-index ranges before workers execute; or +3. A previously materialized **canonical input plan** produced by LitData. -- `chunk_bytes` -- `weights` -- shared dynamic work queue -- `keep_data_ordered=False` -- `reorder_files=True` -- filtering with `None` -- variable-yield generators -- append into an existing table version -- overwrite of a published table version -- a `key_name` / schema-field argument +**Policy:** distributed `join_optimize` / `join_update_optimize` requires deterministic static inputs. Arbitrary generators / streaming iterables are rejected unless running locally with `num_workers=1` and `num_nodes=1`. + +Optional helpers (same semantics): + +```python +plan = create_canonical_input_plan( + inputs=inputs, + key_fn=lambda item: item["entity_id"], + output_dir=root, +) +join_optimize(inputs=inputs, input_plan=plan, ...) +``` -`align_chunking=True` is important because it makes logical chunk boundaries independent of the number of optimize workers or nodes. Workers receive complete item-count chunks, and rank-index merge order reconstructs the canonical sequence. +Without this contract, preassigning `work unit 0 -> positions 0..16383` remains aspirational rather than enforceable. -A small table may therefore be rebuilt with 8 workers while a large table was originally built with 256 workers, provided both consume the same canonical ordered inputs and logical `chunk_size`. Because each call is a normal optimize-style job, `num_nodes` / Studio multi-node execution works the same way as `optimize`. +#### Canonical distributed indexing (hard invariant) + +Before table builders execute, `join_optimize` / `join_update_optimize` deterministically assigns **canonical global-index ranges and complete logical chunks** to work units from the input plan. A retry retains the same assigned range. Worker completion order never determines canonical layout order or chunk numbering. + +```text +canonical input positions: + work unit 0 -> chunks 0..7 + work unit 1 -> chunks 8..15 + work unit 2 -> chunks 16..23 +``` + +Not: + +```text +worker 0 -> whichever inputs it happens to receive next +worker 1 -> whichever inputs it happens to receive next +``` + +Each distributed work unit owns fixed ranges of global positions. Within a range, the worker invokes every selected table builder for each input item and writes each table contribution at the same logical position. Outputs are merged by canonical chunk index, never task completion order. + +V1 must not rely on workers independently calling `enumerate(inputs)`, queue scheduling, task stealing, shard-discovery order, or completion order to define layout positions. + +#### `join_optimize` (initial multi-table build) + +One optimize-style job opens one writer per table under `tables/{table}/{version}/` and advances them together: + +- Deterministic partition of planned canonical input positions into complete chunk ranges. +- Shared `chunk_size` flush points: when the item count hits `chunk_size`, every table writer closes its current chunk. +- Shared `key_fn(input_item)` stream recorded once into the layout. +- Scalable uniqueness validation before first publish (see 8.3.1). +- `align_chunking=True`, `reorder_files=False`, `keep_data_ordered=True`. +- `num_nodes` / Studio multi-node execution works like `optimize`. +- Conditional root initialization (exclusive create / lease / generation precondition). + +Because chunk boundaries and indexes are assigned before builders run, tables cannot drift. + +#### `join_update_optimize` (single-table rewrite or add) + +A later job rewrites or adds one table subpath: + +- Loads the existing layout's length, `chunk_size`, and alignment metadata. +- Uses the same canonical range assignment model as create. +- Streams `layout.iter_keys(start, stop)` for each assigned range and compares against `key_fn(input_item)` sequentially (fail immediately on first mismatch). +- Writes `tables/{table}/{new_version}/` with the same chunk item counts. +- Rejects the publish if any position, digest, or length differs from the layout. +- Does not recompute global uniqueness against the base layout. + +This is how a schema change on one table avoids rewriting siblings while still keeping native bucket sampling. Adding a table uses the same validation path and snapshot-merge publication. + +#### Rejected / hidden settings + +- `chunk_bytes` +- Independent per-table input lists as the primary create API +- Dual create APIs (`fn` returning a mapping **and** `fns=...`) — V1 uses only `tables={...}` +- Distributed jobs over arbitrary generators / non-plannable streaming iterables +- `weights` / shared dynamic work queues that reorder items +- `keep_data_ordered=False`, `reorder_files=True` +- filtering / variable-yield generators / builders returning `None` +- append into an existing table version / overwrite of a published version +- a `key_name` / schema-field argument +- `base_snapshot != expected_snapshot` when both are set ### 10.4 V1 read path @@ -719,7 +935,7 @@ Iteration: 3. Existing in-chunk shuffle computes the same item permutation. 4. Each child `BinaryReader` prefetches its corresponding chunks. 5. `ParallelStreamingDataset` pulls one aligned value from each child. -6. `MultiJoinStreamingDataset` creates the named table mapping. +6. `TablesStreamingDataset` creates the named table mapping. 7. The optional transform creates the final training sample. No key lookup occurs on this path. @@ -784,7 +1000,7 @@ cache/ This prevents same-named LitData chunk files from different tables from colliding. -`max_cache_size` is defined as the aggregate user budget. V1 approximates this by assigning per-table budgets in proportion to average chunk bytes, with a documented minimum and an optional per-table override. +In V1, `max_cache_size` is an **aggregate target** rather than a strict global limit. The wrapper derives per-table child budgets from table chunk-size estimates. Temporary overage is possible because existing child readers retain independent prefetch, refcount, and eviction behavior. V2 introduces strict shared aggregate cache ownership. Peak in-flight storage is approximately: @@ -814,11 +1030,11 @@ These limitations are the primary motivation for V2. Introduce specific public exceptions: -- `MultiJoinError` -- `MultiJoinManifestError` -- `MultiJoinAlignmentError` -- `MultiJoinSnapshotMismatchError` -- `MultiJoinCommitConflictError` +- `TablesDatasetError` +- `TablesManifestError` +- `TablesAlignmentError` +- `TablesSnapshotMismatchError` +- `TablesCommitConflictError` An alignment error includes: @@ -849,29 +1065,34 @@ Do not add a heavy schema dependency solely for these small metadata models. #### Write API -Add `joint_optimize` in `src/litdata/processing/functions.py` (or a thin wrapper module re-exported from `__init__.py`): +Add `join_optimize`, `join_update_optimize`, and `TableBuild` in `src/litdata/processing/functions.py` (re-exported from `__init__.py`): - Same distributed execution path as `optimize` (`DataProcessor`, `num_nodes`, uploaders, index merge). -- Extra args: `table`, optional `version`, `snapshot`, `expected_snapshot`. +- `join_optimize`: `tables={name: TableBuild|callable}`, deterministic canonical chunk-range assignment, N synchronized writers, shared layout publish. +- `join_update_optimize`: one table writer (replace or add), position-by-position layout key checks, digest checks, snapshot merge publish. - Force aligned write settings listed above. -- After merge, validate against the join layout and publish/update the snapshot atomically. -- No `MultiJoinWriter` context manager and no process-local uncommitted session. +- No writer context manager. Extend the optimize internals where necessary to: -- Stream ordered `(global_index, key)` metadata from `key_fn`. -- Compute per-chunk key digests. -- Avoid materializing all keys during alignment generation. -- Attach the backward-compatible `multi_join` section to `index.json`. -- Create the shared layout on first publish into an empty join root. +- Assign canonical global-index / chunk ranges to work units before builders run; retries keep the same ranges; merge by chunk index. +- Fan out one input item to multiple table caches / writers. +- Synchronize chunk closes across table writers on item-count boundaries. +- Call `key_fn` once per input item before builders; stream ordered `(global_index, key)` into the layout. +- On update, compare keys at every global index and fail immediately on mismatch. +- Attach the backward-compatible `tables_dataset` section to each table `index.json`. +- Create the shared layout during the initial `join_optimize`. + +Public modules may keep internal filenames such as `multi_join.py` initially; public names use the tables vocabulary. Likely integration points: - `src/litdata/processing/functions.py` - `src/litdata/processing/data_processor.py` - `src/litdata/streaming/writer.py` +- `src/litdata/streaming/cache.py` - `src/litdata/utilities/keys_index.py` -- `src/litdata/utilities/multi_join.py` +- `src/litdata/utilities/multi_join.py` (internal; public API is tables-named) #### Read API @@ -928,7 +1149,7 @@ Cover: - Full and partial final chunks. - Different table payload byte sizes. - Different schemas and serializers. -- Different compression settings. +- Different compression settings via `TableBuild`. - Different optimize worker counts. - Different optimize node counts. - Stable boundaries with `align_chunking=True`. @@ -936,7 +1157,7 @@ Cover: - Missing key. - Duplicate key. - Extra key. -- `fn` returning `None`. +- Builder returning `None`. - Variable-yield generator rejection. - Corrupt layout ID. - Corrupt per-chunk digest. @@ -944,6 +1165,32 @@ Cover: - Incomplete table version. - Manifest path traversal. - Unsupported manifest format version. +- Add-table snapshot merge; old snapshot cannot select the new table. +- Explicit `table_order` vs selection order / state incompatibility. +- Initial duplicate key detection (external-sort / hash-partition path). +- Distributed rejection of non-plannable generators. +- `base_snapshot != expected_snapshot` rejection. +- Concurrent initializer conflict. + +Deterministic distributed write ordering: + +- Workers complete canonical chunks in reverse order. +- A failed chunk is retried after later chunks complete. +- Different node counts produce identical layout digest and aligned chunk numbering. +- Worker assignment differs between create and update jobs. +- Source shard listing order differs but explicit canonical input position assignment remains stable. + +Early update validation: + +- First key mismatches at global index 0. +- Key mismatches at a chunk boundary. +- Key mismatches in the final partial chunk. +- Input ends before layout end. +- Input has one extra entity after layout end. +- Repeated key appears where a later canonical key is expected. +- Update builder emits `None`. +- Update builder raises midway through a chunk. +- Failure leaves no completion marker and no active-manifest change. Read-path matrix: @@ -969,7 +1216,7 @@ Snapshot tests: - A reader opened after commit sees the new snapshot. - `table_0` objects are not rewritten by a `table_1`-only update. - Rollback activates the previous immutable snapshot. -- Concurrent stale `joint_optimize` publish fails. +- Concurrent stale `join_optimize` / `join_update_optimize` publish fails. - Failure before active-manifest publication leaves the old snapshot active. Remote tests: @@ -1011,8 +1258,8 @@ No production throughput claim should be made before this benchmark runs on a re V1 is complete when: -1. A two-table and a many-table dataset can be created through independent `joint_optimize` jobs. -2. Re-optimizing one table does not write under unchanged table-version prefixes. +1. A two-table and a many-table dataset can be created through one `join_optimize` job. +2. Updating one table with `join_update_optimize` (or adding a table) does not write under unchanged sibling table-version prefixes. 3. A new table snapshot is atomically activated. 4. Old readers continue without observing mixed versions. 5. Every tested shuffle, worker, DDP, and resume configuration preserves key alignment. @@ -1055,7 +1302,7 @@ V2 allows many small-table buckets to be packed into one appropriately sized obj ### 11.3 V2 storage model ```text -multi-join-dataset/ +tables-dataset/ join.json snapshots/ layouts/ @@ -1128,9 +1375,9 @@ V2 replaces the V1 `ParallelStreamingDataset` composition internally with a shar Potential internal components: -- `MultiJoinChunksConfig` -- `MultiJoinShuffle` -- `MultiJoinReader` +- `TablesChunksConfig` +- `TablesShuffle` +- `TablesReader` - `PrepareJoinPartitionsThread` - `ColumnFamilyLoader` - `LitDataColumnFamilyLoader` @@ -1152,47 +1399,41 @@ The shared reader owns one aggregate cache budget: This eliminates V1’s approximate per-child budget split. -### 11.7 V2 `joint_optimize` storage extension +### 11.7 V2 storage extension on `join_optimize` / `join_update_optimize` -`joint_optimize` gains optional per-table storage configuration: +Per-table physical packing lives on `TableBuild` / update kwargs while the logical layout stays shared: ```python -from litdata import joint_optimize, TableStorage +from litdata import TableBuild, join_optimize, join_update_optimize, TableStorage -joint_optimize( - fn=build_table_1_v3, - inputs=inputs_1, +join_optimize( + inputs=inputs, output_dir=root, - table="table_1", + key_fn=lambda item: item["entity_id"], chunk_size=2048, - key_fn=get_id, - storage=TableStorage( - format="parquet", - target_chunk_bytes="128MB", - pack_logical_chunks=True, - ), + tables={ + "images": TableBuild( + fn=build_images, + storage=TableStorage(format="litdata", target_chunk_bytes="256MB", pack_logical_chunks=False), + ), + "metadata": TableBuild( + fn=build_metadata, + storage=TableStorage(format="parquet", target_chunk_bytes="128MB", pack_logical_chunks=True), + ), + }, ) -``` -Large table (`table_0`): - -```python -joint_optimize( - fn=build_table_0, - inputs=inputs_0, +join_update_optimize( + inputs=inputs, output_dir=root, - table="table_0", - chunk_size=2048, - key_fn=get_id, - storage=TableStorage( - format="litdata", - target_chunk_bytes="256MB", - pack_logical_chunks=False, - ), + table="metadata", + fn=build_metadata_v3, + key_fn=lambda item: item["entity_id"], + storage=TableStorage(format="parquet", target_chunk_bytes="128MB", pack_logical_chunks=True), ) ``` -The normal `MultiJoinStreamingDataset` construction remains unchanged. +The normal `TablesStreamingDataset` construction remains unchanged. ### 11.8 V2 variable-cardinality column families @@ -1338,21 +1579,20 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam ### One-time migration -01. Choose the entity key extraction (`key_fn`). -02. Produce unique ordered inputs for every table (same entity order). -03. Intentionally choose that training order before the first `joint_optimize`. +01. Choose `key_fn` over the **source input item** (`entity_id` or equivalent). +02. Build one ordered `inputs` stream that can feed every table builder. +03. Intentionally shuffle / curate that order before `join_optimize` when source order has structure. 04. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. -05. Run `joint_optimize` for the first / driver table to create the join root and layout. -06. Run independent `joint_optimize` jobs for the remaining tables (local or multi-node). -07. Validate the published snapshot. -08. Benchmark against the current baked dataset. +05. Run one `join_optimize` job with `tables={name: TableBuild(...)}` (local or multi-node). +06. Validate the published snapshot with `validate_tables_dataset`. +07. Benchmark against the current baked dataset. -### Small-table schema update +### Small-table schema update or add-table 01. Keep existing sibling table versions as-is. -02. Run `joint_optimize(..., table=..., version=..., expected_snapshot=...)` for the changed table only. +02. Run `join_update_optimize(..., table=..., version=..., base_snapshot=..., expected_snapshot=...)` with inputs in layout order and `key_fn(input_item)` checks. 03. Start new training jobs on the newly published snapshot. -04. Leave existing jobs pinned to their original snapshot. +04. Leave existing jobs pinned to their original snapshot (they cannot see newly added tables). ### Rollback @@ -1387,7 +1627,7 @@ A future append-only layout extension could preserve complete existing buckets a Initial behavior: -- Successful `joint_optimize` jobs never delete previous table versions. +- Successful `join_optimize` / `join_update_optimize` jobs never delete previous table versions. - Rollback remains possible while snapshots and versions are retained. - Failed staging outputs are recorded as unreachable. @@ -1399,7 +1639,7 @@ Later garbage collection: 4. Delete only unreferenced immutable prefixes. 5. Support dry-run output before deletion. -Garbage collection is never part of `joint_optimize` publication. +Garbage collection is never part of `join_optimize` / `join_update_optimize` publication. ### Observability @@ -1448,20 +1688,21 @@ Risk: independent V1 child readers each reserve and prefetch data. Mitigation: - Namespace caches. -- Interpret the public cache budget as aggregate. +- Treat the public cache budget as an aggregate target in V1; strict shared ownership is V2. - Allocate proportional child budgets. - Move to one shared cache coordinator in V2. ### Concurrent publisher race -Risk: two updates overwrite each other’s active snapshot. +Risk: two updates overwrite each other’s active snapshot, or two initializers create different layouts. Mitigation: -- `expected_snapshot`. -- Conditional object-store publication. +- `base_snapshot == expected_snapshot` in V1 (no automatic sibling merge). +- Conditional object-store publication / CAS. +- Exclusive create or lease for first layout initialization. - Immutable snapshots. -- Explicit commit-conflict error. +- Explicit `TablesCommitConflictError`. ### Metadata size at billion-row scale @@ -1477,15 +1718,15 @@ Mitigation: ### Optimize worker-count dependence -Risk: independent table runs use different worker or node counts. +Risk: a `join_update_optimize` run uses different worker or node counts than the original `join_optimize`. Mitigation: -- Force item-count chunking. -- Force ordered static inputs. +- Force item-count chunking from the layout. +- Force ordered static inputs and `key_fn` position checks. - Use `align_chunking=True`. -- Verify merged chunk counts and digests. -- Test all worker-count combinations. +- Verify merged chunk counts and ordered-key digests against the layout. +- Test update jobs with different worker / node counts than the create job. ### Manifest or path injection @@ -1509,10 +1750,21 @@ Mitigation: 07. V2 introduces one shared sampler and coordinated reader. 08. Table versions and snapshots are immutable. 09. `join.json` is published last and readers pin one snapshot. -10. Each `joint_optimize` job validates and publishes atomically; no writer context / deferred `commit()`. -11. Training iteration performs no per-sample key lookup. -12. The public `MultiJoinStreamingDataset` API remains stable across V1 and V2. -13. Keys come from `key_fn`, not from a `key_name` schema-field argument. +10. `join_optimize` writes all selected tables from one input stream with deterministic canonical chunk-range assignment so alignment is by construction. +11. `join_update_optimize` replaces or adds one table and validates against the persisted layout (position checks + digests) before publishing. +12. Each job validates and publishes atomically; no writer context / deferred `commit()`. +13. Training iteration performs no per-sample key lookup. +14. The public `TablesStreamingDataset` API remains stable across V1 and V2. +15. `key_fn` always receives the original input item; table outputs need not contain the key. +16. Public naming uses the tables vocabulary (`join_update_optimize`, `validate_tables_dataset`, `TableBuild`, `Tables*Error`). +17. Snapshots store explicit `table_order`. +18. Builder `None` is always rejected; empty contributions use explicit empty values. +19. Distributed V1 writes require deterministic static input planning; arbitrary streaming iterables cannot define canonical positions across workers. +20. Update validation streams expected canonical keys sequentially by assigned global-index range; it never performs per-item remote key-index lookup. +21. Initial layout publication proves normalized entity-key uniqueness using a scalable external-sort or hash-partitioned validation pass. +22. In V1, `base_snapshot` and `expected_snapshot` identify the same parent snapshot; stale updates fail rather than automatically merging snapshot changes. +23. `input_item` is shared read-only across builders; mutation inside a builder is unsupported. +24. Initial root creation is exclusive/conditional; losing initializers raise `TablesCommitConflictError`. ## 17. Recommended delivery order @@ -1520,8 +1772,8 @@ Mitigation: 01. Freeze manifest, layout, digest, and API specifications. 02. Implement metadata models and validation. -03. Implement scalable canonical layout creation. -04. Implement `joint_optimize` with strict aligned optimize settings and multi-node support. +03. Implement scalable canonical layout creation, input planning, uniqueness validation, and `layout.iter_keys(start, stop)`. +04. Implement `TableBuild`, `join_optimize` (canonical ranges + synchronized writers), and `join_update_optimize` (replace/add with layout validation). 05. Implement atomic immutable snapshots. 06. Implement the validated `ParallelStreamingDataset` wrapper. 07. Add snapshot-aware state and resume. From 78cbe825f0226ba8c1fa7ee5297495e9f90ea802 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:27:30 +0000 Subject: [PATCH 9/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .claude/plans/multi_join_streaming_dataset.md | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/.claude/plans/multi_join_streaming_dataset.md b/.claude/plans/multi_join_streaming_dataset.md index ccae6490..5415c33d 100644 --- a/.claude/plans/multi_join_streaming_dataset.md +++ b/.claude/plans/multi_join_streaming_dataset.md @@ -279,10 +279,10 @@ Separate input lists or separate distributed schedules can diverge in order, fil ### 7.3 Update or add one table with `join_update_optimize` -| Operation | Alignment property | -|---|---| -| `join_optimize` initial build | Guaranteed by one canonical traversal and synchronized writers | -| `join_update_optimize` | Validated against the persisted canonical layout | +| Operation | Alignment property | +| ------------------------------------- | ----------------------------------------------------------------- | +| `join_optimize` initial build | Guaranteed by one canonical traversal and synchronized writers | +| `join_update_optimize` | Validated against the persisted canonical layout | | Future layout-driven update (post-V1) | Guaranteed by streaming persisted layout keys through `lookup_fn` | `join_update_optimize` does **not** infer or repair entity order. It validates that the provided input stream reproduces the persisted layout exactly. A key mismatch, missing entity, duplicate displacement, extra entity, or different final length fails the job before snapshot publication. @@ -1579,20 +1579,20 @@ V2 does not accept arbitrary unrelated table sharding and repair it with per-sam ### One-time migration -01. Choose `key_fn` over the **source input item** (`entity_id` or equivalent). -02. Build one ordered `inputs` stream that can feed every table builder. -03. Intentionally shuffle / curate that order before `join_optimize` when source order has structure. -04. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. -05. Run one `join_optimize` job with `tables={name: TableBuild(...)}` (local or multi-node). -06. Validate the published snapshot with `validate_tables_dataset`. -07. Benchmark against the current baked dataset. +1. Choose `key_fn` over the **source input item** (`entity_id` or equivalent). +2. Build one ordered `inputs` stream that can feed every table builder. +3. Intentionally shuffle / curate that order before `join_optimize` when source order has structure. +4. Choose the logical bucket item count based primarily on the largest table's bytes and desired bucket sampling. +5. Run one `join_optimize` job with `tables={name: TableBuild(...)}` (local or multi-node). +6. Validate the published snapshot with `validate_tables_dataset`. +7. Benchmark against the current baked dataset. ### Small-table schema update or add-table -01. Keep existing sibling table versions as-is. -02. Run `join_update_optimize(..., table=..., version=..., base_snapshot=..., expected_snapshot=...)` with inputs in layout order and `key_fn(input_item)` checks. -03. Start new training jobs on the newly published snapshot. -04. Leave existing jobs pinned to their original snapshot (they cannot see newly added tables). +1. Keep existing sibling table versions as-is. +2. Run `join_update_optimize(..., table=..., version=..., base_snapshot=..., expected_snapshot=...)` with inputs in layout order and `key_fn(input_item)` checks. +3. Start new training jobs on the newly published snapshot. +4. Leave existing jobs pinned to their original snapshot (they cannot see newly added tables). ### Rollback