diff --git a/dpsynth/relational/domain.py b/dpsynth/relational/domain.py index edd4cb4..e4235cc 100644 --- a/dpsynth/relational/domain.py +++ b/dpsynth/relational/domain.py @@ -20,16 +20,14 @@ 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 -# pylint: disable=unused-import -_YAML_UNUSED = yaml -# pylint: enable=unused-import - @dataclasses.dataclass(frozen=True) class ForeignKeyRelation: @@ -80,10 +78,90 @@ def topological_sort_hierarchy( ValueError: If foreign keys contain cycles, missing tables, or if a child table references more than one parent table (in-degree > 1). """ - del tables, foreign_keys - raise NotImplementedError( - 'topological_sort_hierarchy is not yet implemented.' + 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 + + +_ATTRIBUTE_TYPE_MAP: Mapping[str, type[domain.AttributeType]] = { + 'CategoricalAttribute': domain.CategoricalAttribute, + 'NumericalAttribute': domain.NumericalAttribute, + 'OpenSetCategoricalAttribute': domain.OpenSetCategoricalAttribute, + 'FreeFormTextAttribute': domain.FreeFormTextAttribute, +} + + +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) + if 'type' not in attr_data: + raise ValueError( + f'Attribute specification for {table_name}.{col_name} is missing' + " required 'type' field." + ) + attr_type_name = attr_data.pop('type') + attr_cls = _ATTRIBUTE_TYPE_MAP.get(attr_type_name) + if attr_cls is None: + raise ValueError( + f'Unknown attribute type {attr_type_name!r} for' + f' {table_name}.{col_name}. Expected one of' + f' {list(_ATTRIBUTE_TYPE_MAP.keys())}.' + ) + return attr_cls(**attr_data) def from_dict( @@ -100,8 +178,37 @@ def from_dict( Raises: ValueError: If configuration format or attribute specifications are invalid. """ - del config - raise NotImplementedError('from_dict is not yet implemented.') + 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( @@ -115,5 +222,62 @@ def from_yaml_file( Returns: A tuple of (table_domains, foreign_keys). """ - del filepath - raise NotImplementedError('from_yaml_file is not yet implemented.') + 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) diff --git a/tests/relational/domain_test.py b/tests/relational/domain_test.py index bbd620d..20aaf3a 100644 --- a/tests/relational/domain_test.py +++ b/tests/relational/domain_test.py @@ -14,7 +14,10 @@ """Unit tests for dpsynth.relational.domain.""" +import textwrap + from absl.testing import absltest +from dpsynth import domain as base_domain from dpsynth.relational import domain @@ -46,6 +49,312 @@ def test_foreign_key_relation_invalid_capacity(self): max_children_per_parent=0, ) + def test_topological_sort_linear_chain(self): + tables = ['households', 'persons', 'activities'] + fk1 = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk2 = domain.ForeignKeyRelation('persons', 'pid', 'activities', 'pid', 2) + order = domain.topological_sort_hierarchy(tables, [fk1, fk2]) + self.assertEqual( + order, + [ + (0, 'households', None), + (1, 'persons', fk1), + (2, 'activities', fk2), + ], + ) + + def test_topological_sort_branching(self): + tables = ['households', 'persons', 'vehicles'] + fk_p = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk_v = domain.ForeignKeyRelation('households', 'hid', 'vehicles', 'hid', 2) + order = domain.topological_sort_hierarchy(tables, [fk_p, fk_v]) + self.assertEqual(order[0], (0, 'households', None)) + self.assertCountEqual( + order[1:], [(1, 'persons', fk_p), (1, 'vehicles', fk_v)] + ) + + def test_topological_sort_forest(self): + tables = ['households', 'persons', 'companies', 'departments'] + fk_h = domain.ForeignKeyRelation('households', 'hid', 'persons', 'hid', 3) + fk_c = domain.ForeignKeyRelation( + 'companies', 'cid', 'departments', 'cid', 10 + ) + order = domain.topological_sort_hierarchy(tables, [fk_h, fk_c]) + depths = {t: d for d, t, _ in order} + self.assertEqual(depths['households'], 0) + self.assertEqual(depths['companies'], 0) + self.assertEqual(depths['persons'], 1) + self.assertEqual(depths['departments'], 1) + + def test_topological_sort_single_table_and_empty(self): + self.assertEqual( + domain.topological_sort_hierarchy(['households'], []), + [(0, 'households', None)], + ) + self.assertEqual(domain.topological_sort_hierarchy([], []), []) + + def test_topological_sort_cycle_raises_error(self): + tables = ['a', 'b'] + fk1 = domain.ForeignKeyRelation('a', 'id', 'b', 'id', 1) + fk2 = domain.ForeignKeyRelation('b', 'id', 'a', 'id', 1) + with self.assertRaisesRegex(ValueError, 'Cycle detected'): + domain.topological_sort_hierarchy(tables, [fk1, fk2]) + + def test_topological_sort_self_cycle_raises_error(self): + tables = ['a'] + fk = domain.ForeignKeyRelation('a', 'id', 'a', 'id', 1) + with self.assertRaisesRegex(ValueError, 'Self-referential cycle'): + domain.topological_sort_hierarchy(tables, [fk]) + + def test_topological_sort_multiple_parents_raises_error(self): + tables = ['p1', 'p2', 'c'] + fk1 = domain.ForeignKeyRelation('p1', 'id', 'c', 'p1_id', 1) + fk2 = domain.ForeignKeyRelation('p2', 'id', 'c', 'p2_id', 1) + with self.assertRaisesRegex(ValueError, 'multiple parents'): + domain.topological_sort_hierarchy(tables, [fk1, fk2]) + + def test_topological_sort_unknown_table_raises_error(self): + tables = ['households'] + fk = domain.ForeignKeyRelation('households', 'hid', 'unknown', 'hid', 1) + with self.assertRaisesRegex(ValueError, 'unknown table'): + domain.topological_sort_hierarchy(tables, [fk]) + + def test_from_dict_valid_3tier_schema(self): + config = { + 'tables': { + 'households': { + 'income': { + 'type': 'NumericalAttribute', + 'min_value': 0.0, + 'max_value': 200000.0, + }, + 'region': { + 'type': 'CategoricalAttribute', + 'possible_values': ['Urban', 'Rural'], + }, + }, + 'persons': { + 'age': { + 'type': 'NumericalAttribute', + 'min_value': 0, + 'max_value': 100, + 'dtype': 'int', + }, + 'gender': { + 'type': 'CategoricalAttribute', + 'possible_values': ['M', 'F'], + }, + }, + }, + 'foreign_keys': [{ + 'parent_table': 'households', + 'parent_primary_key': 'household_id', + 'child_table': 'persons', + 'child_foreign_key': 'household_id', + 'max_children_per_parent': 3, + }], + } + table_domains, fks = domain.from_dict(config) + self.assertIn('households', table_domains) + self.assertIn('persons', table_domains) + self.assertIsInstance( + table_domains['households']['income'], base_domain.NumericalAttribute + ) + self.assertIsInstance( + table_domains['households']['region'], + base_domain.CategoricalAttribute, + ) + self.assertLen(fks, 1) + self.assertEqual(fks[0].parent_table, 'households') + self.assertEqual(fks[0].parent_primary_key, 'household_id') + self.assertEqual(fks[0].child_table, 'persons') + self.assertEqual(fks[0].child_foreign_key, 'household_id') + self.assertEqual(fks[0].max_children_per_parent, 3) + + def test_from_dict_missing_tables_block_raises_error(self): + with self.assertRaisesRegex(ValueError, "'tables' block missing"): + domain.from_dict({}) + + def test_from_dict_missing_type_field_raises_error(self): + config = { + 'tables': { + 'households': {'income': {'min_value': 0.0, 'max_value': 200000.0}} + } + } + with self.assertRaisesRegex(ValueError, "missing required 'type' field"): + domain.from_dict(config) + + def test_from_dict_unknown_type_field_raises_error(self): + config = { + 'tables': { + 'households': { + 'income': { + 'type': 'InvalidType', + 'min_value': 0.0, + 'max_value': 100.0, + } + } + } + } + with self.assertRaisesRegex( + ValueError, "Unknown attribute type 'InvalidType'" + ): + domain.from_dict(config) + + def test_from_dict_invalid_table_schema_raises_error(self): + with self.assertRaisesRegex(ValueError, 'must be a mapping'): + domain.from_dict({'tables': {'households': 'invalid'}}) + + def test_from_dict_invalid_attribute_spec_raises_error(self): + with self.assertRaisesRegex(ValueError, 'Invalid attribute specification'): + domain.from_dict({'tables': {'households': {'income': 123}}}) + + def test_from_dict_invalid_foreign_key_spec_raises_error(self): + with self.assertRaisesRegex(ValueError, 'Invalid foreign key'): + domain.from_dict({'tables': {'h': {}}, 'foreign_keys': [123]}) + + def test_from_yaml_file_roundtrip(self): + yaml_content = textwrap.dedent("""\ + tables: + households: + income: + type: NumericalAttribute + min_value: 0.0 + max_value: 100000.0 + region: + type: CategoricalAttribute + possible_values: ["Urban", "Rural"] + persons: + age: + type: NumericalAttribute + min_value: 0 + max_value: 100 + dtype: int + gender: + type: CategoricalAttribute + possible_values: ["M", "F"] + foreign_keys: + - parent_table: households + parent_primary_key: hid + child_table: persons + child_foreign_key: hid + max_children_per_parent: 4 + """) + tmp_path = self.create_tempfile(content=yaml_content).full_path + table_domains, fks = domain.from_yaml_file(tmp_path) + self.assertIn('households', table_domains) + self.assertIn('persons', table_domains) + self.assertIsInstance( + table_domains['persons']['age'], base_domain.NumericalAttribute + ) + self.assertIsInstance( + table_domains['persons']['gender'], base_domain.CategoricalAttribute + ) + self.assertLen(fks, 1) + self.assertEqual(fks[0].parent_table, 'households') + self.assertEqual(fks[0].parent_primary_key, 'hid') + self.assertEqual(fks[0].child_table, 'persons') + self.assertEqual(fks[0].child_foreign_key, 'hid') + self.assertEqual(fks[0].max_children_per_parent, 4) + + def test_to_dict_and_roundtrip(self): + table_domains = { + 'households': { + 'income': base_domain.NumericalAttribute( + min_value=0.0, max_value=200000.0, dtype='float' + ), + 'region': base_domain.CategoricalAttribute( + possible_values=['Urban', 'Rural'] + ), + }, + 'persons': { + 'age': base_domain.NumericalAttribute( + min_value=0, max_value=100, dtype='int' + ), + 'gender': base_domain.CategoricalAttribute( + possible_values=['M', 'F'] + ), + }, + } + fks = [ + domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='hid', + child_table='persons', + child_foreign_key='hid', + max_children_per_parent=3, + ) + ] + serialized = domain.to_dict(table_domains, fks) + self.assertIn('tables', serialized) + self.assertIn('foreign_keys', serialized) + self.assertLen(serialized['foreign_keys'], 1) + self.assertEqual( + serialized['tables']['households']['income']['type'], + 'NumericalAttribute', + ) + + # Roundtrip verification + rt_domains, rt_fks = domain.from_dict(serialized) + self.assertEqual( + rt_domains['households']['income'].min_value, + table_domains['households']['income'].min_value, + ) + self.assertEqual( + rt_domains['persons']['gender'].possible_values, + table_domains['persons']['gender'].possible_values, + ) + self.assertEqual(rt_fks, fks) + + def test_to_yaml_file_and_roundtrip(self): + table_domains = { + 'households': { + 'income': base_domain.NumericalAttribute( + min_value=0.0, max_value=150000.0 + ), + 'region': base_domain.CategoricalAttribute( + possible_values=['Urban', 'Suburban'] + ), + }, + 'persons': { + 'age': base_domain.NumericalAttribute( + min_value=0, max_value=120, dtype='int' + ), + }, + } + fks = [ + domain.ForeignKeyRelation( + parent_table='households', + parent_primary_key='household_id', + child_table='persons', + child_foreign_key='household_id', + max_children_per_parent=5, + ) + ] + tmp_path = self.create_tempfile().full_path + domain.to_yaml_file(table_domains, fks, tmp_path) + + rt_domains, rt_fks = domain.from_yaml_file(tmp_path) + self.assertIn('households', rt_domains) + self.assertIn('persons', rt_domains) + self.assertEqual(rt_domains['households']['income'].max_value, 150000.0) + self.assertEqual(rt_fks, fks) + + def test_to_dict_without_foreign_keys(self): + table_domains = { + 'single_table': { + 'col_a': base_domain.CategoricalAttribute( + possible_values=['x', 'y'] + ) + } + } + serialized = domain.to_dict(table_domains) + self.assertIn('tables', serialized) + self.assertNotIn('foreign_keys', serialized) + rt_domains, rt_fks = domain.from_dict(serialized) + self.assertIn('single_table', rt_domains) + self.assertEmpty(rt_fks) + if __name__ == '__main__': absltest.main()