From c788979e273b86c71e0f9a34fcb3ac881165a360 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:25:47 +0530 Subject: [PATCH 1/2] fix(compiler): support columns_added and columns_removed in processor config (#394) Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../src/data_designer/config/base.py | 8 ++ .../tests/config/test_processors.py | 27 ++++ .../src/data_designer/engine/compiler.py | 28 +++++ .../tests/engine/test_compiler.py | 119 ++++++++++++++++++ 4 files changed, 182 insertions(+) diff --git a/packages/data-designer-config/src/data_designer/config/base.py b/packages/data-designer-config/src/data_designer/config/base.py index 8910e4653..dcc00e572 100644 --- a/packages/data-designer-config/src/data_designer/config/base.py +++ b/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.", + ) diff --git a/packages/data-designer-config/tests/config/test_processors.py b/packages/data-designer-config/tests/config/test_processors.py index e688be15e..2d813da31 100644 --- a/packages/data-designer-config/tests/config/test_processors.py +++ b/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"] diff --git a/packages/data-designer-engine/src/data_designer/engine/compiler.py b/packages/data-designer-engine/src/data_designer/engine/compiler.py index bef430f92..ea901b887 100644 --- a/packages/data-designer-engine/src/data_designer/engine/compiler.py +++ b/packages/data-designer-engine/src/data_designer/engine/compiler.py @@ -18,11 +18,39 @@ 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)) + + 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. diff --git a/packages/data-designer-engine/tests/engine/test_compiler.py b/packages/data-designer-engine/tests/engine/test_compiler.py index 4c2ce179b..9d6a6476a 100644 --- a/packages/data-designer-engine/tests/engine/test_compiler.py +++ b/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,121 @@ 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_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) From 7f58dd1c947286684ebf745a2a347eb41415ff38 Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:53:23 +0530 Subject: [PATCH 2/2] fix(compiler): prevent duplicate columns within processor columns_added Signed-off-by: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> --- .../src/data_designer/engine/compiler.py | 1 + .../tests/engine/test_compiler.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/packages/data-designer-engine/src/data_designer/engine/compiler.py b/packages/data-designer-engine/src/data_designer/engine/compiler.py index ea901b887..a66dd9096 100644 --- a/packages/data-designer-engine/src/data_designer/engine/compiler.py +++ b/packages/data-designer-engine/src/data_designer/engine/compiler.py @@ -49,6 +49,7 @@ def _apply_processor_column_modifications(config: DataDesignerConfig) -> None: 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: diff --git a/packages/data-designer-engine/tests/engine/test_compiler.py b/packages/data-designer-engine/tests/engine/test_compiler.py index 9d6a6476a..8287d7ae1 100644 --- a/packages/data-designer-engine/tests/engine/test_compiler.py +++ b/packages/data-designer-engine/tests/engine/test_compiler.py @@ -263,6 +263,22 @@ def test_compile_processor_columns_added_collision(resource_provider: ResourcePr 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()