data-designer 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

data_designer/engine/processing/processors/schema_transform.py

@@ -1,29 +1,26 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 import logging
-
-import pandas as pd
+from typing import TYPE_CHECKING
 
 from data_designer.config.processors import SchemaTransformProcessorConfig
-from data_designer.engine.configurable_task import ConfigurableTaskMetadata
 from data_designer.engine.dataset_builders.artifact_storage import BatchStage
 from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.utils import deserialize_json_values
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
 
 class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[SchemaTransformProcessorConfig]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="schema_transform_processor",
-            description="Generate dataset with transformed schema using a Jinja2 template.",
-        )
-
     @property
     def template_as_str(self) -> str:
         return json.dumps(self.config.template)

data_designer/engine/processing/utils.py

@@ -1,13 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import ast
 import json
 import logging
 import re
-from typing import Any, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeVar, overload
+
+from data_designer.lazy_heavy_imports import pd
 
-import pandas as pd
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +57,6 @@ def deserialize_json_values(data):
             - Dictionary (potentially with nested JSON strings to deserialize)
             - Some other object that can't be deserialized.
 
-
     Returns:
         Deserialized data in the corresponding format:
             - Dictionary (when input is a single string)

data_designer/engine/registry/base.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import threading
 from typing import Any, Generic, TypeVar
 

data_designer/engine/registry/data_designer_registry.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.analysis.column_profilers.registry import (
     ColumnProfilerRegistry,
     create_default_column_profiler_registry,

data_designer/engine/registry/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.errors import DataDesignerError
 
 

data_designer/engine/resources/managed_dataset_generator.py

@@ -1,11 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any
+from __future__ import annotations
 
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 from data_designer.engine.resources.managed_dataset_repository import ManagedDatasetRepository
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class ManagedDatasetGenerator:

data_designer/engine/resources/managed_dataset_repository.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 import tempfile
 import threading
@@ -9,13 +11,15 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import Any
-
-import duckdb
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
 from data_designer.engine.resources.managed_storage import LocalBlobStorageProvider, ManagedBlobStorage
+from data_designer.lazy_heavy_imports import duckdb, pd
+
+if TYPE_CHECKING:
+    import duckdb
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +56,6 @@ class Table:
 
 DataCatalog = list[Table]
 
-
 # For now we hardcode the remote data catalog in code. This make it easier
 # initialize the data catalog. Eventually we can make this work more
 # dynamically once this data catalog pattern becomes more widely adopted.

data_designer/engine/resources/managed_storage.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Iterator

data_designer/engine/resources/resource_provider.py

@@ -1,14 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.config.base import ConfigBase
+from data_designer.config.dataset_metadata import DatasetMetadata
 from data_designer.config.models import ModelConfig
 from data_designer.config.run_config import RunConfig
 from data_designer.config.seed_source import SeedSource
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.model_provider import ModelProviderRegistry
-from data_designer.engine.models.registry import ModelRegistry, create_model_registry
+from data_designer.engine.models.factory import create_model_registry
+from data_designer.engine.models.registry import ModelRegistry
 from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
 from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry
 from data_designer.engine.secret_resolver import SecretResolver
@@ -27,6 +31,17 @@ class ResourceProvider(ConfigBase):
     run_config: RunConfig = RunConfig()
     seed_reader: SeedReader | None = None
 
+    def get_dataset_metadata(self) -> DatasetMetadata:
+        """Get metadata about the dataset being generated.
+
+        Returns:
+            DatasetMetadata with seed column names and other metadata.
+        """
+        seed_column_names = []
+        if self.seed_reader is not None:
+            seed_column_names = self.seed_reader.get_column_names()
+        return DatasetMetadata(seed_column_names=seed_column_names)
+
 
 def create_resource_provider(
     *,
@@ -39,12 +54,16 @@ def create_resource_provider(
     seed_dataset_source: SeedSource | None = None,
     run_config: RunConfig | None = None,
 ) -> ResourceProvider:
+    """Factory function for creating a ResourceProvider instance.
+    This function triggers lazy loading of heavy dependencies like litellm.
+    """
     seed_reader = None
     if seed_dataset_source:
         seed_reader = seed_reader_registry.get_reader(
             seed_dataset_source,
             secret_resolver,
         )
+
     return ResourceProvider(
         artifact_storage=artifact_storage,
         model_registry=create_model_registry(

data_designer/engine/resources/seed_reader.py

@@ -1,11 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 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 Sequence
-from typing import Generic, TypeVar, get_args, get_origin
+from typing import TYPE_CHECKING, Generic, TypeVar, get_args, get_origin
 
-import duckdb
 from huggingface_hub import HfFileSystem
 from typing_extensions import Self
 
@@ -17,6 +18,10 @@ from data_designer.config.seed_source import (
 )
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.errors import DataDesignerError
+from data_designer.lazy_heavy_imports import duckdb
+
+if TYPE_CHECKING:
+    import duckdb
 
 
 class SeedReaderError(DataDesignerError): ...

data_designer/engine/sampling_gen/column.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import Any
 
 from pydantic import field_serializer, model_validator

data_designer/engine/sampling_gen/constraints.py

@@ -1,10 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
-import numpy as np
-import pandas as pd
 from numpy.typing import NDArray
 
 from data_designer.config.base import ConfigBase
@@ -15,6 +16,11 @@ from data_designer.config.sampler_constraints import (
     InequalityOperator,
     ScalarInequalityConstraint,
 )
+from data_designer.lazy_heavy_imports import np, pd
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
 
 
 class ConstraintChecker(ConfigBase, ABC):

data_designer/engine/sampling_gen/data_sources/base.py

@@ -1,24 +1,27 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
-import numpy as np
-import pandas as pd
 from numpy.typing import NDArray
-from scipy import stats
 
 from data_designer.config.sampler_params import SamplerParamsT
 from data_designer.engine.sampling_gen.utils import check_random_state
+from data_designer.lazy_heavy_imports import np, pd, scipy
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+    import scipy
 
 NumpyArray1dT = NDArray[Any]
 RadomStateT = int | np.random.RandomState
 
-
 GenericParamsT = TypeVar("GenericParamsT", bound=SamplerParamsT)
 
-
 ###########################################################
 # Processing Mixins
 # -----------------
@@ -208,7 +211,7 @@ class Sampler(DataSource[GenericParamsT], ABC):
 class ScipyStatsSampler(Sampler[GenericParamsT], ABC):
     @property
     @abstractmethod
-    def distribution(self) -> stats.rv_continuous | stats.rv_discrete: ...
+    def distribution(self) -> scipy.stats.rv_continuous | scipy.stats.rv_discrete: ...
 
     def sample(self, num_samples: int) -> NumpyArray1dT:
         return self.distribution.rvs(size=num_samples, random_state=self.rng)

data_designer/engine/sampling_gen/data_sources/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.sampling_gen.errors import SamplingGenError
 
 

data_designer/engine/sampling_gen/data_sources/sources.py

@@ -1,11 +1,10 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import uuid
+from __future__ import annotations
 
-import numpy as np
-import pandas as pd
-from scipy import stats
+import uuid
+from typing import TYPE_CHECKING
 
 from data_designer.config.sampler_params import (
     BernoulliMixtureSamplerParams,
@@ -40,6 +39,12 @@ from data_designer.engine.sampling_gen.data_sources.errors import (
 )
 from data_designer.engine.sampling_gen.entities.dataset_based_person_fields import PERSONA_FIELDS, PII_FIELDS
 from data_designer.engine.sampling_gen.people_gen import PeopleGen
+from data_designer.lazy_heavy_imports import np, pd, scipy
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+    import scipy
 
 ONE_BILLION = 10**9
 
@@ -264,8 +269,8 @@ class ScipySampler(TypeConversionMixin, ScipyStatsSampler[ScipySamplerParams]):
     """Escape hatch sampler to give users access to any scipy.stats distribution."""
 
     @property
-    def distribution(self) -> stats.rv_continuous | stats.rv_discrete:
-        return getattr(stats, self.params.dist_name)(**self.params.dist_params)
+    def distribution(self) -> scipy.stats.rv_continuous | scipy.stats.rv_discrete:
+        return getattr(scipy.stats, self.params.dist_name)(**self.params.dist_params)
 
     def _validate(self) -> None:
         _validate_scipy_distribution(self.params.dist_name, self.params.dist_params)
@@ -274,16 +279,16 @@ class ScipySampler(TypeConversionMixin, ScipyStatsSampler[ScipySamplerParams]):
 @SamplerRegistry.register(SamplerType.BERNOULLI)
 class BernoulliSampler(TypeConversionMixin, ScipyStatsSampler[BernoulliSamplerParams]):
     @property
-    def distribution(self) -> stats.rv_discrete:
-        return stats.bernoulli(p=self.params.p)
+    def distribution(self) -> scipy.stats.rv_discrete:
+        return scipy.stats.bernoulli(p=self.params.p)
 
 
 @SamplerRegistry.register(SamplerType.BERNOULLI_MIXTURE)
 class BernoulliMixtureSampler(TypeConversionMixin, Sampler[BernoulliMixtureSamplerParams]):
     def sample(self, num_samples: int) -> NumpyArray1dT:
-        return stats.bernoulli(p=self.params.p).rvs(size=num_samples) * getattr(stats, self.params.dist_name)(
-            **self.params.dist_params
-        ).rvs(size=num_samples)
+        return scipy.stats.bernoulli(p=self.params.p).rvs(size=num_samples) * getattr(
+            scipy.stats, self.params.dist_name
+        )(**self.params.dist_params).rvs(size=num_samples)
 
     def _validate(self) -> None:
         _validate_scipy_distribution(self.params.dist_name, self.params.dist_params)
@@ -292,29 +297,29 @@ class BernoulliMixtureSampler(TypeConversionMixin, Sampler[BernoulliMixtureSampl
 @SamplerRegistry.register(SamplerType.BINOMIAL)
 class BinomialSampler(TypeConversionMixin, ScipyStatsSampler[BinomialSamplerParams]):
     @property
-    def distribution(self) -> stats.rv_discrete:
-        return stats.binom(n=self.params.n, p=self.params.p)
+    def distribution(self) -> scipy.stats.rv_discrete:
+        return scipy.stats.binom(n=self.params.n, p=self.params.p)
 
 
 @SamplerRegistry.register(SamplerType.GAUSSIAN)
 class GaussianSampler(TypeConversionMixin, ScipyStatsSampler[GaussianSamplerParams]):
     @property
-    def distribution(self) -> stats.rv_continuous:
-        return stats.norm(loc=self.params.mean, scale=self.params.stddev)
+    def distribution(self) -> scipy.stats.rv_continuous:
+        return scipy.stats.norm(loc=self.params.mean, scale=self.params.stddev)
 
 
 @SamplerRegistry.register(SamplerType.POISSON)
 class PoissonSampler(TypeConversionMixin, ScipyStatsSampler[PoissonSamplerParams]):
     @property
-    def distribution(self) -> stats.rv_discrete:
-        return stats.poisson(mu=self.params.mean)
+    def distribution(self) -> scipy.stats.rv_discrete:
+        return scipy.stats.poisson(mu=self.params.mean)
 
 
 @SamplerRegistry.register(SamplerType.UNIFORM)
 class UniformSampler(TypeConversionMixin, ScipyStatsSampler[UniformSamplerParams]):
     @property
-    def distribution(self) -> stats.rv_continuous:
-        return stats.uniform(loc=self.params.low, scale=self.params.high - self.params.low)
+    def distribution(self) -> scipy.stats.rv_continuous:
+        return scipy.stats.uniform(loc=self.params.low, scale=self.params.high - self.params.low)
 
 
 ###################################################
@@ -328,14 +333,14 @@ def load_sampler(sampler_type: SamplerType, **params) -> DataSource:
 
 
 def _validate_scipy_distribution(dist_name: str, dist_params: dict) -> None:
-    if not hasattr(stats, dist_name):
+    if not hasattr(scipy.stats, dist_name):
         raise InvalidSamplerParamsError(f"Distribution {dist_name} not found in scipy.stats")
-    if not hasattr(getattr(stats, dist_name), "rvs"):
+    if not hasattr(getattr(scipy.stats, dist_name), "rvs"):
         raise InvalidSamplerParamsError(
             f"Distribution {dist_name} does not have a `rvs` method, which is required for sampling."
         )
     try:
-        getattr(stats, dist_name)(**dist_params)
+        getattr(scipy.stats, dist_name)(**dist_params)
     except Exception:
         raise InvalidSamplerParamsError(
             f"Distribution parameters {dist_params} are not a valid for distribution '{dist_name}'"

data_designer/engine/sampling_gen/entities/dataset_based_person_fields.py

@@ -10,8 +10,9 @@ This file contains all possible fields that:
 Do not add any other code or logic in this file.
 """
 
-REQUIRED_FIELDS = {"first_name", "last_name", "age", "locale"}
+from __future__ import annotations
 
+REQUIRED_FIELDS = {"first_name", "last_name", "age", "locale"}
 
 PII_FIELDS = [
     # Core demographic fields
@@ -52,7 +53,6 @@ PII_FIELDS = [
     "third_language",
 ]
 
-
 PERSONA_FIELDS = [
     # Core persona fields
     "persona",

data_designer/engine/sampling_gen/entities/email_address_utils.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import random
 import re
 from datetime import date

data_designer/engine/sampling_gen/entities/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.errors import DataDesignerError
 
 

data_designer/engine/sampling_gen/entities/national_id_utils.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import random
 from datetime import date
 

data_designer/engine/sampling_gen/entities/person.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import random
 from datetime import date, timedelta
 from typing import Any, Literal, TypeAlias

data_designer/engine/sampling_gen/entities/phone_number.py

@@ -1,12 +1,19 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import random
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-import pandas as pd
 from pydantic import BaseModel, Field, field_validator
 
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
 ZIP_AREA_CODE_DATA = pd.read_parquet(Path(__file__).parent / "assets" / "zip_area_code_map.parquet")
 ZIPCODE_AREA_CODE_MAP = dict(zip(ZIP_AREA_CODE_DATA["zipcode"], ZIP_AREA_CODE_DATA["area_code"]))
 ZIPCODE_POPULATION_MAP = dict(zip(ZIP_AREA_CODE_DATA["zipcode"], ZIP_AREA_CODE_DATA["count"]))

data_designer/engine/sampling_gen/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.engine.errors import DataDesignerError
 
 

data_designer/engine/sampling_gen/generator.py

@@ -6,18 +6,19 @@ from __future__ import annotations
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
-import networkx as nx
-import numpy as np
-import pandas as pd
-
 from data_designer.engine.sampling_gen.data_sources.base import RadomStateT
 from data_designer.engine.sampling_gen.errors import RejectionSamplingError
 from data_designer.engine.sampling_gen.jinja_utils import JinjaDataFrame
 from data_designer.engine.sampling_gen.people_gen import create_people_gen_resource
 from data_designer.engine.sampling_gen.schema import DataSchema
 from data_designer.engine.sampling_gen.utils import check_random_state
+from data_designer.lazy_heavy_imports import np, nx, pd
 
 if TYPE_CHECKING:
+    import networkx as nx
+    import numpy as np
+    import pandas as pd
+
     from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
     from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
     from data_designer.engine.sampling_gen.column import ConditionalDataColumn

data_designer/engine/sampling_gen/jinja_utils.py

@@ -1,15 +1,19 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import ast
-from typing import Any
+from __future__ import annotations
 
-import pandas as pd
+import ast
+from typing import TYPE_CHECKING, Any
 
 from data_designer.engine.processing.ginja.environment import (
     UserTemplateSandboxEnvironment,
     WithJinja2UserTemplateRendering,
 )
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class JinjaDataFrame(WithJinja2UserTemplateRendering):

data_designer/engine/sampling_gen/people_gen.py

@@ -10,9 +10,6 @@ from collections.abc import Callable
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, TypeAlias
 
-import pandas as pd
-from faker import Faker
-
 from data_designer.config.utils.constants import 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
@@ -22,12 +19,15 @@ from data_designer.engine.sampling_gen.entities.person import (
 )
 from data_designer.engine.sampling_gen.errors import ManagedDatasetGeneratorError
 from data_designer.engine.sampling_gen.person_constants import faker_constants
+from data_designer.lazy_heavy_imports import faker, pd
 
 if TYPE_CHECKING:
-    from data_designer.engine.sampling_gen.schema import DataSchema
+    import faker
+    import pandas as pd
 
+    from data_designer.engine.sampling_gen.schema import DataSchema
 
-EngineT: TypeAlias = Faker | ManagedDatasetGenerator
+EngineT: TypeAlias = faker.Faker | ManagedDatasetGenerator
 
 
 class PeopleGen(ABC):
@@ -46,7 +46,7 @@ class PeopleGen(ABC):
 
 class PeopleGenFaker(PeopleGen):
     @property
-    def _fake(self) -> Faker:
+    def _fake(self) -> faker.Faker:
         return self._engine
 
     def try_fake_else_none(self, attr_name: str, none_fill: Any | None = None) -> type:
@@ -193,7 +193,7 @@ def create_people_gen_resource(
         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
+                    engine=faker.Faker(params.locale), locale=params.locale
                 )
 
     return people_gen_resource

data_designer/engine/sampling_gen/person_constants.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import NamedTuple
 
 

data_designer/engine/sampling_gen/schema.py

@@ -4,8 +4,8 @@
 from __future__ import annotations
 
 from functools import cached_property
+from typing import TYPE_CHECKING
 
-import networkx as nx
 from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self
 
@@ -14,6 +14,10 @@ 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
+from data_designer.lazy_heavy_imports import nx
+
+if TYPE_CHECKING:
+    import networkx as nx
 
 
 class Dag(BaseModel):