refactor: callback-based processor design

docs/concepts/processors.md

@@ -13,7 +13,18 @@ Each processor:
 - Applies its transformation
 - Passes the result to the next processor (or to output)
 
-Currently, processors run only at the `POST_BATCH` stage, i.e., after column generation completes for each batch.
+Processors can run at three stages, determined by which callback methods they implement:
+
+| Stage | When it runs | Callback method | Use cases |
+|-------|--------------|-----------------|-----------|
+| Pre-batch | After seed columns, before dependent columns | `process_before_batch()` | Transform seed data before other columns are generated |
+| Post-batch | After each batch completes | `process_after_batch()` | Drop columns, transform schema per batch |
+| After generation | Once, on final dataset after all batches | `process_after_generation()` | Deduplicate, aggregate statistics, final cleanup |
+
+!!! info "Full Schema Available During Generation"
+    Each batch carries the full dataset schema during generation. Post-batch schema changes such as column dropping only alter past batches, so all columns remain accessible to generators while building follow-up batches.
+
+A processor can implement any combination of these callbacks. The built-in processors use `process_after_batch()` by default.
 
 ## Processor Types
 
@@ -134,7 +145,6 @@ Processors execute in the order they're added. Plan accordingly when one process
 | Parameter | Type | Description |
 |-----------|------|-------------|
 | `name` | str | Identifier for the processor, used in output directory names |
-| `build_stage` | BuildStage | When to run (default: `POST_BATCH`) |
 
 ### DropColumnsProcessorConfig
 

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

@@ -30,7 +30,6 @@
     from data_designer.config.config_builder import DataDesignerConfigBuilder  # noqa: F401
     from data_designer.config.custom_column import custom_column_generator  # noqa: F401
     from data_designer.config.data_designer_config import DataDesignerConfig  # noqa: F401
-    from data_designer.config.dataset_builders import BuildStage  # noqa: F401
     from data_designer.config.mcp import (  # noqa: F401
         LocalStdioMCPProvider,
         MCPProvider,
@@ -141,8 +140,6 @@
     "custom_column_generator": (f"{_MOD_BASE}.custom_column", "custom_column_generator"),
     # data_designer_config
     "DataDesignerConfig": (f"{_MOD_BASE}.data_designer_config", "DataDesignerConfig"),
-    # dataset_builders
-    "BuildStage": (f"{_MOD_BASE}.dataset_builders", "BuildStage"),
     # mcp
     "LocalStdioMCPProvider": (_MOD_MCP, "LocalStdioMCPProvider"),
     "MCPProvider": (_MOD_MCP, "MCPProvider"),

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

@@ -22,7 +22,6 @@
     get_column_display_order,
 )
 from data_designer.config.data_designer_config import DataDesignerConfig
-from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.default_model_settings import get_default_model_configs
 from data_designer.config.errors import BuilderConfigurationError, BuilderSerializationError, InvalidColumnTypeError
 from data_designer.config.exportable_config import ExportableConfigBase
@@ -572,7 +571,7 @@ def get_columns_excluding_type(self, column_type: DataDesignerColumnType) -> lis
         column_type = resolve_string_enum(column_type, DataDesignerColumnType)
         return [c for c in self._column_configs.values() if c.column_type != column_type]
 
-    def get_processor_configs(self) -> dict[BuildStage, list[ProcessorConfigT]]:
+    def get_processor_configs(self) -> list[ProcessorConfigT]:
         """Get processor configuration objects.
 
         Returns:

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

@@ -12,11 +12,8 @@
 from typing_extensions import TypeAlias
 
 from data_designer.config.base import ConfigBase
-from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.errors import InvalidConfigError
 
-SUPPORTED_STAGES = [BuildStage.POST_BATCH]
-
 
 class ProcessorType(str, Enum):
     """Enumeration of available processor types.
@@ -33,33 +30,22 @@ class ProcessorType(str, Enum):
 class ProcessorConfig(ConfigBase, ABC):
     """Abstract base class for all processor configuration types.
 
-    Processors are transformations that run before or after columns are generated.
-    They can modify, reshape, or augment the dataset before it's saved.
+    Processors are transformations that run at different stages of the generation
+    pipeline. They can modify, reshape, or augment the dataset.
+
+    The processor implementation determines which stages it handles by overriding
+    the appropriate callback methods (process_before_batch, process_after_batch, process_after_generation).
 
     Attributes:
         name: Unique name of the processor, used to identify the processor in results
             and to name output artifacts on disk.
-        build_stage: The stage at which the processor runs. Currently only `POST_BATCH`
-            is supported, meaning processors run after each batch of columns is generated.
     """
 
     name: str = Field(
         description="The name of the processor, used to identify the processor in the results and to write the artifacts to disk.",
     )
-    build_stage: BuildStage = Field(
-        default=BuildStage.POST_BATCH,
-        description=f"The stage at which the processor will run. Supported stages: {', '.join(SUPPORTED_STAGES)}",
-    )
     processor_type: str
 
-    @field_validator("build_stage")
-    def validate_build_stage(cls, v: BuildStage) -> BuildStage:
-        if v not in SUPPORTED_STAGES:
-            raise ValueError(
-                f"Invalid dataset builder stage: {v}. Only these stages are supported: {', '.join(SUPPORTED_STAGES)}"
-            )
-        return v
-
 
 def get_processor_config_from_kwargs(processor_type: ProcessorType, **kwargs: Any) -> ProcessorConfig:
     """Create a processor configuration from a processor type and keyword arguments.

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

@@ -4,7 +4,6 @@
 import pytest
 from pydantic import ValidationError
 
-from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.errors import InvalidConfigError
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
@@ -16,92 +15,64 @@
 
 
 def test_drop_columns_processor_config_creation():
-    config = DropColumnsProcessorConfig(
-        name="drop_columns_processor", build_stage=BuildStage.POST_BATCH, column_names=["col1", "col2"]
-    )
+    config = DropColumnsProcessorConfig(name="drop_columns_processor", column_names=["col1", "col2"])
 
-    assert config.build_stage == BuildStage.POST_BATCH
     assert config.column_names == ["col1", "col2"]
     assert config.processor_type == ProcessorType.DROP_COLUMNS
     assert isinstance(config, ProcessorConfig)
 
 
 def test_drop_columns_processor_config_validation():
-    # Test unsupported stage raises error
-    with pytest.raises(ValidationError, match="Invalid dataset builder stage"):
-        DropColumnsProcessorConfig(
-            name="drop_columns_processor", build_stage=BuildStage.PRE_BATCH, column_names=["col1"]
-        )
-
     # Test missing required field raises error
     with pytest.raises(ValidationError, match="Field required"):
-        DropColumnsProcessorConfig(name="drop_columns_processor", build_stage=BuildStage.POST_BATCH)
+        DropColumnsProcessorConfig(name="drop_columns_processor")
 
 
 def test_drop_columns_processor_config_serialization():
-    config = DropColumnsProcessorConfig(
-        name="drop_columns_processor", build_stage=BuildStage.POST_BATCH, column_names=["col1", "col2"]
-    )
+    config = DropColumnsProcessorConfig(name="drop_columns_processor", column_names=["col1", "col2"])
 
     # Serialize to dict
     config_dict = config.model_dump()
-    assert config_dict["build_stage"] == "post_batch"
     assert config_dict["column_names"] == ["col1", "col2"]
 
     # Deserialize from dict
     config_restored = DropColumnsProcessorConfig.model_validate(config_dict)
-    assert config_restored.build_stage == config.build_stage
     assert config_restored.column_names == config.column_names
 
 
 def test_schema_transform_processor_config_creation():
     config = SchemaTransformProcessorConfig(
         name="output_format_processor",
-        build_stage=BuildStage.POST_BATCH,
         template={"text": "{{ col1 }}"},
     )
 
-    assert config.build_stage == BuildStage.POST_BATCH
     assert config.template == {"text": "{{ col1 }}"}
     assert config.processor_type == ProcessorType.SCHEMA_TRANSFORM
     assert isinstance(config, ProcessorConfig)
 
 
 def test_schema_transform_processor_config_validation():
-    # Test unsupported stage raises error
-    with pytest.raises(ValidationError, match="Invalid dataset builder stage"):
-        SchemaTransformProcessorConfig(
-            name="schema_transform_processor",
-            build_stage=BuildStage.PRE_BATCH,
-            template={"text": "{{ col1 }}"},
-        )
-
     # Test missing required field raises error
     with pytest.raises(ValidationError, match="Field required"):
-        SchemaTransformProcessorConfig(name="schema_transform_processor", build_stage=BuildStage.POST_BATCH)
+        SchemaTransformProcessorConfig(name="schema_transform_processor")
 
     # Test invalid template raises error
     with pytest.raises(InvalidConfigError, match="Template must be JSON serializable"):
-        SchemaTransformProcessorConfig(
-            name="schema_transform_processor", build_stage=BuildStage.POST_BATCH, template={"text": {1, 2, 3}}
-        )
+        SchemaTransformProcessorConfig(name="schema_transform_processor", template={"text": {1, 2, 3}})
 
 
 def test_schema_transform_processor_config_serialization():
     config = SchemaTransformProcessorConfig(
         name="schema_transform_processor",
-        build_stage=BuildStage.POST_BATCH,
         template={"text": "{{ col1 }}"},
     )
 
     # Serialize to dict
     config_dict = config.model_dump()
-    assert config_dict["build_stage"] == "post_batch"
     assert config_dict["template"] == {"text": "{{ col1 }}"}
 
     # Deserialize from dict
     config_restored = SchemaTransformProcessorConfig.model_validate(config_dict)
-    assert config_restored.build_stage == config.build_stage
     assert config_restored.template == config.template
 
 
@@ -110,7 +81,6 @@ def test_get_processor_config_from_kwargs():
     config_drop_columns = get_processor_config_from_kwargs(
         ProcessorType.DROP_COLUMNS,
         name="drop_columns_processor",
-        build_stage=BuildStage.POST_BATCH,
         column_names=["col1"],
     )
     assert isinstance(config_drop_columns, DropColumnsProcessorConfig)
@@ -120,7 +90,6 @@ def test_get_processor_config_from_kwargs():
     config_schema_transform = get_processor_config_from_kwargs(
         ProcessorType.SCHEMA_TRANSFORM,
         name="output_format_processor",
-        build_stage=BuildStage.POST_BATCH,
         template={"text": "{{ col1 }}"},
     )
     assert isinstance(config_schema_transform, SchemaTransformProcessorConfig)
@@ -134,6 +103,6 @@ class UnknownProcessorType(str, Enum):
         UNKNOWN = "unknown"
 
     result = get_processor_config_from_kwargs(
-        UnknownProcessorType.UNKNOWN, name="unknown_processor", build_stage=BuildStage.POST_BATCH, column_names=["col1"]
+        UnknownProcessorType.UNKNOWN, name="unknown_processor", column_names=["col1"]
     )
     assert result is None