data-designer 0.1.0__py3-none-any.whl

data_designer/engine/analysis/column_profilers/judge_score_profiler.py

@@ -0,0 +1,160 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import random
+from typing import Union
+
+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.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
+from data_designer.engine.resources.resource_provider import ResourceType
+
+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.",
+            required_resources=[ResourceType.MODEL_REGISTRY],
+            applicable_column_types=[DataDesignerColumnType.LLM_JUDGE],
+        )
+
+    def get_model(self, model_alias: str) -> ModelFacade:
+        return self.resource_provider.model_registry.get_model(model_alias=model_alias)
+
+    def profile(self, column_config_with_df: ColumnConfigWithDataFrame) -> JudgeScoreProfilerResults:
+        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}'"
+        )
+
+        score_summaries = {}
+        score_distributions = extract_judge_score_distributions(column_config, df)
+
+        if self.config.summary_score_sample_size is None or isinstance(score_distributions, MissingValue):
+            return JudgeScoreProfilerResults(
+                summaries={},
+                column_name=column_config.name,
+                score_distributions=score_distributions,
+            )
+
+        for score in column_config.scores:
+            score_name = score.name.lower()
+            logger.info(f"{random.choice(['👩‍⚖️', '👨‍⚖️'])} Summarizing LLM-as-judge score: '{score_name}'")
+            score_sample = sample_scores_and_reasoning(
+                scores=score_distributions.scores[score_name],
+                reasoning=score_distributions.reasoning[score_name],
+                num_samples=self.config.summary_score_sample_size,
+            )
+
+            score_summaries[score_name] = self._summarize_score_sample(
+                name=score_name,
+                sample=score_sample,
+                histogram=score_distributions.histograms[score_name],
+                distribution=score_distributions.distributions[score_name],
+                distribution_type=score_distributions.distribution_types[score_name],
+            )
+
+        return JudgeScoreProfilerResults(
+            column_name=column_config.name,
+            summaries=score_summaries,
+            score_distributions=score_distributions,
+        )
+
+    def _summarize_score_sample(
+        self,
+        name: str,
+        sample: list[JudgeScoreSample],
+        histogram: CategoricalHistogramData,
+        distribution: Union[CategoricalDistribution, NumericalDistribution, MissingValue],
+        distribution_type: ColumnDistributionType,
+    ) -> JudgeScoreSummary:
+        if isinstance(distribution, MissingValue) or not sample:
+            return JudgeScoreSummary(
+                score_name=name,
+                summary="No judge score information available to summarize.",
+                score_samples=sample,
+            )
+
+        category_info = []
+        total_count = sum(histogram.counts)
+        for cat, count in zip(histogram.categories, histogram.counts):
+            percentage = (count / total_count) * 100
+            category_info.append(f"{cat}: {count} records ({percentage:.1f}%)")
+
+        distribution_context = f"Score distribution - {', '.join(category_info)}, "
+        if distribution_type == ColumnDistributionType.CATEGORICAL:
+            distribution_context += f"Most common value: {distribution.most_common_value}. "
+        if distribution_type == ColumnDistributionType.NUMERICAL:
+            distribution_context += f"Mean score: {distribution.mean:.2f}. "
+
+        logger.info(f"  |-- number of score samples: {len(sample)}")
+        logger.info(f"  |-- {distribution_context.lower()}")
+
+        combined_reasoning = "\n".join([r.reasoning for r in sample])
+        prompt = (
+            f"Based on the following evaluator reasoning for the '{name}' criterion, "
+            "provide a concise summary that captures both the strengths and areas for improvement mentioned. "
+            "Be specific about what worked well and what needs improvement.\n\n"
+            f"Overall distribution of scores: {distribution_context}"
+            f"\nA sample of reasoning:\n{combined_reasoning}\n\n"
+            "Do not include any titles like `Summary` or `Summary:`. "
+            "Do not wrap the summary in quotation marks. "
+            "YOU WILL PRODUCE LESS THAN 75 WORDS in a readable sentence format. "
+            "No need to use bullets or headers. Write naturally."
+        )
+
+        system_prompt = (
+            "You are an expert at distilling complex feedback into concise summaries. "
+            "Focus on specificity and balance, incorporating both the distribution context and individual reasoning examples."
+        )
+
+        try:
+            model = self.get_model(self.config.model_alias)
+            recipe = TextResponseRecipe()
+            summary, _ = model.generate(
+                prompt=recipe.apply_recipe_to_user_prompt(prompt),
+                system_prompt=recipe.apply_recipe_to_system_prompt(system_prompt),
+                parser=recipe.parse,
+            )
+            return JudgeScoreSummary(
+                score_name=name,
+                summary=summary.strip(),
+                score_samples=sample,
+            )
+        except Exception as e:
+            return JudgeScoreSummary(
+                score_name=name,
+                summary=f"Score summarization failed: {e}",
+                score_samples=sample,
+            )

data_designer/engine/analysis/column_profilers/registry.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+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
+from data_designer.engine.analysis.column_profilers.judge_score_profiler import (
+    JudgeScoreProfiler,
+    JudgeScoreProfilerConfig,
+)
+from data_designer.engine.registry.base import TaskRegistry
+
+
+class ColumnProfilerRegistry(TaskRegistry[ColumnProfilerType, ColumnProfiler, ConfigBase]): ...
+
+
+def create_default_column_profiler_registry() -> ColumnProfilerRegistry:
+    registry = ColumnProfilerRegistry()
+    registry.register(ColumnProfilerType.JUDGE_SCORE, JudgeScoreProfiler, JudgeScoreProfilerConfig, False)
+    return registry

data_designer/engine/analysis/column_statistics.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Type, TypeAlias, Union
+
+import pandas as pd
+from pydantic import BaseModel
+from typing_extensions import Self
+
+from data_designer.config.analysis.column_statistics import (
+    DEFAULT_COLUMN_STATISTICS_MAP,
+    ColumnStatisticsT,
+    GeneralColumnStatistics,
+)
+from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType
+from data_designer.config.sampler_params import SamplerType, is_numerical_sampler_type
+from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame
+from data_designer.engine.analysis.utils.column_statistics_calculations import (
+    ColumnDistributionType,
+    calculate_column_distribution,
+    calculate_general_column_info,
+    calculate_token_stats,
+    calculate_validation_column_info,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralColumnStatisticsCalculator(BaseModel):
+    column_config_with_df: ColumnConfigWithDataFrame
+
+    @property
+    def column_config(self) -> ColumnConfigT:
+        return self.column_config_with_df.column_config
+
+    @property
+    def df(self) -> pd.DataFrame:
+        return self.column_config_with_df.df
+
+    @property
+    def column_statistics_type(self) -> Type[ColumnStatisticsT]:
+        return DEFAULT_COLUMN_STATISTICS_MAP.get(self.column_config.column_type, GeneralColumnStatistics)
+
+    def calculate(self) -> Self:
+        """Calculate all the column statistics fields for the given column configuration and dataset profiler.
+
+        This method dynamically collects all class methods prefixed with 'calculate_' and invokes them to
+        compute various column statistics, aggregating their results into a single statistics object.
+        """
+        calculate_methods = [
+            name for name in dir(self) if name.startswith("calculate_") and callable(getattr(self, name))
+        ]
+        return self.column_statistics_type(
+            column_name=self.column_config.name,
+            **{k: v for name in calculate_methods for k, v in getattr(self, name)().items()},
+        )
+
+    def calculate_general_column_info(self) -> dict[str, Any]:
+        return calculate_general_column_info(self.column_config, self.df)
+
+    def __repr__(self) -> str:
+        params = []
+        for field, value in self.model_dump(mode="json").items():
+            params.append(f"    {field}: {value}")
+        params_str = "\n".join(params)
+        return f"{self.__class__.__name__}(\n{params_str}\n)"
+
+
+class LLMTextColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
+    def calculate_token_stats(self) -> dict[str, Any]:
+        return calculate_token_stats(self.column_config, self.df)
+
+
+class LLMCodeColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class LLMStructuredColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class LLMJudgedColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...
+
+
+class SamplerColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
+    def calculate_sampler_distribution(self) -> dict[str, Any]:
+        make_dist, dist_type = False, ColumnDistributionType.OTHER
+        if self.column_config.sampler_type in [SamplerType.CATEGORY, SamplerType.SUBCATEGORY]:
+            make_dist, dist_type = True, ColumnDistributionType.CATEGORICAL
+        elif is_numerical_sampler_type(self.column_config.sampler_type):
+            make_dist, dist_type = True, ColumnDistributionType.NUMERICAL
+        return (
+            {
+                "sampler_type": SamplerType(self.column_config.sampler_type),
+                **calculate_column_distribution(self.column_config, self.df, dist_type),
+            }
+            if make_dist
+            else {
+                "sampler_type": SamplerType(self.column_config.sampler_type),
+                "distribution_type": dist_type,
+                "distribution": None,
+            }
+        )
+
+
+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)
+
+
+class ExpressionColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
+
+
+ColumnStatisticsCalculatorT: TypeAlias = Union[
+    ExpressionColumnStatisticsCalculator,
+    ValidationColumnStatisticsCalculator,
+    GeneralColumnStatisticsCalculator,
+    LLMCodeColumnStatisticsCalculator,
+    LLMJudgedColumnStatisticsCalculator,
+    LLMStructuredColumnStatisticsCalculator,
+    LLMTextColumnStatisticsCalculator,
+    SamplerColumnStatisticsCalculator,
+    SeedDatasetColumnStatisticsCalculator,
+]
+DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP = {
+    DataDesignerColumnType.EXPRESSION: ExpressionColumnStatisticsCalculator,
+    DataDesignerColumnType.VALIDATION: ValidationColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_CODE: LLMCodeColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_JUDGE: LLMJudgedColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnStatisticsCalculator,
+    DataDesignerColumnType.LLM_TEXT: LLMTextColumnStatisticsCalculator,
+    DataDesignerColumnType.SAMPLER: SamplerColumnStatisticsCalculator,
+    DataDesignerColumnType.SEED_DATASET: SeedDatasetColumnStatisticsCalculator,
+}
+
+
+def get_column_statistics_calculator(column_type: DataDesignerColumnType) -> ColumnStatisticsCalculatorT:
+    return DEFAULT_COLUMN_STATISTICS_CALCULATOR_MAP.get(column_type, GeneralColumnStatisticsCalculator)

data_designer/engine/analysis/dataset_profiler.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+from functools import cached_property
+import logging
+
+import pandas as pd
+from pydantic import Field, field_validator
+
+from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
+from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.column_types import (
+    COLUMN_TYPE_EMOJI_MAP,
+    ColumnConfigT,
+)
+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.registry.data_designer_registry import DataDesignerRegistry
+from data_designer.engine.resources.resource_provider import ResourceProvider
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetProfilerConfig(ConfigBase):
+    column_configs: Sequence[DatasetBuilderColumnConfigT] = Field(..., min_length=1)
+    column_profiler_configs: Sequence[ColumnProfilerConfigT] | None = None
+
+    @field_validator("column_configs")
+    def flatten_and_validate_column_configs(cls, v: list[DatasetBuilderColumnConfigT]) -> list[ColumnConfigT]:
+        column_configs = []
+        for config in v:
+            if isinstance(config, SingleColumnConfig) and not config.drop:
+                column_configs.append(config)
+            elif isinstance(config, MultiColumnConfig):
+                column_configs.extend([c for c in config.columns if not c.drop])
+        if len(column_configs) == 0:
+            raise DatasetProfilerConfigurationError("All columns were dropped!")
+        return column_configs
+
+
+class DataDesignerDatasetProfiler:
+    def __init__(self, config: DatasetProfilerConfig, resource_provider: ResourceProvider):
+        self.config = config
+        self.resource_provider = resource_provider
+        self._validate_column_profiler_configs()
+
+    @cached_property
+    def column_names_from_configs(self) -> list[str]:
+        return [c.name for c in self.config.column_configs]
+
+    @cached_property
+    def registry(self) -> DataDesignerRegistry:
+        return DataDesignerRegistry()
+
+    def profile_dataset(
+        self,
+        target_num_records: int,
+        dataset: pd.DataFrame,
+    ) -> DatasetProfilerResults:
+        logger.info("📐 Measuring dataset column statistics:")
+
+        self._validate_schema_consistency(list(dataset.columns))
+
+        column_statistics = []
+        for c in self.config.column_configs:
+            logger.info(f"  |-- {COLUMN_TYPE_EMOJI_MAP[c.column_type]} column: '{c.name}'")
+            column_statistics.append(
+                get_column_statistics_calculator(c.column_type)(
+                    column_config_with_df=ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                ).calculate()
+            )
+
+        column_profiles = []
+        for profiler_config in self.config.column_profiler_configs or []:
+            profiler = self._create_column_profiler(profiler_config)
+            applicable_column_types = profiler.metadata().applicable_column_types
+            for c in self.config.column_configs:
+                if c.column_type in applicable_column_types:
+                    params = ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                    column_profiles.append(profiler.profile(params))
+            if len(column_profiles) == 0:
+                logger.warning(
+                    f"⚠️ No applicable column types found for the '{profiler.metadata().name}' profiler. "
+                    f"This profiler is applicable to the following column types: {applicable_column_types}"
+                )
+
+        return DatasetProfilerResults(
+            num_records=len(dataset),
+            target_num_records=target_num_records,
+            side_effect_column_names=list(set(dataset.columns) - set(self.column_names_from_configs)),
+            column_statistics=column_statistics,
+            column_profiles=column_profiles if column_profiles else None,
+        )
+
+    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
+        )
+
+    def _validate_column_profiler_configs(self) -> None:
+        if self.config.column_profiler_configs:
+            if self.resource_provider.model_registry is None:
+                raise DatasetProfilerConfigurationError("Model registry is required for column profiler configs")
+            self._validate_model_configs()
+
+    def _validate_model_configs(self) -> None:
+        aliases = [alias for alias in self.resource_provider.model_registry.model_configs.keys()]
+        for column_config in self.config.column_configs:
+            if hasattr(column_config, "model_alias") and column_config.model_alias not in aliases:
+                raise DatasetProfilerConfigurationError(
+                    f"Model config '{column_config.model_alias}' not found in model configs"
+                )
+
+    def _validate_schema_consistency(self, dataset_column_names: list[str]) -> None:
+        for column_name in self.column_names_from_configs:
+            if column_name not in dataset_column_names:
+                raise DatasetProfilerConfigurationError(f"Column '{column_name}' not found in dataset")

data_designer/engine/analysis/errors.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.errors import DataDesignerError
+
+
+class DatasetProfilerConfigurationError(DataDesignerError): ...

data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from numbers import Number
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import tiktoken
+
+from data_designer.config.analysis.column_statistics import (
+    CategoricalDistribution,
+    ColumnDistributionType,
+    MissingValue,
+    NumericalDistribution,
+)
+from data_designer.config.column_configs import (
+    LLMTextColumnConfig,
+    SingleColumnConfig,
+    ValidationColumnConfig,
+)
+from data_designer.engine.column_generators.generators.llm_generators import (
+    PromptType,
+    RecordBasedPromptRenderer,
+    create_response_recipe,
+)
+
+RANDOM_SEED = 42
+MAX_PROMPT_SAMPLE_SIZE = 1000
+TOKENIZER = tiktoken.get_encoding("cl100k_base")
+WARNING_PREFIX = "⚠️ Error during column profile calculation: "
+TEXT_FIELD_AVG_SPACE_COUNT_THRESHOLD = 0.1
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_column_distribution(
+    column_config: SingleColumnConfig, df: pd.DataFrame, distribution_type: ColumnDistributionType
+) -> dict[str, CategoricalDistribution | NumericalDistribution | MissingValue | None]:
+    distribution_type = ColumnDistributionType(distribution_type)
+    try:
+        if distribution_type == ColumnDistributionType.CATEGORICAL:
+            return {
+                "distribution_type": ColumnDistributionType.CATEGORICAL,
+                "distribution": CategoricalDistribution.from_series(df[column_config.name]),
+            }
+
+        if distribution_type == ColumnDistributionType.NUMERICAL:
+            return {
+                "distribution_type": ColumnDistributionType.NUMERICAL,
+                "distribution": NumericalDistribution.from_series(df[column_config.name]),
+            }
+    except Exception as e:
+        logger.warning(f"{WARNING_PREFIX} failed to calculate column distribution for '{column_config.name}' {e}")
+        return {
+            "distribution_type": ColumnDistributionType.UNKNOWN,
+            "distribution": MissingValue.CALCULATION_FAILED,
+        }
+
+
+def calculate_general_column_info(column_config: SingleColumnConfig, df: pd.DataFrame) -> dict[str, Any]:
+    try:
+        _df = pd.DataFrame(df[column_config.name].apply(ensure_hashable))
+        return {
+            "pyarrow_dtype": str(df[column_config.name].dtype.pyarrow_dtype),
+            "simple_dtype": convert_pyarrow_dtype_to_simple_dtype(df[column_config.name].dtype.pyarrow_dtype),
+            "num_records": len(_df[column_config.name]),
+            "num_null": _df[column_config.name].isnull().sum(),
+            "num_unique": _df[column_config.name].nunique(),
+        }
+    except Exception as e:
+        logger.warning(f"{WARNING_PREFIX} failed to calculate general column info for '{column_config.name}': {e}")
+        return {
+            "pyarrow_dtype": MissingValue.CALCULATION_FAILED,
+            "simple_dtype": MissingValue.CALCULATION_FAILED,
+            "num_records": MissingValue.CALCULATION_FAILED,
+            "num_null": MissingValue.CALCULATION_FAILED,
+            "num_unique": MissingValue.CALCULATION_FAILED,
+        }
+
+
+def calculate_prompt_token_stats(
+    column_config: LLMTextColumnConfig, df: pd.DataFrame
+) -> dict[str, float | MissingValue]:
+    try:
+        num_tokens = []
+        num_samples = min(MAX_PROMPT_SAMPLE_SIZE, len(df))
+        renderer = RecordBasedPromptRenderer(response_recipe=create_response_recipe(column_config))
+        for record in df.sample(num_samples, random_state=RANDOM_SEED).to_dict(orient="records"):
+            system_prompt = renderer.render(
+                prompt_template=column_config.system_prompt, record=record, prompt_type=PromptType.SYSTEM_PROMPT
+            )
+            prompt = renderer.render(
+                prompt_template=column_config.prompt, record=record, prompt_type=PromptType.USER_PROMPT
+            )
+            concatenated_prompt = str(system_prompt + "\n\n" + prompt)
+            num_tokens.append(len(TOKENIZER.encode(concatenated_prompt, disallowed_special=())))
+    except Exception as e:
+        logger.warning(
+            f"{WARNING_PREFIX} failed to calculate prompt token stats for column {column_config.name!r}: {e}"
+        )
+        return {
+            "prompt_tokens_mean": MissingValue.CALCULATION_FAILED,
+            "prompt_tokens_median": MissingValue.CALCULATION_FAILED,
+            "prompt_tokens_stddev": MissingValue.CALCULATION_FAILED,
+        }
+    return {
+        "prompt_tokens_mean": np.mean(num_tokens),
+        "prompt_tokens_median": np.median(num_tokens),
+        "prompt_tokens_stddev": np.std(num_tokens),
+    }
+
+
+def calculate_completion_token_stats(
+    column_config: LLMTextColumnConfig, df: pd.DataFrame
+) -> dict[str, float | MissingValue]:
+    try:
+        tokens_per_record = df[column_config.name].apply(
+            lambda value: len(TOKENIZER.encode(str(value), disallowed_special=()))
+        )
+        return {
+            "completion_tokens_mean": tokens_per_record.mean(),
+            "completion_tokens_median": tokens_per_record.median(),
+            "completion_tokens_stddev": tokens_per_record.std(),
+        }
+    except Exception as e:
+        logger.warning(
+            f"{WARNING_PREFIX} failed to calculate completion token stats for column {column_config.name}: {e}"
+        )
+        return {
+            "completion_tokens_mean": MissingValue.CALCULATION_FAILED,
+            "completion_tokens_median": MissingValue.CALCULATION_FAILED,
+            "completion_tokens_stddev": MissingValue.CALCULATION_FAILED,
+        }
+
+
+def calculate_token_stats(column_config: LLMTextColumnConfig, df: pd.DataFrame) -> dict[str, float | MissingValue]:
+    return {
+        **calculate_prompt_token_stats(column_config, df),
+        **calculate_completion_token_stats(column_config, df),
+    }
+
+
+def calculate_validation_column_info(column_config: ValidationColumnConfig, df: pd.DataFrame) -> dict[str, Any]:
+    try:
+        return {"num_valid_records": df[column_config.name].apply(lambda x: ensure_boolean(x["is_valid"])).sum()}
+    except Exception as e:
+        logger.warning(
+            f"{WARNING_PREFIX} failed to calculate code validation column info for column {column_config.name}: {e}"
+        )
+        return {"num_valid_records": MissingValue.CALCULATION_FAILED}
+
+
+def convert_pyarrow_dtype_to_simple_dtype(pyarrow_dtype: pa.DataType) -> str:
+    if isinstance(pyarrow_dtype, pa.ListType):
+        return f"list[{convert_pyarrow_dtype_to_simple_dtype(pyarrow_dtype.value_type)}]"
+    if isinstance(pyarrow_dtype, pa.StructType):
+        return "dict"
+    pyarrow_dtype_str = str(pyarrow_dtype)
+    if "int" in pyarrow_dtype_str:
+        return "int"
+    if "double" in pyarrow_dtype_str:
+        return "float"
+    if "float" in pyarrow_dtype_str:
+        return "float"
+    if "string" in pyarrow_dtype_str:
+        return "string"
+    if "timestamp" in pyarrow_dtype_str:
+        return "timestamp"
+    if "time" in pyarrow_dtype_str:
+        return "time"
+    if "date" in pyarrow_dtype_str:
+        return "date"
+    return pyarrow_dtype_str
+
+
+def ensure_hashable(x: Any) -> str:
+    """
+    Makes a best effort turn known unhashable types to a hashable
+    string representation that preserves both structure and values.
+    """
+    if isinstance(x, (Number, bool)) or x is None:
+        return x
+
+    if isinstance(x, dict):
+        # Sort by keys and convert key-value pairs to tuples
+        return str(sorted([(str(k), ensure_hashable(v)) for k, v in x.items()]))
+
+    if isinstance(x, (list, tuple, set, np.ndarray)):
+        # Recursively make all elements hashable
+        return str(sorted([ensure_hashable(e) for e in x]))
+
+    return str(x)
+
+
+def ensure_boolean(v: bool | str | int | None) -> bool:
+    if isinstance(v, (bool, np.bool_)):
+        return bool(v)
+    if isinstance(v, (int, float, np.integer, np.floating)) and v in [0, 1, 0.0, 1.0]:
+        return bool(v)
+    if isinstance(v, (str, np.str_)) and v.lower() in ["true", "false"]:
+        return v.lower() == "true"
+    if v is None:
+        return False
+    raise ValueError(f"Invalid boolean value: {v}")