Fix issue with full column generating messing up order of skipped rows

Author: nabinchha

packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py

@@ -42,6 +42,7 @@
 from data_designer.engine.dataset_builders.utils.skip_tracker import (
     SKIPPED_COLUMNS_RECORD_KEY,
     apply_skip_to_record,
+    prepare_records_for_skip_metadata_round_trip,
     restore_skip_metadata,
     strip_skip_metadata_from_records,
 )
@@ -579,11 +580,19 @@ def _run_full_column_generator_without_skip(self, generator: ColumnGenerator) ->
         original_count = self.batch_manager.num_records_in_buffer
         allow_resize = generator.config.allow_resize if not isinstance(generator.config, MultiColumnConfig) else False
         old_records = [record for _, record in self.batch_manager.iter_current_batch()]
+        input_records, restore_context = prepare_records_for_skip_metadata_round_trip(old_records)
 
-        df = generator.generate(self.batch_manager.get_current_batch(as_dataframe=True))
+        df = generator.generate(lazy.pd.DataFrame(input_records))
         self._log_resize_if_changed(self._column_display_name(generator.config), original_count, len(df), allow_resize)
         new_records = df.to_dict(orient="records")
-        restore_skip_metadata(old_records, new_records)
+        if restore_context is not None:
+            try:
+                restore_skip_metadata(new_records, context=restore_context, allow_resize=allow_resize)
+            except ValueError as exc:
+                raise DatasetGenerationError(
+                    f"Unable to restore skip provenance after FULL_COLUMN generation for "
+                    f"{self._column_display_name(generator.config)}: {exc}"
+                ) from exc
         self.batch_manager.replace_buffer(new_records, allow_resize=allow_resize)
 
     def _run_full_column_generator_with_skip(self, generator: ColumnGenerator, column_name: str) -> None:

packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/skip_tracker.py

@@ -10,9 +10,20 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from dataclasses import dataclass
 from typing import Final
 
 SKIPPED_COLUMNS_RECORD_KEY: Final[str] = "__internal_skipped_columns"
+SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX: Final[str] = "__internal_skip_restore_id"
+
+
+@dataclass(frozen=True, slots=True)
+class SkipMetadataRestoreContext:
+    """Metadata needed to restore skip provenance after a DataFrame round-trip."""
+
+    restore_id_column: str
+    source_ids: set[str]
+    skipped_columns_by_source_id: dict[str, set[str]]
 
 
 def apply_skip_to_record(
@@ -46,16 +57,82 @@ def strip_skip_metadata_from_records(records: Sequence[dict]) -> list[dict]:
     return [strip_skip_metadata_for_dataframe_row(r) for r in records]
 
 
-def restore_skip_metadata(old_records: Sequence[dict], new_records: Sequence[dict]) -> None:
-    """Copy ``SKIPPED_COLUMNS_RECORD_KEY`` from *old_records* into *new_records* in-place.
+def prepare_records_for_skip_metadata_round_trip(
+    records: Sequence[dict],
+) -> tuple[list[dict], SkipMetadataRestoreContext | None]:
+    """Prepare records for a DataFrame round-trip while preserving skip metadata.
 
-    ``pd.DataFrame`` construction drops non-column keys, so skip metadata is
-    lost when records round-trip through a DataFrame.  Call this after
-    ``df.to_dict(orient="records")`` to re-attach the metadata before passing
-    the records to ``replace_buffer``.  When lengths differ (e.g.
-    ``allow_resize``), only positionally matched rows are restored.
+    Returns stripped records ready for ``pd.DataFrame(...)``. If any record has
+    skip metadata, injects a hidden restore-ID column and returns a context that
+    can later be passed to :func:`restore_skip_metadata`.
     """
-    for i in range(min(len(old_records), len(new_records))):
-        meta = old_records[i].get(SKIPPED_COLUMNS_RECORD_KEY)
+    if not any(SKIPPED_COLUMNS_RECORD_KEY in record for record in records):
+        return strip_skip_metadata_from_records(records), None
+
+    restore_id_column = _choose_restore_id_column(records)
+    prepared_records: list[dict] = []
+    source_ids: set[str] = set()
+    skipped_columns_by_source_id: dict[str, set[str]] = {}
+
+    for index, record in enumerate(records):
+        source_id = str(index)
+        source_ids.add(source_id)
+        prepared_record = strip_skip_metadata_for_dataframe_row(record)
+        prepared_record[restore_id_column] = source_id
+        prepared_records.append(prepared_record)
+
+        meta = record.get(SKIPPED_COLUMNS_RECORD_KEY)
         if meta is not None:
-            new_records[i][SKIPPED_COLUMNS_RECORD_KEY] = meta
+            skipped_columns_by_source_id[source_id] = set(meta)
+
+    return prepared_records, SkipMetadataRestoreContext(
+        restore_id_column=restore_id_column,
+        source_ids=source_ids,
+        skipped_columns_by_source_id=skipped_columns_by_source_id,
+    )
+
+
+def restore_skip_metadata(
+    records: Sequence[dict],
+    *,
+    context: SkipMetadataRestoreContext,
+    allow_resize: bool,
+) -> None:
+    """Restore skip provenance using hidden restore IDs instead of row position."""
+    restored_source_ids: list[str] = []
+    for record in records:
+        if context.restore_id_column not in record:
+            raise ValueError(
+                f"Records returned from the DataFrame round-trip must preserve "
+                f"the internal column {context.restore_id_column!r} so skip "
+                "provenance can be restored."
+            )
+
+        source_id = str(record.pop(context.restore_id_column))
+        if source_id not in context.source_ids:
+            raise ValueError(
+                f"Record returned unknown restore ID {source_id!r}. Skip provenance "
+                "can only be restored for rows derived from the original input."
+            )
+
+        restored_source_ids.append(source_id)
+        meta = context.skipped_columns_by_source_id.get(source_id)
+        if meta is not None:
+            record[SKIPPED_COLUMNS_RECORD_KEY] = set(meta)
+
+    if not allow_resize:
+        if len(restored_source_ids) != len(context.source_ids) or set(restored_source_ids) != context.source_ids:
+            raise ValueError(
+                "Full-column generation changed the row identity mapping while "
+                "allow_resize=False. Returned rows must preserve a 1:1 mapping "
+                "to the original input so skip provenance can be restored."
+            )
+
+
+def _choose_restore_id_column(records: Sequence[dict]) -> str:
+    candidate = SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX
+    suffix = 0
+    while any(candidate in record for record in records):
+        suffix += 1
+        candidate = f"{SKIP_METADATA_RESTORE_ID_COLUMN_PREFIX}_{suffix}"
+    return candidate

packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py

@@ -805,6 +805,12 @@ def _resize_full_keep_first(df: pd.DataFrame) -> pd.DataFrame:
     return df.drop_duplicates(subset="seed_id").assign(filtered=True)
 
 
+@custom_column_generator(required_columns=["seed_id"])
+def _resize_full_drop_seed_one(df: pd.DataFrame) -> pd.DataFrame:
+    """FULL_COLUMN: drop the row with seed_id == 1."""
+    return df[df["seed_id"] != 1].reset_index(drop=True).assign(filtered=True)
+
+
 @custom_column_generator(required_columns=["seed_id"])
 def _resize_cell_expand(row: dict) -> list[dict]:
     """CELL_BY_CELL: one row -> two rows (doubled)."""
@@ -1128,6 +1134,49 @@ def test_skip_propagation_resolves_side_effect_dependencies_in_sync_builder(
             assert row["analysis"] == "generated_analysis", f"seed_id={row['seed_id']}: analysis should be generated"
 
 
+def test_skip_metadata_restore_preserves_row_identity_across_allow_resize_full_column(
+    stub_resource_provider, stub_model_configs, seed_data_setup
+):
+    """Filtering out a skipped row must not transfer its skip provenance to surviving rows."""
+    config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs)
+    config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"])))
+
+    config_builder.add_column(
+        CustomColumnConfig(
+            name="review",
+            generator_function=_make_label_generator("review", "seed_id"),
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            skip=SkipConfig(when="{{ seed_id == 1 }}"),
+        )
+    )
+    config_builder.add_column(
+        CustomColumnConfig(
+            name="filtered",
+            generator_function=_resize_full_drop_seed_one,
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            allow_resize=True,
+            propagate_skip=False,
+        )
+    )
+    config_builder.add_column(
+        CustomColumnConfig(
+            name="analysis",
+            generator_function=_make_label_generator("analysis", "review"),
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            propagate_skip=True,
+        )
+    )
+
+    builder = DatasetBuilder(
+        data_designer_config=config_builder.build(),
+        resource_provider=stub_resource_provider,
+    )
+    result = builder.build_preview(num_records=5)
+
+    assert result["seed_id"].tolist() == [2, 3, 4, 5]
+    assert result["analysis"].tolist() == ["generated_analysis"] * 4
+
+
 def test_allow_resize_column_not_blocked_by_upstream_skip(stub_resource_provider, stub_model_configs, seed_data_setup):
     """An allow_resize=True column depending on a skippable upstream must not
     enter the skip-aware branch (which enforces 1:1 row counts).

packages/data-designer-engine/tests/engine/dataset_builders/utils/test_skip_tracker.py

@@ -9,6 +9,7 @@
 from data_designer.engine.dataset_builders.utils.skip_tracker import (
     SKIPPED_COLUMNS_RECORD_KEY,
     apply_skip_to_record,
+    prepare_records_for_skip_metadata_round_trip,
     restore_skip_metadata,
     strip_skip_metadata_for_dataframe_row,
     strip_skip_metadata_from_records,
@@ -131,32 +132,75 @@ def test_strip_skip_metadata_from_records(rows: list[dict], expected: list[dict]
     assert strip_skip_metadata_from_records(rows) == expected
 
 
-def test_restore_skip_metadata_copies_metadata() -> None:
-    old = [
+def test_prepare_records_for_skip_metadata_round_trip_without_metadata() -> None:
+    rows = [{"a": 1}, {"a": 2}]
+    prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(rows)
+    assert restore_context is None
+    assert prepared_rows == rows
+    assert prepared_rows is not rows
+
+
+def test_prepare_records_for_skip_metadata_round_trip_injects_restore_ids() -> None:
+    rows = [
         {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}},
         {"a": 2},
         {"a": 3, SKIPPED_COLUMNS_RECORD_KEY: {"col_y", "col_z"}},
     ]
-    new = [{"a": 10}, {"a": 20}, {"a": 30}]
-    restore_skip_metadata(old, new)
-    assert new[0][SKIPPED_COLUMNS_RECORD_KEY] == {"col_x"}
-    assert SKIPPED_COLUMNS_RECORD_KEY not in new[1]
-    assert new[2][SKIPPED_COLUMNS_RECORD_KEY] == {"col_y", "col_z"}
+    prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(rows)
+    assert restore_context is not None
+    assert SKIPPED_COLUMNS_RECORD_KEY not in prepared_rows[0]
+    assert restore_context.restore_id_column in prepared_rows[0]
+    assert restore_context.skipped_columns_by_source_id == {
+        "0": {"col_x"},
+        "2": {"col_y", "col_z"},
+    }
 
 
-def test_restore_skip_metadata_handles_length_mismatch() -> None:
+def test_restore_skip_metadata_uses_restore_ids_after_reorder() -> None:
     old = [
         {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}},
-        {"a": 2, SKIPPED_COLUMNS_RECORD_KEY: {"col_y"}},
+        {"a": 2},
+        {"a": 3, SKIPPED_COLUMNS_RECORD_KEY: {"col_z"}},
     ]
-    new = [{"a": 10}]
-    restore_skip_metadata(old, new)
-    assert new[0][SKIPPED_COLUMNS_RECORD_KEY] == {"col_x"}
+    prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old)
+    assert restore_context is not None
+    restore_id_column = restore_context.restore_id_column
+
+    new = [
+        {"a": 30, restore_id_column: prepared_rows[2][restore_id_column]},
+        {"a": 10, restore_id_column: prepared_rows[0][restore_id_column]},
+        {"a": 20, restore_id_column: prepared_rows[1][restore_id_column]},
+    ]
+    restore_skip_metadata(new, context=restore_context, allow_resize=False)
+
+    assert new[0][SKIPPED_COLUMNS_RECORD_KEY] == {"col_z"}
+    assert new[1][SKIPPED_COLUMNS_RECORD_KEY] == {"col_x"}
+    assert SKIPPED_COLUMNS_RECORD_KEY not in new[2]
 
 
-def test_restore_skip_metadata_no_metadata() -> None:
+def test_restore_skip_metadata_allow_resize_handles_filtered_rows() -> None:
     old = [{"a": 1}, {"a": 2}]
-    new = [{"a": 10}, {"a": 20}]
-    restore_skip_metadata(old, new)
+    prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old)
+    assert restore_context is None
+
+    old = [
+        {"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}},
+        {"a": 2},
+    ]
+    prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old)
+    assert restore_context is not None
+    restore_id_column = restore_context.restore_id_column
+
+    new = [{"a": 20, restore_id_column: prepared_rows[1][restore_id_column]}]
+    restore_skip_metadata(new, context=restore_context, allow_resize=True)
+
     assert SKIPPED_COLUMNS_RECORD_KEY not in new[0]
-    assert SKIPPED_COLUMNS_RECORD_KEY not in new[1]
+
+
+def test_restore_skip_metadata_rejects_missing_restore_id_column() -> None:
+    old = [{"a": 1, SKIPPED_COLUMNS_RECORD_KEY: {"col_x"}}]
+    _prepared_rows, restore_context = prepare_records_for_skip_metadata_round_trip(old)
+    assert restore_context is not None
+
+    with pytest.raises(ValueError, match="must preserve the internal column"):
+        restore_skip_metadata([{"a": 10}], context=restore_context, allow_resize=False)