diff --git a/src/data_designer/engine/compiler.py b/src/data_designer/engine/compiler.py index e9690fb86..fa3917269 100644 --- a/src/data_designer/engine/compiler.py +++ b/src/data_designer/engine/compiler.py @@ -5,9 +5,10 @@ import logging -from data_designer.config.column_configs import SeedDatasetColumnConfig +from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig from data_designer.config.data_designer_config import DataDesignerConfig from data_designer.config.errors import InvalidConfigError +from data_designer.config.sampler_params import UUIDSamplerParams from data_designer.engine.resources.resource_provider import ResourceProvider from data_designer.engine.resources.seed_reader import SeedReader from data_designer.engine.validation import ViolationLevel, rich_print_violations, validate_data_designer_config @@ -17,6 +18,7 @@ def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig: _resolve_and_add_seed_columns(config, resource_provider.seed_reader) + _add_internal_row_id_column_if_needed(config) _validate(config) return config @@ -41,6 +43,35 @@ def _resolve_and_add_seed_columns(config: DataDesignerConfig, seed_reader: SeedR config.columns.extend([SeedDatasetColumnConfig(name=col_name) for col_name in seed_col_names]) +def _add_internal_row_id_column_if_needed(config: DataDesignerConfig) -> None: + """Adds a UUID sampler column named '_internal_row_id' (set to drop) if needed to enable generation. + + Generation requires either: + - At least one sampler column (which can generate data from scratch), OR + - A seed dataset (which provides initial data rows) + + If neither exists, a UUID sampler column '_internal_row_id' is automatically added and marked for drop + to enable the generation process to start. + + Args: + config: The DataDesigner configuration to potentially modify. + """ + has_sampler_column = any(isinstance(col, SamplerColumnConfig) for col in config.columns) + has_seed_dataset_column = any(isinstance(col, SeedDatasetColumnConfig) for col in config.columns) + + if not has_sampler_column and not has_seed_dataset_column: + logger.warning( + "🔔 No sampler column or seed dataset detected. Adding UUID column '_internal_row_id' (marked for drop) to enable generation." + ) + id_column = SamplerColumnConfig( + name="_internal_row_id", + sampler_type="uuid", + params=UUIDSamplerParams(), + drop=True, + ) + config.columns.insert(0, id_column) + + def _validate(config: DataDesignerConfig) -> None: allowed_references = _get_allowed_references(config) violations = validate_data_designer_config( diff --git a/tests/engine/dataset_builders/test_column_wise_builder.py b/tests/engine/dataset_builders/test_column_wise_builder.py index 08c3545f8..6aab3d5e8 100644 --- a/tests/engine/dataset_builders/test_column_wise_builder.py +++ b/tests/engine/dataset_builders/test_column_wise_builder.py @@ -13,6 +13,7 @@ from data_designer.config.dataset_builders import BuildStage from data_designer.config.processors import DropColumnsProcessorConfig from data_designer.config.run_config import RunConfig +from data_designer.config.sampler_params import SamplerType, UUIDSamplerParams from data_designer.engine.column_generators.generators.base import GenerationStrategy from data_designer.engine.dataset_builders.column_wise_builder import ( MAX_CONCURRENCY_PER_NON_LLM_GENERATOR, @@ -31,6 +32,7 @@ @pytest.fixture def stub_test_column_configs(): return [ + SamplerColumnConfig(name="some_id", sampler_type=SamplerType.UUID, params=UUIDSamplerParams()), LLMTextColumnConfig(name="test_column", prompt="Test prompt", model_alias="test_model"), LLMTextColumnConfig(name="column_to_drop", prompt="Test prompt", model_alias="test_model"), ] @@ -92,8 +94,7 @@ def test_column_wise_dataset_builder_creation(stub_resource_provider, stub_test_ data_designer_config=stub_test_config_builder.build(), resource_provider=stub_resource_provider, ) - - assert len(builder._column_configs) == 2 + assert len(builder._column_configs) == 3 assert builder._resource_provider == stub_resource_provider assert isinstance(builder._registry, DataDesignerRegistry) @@ -136,26 +137,33 @@ def test_column_wise_dataset_builder_batch_manager_initialization(stub_column_wi def test_column_wise_dataset_builder_single_column_configs_property( stub_resource_provider, stub_model_configs, config_type, expected_single_configs ): + config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + if config_type == "single": - single_config = LLMTextColumnConfig(name="test_column", prompt="Test prompt", model_alias="test_model") - config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + # Add an LLM text column - these don't get grouped into MultiColumnConfigs + single_config = expected_single_configs[0] config_builder.add_column(single_config) + builder = ColumnWiseDatasetBuilder( data_designer_config=config_builder.build(), resource_provider=stub_resource_provider, ) - assert builder.single_column_configs == [single_config] + + # Since there's no sampler, _internal_row_id is auto-added, plus the LLM column + configs = builder.single_column_configs + assert len(configs) == 2 + assert configs[0].name == "_internal_row_id" + assert configs[1] == single_config + else: - sampler_config = SamplerColumnConfig( - name="sampler_col", sampler_type="category", params={"values": ["A", "B", "C"]} - ) - config_builder = DataDesignerConfigBuilder(model_configs=stub_model_configs) + sampler_config = expected_single_configs[0] config_builder.add_column(sampler_config) + builder = ColumnWiseDatasetBuilder( data_designer_config=config_builder.build(), resource_provider=stub_resource_provider, ) - assert builder.single_column_configs == [sampler_config] + assert builder.single_column_configs == expected_single_configs def test_column_wise_dataset_builder_build_method_basic_flow( diff --git a/tests/engine/test_compiler.py b/tests/engine/test_compiler.py index 92663dd94..fb3f9dbc0 100644 --- a/tests/engine/test_compiler.py +++ b/tests/engine/test_compiler.py @@ -5,10 +5,10 @@ import pytest -from data_designer.config.column_configs import SamplerColumnConfig +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.sampler_params import CategorySamplerParams, SamplerType +from data_designer.config.sampler_params import CategorySamplerParams, SamplerType, UUIDSamplerParams from data_designer.config.seed_source import HuggingFaceSeedSource from data_designer.engine.compiler import compile_data_designer_config from data_designer.engine.resources.resource_provider import ResourceProvider @@ -78,3 +78,69 @@ def test_validation_errors(resource_provider: ResourceProvider): compile_data_designer_config(builder.build(), resource_provider) assert "validation errors" in str(excinfo) + + +def test_adds_id_column_when_no_sampler_and_no_seed_dataset(stub_resource_provider: ResourceProvider): + """Test that a UUID '_internal_row_id' column is automatically added when there's no sampler column or seed dataset.""" + builder = DataDesignerConfigBuilder() + builder.add_column( + ExpressionColumnConfig( + name="derived_value", + expr="'constant_value'", + ) + ) + stub_resource_provider.seed_reader = None + + config = compile_data_designer_config(builder.build(), stub_resource_provider) + + assert len(config.columns) == 2 + assert config.columns[0].name == "_internal_row_id" + assert isinstance(config.columns[0], SamplerColumnConfig) + assert config.columns[0].sampler_type == "uuid" + assert isinstance(config.columns[0].params, UUIDSamplerParams) + assert config.columns[0].drop is True + + +def test_does_not_add_id_column_when_sampler_exists(stub_resource_provider: ResourceProvider): + """Test that no '_internal_row_id' column is added when a sampler column already exists.""" + builder = DataDesignerConfigBuilder() + builder.add_column( + SamplerColumnConfig( + name="category", + sampler_type=SamplerType.CATEGORY, + params=CategorySamplerParams(values=["a", "b", "c"]), + ) + ) + builder.add_column( + ExpressionColumnConfig( + name="derived_value", + expr="{{ category }}_suffix", + ) + ) + stub_resource_provider.seed_reader = None + + config = compile_data_designer_config(builder.build(), stub_resource_provider) + + assert len(config.columns) == 2 + assert config.columns[0].name == "category" + assert config.columns[1].name == "derived_value" + assert not any(col.name == "_internal_row_id" for col in config.columns) + + +def test_does_not_add_id_column_when_seed_dataset_exists(resource_provider: ResourceProvider): + """Test that no '_internal_row_id' column is added when a seed dataset is configured.""" + builder = DataDesignerConfigBuilder() + builder.add_column( + ExpressionColumnConfig( + name="derived_value", + expr="{{ city }}_derived", + ) + ) + builder.with_seed_dataset(HuggingFaceSeedSource(path="hf://datasets/test/data.csv")) + + config = compile_data_designer_config(builder.build(), resource_provider) + + # Should have the expression column + 2 seed columns (city, country) from the fixture + 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)