fix(compiler): support columns_added and columns_removed in processor config (#394)

packages/data-designer-config/src/data_designer/config/base.py

@@ -196,3 +196,11 @@ class ProcessorConfig(ConfigBase, ABC):
         description="The name of the processor, used to identify the processor in the results and to write the artifacts to disk.",
     )
     processor_type: str
+    columns_added: list[str] = Field(
+        default_factory=list,
+        description="List of column names added to the dataset by this processor.",
+    )
+    columns_removed: list[str] = Field(
+        default_factory=list,
+        description="List of column names removed from the dataset by this processor.",
+    )

packages/data-designer-config/tests/config/test_processors.py

@@ -106,3 +106,30 @@ class UnknownProcessorType(str, Enum):
         UnknownProcessorType.UNKNOWN, name="unknown_processor", column_names=["col1"]
     )
     assert result is None
+
+
+def test_processor_config_columns_added_and_removed_defaults():
+    config = DropColumnsProcessorConfig(name="drop_proc", column_names=["col1"])
+    assert config.columns_added == []
+    assert config.columns_removed == []
+
+
+def test_processor_config_columns_added_and_removed_custom():
+    class CustomProcessorConfig(ProcessorConfig):
+        processor_type: str = "custom"
+
+    config = CustomProcessorConfig(
+        name="custom_proc",
+        columns_added=["col_new"],
+        columns_removed=["col_old"],
+    )
+    assert config.columns_added == ["col_new"]
+    assert config.columns_removed == ["col_old"]
+
+    data = config.model_dump()
+    assert data["columns_added"] == ["col_new"]
+    assert data["columns_removed"] == ["col_old"]
+
+    restored = CustomProcessorConfig.model_validate(data)
+    assert restored.columns_added == ["col_new"]
+    assert restored.columns_removed == ["col_old"]

packages/data-designer-engine/src/data_designer/engine/compiler.py

@@ -18,11 +18,40 @@
 
 def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig:
     _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
+    _apply_processor_column_modifications(config)
     _add_internal_row_id_column_if_needed(config)
     _validate(config)
     return config
 
 
+def _apply_processor_column_modifications(config: DataDesignerConfig) -> None:
+    """Adjusts columns according to columns_added and columns_removed declared by processors."""
+    for processor in config.processors or []:
+        if processor.columns_removed:
+            current_columns = {col.name for col in config.columns}
+            for col_name in processor.columns_removed:
+                if col_name not in current_columns:
+                    raise InvalidConfigError(
+                        f"🛑 Processor '{processor.name}' cannot remove column '{col_name}' because it does not exist."
+                    )
+            removed_set = set(processor.columns_removed)
+            config.columns = [col for col in config.columns if col.name not in removed_set]
+
+        if processor.columns_added:
+            if config.seed_config is None:
+                raise InvalidConfigError(
+                    f"🛑 Processor '{processor.name}' specifies 'columns_added', but no seed dataset is configured."
+                )
+            existing_columns = {col.name for col in config.columns}
+            for col_name in processor.columns_added:
+                if col_name in existing_columns:
+                    raise InvalidConfigError(
+                        f"🛑 Processor '{processor.name}' adds column '{col_name}' which collides with an existing column."
+                    )
+                config.columns.append(SeedDatasetColumnConfig(name=col_name))
+                existing_columns.add(col_name)
+
+
 def _resolve_and_add_seed_columns(config: DataDesignerConfig, seed_reader: SeedReader | None) -> None:
     """Fetches the seed dataset column names, ensures there are no conflicts
     with other columns, and adds seed column configs to the DataDesignerConfig.

packages/data-designer-engine/tests/engine/test_compiler.py

@@ -8,6 +8,7 @@
 from data_designer.config.column_configs import ExpressionColumnConfig, SamplerColumnConfig
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.errors import InvalidConfigError
+from data_designer.config.processors import DropColumnsProcessorConfig
 from data_designer.config.sampler_params import CategorySamplerParams, SamplerType, UUIDSamplerParams
 from data_designer.config.seed_source import FileContentsSeedSource, HuggingFaceSeedSource
 from data_designer.engine.compiler import compile_data_designer_config
@@ -174,3 +175,137 @@ def test_does_not_add_id_column_when_seed_dataset_exists(resource_provider: Reso
     assert len(config.columns) == 3
     assert config.columns[0].name == "derived_value"
     assert not any(col.name == "_internal_row_id" for col in config.columns)
+
+
+def test_compile_applies_processor_columns_added(resource_provider: ResourceProvider):
+    """Test that columns declared in columns_added can be referenced by downstream expressions/templates."""
+    builder = DataDesignerConfigBuilder()
+    builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_add",
+            column_names=[],
+            columns_added=["state"],
+        )
+    )
+    builder.add_column(
+        ExpressionColumnConfig(
+            name="derived_value",
+            expr="{{ state }}_processed",
+        )
+    )
+
+    config = compile_data_designer_config(builder.build(), resource_provider)
+
+    column_names = [col.name for col in config.columns]
+    assert "state" in column_names
+    assert "city" in column_names
+    assert "age" in column_names
+    assert "derived_value" in column_names
+
+
+def test_compile_applies_processor_columns_removed(resource_provider: ResourceProvider):
+    """Test that columns declared in columns_removed are removed and cannot be referenced downstream."""
+    builder = DataDesignerConfigBuilder()
+    builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_drop",
+            column_names=[],
+            columns_removed=["city"],
+        )
+    )
+    builder.add_column(
+        ExpressionColumnConfig(
+            name="derived_value",
+            expr="{{ age }}_processed",
+        )
+    )
+
+    config = compile_data_designer_config(builder.build(), resource_provider)
+    column_names = [col.name for col in config.columns]
+    assert "city" not in column_names
+    assert "age" in column_names
+
+    # If downstream references the removed column, compilation/validation should fail
+    builder_invalid = DataDesignerConfigBuilder()
+    builder_invalid.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder_invalid.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_drop",
+            column_names=[],
+            columns_removed=["city"],
+        )
+    )
+    builder_invalid.add_column(
+        ExpressionColumnConfig(
+            name="derived_value",
+            expr="{{ city }}_processed",
+        )
+    )
+    with pytest.raises(InvalidConfigError, match="validation errors"):
+        compile_data_designer_config(builder_invalid.build(), resource_provider)
+
+
+def test_compile_processor_columns_added_collision(resource_provider: ResourceProvider):
+    """Test that adding an already existing column via columns_added raises InvalidConfigError."""
+    builder = DataDesignerConfigBuilder()
+    builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_add",
+            column_names=[],
+            columns_added=["city"],
+        )
+    )
+
+    with pytest.raises(InvalidConfigError, match="collides with an existing column"):
+        compile_data_designer_config(builder.build(), resource_provider)
+
+
+def test_compile_processor_columns_added_duplicate(resource_provider: ResourceProvider):
+    """Test that specifying duplicate columns in columns_added raises InvalidConfigError."""
+    builder = DataDesignerConfigBuilder()
+    builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_add",
+            column_names=[],
+            columns_added=["state", "state"],
+        )
+    )
+
+    with pytest.raises(InvalidConfigError, match="collides with an existing column"):
+        compile_data_designer_config(builder.build(), resource_provider)
+
+
+def test_compile_processor_columns_removed_nonexistent(resource_provider: ResourceProvider):
+    """Test that removing a non-existent column via columns_removed raises InvalidConfigError."""
+    builder = DataDesignerConfigBuilder()
+    builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv"))
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_drop",
+            column_names=[],
+            columns_removed=["non_existent"],
+        )
+    )
+
+    with pytest.raises(InvalidConfigError, match="cannot remove column 'non_existent' because it does not exist"):
+        compile_data_designer_config(builder.build(), resource_provider)
+
+
+def test_compile_processor_columns_added_without_seed_dataset(stub_resource_provider: ResourceProvider):
+    """Test that columns_added without a seed dataset raises InvalidConfigError."""
+    builder = DataDesignerConfigBuilder()
+    builder.add_processor(
+        DropColumnsProcessorConfig(
+            name="pre_batch_add",
+            column_names=[],
+            columns_added=["state"],
+        )
+    )
+    stub_resource_provider.seed_reader = None
+
+    with pytest.raises(InvalidConfigError, match="specifies 'columns_added', but no seed dataset is configured"):
+        compile_data_designer_config(builder.build(), stub_resource_provider)