Fix for use case when generator skips input rows in checkpoints (#1609)

Fix for use case when generator skips input rows in checkpoints

Author: ilongin

src/datachain/lib/udf.py

@@ -564,16 +564,23 @@ def _prepare_rows(udf_inputs) -> "abc.Generator[Sequence[Any], None, None]":
 
         def _process_row(row):
             row_id, *row = row
+            has_output = False
             with safe_closing(self.process(*row)) as result_objs:
                 for result_obj, is_last in with_last_flag(result_objs):
+                    has_output = True
                     udf_output = self._flatten_row(result_obj)
                     udf_output = dict(zip(self.signal_names, udf_output, strict=False))
-                    # Include sys__input_id to track which input generated this
-                    # output.
-                    udf_output["sys__input_id"] = row_id  # input id
-                    # Mark as partial=True unless it's the last output
+                    udf_output["sys__input_id"] = row_id
                     udf_output["sys__partial"] = not is_last
+                    udf_output["sys__empty"] = None
                     yield udf_output
+            if not has_output:
+                # Marker: records that this input was processed but yielded nothing.
+                yield {
+                    "sys__input_id": row_id,
+                    "sys__partial": False,
+                    "sys__empty": True,
+                }
 
         prepared_inputs = _prepare_rows(udf_inputs)
         prepared_inputs = _prefetch_inputs(
@@ -584,12 +591,6 @@ def _process_row(row):
         )
 
         with closing(prepared_inputs):
-            # TODO: Fix limitation where inputs yielding nothing are not tracked in
-            # processed table. Currently, if process() yields nothing for an input,
-            # that input's sys__id is never added to the processed table, causing it
-            # to be re-processed on checkpoint recovery. Solution: yield a marker
-            # row with sys__input_id when process() yields nothing, then filter
-            # these marker rows before inserting to output table.
             for row in prepared_inputs:
                 yield _process_row(row)
                 processed_cb.relative_update(1)

src/datachain/query/dataset.py

@@ -538,6 +538,7 @@ def _checkpoint_tracking_columns(self) -> list["sqlalchemy.Column"]:
         return [
             sa.Column("sys__input_id", sa.Integer, nullable=True),
             sa.Column("sys__partial", sa.Boolean, nullable=True),
+            sa.Column("sys__empty", sa.Boolean, nullable=True),
         ]
 
     def get_input_query(self, input_table_name: str, original_query: Select) -> Select:
@@ -1591,13 +1592,21 @@ def create_output_table(self, name: str) -> "Table":
     def create_result_query(
         self, udf_table, query: Select
     ) -> tuple[QueryGeneratorFunc, list["sqlalchemy.Column"]]:
-        udf_table_query = udf_table.select().subquery()
-        # Exclude sys__input_id and sys__partial - they're only needed for tracking
-        # during UDF execution and checkpoint recovery
+        # Filter out empty-input marker rows (inputs that yielded nothing)
+        udf_table_query = (
+            udf_table.select()
+            .where(
+                sa.or_(
+                    udf_table.c.sys__empty.is_(None),
+                    udf_table.c.sys__empty == sa.false(),
+                )
+            )
+            .subquery()
+        )
+        # Exclude checkpoint tracking columns from the result
+        excluded = {c.name for c in self._checkpoint_tracking_columns()}
         udf_table_cols: list[sqlalchemy.Label[Any]] = [
-            label(c.name, c)
-            for c in udf_table_query.columns
-            if c.name not in ("sys__input_id", "sys__partial")
+            label(c.name, c) for c in udf_table_query.columns if c.name not in excluded
         ]
 
         def q(*columns):
@@ -1606,11 +1615,7 @@ def q(*columns):
             cols = [c for c in udf_table_cols if c.name in names]
             return sqlalchemy.select(*cols).select_from(udf_table_query)
 
-        return q, [
-            c
-            for c in udf_table_query.columns
-            if c.name not in ("sys__input_id", "sys__partial")
-        ]
+        return q, [c for c in udf_table_query.columns if c.name not in excluded]
 
 
 @frozen

tests/func/checkpoints/test_checkpoint_recovery.py

@@ -254,22 +254,29 @@ def gen_multiple(num) -> Iterator[int]:
     assert len(all_results) == len(set(all_results)), "Should have no duplicate results"
 
 
-@pytest.mark.xfail(
-    reason="Known limitation: inputs that yield nothing are not tracked "
-    "in processed table"
-)
 def test_generator_yielding_nothing(test_session, monkeypatch, nums_dataset):
-    """Test that generator correctly handles inputs that yield zero outputs."""
+    """Test that generator correctly handles inputs that yield zero outputs.
+
+    Uses fail_after_count=4 so that regardless of DB row ordering, at least
+    one odd input (which yields nothing) is processed before the failure.
+    With 6 inputs (3 odd, 3 even), any 4 must include at least 1 odd.
+    """
+    fail_after_count = 4
     processed = []
+    skipped = []
+    call_count = [0]
 
     def selective_generator(num) -> Iterator[int]:
         processed.append(num)
-        if num == 3:
+        call_count[0] += 1
+        if call_count[0] > fail_after_count:
             raise Exception("Simulated failure")
         if num % 2 == 0:  # Only even numbers yield outputs
             yield num * 10
+        else:
+            skipped.append(num)
 
-    # First run - fails on num=3
+    # First run - fails after processing 4 inputs
     reset_session_job_state()
     chain = dc.read_dataset("nums", session=test_session).gen(
         value=selective_generator, output=int
@@ -278,13 +285,24 @@ def selective_generator(num) -> Iterator[int]:
     with pytest.raises(Exception, match="Simulated failure"):
         chain.save("results")
 
+    assert len(processed) == fail_after_count + 1  # 4 succeeded + 1 failed
+    skipped_first_run = list(skipped)
+    assert skipped_first_run, "Expected at least one empty-yield input"
+
     # Second run - should continue from checkpoint
     reset_session_job_state()
     processed.clear()
+    skipped.clear()
+    call_count[0] = 0
     chain.save("results")
 
-    # Only inputs 3,4,5,6 should be processed (1,2 were already done)
-    assert processed == [3, 4, 5, 6]
+    # Empty-yield inputs from first run must not be re-processed
+    assert not set(skipped_first_run) & set(processed), (
+        f"Empty-yield inputs {set(skipped_first_run) & set(processed)} "
+        f"were re-processed despite being checkpointed"
+    )
+    assert len(processed) == 2
+    # Final result: all even numbers yield output
     result = sorted(dc.read_dataset("results", session=test_session).to_list("value"))
     assert result == [(20,), (40,), (60,)]