diff --git a/packages/data-designer-config/src/data_designer/config/analysis/column_statistics.py b/packages/data-designer-config/src/data_designer/config/analysis/column_statistics.py index 99e3a388d..a55571164 100644 --- a/packages/data-designer-config/src/data_designer/config/analysis/column_statistics.py +++ b/packages/data-designer-config/src/data_designer/config/analysis/column_statistics.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, ConfigDict, create_model, field_validator, model_validator +import data_designer.lazy_heavy_imports as lazy from data_designer.config.column_types import DataDesignerColumnType from data_designer.config.sampler_params import SamplerType from data_designer.config.utils.constants import EPSILON @@ -347,6 +348,15 @@ def from_series(cls, series: pd.Series) -> Self: ) +def _undefined_as_nan(value: Any) -> Any: + """Report a statistic that pandas leaves undefined as NaN. + + The sample standard deviation of a single value is NaN for NumPy-backed series but `pd.NA` for + nullable and Arrow-backed ones, and `pd.NA` cannot be converted to a float. + """ + return float("nan") if value is lazy.pd.NA else value + + class NumericalDistribution(BaseModel): """Container for computed numerical distribution statistics. @@ -374,7 +384,7 @@ def from_series(cls, series: pd.Series) -> Self: min=series.min(skipna=True), max=series.max(skipna=True), mean=series.mean(skipna=True), - stddev=series.std(skipna=True), + stddev=_undefined_as_nan(series.std(skipna=True)), median=series.median(skipna=True), ) diff --git a/packages/data-designer-config/tests/config/analysis/test_column_statistics.py b/packages/data-designer-config/tests/config/analysis/test_column_statistics.py index 7f709585e..b063f7aac 100644 --- a/packages/data-designer-config/tests/config/analysis/test_column_statistics.py +++ b/packages/data-designer-config/tests/config/analysis/test_column_statistics.py @@ -300,3 +300,31 @@ def test_numerical_distribution(): assert numerical_distribution.mean == 3.0 assert numerical_distribution.stddev == 1.58 assert numerical_distribution.median == 3.0 + + +@pytest.mark.parametrize( + "series", + [ + pytest.param(lazy.pd.Series([4.5]), id="numpy-float"), + pytest.param(lazy.pd.Series([4.5], dtype="Float64"), id="nullable-float"), + pytest.param(lazy.pd.Series([4.5], dtype=lazy.pd.ArrowDtype(lazy.pa.float64())), id="arrow-float"), + ], +) +def test_numerical_distribution_from_single_value_keeps_the_defined_statistics(series): + # The sample standard deviation of one value is undefined: NaN for NumPy-backed series but pd.NA for + # nullable and Arrow-backed ones. Only that field is undefined; the others must survive. + distribution = NumericalDistribution.from_series(series) + + assert distribution.min == 4.5 + assert distribution.max == 4.5 + assert distribution.mean == 4.5 + assert distribution.median == 4.5 + assert lazy.np.isnan(distribution.stddev) + + +def test_numerical_distribution_from_arrow_series_with_several_values_has_a_defined_stddev(): + series = lazy.pd.Series([1, 2, 3, 4, 5], dtype=lazy.pd.ArrowDtype(lazy.pa.int64())) + + distribution = NumericalDistribution.from_series(series) + + assert distribution.stddev == pytest.approx(1.58, abs=0.01) diff --git a/packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py b/packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py index 2c85e5d25..d592e946c 100644 --- a/packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py +++ b/packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py @@ -61,7 +61,11 @@ def calculate_column_distribution( def calculate_general_column_info(column_name: str, df: pd.DataFrame) -> dict[str, Any]: try: - _df = lazy.pd.DataFrame(df[column_name].apply(ensure_hashable)) + column = df[column_name] + # Count missing values on the original column and leave them out of the unique count: `ensure_hashable` + # turns `pd.NA` into the string "", which would otherwise read as a distinct non-null value. + num_null = column.isna().sum() + num_unique = column.dropna().apply(ensure_hashable).nunique() if has_pyarrow_backend(df): pyarrow_dtype = str(df[column_name].dtype.pyarrow_dtype) @@ -78,9 +82,9 @@ def calculate_general_column_info(column_name: str, df: pd.DataFrame) -> dict[st return { "pyarrow_dtype": pyarrow_dtype, "simple_dtype": simple_dtype, - "num_records": len(_df[column_name]), - "num_null": _df[column_name].isnull().sum(), - "num_unique": _df[column_name].nunique(), + "num_records": len(column), + "num_null": num_null, + "num_unique": num_unique, } except Exception as e: logger.warning(f"{WARNING_PREFIX} failed to calculate general column info for '{column_name}': {e}") @@ -133,7 +137,9 @@ def calculate_output_token_stats( column_config: LLMTextColumnConfig, df: pd.DataFrame ) -> dict[str, float | MissingValue]: try: - tokens_per_record = df[column_config.name].apply(lambda value: count_text_tokens(str(value))) + # Arrow-backed columns keep their backing through `apply` (pandas 3.1+), and the standard deviation of + # one value is then `pd.NA` instead of NaN. Plain float64 keeps the result a number on every pandas. + tokens_per_record = df[column_config.name].apply(lambda value: count_text_tokens(str(value))).astype("float64") return { "output_tokens_mean": tokens_per_record.mean(), "output_tokens_median": tokens_per_record.median(), diff --git a/packages/data-designer-engine/tests/engine/analysis/utils/test_column_statistics_calculations.py b/packages/data-designer-engine/tests/engine/analysis/utils/test_column_statistics_calculations.py index 23bf6cb1b..609ffc64d 100644 --- a/packages/data-designer-engine/tests/engine/analysis/utils/test_column_statistics_calculations.py +++ b/packages/data-designer-engine/tests/engine/analysis/utils/test_column_statistics_calculations.py @@ -13,6 +13,7 @@ CategoricalDistribution, CategoricalHistogramData, ColumnDistributionType, + LLMTextColumnStatistics, MissingValue, NumericalDistribution, ) @@ -392,3 +393,81 @@ def test_calculate_general_column_info_edge_cases(): assert result["num_records"] == 0 assert result["num_null"] == 0 assert result["simple_dtype"] == MissingValue.CALCULATION_FAILED + + +def _arrow_structured_column(rows: list[dict[str, str]]) -> lazy.pd.DataFrame: + """A column shaped like structured LLM output: Arrow-backed, as it is after generation.""" + dtype = lazy.pd.ArrowDtype(lazy.pa.struct([("answer", lazy.pa.string())])) + return lazy.pd.DataFrame({"test_column": lazy.pd.Series(rows, dtype=dtype)}) + + +def test_calculate_output_token_stats_single_record_arrow_backed_column(stub_column_config): + # The standard deviation of one value is undefined. On pandas 3.1+ it comes back as pd.NA for Arrow-backed + # data, which the statistics model cannot convert to a float. + df = _arrow_structured_column([{"answer": "hello world"}]) + + result = calculate_output_token_stats(stub_column_config, df) + + assert result["output_tokens_mean"] > 0 + assert result["output_tokens_median"] == result["output_tokens_mean"] + assert isinstance(result["output_tokens_stddev"], float) + assert lazy.np.isnan(result["output_tokens_stddev"]) + # The numbers must also be accepted by the statistics model, which is where pd.NA used to fail. + statistics = LLMTextColumnStatistics( + column_name="test_column", + num_records=1, + num_null=0, + num_unique=1, + pyarrow_dtype="struct", + simple_dtype="dict", + input_tokens_mean=1.0, + input_tokens_median=1.0, + input_tokens_stddev=0.0, + **result, + ) + assert lazy.np.isnan(statistics.output_tokens_stddev) + + +def test_calculate_output_token_stats_multi_record_arrow_backed_column(stub_column_config): + df = _arrow_structured_column([{"answer": "short"}, {"answer": "a much longer answer than the first"}]) + + result = calculate_output_token_stats(stub_column_config, df) + + assert isinstance(result["output_tokens_stddev"], float) + assert result["output_tokens_stddev"] > 0 + + +@pytest.mark.parametrize( + "series, expected_null, expected_unique", + [ + pytest.param(lazy.pd.Series([7, None], dtype="Int64"), 1, 1, id="nullable-int"), + pytest.param(lazy.pd.Series([1.5, None, 1.5], dtype="Float64"), 1, 1, id="nullable-float"), + pytest.param(lazy.pd.Series(["a", None, "b"], dtype="string"), 1, 2, id="nullable-string"), + pytest.param(lazy.pd.Series(["a", None], dtype="string[pyarrow]"), 1, 1, id="arrow-string"), + pytest.param(lazy.pd.Series([True, None, False], dtype="boolean"), 1, 2, id="nullable-boolean"), + pytest.param(lazy.pd.Series([None, None], dtype="string[pyarrow]"), 2, 0, id="all-null-string"), + pytest.param(lazy.pd.Series([None, None], dtype="Int64"), 2, 0, id="all-null-int"), + pytest.param(lazy.pd.Series([None, "x", None], dtype=object), 2, 1, id="object-none"), + pytest.param(lazy.pd.Series([3, 4, 3], dtype="Int64"), 0, 2, id="no-nulls"), + ], +) +def test_calculate_general_column_info_counts_missing_values_from_the_original_column( + series, expected_null, expected_unique +): + # A missing value must not be counted as a distinct "" string. + result = calculate_general_column_info("value", lazy.pd.DataFrame({"value": series})) + + assert result["num_records"] == len(series) + assert result["num_null"] == expected_null + assert result["num_unique"] == expected_unique + + +def test_calculate_general_column_info_still_hashes_nested_values_for_uniqueness(): + df = lazy.pd.DataFrame( + {"value": [{"a": 1}, {"a": 1}, {"a": 2}, [1, 2], [2, 1], None]}, + ) + + result = calculate_general_column_info("value", df) + + assert result["num_null"] == 1 + assert result["num_unique"] == 3 diff --git a/packages/data-designer/tests/interface/test_profiling_statistics.py b/packages/data-designer/tests/interface/test_profiling_statistics.py new file mode 100644 index 000000000..97697bd98 --- /dev/null +++ b/packages/data-designer/tests/interface/test_profiling_statistics.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +import pytest + +from data_designer.config.analysis.column_statistics import ColumnDistributionType, NumericalDistribution +from data_designer.config.column_configs import ExpressionColumnConfig, SamplerColumnConfig +from data_designer.config.config_builder import DataDesignerConfigBuilder +from data_designer.config.models import ModelProvider +from data_designer.config.run_config import ResumeMode, RunConfig +from data_designer.config.sampler_params import SamplerType, UniformSamplerParams +from data_designer.config.seed_source_dataframe import DataFrameSeedSource +from data_designer.engine.secret_resolver import PlaintextResolver +from data_designer.interface.data_designer import DataDesigner + + +@pytest.fixture +def designer(tmp_path: Path, stub_model_providers: list[ModelProvider]) -> DataDesigner: + managed = tmp_path / "managed" + managed.mkdir() + designer = DataDesigner( + artifact_path=tmp_path / "artifacts", + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=managed, + ) + designer.set_run_config(RunConfig(otel_metrics_port=None)) + return designer + + +def _statistics(analysis, column_name: str): + return next(stat for stat in analysis.column_statistics if stat.column_name == column_name) + + +@pytest.mark.parametrize("run", ["preview", "create"]) +@pytest.mark.parametrize( + "values", + [ + pytest.param(pd.Series([7, None], dtype="Int64"), id="nullable-int"), + pytest.param(pd.Series(["a", None], dtype="string"), id="nullable-string"), + pytest.param(pd.Series([True, None], dtype="boolean"), id="nullable-boolean"), + ], +) +def test_profile_counts_a_missing_seed_value_as_null_not_as_a_distinct_value( + designer: DataDesigner, values: pd.Series, run: str +) -> None: + config = DataDesignerConfigBuilder().with_seed_dataset(DataFrameSeedSource(df=pd.DataFrame({"value": values}))) + config.add_column( + ExpressionColumnConfig(name="rendered", expr="{{ 'MISSING' if value is none else 'PRESENT' }}", dtype="str") + ) + + if run == "preview": + analysis = designer.preview(config, num_records=2).analysis + else: + analysis = designer.create(config, num_records=2, dataset_name="profile_nulls").load_analysis() + + statistics = _statistics(analysis, "value") + assert statistics.num_null == 1 + assert statistics.num_unique == 1 + + +@pytest.mark.parametrize("run", ["preview", "create"]) +def test_profile_keeps_the_distribution_of_a_single_record_numerical_sampler(designer: DataDesigner, run: str) -> None: + config = DataDesignerConfigBuilder() + config.add_column( + SamplerColumnConfig(name="value", sampler_type=SamplerType.UNIFORM, params=UniformSamplerParams(low=0, high=10)) + ) + + if run == "preview": + result = designer.preview(config, num_records=1) + analysis, dataset = result.analysis, result.dataset + else: + created = designer.create(config, num_records=1, dataset_name="profile_single") + analysis, dataset = created.load_analysis(), created.load_dataset() + + statistics = _statistics(analysis, "value") + assert statistics.distribution_type == ColumnDistributionType.NUMERICAL + sampled = float(dataset["value"].iloc[0]) + assert statistics.distribution.min == pytest.approx(sampled, abs=0.01) + assert statistics.distribution.max == pytest.approx(sampled, abs=0.01) + assert statistics.distribution.mean == pytest.approx(sampled, abs=0.01) + assert statistics.distribution.median == pytest.approx(sampled, abs=0.01) + assert statistics.distribution.stddev != statistics.distribution.stddev # NaN: undefined for one record + + +def test_single_record_distribution_survives_saved_metadata_and_resume(designer: DataDesigner) -> None: + config = DataDesignerConfigBuilder() + config.add_column( + SamplerColumnConfig(name="value", sampler_type=SamplerType.UNIFORM, params=UniformSamplerParams(low=0, high=10)) + ) + + created = designer.create(config, num_records=1, dataset_name="profile_resume") + saved = created.artifact_storage.read_metadata()["column_statistics"][0]["distribution"] + reloaded = NumericalDistribution.model_validate(saved) + sampled = float(created.load_dataset()["value"].iloc[0]) + assert reloaded.mean == pytest.approx(sampled, abs=0.01) + assert reloaded.stddev != reloaded.stddev # NaN survives the JSON round trip + + resumed = designer.create(config, num_records=1, dataset_name="profile_resume", resume=ResumeMode.ALWAYS) + statistics = _statistics(resumed.load_analysis(), "value") + assert statistics.distribution.mean == pytest.approx(sampled, abs=0.01) + assert statistics.distribution.stddev != statistics.distribution.stddev + resaved = NumericalDistribution.model_validate( + resumed.artifact_storage.read_metadata()["column_statistics"][0]["distribution"] + ) + assert resaved.stddev != resaved.stddev