data-designer 0.1.0__py3-none-any.whl

data_designer/engine/sampling_gen/jinja_utils.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import ast
+from typing import Any
+
+import pandas as pd
+
+from data_designer.engine.processing.ginja.environment import (
+    UserTemplateSandboxEnvironment,
+    WithJinja2UserTemplateRendering,
+)
+
+
+class JinjaDataFrame(WithJinja2UserTemplateRendering):
+    def __init__(self, expr: str):
+        self.expr = expr
+
+    def _jsonify(self, record) -> dict[str, Any]:
+        for key, value in record.items():
+            if isinstance(value, pd.Timestamp):
+                record[key] = value.isoformat()
+        return record
+
+    @property
+    def _expr(self) -> str:
+        return "{{ " + self.expr + " }}"
+
+    def select_index(self, dataframe: pd.DataFrame) -> pd.Index:
+        if dataframe.empty or self.expr == "...":
+            return dataframe.index
+
+        self.prepare_jinja2_template_renderer(self._expr, list(dataframe))
+
+        where = dataframe.apply(
+            lambda row: self.render_template(self._jsonify(row.to_dict())) == "True",
+            axis=1,
+        ).to_numpy()
+
+        return dataframe[where].index
+
+    def to_column(self, dataframe: pd.DataFrame) -> list[Any]:
+        self.prepare_jinja2_template_renderer(self._expr, list(dataframe))
+
+        expr_values = []
+        for record in dataframe.to_dict(orient="records"):
+            rendered = self.render_template(self._jsonify(record))
+            try:
+                # Non-string expressions are evaluated as literals.
+                expr_values.append(ast.literal_eval(rendered))
+            except (ValueError, SyntaxError):
+                # Strings throw an error and are appended directly.
+                expr_values.append(rendered)
+
+        return expr_values
+
+
+def extract_column_names_from_expression(expr: str) -> set[str]:
+    """Extract valid column names from the given expression."""
+    return UserTemplateSandboxEnvironment().get_references("{{ " + expr + " }}")

data_designer/engine/sampling_gen/people_gen.py

@@ -0,0 +1,203 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from copy import deepcopy
+import random
+from typing import TYPE_CHECKING, Any, Union
+import uuid
+
+from faker import Faker
+import pandas as pd
+
+from data_designer.config.utils.constants import AVAILABLE_LOCALES, DEFAULT_AGE_RANGE
+from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
+from data_designer.engine.sampling_gen.entities.dataset_based_person_fields import PERSONA_FIELDS, PII_FIELDS
+from data_designer.engine.sampling_gen.entities.person import (
+    convert_age_to_birth_date,
+    generate_and_insert_derived_fields,
+)
+from data_designer.engine.sampling_gen.errors import ManagedDatasetGeneratorError
+from data_designer.engine.sampling_gen.person_constants import faker_constants
+
+if TYPE_CHECKING:
+    from data_designer.engine.sampling_gen.schema import DataSchema
+
+
+EngineT = Union[Faker, ManagedDatasetGenerator]
+
+
+class PeopleGen(ABC):
+    """Unified interface for generating people data."""
+
+    def __init__(self, engine: EngineT, locale: str):
+        if locale not in AVAILABLE_LOCALES:
+            raise ValueError(
+                f"Locale {locale} is not a supported locale.Supported locales: {', '.join(AVAILABLE_LOCALES)}"
+            )
+        self.locale = locale
+        self._engine = engine
+
+    def set_engine(self, engine: EngineT) -> None:
+        self._engine = engine
+
+    @abstractmethod
+    def generate(self, n: int, **kwargs) -> list[dict[str, Any]]: ...
+
+
+class PeopleGenFaker(PeopleGen):
+    @property
+    def _fake(self) -> Faker:
+        return self._engine
+
+    def try_fake_else_none(self, attr_name: str, none_fill: Any | None = None) -> type:
+        try:
+            return getattr(self._fake, attr_name)()
+        except AttributeError:
+            return none_fill
+
+    def _generate_name_and_sex(self, **kwargs) -> dict[str, str]:
+        options = faker_constants.sex
+        if "sex" in kwargs and kwargs["sex"] in [*options, *[[o] for o in options]]:
+            sex = random.choice(kwargs["sex"]) if isinstance(kwargs["sex"], list) else kwargs["sex"]
+        else:
+            sex = random.choice(options)
+
+        return {
+            "first_name": getattr(self._fake, f"first_name_{sex.lower()}")(),
+            "last_name": getattr(self._fake, f"last_name_{sex.lower()}")(),
+            "middle_name": None,
+            "sex": sex,
+        }
+
+    def _generate_address_fields(self, **kwargs) -> dict[str, str]:
+        address = {
+            "street_number": self.try_fake_else_none(faker_constants.attr_map["street_number"]),
+            "street_name": self.try_fake_else_none("street_name"),
+        }
+
+        # Location fields can be filtered using the fixed_kwargs.
+        for attr in faker_constants.location:
+            if attr in kwargs:
+                address[attr] = random.choice(kwargs[attr]) if isinstance(kwargs[attr], list) else kwargs[attr]
+            else:
+                address[attr] = self.try_fake_else_none(attr)
+
+        return address
+
+    def _generate_age(self, **kwargs) -> int:
+        return random.randint(*kwargs.get("age_range", DEFAULT_AGE_RANGE))
+
+    def _generate_marital_status(self, **kwargs) -> str:
+        return random.choice(faker_constants.marital_status)
+
+    def _generate_bachelors_field(self, **kwargs) -> str:
+        return random.choice(faker_constants.bachelors)
+
+    def _generate_education_level(self, **kwargs) -> str:
+        return random.choice(faker_constants.education_level)
+
+    def make_person(self, **kwargs) -> dict[str, Any]:
+        person = {"uuid": str(uuid.uuid4()), "locale": self.locale}
+        person.update(self._generate_name_and_sex(**kwargs))
+        person.update(self._generate_address_fields(**kwargs))
+        person.update({"age": self._generate_age(**kwargs)})
+        person.update({"birth_date": convert_age_to_birth_date(person["age"]).isoformat()})
+        person.update({"country": self.try_fake_else_none("country")})
+        person.update({"marital_status": self._generate_marital_status(**kwargs)})
+        person.update({"education_level": self._generate_education_level(**kwargs)})
+        person.update({"unit": ""})
+        person.update({"occupation": self.try_fake_else_none(faker_constants.attr_map["occupation"])})
+        person.update({"phone_number": (self.try_fake_else_none("phone_number") if person["age"] >= 18 else None)})
+        if person["education_level"] in faker_constants.college_level:
+            person.update({"bachelors_field": self._generate_bachelors_field(**kwargs)})
+        else:
+            person.update({"bachelors_field": "no_degree"})
+        return person
+
+    def generate(self, n: int, **kwargs) -> list[dict[str, Any]]:
+        return [self.make_person(**kwargs) for _ in range(n)]
+
+
+class PeopleGenFromDataset(PeopleGen):
+    def _get_ages(self, age_range: tuple[int, int]) -> list[int]:
+        return list(range(age_range[0], age_range[1] + 1))
+
+    def _generate_from_dataset(self, n: int, **kwargs) -> pd.DataFrame:
+        kw = deepcopy(kwargs)
+        with_synthetic_personas = kw.pop("with_synthetic_personas", False)
+        kw["age"] = self._get_ages(kw.pop("age_range", DEFAULT_AGE_RANGE))
+
+        # Generate samples and drop columns where all rows are null.
+        df = self._engine.generate_samples(size=n, evidence=kw).dropna(axis=1, how="all")
+
+        # We need this for derived fields.
+        df["locale"] = self.locale
+
+        # Only keep columns that are listed in the schema.
+        fields = [field for field in PII_FIELDS if field in df.columns]
+        if with_synthetic_personas:
+            fields.extend([field for field in PERSONA_FIELDS if field in df.columns])
+
+        return df[fields]
+
+    def generate(self, n: int, **kwargs) -> list[dict[str, Any]]:
+        return [
+            generate_and_insert_derived_fields(p)
+            for p in self._generate_from_dataset(n, **kwargs).to_dict(orient="records")
+        ]
+
+
+def create_people_gen_resource(
+    schema: DataSchema,
+    person_generator_loader: Callable[[bool], ManagedDatasetGenerator] | None = None,
+) -> dict[str, PeopleGen]:
+    """Creates resource of unique people generators needed to generate the dataset.
+
+    The resource is a dictionary of person generators, where the keys are the following:
+        - {locale} for dataset-based person generators
+        - {locale}_with_personas for dataset-based person generators with synthetic personas
+        - {locale}_faker for faker-based person generators
+
+    Args:
+        schema: Schema of the dataset that we will generate.
+        person_generator_loader: Function that loads a managed dataset generator.
+
+    Returns:
+        Dictionary of unique people generators needed to generate the dataset.
+    """
+    people_gen_resource = {}
+
+    # ------------------------------------------------------------
+    # Preload dataset-based person generators
+    # ------------------------------------------------------------
+
+    for column in schema.get_columns_by_sampler_type("person"):
+        for params in [column.params, *list(column.conditional_params.values())]:
+            if params.people_gen_key not in people_gen_resource:
+                try:
+                    engine = person_generator_loader(locale=params.locale)
+                    people_gen_resource[params.people_gen_key] = PeopleGenFromDataset(
+                        engine=engine, locale=params.locale
+                    )
+                except Exception as e:
+                    raise ManagedDatasetGeneratorError(
+                        f"🛑 Failed to load dataset-based person generator for locale {params.locale}. "
+                        "Please check if you have access to person data for this locale. "
+                    ) from e
+
+    # ------------------------------------------------------------
+    # Preload faker-based person generators
+    # ------------------------------------------------------------
+
+    for column in schema.get_columns_by_sampler_type("person_from_faker"):
+        for params in [column.params, *list(column.conditional_params.values())]:
+            if params.people_gen_key not in people_gen_resource:
+                people_gen_resource[params.people_gen_key] = PeopleGenFaker(
+                    engine=Faker(params.locale), locale=params.locale
+                )
+
+    return people_gen_resource

data_designer/engine/sampling_gen/person_constants.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import NamedTuple
+
+
+class FakerPersonData(NamedTuple):
+    sex: list[str] = ["Male", "Female"]
+
+    us_locale_only: list[str] = [
+        "state",
+        "county",
+        "unit",
+        "middle_name",
+        "ethnic_background",
+        "ssn",
+    ]
+
+    location: list[str] = ["city", "state", "postcode"]
+
+    bachelors: list[str] = [
+        "stem",
+        "business",
+        "education",
+        "arts_humanities",
+        "stem_related",
+    ]
+
+    education_level: list[str] = [
+        "secondary_education",
+        "some_college",
+        "bachelors",
+        "associates",
+        "graduate",
+        "doctorate",
+    ]
+
+    marital_status: list[str] = [
+        "married_present",
+        "divorced",
+        "never_married",
+        "separated",
+        "widowed",
+    ]
+
+    college_level: list[str] = ["bachelors", "graduate", "doctorate"]
+
+    attr_map: dict[str, str] = {
+        "street_number": "building_number",
+        "occupation": "job",
+    }
+
+
+faker_constants = FakerPersonData()

data_designer/engine/sampling_gen/schema.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from functools import cached_property
+
+import networkx as nx
+from pydantic import BaseModel, Field, field_validator, model_validator
+from typing_extensions import Self
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.sampler_constraints import ColumnConstraintT
+from data_designer.config.sampler_params import SamplerType
+from data_designer.engine.sampling_gen.column import ConditionalDataColumn
+from data_designer.engine.sampling_gen.constraints import ConstraintChecker, get_constraint_checker
+
+
+class Dag(BaseModel):
+    nodes: set[str]
+    edges: set[tuple[str, str]]
+
+    @model_validator(mode="after")
+    def validate_is_dag(self) -> Self:
+        if not nx.is_directed_acyclic_graph(self.to_networkx()):
+            raise ValueError("There are circular dependencies in the definitions of your sampler columns.")
+        return self
+
+    def to_networkx(self) -> nx.DiGraph:
+        dag = nx.DiGraph()
+        for node in self.nodes:
+            dag.add_node(node)
+        for edge in self.edges:
+            dag.add_edge(*edge)
+        return dag
+
+
+class DataSchema(ConfigBase):
+    """Defines the data schema for synthetic data generation.
+
+    A DataSchema represents a collection of columns and their relationships through
+    conditional parameters and/or constraints. Upon initialization, the schema validates
+    that column dependencies form a DAG and that all constraints reference valid columns.
+    """
+
+    columns: list[ConditionalDataColumn] = Field(..., min_length=1)
+    constraints: list[ColumnConstraintT] = []
+
+    @cached_property
+    def constraint_checkers(self) -> list[ConstraintChecker]:
+        return [get_constraint_checker(c.constraint_type)(constraint=c) for c in self.constraints]
+
+    @property
+    def column_names(self) -> list[str]:
+        return [column.name for column in self.columns]
+
+    @property
+    def dag(self) -> Dag:
+        nodes = set()
+        edges = set()
+
+        for column in self.columns:
+            nodes.add(column.name)
+
+            # Add edges for the conditional columns.
+            for conditional_column in column.conditional_column_names:
+                edges.add((conditional_column, column.name))
+
+            # Add edges if the source has required columns.
+            for condition in column.conditions:
+                source = column.get_sampler(condition)
+                for required_column in source.get_required_column_names():
+                    edges.add((required_column, column.name))
+
+        for checker in self.constraint_checkers:
+            column_names = checker.get_required_column_names()
+            if len(column_names) == 2:
+                edges.add((column_names[1], column_names[0]))
+        return Dag(nodes=nodes, edges=edges)
+
+    @field_validator("columns", mode="after")
+    def check_unique_column_names(cls, columns: list[ConditionalDataColumn]) -> list[ConditionalDataColumn]:
+        column_names = [column.name for column in columns]
+        if len(column_names) != len(set(column_names)):
+            raise ValueError("Column names must be unique")
+        return columns
+
+    @model_validator(mode="after")
+    def validate_constraints(self) -> Self:
+        column_names = [column.name for column in self.columns]
+
+        # Check if all columns required by constraints are present in the schema.
+        for checker in self.constraint_checkers:
+            constrained_column_names = checker.get_required_column_names()
+            if not set(constrained_column_names).issubset(column_names):
+                missing = set(constrained_column_names) - set(column_names)
+                raise ValueError(
+                    f"These constrained columns are missing in the definitions of your sampler columns: {missing}"
+                )
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_dag(self) -> Self:
+        self.dag
+        return self
+
+    @model_validator(mode="after")
+    def validate_subcategory_columns_if_present(self) -> Self:
+        for sub in self.get_columns_by_sampler_type(SamplerType.SUBCATEGORY):
+            cat = self.get_column(sub.params.category)
+            if cat.sampler_type != SamplerType.CATEGORY:
+                raise ValueError(
+                    f"The parent of subcategory column '{sub.name}' must be a category "
+                    f"source type, but '{cat.name}' is of type '{cat.sampler_type}'."
+                )
+            cat_vals = set(cat.params.values)
+            for params in cat.conditional_params.values():
+                cat_vals.update(params.values)
+            sub_vals = set(sub.params.values.keys())
+            if cat_vals.symmetric_difference(sub_vals):
+                raise ValueError(
+                    f"Subcategory column '{sub.name}' must have values for each value of "
+                    f"its parent category '{sub.params.category}'. The following "
+                    f"values need attention: {cat_vals.symmetric_difference(sub_vals)}"
+                )
+            if not all(len(v) > 0 for v in sub.params.values.values()):
+                raise ValueError(
+                    f"Subcategory column '{sub.name}' must have non-empty values for "
+                    f"each value of its parent category '{sub.params.category}'."
+                )
+        return self
+
+    def get_column(self, column_name: str) -> ConditionalDataColumn:
+        if column_name not in self.column_names:
+            raise ValueError(f"Column '{column_name}' not found in schema")
+        return next(column for column in self.columns if column.name == column_name)
+
+    def get_columns_by_sampler_type(self, sampler_type: SamplerType) -> list[ConditionalDataColumn]:
+        return [c for c in self.columns if c.sampler_type == sampler_type]
+
+    def get_constraint_checkers(self, column_name: str) -> list[ConstraintChecker]:
+        return [c for c in self.constraint_checkers if column_name == c.constraint.target_column]

data_designer/engine/sampling_gen/schema_builder.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from copy import deepcopy
+
+from data_designer.config.column_configs import SamplerColumnConfig
+from data_designer.config.sampler_constraints import ColumnConstraintT
+from data_designer.config.sampler_params import SamplerParamsT
+from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
+from data_designer.engine.sampling_gen.column import ConditionalDataColumn
+from data_designer.engine.sampling_gen.schema import DataSchema
+
+
+class SchemaBuilder:
+    """Builder class for DataSchema objects.
+
+    This class is meant to be a helper for internal usage and experimentation. It
+    provides a simple interface for constructing a DataSchema object via `add_column`
+    and `add_constraint` methods similar.
+    """
+
+    def __init__(
+        self,
+        columns: list[ConditionalDataColumn] | None = None,
+        constraints: list[ColumnConstraintT] | None = None,
+    ):
+        self._columns = columns or []
+        self._constraints = constraints or []
+
+    def add_column(
+        self,
+        name: str,
+        sampler_type: str | None,
+        params: dict | SamplerParamsT | None,
+        conditional_params: dict[str, SamplerParamsT] | None = None,
+        convert_to: str | None = None,
+    ) -> None:
+        self._columns.append(
+            ConditionalDataColumn(
+                name=name,
+                sampler_type=sampler_type,
+                params=params,
+                conditional_params=conditional_params or {},
+                convert_to=convert_to,
+            )
+        )
+
+    def add_constraint(self, constraint: ColumnConstraintT) -> None:
+        self._constraints.append(constraint)
+
+    def to_sampler_columns(self, max_rejections_factor: int = 5) -> SamplerMultiColumnConfig:
+        return SamplerMultiColumnConfig(
+            columns=[SamplerColumnConfig(**c.model_dump(mode="json")) for c in self._columns],
+            constraints=self._constraints,
+            max_rejections_factor=max_rejections_factor,
+        )
+
+    def build(self) -> DataSchema:
+        return DataSchema(columns=deepcopy(self._columns), constraints=deepcopy(self._constraints))

data_designer/engine/sampling_gen/utils.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import numbers
+
+import numpy as np
+
+
+def check_random_state(seed):
+    """Turn seed into a np.random.RandomState instance.
+
+    This function was taken from scikit-learn's utils module.
+    Source GitHub: https://github.com/scikit-learn/scikit-learn
+
+    Parameters
+    ----------
+    seed : None, int or instance of RandomState
+        If seed is None, return the RandomState singleton used by np.random.
+        If seed is an int, return a new RandomState instance seeded with seed.
+        If seed is already a RandomState instance, return it.
+        Otherwise raise ValueError.
+
+    Returns
+    -------
+    :class:`numpy:numpy.random.RandomState`
+        The random state object based on `seed` parameter.
+
+    Examples
+    --------
+    >>> from data_designer.engine.sampling_gen.utils import check_random_state
+    >>> check_random_state(42)
+    RandomState(MT19937) at 0x...
+    """
+    if seed is None or seed is np.random:
+        return np.random.mtrand._rand
+    if isinstance(seed, numbers.Integral):
+        return np.random.RandomState(seed)
+    if isinstance(seed, np.random.RandomState):
+        return seed
+    raise ValueError("%r cannot be used to seed a numpy.random.RandomState instance" % seed)

data_designer/engine/secret_resolver.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+import json
+import logging
+import os
+from pathlib import Path
+from typing import Protocol
+
+from data_designer.engine.errors import SecretResolutionError
+
+logger = logging.getLogger(__name__)
+
+
+class SecretResolver(Protocol):
+    def resolve(self, secret: str) -> str: ...
+
+
+class SecretsFileResolver(SecretResolver):
+    _secrets: dict[str, str]
+
+    def __init__(self, filepath: Path):
+        if not filepath.exists():
+            self._secrets = {}
+        else:
+            with open(filepath) as f:
+                self._secrets = json.load(f)
+
+    def resolve(self, secret: str) -> str:
+        try:
+            return self._secrets[secret]
+        except KeyError:
+            raise SecretResolutionError(f"No secret found in secrets file with key {secret!r}")
+
+
+class EnvironmentResolver(SecretResolver):
+    def resolve(self, secret: str) -> str:
+        try:
+            return os.environ[secret]
+        except KeyError:
+            raise SecretResolutionError(
+                f"Environment variable with name {secret!r} is required but not set. Please set it in your environment and try again."
+            )
+
+
+class PlaintextResolver(SecretResolver):
+    def resolve(self, secret: str) -> str:
+        return secret
+
+
+class CompositeResolver(SecretResolver):
+    _resolvers: Sequence[SecretResolver]
+
+    def __init__(self, resolvers: Sequence[SecretResolver]):
+        if len(resolvers) == 0:
+            raise SecretResolutionError("Must provide at least one SecretResolver to CompositeResolver")
+        self._resolvers = resolvers
+
+    @property
+    def resolvers(self) -> Sequence[SecretResolver]:
+        """Get the sequence of resolvers in this composite resolver.
+
+        Returns:
+            Sequence of SecretResolver instances used to resolve secrets.
+        """
+        return self._resolvers
+
+    def resolve(self, secret: str) -> str:
+        errors = []
+        for resolver in self._resolvers:
+            try:
+                return resolver.resolve(secret)
+            except SecretResolutionError as err:
+                errors.append(str(err))
+                continue
+
+        raise SecretResolutionError(
+            f"No configured resolvers were able to resolve secret {secret!r}: {', '.join(errors)}"
+        )

data_designer/engine/validators/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.engine.validators.base import BaseValidator, ValidationResult
+from data_designer.engine.validators.local_callable import LocalCallableValidator
+from data_designer.engine.validators.python import PythonValidator
+from data_designer.engine.validators.remote import RemoteValidator
+from data_designer.engine.validators.sql import SQLValidator
+
+__all__ = [
+    "BaseValidator",
+    "LocalCallableValidator",
+    "RemoteValidator",
+    "ValidationResult",
+    "PythonValidator",
+    "SQLValidator",
+]