From e56150568c428bd1d7d17c2606bb38b19ab1e3fb Mon Sep 17 00:00:00 2001 From: importcpp Date: Sat, 31 Oct 2026 12:32:16 +0800 Subject: [PATCH] feat(python): fill failed streaming transform rows Add `on_transform_error="fill"` to `StreamingDataset`, replacing failed transform rows at their original positions with a deep-copied fill value or the result of `repair(exc, row_batch)`. Fill mode preserves split length, elastic step layout, worker accounting, and checkpoint progress. Repair failures remain fatal, and `rows_filled` is shared across workers. Validated with focused fill and adjacent dataloader tests (10 and 6 passed) and the complete applicable CI run for 6b51b33e560ca9f4b6addd9713804796074fcf2e: four workflows and fifteen checks all succeeded. Closes #3916 --- docs/src/python/python.md | 21 ++ python/python/lancedb/streaming.py | 85 +++++++-- .../python/tests/test_elastic_dataloader.py | 179 +++++++++++++++++- 3 files changed, 267 insertions(+), 18 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 60746a0f0..b3d4efa95 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -778,6 +778,27 @@ so the fetch APIs cannot read their bytes yet. ::: lancedb.streaming.StreamingDataset +When a transform fails for some rows, fill mode replaces each failed row in place: + +```python +dataset = StreamingDataset( + table, + num_splits=8, + transform=transform, + on_transform_error="fill", + fill_value={"id": -1, "vector": [0.0, 0.0]}, +) +for row in dataset: + train(row) +print(dataset.rows_filled) +``` + +The placeholder must have the same structure as a successful transform output. +Pass `repair(exc, row_batch) -> row` instead of `fill_value` to construct a +replacement from the failed one-row Arrow batch. Keep the repair deterministic +when using elastic training or checkpoints. Fill mode preserves split lengths +and sample positions. + ::: lancedb.streaming.StreamingDataLoader ::: lancedb.permutation.permutation_builder diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index db71840c6..810911e28 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -13,7 +13,7 @@ Provides StreamingDataset, a PyTorch IterableDataset that guarantees: distributed topology changes between runs. Transform failures on bad rows (e.g. nulls or NaNs from incomplete data) can -be tolerated with ``on_transform_error="skip"``; see the parameter +be tolerated with ``on_transform_error="skip"`` or ``"fill"``; see the parameter documentation on StreamingDataset for how this interacts with the guarantees above. """ @@ -30,7 +30,7 @@ import warnings from collections import deque from concurrent.futures import ThreadPoolExecutor from copy import deepcopy -from multiprocessing import RawArray +from multiprocessing import RawArray, get_context from typing import Any, Callable, cast, Iterator, Literal, NamedTuple, Optional, Union import pyarrow as pa @@ -57,6 +57,7 @@ _EPOCH_PRIME = 100003 DEFAULT_READ_BATCH_SIZE = 64 DEFAULT_PREFETCH_BATCHES = 4 +_UNSET = object() class _WorkerSample(NamedTuple): @@ -350,6 +351,10 @@ class StreamingDataset(IterableDataset): - ``"skip"``: the failing rows are dropped and iteration continues. - ``"warn"``: like ``"skip"``, but a warning is logged for each failing batch. + - ``"fill"``: replace each failing row with ``fill_value`` or the + result of ``repair(exc, row_batch)``. Exactly one of these must be + provided. Filled rows retain their original positions, so split + lengths, elastic determinism, and checkpoint offsets are preserved. - a callable ``handler(exc) -> bool``: called with the exception; return ``True`` to skip the failing rows or ``False`` to re-raise. Useful to skip only expected error types (compatible with @@ -357,9 +362,9 @@ class StreamingDataset(IterableDataset): When a batch fails, the transform is re-invoked on each single-row slice of the batch so that only the rows that actually fail are - dropped. Transforms should therefore be deterministic and accept - batches of any size (including one row). Skipped rows are counted in - ``rows_skipped``. + skipped or filled. Transforms should therefore be deterministic and + accept batches of any size (including one row). Skipped and filled + rows are counted in ``rows_skipped`` and ``rows_filled`` respectively. Skipping weakens the elastic-determinism guarantee at the end of the epoch: splits that lose more rows than others run dry earlier, and @@ -372,8 +377,9 @@ class StreamingDataset(IterableDataset): but synchronous distributed training (e.g. ranks that call ``all_reduce`` every step) can hang or deadlock if one rank's iterator is exhausted while others are still stepping; callers doing - synchronous multi-rank training with ``on_transform_error != "raise"`` - are responsible for their own cross-rank stopping mechanism (e.g. + synchronous multi-rank training with ``on_transform_error="skip"`` or + ``"warn"`` (or a skipping callable) are responsible for their own + cross-rank stopping mechanism (e.g. broadcasting a stop signal on ``StopIteration``). The final few global steps can also differ across topologies (bounded by the skew in bad-row counts across splits). The sequence of samples yielded @@ -385,6 +391,18 @@ class StreamingDataset(IterableDataset): Prefer the ``filter`` parameter when bad rows can be expressed as a SQL predicate (e.g. ``"col IS NOT NULL"``) — filtering happens before splits are built, so every guarantee is fully preserved. + fill_value: + Constant output row used with ``on_transform_error="fill"``. A copy + is made for each failed row so mutable placeholders are independent. + For example, ``fill_value={"vector": [0.0, 0.0]}``. Pass exactly one + of ``fill_value`` or ``repair`` in fill mode. + repair: + Optional function ``repair(exc, row_batch) -> row`` used with + ``on_transform_error="fill"``. ``row_batch`` is a one-row + ``pyarrow.RecordBatch``. The returned value replaces that row's + transformed output. Replacements should have the same structure as + successful transform outputs. The repair must be deterministic for + reproducible checkpoints and elastic steps. transform_queue_depth: Number of transform-result batches to buffer per split in the post-transform queue before backpressure is applied to the transform @@ -438,6 +456,8 @@ class StreamingDataset(IterableDataset): pad_id: Optional[int] = None, blocks_per_epoch: Optional[Union[int, Literal["auto"]]] = None, on_transform_error: Union[str, Callable[[Exception], bool]] = "raise", + fill_value: Any = _UNSET, + repair: Optional[Callable[[Exception, pa.RecordBatch], Any]] = None, transform_queue_depth: Optional[int] = None, connection_factory: Optional[Callable[[str], Any]] = None, worker_info_override=None, @@ -513,13 +533,20 @@ class StreamingDataset(IterableDataset): ) elif blocks_per_epoch is not None: raise ValueError("blocks_per_epoch requires pack_sequences") - if on_transform_error not in ("raise", "skip", "warn") and not callable( + if on_transform_error not in ("raise", "skip", "warn", "fill") and not callable( on_transform_error ): raise ValueError( - "on_transform_error must be 'raise', 'skip', 'warn', or a " + "on_transform_error must be 'raise', 'skip', 'warn', 'fill', or a " f"callable, got {on_transform_error!r}" ) + if on_transform_error == "fill": + if (fill_value is _UNSET) == (repair is None): + raise ValueError( + "fill mode requires exactly one of fill_value or repair" + ) + elif fill_value is not _UNSET or repair is not None: + raise ValueError("fill_value and repair require on_transform_error='fill'") if transform_queue_depth is not None and transform_queue_depth <= 0: raise ValueError("transform_queue_depth must be greater than 0") @@ -542,6 +569,8 @@ class StreamingDataset(IterableDataset): self._pad_id = pad_id self._blocks_per_epoch = blocks_per_epoch self._on_transform_error = on_transform_error + self._fill_value = None if fill_value is _UNSET else fill_value + self._repair = repair self._transform_queue_depth = transform_queue_depth self._connection_factory = connection_factory self._worker_info_override = worker_info_override @@ -568,6 +597,9 @@ class StreamingDataset(IterableDataset): # bytes_loaded, fetch_time_us, transform_time_us, # rows_skipped] self._worker_stats: RawArray = RawArray(ctypes.c_int64, 8) + # A locked shared counter keeps fills from separate DataLoader workers + # from overwriting one another's totals. + self._filled_count = get_context("spawn").Value(ctypes.c_int64, 0) # A standard multi-process DataLoader cannot report which prefetched # batches were actually returned to its consumer. Workers set this @@ -951,12 +983,12 @@ class StreamingDataset(IterableDataset): on_error = self._on_transform_error - def _should_skip(exc: Exception) -> bool: + def _should_continue(exc: Exception) -> bool: if on_error == "raise": return False if callable(on_error): return bool(on_error(exc)) - return True # "skip" or "warn" + return True # "skip", "warn", or "fill" def _check_row_count(rows: list, num_rows: int) -> None: if len(rows) != num_rows: @@ -968,16 +1000,27 @@ class StreamingDataset(IterableDataset): ) def _transform_isolated(abs_start, batch, batch_exc): - """Re-run the transform on single-row slices, dropping failures.""" + """Re-run the transform on single-row slices to isolate failures.""" out = [] skipped = 0 + filled = 0 first_exc = None for j in range(batch.num_rows): + row_batch = batch.slice(j, 1) try: - rows = list(final_transform(batch.slice(j, 1))) + rows = list(final_transform(row_batch)) except Exception as exc: - if not _should_skip(exc): + if not _should_continue(exc): raise + if on_error == "fill": + row = ( + self._repair(exc, row_batch) + if self._repair is not None + else deepcopy(self._fill_value) + ) + out.append((abs_start + j, row)) + filled += 1 + continue skipped += 1 if first_exc is None: first_exc = exc @@ -985,6 +1028,9 @@ class StreamingDataset(IterableDataset): _check_row_count(rows, 1) out.append((abs_start + j, rows[0])) self._rows_skipped += skipped + if filled: + with self._filled_count.get_lock(): + self._filled_count.value += filled if skipped and on_error == "warn": logger.warning( "Skipped %d of %d rows whose transform failed (first error: %r)", @@ -999,7 +1045,7 @@ class StreamingDataset(IterableDataset): try: rows = list(final_transform(batch)) except Exception as exc: - if not _should_skip(exc): + if not _should_continue(exc): raise return _transform_isolated(abs_start, batch, exc) _check_row_count(rows, batch.num_rows) @@ -1370,6 +1416,15 @@ class StreamingDataset(IterableDataset): return self._rows_skipped return int(self._worker_stats[7]) + @property + def rows_filled(self) -> int: + """Number of failed transform rows replaced by fill mode. + + Accumulates across iterations of the same dataset instance and is + never reset automatically. + """ + return int(self._filled_count.value) + @property def consumed_rows(self) -> int: """Number of rows already yielded to the caller across all splits. diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 5eb3f095c..e1ded71cc 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -2337,6 +2337,13 @@ def _failing_transform(bad_ids: set): return transform +def _fill_workers_transform(batch: pa.RecordBatch) -> list[dict]: + ids = batch.column("id").to_pylist() + if 0 in ids or 60 in ids: + raise BadRowError("bad row") + return [{"id": row_id} for row_id in ids] + + def _sequential_split_members(table) -> list[list[int]]: """Return each split's ids in yield order for shuffle=False. @@ -2367,6 +2374,115 @@ def test_on_transform_error_invalid_value(lance_table): StreamingDataset(lance_table, num_splits=NUM_SPLITS, on_transform_error="bogus") +def test_on_transform_error_fill_requires_one_replacement(lance_table): + kwargs = dict(num_splits=NUM_SPLITS, on_transform_error="fill") + with pytest.raises(ValueError, match="exactly one"): + StreamingDataset(lance_table, **kwargs) + with pytest.raises(ValueError, match="exactly one"): + StreamingDataset( + lance_table, **kwargs, fill_value={"id": -1}, repair=lambda exc, row: row + ) + with pytest.raises(ValueError, match="require on_transform_error='fill'"): + StreamingDataset(lance_table, num_splits=NUM_SPLITS, fill_value=None) + StreamingDataset(lance_table, **kwargs, fill_value=None) + + +def test_on_transform_error_fill_preserves_rows_and_positions(lance_table): + """Uneven failures are replaced in place, including an entirely bad split.""" + members = _sequential_split_members(lance_table) + bad_ids = set(members[0]) | {members[5][7]} + clean = [ + row["id"] + for row in StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle=False) + ] + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="fill", + fill_value={"id": -1, "vector": [0.0, 0.0]}, + ) + assert ds.rows_filled == 0 + rows = list(ds) + assert [row["id"] for row in rows] == [ + -1 if row_id in bad_ids else row_id for row_id in clean + ] + assert len(rows) == NUM_ROWS + assert ds.rows_filled == len(bad_ids) + assert ds.rows_skipped == 0 + + placeholders = [row for row in rows if row["id"] == -1] + placeholders[0]["vector"][0] = 1.0 + assert all(row["vector"] == [0.0, 0.0] for row in placeholders[1:]) + + +def test_on_transform_error_fill_repair_receives_failed_row(lance_table): + bad_ids = {5, 17, 46} + repaired = [] + + def repair(exc: Exception, row_batch: pa.RecordBatch) -> dict: + assert isinstance(exc, BadRowError) + assert row_batch.num_rows == 1 + row_id = row_batch.column("id")[0].as_py() + repaired.append(row_id) + return {"id": row_id, "repaired": True} + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="fill", + repair=repair, + ) + rows = list(ds) + assert sorted(row["id"] for row in rows) == list(range(NUM_ROWS)) + assert set(repaired) == bad_ids + assert {row["id"] for row in rows if row.get("repaired")} == bad_ids + assert ds.rows_filled == len(bad_ids) + + +def test_on_transform_error_fill_repair_error_propagates(lance_table): + def broken_repair(exc: Exception, row_batch: pa.RecordBatch): + raise RuntimeError("repair failed") + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + transform=_failing_transform({5}), + on_transform_error="fill", + repair=broken_repair, + ) + with pytest.raises(RuntimeError, match="repair failed"): + list(ds) + + +def test_on_transform_error_fill_counts_across_workers(lance_table): + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_fill_workers_transform, + on_transform_error="fill", + fill_value={"id": -1}, + ) + loader = StreamingDataLoader( + ds, + batch_size=6, + num_workers=2, + multiprocessing_context="spawn", + ) + iterator = iter(loader) + try: + ids = [row_id for batch in iterator for row_id in batch["id"].tolist()] + finally: + iterator._shutdown_workers() + assert len(ids) == NUM_ROWS + assert ids.count(-1) == 2 + assert ds.rows_filled == 2 + + def test_on_transform_error_skip_drops_bad_rows(lance_table): """With one bad row per split, 'skip' yields every good row exactly once and counts the dropped rows in rows_skipped.""" @@ -2471,9 +2587,10 @@ def test_on_transform_error_callable_selective(lance_table): list(ds2) -def test_transform_wrong_row_count_raises(lance_table): +@pytest.mark.parametrize("on_error", ["skip", "fill"]) +def test_transform_wrong_row_count_raises(lance_table, on_error): """A transform that returns the wrong number of rows is an error even with - on_transform_error='skip' — silent shrinkage would corrupt accounting.""" + error recovery enabled — silent shrinkage would corrupt accounting.""" def drops_rows(batch: pa.RecordBatch) -> list: return batch.column("id").to_pylist()[:-1] @@ -2483,7 +2600,8 @@ def test_transform_wrong_row_count_raises(lance_table): num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED, transform=drops_rows, - on_transform_error="skip", + on_transform_error=on_error, + **({"fill_value": {"id": -1}} if on_error == "fill" else {}), ) with pytest.raises(ValueError, match="one output row per input row"): list(ds) @@ -2561,6 +2679,61 @@ def test_skip_elastic_det_across_world_sizes(lance_table): assert collect(ws) == reference, f"world_size={ws} diverged" +def test_fill_elastic_determinism_and_checkpoint(lance_table): + """Uneven fills keep every global step and resume across world sizes.""" + members = _sequential_split_members(lance_table) + bad_ids = {members[0][0], members[0][2], members[6][1]} + kwargs = dict( + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="fill", + fill_value={"id": -1}, + ) + reference = [row["id"] for row in StreamingDataset(lance_table, **kwargs)] + ref_steps = [ + sorted(reference[start : start + GLOBAL_BATCH_SIZE]) + for start in range(0, len(reference), GLOBAL_BATCH_SIZE) + ] + assert len(ref_steps) == STEPS_PER_EPOCH + + for world_size in (1, 2, 3, 4): + micro = GLOBAL_BATCH_SIZE // world_size + datasets = [ + StreamingDataset(lance_table, rank=rank, world_size=world_size, **kwargs) + for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + steps = [ + sorted(next(it)["id"] for it in iters for _ in range(micro)) + for _ in range(STEPS_PER_EPOCH) + ] + assert steps == ref_steps + assert all(next(it, None) is None for it in iters) + + world_size = 2 + micro = GLOBAL_BATCH_SIZE // world_size + datasets = [ + StreamingDataset(lance_table, rank=rank, world_size=world_size, **kwargs) + for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + for _ in range(3): + for it in iters: + for _ in range(micro): + next(it) + states = [ds.state_dict() for ds in datasets] + for it in iters: + it.close() + merged = StreamingDataset.merge_state_dicts(states) + assert merged["positions_consumed_per_split"] == [3] * NUM_SPLITS + assert merged["samples_consumed_per_split"] == [3] * NUM_SPLITS + + resumed = StreamingDataset(lance_table, **kwargs) + resumed.load_state_dict(merged) + assert [row["id"] for row in resumed] == reference[3 * NUM_SPLITS :] + + def test_resumability_with_skips_same_topology(lance_table): """Checkpointing mid-epoch with skipped rows resumes exactly: no sample repeated, no sample lost, skipped rows stay skipped."""