address cr feedback

Author: nabinchha

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

@@ -74,7 +74,6 @@ def create(
                 for se_col in sub.side_effect_columns:
                     graph.set_side_effect(se_col, name)
 
-                graph.set_required_columns(name, list(sub.required_columns))
                 graph.set_propagate_skip(name, sub.propagate_skip)
                 if sub.skip is not None:
                     graph.set_skip_config(name, sub.skip)
@@ -90,6 +89,7 @@ def create(
 
             for sub in sub_configs:
                 name = sub.name
+                resolved_required: list[str] = []
                 for req in sub.required_columns:
                     resolved = graph.resolve_side_effect(req)
                     if resolved not in known_columns:
@@ -98,7 +98,10 @@ def create(
                         )
                     if resolved == name:
                         continue
+                    if resolved not in resolved_required:
+                        resolved_required.append(resolved)
                     graph.add_edge(upstream=resolved, downstream=name)
+                graph.set_required_columns(name, resolved_required)
 
                 if sub.skip is not None:
                     for skip_col in sub.skip.columns:
@@ -135,7 +138,7 @@ def set_side_effect(self, side_effect_col: str, producer: str) -> None:
         self._producer_to_side_effect_map.setdefault(producer, []).append(side_effect_col)
 
     def set_required_columns(self, column: str, required: list[str]) -> None:
-        """Store the config-level ``required_columns`` for *column*."""
+        """Store producer-resolved ``required_columns`` for skip propagation."""
         self._required_columns[column] = required
 
     def set_propagate_skip(self, column: str, propagate: bool) -> None:
@@ -165,7 +168,7 @@ def get_downstream_columns(self, column: str) -> set[str]:
         return set(self._downstream.get(column, set()))
 
     def get_required_columns(self, column: str) -> list[str]:
-        """Config-level ``required_columns`` for *column* (data dependencies only)."""
+        """Producer-resolved ``required_columns`` for *column* (data dependencies only)."""
         return list(self._required_columns.get(column, []))
 
     def get_skip_config(self, column: str) -> SkipConfig | None:

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

@@ -10,6 +10,7 @@
 import pytest
 
 import data_designer.lazy_heavy_imports as lazy
+from data_designer.config.base import SkipConfig
 from data_designer.config.column_configs import (
     CustomColumnConfig,
     ExpressionColumnConfig,
@@ -1563,8 +1564,6 @@ async def test_scheduler_skip_cell_by_cell_with_propagation() -> None:
     Pipeline: seed(sampler) -> review(cell, skip.when seed<2) -> complaint(cell, propagate_skip)
     Rows with seed < 2 should be skipped for review and propagated to complaint.
     """
-    from data_designer.config.base import SkipConfig
-
     provider = _mock_provider()
     num_records = 4
 
@@ -1640,15 +1639,97 @@ def generate_from_scratch(self, num_records: int) -> lazy.pd.DataFrame:
             assert row.get("complaint") is not None, f"row {ri}: complaint should be generated (seed={seed_val})"
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_scheduler_skip_propagates_through_side_effect_dependency() -> None:
+    """A downstream dependency on a skipped side-effect should auto-skip.
+
+    Pipeline: seed(sampler) -> review(cell, skip.when seed<2, produces
+    review__trace) -> complaint(cell, depends on review__trace,
+    propagate_skip=True).
+    """
+    provider = _mock_provider()
+    num_records = 4
+
+    configs = [
+        SamplerColumnConfig(name="seed", sampler_type=SamplerType.CATEGORY, params={"values": ["A"]}),
+        LLMTextColumnConfig(
+            name="review",
+            prompt="{{ seed }}",
+            model_alias=MODEL_ALIAS,
+            with_trace="last_message",
+            skip=SkipConfig(when="{{ seed < 2 }}"),
+        ),
+        LLMTextColumnConfig(
+            name="complaint",
+            prompt="{{ review__trace }}",
+            model_alias=MODEL_ALIAS,
+            propagate_skip=True,
+        ),
+    ]
+    strategies = {
+        "seed": GenerationStrategy.FULL_COLUMN,
+        "review": GenerationStrategy.CELL_BY_CELL,
+        "complaint": GenerationStrategy.CELL_BY_CELL,
+    }
+
+    class IntSeedGenerator(FromScratchColumnGenerator[ExpressionColumnConfig]):
+        @staticmethod
+        def get_generation_strategy() -> GenerationStrategy:
+            return GenerationStrategy.FULL_COLUMN
+
+        def generate(self, data: lazy.pd.DataFrame) -> lazy.pd.DataFrame:
+            return data
+
+        def generate_from_scratch(self, num_records: int) -> lazy.pd.DataFrame:
+            return lazy.pd.DataFrame({"seed": list(range(num_records))})
+
+    generators: dict[str, ColumnGenerator] = {
+        "seed": IntSeedGenerator(config=_expr_config("seed"), resource_provider=provider),
+        "review": MockCellGenerator(config=_expr_config("review"), resource_provider=provider),
+        "complaint": MockCellGenerator(config=_expr_config("complaint"), resource_provider=provider),
+    }
+
+    storage = MagicMock()
+    storage.dataset_name = "test"
+    storage.get_file_paths.return_value = {}
+    buffer_mgr = RowGroupBufferManager(storage)
+
+    graph = ExecutionGraph.create(configs, strategies)
+    row_groups = [(0, num_records)]
+    tracker = CompletionTracker.with_graph(graph, row_groups)
+
+    scheduler = AsyncTaskScheduler(
+        generators=generators,
+        graph=graph,
+        tracker=tracker,
+        row_groups=row_groups,
+        buffer_manager=buffer_mgr,
+        trace=True,
+        num_records=num_records,
+        buffer_size=num_records,
+    )
+    await asyncio.wait_for(scheduler.run(), timeout=10.0)
+
+    assert tracker.is_row_group_complete(0, num_records, ["seed", "review", "complaint"])
+
+    for ri in range(num_records):
+        row = buffer_mgr.get_row(0, ri)
+        seed_val = row["seed"]
+        if seed_val < 2:
+            assert row.get("review") is None, f"row {ri}: review should be skipped (seed={seed_val})"
+            assert row.get("review__trace") is None, f"row {ri}: review__trace should be cleared on skip"
+            assert row.get("complaint") is None, f"row {ri}: complaint should propagate skip (seed={seed_val})"
+        else:
+            assert row.get("complaint") is not None, f"row {ri}: complaint should be generated (seed={seed_val})"
+
+
 @pytest.mark.asyncio(loop_scope="session")
 async def test_scheduler_skip_full_column_batch() -> None:
     """Full-column (batch) generator skips rows via expression gate.
 
     Pipeline: seed(sampler) -> review(full_column, skip.when seed<2)
     Only active (non-skipped) rows should be passed to the generator.
     """
-    from data_designer.config.base import SkipConfig
-
     provider = _mock_provider()
     num_records = 4
 

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

@@ -11,6 +11,7 @@
 
 import data_designer.engine.dataset_builders.dataset_builder as builder_mod
 import data_designer.lazy_heavy_imports as lazy
+from data_designer.config.base import SkipConfig
 from data_designer.config.column_configs import CustomColumnConfig, LLMTextColumnConfig, SamplerColumnConfig
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.custom_column import custom_column_generator
@@ -962,6 +963,21 @@ def fn(df: pd.DataFrame) -> pd.DataFrame:
     return fn
 
 
+def _make_label_generator_with_side_effect(label: str, side_effect_label: str, *required: str):
+    """FULL_COLUMN generator that adds a column plus one side-effect column."""
+
+    @custom_column_generator(required_columns=list(required), side_effect_columns=[side_effect_label])
+    def fn(df: pd.DataFrame) -> pd.DataFrame:
+        return df.assign(
+            **{
+                label: f"generated_{label}",
+                side_effect_label: f"generated_{side_effect_label}",
+            }
+        )
+
+    return fn
+
+
 def test_skip_metadata_preserved_across_non_skip_aware_full_column(
     stub_resource_provider, stub_model_configs, seed_data_setup
 ):
@@ -972,8 +988,6 @@ def test_skip_metadata_preserved_across_non_skip_aware_full_column(
     Before the fix, summary's replace_buffer erased __internal_skipped_columns,
     causing complaint to generate for rows that should have been skipped.
     """
-    from data_designer.config.base import SkipConfig
-
     config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs)
     config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"])))
 
@@ -1031,8 +1045,6 @@ def test_skip_metadata_preserved_when_no_rows_skipped_for_current_column(
     own expression (it has none). The has_skipped=False fallthrough must still
     preserve review's skip metadata so propagation works.
     """
-    from data_designer.config.base import SkipConfig
-
     config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs)
     config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"])))
 
@@ -1069,6 +1081,53 @@ def test_skip_metadata_preserved_when_no_rows_skipped_for_current_column(
             assert row["analysis"] == "generated_analysis", f"seed_id={row['seed_id']}: analysis should be generated"
 
 
+def test_skip_propagation_resolves_side_effect_dependencies_in_sync_builder(
+    stub_resource_provider, stub_model_configs, seed_data_setup
+):
+    """A downstream dependency on a skipped side-effect should auto-skip.
+
+    Scenario: review(skip.when, produces review_side_effect) ->
+    analysis(required_columns=[review_side_effect], propagate_skip=True).
+    """
+    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_with_side_effect("review", "review_side_effect", "seed_id"),
+            generation_strategy=GenerationStrategy.FULL_COLUMN,
+            skip=SkipConfig(when="{{ seed_id < 3 }}"),
+        )
+    )
+    config_builder.add_column(
+        CustomColumnConfig(
+            name="analysis",
+            generator_function=_make_label_generator("analysis", "review_side_effect"),
+            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)
+
+    skipped_ids = {1, 2}
+    for _, row in result.iterrows():
+        if row["seed_id"] in skipped_ids:
+            assert row["review_side_effect"] is None or lazy.pd.isna(row["review_side_effect"]), (
+                f"seed_id={row['seed_id']}: review_side_effect should be cleared when review is skipped"
+            )
+            assert row["analysis"] is None or lazy.pd.isna(row["analysis"]), (
+                f"seed_id={row['seed_id']}: analysis should propagate skip from review"
+            )
+        else:
+            assert row["analysis"] == "generated_analysis", f"seed_id={row['seed_id']}: analysis should be generated"
+
+
 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).
@@ -1077,8 +1136,6 @@ def test_allow_resize_column_not_blocked_by_upstream_skip(stub_resource_provider
     with propagate_skip=True and required_columns pointing to a skippable
     upstream, causing a DatasetGenerationError on the row-count check.
     """
-    from data_designer.config.base import SkipConfig
-
     config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs)
     config_builder.with_seed_dataset(LocalFileSeedSource(path=str(seed_data_setup["seed_path"])))
 

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

@@ -112,6 +112,7 @@ def test_side_effect_column_resolution() -> None:
 
     assert graph.get_upstream_columns("trace_len") == {"summary"}
     assert graph.get_downstream_columns("summary") == {"trace_len"}
+    assert graph.get_required_columns("trace_len") == ["summary"]
 
 
 def test_reasoning_content_side_effect() -> None: