data-designer 0.1.5__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_designer/config/utils/code_lang.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Union
 
 
 class CodeLang(str, Enum):
@@ -26,17 +25,17 @@ class CodeLang(str, Enum):
     SQL_ANSI = "sql:ansi"
 
     @staticmethod
-    def parse(value: Union[str, CodeLang]) -> tuple[str, Union[str, None]]:
+    def parse(value: str | CodeLang) -> tuple[str, str | None]:
         value = value.value if isinstance(value, CodeLang) else value
         split_vals = value.split(":")
         return (split_vals[0], split_vals[1] if len(split_vals) > 1 else None)
 
     @staticmethod
-    def parse_lang(value: Union[str, CodeLang]) -> str:
+    def parse_lang(value: str | CodeLang) -> str:
         return CodeLang.parse(value)[0]
 
     @staticmethod
-    def parse_dialect(value: Union[str, CodeLang]) -> Union[str, None]:
+    def parse_dialect(value: str | CodeLang) -> str | None:
         return CodeLang.parse(value)[1]
 
     @staticmethod
@@ -58,7 +57,7 @@ SQL_DIALECTS: set[CodeLang] = {
 ##########################################################
 
 
-def code_lang_to_syntax_lexer(code_lang: Union[CodeLang, str]) -> str:
+def code_lang_to_syntax_lexer(code_lang: CodeLang | str) -> str:
     """Convert the code language to a syntax lexer for Pygments.
 
     Reference: https://pygments.org/docs/lexers/

data_designer/config/utils/constants.py

@@ -97,8 +97,6 @@ DEFAULT_AGE_RANGE = [18, 114]
 MIN_AGE = 0
 MAX_AGE = 114
 
-LOCALES_WITH_MANAGED_DATASETS = ["en_US", "ja_JP", "en_IN", "hi_IN"]
-
 US_STATES_AND_MAJOR_TERRITORIES = {
     # States
     "AK",
@@ -299,15 +297,40 @@ PREDEFINED_PROVIDERS = [
     },
 ]
 
+
+DEFAULT_TEXT_INFERENCE_PARAMS = {"temperature": 0.85, "top_p": 0.95}
+DEFAULT_REASONING_INFERENCE_PARAMS = {"temperature": 0.35, "top_p": 0.95}
+DEFAULT_VISION_INFERENCE_PARAMS = {"temperature": 0.85, "top_p": 0.95}
+DEFAULT_EMBEDDING_INFERENCE_PARAMS = {"encoding_format": "float"}
+
+
 PREDEFINED_PROVIDERS_MODEL_MAP = {
     NVIDIA_PROVIDER_NAME: {
-        "text": "nvidia/nvidia-nemotron-nano-9b-v2",
-        "reasoning": "openai/gpt-oss-20b",
-        "vision": "nvidia/nemotron-nano-12b-v2-vl",
+        "text": {"model": "nvidia/nemotron-3-nano-30b-a3b", "inference_parameters": {"temperature": 1.0, "top_p": 1.0}},
+        "reasoning": {"model": "openai/gpt-oss-20b", "inference_parameters": DEFAULT_REASONING_INFERENCE_PARAMS},
+        "vision": {"model": "nvidia/nemotron-nano-12b-v2-vl", "inference_parameters": DEFAULT_VISION_INFERENCE_PARAMS},
+        "embedding": {
+            "model": "nvidia/llama-3.2-nv-embedqa-1b-v2",
+            "inference_parameters": DEFAULT_EMBEDDING_INFERENCE_PARAMS | {"extra_body": {"input_type": "query"}},
+        },
     },
     OPENAI_PROVIDER_NAME: {
-        "text": "gpt-4.1",
-        "reasoning": "gpt-5",
-        "vision": "gpt-5",
+        "text": {"model": "gpt-4.1", "inference_parameters": DEFAULT_TEXT_INFERENCE_PARAMS},
+        "reasoning": {"model": "gpt-5", "inference_parameters": DEFAULT_REASONING_INFERENCE_PARAMS},
+        "vision": {"model": "gpt-5", "inference_parameters": DEFAULT_VISION_INFERENCE_PARAMS},
+        "embedding": {"model": "text-embedding-3-large", "inference_parameters": DEFAULT_EMBEDDING_INFERENCE_PARAMS},
     },
 }
+
+# Persona locale metadata - used by the CLI and the person sampler.
+NEMOTRON_PERSONAS_DATASET_SIZES = {
+    "en_US": "1.24 GB",
+    "en_IN": "2.39 GB",
+    "hi_Deva_IN": "4.14 GB",
+    "hi_Latn_IN": "2.7 GB",
+    "ja_JP": "1.69 GB",
+}
+
+LOCALES_WITH_MANAGED_DATASETS = list[str](NEMOTRON_PERSONAS_DATASET_SIZES.keys())
+
+NEMOTRON_PERSONAS_DATASET_PREFIX = "nemotron-personas-dataset-"

data_designer/config/utils/io_helpers.py

@@ -8,7 +8,7 @@ from datetime import date, datetime, timedelta
 from decimal import Decimal
 from numbers import Number
 from pathlib import Path
-from typing import Any, Union
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -128,7 +128,7 @@ def write_seed_dataset(dataframe: pd.DataFrame, file_path: Path) -> None:
         dataframe.to_json(file_path, orient="records", lines=True)
 
 
-def validate_dataset_file_path(file_path: Union[str, Path], should_exist: bool = True) -> Path:
+def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path:
     """Validate that a dataset file path has a valid extension and optionally exists.
 
     Args:
@@ -165,7 +165,7 @@ def validate_path_contains_files_of_type(path: str | Path, file_extension: str)
         raise InvalidFilePathError(f"🛑 Path {path!r} does not contain files of type {file_extension!r}.")
 
 
-def smart_load_dataframe(dataframe: Union[str, Path, pd.DataFrame]) -> pd.DataFrame:
+def smart_load_dataframe(dataframe: str | Path | pd.DataFrame) -> pd.DataFrame:
     """Load a dataframe from file if a path is given, otherwise return the dataframe.
 
     Args:
@@ -197,7 +197,7 @@ def smart_load_dataframe(dataframe: Union[str, Path, pd.DataFrame]) -> pd.DataFr
         raise ValueError(f"Unsupported file format: {dataframe}")
 
 
-def smart_load_yaml(yaml_in: Union[str, Path, dict]) -> dict:
+def smart_load_yaml(yaml_in: str | Path | dict) -> dict:
     """Return the yaml config as a dict given flexible input types.
 
     Args:
@@ -227,7 +227,7 @@ def smart_load_yaml(yaml_in: Union[str, Path, dict]) -> dict:
     return yaml_out
 
 
-def serialize_data(data: Union[dict, list, str, Number], **kwargs) -> str:
+def serialize_data(data: dict | list | str | Number, **kwargs) -> str:
     if isinstance(data, dict):
         return json.dumps(data, ensure_ascii=False, default=_convert_to_serializable, **kwargs)
     elif isinstance(data, list):

data_designer/config/utils/misc.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import json
 from contextlib import contextmanager
-from typing import Optional, Union
 
 from jinja2 import TemplateSyntaxError, meta
 from jinja2.sandbox import ImmutableSandboxedEnvironment
@@ -58,9 +57,7 @@ def get_prompt_template_keywords(template: str) -> set[str]:
     return keywords
 
 
-def json_indent_list_of_strings(
-    column_names: list[str], *, indent: Optional[Union[int, str]] = None
-) -> Optional[Union[list[str], str]]:
+def json_indent_list_of_strings(column_names: list[str], *, indent: int | str | None = None) -> list[str] | str | None:
     """Convert a list of column names to a JSON string if the list is long.
 
     This function helps keep Data Designer's __repr__ output clean and readable.

data_designer/config/utils/numerical_helpers.py

@@ -3,7 +3,7 @@
 
 import numbers
 from numbers import Number
-from typing import Any, Type
+from typing import Any
 
 from data_designer.config.utils.constants import REPORTING_PRECISION
 
@@ -18,7 +18,7 @@ def is_float(val: Any) -> bool:
 
 def prepare_number_for_reporting(
     value: Number,
-    target_type: Type[Number],
+    target_type: type[Number],
     precision: int = REPORTING_PRECISION,
 ) -> Number:
     """Ensure native python types and round to `precision` decimal digits."""

data_designer/config/utils/type_helpers.py

@@ -3,7 +3,7 @@
 
 import inspect
 from enum import Enum
-from typing import Any, Literal, Type, get_args, get_origin
+from typing import Any, Literal, get_args, get_origin
 
 from pydantic import BaseModel
 
@@ -56,7 +56,7 @@ def create_str_enum_from_discriminated_type_union(
     return StrEnum(enum_name, {v.replace("-", "_").upper(): v for v in set(discriminator_field_values)})
 
 
-def get_sampler_params() -> dict[str, Type[BaseModel]]:
+def get_sampler_params() -> dict[str, type[BaseModel]]:
     """Returns a dictionary of sampler parameter classes."""
     params_cls_list = [
         params_cls
@@ -83,7 +83,7 @@ def get_sampler_params() -> dict[str, Type[BaseModel]]:
     return params_cls_dict
 
 
-def resolve_string_enum(enum_instance: Any, enum_type: Type[Enum]) -> Enum:
+def resolve_string_enum(enum_instance: Any, enum_type: type[Enum]) -> Enum:
     if not issubclass(enum_type, Enum):
         raise InvalidEnumValueError(f"🛑 `enum_type` must be a subclass of Enum. You provided: {enum_type}")
     invalid_enum_value_error = InvalidEnumValueError(

data_designer/config/utils/validation.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 from enum import Enum
 from string import Formatter
-from typing import Optional
 
 from jinja2 import meta
 from jinja2.sandbox import ImmutableSandboxedEnvironment
@@ -15,8 +14,8 @@ from rich.console import Console, Group
 from rich.padding import Padding
 from rich.panel import Panel
 
-from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType, column_type_is_llm_generated
-from data_designer.config.processors import ProcessorConfig, ProcessorType
+from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType, column_type_is_model_generated
+from data_designer.config.processors import ProcessorConfigT, ProcessorType
 from data_designer.config.utils.constants import RICH_CONSOLE_THEME
 from data_designer.config.utils.misc import (
     can_run_data_designer_locally,
@@ -45,7 +44,7 @@ class ViolationLevel(str, Enum):
 
 
 class Violation(BaseModel):
-    column: Optional[str] = None
+    column: str | None = None
     type: ViolationType
     message: str
     level: ViolationLevel
@@ -57,7 +56,7 @@ class Violation(BaseModel):
 
 def validate_data_designer_config(
     columns: list[ColumnConfigT],
-    processor_configs: list[ProcessorConfig],
+    processor_configs: list[ProcessorConfigT],
     allowed_references: list[str],
 ) -> list[Violation]:
     violations = []
@@ -119,7 +118,7 @@ def validate_prompt_templates(
 ) -> list[Violation]:
     env = ImmutableSandboxedEnvironment()
 
-    columns_with_prompts = [c for c in columns if column_type_is_llm_generated(c.column_type)]
+    columns_with_prompts = [c for c in columns if column_type_is_model_generated(c.column_type)]
 
     violations = []
     for column in columns_with_prompts:
@@ -273,7 +272,7 @@ def validate_columns_not_all_dropped(
 
 def validate_drop_columns_processor(
     columns: list[ColumnConfigT],
-    processor_configs: list[ProcessorConfig],
+    processor_configs: list[ProcessorConfigT],
 ) -> list[Violation]:
     all_column_names = {c.name for c in columns}
     for processor_config in processor_configs:
@@ -294,7 +293,7 @@ def validate_drop_columns_processor(
 
 def validate_schema_transform_processor(
     columns: list[ColumnConfigT],
-    processor_configs: list[ProcessorConfig],
+    processor_configs: list[ProcessorConfigT],
 ) -> list[Violation]:
     violations = []
 

data_designer/config/utils/visualization.py

@@ -8,7 +8,7 @@ import os
 from collections import OrderedDict
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
@@ -36,11 +36,11 @@ if TYPE_CHECKING:
 console = Console()
 
 
-def get_nvidia_api_key() -> Optional[str]:
+def get_nvidia_api_key() -> str | None:
     return os.getenv(NVIDIA_API_KEY_ENV_VAR_NAME)
 
 
-def get_openai_api_key() -> Optional[str]:
+def get_openai_api_key() -> str | None:
     return os.getenv(OPENAI_API_KEY_ENV_VAR_NAME)
 
 
@@ -77,12 +77,12 @@ class WithRecordSamplerMixin:
 
     def display_sample_record(
         self,
-        index: Optional[int] = None,
+        index: int | None = None,
         *,
         hide_seed_columns: bool = False,
         syntax_highlighting_theme: str = "dracula",
-        background_color: Optional[str] = None,
-        processors_to_display: Optional[list[str]] = None,
+        background_color: str | None = None,
+        processors_to_display: list[str] | None = None,
     ) -> None:
         """Display a sample record from the Data Designer dataset preview.
 
@@ -134,11 +134,11 @@ class WithRecordSamplerMixin:
 
 
 def create_rich_histogram_table(
-    data: dict[str, Union[int, float]],
+    data: dict[str, int | float],
     column_names: tuple[int, int],
     name_style: str = ColorPalette.BLUE.value,
     value_style: str = ColorPalette.TEAL.value,
-    title: Optional[str] = None,
+    title: str | None = None,
     **kwargs,
 ) -> Table:
     table = Table(title=title, **kwargs)
@@ -154,12 +154,12 @@ def create_rich_histogram_table(
 
 
 def display_sample_record(
-    record: Union[dict, pd.Series, pd.DataFrame],
+    record: dict | pd.Series | pd.DataFrame,
     config_builder: DataDesignerConfigBuilder,
-    processor_data_to_display: Optional[dict[str, Union[list[str], str]]] = None,
-    background_color: Optional[str] = None,
+    processor_data_to_display: dict[str, list[str] | str] | None = None,
+    background_color: str | None = None,
     syntax_highlighting_theme: str = "dracula",
-    record_index: Optional[int] = None,
+    record_index: int | None = None,
     hide_seed_columns: bool = False,
 ):
     if isinstance(record, (dict, pd.Series)):
@@ -194,6 +194,7 @@ def display_sample_record(
         + config_builder.get_columns_of_type(DataDesignerColumnType.EXPRESSION)
         + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_TEXT)
         + config_builder.get_columns_of_type(DataDesignerColumnType.LLM_STRUCTURED)
+        + config_builder.get_columns_of_type(DataDesignerColumnType.EMBEDDING)
     )
     if len(non_code_columns) > 0:
         table = Table(title="Generated Columns", **table_kws)
@@ -201,6 +202,10 @@ def display_sample_record(
         table.add_column("Value")
         for col in non_code_columns:
             if not col.drop:
+                if col.column_type == DataDesignerColumnType.EMBEDDING:
+                    record[col.name]["embeddings"] = [
+                        get_truncated_list_as_string(embd) for embd in record[col.name].get("embeddings")
+                    ]
                 table.add_row(col.name, convert_to_row_element(record[col.name]))
         render_list.append(pad_console_element(table))
 
@@ -269,9 +274,19 @@ def display_sample_record(
     console.print(Group(*render_list), markup=False)
 
 
+def get_truncated_list_as_string(long_list: list[Any], max_items: int = 2) -> str:
+    if max_items <= 0:
+        raise ValueError("max_items must be greater than 0")
+    if len(long_list) > max_items:
+        truncated_part = long_list[:max_items]
+        return f"[{', '.join(str(x) for x in truncated_part)}, ...]"
+    else:
+        return str(long_list)
+
+
 def display_sampler_table(
     sampler_params: dict[SamplerType, ConfigBase],
-    title: Optional[str] = None,
+    title: str | None = None,
 ) -> None:
     table = Table(expand=True)
     table.add_column("Type")
@@ -306,15 +321,15 @@ def display_model_configs_table(model_configs: list[ModelConfig]) -> None:
     table_model_configs.add_column("Alias")
     table_model_configs.add_column("Model")
     table_model_configs.add_column("Provider")
-    table_model_configs.add_column("Temperature")
-    table_model_configs.add_column("Top P")
+    table_model_configs.add_column("Inference Parameters")
     for model_config in model_configs:
+        params_display = model_config.inference_parameters.format_for_display()
+
         table_model_configs.add_row(
             model_config.alias,
             model_config.model,
             model_config.provider,
-            str(model_config.inference_parameters.temperature),
-            str(model_config.inference_parameters.top_p),
+            params_display,
         )
     group_args: list = [Rule(title="Model Configs"), table_model_configs]
     if len(model_configs) == 0:

data_designer/config/validator_params.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
-from typing import Any, Optional, Union
+from typing import Any
 
 from pydantic import Field, field_serializer, model_validator
 from typing_extensions import Self, TypeAlias
@@ -51,7 +51,7 @@ class LocalCallableValidatorParams(ConfigBase):
     validation_function: Any = Field(
         description="Function (Callable[[pd.DataFrame], pd.DataFrame]) to validate the data"
     )
-    output_schema: Optional[dict[str, Any]] = Field(
+    output_schema: dict[str, Any] | None = Field(
         default=None, description="Expected schema for local callable validator's output"
     )
 
@@ -80,7 +80,7 @@ class RemoteValidatorParams(ConfigBase):
     """
 
     endpoint_url: str = Field(description="URL of the remote endpoint")
-    output_schema: Optional[dict[str, Any]] = Field(
+    output_schema: dict[str, Any] | None = Field(
         default=None, description="Expected schema for remote validator's output"
     )
     timeout: float = Field(default=30.0, gt=0, description="The timeout for the HTTP request")
@@ -89,8 +89,4 @@ class RemoteValidatorParams(ConfigBase):
     max_parallel_requests: int = Field(default=4, ge=1, description="The maximum number of parallel requests to make")
 
 
-ValidatorParamsT: TypeAlias = Union[
-    CodeValidatorParams,
-    LocalCallableValidatorParams,
-    RemoteValidatorParams,
-]
+ValidatorParamsT: TypeAlias = CodeValidatorParams | LocalCallableValidatorParams | RemoteValidatorParams

data_designer/engine/analysis/column_profilers/base.py

@@ -7,7 +7,6 @@ import logging
 from abc import ABC, abstractmethod
 
 import pandas as pd
-import pyarrow as pa
 from pydantic import BaseModel, model_validator
 from typing_extensions import Self
 
@@ -29,12 +28,6 @@ class ColumnConfigWithDataFrame(ConfigBase):
             raise ValueError(f"Column {self.column_config.name!r} not found in DataFrame")
         return self
 
-    @model_validator(mode="after")
-    def ensure_pyarrow_backend(self) -> Self:
-        if not all(isinstance(dtype, pd.ArrowDtype) for dtype in self.df.dtypes):
-            self.df = pa.Table.from_pandas(self.df).to_pandas(types_mapper=pd.ArrowDtype)
-        return self
-
     def as_tuple(self) -> tuple[SingleColumnConfig, pd.DataFrame]:
         return (self.column_config, self.df)
 

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import logging
 import random
-from typing import Union
 
 from data_designer.config.analysis.column_profilers import (
     JudgeScoreProfilerConfig,
@@ -69,7 +68,7 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
             )
 
         for score in column_config.scores:
-            score_name = score.name.lower()
+            score_name = score.name
             logger.info(f"{random.choice(['👩‍⚖️', '👨‍⚖️'])} Summarizing LLM-as-judge score: '{score_name}'")
             score_sample = sample_scores_and_reasoning(
                 scores=score_distributions.scores[score_name],
@@ -96,7 +95,7 @@ class JudgeScoreProfiler(ColumnProfiler[JudgeScoreProfilerConfig]):
         name: str,
         sample: list[JudgeScoreSample],
         histogram: CategoricalHistogramData,
-        distribution: Union[CategoricalDistribution, NumericalDistribution, MissingValue],
+        distribution: CategoricalDistribution | NumericalDistribution | MissingValue,
         distribution_type: ColumnDistributionType,
     ) -> JudgeScoreSummary:
         if isinstance(distribution, MissingValue) or not sample:

data_designer/engine/analysis/column_statistics.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Type, TypeAlias, Union
+from typing import Any, TypeAlias
 
 import pandas as pd
 from pydantic import BaseModel
@@ -41,7 +41,7 @@ class GeneralColumnStatisticsCalculator(BaseModel):
         return self.column_config_with_df.df
 
     @property
-    def column_statistics_type(self) -> Type[ColumnStatisticsT]:
+    def column_statistics_type(self) -> type[ColumnStatisticsT]:
         return DEFAULT_COLUMN_STATISTICS_MAP.get(self.column_config.column_type, GeneralColumnStatistics)
 
     def calculate(self) -> Self:
@@ -59,7 +59,7 @@ class GeneralColumnStatisticsCalculator(BaseModel):
         )
 
     def calculate_general_column_info(self) -> dict[str, Any]:
-        return calculate_general_column_info(self.column_config, self.df)
+        return calculate_general_column_info(self.column_config.name, self.df)
 
     def __repr__(self) -> str:
         params = []
@@ -93,7 +93,7 @@ class SamplerColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
         return (
             {
                 "sampler_type": SamplerType(self.column_config.sampler_type),
-                **calculate_column_distribution(self.column_config, self.df, dist_type),
+                **calculate_column_distribution(self.column_config.name, self.df, dist_type),
             }
             if make_dist
             else {
@@ -109,23 +109,23 @@ class SeedDatasetColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
 
 class ValidationColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
     def calculate_validation_column_info(self) -> dict[str, Any]:
-        return calculate_validation_column_info(self.column_config, self.df)
+        return calculate_validation_column_info(self.column_config.name, self.df)
 
 
 class ExpressionColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
 
 
-ColumnStatisticsCalculatorT: TypeAlias = Union[
-    ExpressionColumnStatisticsCalculator,
-    ValidationColumnStatisticsCalculator,
-    GeneralColumnStatisticsCalculator,
-    LLMCodeColumnStatisticsCalculator,
-    LLMJudgedColumnStatisticsCalculator,
-    LLMStructuredColumnStatisticsCalculator,
-    LLMTextColumnStatisticsCalculator,
-    SamplerColumnStatisticsCalculator,
-    SeedDatasetColumnStatisticsCalculator,
-]
+ColumnStatisticsCalculatorT: TypeAlias = (
+    ExpressionColumnStatisticsCalculator
+    | ValidationColumnStatisticsCalculator
+    | GeneralColumnStatisticsCalculator
+    | LLMCodeColumnStatisticsCalculator
+    | LLMJudgedColumnStatisticsCalculator
+    | LLMStructuredColumnStatisticsCalculator
+    | LLMTextColumnStatisticsCalculator
+    | SamplerColumnStatisticsCalculator
+    | SeedDatasetColumnStatisticsCalculator
+)
 DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP = {
     DataDesignerColumnType.EXPRESSION: ExpressionColumnStatisticsCalculator,
     DataDesignerColumnType.VALIDATION: ValidationColumnStatisticsCalculator,

data_designer/engine/analysis/dataset_profiler.py

@@ -6,6 +6,7 @@ from collections.abc import Sequence
 from functools import cached_property
 
 import pandas as pd
+import pyarrow as pa
 from pydantic import Field, field_validator
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
@@ -19,10 +20,8 @@ from data_designer.config.column_types import (
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator
 from data_designer.engine.analysis.errors import DatasetProfilerConfigurationError
-from data_designer.engine.dataset_builders.multi_column_configs import (
-    DatasetBuilderColumnConfigT,
-    MultiColumnConfig,
-)
+from data_designer.engine.analysis.utils.column_statistics_calculations import has_pyarrow_backend
+from data_designer.engine.dataset_builders.multi_column_configs import DatasetBuilderColumnConfigT, MultiColumnConfig
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
 from data_designer.engine.resources.resource_provider import ResourceProvider
 
@@ -68,6 +67,7 @@ class DataDesignerDatasetProfiler:
         logger.info("📐 Measuring dataset column statistics:")
 
         self._validate_schema_consistency(list(dataset.columns))
+        dataset = self._convert_to_pyarrow_backend_if_needed(dataset)
 
         column_statistics = []
         for c in self.config.column_configs:
@@ -100,6 +100,27 @@ class DataDesignerDatasetProfiler:
             column_profiles=column_profiles if column_profiles else None,
         )
 
+    def _convert_to_pyarrow_backend_if_needed(self, dataset: pd.DataFrame) -> pd.DataFrame:
+        if not has_pyarrow_backend(dataset):
+            try:
+                dataset = pa.Table.from_pandas(dataset).to_pandas(types_mapper=pd.ArrowDtype)
+            except Exception as e:
+                # For ArrowTypeError, the second arg contains the more informative message
+                if isinstance(e, pa.lib.ArrowTypeError) and len(e.args) > 1:
+                    error_msg = str(e.args[1])
+                else:
+                    error_msg = str(e)
+                for col in dataset.columns:
+                    # Make sure column names are clear in the error message
+                    error_msg = error_msg.replace(col, f"'{col}'")
+                logger.warning("⚠️ Unable to convert the dataset to a PyArrow backend")
+                logger.warning(f"  |-- Conversion Error Message: {error_msg}")
+                logger.warning("  |-- This is often due to at least one column having mixed data types")
+                logger.warning(
+                    "  |-- Note: Reported data types will be inferred from the first non-null value of each column"
+                )
+        return dataset
+
     def _create_column_profiler(self, profiler_config: ColumnProfilerConfigT) -> ColumnProfiler:
         return self.registry.column_profilers.get_for_config_type(type(profiler_config))(
             config=profiler_config, resource_provider=self.resource_provider