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 7bcaa4ec7..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 ( @@ -145,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-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"] 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"]