data-designer 0.1.1__tar.gz → 0.1.2__tar.gz

{data_designer-0.1.1 → data_designer-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer
-Version: 0.1.1
+Version: 0.1.2
 Summary: General framework for synthetic data generation
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -97,8 +97,7 @@ export NVIDIA_API_KEY="your-api-key-here"
 export OPENAI_API_KEY="your-openai-api-key-here"
 ```
 
-### 3. Generate your first dataset
-
+### 3. Start generating data!
 ```python
 from data_designer.essentials import (
     CategorySamplerParams,
@@ -139,8 +138,6 @@ preview = data_designer.preview(config_builder=config_builder)
 preview.display_sample_record()
 ```
 
-**That's it!** You've created a dataset.
-
 ---
 
 ## What's next?
@@ -148,7 +145,7 @@ preview.display_sample_record()
 ### 📚 Learn more
 
 - **[Quick Start Guide](https://nvidia-nemo.github.io/DataDesigner/quick-start/)** – Detailed walkthrough with more examples
-- **[Tutorial Notebooks](https://nvidia-nemo.github.io/DataDesigner/notebooks/intro/)** – Step-by-step interactive tutorials
+- **[Tutorial Notebooks](https://nvidia-nemo.github.io/DataDesigner/notebooks/)** – Step-by-step interactive tutorials
 - **[Column Types](https://nvidia-nemo.github.io/DataDesigner/concepts/columns/)** – Explore samplers, LLM columns, validators, and more
 - **[Validators](https://nvidia-nemo.github.io/DataDesigner/concepts/validators/)** – Learn how to validate generated data with Python, SQL, and remote validators
 - **[Model Configuration](https://nvidia-nemo.github.io/DataDesigner/models/model-configs/)** – Configure custom models and providers

{data_designer-0.1.1 → data_designer-0.1.2}/README.md

@@ -48,8 +48,7 @@ export NVIDIA_API_KEY="your-api-key-here"
 export OPENAI_API_KEY="your-openai-api-key-here"
 ```
 
-### 3. Generate your first dataset
-
+### 3. Start generating data!
 ```python
 from data_designer.essentials import (
     CategorySamplerParams,
@@ -90,8 +89,6 @@ preview = data_designer.preview(config_builder=config_builder)
 preview.display_sample_record()
 ```
 
-**That's it!** You've created a dataset.
-
 ---
 
 ## What's next?
@@ -99,7 +96,7 @@ preview.display_sample_record()
 ### 📚 Learn more
 
 - **[Quick Start Guide](https://nvidia-nemo.github.io/DataDesigner/quick-start/)** – Detailed walkthrough with more examples
-- **[Tutorial Notebooks](https://nvidia-nemo.github.io/DataDesigner/notebooks/intro/)** – Step-by-step interactive tutorials
+- **[Tutorial Notebooks](https://nvidia-nemo.github.io/DataDesigner/notebooks/)** – Step-by-step interactive tutorials
 - **[Column Types](https://nvidia-nemo.github.io/DataDesigner/concepts/columns/)** – Explore samplers, LLM columns, validators, and more
 - **[Validators](https://nvidia-nemo.github.io/DataDesigner/concepts/validators/)** – Learn how to validate generated data with Python, SQL, and remote validators
 - **[Model Configuration](https://nvidia-nemo.github.io/DataDesigner/models/model-configs/)** – Configure custom models and providers

{data_designer-0.1.1 → data_designer-0.1.2}/docs/index.md

@@ -34,11 +34,11 @@ Data Designer helps you create datasets through an intuitive, **iterative** proc
 3.  **🔁 Preview** your results and iterate
     - Generate a preview dataset stored in memory for fast iteration
     - Inspect sample records and analysis results to refine your configuration
-    - Try for yourself by running the [tutorial notebooks](notebooks/intro.md)
+    - Try for yourself by running the [tutorial notebooks](notebooks/README.md)
 4.  **🖼️ Create** your dataset
     - Generate your full dataset and save results to disk
     - Access the generated dataset and associated artifacts for downstream use
-    - Give it a try by running the [tutorial notebooks](notebooks/intro.md)
+    - Give it a try by running the [tutorial notebooks](notebooks/README.md)
 
 ## Library and Microservice
 

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.1.2'
+__version_tuple__ = version_tuple = (0, 1, 2)
 
 __commit_id__ = commit_id = None

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/config/column_configs.py

@@ -2,9 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC
-from typing import Literal, Optional, Type, Union
+from typing import Annotated, Literal, Optional, Type, Union
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Discriminator, Field, model_validator
 from typing_extensions import Self
 
 from .base import ConfigBase
@@ -89,11 +89,36 @@ class SamplerColumnConfig(SingleColumnConfig):
     """
 
     sampler_type: SamplerType
-    params: SamplerParamsT
-    conditional_params: dict[str, SamplerParamsT] = {}
+    params: Annotated[SamplerParamsT, Discriminator("sampler_type")]
+    conditional_params: dict[str, Annotated[SamplerParamsT, Discriminator("sampler_type")]] = {}
     convert_to: Optional[str] = None
     column_type: Literal["sampler"] = "sampler"
 
+    @model_validator(mode="before")
+    @classmethod
+    def inject_sampler_type_into_params(cls, data: dict) -> dict:
+        """Inject sampler_type into params dict to enable discriminated union resolution.
+
+        This allows users to pass params as a simple dict without the sampler_type field,
+        which will be automatically added based on the outer sampler_type field.
+        """
+        if isinstance(data, dict):
+            sampler_type = data.get("sampler_type")
+            params = data.get("params")
+
+            # If params is a dict and doesn't have sampler_type, inject it
+            if sampler_type and isinstance(params, dict) and "sampler_type" not in params:
+                data["params"] = {"sampler_type": sampler_type, **params}
+
+            # Handle conditional_params similarly
+            conditional_params = data.get("conditional_params")
+            if conditional_params and isinstance(conditional_params, dict):
+                for condition, cond_params in conditional_params.items():
+                    if isinstance(cond_params, dict) and "sampler_type" not in cond_params:
+                        data["conditional_params"][condition] = {"sampler_type": sampler_type, **cond_params}
+
+        return data
+
 
 class LLMTextColumnConfig(SingleColumnConfig):
     """Configuration for text generation columns using Large Language Models.

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/config/datastore.py

@@ -31,34 +31,37 @@ class DatastoreSettings(BaseModel):
     token: Optional[str] = Field(default=None, description="If needed, token to use for authentication.")
 
 
-def get_file_column_names(file_path: Union[str, Path], file_type: str) -> list[str]:
-    """Extract column names based on file type. Supports glob patterns like '../path/*.parquet'."""
-    file_path = Path(file_path)
-    if "*" in str(file_path):
-        matching_files = sorted(file_path.parent.glob(file_path.name))
-        if not matching_files:
-            raise InvalidFilePathError(f"🛑 No files found matching pattern: {str(file_path)!r}")
-        logger.debug(f"0️⃣ Using the first matching file in {str(file_path)!r} to determine column names in seed dataset")
-        file_path = matching_files[0]
+def get_file_column_names(file_reference: Union[str, Path, HfFileSystem], file_type: str) -> list[str]:
+    """Get column names from a dataset file.
+
+    Args:
+        file_reference: Path to the dataset file, or an HfFileSystem object.
+        file_type: Type of the dataset file. Must be one of: 'parquet', 'json', 'jsonl', 'csv'.
 
+    Raises:
+        InvalidFilePathError: If the file type is not supported.
+
+    Returns:
+        List of column names.
+    """
     if file_type == "parquet":
         try:
-            schema = pq.read_schema(file_path)
+            schema = pq.read_schema(file_reference)
             if hasattr(schema, "names"):
                 return schema.names
             else:
                 return [field.name for field in schema]
         except Exception as e:
-            logger.warning(f"Failed to process parquet file {file_path}: {e}")
+            logger.warning(f"Failed to process parquet file {file_reference}: {e}")
             return []
     elif file_type in ["json", "jsonl"]:
-        return pd.read_json(file_path, orient="records", lines=True, nrows=1).columns.tolist()
+        return pd.read_json(file_reference, orient="records", lines=True, nrows=1).columns.tolist()
     elif file_type == "csv":
         try:
-            df = pd.read_csv(file_path, nrows=1)
+            df = pd.read_csv(file_reference, nrows=1)
             return df.columns.tolist()
         except (pd.errors.EmptyDataError, pd.errors.ParserError) as e:
-            logger.warning(f"Failed to process CSV file {file_path}: {e}")
+            logger.warning(f"Failed to process CSV file {file_reference}: {e}")
             return []
     else:
         raise InvalidFilePathError(f"🛑 Unsupported file type: {file_type!r}")
@@ -66,12 +69,36 @@ def get_file_column_names(file_path: Union[str, Path], file_type: str) -> list[s
 
 def fetch_seed_dataset_column_names(seed_dataset_reference: SeedDatasetReference) -> list[str]:
     if hasattr(seed_dataset_reference, "datastore_settings"):
-        return _fetch_seed_dataset_column_names_from_datastore(
+        return fetch_seed_dataset_column_names_from_datastore(
             seed_dataset_reference.repo_id,
             seed_dataset_reference.filename,
             seed_dataset_reference.datastore_settings,
         )
-    return _fetch_seed_dataset_column_names_from_local_file(seed_dataset_reference.dataset)
+    return fetch_seed_dataset_column_names_from_local_file(seed_dataset_reference.dataset)
+
+
+def fetch_seed_dataset_column_names_from_datastore(
+    repo_id: str,
+    filename: str,
+    datastore_settings: Optional[Union[DatastoreSettings, dict]] = None,
+) -> list[str]:
+    file_type = filename.split(".")[-1]
+    if f".{file_type}" not in VALID_DATASET_FILE_EXTENSIONS:
+        raise InvalidFileFormatError(f"🛑 Unsupported file type: {filename!r}")
+
+    datastore_settings = resolve_datastore_settings(datastore_settings)
+    fs = HfFileSystem(endpoint=datastore_settings.endpoint, token=datastore_settings.token, skip_instance_cache=True)
+
+    file_path = _extract_single_file_path_from_glob_pattern_if_present(f"datasets/{repo_id}/{filename}", fs=fs)
+
+    with fs.open(file_path) as f:
+        return get_file_column_names(f, file_type)
+
+
+def fetch_seed_dataset_column_names_from_local_file(dataset_path: str | Path) -> list[str]:
+    dataset_path = _validate_dataset_path(dataset_path, allow_glob_pattern=True)
+    dataset_path = _extract_single_file_path_from_glob_pattern_if_present(dataset_path)
+    return get_file_column_names(dataset_path, str(dataset_path).split(".")[-1])
 
 
 def resolve_datastore_settings(datastore_settings: DatastoreSettings | dict | None) -> DatastoreSettings:
@@ -114,25 +141,34 @@ def upload_to_hf_hub(
     return f"{repo_id}/{filename}"
 
 
-def _fetch_seed_dataset_column_names_from_datastore(
-    repo_id: str,
-    filename: str,
-    datastore_settings: Optional[Union[DatastoreSettings, dict]] = None,
-) -> list[str]:
-    file_type = filename.split(".")[-1]
-    if f".{file_type}" not in VALID_DATASET_FILE_EXTENSIONS:
-        raise InvalidFileFormatError(f"🛑 Unsupported file type: {filename!r}")
-
-    datastore_settings = resolve_datastore_settings(datastore_settings)
-    fs = HfFileSystem(endpoint=datastore_settings.endpoint, token=datastore_settings.token, skip_instance_cache=True)
-
-    with fs.open(f"datasets/{repo_id}/{filename}") as f:
-        return get_file_column_names(f, file_type)
-
+def _extract_single_file_path_from_glob_pattern_if_present(
+    file_path: str | Path,
+    fs: HfFileSystem | None = None,
+) -> Path:
+    file_path = Path(file_path)
 
-def _fetch_seed_dataset_column_names_from_local_file(dataset_path: str | Path) -> list[str]:
-    dataset_path = _validate_dataset_path(dataset_path, allow_glob_pattern=True)
-    return get_file_column_names(dataset_path, str(dataset_path).split(".")[-1])
+    # no glob pattern
+    if "*" not in str(file_path):
+        return file_path
+
+    # glob pattern with HfFileSystem
+    if fs is not None:
+        file_to_check = None
+        file_extension = file_path.name.split(".")[-1]
+        for file in fs.ls(str(file_path.parent)):
+            filename = file["name"]
+            if filename.endswith(f".{file_extension}"):
+                file_to_check = filename
+        if file_to_check is None:
+            raise InvalidFilePathError(f"🛑 No files found matching pattern: {str(file_path)!r}")
+        logger.debug(f"Using the first matching file in {str(file_path)!r} to determine column names in seed dataset")
+        return Path(file_to_check)
+
+    # glob pattern with local file system
+    if not (matching_files := sorted(file_path.parent.glob(file_path.name))):
+        raise InvalidFilePathError(f"🛑 No files found matching pattern: {str(file_path)!r}")
+    logger.debug(f"Using the first matching file in {str(file_path)!r} to determine column names in seed dataset")
+    return matching_files[0]
 
 
 def _validate_dataset_path(dataset_path: Union[str, Path], allow_glob_pattern: bool = False) -> Path:

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/config/default_model_settings.py

@@ -78,7 +78,7 @@ def get_default_model_configs() -> list[ModelConfig]:
     return []
 
 
-def get_defaul_model_providers_missing_api_keys() -> list[str]:
+def get_default_model_providers_missing_api_keys() -> list[str]:
     missing_api_keys = []
     for predefined_provider in PREDEFINED_PROVIDERS:
         if os.environ.get(predefined_provider["api_key"]) is None:

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/config/sampler_params.py

@@ -66,6 +66,7 @@ class CategorySamplerParams(ConfigBase):
             "Larger values will be sampled with higher probability."
         ),
     )
+    sampler_type: Literal[SamplerType.CATEGORY] = SamplerType.CATEGORY
 
     @model_validator(mode="after")
     def _normalize_weights_if_needed(self) -> Self:
@@ -106,6 +107,7 @@ class DatetimeSamplerParams(ConfigBase):
         default="D",
         description="Sampling units, e.g. the smallest possible time interval between samples.",
     )
+    sampler_type: Literal[SamplerType.DATETIME] = SamplerType.DATETIME
 
     @field_validator("start", "end")
     @classmethod
@@ -136,6 +138,7 @@ class SubcategorySamplerParams(ConfigBase):
         ...,
         description="Mapping from each value of parent category to a list of subcategory values.",
     )
+    sampler_type: Literal[SamplerType.SUBCATEGORY] = SamplerType.SUBCATEGORY
 
 
 class TimeDeltaSamplerParams(ConfigBase):
@@ -187,6 +190,7 @@ class TimeDeltaSamplerParams(ConfigBase):
         default="D",
         description="Sampling units, e.g. the smallest possible time interval between samples.",
     )
+    sampler_type: Literal[SamplerType.TIMEDELTA] = SamplerType.TIMEDELTA
 
     @model_validator(mode="after")
     def _validate_min_less_than_max(self) -> Self:
@@ -219,6 +223,7 @@ class UUIDSamplerParams(ConfigBase):
         default=False,
         description="If true, all letters in the UUID will be capitalized.",
     )
+    sampler_type: Literal[SamplerType.UUID] = SamplerType.UUID
 
     @property
     def last_index(self) -> int:
@@ -257,6 +262,7 @@ class ScipySamplerParams(ConfigBase):
     decimal_places: Optional[int] = Field(
         default=None, description="Number of decimal places to round the sampled values to."
     )
+    sampler_type: Literal[SamplerType.SCIPY] = SamplerType.SCIPY
 
 
 class BinomialSamplerParams(ConfigBase):
@@ -273,6 +279,7 @@ class BinomialSamplerParams(ConfigBase):
 
     n: int = Field(..., description="Number of trials.")
     p: float = Field(..., description="Probability of success on each trial.", ge=0.0, le=1.0)
+    sampler_type: Literal[SamplerType.BINOMIAL] = SamplerType.BINOMIAL
 
 
 class BernoulliSamplerParams(ConfigBase):
@@ -288,6 +295,7 @@ class BernoulliSamplerParams(ConfigBase):
     """
 
     p: float = Field(..., description="Probability of success.", ge=0.0, le=1.0)
+    sampler_type: Literal[SamplerType.BERNOULLI] = SamplerType.BERNOULLI
 
 
 class BernoulliMixtureSamplerParams(ConfigBase):
@@ -327,6 +335,7 @@ class BernoulliMixtureSamplerParams(ConfigBase):
         ...,
         description="Parameters of the scipy.stats distribution given in `dist_name`.",
     )
+    sampler_type: Literal[SamplerType.BERNOULLI_MIXTURE] = SamplerType.BERNOULLI_MIXTURE
 
 
 class GaussianSamplerParams(ConfigBase):
@@ -350,6 +359,7 @@ class GaussianSamplerParams(ConfigBase):
     decimal_places: Optional[int] = Field(
         default=None, description="Number of decimal places to round the sampled values to."
     )
+    sampler_type: Literal[SamplerType.GAUSSIAN] = SamplerType.GAUSSIAN
 
 
 class PoissonSamplerParams(ConfigBase):
@@ -369,6 +379,7 @@ class PoissonSamplerParams(ConfigBase):
     """
 
     mean: float = Field(..., description="Mean number of events in a fixed interval.")
+    sampler_type: Literal[SamplerType.POISSON] = SamplerType.POISSON
 
 
 class UniformSamplerParams(ConfigBase):
@@ -390,6 +401,7 @@ class UniformSamplerParams(ConfigBase):
     decimal_places: Optional[int] = Field(
         default=None, description="Number of decimal places to round the sampled values to."
     )
+    sampler_type: Literal[SamplerType.UNIFORM] = SamplerType.UNIFORM
 
 
 #########################################
@@ -470,11 +482,12 @@ class PersonSamplerParams(ConfigBase):
         default=False,
         description="If True, then append synthetic persona columns to each generated person.",
     )
+    sampler_type: Literal[SamplerType.PERSON] = SamplerType.PERSON
 
     @property
     def generator_kwargs(self) -> list[str]:
         """Keyword arguments to pass to the person generator."""
-        return [f for f in list(PersonSamplerParams.model_fields) if f != "locale"]
+        return [f for f in list(PersonSamplerParams.model_fields) if f not in ("locale", "sampler_type")]
 
     @property
     def people_gen_key(self) -> str:
@@ -533,11 +546,12 @@ class PersonFromFakerSamplerParams(ConfigBase):
         min_length=2,
         max_length=2,
     )
+    sampler_type: Literal[SamplerType.PERSON_FROM_FAKER] = SamplerType.PERSON_FROM_FAKER
 
     @property
     def generator_kwargs(self) -> list[str]:
         """Keyword arguments to pass to the person generator."""
-        return [f for f in list(PersonFromFakerSamplerParams.model_fields) if f != "locale"]
+        return [f for f in list(PersonFromFakerSamplerParams.model_fields) if f not in ("locale", "sampler_type")]
 
     @property
     def people_gen_key(self) -> str:

{data_designer-0.1.1 → data_designer-0.1.2}/src/data_designer/interface/data_designer.py

@@ -9,8 +9,8 @@ import pandas as pd
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.default_model_settings import (
-    get_defaul_model_providers_missing_api_keys,
     get_default_model_configs,
+    get_default_model_providers_missing_api_keys,
     get_default_provider_name,
     get_default_providers,
     resolve_seed_default_model_settings,
@@ -313,7 +313,7 @@ class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
         if model_providers is None:
             if can_run_data_designer_locally():
                 model_providers = get_default_providers()
-                missing_api_keys = get_defaul_model_providers_missing_api_keys()
+                missing_api_keys = get_default_model_providers_missing_api_keys()
                 if len(missing_api_keys) == len(PREDEFINED_PROVIDERS):
                     logger.warning(
                         "🚨 You are trying to use a default model provider but your API keys are missing."

{data_designer-0.1.1 → data_designer-0.1.2}/tests/config/test_columns.py

@@ -23,7 +23,15 @@ from data_designer.config.column_types import (
     get_column_display_order,
 )
 from data_designer.config.errors import InvalidConfigError
-from data_designer.config.sampler_params import SamplerType, UUIDSamplerParams
+from data_designer.config.sampler_params import (
+    CategorySamplerParams,
+    GaussianSamplerParams,
+    PersonFromFakerSamplerParams,
+    PersonSamplerParams,
+    SamplerType,
+    UniformSamplerParams,
+    UUIDSamplerParams,
+)
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.errors import UserJinjaTemplateSyntaxError
 from data_designer.config.validator_params import CodeValidatorParams
@@ -324,3 +332,114 @@ def test_get_column_config_from_kwargs():
         ),
         SeedDatasetColumnConfig,
     )
+
+
+def test_sampler_column_config_discriminated_union_with_dict_params():
+    """Test that sampler_type field is automatically injected into params dict."""
+    config = SamplerColumnConfig(
+        name="test_uniform",
+        sampler_type=SamplerType.UNIFORM,
+        params={"low": 0.0, "high": 1.0, "decimal_places": 2},
+    )
+    assert config.name == "test_uniform"
+    assert config.sampler_type == SamplerType.UNIFORM
+    assert isinstance(config.params, UniformSamplerParams)
+    assert config.params.sampler_type == SamplerType.UNIFORM
+    assert config.params.low == 0.0
+    assert config.params.high == 1.0
+    assert config.params.decimal_places == 2
+
+
+def test_sampler_column_config_discriminated_union_with_explicit_sampler_type():
+    """Test that explicit sampler_type in params dict is preserved."""
+    config = SamplerColumnConfig(
+        name="test_category",
+        sampler_type=SamplerType.CATEGORY,
+        params={"sampler_type": "category", "values": ["A", "B", "C"], "weights": [0.5, 0.3, 0.2]},
+    )
+    assert config.name == "test_category"
+    assert config.sampler_type == SamplerType.CATEGORY
+    assert isinstance(config.params, CategorySamplerParams)
+    assert config.params.sampler_type == SamplerType.CATEGORY
+    assert config.params.values == ["A", "B", "C"]
+
+
+def test_sampler_column_config_discriminated_union_serialization():
+    """Test that discriminated union works correctly with serialization/deserialization."""
+    config = SamplerColumnConfig(
+        name="test_person",
+        sampler_type=SamplerType.PERSON,
+        params={"locale": "en_US", "sex": "Female", "age_range": [25, 45]},
+    )
+
+    # Serialize
+    serialized = config.model_dump()
+    assert "sampler_type" in serialized["params"]
+    assert serialized["params"]["sampler_type"] == "person"
+
+    # Deserialize
+    deserialized = SamplerColumnConfig(**serialized)
+    assert isinstance(deserialized.params, PersonSamplerParams)
+    assert deserialized.params.locale == "en_US"
+    assert deserialized.params.sex == "Female"
+    assert deserialized.params.age_range == [25, 45]
+
+
+def test_sampler_column_config_discriminated_union_person_vs_person_from_faker():
+    """Test that discriminated union correctly distinguishes between person and person_from_faker."""
+    # Test person sampler (managed datasets)
+    person_config = SamplerColumnConfig(
+        name="test_person",
+        sampler_type=SamplerType.PERSON,
+        params={"locale": "en_US", "sex": "Male", "age_range": [30, 50]},
+    )
+    assert isinstance(person_config.params, PersonSamplerParams)
+    assert person_config.params.sampler_type == SamplerType.PERSON
+    assert person_config.params.locale == "en_US"
+
+    # Test person_from_faker sampler (Faker-based)
+    person_faker_config = SamplerColumnConfig(
+        name="test_person_faker",
+        sampler_type=SamplerType.PERSON_FROM_FAKER,
+        params={"locale": "en_GB", "sex": "Female", "age_range": [20, 40]},
+    )
+    assert isinstance(person_faker_config.params, PersonFromFakerSamplerParams)
+    assert person_faker_config.params.sampler_type == SamplerType.PERSON_FROM_FAKER
+    assert person_faker_config.params.locale == "en_GB"
+
+    # Verify they are different types
+    assert type(person_config.params) is not type(person_faker_config.params)
+    assert isinstance(person_config.params, PersonSamplerParams)
+    assert isinstance(person_faker_config.params, PersonFromFakerSamplerParams)
+
+
+def test_sampler_column_config_discriminated_union_with_conditional_params():
+    """Test that sampler_type is injected into conditional_params as well."""
+    config = SamplerColumnConfig(
+        name="test_gaussian",
+        sampler_type=SamplerType.GAUSSIAN,
+        params={"mean": 0.0, "stddev": 1.0},
+        conditional_params={"age > 21": {"mean": 5.0, "stddev": 2.0}},
+    )
+
+    assert isinstance(config.params, GaussianSamplerParams)
+    assert config.params.mean == 0.0
+    assert config.params.stddev == 1.0
+
+    # Check conditional params
+    assert "age > 21" in config.conditional_params
+    cond_param = config.conditional_params["age > 21"]
+    assert isinstance(cond_param, GaussianSamplerParams)
+    assert cond_param.sampler_type == SamplerType.GAUSSIAN
+    assert cond_param.mean == 5.0
+    assert cond_param.stddev == 2.0
+
+
+def test_sampler_column_config_discriminated_union_wrong_params_type():
+    """Test that discriminated union rejects params that don't match the sampler_type."""
+    with pytest.raises(ValidationError):
+        SamplerColumnConfig(
+            name="test_wrong_params",
+            sampler_type=SamplerType.UNIFORM,
+            params={"values": ["A", "B"]},  # Category params for uniform sampler
+        )

{data_designer-0.1.1 → data_designer-0.1.2}/tests/config/test_datastore.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -13,6 +12,7 @@ import pytest
 from data_designer.config.datastore import (
     DatastoreSettings,
     fetch_seed_dataset_column_names,
+    fetch_seed_dataset_column_names_from_local_file,
     get_file_column_names,
     resolve_datastore_settings,
     upload_to_hf_hub,
@@ -127,22 +127,6 @@ def test_get_file_column_names_unicode(tmp_path, file_type):
     assert get_file_column_names(str(unicode_path), file_type) == df_unicode.columns.tolist()
 
 
-@pytest.mark.parametrize("file_type", ["parquet", "csv", "json", "jsonl"])
-def test_get_file_column_names_with_glob_pattern(tmp_path, file_type):
-    df = pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6]})
-    for i in range(5):
-        _write_file(df, tmp_path / f"{i}.{file_type}", file_type)
-    assert get_file_column_names(f"{tmp_path}/*.{file_type}", file_type) == ["col1", "col2"]
-
-
-def test_get_file_column_names_with_glob_pattern_error(tmp_path):
-    df = pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6]})
-    for i in range(5):
-        _write_file(df, tmp_path / f"{i}.parquet", "parquet")
-    with pytest.raises(InvalidFilePathError, match="No files found matching pattern"):
-        get_file_column_names(f"{tmp_path}/*.csv", "csv")
-
-
 def test_get_file_column_names_with_filesystem_parquet():
     """Test get_file_column_names with filesystem parameter for parquet files."""
     mock_schema = MagicMock()
@@ -153,7 +137,7 @@ def test_get_file_column_names_with_filesystem_parquet():
         result = get_file_column_names("datasets/test/file.parquet", "parquet")
 
         assert result == ["col1", "col2", "col3"]
-        mock_read_schema.assert_called_once_with(Path("datasets/test/file.parquet"))
+        mock_read_schema.assert_called_once_with("datasets/test/file.parquet")
 
 
 @pytest.mark.parametrize("file_type", ["json", "jsonl", "csv"])
@@ -274,3 +258,29 @@ def test_upload_to_hf_hub_error_handling(datastore_settings):
         with patch("data_designer.config.datastore.Path.is_file", autospec=True) as mock_is_file:
             mock_is_file.return_value = True
             upload_to_hf_hub("test.text", "test.txt", "test/repo", datastore_settings)
+
+
+@pytest.mark.parametrize("file_type", ["parquet", "json", "jsonl", "csv"])
+def test_fetch_seed_dataset_column_names_from_local_file_with_glob(tmp_path, file_type):
+    """Test fetch_seed_dataset_column_names_from_local_file with glob pattern matching multiple files."""
+    test_data = pd.DataFrame({"col1": [1, 2], "col2": [3, 4], "col3": [5, 6]})
+
+    # Create multiple files with the same schema
+    for i in range(3):
+        file_path = tmp_path / f"data_{i}.{file_type}"
+        _write_file(test_data, file_path, file_type)
+
+    # Test glob pattern that matches all files
+    glob_pattern = str(tmp_path / f"*.{file_type}")
+    result = fetch_seed_dataset_column_names_from_local_file(glob_pattern)
+
+    assert result == ["col1", "col2", "col3"]
+
+
+@pytest.mark.parametrize("file_type", ["parquet", "csv"])
+def test_fetch_seed_dataset_column_names_from_local_file_with_glob_no_matches(tmp_path, file_type):
+    """Test fetch_seed_dataset_column_names_from_local_file with glob pattern that matches no files."""
+    glob_pattern = str(tmp_path / f"nonexistent_*.{file_type}")
+
+    with pytest.raises(InvalidFilePathError, match="does not contain files of type"):
+        fetch_seed_dataset_column_names_from_local_file(glob_pattern)

{data_designer-0.1.1 → data_designer-0.1.2}/tests/config/test_default_model_settings.py

@@ -11,9 +11,9 @@ import yaml
 from data_designer.config.default_model_settings import (
     get_builtin_model_configs,
     get_builtin_model_providers,
-    get_defaul_model_providers_missing_api_keys,
     get_default_inference_parameters,
     get_default_model_configs,
+    get_default_model_providers_missing_api_keys,
     get_default_provider_name,
     get_default_providers,
     resolve_seed_default_model_settings,
@@ -152,4 +152,4 @@ def test_resolve_seed_default_model_settings(tmp_path: Path):
 @patch("data_designer.config.default_model_settings.os.environ.get")
 def test_get_default_model_providers_missing_api_keys(mock_environ_get):
     mock_environ_get.return_value = None
-    assert get_defaul_model_providers_missing_api_keys() == ["NVIDIA_API_KEY", "OPENAI_API_KEY"]
+    assert get_default_model_providers_missing_api_keys() == ["NVIDIA_API_KEY", "OPENAI_API_KEY"]