fix: handle missing values in column profiling

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),
         )
 

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)

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 "<NA>", 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(),

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 "<NA>" 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

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