data-designer 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

data_designer/config/column_types.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
 
 from typing_extensions import TypeAlias
 
@@ -15,7 +16,7 @@ from data_designer.config.column_configs import (
     SeedDatasetColumnConfig,
     ValidationColumnConfig,
 )
-from data_designer.config.errors import InvalidColumnTypeError, InvalidConfigError
+from data_designer.config.errors import InvalidConfigError
 from data_designer.config.sampler_params import SamplerType
 from data_designer.config.utils.type_helpers import (
     SAMPLER_PARAMS,
@@ -45,22 +46,6 @@ DataDesignerColumnType = create_str_enum_from_discriminated_type_union(
     discriminator_field_name="column_type",
 )
 
-COLUMN_TYPE_EMOJI_MAP = {
-    "general": "⚛️",  # possible analysis column type
-    DataDesignerColumnType.EXPRESSION: "🧩",
-    DataDesignerColumnType.LLM_CODE: "💻",
-    DataDesignerColumnType.LLM_JUDGE: "⚖️",
-    DataDesignerColumnType.LLM_STRUCTURED: "🗂️",
-    DataDesignerColumnType.LLM_TEXT: "📝",
-    DataDesignerColumnType.SEED_DATASET: "🌱",
-    DataDesignerColumnType.SAMPLER: "🎲",
-    DataDesignerColumnType.VALIDATION: "🔍",
-    DataDesignerColumnType.EMBEDDING: "🧬",
-}
-COLUMN_TYPE_EMOJI_MAP.update(
-    {DataDesignerColumnType(p.name): p.emoji for p in plugin_manager.get_column_generator_plugins()}
-)
-
 
 def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
     """Create a Data Designer column config object from kwargs.
@@ -74,27 +59,20 @@ def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType
         Data Designer column object of the appropriate type.
     """
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    if column_type == DataDesignerColumnType.LLM_TEXT:
-        return LLMTextColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_CODE:
-        return LLMCodeColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_STRUCTURED:
-        return LLMStructuredColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.LLM_JUDGE:
-        return LLMJudgeColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.VALIDATION:
-        return ValidationColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.EXPRESSION:
-        return ExpressionColumnConfig(name=name, **kwargs)
+    config_cls = get_column_config_cls_from_type(column_type)
     if column_type == DataDesignerColumnType.SAMPLER:
-        return SamplerColumnConfig(name=name, **_resolve_sampler_kwargs(name, kwargs))
-    if column_type == DataDesignerColumnType.SEED_DATASET:
-        return SeedDatasetColumnConfig(name=name, **kwargs)
-    if column_type == DataDesignerColumnType.EMBEDDING:
-        return EmbeddingColumnConfig(name=name, **kwargs)
+        kwargs = _resolve_sampler_kwargs(name, kwargs)
+    return config_cls(name=name, **kwargs)
+
+
+def get_column_config_cls_from_type(column_type: DataDesignerColumnType) -> type[ColumnConfigT]:
+    """Get the column config class for a column type."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    if column_type in _COLUMN_TYPE_CONFIG_CLS_MAP:
+        return _COLUMN_TYPE_CONFIG_CLS_MAP[column_type]
     if plugin := plugin_manager.get_column_generator_plugin_if_exists(column_type.value):
-        return plugin.config_cls(name=name, **kwargs)
-    raise InvalidColumnTypeError(f"🛑 {column_type} is not a valid column type.")  # pragma: no cover
+        return plugin.config_cls
+    raise InvalidConfigError(f"🛑 {column_type} is not a valid column type.")
 
 
 def get_column_display_order() -> list[DataDesignerColumnType]:
@@ -114,6 +92,12 @@ def get_column_display_order() -> list[DataDesignerColumnType]:
     return display_order
 
 
+def get_column_emoji_from_type(column_type: DataDesignerColumnType) -> str:
+    """Get the emoji for a column type."""
+    config_cls = get_column_config_cls_from_type(resolve_string_enum(column_type, DataDesignerColumnType))
+    return config_cls.get_column_emoji()
+
+
 def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
     if "sampler_type" not in kwargs:
         raise InvalidConfigError(f"🛑 `sampler_type` is required for sampler column '{name}'.")
@@ -142,3 +126,16 @@ def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
         "params": params,
         **{k: v for k, v in kwargs.items() if k not in ["sampler_type", "params"]},
     }
+
+
+_COLUMN_TYPE_CONFIG_CLS_MAP = {
+    DataDesignerColumnType.LLM_TEXT: LLMTextColumnConfig,
+    DataDesignerColumnType.LLM_CODE: LLMCodeColumnConfig,
+    DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnConfig,
+    DataDesignerColumnType.LLM_JUDGE: LLMJudgeColumnConfig,
+    DataDesignerColumnType.VALIDATION: ValidationColumnConfig,
+    DataDesignerColumnType.EXPRESSION: ExpressionColumnConfig,
+    DataDesignerColumnType.SAMPLER: SamplerColumnConfig,
+    DataDesignerColumnType.SEED_DATASET: SeedDatasetColumnConfig,
+    DataDesignerColumnType.EMBEDDING: EmbeddingColumnConfig,
+}

data_designer/config/dataset_builders.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 enum import Enum
 
 

data_designer/config/default_model_settings.py

@@ -1,6 +1,7 @@
 # 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 os

data_designer/config/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/config/exports.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.config.analysis.column_profilers import JudgeScoreProfilerConfig
 from data_designer.config.column_configs import (
     EmbeddingColumnConfig,

data_designer/config/interface.py

@@ -6,13 +6,14 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Generic, Protocol, TypeVar
 
-import pandas as pd
-
 from data_designer.config.models import ModelConfig, ModelProvider
 from data_designer.config.utils.constants import DEFAULT_NUM_RECORDS
 from data_designer.config.utils.info import InterfaceInfo
+from data_designer.lazy_heavy_imports import pd
 
 if TYPE_CHECKING:
+    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.preview_results import PreviewResults

data_designer/config/models.py

@@ -1,13 +1,14 @@
 # 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 enum import Enum
 from pathlib import Path
-from typing import Annotated, Any, Generic, Literal, TypeVar
+from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal, TypeVar
 
-import numpy as np
 from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self, TypeAlias
 
@@ -20,6 +21,10 @@ from data_designer.config.utils.constants import (
     MIN_TOP_P,
 )
 from data_designer.config.utils.io_helpers import smart_load_yaml
+from data_designer.lazy_heavy_imports import np
+
+if TYPE_CHECKING:
+    import numpy as np
 
 logger = logging.getLogger(__name__)
 

data_designer/config/preview_results.py

@@ -3,12 +3,16 @@
 
 from __future__ import annotations
 
-import pandas as pd
+from typing import TYPE_CHECKING
 
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.dataset_metadata import DatasetMetadata
 from data_designer.config.utils.visualization import WithRecordSamplerMixin
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class PreviewResults(WithRecordSamplerMixin):
@@ -16,7 +20,7 @@ class PreviewResults(WithRecordSamplerMixin):
         self,
         *,
         config_builder: DataDesignerConfigBuilder,
-        dataset_metadata: DatasetMetadata,
+        dataset_metadata: DatasetMetadata | None = None,
         dataset: pd.DataFrame | None = None,
         analysis: DatasetProfilerResults | None = None,
         processor_artifacts: dict[str, list[str] | str] | None = None,
@@ -33,5 +37,5 @@ class PreviewResults(WithRecordSamplerMixin):
         self.dataset: pd.DataFrame | None = dataset
         self.analysis: DatasetProfilerResults | None = analysis
         self.processor_artifacts: dict[str, list[str] | str] | None = processor_artifacts
-        self.dataset_metadata = dataset_metadata
+        self.dataset_metadata: DatasetMetadata | None = dataset_metadata
         self._config_builder = config_builder

data_designer/config/processors.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 json
 from abc import ABC
 from enum import Enum

data_designer/config/run_config.py

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

data_designer/config/sampler_constraints.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 abc import ABC, abstractmethod
 from enum import Enum
 

data_designer/config/sampler_params.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 enum import Enum
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
-import pandas as pd
 from pydantic import Field, field_validator, model_validator
 from typing_extensions import Self, TypeAlias
 
@@ -16,6 +17,10 @@ from data_designer.config.utils.constants import (
     MAX_AGE,
     MIN_AGE,
 )
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class SamplerType(str, Enum):

data_designer/config/seed.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 enum import Enum
 
 from pydantic import Field, model_validator

data_designer/config/seed_source.py

@@ -1,10 +1,11 @@
 # 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
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
-import pandas as pd
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from pydantic.json_schema import SkipJsonSchema
 from typing_extensions import Self
@@ -14,6 +15,10 @@ from data_designer.config.utils.io_helpers import (
     validate_dataset_file_path,
     validate_path_contains_files_of_type,
 )
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class SeedSource(BaseModel, ABC):

data_designer/config/seed_source_types.py

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

data_designer/config/utils/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
+
 import os
 from enum import Enum
 from pathlib import Path

data_designer/config/utils/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/config/utils/info.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 abc import ABC, abstractmethod
 from enum import Enum
 from typing import Literal, TypeVar

data_designer/config/utils/io_helpers.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 json
 import logging
 import os
@@ -8,13 +10,16 @@ from datetime import date, datetime, timedelta
 from decimal import Decimal
 from numbers import Number
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-import numpy as np
-import pandas as pd
 import yaml
 
 from data_designer.config.errors import InvalidFileFormatError, InvalidFilePathError
+from data_designer.lazy_heavy_imports import np, pd
+
+if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 

data_designer/config/utils/misc.py

@@ -48,8 +48,8 @@ def can_run_data_designer_locally() -> bool:
     return True
 
 
-def get_prompt_template_keywords(template: str) -> set[str]:
-    """Extract all keywords from a valid string template."""
+def extract_keywords_from_jinja2_template(template: str) -> set[str]:
+    """Extract all keywords from a valid Jinja2 template."""
     with template_error_handler():
         ast = ImmutableSandboxedEnvironment().parse(template)
         keywords = set(meta.find_undeclared_variables(ast))

data_designer/config/utils/numerical_helpers.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 numbers
 from numbers import Number
 from typing import Any

data_designer/config/utils/type_helpers.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 inspect
 from enum import Enum
 from typing import Any, Literal, get_args, get_origin

data_designer/config/utils/visualization.py

@@ -10,8 +10,6 @@ from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-import numpy as np
-import pandas as pd
 from rich.console import Console, Group
 from rich.padding import Padding
 from rich.panel import Panel
@@ -28,8 +26,12 @@ from data_designer.config.sampler_params import SamplerType
 from data_designer.config.utils.code_lang import code_lang_to_syntax_lexer
 from data_designer.config.utils.constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
 from data_designer.config.utils.errors import DatasetSampleDisplayError
+from data_designer.lazy_heavy_imports import np, pd
 
 if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
     from data_designer.config.config_builder import DataDesignerConfigBuilder
     from data_designer.config.dataset_metadata import DatasetMetadata
 
@@ -58,7 +60,7 @@ class ColorPalette(str, Enum):
 
 class WithRecordSamplerMixin:
     _display_cycle_index: int = 0
-    dataset_metadata: DatasetMetadata
+    dataset_metadata: DatasetMetadata | None
 
     @cached_property
     def _record_sampler_dataset(self) -> pd.DataFrame:
@@ -122,7 +124,9 @@ class WithRecordSamplerMixin:
                     else:
                         processor_data_to_display[processor] = self.processor_artifacts[processor]
 
-        seed_column_names = None if hide_seed_columns else self.dataset_metadata.seed_column_names
+        seed_column_names = (
+            None if hide_seed_columns or self.dataset_metadata is None else self.dataset_metadata.seed_column_names
+        )
 
         display_sample_record(
             record=record,

data_designer/config/validator_params.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 enum import Enum
 from typing import Any
 

data_designer/engine/analysis/column_profilers/base.py

@@ -5,15 +5,19 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
-import pandas as pd
 from pydantic import BaseModel, model_validator
 from typing_extensions import Self
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
-from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, TaskConfigT
+from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -32,17 +36,14 @@ class ColumnConfigWithDataFrame(ConfigBase):
         return (self.column_config, self.df)
 
 
-class ColumnProfilerMetadata(ConfigurableTaskMetadata):
-    applicable_column_types: list[DataDesignerColumnType]
-
-
 class ColumnProfiler(ConfigurableTask[TaskConfigT], ABC):
     @staticmethod
     @abstractmethod
-    def metadata() -> ColumnProfilerMetadata: ...
+    def get_applicable_column_types() -> list[DataDesignerColumnType]:
+        """Returns a list of column types that this profiler can be applied to during dataset profiling."""
 
     @abstractmethod
     def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> BaseModel: ...
 
     def _initialize(self) -> None:
-        logger.info(f"💫 Initializing column profiler: '{self.metadata().name}'")
+        logger.info(f"💫 Initializing column profiler: '{self.name}'")

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -5,44 +5,41 @@ from __future__ import annotations
 
 import logging
 import random
+from typing import TYPE_CHECKING
 
 from data_designer.config.analysis.column_profilers import (
     JudgeScoreProfilerConfig,
     JudgeScoreProfilerResults,
-    JudgeScoreSample,
     JudgeScoreSummary,
 )
 from data_designer.config.analysis.column_statistics import (
-    CategoricalDistribution,
-    CategoricalHistogramData,
     ColumnDistributionType,
     MissingValue,
-    NumericalDistribution,
-)
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP, DataDesignerColumnType
-from data_designer.engine.analysis.column_profilers.base import (
-    ColumnConfigWithDataFrame,
-    ColumnProfiler,
-    ColumnProfilerMetadata,
 )
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.utils.judge_score_processing import (
     extract_judge_score_distributions,
     sample_scores_and_reasoning,
 )
-from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.recipes.response_recipes import TextResponseRecipe
 
+if TYPE_CHECKING:
+    from data_designer.config.analysis.column_profilers import JudgeScoreSample
+    from data_designer.config.analysis.column_statistics import (
+        CategoricalDistribution,
+        CategoricalHistogramData,
+        NumericalDistribution,
+    )
+    from data_designer.engine.models.facade import ModelFacade
+
 logger = logging.getLogger(__name__)
 
 
 class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
     @staticmethod
-    def metadata() -> ColumnProfilerMetadata:
-        return ColumnProfilerMetadata(
-            name="judge_score_profiler",
-            description="Analyzes LLM-as-judge score distributions in a Data Designer dataset.",
-            applicable_column_types=[DataDesignerColumnType.LLM_JUDGE],
-        )
+    def get_applicable_column_types() -> list[DataDesignerColumnType]:
+        return [DataDesignerColumnType.LLM_JUDGE]
 
     def get_model(self, model_alias: str) -> ModelFacade:
         return self.resource_provider.model_registry.get_model(model_alias=model_alias)
@@ -51,8 +48,7 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
         column_config, df = column_config_with_df.as_tuple()
 
         logger.info(
-            f"{COLUMN_TYPE_EMOJI_MAP[column_config.column_type]} Analyzing LLM-as-judge "
-            f"scores for column: '{column_config.name}'"
+            f"{column_config.get_column_emoji()} Analyzing LLM-as-judge scores for column: '{column_config.name}'"
         )
 
         score_summaries = {}

data_designer/engine/analysis/column_profilers/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.config.analysis.column_profilers import ColumnProfilerType
 from data_designer.config.base import ConfigBase
 from data_designer.engine.analysis.column_profilers.base import ColumnProfiler

data_designer/engine/analysis/column_statistics.py

@@ -4,9 +4,8 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, TypeAlias
+from typing import TYPE_CHECKING, Any, TypeAlias
 
-import pandas as pd
 from pydantic import BaseModel
 from typing_extensions import Self
 
@@ -25,6 +24,10 @@ from data_designer.engine.analysis.utils.column_statistics_calculations import (
     calculate_token_stats,
     calculate_validation_column_info,
 )
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)