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

data_designer/cli/ui.py

@@ -182,6 +182,137 @@ def select_with_arrows(
         return None
 
 
+def select_multiple_with_arrows(
+    options: dict[str, str],
+    prompt_text: str,
+    default_keys: list[str] | None = None,
+    allow_empty: bool = False,
+) -> list[str] | None:
+    """Interactive multi-selection with arrow key navigation and space to toggle.
+
+    Uses prompt_toolkit's Application for an inline checkbox-style menu experience.
+
+    Args:
+        options: Dictionary of {key: display_text} options
+        prompt_text: Prompt to display above options
+        default_keys: List of keys that should be pre-selected
+        allow_empty: If True, allows user to submit with no selections
+
+    Returns:
+        List of selected keys, or None if cancelled
+    """
+    if not options:
+        return None
+
+    # Build list of keys and track selected state
+    keys = list(options.keys())
+    selected_set = set(default_keys) if default_keys else set()
+    current_index = 0
+
+    # Store result
+    result = {"value": None, "cancelled": False}
+
+    def get_formatted_text() -> list[tuple[str, str]]:
+        """Generate the formatted text for the multi-select menu."""
+        text = []
+        # Add prompt with padding
+        padding = " " * LEFT_PADDING
+        text.append(("", f"{padding}{prompt_text}\n"))
+
+        # Add options with checkboxes
+        for i, key in enumerate(keys):
+            display = options[key]
+            checkbox = "[✓]" if key in selected_set else "[ ]"
+
+            if i == current_index:
+                # Highlighted item with Nord8 color
+                text.append((f"fg:{NordColor.NORD8.value} bold", f"{padding}  → {checkbox} {display}\n"))
+            else:
+                # Unselected item
+                text.append(("", f"{padding}    {checkbox} {display}\n"))
+
+        # Add hint
+        count = len(selected_set)
+        text.append(
+            (
+                "fg:#666666",
+                f"{padding}  (↑/↓: navigate, Space: toggle, Enter: confirm ({count} selected), Esc: cancel)\n",
+            )
+        )
+        return text
+
+    # Create key bindings
+    kb = KeyBindings()
+
+    @kb.add("up")
+    @kb.add("c-p")  # Ctrl+P
+    def _move_up(event) -> None:
+        nonlocal current_index
+        current_index = (current_index - 1) % len(keys)
+
+    @kb.add("down")
+    @kb.add("c-n")  # Ctrl+N
+    def _move_down(event) -> None:
+        nonlocal current_index
+        current_index = (current_index + 1) % len(keys)
+
+    @kb.add("c-h")  # Ctrl+H as alternative
+    @kb.add(" ", eager=True)  # Space key - eager to capture immediately
+    def _toggle(event) -> None:
+        key = keys[current_index]
+        if key in selected_set:
+            selected_set.remove(key)
+        else:
+            selected_set.add(key)
+
+    @kb.add("enter")
+    def _confirm(event) -> None:
+        if not allow_empty and not selected_set:
+            # Don't allow empty selection if not permitted
+            return
+        result["value"] = list(selected_set)
+        event.app.exit()
+
+    @kb.add("escape")
+    @kb.add("c-c")  # Ctrl+C
+    def _cancel(event) -> None:
+        result["cancelled"] = True
+        event.app.exit()
+
+    # Create the application
+    app = Application(
+        layout=Layout(
+            HSplit(
+                [
+                    Window(
+                        content=FormattedTextControl(get_formatted_text),
+                        dont_extend_height=True,
+                        always_hide_cursor=True,
+                    )
+                ]
+            )
+        ),
+        key_bindings=kb,
+        full_screen=False,
+        mouse_support=False,
+    )
+
+    try:
+        # Run the application
+        app.run()
+
+        # Handle the result
+        if result["cancelled"]:
+            print_warning("Cancelled")
+            return None
+        else:
+            return result["value"]
+
+    except (KeyboardInterrupt, EOFError):
+        print_warning("Cancelled")
+        return None
+
+
 def prompt_text_input(
     prompt_msg: str,
     default: str | None = None,

data_designer/cli/utils.py

@@ -1,6 +1,40 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import shutil
+import subprocess
+
+
+def check_ngc_cli_available() -> bool:
+    """Check if NGC CLI is installed and available.
+
+    Returns:
+        True if NGC CLI is in PATH and executable, False otherwise.
+    """
+    if shutil.which("ngc") is None:
+        return False
+
+    return get_ngc_version() is not None
+
+
+def get_ngc_version() -> str | None:
+    """Get the NGC CLI version if available.
+
+    Returns:
+        NGC CLI version string if available, None otherwise.
+    """
+    try:
+        result = subprocess.run(
+            ["ngc", "--version"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=5,
+        )
+        return result.stdout.strip()
+    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
+        return None
+
 
 def validate_url(url: str) -> bool:
     """Validate that a string is a valid URL.

data_designer/config/analysis/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

data_designer/config/analysis/column_profilers.py

@@ -3,7 +3,6 @@
 
 from abc import ABC
 from enum import Enum
-from typing import Optional, Union
 
 from pydantic import BaseModel, Field
 from rich.panel import Panel
@@ -27,7 +26,20 @@ class ColumnProfilerType(str, Enum):
 
 
 class ColumnProfilerResults(BaseModel, ABC):
+    """Abstract base class for column profiler results.
+
+    Stores results from column profiling operations. Subclasses hold profiler-specific
+    analysis results and provide methods for generating formatted report sections for display.
+    """
+
     def create_report_section(self) -> Panel:
+        """Creates a Rich Panel containing the formatted profiler results for display.
+
+        Returns:
+            A Rich Panel containing the formatted profiler results. Default implementation
+            returns a "Not Implemented" message; subclasses should override to provide
+            specific formatting.
+        """
         return Panel(
             f"Report section generation not implemented for '{self.__class__.__name__}'.",
             title="Not Implemented",
@@ -37,33 +49,89 @@ class ColumnProfilerResults(BaseModel, ABC):
 
 
 class JudgeScoreProfilerConfig(ConfigBase):
+    """Configuration for the LLM-as-a-judge score profiler.
+
+    Attributes:
+        model_alias: Alias of the LLM model to use for generating score distribution summaries.
+            Must match a model alias defined in the Data Designer configuration.
+        summary_score_sample_size: Number of score samples to include when prompting the LLM
+            to generate summaries. Larger sample sizes provide more context but increase
+            token usage. Must be at least 1. Defaults to 20.
+    """
+
     model_alias: str
-    summary_score_sample_size: Optional[int] = Field(default=20, ge=1)
+    summary_score_sample_size: int | None = Field(default=20, ge=1)
 
 
 class JudgeScoreSample(BaseModel):
-    score: Union[int, str]
+    """Container for a single judge score and its associated reasoning.
+
+    Stores a paired score-reasoning sample extracted from an LLM-as-a-judge column.
+    Used when generating summaries to provide the LLM with examples of scoring patterns.
+
+    Attributes:
+        score: The score value assigned by the judge. Can be numeric (int) or categorical (str).
+        reasoning: The reasoning or explanation provided by the judge for this score.
+    """
+
+    score: int | str
     reasoning: str
 
 
 class JudgeScoreDistributions(BaseModel):
-    scores: dict[str, list[Union[int, str]]]
+    """Container for computed distributions across all judge score dimensions.
+
+    Stores the complete distribution analysis for all score dimensions in an LLM-as-a-judge
+    column. Each score dimension (e.g., "relevance", "fluency") has its own distribution
+    computed from the generated data.
+
+    Attributes:
+        scores: Mapping of each score dimension name to its list of score values.
+        reasoning: Mapping of each score dimension name to its list of reasoning texts.
+        distribution_types: Mapping of each score dimension name to its classification.
+        distributions: Mapping of each score dimension name to its computed distribution statistics.
+        histograms: Mapping of each score dimension name to its histogram data.
+    """
+
+    scores: dict[str, list[int | str]]
     reasoning: dict[str, list[str]]
     distribution_types: dict[str, ColumnDistributionType]
-    distributions: dict[str, Union[CategoricalDistribution, NumericalDistribution, MissingValue]]
-    histograms: dict[str, Union[CategoricalHistogramData, MissingValue]]
+    distributions: dict[str, CategoricalDistribution | NumericalDistribution | MissingValue]
+    histograms: dict[str, CategoricalHistogramData | MissingValue]
 
 
 class JudgeScoreSummary(BaseModel):
+    """Container for an LLM-generated summary of a judge score dimension.
+
+    Stores the natural language summary and sample data for a single score dimension
+    generated by the judge score profiler. The summary is created by an LLM analyzing
+    the distribution and patterns in the score-reasoning pairs.
+
+    Attributes:
+        score_name: Name of the score dimension being summarized (e.g., "relevance", "fluency").
+        summary: LLM-generated natural language summary describing the scoring patterns,
+            distribution characteristics, and notable trends for this score dimension.
+        score_samples: List of score-reasoning pairs that were used to generate the summary.
+            These are the examples of the scoring behavior that were used to generate the summary.
+    """
+
     score_name: str
     summary: str
     score_samples: list[JudgeScoreSample]
 
 
 class JudgeScoreProfilerResults(ColumnProfilerResults):
+    """Container for complete judge score profiler analysis results.
+
+    Attributes:
+        column_name: Name of the judge column that was profiled.
+        summaries: Mapping of each score dimension name to its LLM-generated summary.
+        score_distributions: Complete distribution analysis across all score dimensions.
+    """
+
     column_name: str
     summaries: dict[str, JudgeScoreSummary]
-    score_distributions: Union[JudgeScoreDistributions, MissingValue]
+    score_distributions: JudgeScoreDistributions | MissingValue
 
     def create_report_section(self) -> Panel:
         layout = Table.grid(Column(), expand=True, padding=(2, 0))

data_designer/config/analysis/column_statistics.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal
 
 from pandas import Series
 from pydantic import BaseModel, ConfigDict, create_model, field_validator, model_validator
@@ -32,27 +32,56 @@ class ColumnDistributionType(str, Enum):
 
 
 class BaseColumnStatistics(BaseModel, ABC):
+    """Abstract base class for all column statistics types.
+
+    Serves as a container for computed statistics across different column types in
+    Data-Designer-generated datasets. Subclasses hold column-specific statistical results
+    and provide methods for formatting these results for display in reports.
+    """
+
     model_config = ConfigDict(use_enum_values=True)
 
     @abstractmethod
-    def create_report_row_data(self) -> dict[str, str]: ...
+    def create_report_row_data(self) -> dict[str, str]:
+        """Creates a formatted dictionary of statistics for display in reports.
+
+        Returns:
+            Dictionary mapping display labels to formatted statistic values.
+        """
+        ...
 
 
 class GeneralColumnStatistics(BaseColumnStatistics):
+    """Container for general statistics applicable to all column types.
+
+    Holds core statistical measures that apply universally across all column types,
+    including null counts, unique values, and data type information. Serves as the base
+    for more specialized column statistics classes that store additional column-specific metrics.
+
+    Attributes:
+        column_name: Name of the column being analyzed.
+        num_records: Total number of records in the column.
+        num_null: Number of null/missing values in the column.
+        num_unique: Number of distinct values in the column. If a value is not hashable, it is converted to a string.
+        pyarrow_dtype: PyArrow data type of the column as a string.
+        simple_dtype: Simplified human-readable data type label.
+        column_type: Discriminator field, always "general" for this statistics type.
+    """
+
     column_name: str
-    num_records: Union[int, MissingValue]
-    num_null: Union[int, MissingValue]
-    num_unique: Union[int, MissingValue]
+    num_records: int | MissingValue
+    num_null: int | MissingValue
+    num_unique: int | MissingValue
     pyarrow_dtype: str
     simple_dtype: str
     column_type: Literal["general"] = "general"
 
     @field_validator("num_null", "num_unique", "num_records", mode="before")
-    def general_statistics_ensure_python_integers(cls, v: Union[int, MissingValue]) -> Union[int, MissingValue]:
+    def general_statistics_ensure_python_integers(cls, v: int | MissingValue) -> int | MissingValue:
         return v if isinstance(v, MissingValue) else prepare_number_for_reporting(v, int)
 
     @property
-    def percent_null(self) -> Union[float, MissingValue]:
+    def percent_null(self) -> float | MissingValue:
         return (
             self.num_null
             if self._is_missing_value(self.num_null)
@@ -60,7 +89,7 @@ class GeneralColumnStatistics(BaseColumnStatistics):
         )
 
     @property
-    def percent_unique(self) -> Union[float, MissingValue]:
+    def percent_unique(self) -> float | MissingValue:
         return (
             self.num_unique
             if self._is_missing_value(self.num_unique)
@@ -79,40 +108,55 @@ class GeneralColumnStatistics(BaseColumnStatistics):
     def create_report_row_data(self) -> dict[str, str]:
         return self._general_display_row
 
-    def _is_missing_value(self, v: Union[float, int, MissingValue]) -> bool:
+    def _is_missing_value(self, v: float | int | MissingValue) -> bool:
         return v in set(MissingValue)
 
 
 class LLMTextColumnStatistics(GeneralColumnStatistics):
-    completion_tokens_mean: Union[float, MissingValue]
-    completion_tokens_median: Union[float, MissingValue]
-    completion_tokens_stddev: Union[float, MissingValue]
-    prompt_tokens_mean: Union[float, MissingValue]
-    prompt_tokens_median: Union[float, MissingValue]
-    prompt_tokens_stddev: Union[float, MissingValue]
+    """Container for statistics on LLM-generated text columns.
+
+    Inherits general statistics plus token usage metrics specific to LLM text generation.
+    Stores both prompt and completion token consumption data.
+
+    Attributes:
+        output_tokens_mean: Mean number of output tokens generated per record.
+        output_tokens_median: Median number of output tokens generated per record.
+        output_tokens_stddev: Standard deviation of output tokens per record.
+        input_tokens_mean: Mean number of input tokens used per record.
+        input_tokens_median: Median number of input tokens used per record.
+        input_tokens_stddev: Standard deviation of input tokens per record.
+        column_type: Discriminator field, always "llm-text" for this statistics type.
+    """
+
+    output_tokens_mean: float | MissingValue
+    output_tokens_median: float | MissingValue
+    output_tokens_stddev: float | MissingValue
+    input_tokens_mean: float | MissingValue
+    input_tokens_median: float | MissingValue
+    input_tokens_stddev: float | MissingValue
     column_type: Literal[DataDesignerColumnType.LLM_TEXT.value] = DataDesignerColumnType.LLM_TEXT.value
 
     @field_validator(
-        "completion_tokens_mean",
-        "completion_tokens_median",
-        "completion_tokens_stddev",
-        "prompt_tokens_mean",
-        "prompt_tokens_median",
-        "prompt_tokens_stddev",
+        "output_tokens_mean",
+        "output_tokens_median",
+        "output_tokens_stddev",
+        "input_tokens_mean",
+        "input_tokens_median",
+        "input_tokens_stddev",
         mode="before",
     )
-    def llm_column_ensure_python_floats(cls, v: Union[float, int, MissingValue]) -> Union[float, int, MissingValue]:
+    def llm_column_ensure_python_floats(cls, v: float | int | MissingValue) -> float | int | MissingValue:
         return v if isinstance(v, MissingValue) else prepare_number_for_reporting(v, float)
 
     def create_report_row_data(self) -> dict[str, Any]:
         prompt_tokens_str = (
-            f"{self.prompt_tokens_median:.1f} +/- {self.prompt_tokens_stddev:.1f}"
-            if not self._is_missing_value(self.prompt_tokens_median)
+            f"{self.input_tokens_median:.1f} +/- {self.input_tokens_stddev:.1f}"
+            if not self._is_missing_value(self.input_tokens_median)
             else "--"
         )
         completion_tokens_str = (
-            f"{self.completion_tokens_median:.1f} +/- {self.completion_tokens_stddev:.1f}"
-            if not self._is_missing_value(self.completion_tokens_median)
+            f"{self.output_tokens_median:.1f} +/- {self.output_tokens_stddev:.1f}"
+            if not self._is_missing_value(self.output_tokens_median)
             else "--"
         )
         return {
@@ -123,21 +167,65 @@ class LLMTextColumnStatistics(GeneralColumnStatistics):
 
 
 class LLMCodeColumnStatistics(LLMTextColumnStatistics):
+    """Container for statistics on LLM-generated code columns.
+
+    Inherits all token usage metrics from LLMTextColumnStatistics. Stores
+    statistics from columns that generate code snippets in specific programming languages.
+
+    Attributes:
+        column_type: Discriminator field, always "llm-code" for this statistics type.
+    """
+
     column_type: Literal[DataDesignerColumnType.LLM_CODE.value] = DataDesignerColumnType.LLM_CODE.value
 
 
 class LLMStructuredColumnStatistics(LLMTextColumnStatistics):
+    """Container for statistics on LLM-generated structured JSON columns.
+
+    Inherits all token usage metrics from LLMTextColumnStatistics. Stores statistics from
+    columns that generate structured data conforming to JSON schemas or Pydantic models.
+
+    Attributes:
+        column_type: Discriminator field, always "llm-structured" for this statistics type.
+    """
+
     column_type: Literal[DataDesignerColumnType.LLM_STRUCTURED.value] = DataDesignerColumnType.LLM_STRUCTURED.value
 
 
 class LLMJudgedColumnStatistics(LLMTextColumnStatistics):
+    """Container for statistics on LLM-as-a-judge quality assessment columns.
+
+    Inherits all token usage metrics from LLMTextColumnStatistics. Stores statistics from
+    columns that evaluate and score other generated content based on defined criteria.
+
+    Attributes:
+        column_type: Discriminator field, always "llm-judge" for this statistics type.
+    """
+
     column_type: Literal[DataDesignerColumnType.LLM_JUDGE.value] = DataDesignerColumnType.LLM_JUDGE.value
 
 
 class SamplerColumnStatistics(GeneralColumnStatistics):
+    """Container for statistics on sampler-generated columns.
+
+    Inherits general statistics plus sampler-specific information including the sampler type
+    used and the empirical distribution of generated values. Stores both categorical and
+    numerical distribution results.
+
+    Attributes:
+        sampler_type: Type of sampler used to generate this column (e.g., "uniform", "category",
+            "gaussian", "person").
+        distribution_type: Classification of the column's distribution (categorical, numerical,
+            text, other, or unknown).
+        distribution: Empirical distribution statistics for the generated values. Can be
+            CategoricalDistribution (for discrete values), NumericalDistribution (for continuous
+            values), or MissingValue if distribution could not be computed.
+        column_type: Discriminator field, always "sampler" for this statistics type.
+    """
+
     sampler_type: SamplerType
     distribution_type: ColumnDistributionType
-    distribution: Optional[Union[CategoricalDistribution, NumericalDistribution, MissingValue]]
+    distribution: CategoricalDistribution | NumericalDistribution | MissingValue | None
     column_type: Literal[DataDesignerColumnType.SAMPLER.value] = DataDesignerColumnType.SAMPLER.value
 
     def create_report_row_data(self) -> dict[str, str]:
@@ -148,23 +236,52 @@ class SamplerColumnStatistics(GeneralColumnStatistics):
 
 
 class SeedDatasetColumnStatistics(GeneralColumnStatistics):
+    """Container for statistics on columns sourced from seed datasets.
+
+    Inherits general statistics and stores statistics computed from columns that originate
+    from existing data provided via the seed dataset functionality.
+
+    Attributes:
+        column_type: Discriminator field, always "seed-dataset" for this statistics type.
+    """
+
     column_type: Literal[DataDesignerColumnType.SEED_DATASET.value] = DataDesignerColumnType.SEED_DATASET.value
 
 
 class ExpressionColumnStatistics(GeneralColumnStatistics):
+    """Container for statistics on expression-based derived columns.
+
+    Inherits general statistics and stores statistics computed from columns that are derived
+    from columns that are derived from Jinja2 expressions referencing other column values.
+
+    Attributes:
+        column_type: Discriminator field, always "expression" for this statistics type.
+    """
+
     column_type: Literal[DataDesignerColumnType.EXPRESSION.value] = DataDesignerColumnType.EXPRESSION.value
 
 
 class ValidationColumnStatistics(GeneralColumnStatistics):
-    num_valid_records: Union[int, MissingValue]
+    """Container for statistics on validation result columns.
+
+    Inherits general statistics plus validation-specific metrics including the count and
+    percentage of records that passed validation. Stores results from validation logic
+    (Python, SQL, or remote) executed against target columns.
+
+    Attributes:
+        num_valid_records: Number of records that passed validation.
+        column_type: Discriminator field, always "validation" for this statistics type.
+    """
+
+    num_valid_records: int | MissingValue
     column_type: Literal[DataDesignerColumnType.VALIDATION.value] = DataDesignerColumnType.VALIDATION.value
 
     @field_validator("num_valid_records", mode="before")
-    def code_validation_column_ensure_python_integers(cls, v: Union[int, MissingValue]) -> Union[int, MissingValue]:
+    def code_validation_column_ensure_python_integers(cls, v: int | MissingValue) -> int | MissingValue:
         return v if isinstance(v, MissingValue) else prepare_number_for_reporting(v, int)
 
     @property
-    def percent_valid(self) -> Union[float, MissingValue]:
+    def percent_valid(self) -> float | MissingValue:
         return (
             self.num_valid_records
             if self._is_missing_value(self.num_valid_records)
@@ -177,7 +294,16 @@ class ValidationColumnStatistics(GeneralColumnStatistics):
 
 
 class CategoricalHistogramData(BaseModel):
-    categories: list[Union[float, int, str]]
+    """Container for categorical distribution histogram data.
+
+    Stores the computed frequency distribution of categorical values.
+
+    Attributes:
+        categories: List of unique category values that appear in the data.
+        counts: List of occurrence counts for each category.
+    """
+
+    categories: list[float | int | str]
     counts: list[int]
 
     @model_validator(mode="after")
@@ -194,12 +320,20 @@ class CategoricalHistogramData(BaseModel):
 
 
 class CategoricalDistribution(BaseModel):
-    most_common_value: Union[str, int]
-    least_common_value: Union[str, int]
+    """Container for computed categorical distribution statistics.
+
+    Attributes:
+        most_common_value: The category value that appears most frequently in the data.
+        least_common_value: The category value that appears least frequently in the data.
+        histogram: Complete frequency distribution showing all categories and their counts.
+    """
+
+    most_common_value: str | int
+    least_common_value: str | int
     histogram: CategoricalHistogramData
 
     @field_validator("most_common_value", "least_common_value", mode="before")
-    def ensure_python_types(cls, v: Union[str, int]) -> Union[str, int]:
+    def ensure_python_types(cls, v: str | int) -> str | int:
         return str(v) if not is_int(v) else prepare_number_for_reporting(v, int)
 
     @classmethod
@@ -213,14 +347,24 @@ class CategoricalDistribution(BaseModel):
 
 
 class NumericalDistribution(BaseModel):
-    min: Union[float, int]
-    max: Union[float, int]
+    """Container for computed numerical distribution statistics.
+
+    Attributes:
+        min: Minimum value in the distribution.
+        max: Maximum value in the distribution.
+        mean: Arithmetic mean (average) of all values.
+        stddev: Standard deviation measuring the spread of values around the mean.
+        median: Median value of the distribution.
+    """
+
+    min: float | int
+    max: float | int
     mean: float
     stddev: float
     median: float
 
     @field_validator("min", "max", "mean", "stddev", "median", mode="before")
-    def ensure_python_types(cls, v: Union[float, int]) -> Union[float, int]:
+    def ensure_python_types(cls, v: float | int) -> float | int:
         return prepare_number_for_reporting(v, int if is_int(v) else float)
 
     @classmethod
@@ -234,17 +378,17 @@ class NumericalDistribution(BaseModel):
         )
 
 
-ColumnStatisticsT: TypeAlias = Union[
-    GeneralColumnStatistics,
-    LLMTextColumnStatistics,
-    LLMCodeColumnStatistics,
-    LLMStructuredColumnStatistics,
-    LLMJudgedColumnStatistics,
-    SamplerColumnStatistics,
-    SeedDatasetColumnStatistics,
-    ValidationColumnStatistics,
-    ExpressionColumnStatistics,
-]
+ColumnStatisticsT: TypeAlias = (
+    GeneralColumnStatistics
+    | LLMTextColumnStatistics
+    | LLMCodeColumnStatistics
+    | LLMStructuredColumnStatistics
+    | LLMJudgedColumnStatistics
+    | SamplerColumnStatistics
+    | SeedDatasetColumnStatistics
+    | ValidationColumnStatistics
+    | ExpressionColumnStatistics
+)
 
 
 DEFAULT_COLUMN_STATISTICS_MAP = {