mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
Sravan Avvaru
a615306f39
Closes #3704 ## Problem Transforms can fail on bad data (e.g. nulls/NaNs from incomplete user surveys). Today any transform exception aborts iteration, and there is no way to skip invalid rows during loading. ## Solution New `on_transform_error` parameter on `StreamingDataset`: - `"raise"` (default, matches current behavior and the convention in tf.data / WebDataset / Ray Data) - `"skip"` — drop the failing rows and continue - `"warn"` — like skip, plus a logged warning per failing batch - a WebDataset-style callable `handler(exc) -> bool`, so users can skip only expected error types Key design points: - **Row-granular skipping**: when a batch fails, the transform is re-run on single-row slices so only the rows that actually fail are dropped (avoids Ray-style whole-block loss). Skips are counted in a new `rows_skipped` property. - **No crash on uneven skips**: the round-robin loop now ends the epoch at the last cycle where every split still has a row, instead of hitting `IndexError` when a split runs dry early. - **Exact resumability under skips**: checkpoints are now position-based. `state_dict` gains `positions_consumed_per_split` (exact for owned splits), and a new `merge_state_dicts` static method combines per-rank states via elementwise max for elastic resume across topology changes. Old checkpoints without the new key still load. Positions equal sample counts when nothing is skipped, so existing behavior is unchanged. - **Guardrail**: transforms returning the wrong number of rows now raise a clear `ValueError` instead of silently corrupting split accounting. ### Answers to the issue's open questions - *Can we do this?* Yes — all transforms funnel through one guarded call in the Stage 2 pipeline. - *What do other libraries do?* tf.data `ignore_errors()`, WebDataset `handler=`, Ray `max_errored_blocks`; MosaicML StreamingDataset offers nothing (skipping conflicts with its determinism model). This design follows the common conventions: raise by default, opt-in skipping, count/log drops. - *Error handling or pre-filtering?* Both: the existing `filter=` remains the recommended tool for predictable bad data (splits are built post-filter, so all guarantees hold — now documented); `on_transform_error` covers failures not expressible as a predicate. - *Impact on splits / elastic determinism?* Per-split sample sequences stay deterministic (skips are data-dependent, not topology-dependent). With unequal bad-row counts across splits the last few global steps of an epoch can differ across topologies (bounded by the skew), which is documented on the parameter. With equal counts per split, full determinism is preserved — covered by a test. ## Testing 15 new tests in `test_elastic_dataloader.py` covering: default raise, invalid values, uniform and uneven skips (including epoch-end truncation), warn logging, selective callable handlers, wrong-row-count guardrail, determinism across runs and across world sizes (1/2/3/4) with skips, exact mid-epoch resume with skips on the same topology, elastic resume via `merge_state_dicts` (ws=2 → ws=1), merge validation, and backward-compat loading of old checkpoints. Note: relying on CI for the test run — my local machine OOMs during the final link of the native extension. The change itself is pure Python. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>