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
6b51b33e56: four workflows and fifteen checks
all succeeded.

Closes #3916

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.

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."""