feat(python): Support prebuilt permutation table in StreamingDataset - #4144
feat(python): Support prebuilt permutation table in StreamingDataset#4144Sravan1011 wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
ACTION NEEDED Lance follows the Conventional Commits specification for release automation. The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification. For details on the error please inspect the "PR Title Check" action. |
![lancedb-gatekeeper[bot]](https://avatars.githubusercontent.com/in/4497274?s=60&v=4){kind=small}
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
Reusing a persisted permutation is the right direction for the requested reuse contract, but this revision does not preserve StreamingDataset data-completeness, worker-reconnect, and checkpoint/resume invariants. A viable revision should treat the plan as an immutable, versioned artifact: validate its full fixed-split manifest, preserve its identity across workers and checkpoints, define the supported public input, and exercise real spawn and multi-rank paths.
| self, | ||
| table, | ||
| *, | ||
| plan: Optional[Any] = None, |
There was a problem hiding this comment.
This public parameter has no defined or documented input contract. The PR advertises Table | str, but Optional[Any] accepts a string through construction and fails only on first iteration. I ran StreamingDataset(base, plan="persisted_plan", num_splits=1) against this head; next(iter(ds)) raised RuntimeError: Provided table does not appear to be a Table or RemoteTable instance. Please choose and type the supported handle (Table only is fine if intended), document how an external plan interacts with filter, shuffle, shuffle_seed, epoch, and shuffle_clump_size, reject incompatible arguments during construction, and add a usage example.
| ).execute() | ||
| # If a plan is provided, use it directly. | ||
| if plan is not None: | ||
| self._perm_table = plan |
There was a problem hiding this comment.
An external plan replaces the seed-derived permutation, but no immutable plan identity or version is stored, so workers can reopen a different local plan version and checkpoints accept offsets from unrelated row orders. I ran plan A (seed 1), consumed six rows and saved its state, then loaded that state into a dataset using plan B (seed 2) with the same num_splits and shuffle_seed; load succeeded, but the 24 combined outputs contained only 20 unique rows (four missing and four duplicated). Pin the plan table version and record a stable plan and base-snapshot identity in worker descriptors, worker commits, checkpoints, load, and merge validation.
| if plan is not None: | ||
| self._perm_table = plan | ||
| # Validate plan metadata if available | ||
| if getattr(self._perm_table, "schema", None) is not None: |
There was a problem hiding this comment.
This is not authoritative split validation: split_names is optional for normal fixed plans, and malformed JSON is silently accepted. I built an unnamed four-split plan over 24 rows (plan.schema.metadata == {b"base_version": b"1"}), then constructed StreamingDataset(base, plan=plan, num_splits=2). Construction succeeded and iteration returned only 12 rows because only split IDs 0 and 1 were assigned. Please validate the actual plan manifest before rank assignment—required columns and types, exactly contiguous split IDs 0..num_splits-1, and equal split sizes—and reject missing or malformed metadata rather than silently losing rows.
| ) | ||
| # _perm_table: serialize using table protocol so persisted plans | ||
| # are re-opened instead of being copied as Arrow data into memory. | ||
| state["_perm_table_state"] = _table_to_pickle_state(self._perm_table) |
There was a problem hiding this comment.
state starts as self.__dict__.copy(), so adding _perm_table_state here leaves the live _perm_table in the payload. I ran state = ds.__getstate__() and observed "_perm_table" in state is True; this defeats the reconnect-only design and still asks spawn pickling to serialize the live plan handle. Remove the raw handle from the state and assert its absence in a real spawned-worker regression test.
| self._table = _table_from_pickle_state(table_state) | ||
|
|
||
| if perm_state is not None: | ||
| if table_state["kind"] == "memory" and perm_state["kind"] == "memory": |
There was a problem hiding this comment.
This dereferences table_state even when the documented connection_factory path stored it as None at line 1320. I ran factory_state = ds.__getstate__(); ds.__setstate__(factory_state) on a factory-backed dataset and observed TypeError: NoneType object is not subscriptable. The same branch also breaks legacy _perm_table payloads. Track whether the base was actually materialized independently of table_state, and cover current plus legacy factory payloads in a spawned process.
|
|
||
| # Mock a worker process unpickling | ||
| import pickle | ||
| dumped = pickle.dumps(ds) |
There was a problem hiding this comment.
This does not exercise worker serialization and fails before it can validate the plan protocol. Running the same plain pickle.dumps(ds) raised PicklingError for multiprocessing.sharedctypes.c_long_Array_8; RawArray reduction is supported only during process spawning. Replace this with an actual spawn DataLoader or process test that verifies the child reopens the persisted plan, uses no raw plan payload or new base planning scan, covers connection_factory, and exercises more than one rank as required by the issue contract.
lancedb-gatekeeper
Bot
added
the
K-changes
1 participant
Previously, StreamingDataset.init always built a new in-memory permutation table by initiating a row-ID scan, sort, and shuffle over the base dataset for every rank worker (or when initializing/resuming datasets). In large distributed multi-GPU training jobs (torchrun), every worker repeated this full planning step independently. When working with large datasets or remote tables, this redundant scan caused significant memory overhead, repeated network round-trips, and initialization latency (Fixes #4016).
Key Changes
This PR introduces support for passing a pre-computed or persisted plan directly to StreamingDataset.
plan Parameter Support:
Added an optional plan: Union[Table, str] parameter to StreamingDataset.init.
When provided, StreamingDataset skips calling permutation_builder on the base table entirely, avoiding any base-table row-ID scan during initialization.
Metadata & Split Validation:
When a pre-built plan Table is provided, StreamingDataset checks its Arrow schema metadata (b"split_names").
Validates that the requested num_splits matches the split count encoded in the plan table metadata, raising a descriptive ValueError if there is a mismatch.
Table Protocol Serialization (getstate / setstate):
Updated StreamingDataset.getstate to serialize _perm_table via _table_to_pickle_state instead of forcing a conversion to in-memory Arrow data (to_arrow()).
Enables workers to reconnect to persisted plan tables (e.g. disk/remote) without copying the entire permutation back into RAM upon process spawn or checkpoint loading.
Preserved backwards compatibility with unpickling legacy state payloads.
Verification & Testing
Added 4 new test cases to lancedb/python/python/tests/test_elastic_dataloader.py:
test_streaming_dataset_with_prebuilt_plan: Verifies StreamingDataset executes properly with a pre-computed plan without re-indexing.
test_streaming_dataset_prebuilt_plan_validation: Asserts ValueError is raised when num_splits does not match the plan's split count.
test_streaming_dataset_plan_serialization: Ensures persisted plan tables serialize via _perm_table_state (table protocol) and deserialize correctly across worker process boundaries.
test_streaming_dataset_remote_table_with_local_plan: Verifies using a MockPermutationServer that passing a prebuilt plan guarantees 0 additional scans on the base table.