diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index cc75377..79c1477 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -8,5 +8,6 @@ # Plugins /plugins/data-designer-github/ @NVIDIA-NeMo/data_designer_reviewers @eric-tramel +/plugins/data-designer-group-consistent/ @andreatnvidia /plugins/data-designer-retrieval-sdg/ @NVIDIA-NeMo/data_designer_reviewers @shan-nvidia @oliverholworthy /plugins/data-designer-template/ @NVIDIA-NeMo/data_designer_reviewers diff --git a/docs/plugins/data-designer-group-consistent/index.md b/docs/plugins/data-designer-group-consistent/index.md new file mode 100644 index 0000000..67ad068 --- /dev/null +++ b/docs/plugins/data-designer-group-consistent/index.md @@ -0,0 +1,57 @@ +# Group-consistent generation + +Use `group-consistent` to select one correlated candidate record for each +logical group. The selection is deterministic for a given combination of group +values, role, seed, and candidate pool. + +## Installation + +```bash +uv add data-designer data-designer-group-consistent +``` + +## Column type + +```python +from data_designer_group_consistent.config import GroupConsistentColumnConfig + +builder.add_column( + GroupConsistentColumnConfig( + name="synthetic_first_name", + group_by=["patient_id"], + role="patient", + seed=7, + records=[ + {"first_name": "Amina", "last_name": "Diallo", "email": "amina@example.test"}, + {"first_name": "Carlos", "last_name": "Silva", "email": "carlos@example.test"}, + ], + field_mapping={ + "synthetic_first_name": "first_name", + "synthetic_last_name": "last_name", + "synthetic_email": "email", + }, + ), +) +``` + +All mapped fields come from the same candidate record. Rows with the same +`patient_id` therefore receive a coherent first name, last name, and email even +when those rows are generated in different batches. + +| Field | Required | Description | +| --- | --- | --- | +| `name` | Yes | Primary output column. It must be a key in `field_mapping`. | +| `group_by` | Yes | Ordered list of upstream columns defining a logical group. | +| `records` | Yes | Non-empty candidate record pool. | +| `field_mapping` | Yes | Mapping from output column names to candidate record fields. | +| `role` | No | Namespace for independent entities in one group. Defaults to `default`. | +| `seed` | No | Deterministic selection seed. Defaults to `0`. | + +## Implementation notes + +The plugin uses SHA-256 rather than Python's process-dependent `hash()` function. +It does not persist a mapping table or call an LLM. Candidate order and pool size +are part of the effective configuration: changing them may change prior choices. + +For the full plugin authoring guide, see the +[main repository docs](https://nvidia-nemo.github.io/DataDesignerPlugins/authoring/). diff --git a/docs/plugins/index.md b/docs/plugins/index.md index 5fc1a9c..ee82842 100644 --- a/docs/plugins/index.md +++ b/docs/plugins/index.md @@ -16,6 +16,17 @@ Browse available Data Designer plugins by what they add to your data generation github + + + data-designer-group-consistent + v0.1.0 + + Deterministic group-scoped record generation for Data Designer + + Entry points + group-consistent + + data-designer-retrieval-sdg diff --git a/plugins/data-designer-group-consistent/CODEOWNERS b/plugins/data-designer-group-consistent/CODEOWNERS new file mode 100644 index 0000000..bc91e78 --- /dev/null +++ b/plugins/data-designer-group-consistent/CODEOWNERS @@ -0,0 +1,3 @@ +# Owner(s) of this plugin — used to generate the root CODEOWNERS file. +# GitHub accepts @username, @org/team, or email format. +* @andreatnvidia diff --git a/plugins/data-designer-group-consistent/README.md b/plugins/data-designer-group-consistent/README.md new file mode 100644 index 0000000..401bf0e --- /dev/null +++ b/plugins/data-designer-group-consistent/README.md @@ -0,0 +1,50 @@ +# data-designer-group-consistent + +Deterministic group-scoped record generation for Data Designer. + +## Installation + +```bash +uv add data-designer data-designer-group-consistent +``` + +## Usage + +The `group-consistent` column type selects one candidate record for each logical +group and writes one or more correlated output columns. Selection is derived from +the group values, role, and seed, so it remains stable across row order, batches, +retries, and resumed runs with the same configuration. + +```python +from data_designer.config import DataDesignerConfigBuilder +from data_designer_group_consistent.config import GroupConsistentColumnConfig + +builder = DataDesignerConfigBuilder() +builder.add_column( + GroupConsistentColumnConfig( + name="synthetic_first_name", + group_by=["patient_id"], + role="patient", + seed=7, + records=[ + {"first_name": "Amina", "last_name": "Diallo", "email": "amina@example.test"}, + {"first_name": "Carlos", "last_name": "Silva", "email": "carlos@example.test"}, + ], + field_mapping={ + "synthetic_first_name": "first_name", + "synthetic_last_name": "last_name", + "synthetic_email": "email", + }, + ), +) +``` + +Every row with the same `patient_id` receives fields from the same candidate +record. Use a different `role` to create an independent identity, such as a +doctor or emergency contact, within the same group. + +For the full plugin authoring guide, see the +[main repository docs](https://nvidia-nemo.github.io/DataDesignerPlugins/authoring/). + +Plugin documentation for the repository site lives in this package's `docs/` +directory. diff --git a/plugins/data-designer-group-consistent/docs/index.md b/plugins/data-designer-group-consistent/docs/index.md new file mode 100644 index 0000000..67ad068 --- /dev/null +++ b/plugins/data-designer-group-consistent/docs/index.md @@ -0,0 +1,57 @@ +# Group-consistent generation + +Use `group-consistent` to select one correlated candidate record for each +logical group. The selection is deterministic for a given combination of group +values, role, seed, and candidate pool. + +## Installation + +```bash +uv add data-designer data-designer-group-consistent +``` + +## Column type + +```python +from data_designer_group_consistent.config import GroupConsistentColumnConfig + +builder.add_column( + GroupConsistentColumnConfig( + name="synthetic_first_name", + group_by=["patient_id"], + role="patient", + seed=7, + records=[ + {"first_name": "Amina", "last_name": "Diallo", "email": "amina@example.test"}, + {"first_name": "Carlos", "last_name": "Silva", "email": "carlos@example.test"}, + ], + field_mapping={ + "synthetic_first_name": "first_name", + "synthetic_last_name": "last_name", + "synthetic_email": "email", + }, + ), +) +``` + +All mapped fields come from the same candidate record. Rows with the same +`patient_id` therefore receive a coherent first name, last name, and email even +when those rows are generated in different batches. + +| Field | Required | Description | +| --- | --- | --- | +| `name` | Yes | Primary output column. It must be a key in `field_mapping`. | +| `group_by` | Yes | Ordered list of upstream columns defining a logical group. | +| `records` | Yes | Non-empty candidate record pool. | +| `field_mapping` | Yes | Mapping from output column names to candidate record fields. | +| `role` | No | Namespace for independent entities in one group. Defaults to `default`. | +| `seed` | No | Deterministic selection seed. Defaults to `0`. | + +## Implementation notes + +The plugin uses SHA-256 rather than Python's process-dependent `hash()` function. +It does not persist a mapping table or call an LLM. Candidate order and pool size +are part of the effective configuration: changing them may change prior choices. + +For the full plugin authoring guide, see the +[main repository docs](https://nvidia-nemo.github.io/DataDesignerPlugins/authoring/). diff --git a/plugins/data-designer-group-consistent/pyproject.toml b/plugins/data-designer-group-consistent/pyproject.toml new file mode 100644 index 0000000..d9de34b --- /dev/null +++ b/plugins/data-designer-group-consistent/pyproject.toml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[project] +name = "data-designer-group-consistent" +version = "0.1.0" +description = "Deterministic group-scoped record generation for Data Designer" +requires-python = ">=3.10" +dependencies = [ + "data-designer>=0.5.7", +] +license = "Apache-2.0" +readme = "README.md" +authors = [ + {name = "NVIDIA Corporation"}, +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: Python :: 3", +] + +[project.entry-points."data_designer.plugins"] +group-consistent = "data_designer_group_consistent.plugin:plugin" + +[project.urls] +Repository = "https://github.com/NVIDIA-NeMo/DataDesignerPlugins" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/data_designer_group_consistent"] + +[tool.ruff] +extend = "../../pyproject.toml" diff --git a/plugins/data-designer-group-consistent/src/data_designer_group_consistent/__init__.py b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/__init__.py new file mode 100644 index 0000000..52a7a9d --- /dev/null +++ b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/plugins/data-designer-group-consistent/src/data_designer_group_consistent/config.py b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/config.py new file mode 100644 index 0000000..6cefea1 --- /dev/null +++ b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/config.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Literal + +from data_designer.config.base import SingleColumnConfig +from pydantic import Field, JsonValue, model_validator +from typing_extensions import Self + + +class GroupConsistentColumnConfig(SingleColumnConfig): + """Selects one correlated record deterministically for each logical group.""" + + column_type: Literal["group-consistent"] = "group-consistent" + group_by: list[str] = Field(min_length=1) + records: list[dict[str, JsonValue]] = Field(min_length=1) + field_mapping: dict[str, str] = Field(min_length=1) + role: str = Field(default="default", min_length=1) + seed: int = 0 + + @model_validator(mode="after") + def validate_generation_contract(self) -> Self: + """Validate group keys, output columns, and candidate record fields. + + Returns: + The validated configuration. + + Raises: + ValueError: If the configuration cannot generate every declared output. + """ + if len(set(self.group_by)) != len(self.group_by): + raise ValueError("group_by columns must be unique") + if self.name not in self.field_mapping: + raise ValueError(f"field_mapping must include the primary output column {self.name!r}") + if not self.role.strip(): + raise ValueError("role must not be blank") + + required_fields = set(self.field_mapping.values()) + for index, record in enumerate(self.records): + missing_fields = required_fields - record.keys() + if missing_fields: + raise ValueError(f"records[{index}] is missing mapped fields: {sorted(missing_fields)}") + return self + + @staticmethod + def get_column_emoji() -> str: + return "🔗" + + @property + def required_columns(self) -> list[str]: + return self.group_by + + @property + def side_effect_columns(self) -> list[str]: + return [column for column in self.field_mapping if column != self.name] diff --git a/plugins/data-designer-group-consistent/src/data_designer_group_consistent/impl.py b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/impl.py new file mode 100644 index 0000000..262445c --- /dev/null +++ b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/impl.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +import math +from datetime import date, datetime, time +from typing import TYPE_CHECKING + +from data_designer.engine.column_generators.generators.base import ColumnGeneratorFullColumn +from pydantic import JsonValue + +from data_designer_group_consistent.config import GroupConsistentColumnConfig + +if TYPE_CHECKING: + import pandas as pd + + +def normalize_key_component(value: object) -> dict[str, JsonValue]: + """Convert a group-key value to a stable JSON representation. + + Args: + value: Scalar value from a group-key column. + + Returns: + Type-tagged JSON data suitable for deterministic hashing. + """ + item_method = getattr(value, "item", None) + if callable(item_method): + value = item_method() + + if value is None: + return {"type": "null", "value": None} + if isinstance(value, float) and math.isnan(value): + return {"type": "null", "value": None} + if isinstance(value, float) and math.isinf(value): + return {"type": "float", "value": "infinity" if value > 0 else "-infinity"} + if isinstance(value, (date, datetime, time)): + return {"type": type(value).__name__, "value": value.isoformat()} + if isinstance(value, (bool, int, float, str)): + return {"type": type(value).__name__, "value": value} + return { + "type": f"{type(value).__module__}.{type(value).__qualname__}", + "value": str(value), + } + + +def select_record_index(group_values: tuple[object, ...], role: str, seed: int, record_count: int) -> int: + """Select a deterministic candidate record for a logical group. + + Args: + group_values: Ordered values from the configured group-key columns. + role: Namespace for independently generated entities in the same group. + seed: User-controlled generation seed. + record_count: Number of candidate records. + + Returns: + An index into the configured candidate record list. + """ + payload = { + "version": 1, + "group": [normalize_key_component(value) for value in group_values], + "role": role, + "seed": seed, + } + encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8") + return int.from_bytes(hashlib.sha256(encoded).digest(), byteorder="big") % record_count + + +class GroupConsistentColumnGenerator(ColumnGeneratorFullColumn[GroupConsistentColumnConfig]): + """Generate correlated fields that remain stable within each logical group.""" + + def generate(self, data: pd.DataFrame) -> pd.DataFrame: + """Add fields from one deterministic candidate record per group. + + Args: + data: Batch containing the configured group-key columns. + + Returns: + The batch with the configured output columns added. + """ + records = self.config.records + indexes = [ + select_record_index(group_values, self.config.role, self.config.seed, len(records)) + for group_values in data[self.config.group_by].itertuples(index=False, name=None) + ] + selected_records = [records[index] for index in indexes] + for output_column, record_field in self.config.field_mapping.items(): + data[output_column] = [record[record_field] for record in selected_records] + return data diff --git a/plugins/data-designer-group-consistent/src/data_designer_group_consistent/plugin.py b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/plugin.py new file mode 100644 index 0000000..41cf232 --- /dev/null +++ b/plugins/data-designer-group-consistent/src/data_designer_group_consistent/plugin.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from data_designer.plugins.plugin import Plugin, PluginType + +plugin = Plugin( + config_qualified_name="data_designer_group_consistent.config.GroupConsistentColumnConfig", + impl_qualified_name="data_designer_group_consistent.impl.GroupConsistentColumnGenerator", + plugin_type=PluginType.COLUMN_GENERATOR, +) diff --git a/plugins/data-designer-group-consistent/tests/test_plugin.py b/plugins/data-designer-group-consistent/tests/test_plugin.py new file mode 100644 index 0000000..ab19564 --- /dev/null +++ b/plugins/data-designer-group-consistent/tests/test_plugin.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pandas as pd +import pytest +from data_designer.config.config_builder import DataDesignerConfigBuilder +from data_designer.config.seed_source_dataframe import DataFrameSeedSource +from data_designer.engine.testing.utils import assert_valid_plugin +from data_designer.interface.data_designer import DataDesigner +from pydantic import ValidationError + +from data_designer_group_consistent.config import GroupConsistentColumnConfig +from data_designer_group_consistent.impl import GroupConsistentColumnGenerator +from data_designer_group_consistent.plugin import plugin + +PERSONAS = [ + {"first_name": "Amina", "last_name": "Diallo", "email": "amina@example.test"}, + {"first_name": "Carlos", "last_name": "Silva", "email": "carlos@example.test"}, + {"first_name": "Mei", "last_name": "Chen", "email": "mei@example.test"}, +] + + +def test_valid_plugin() -> None: + assert_valid_plugin(plugin) + + +def make_config(**overrides: object) -> GroupConsistentColumnConfig: + """Create a valid test configuration with optional field overrides.""" + values = { + "name": "synthetic_first_name", + "group_by": ["patient_id"], + "records": PERSONAS, + "field_mapping": { + "synthetic_first_name": "first_name", + "synthetic_last_name": "last_name", + "synthetic_email": "email", + }, + "role": "patient", + "seed": 7, + } + values.update(overrides) + return GroupConsistentColumnConfig(**values) + + +def make_generator(config: GroupConsistentColumnConfig | None = None) -> GroupConsistentColumnGenerator: + """Create a generator without requiring a resource provider.""" + generator = GroupConsistentColumnGenerator.__new__(GroupConsistentColumnGenerator) + generator._config = config or make_config() + return generator + + +class TestGroupConsistentColumnConfig: + def test_declares_dependencies_and_side_effects(self) -> None: + config = make_config(group_by=["household_id", "patient_id"]) + + assert config.required_columns == ["household_id", "patient_id"] + assert config.side_effect_columns == ["synthetic_last_name", "synthetic_email"] + + @pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"group_by": ["patient_id", "patient_id"]}, "group_by columns must be unique"), + ({"field_mapping": {"synthetic_last_name": "last_name"}}, "must include the primary output"), + ({"role": " "}, "role must not be blank"), + ( + {"records": [{"first_name": "Amina", "last_name": "Diallo"}]}, + "missing mapped fields", + ), + ], + ) + def test_rejects_invalid_generation_contract(self, overrides: dict[str, object], message: str) -> None: + with pytest.raises(ValidationError, match=message): + make_config(**overrides) + + +class TestGroupConsistentColumnGenerator: + def test_reuses_correlated_record_for_noncontiguous_group_rows(self) -> None: + data = pd.DataFrame({"patient_id": ["p1", "p2", "p1", "p3", "p2"]}) + + result = make_generator().generate(data) + + for _, group in result.groupby("patient_id"): + assert group["synthetic_first_name"].nunique() == 1 + assert group["synthetic_last_name"].nunique() == 1 + assert group["synthetic_email"].nunique() == 1 + selected_personas = {(record["first_name"], record["last_name"], record["email"]) for record in PERSONAS} + assert ( + set( + result[["synthetic_first_name", "synthetic_last_name", "synthetic_email"]].itertuples( + index=False, name=None + ) + ) + <= selected_personas + ) + + def test_is_stable_across_batches_and_row_order(self) -> None: + generator = make_generator() + + first_batch = generator.generate(pd.DataFrame({"patient_id": ["p1", "p2"]})) + second_batch = generator.generate(pd.DataFrame({"patient_id": ["p3", "p1"]})) + + first_value = first_batch.loc[first_batch["patient_id"] == "p1", "synthetic_email"].item() + second_value = second_batch.loc[second_batch["patient_id"] == "p1", "synthetic_email"].item() + assert first_value == second_value + + def test_treats_missing_group_values_as_one_stable_group(self) -> None: + data = pd.DataFrame({"patient_id": [None, "p1", float("nan"), None]}) + + result = make_generator().generate(data) + + missing_group = result[result["patient_id"].isna()] + assert missing_group["synthetic_email"].nunique() == 1 + + +class TestGroupConsistentPreviewIntegration: + def test_preview_generates_group_consistent_personas(self, tmp_path: Path) -> None: + seed_df = pd.DataFrame({"patient_id": ["p1", "p2", "p1", "p3"]}) + builder = DataDesignerConfigBuilder() + builder.with_seed_dataset(DataFrameSeedSource(df=seed_df)) + builder.add_column( + GroupConsistentColumnConfig( + name="synthetic_first_name", + group_by=["patient_id"], + records=PERSONAS, + field_mapping={ + "synthetic_first_name": "first_name", + "synthetic_last_name": "last_name", + "synthetic_email": "email", + }, + role="patient", + seed=7, + ), + ) + + result = DataDesigner(artifact_path=tmp_path / "artifacts").preview(builder, num_records=4) + + assert result.dataset is not None + patient_rows = result.dataset[result.dataset["patient_id"] == "p1"] + assert patient_rows["synthetic_first_name"].nunique() == 1 + assert patient_rows["synthetic_last_name"].nunique() == 1 + assert patient_rows["synthetic_email"].nunique() == 1 diff --git a/uv.lock b/uv.lock index 7d1c44c..83c1c3e 100644 --- a/uv.lock +++ b/uv.lock @@ -13,6 +13,7 @@ resolution-markers = [ [manifest] members = [ "data-designer-github", + "data-designer-group-consistent", "data-designer-plugins-workspace", "data-designer-retrieval-sdg", "data-designer-template", @@ -509,6 +510,17 @@ requires-dist = [ { name = "starlette", specifier = ">=1.4.1" }, ] +[[package]] +name = "data-designer-group-consistent" +version = "0.1.0" +source = { editable = "plugins/data-designer-group-consistent" } +dependencies = [ + { name = "data-designer" }, +] + +[package.metadata] +requires-dist = [{ name = "data-designer", specifier = ">=0.5.7" }] + [[package]] name = "data-designer-plugins-workspace" version = "0.0.0" diff --git a/zensical.toml b/zensical.toml index ac6b49c..ca137f5 100644 --- a/zensical.toml +++ b/zensical.toml @@ -25,6 +25,9 @@ nav = [ {"Overview" = "plugins/data-designer-github/index.md"}, {"Usage" = "plugins/data-designer-github/usage.md"}, ]}, + {"data-designer-group-consistent" = [ + {"Overview" = "plugins/data-designer-group-consistent/index.md"}, + ]}, {"data-designer-retrieval-sdg" = [ {"Overview" = "plugins/data-designer-retrieval-sdg/index.md"}, ]},