From 53ef196a16e540cf3601c5a565885ce7000215c7 Mon Sep 17 00:00:00 2001 From: Yev Meyer Date: Mon, 9 Nov 2026 19:00:59 -0500 Subject: [PATCH 1/2] fix: sample missing person fields as None A person field that is empty for only some of the sampled people is None in the derived fields and in the output. pandas 3 holds a missing string as NaN, which is truthy and not valid JSON: an en_SG person without a middle name made the email address fail with a TypeError, and en_US bachelors_field came back as NaN. Fixes #1146 Signed-off-by: Yev Meyer --- .../engine/sampling_gen/people_gen.py | 5 ++++ .../engine/sampling_gen/test_people_gen.py | 25 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py index 7bcaa4ec7..6fb127237 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py @@ -132,6 +132,11 @@ def _generate_from_dataset(self, n: int, **kwargs) -> pd.DataFrame: # Generate samples and drop columns where all rows are null. df = self._engine.generate_samples(size=n, evidence=kw).dropna(axis=1, how="all") + # A value missing from only some rows becomes None: a pandas string column holds it as NaN, + # which is truthy and not JSON-serializable. + for column in df.columns[df.isna().any()]: + df[column] = df[column].astype(object).where(df[column].notna(), None) + # We need this for derived fields. df["locale"] = self.locale diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py index 26a78efc8..2f5ad2606 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py @@ -198,3 +198,28 @@ def test_managed_locale_demographic_output_and_filtering( assert all(person[field] == selected_value for person in people) assert all(person["locale"] == locale for person in people) assert all(("persona" in person) == with_synthetic_personas for person in people) + + +def test_managed_sampling_gives_none_for_a_value_only_some_people_have(tmp_path: Path) -> None: + """A field empty for some sampled people, such as en_SG's middle_name, comes back as None.""" + dataset_dir = tmp_path / "datasets" + dataset_dir.mkdir() + records = [ + {"first_name": "Chin", "middle_name": "Peng", "last_name": "Tan", "age": 67, "industry": "Finance"}, + {"first_name": "Noor", "middle_name": None, "last_name": "Ismail", "age": 60, "industry": None}, + {"first_name": "Lin", "middle_name": "Hui", "last_name": "Chng", "age": 73, "industry": None}, + {"first_name": "Gulzar", "middle_name": None, "last_name": "Subramaniam", "age": 53, "industry": "Education"}, + ] + lazy.pd.DataFrame(records).to_parquet(dataset_dir / "en_SG.parquet") + generator = PeopleGenFromDataset(load_person_data_sampler(LocalPersonReader(tmp_path), "en_SG"), "en_SG") + params = PersonSamplerParams(locale="en_SG") + sampler = PersonSampler(params, people_gen_resource={params.people_gen_key: generator}) + + people = sampler.sample(len(records)) + + expected = {record["first_name"]: record for record in records} + assert sorted(person["first_name"] for person in people) == sorted(expected) + for person in people: + assert person["middle_name"] == expected[person["first_name"]]["middle_name"] + assert person["industry"] == expected[person["first_name"]]["industry"] + assert person["email_address"] From 5ca964d8fb7ceeac359f37dc238d8efba2441d09 Mon Sep 17 00:00:00 2001 From: Yev Meyer Date: Mon, 9 Nov 2026 19:21:27 -0500 Subject: [PATCH 2/2] fix: give templates None for missing text values A row built from a DataFrame for generation holds None for a missing value in a text or object column, as it does under pandas 2. pandas 3 holds that value as NaN (pd.NA for the string dtypes), which Jinja's `is none` test rejects: an expression or prompt that checked a missing seed value took the wrong branch, and a custom column saw NaN. dataframe_to_records does the conversion for seed and batch results, the row buffer, expression columns, the schema transform processor and the person sampler, which now uses it in place of its own loop. Numeric, datetime and nested values are unchanged. Fixes #1148 Signed-off-by: Yev Meyer --- .../generators/expression.py | 6 +-- .../dataset_builders/async_scheduler.py | 6 ++- .../utils/row_group_buffer.py | 3 +- .../processing/processors/schema_transform.py | 4 +- .../data_designer/engine/processing/utils.py | 27 ++++++++++ .../engine/sampling_gen/people_gen.py | 8 +-- .../generators/test_expression.py | 19 +++++++ .../utils/test_row_group_buffer.py | 11 ++++ .../tests/engine/processing/test_utils.py | 40 ++++++++++++++ .../tests/interface/test_missing_values.py | 54 +++++++++++++++++++ 10 files changed, 164 insertions(+), 14 deletions(-) create mode 100644 packages/data-designer/tests/interface/test_missing_values.py diff --git a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/expression.py b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/expression.py index 57bf12d71..cd6a41876 100644 --- a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/expression.py +++ b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/expression.py @@ -14,7 +14,7 @@ from data_designer.engine.context import format_row_group_tag from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering from data_designer.engine.processing.ginja.exceptions import EmptyTemplateRenderError, UserTemplateError -from data_designer.engine.processing.utils import deserialize_json_values +from data_designer.engine.processing.utils import dataframe_to_records, deserialize_json_values if TYPE_CHECKING: import pandas as pd @@ -47,14 +47,14 @@ def generate(self, data: pd.DataFrame) -> pd.DataFrame: retained_indexes: list[object] = [] drop_counts: Counter[str] = Counter() - # ``DataFrame.to_dict(orient="records")`` returns ``[]`` for a 0-column + # ``dataframe_to_records`` returns ``[]`` for a 0-column # DataFrame regardless of its row count, so a root expression dispatched # against a fresh row group buffer (no upstream columns) would otherwise # skip rendering entirely. Synthesize empty per-row dicts in that case so # the loop still fires once per row and the template can render its # constant value (or drop the row uniformly). row_indexes = data.index.to_list() - row_records = data.to_dict(orient="records") + row_records = dataframe_to_records(data) if not row_records and row_indexes: row_records = [{} for _ in row_indexes] diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py index 53130fd85..17a62d0b5 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py @@ -83,6 +83,7 @@ scheduler_event_sink_accepts, ) from data_designer.engine.processing.ginja.exceptions import UserTemplateError +from data_designer.engine.processing.utils import dataframe_to_records from data_designer.engine.progress.reporter import ( DEFAULT_REPORT_INTERVAL, AsyncProgressReporter, @@ -2263,9 +2264,10 @@ async def _run_from_scratch( # Write results to buffer (include side-effect columns) if self._buffer_manager is not None: write_cols = self._gen_instance_to_columns_including_side_effects.get(id(generator), [task.column]) + result_records = dataframe_to_records(result_df) for col in write_cols: if col in result_df.columns: - values = result_df[col].tolist() + values = [record[col] for record in result_records] self._buffer_manager.update_batch(task.row_group, col, values) return result_df @@ -2514,7 +2516,7 @@ def _batch_result_by_row_index( active_row_indices: list[int], supports_row_drops: bool, ) -> dict[int, dict[str, Any]]: - result_records = result_df.to_dict(orient="records") + result_records = dataframe_to_records(result_df) if supports_row_drops: return dict(zip(result_df.index.to_list(), result_records, strict=True)) return dict(zip(active_row_indices, result_records, strict=True)) diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py index 11ad72c66..2c9635f56 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/row_group_buffer.py @@ -9,6 +9,7 @@ from data_designer.config.utils.storage_encoding import create_logical_dataframe from data_designer.engine.dataset_builders.utils.skip_tracker import strip_skip_metadata_from_records +from data_designer.engine.processing.utils import dataframe_to_records if TYPE_CHECKING: import pandas as pd @@ -89,7 +90,7 @@ def replace_dataframe(self, row_group: int, df: pd.DataFrame) -> None: If *df* has fewer rows than active slots, trailing slots are marked as dropped. """ dropped = self._dropped.get(row_group, set()) - records = df.to_dict(orient="records") + records = dataframe_to_records(df) buf_idx = 0 for ri in range(self._row_group_sizes[row_group]): if ri in dropped: diff --git a/packages/data-designer-engine/src/data_designer/engine/processing/processors/schema_transform.py b/packages/data-designer-engine/src/data_designer/engine/processing/processors/schema_transform.py index 84539affe..46b79be9b 100644 --- a/packages/data-designer-engine/src/data_designer/engine/processing/processors/schema_transform.py +++ b/packages/data-designer-engine/src/data_designer/engine/processing/processors/schema_transform.py @@ -12,7 +12,7 @@ from data_designer.config.utils.storage_schema import JsonStoragePlan from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering from data_designer.engine.processing.processors.base import Processor -from data_designer.engine.processing.utils import deserialize_json_values +from data_designer.engine.processing.utils import dataframe_to_records, deserialize_json_values from data_designer.engine.storage.artifact_storage import BatchStage if TYPE_CHECKING: @@ -82,7 +82,7 @@ def _transform(self, data: pd.DataFrame) -> pd.DataFrame: self.template_as_str, data.columns.to_list(), record_str_fn=_escape_value_for_json ) formatted_records = [] - for record in data.to_dict(orient="records"): + for record in dataframe_to_records(data): deserialized = deserialize_json_values(record) rendered = self.render_template(deserialized) formatted_records.append(json.loads(rendered)) diff --git a/packages/data-designer-engine/src/data_designer/engine/processing/utils.py b/packages/data-designer-engine/src/data_designer/engine/processing/utils.py index afdff5a49..632bfbf36 100644 --- a/packages/data-designer-engine/src/data_designer/engine/processing/utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/processing/utils.py @@ -30,6 +30,33 @@ def concat_datasets(datasets: list[pd.DataFrame]) -> pd.DataFrame: return lazy.pd.concat([df for df in datasets], axis=1) +def dataframe_to_records(df: pd.DataFrame) -> list[dict[str, Any]]: + """Convert a DataFrame to one dict per row, with None for each missing text value. + + pandas 3 holds a missing value in a text column as NaN (``pd.NA`` for the ``string`` dtypes). Both are + truthy and fail Jinja's ``is none`` test, so a missing value in a text or object column becomes None. + Values in numeric, datetime and other columns, and values nested inside a cell, are left unchanged. + + Args: + df: DataFrame to convert. + + Returns: + The rows of ``df`` as dicts keyed by column name. + """ + # Columns are addressed by position, so a repeated column name converts like any other. + text_positions_with_missing = [ + position + for position, dtype in enumerate(df.dtypes) + if lazy.pd.api.types.is_string_dtype(dtype) and df.iloc[:, position].hasnans + ] + if text_positions_with_missing: + df = df.copy(deep=False) + for position in text_positions_with_missing: + column = df.iloc[:, position] + df.isetitem(position, column.astype(object).where(column.notna(), None)) + return df.to_dict(orient="records") + + # Overloads to help static type checker better understand # the input/output types of the deserialize_json_values function. @overload diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py index 6fb127237..ffb831dd7 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/people_gen.py @@ -12,6 +12,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.config.utils.constants import DEFAULT_AGE_RANGE +from data_designer.engine.processing.utils import dataframe_to_records from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator from data_designer.engine.sampling_gen.entities.dataset_based_person_fields import PERSONA_FIELDS, PII_FIELDS from data_designer.engine.sampling_gen.entities.person import ( @@ -132,11 +133,6 @@ def _generate_from_dataset(self, n: int, **kwargs) -> pd.DataFrame: # Generate samples and drop columns where all rows are null. df = self._engine.generate_samples(size=n, evidence=kw).dropna(axis=1, how="all") - # A value missing from only some rows becomes None: a pandas string column holds it as NaN, - # which is truthy and not JSON-serializable. - for column in df.columns[df.isna().any()]: - df[column] = df[column].astype(object).where(df[column].notna(), None) - # We need this for derived fields. df["locale"] = self.locale @@ -150,7 +146,7 @@ def _generate_from_dataset(self, n: int, **kwargs) -> pd.DataFrame: def generate(self, n: int, **kwargs) -> list[dict[str, Any]]: return [ generate_and_insert_derived_fields(p) - for p in self._generate_from_dataset(n, **kwargs).to_dict(orient="records") + for p in dataframe_to_records(self._generate_from_dataset(n, **kwargs)) ] diff --git a/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py b/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py index a27f9f7bb..a1cf3468f 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py +++ b/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py @@ -312,3 +312,22 @@ def test_generate_respects_run_config_jinja_rendering_engine() -> None: with pytest.raises(UserTemplateUnsupportedFiltersError): secure_generator.generate(df) + + +@pytest.mark.parametrize("dtype", [None, "string"], ids=["inferred", "string"]) +def test_generate_sees_a_missing_text_value_as_none(dtype: str | None) -> None: + config = _create_test_config("rendered", "{{ 'MISSING' if value is none else 'PRESENT' }}", "str") + df = lazy.pd.DataFrame({"value": lazy.pd.Series(["present", None, ""], dtype=dtype)}) + + result = _create_test_generator(config).generate(df) + + assert result["rendered"].tolist() == ["PRESENT", "MISSING", "PRESENT"] + + +def test_generate_keeps_nested_values() -> None: + config = _create_test_config("rendered", "{{ 'NONE' if payload['key'] is none else payload['key'] }}", "str") + df = lazy.pd.DataFrame({"payload": [{"key": None}, {"key": "set"}]}) + + result = _create_test_generator(config).generate(df) + + assert result["rendered"].tolist() == ["NONE", "set"] diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_row_group_buffer.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_row_group_buffer.py index e460583b7..3fb998bb5 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_row_group_buffer.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_row_group_buffer.py @@ -207,6 +207,17 @@ def test_replace_dataframe_fewer_rows_marks_trailing_dropped() -> None: assert len(result_df) == 2 +def test_replace_dataframe_gives_none_for_missing_text() -> None: + """A missing text value is None in the buffer's rows; numeric values are kept as they are.""" + mgr = RowGroupBufferManager(_mock_artifact_storage()) + mgr.init_row_group(0, 2) + + mgr.replace_dataframe(0, lazy.pd.DataFrame({"col": ["x", None], "score": [1.0, 2.0]})) + + assert mgr.get_row(0, 0) == {"col": "x", "score": 1.0} + assert mgr.get_row(0, 1) == {"col": None, "score": 2.0} + + def test_free_row_group_releases_memory() -> None: """free_row_group releases buffer memory without writing to disk.""" storage = _mock_artifact_storage() diff --git a/packages/data-designer-engine/tests/engine/processing/test_utils.py b/packages/data-designer-engine/tests/engine/processing/test_utils.py index 7ca3bfce2..16d67c47a 100644 --- a/packages/data-designer-engine/tests/engine/processing/test_utils.py +++ b/packages/data-designer-engine/tests/engine/processing/test_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations +import math from unittest.mock import patch import pytest @@ -10,6 +11,7 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.engine.processing.utils import ( concat_datasets, + dataframe_to_records, deserialize_json_values, parse_list_string, ) @@ -135,3 +137,41 @@ def test_deserialize_json_values_scenarios(test_case, input_data, expected_resul def test_parse_list_string_scenarios(input_string, expected_result): result = parse_list_string(input_string) assert result == expected_result + + +@pytest.mark.parametrize( + "missing,dtype", + [(None, None), (None, "string"), (None, "string[pyarrow]"), (float("nan"), object)], + ids=["inferred", "string", "string-pyarrow", "object-nan"], +) +def test_dataframe_to_records_gives_none_for_missing_text(missing: object, dtype: str | type | None) -> None: + df = lazy.pd.DataFrame({"value": lazy.pd.Series(["present", missing, ""], dtype=dtype)}) + + records = dataframe_to_records(df) + + assert [record["value"] for record in records] == ["present", None, ""] + assert df["value"].isna().tolist() == [False, True, False] + + +def test_dataframe_to_records_keeps_numeric_datetime_and_nested_values() -> None: + df = lazy.pd.DataFrame( + { + "score": [1.5, float("nan")], + "when": [lazy.pd.Timestamp("2026-01-01"), lazy.pd.NaT], + "nested": [{"key": None}, ["a", None]], + "text": ["nan", "None"], + } + ) + + records = dataframe_to_records(df) + + assert records[0] == { + "score": 1.5, + "when": lazy.pd.Timestamp("2026-01-01"), + "nested": {"key": None}, + "text": "nan", + } + assert math.isnan(records[1]["score"]) + assert records[1]["when"] is lazy.pd.NaT + assert records[1]["nested"] == ["a", None] + assert records[1]["text"] == "None" diff --git a/packages/data-designer/tests/interface/test_missing_values.py b/packages/data-designer/tests/interface/test_missing_values.py new file mode 100644 index 000000000..75713f5b6 --- /dev/null +++ b/packages/data-designer/tests/interface/test_missing_values.py @@ -0,0 +1,54 @@ +# 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 +from typing import Any + +import pandas as pd +import pytest + +from data_designer.config.column_configs import CustomColumnConfig, ExpressionColumnConfig +from data_designer.config.config_builder import DataDesignerConfigBuilder +from data_designer.config.custom_column import custom_column_generator +from data_designer.config.models import ModelProvider +from data_designer.config.run_config import RunConfig +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 + + +@custom_column_generator(required_columns=["value"]) +def flag_missing_value(row: dict[str, Any]) -> dict[str, Any]: + return {**row, "flag": "MISSING" if row["value"] is None else "PRESENT"} + + +@pytest.mark.parametrize("run", ["preview", "create"]) +def test_a_missing_seed_text_value_reaches_templates_and_generators_as_none( + tmp_path: Path, stub_model_providers: list[ModelProvider], run: str +) -> None: + 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)) + config = DataDesignerConfigBuilder().with_seed_dataset( + DataFrameSeedSource(df=pd.DataFrame({"value": ["present", None, ""]})) + ) + config.add_column( + ExpressionColumnConfig(name="rendered", expr="{{ 'MISSING' if value is none else 'PRESENT' }}", dtype="str") + ) + config.add_column(CustomColumnConfig(name="flag", generator_function=flag_missing_value)) + + if run == "preview": + dataset = designer.preview(config, num_records=3).dataset + else: + dataset = designer.create(config, num_records=3, dataset_name="missing_values").load_dataset() + + assert dataset["rendered"].tolist() == ["PRESENT", "MISSING", "PRESENT"] + assert dataset["flag"].tolist() == ["PRESENT", "MISSING", "PRESENT"]