Add build_permuted_exploration_dataset for relational candidate selection, preserving intra-group relationships (between children under the same parent).

dpsynth/relational/README.md

@@ -0,0 +1,65 @@
+# `dpsynth.relational` (Multi-Table Relational DP Synthesis)
+
+> [!WARNING]
+> **UNDER ACTIVE DEVELOPMENT / EXPERIMENTAL**: This module is in an early
+> experimental stage and is under active development. The APIs, internal
+> interfaces, and algorithms are subject to breaking changes. It is not yet
+> recommended for production use.
+
+---
+
+## Overview
+
+The `dpsynth.relational` package extends DP Synth to support **hierarchical
+multi-table relational databases** under differential privacy.
+
+In relational databases (e.g. `Household -> Person -> Activity`), privacy is
+protected at the **root parent record** (e.g. the household), while valid
+foreign-key linkages and cross-table statistical correlations are preserved
+across child tables without generating orphaned records or requiring
+intractable flat Cartesian joins.
+
+### Core Approach
+
+- **Cascading Relational Synthesis**: Synthesizes databases table-by-table
+  down foreign-key hierarchy trees, ensuring root-entity differential privacy
+  without materializing expensive Cartesian joins.
+- **Permutation Modeling**: Captures rich parent-child and sibling-to-sibling
+  correlations through slot permutation and exchangeability, inspired by
+  PrivPetal ([Cai et al., 2025](https://arxiv.org/abs/2503.22970)).
+- **Built on `dpsynth` & `mbi` Core**: Reuses existing discrete mechanisms (e.g.
+  AIM, MST) and the `mbi` (Private-PGM) graphical model engine under the hood.
+- **Unified Mechanism API**: Integrates directly with the `MechanismConfig` /
+  `CalibratedMechanism` paradigm used throughout `dpsynth`.
+
+---
+
+## Planned API
+
+```python
+import dpsynth
+from dpsynth.relational import ForeignKeyRelation, MultiTableConfig
+
+# Define schemas and foreign key relations
+foreign_keys = [
+    ForeignKeyRelation(
+        parent_table="households",
+        parent_primary_key="household_id",
+        child_table="persons",
+        child_foreign_key="household_id",
+        max_children_per_parent=5,
+    ),
+]
+
+# Configure relational synthesizer
+config = MultiTableConfig(
+    domains=table_domains,
+    foreign_keys=foreign_keys,
+    discrete_mechanism=dpsynth.discrete_mechanisms.AIMConfig(),
+)
+
+# Calibrate and run
+calibrated = config.calibrate(epsilon=1.0, delta=1e-5)
+result = calibrated(rng, {"households": df_h, "persons": df_p})
+synth_tables = result.synthetic_tables
+```

dpsynth/relational/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Public API for multi-table relational differential privacy synthesis."""
+
+# pylint: disable=g-importing-member
+
+from dpsynth.relational.domain import ForeignKeyRelation
+from dpsynth.relational.domain import from_dict
+from dpsynth.relational.domain import from_yaml_file
+from dpsynth.relational.domain import topological_sort_hierarchy
+from dpsynth.relational.synthesizer import MultiDataGenerationResult
+from dpsynth.relational.synthesizer import MultiTableConfig
+from dpsynth.relational.synthesizer import MultiTableMechanism
+
+__all__ = [
+    'ForeignKeyRelation',
+    'MultiDataGenerationResult',
+    'MultiTableConfig',
+    'MultiTableMechanism',
+    'from_dict',
+    'from_yaml_file',
+    'topological_sort_hierarchy',
+]

dpsynth/relational/domain.py

@@ -0,0 +1,275 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Domain representations, schema definitions, and DAG validators for relational data."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+import dataclasses
+from typing import Any
+
+from absl import logging
+from dpsynth import domain
+from etils import epath
+import networkx as nx
+import yaml
+
+PathType = epath.PathLike
+
+
+@dataclasses.dataclass(frozen=True)
+class ForeignKeyRelation:
+  """Defines a directed foreign key relationship between parent and child tables.
+
+  Attributes:
+    parent_table: Name of the parent table (e.g. 'households').
+    parent_primary_key: Name of the parent primary key column (e.g.
+      'household_id').
+    child_table: Name of the child table (e.g. 'persons').
+    child_foreign_key: Name of the child foreign key column referencing parent.
+    max_children_per_parent: Maximum number of children associated with a single
+      parent record (group size capacity bound s). Must be >= 1. Determines the
+      wide MRF generation slot count (s) and directly scales cascading DP
+      sensitivity (Delta_k = prod s_ancestors) for downstream child tables.
+  """
+
+  parent_table: str
+  parent_primary_key: str
+  child_table: str
+  child_foreign_key: str
+  max_children_per_parent: int
+
+  def __post_init__(self):
+    if self.max_children_per_parent < 1:
+      raise ValueError(
+          'max_children_per_parent must be >= 1, got'
+          f' {self.max_children_per_parent}.'
+      )
+
+
+def topological_sort_hierarchy(
+    tables: Sequence[str],
+    foreign_keys: Sequence[ForeignKeyRelation],
+) -> list[tuple[int, str, ForeignKeyRelation | None]]:
+  """Validates DAG tree structure and computes topological synthesis levels.
+
+  Args:
+    tables: Sequence of all table names in the database.
+    foreign_keys: Sequence of foreign key relationships between tables.
+
+  Returns:
+    An ordered list of (depth, table_name, foreign_key_relation) tuples, where
+    depth is 0 for root tables (foreign_key_relation is None) and depth >= 1 for
+    child tables (foreign_key_relation links the table to its immediate parent).
+
+  Raises:
+    ValueError: If foreign keys contain cycles, missing tables, or if a child
+      table references more than one parent table (in-degree > 1).
+  """
+  logging.debug(
+      'Computing topological sort for %d tables with %d foreign keys.',
+      len(tables),
+      len(foreign_keys),
+  )
+  table_set = set(tables)
+  graph = nx.DiGraph()
+  graph.add_nodes_from(tables)
+
+  # Maps child_table -> ForeignKeyRelation (incoming edge, e.g.
+  # 'persons' -> fk_household_person).
+  incoming_fk_map: dict[str, ForeignKeyRelation] = {}
+  for fk in foreign_keys:
+    if fk.parent_table not in table_set or fk.child_table not in table_set:
+      raise ValueError(f'Foreign key references unknown table in {fk}.')
+    if fk.child_table in incoming_fk_map:
+      raise ValueError(
+          f'Child table {fk.child_table!r} has multiple parents;'
+          ' in-degree must be <= 1.'
+      )
+    if fk.parent_table == fk.child_table:
+      raise ValueError(f'Self-referential cycle in table {fk.parent_table!r}.')
+    incoming_fk_map[fk.child_table] = fk
+    graph.add_edge(fk.parent_table, fk.child_table)
+
+  if not nx.is_directed_acyclic_graph(graph):
+    raise ValueError('Cycle detected in foreign keys.')
+
+  roots = [t for t in tables if t not in incoming_fk_map]
+  logging.debug('Identified %d root privacy unit tables: %s', len(roots), roots)
+
+  result: list[tuple[int, str, ForeignKeyRelation | None]] = []
+  for depth, generation in enumerate(nx.topological_generations(graph)):
+    for table in generation:
+      result.append((depth, table, incoming_fk_map.get(table)))
+
+  logging.info(
+      'Computed topological synthesis levels: %s',
+      [(d, t) for d, t, _ in result],
+  )
+  return result
+
+
+def _parse_attribute(
+    table_name: str, col_name: str, spec: Any
+) -> domain.AttributeType:
+  """Parses a single attribute specification."""
+  if isinstance(
+      spec,
+      (
+          domain.CategoricalAttribute,
+          domain.NumericalAttribute,
+          domain.OpenSetCategoricalAttribute,
+          domain.FreeFormTextAttribute,
+      ),
+  ):
+    return spec
+  if not isinstance(spec, Mapping):
+    raise ValueError(
+        f'Invalid attribute specification for {table_name}.{col_name}: {spec}'
+    )
+  attr_data = dict(spec)
+  attr_data.pop('type', None)
+  if 'possible_values' in attr_data:
+    return domain.CategoricalAttribute(**attr_data)
+  elif 'min_value' in attr_data:
+    return domain.NumericalAttribute(**attr_data)
+  elif 'max_tokens' in attr_data:
+    return domain.FreeFormTextAttribute(**attr_data)
+  elif 'default_value' in attr_data or not attr_data:
+    return domain.OpenSetCategoricalAttribute(**attr_data)
+  else:
+    raise ValueError(
+        f'Invalid attribute specification for {table_name}.{col_name}:'
+        f' {attr_data}'
+    )
+
+
+def from_dict(
+    config: Mapping[str, Any],
+) -> tuple[dict[str, domain.Schema], list[ForeignKeyRelation]]:
+  """Parses multi-table schema and foreign keys from a dictionary.
+
+  Args:
+    config: Dictionary with 'tables' and optional 'foreign_keys' blocks.
+
+  Returns:
+    A tuple of (table_domains, foreign_keys).
+
+  Raises:
+    ValueError: If configuration format or attribute specifications are invalid.
+  """
+  if 'tables' not in config or not isinstance(config['tables'], Mapping):
+    raise ValueError("'tables' block missing or invalid in config.")
+
+  logging.debug(
+      'Parsing multi-table schema dictionary for %d tables.',
+      len(config['tables']),
+  )
+  table_domains: dict[str, domain.Schema] = {}
+  for table_name, table_schema in config['tables'].items():
+    if not isinstance(table_schema, Mapping):
+      raise ValueError(f'Table schema for {table_name!r} must be a mapping.')
+    table_domains[table_name] = {
+        col_name: _parse_attribute(table_name, col_name, spec)
+        for col_name, spec in table_schema.items()
+    }
+
+  foreign_keys: list[ForeignKeyRelation] = []
+  for fk in config.get('foreign_keys', []):
+    if isinstance(fk, ForeignKeyRelation):
+      foreign_keys.append(fk)
+    elif isinstance(fk, Mapping):
+      foreign_keys.append(ForeignKeyRelation(**fk))
+    else:
+      raise ValueError(f'Invalid foreign key specification: {fk}')
+
+  logging.info(
+      'Successfully parsed multi-table schema: %d tables, %d foreign keys.',
+      len(table_domains),
+      len(foreign_keys),
+  )
+  return table_domains, foreign_keys
+
+
+def from_yaml_file(
+    filepath: str | PathType,
+) -> tuple[dict[str, domain.Schema], list[ForeignKeyRelation]]:
+  """Reads multi-table schema and foreign keys from a YAML file.
+
+  Args:
+    filepath: Path to the YAML schema file.
+
+  Returns:
+    A tuple of (table_domains, foreign_keys).
+  """
+  logging.info('Loading relational domain schema from YAML file: %s', filepath)
+  path = epath.Path(filepath)
+  with path.open('r') as f:
+    config = yaml.safe_load(f)
+  if not isinstance(config, Mapping):
+    raise ValueError(f'YAML root in {filepath} must be a mapping.')
+  return from_dict(config)
+
+
+def to_dict(
+    table_domains: Mapping[str, domain.Schema],
+    foreign_keys: Sequence[ForeignKeyRelation] = (),
+) -> dict[str, Any]:
+  """Converts multi-table schemas and foreign keys to a dictionary.
+
+  Args:
+    table_domains: Mapping from table name to per-column AttributeType schemas.
+    foreign_keys: Optional sequence of ForeignKeyRelation objects.
+
+  Returns:
+    A dictionary with 'tables' and optional 'foreign_keys' blocks.
+  """
+  tables_dict: dict[str, dict[str, Any]] = {}
+  for table_name, schema in table_domains.items():
+    table_dict: dict[str, Any] = {}
+    for col_name, attr in schema.items():
+      attr_dict = dataclasses.asdict(attr)
+      attr_dict['type'] = attr.__class__.__name__
+      table_dict[col_name] = attr_dict
+    tables_dict[table_name] = table_dict
+
+  result: dict[str, Any] = {'tables': tables_dict}
+  if foreign_keys:
+    result['foreign_keys'] = [dataclasses.asdict(fk) for fk in foreign_keys]
+  return result
+
+
+def to_yaml_file(
+    table_domains: Mapping[str, domain.Schema],
+    foreign_keys: Sequence[ForeignKeyRelation],
+    filepath: str | PathType,
+) -> None:
+  """Writes multi-table schema and foreign keys to a YAML file.
+
+  Args:
+    table_domains: Mapping from table name to per-column AttributeType schemas.
+    foreign_keys: Sequence of ForeignKeyRelation objects.
+    filepath: Destination path for the YAML schema file.
+  """
+  logging.info(
+      'Saving relational domain schema (%d tables, %d foreign keys) to: %s',
+      len(table_domains),
+      len(foreign_keys),
+      filepath,
+  )
+  data = to_dict(table_domains, foreign_keys)
+  path = epath.Path(filepath)
+  with path.open('w') as f:
+    yaml.dump(data, f, default_flow_style=False, sort_keys=False)