Add support for composite primary keys in the metadata

Author: frances-h

sdv/metadata/multi_table.py

@@ -183,9 +183,14 @@ def _validate_new_foreign_key_is_not_reused(
                 and relationship['parent_primary_key'] == parent_primary_key
             )
             if foreign_key_already_used and not parent_matches:
+                child_foreign_key = (
+                    f"('{child_foreign_key}')"
+                    if isinstance(child_foreign_key, str)
+                    else f'({child_foreign_key})'
+                )
                 raise InvalidMetadataError(
                     f'Relationship between tables ({parent_table_name}, {child_table_name}) uses '
-                    f"a foreign key column ('{child_foreign_key}') that is already used in another "
+                    f'a foreign key {child_foreign_key} that is already used in another '
                     'relationship.'
                 )
 
@@ -197,15 +202,23 @@ def _validate_foreign_key_uniqueness_across_relationships(
         child_foreign_key,
         seen_foreign_keys,
     ):
-        key = (child_table_name, child_foreign_key)
+        key = (
+            tuple(_cast_to_iterable(child_table_name)),
+            tuple(_cast_to_iterable(child_foreign_key)),
+        )
         current_relationship = (parent_table_name, parent_primary_key)
 
         if key in seen_foreign_keys:
             existing_relationship = seen_foreign_keys[key]
             if existing_relationship != current_relationship:
+                child_foreign_key = (
+                    f"('{child_foreign_key}')"
+                    if isinstance(child_foreign_key, str)
+                    else f'({child_foreign_key})'
+                )
                 raise InvalidMetadataError(
                     f'Relationship between tables ({parent_table_name}, {child_table_name}) uses '
-                    f"a foreign key column ('{child_foreign_key}') that is already used in another "
+                    f'a foreign key {child_foreign_key} that is already used in another '
                     'relationship.'
                 )
         else:
@@ -296,10 +309,10 @@ def add_relationship(
                 A string representing the name of the parent table.
             child_table_name (str):
                 A string representing the name of the child table.
-            parent_primary_key (str or tuple):
-                A string or tuple of strings representing the primary key of the parent.
-            child_foreign_key (str or tuple):
-                A string or tuple of strings representing the foreign key of the child.
+            parent_primary_key (str or list[str]):
+                A string or list of strings representing the primary key of the parent.
+            child_foreign_key (str or list[str]):
+                A string or list of strings representing the foreign key of the child.
 
         Raises:
             - ``InvalidMetadataError`` if a table is missing.
@@ -1206,9 +1219,19 @@ def _set_metadata_dict(self, metadata):
                 ) from error
 
         for relationship in metadata.get('relationships', []):
+            parent_pk = relationship.get('parent_primary_key')
+            child_fk = relationship.get('child_foreign_key')
+            type_safe_pk = (
+                [str(col) for col in parent_pk] if isinstance(parent_pk, list) else str(parent_pk)
+            )
+            type_safe_fk = (
+                [str(col) for col in child_fk] if isinstance(parent_pk, list) else str(child_fk)
+            )
             type_safe_relationships = {
-                key: str(value) if not isinstance(value, str) else value
-                for key, value in relationship.items()
+                'parent_table_name': str(relationship.get('parent_table_name')),
+                'child_table_name': str(relationship.get('child_table_name')),
+                'parent_primary_key': type_safe_pk,
+                'child_foreign_key': type_safe_fk,
             }
             self.relationships.append(type_safe_relationships)
 

sdv/metadata/single_table.py

@@ -211,14 +211,33 @@ def _validate_pii(column_name, **kwargs):
 
     def __init__(self):
         self.columns = {}
-        self.primary_key = None
+        self._primary_key = None
         self.alternate_keys = []
         self.sequence_key = None
         self.sequence_index = None
         self.column_relationships = []
         self._version = self.METADATA_SPEC_VERSION
         self._updated = False
 
+    @property
+    def _primary_key_is_composite(self):
+        if self.primary_key and isinstance(self.primary_key, list) and len(self.primary_key) > 1:
+            return True
+
+        return False
+
+    @property
+    def primary_key(self):
+        """Property to handle singleton composite key case."""
+        if isinstance(self._primary_key, list) and len(self._primary_key) == 1:
+            return self._primary_key[0]
+
+        return self._primary_key
+
+    @primary_key.setter
+    def primary_key(self, primary_key):
+        self._primary_key = primary_key
+
     def _get_unexpected_kwargs(self, sdtype, **kwargs):
         expected_kwargs = self._SDTYPE_KWARGS.get(sdtype, ['pii'])
         unexpected_kwargs = set(kwargs) - set(expected_kwargs)
@@ -774,29 +793,40 @@ def detect_from_csv(self, filepath, read_csv_parameters=None):
         self.detect_from_dataframe(data)
 
     @staticmethod
-    def _validate_key_datatype(column_name):
+    def _validate_key_datatype(column_name, key_type):
         """Check whether column_name is a string."""
-        return isinstance(column_name, str)
+        is_string = isinstance(column_name, str)
+        is_list_of_strings = isinstance(column_name, list) and all(
+            isinstance(i, str) for i in column_name
+        )
+        return is_string or (key_type == 'primary' and is_list_of_strings)
 
     def _validate_keys_sdtype(self, keys, key_type):
         """Validate that each key is of type 'id' or a valid Faker function."""
-        bad_keys = set()
+        bad_keys = []
         for key in keys:
-            if not (
-                self.columns[key]['sdtype'] == 'id'
-                or is_faker_function(self.columns[key]['sdtype'])
+            if not any(
+                self.columns[key_col]['sdtype'] == 'id'
+                or is_faker_function(self.columns[key_col]['sdtype'])
+                for key_col in _cast_to_iterable(key)
             ):
-                bad_keys.add(key)
+                bad_keys.append(key)
+
         if bad_keys:
             raise InvalidMetadataError(
-                f"The {key_type}_keys {sorted(bad_keys)} must be type 'id' or another PII type."
+                f'The {key_type}_keys {bad_keys} must have a column of '
+                "type 'id' or another PII type."
             )
 
     def _validate_key(self, column_name, key_type):
         """Validate the primary and sequence keys."""
         if column_name is not None:
-            if not self._validate_key_datatype(column_name):
-                raise InvalidMetadataError(f"'{key_type}_key' must be a string.")
+            if not self._validate_key_datatype(column_name, key_type):
+                err_msg = f"'{key_type}_key' must be a string"
+                if key_type == 'primary':
+                    err_msg += ' or a list of strings'
+
+                raise InvalidMetadataError(err_msg + '.')
 
             keys = {column_name} if isinstance(column_name, str) else set(column_name)
             setting_sequence_as_primary = key_type == 'primary' and column_name == self.sequence_key
@@ -814,7 +844,7 @@ def _validate_key(self, column_name, key_type):
                     ' Keys should be columns that exist in the table.'
                 )
 
-            self._validate_keys_sdtype(keys, key_type)
+            self._validate_keys_sdtype([column_name], key_type)
 
     def set_primary_key(self, column_name):
         """Set the metadata primary key.
@@ -866,7 +896,8 @@ def set_sequence_key(self, column_name):
 
     def _validate_alternate_keys(self, column_names):
         if not isinstance(column_names, list) or not all(
-            self._validate_key_datatype(column_name) for column_name in column_names
+            self._validate_key_datatype(column_name, 'alternate_keys')
+            for column_name in column_names
         ):
             raise InvalidMetadataError("'alternate_keys' must be a list of strings.")
 
@@ -1158,7 +1189,10 @@ def _get_primary_and_alternate_keys(self):
         """
         keys = set(self.alternate_keys)
         if self.primary_key:
-            keys.update({self.primary_key})
+            primary_key = (
+                tuple(self.primary_key) if isinstance(self.primary_key, list) else self.primary_key
+            )
+            keys.add(primary_key)
 
         return keys
 
@@ -1181,31 +1215,45 @@ def _validate_keys_dont_have_missing_values(self, data):
         errors = []
         keys = self._get_primary_and_alternate_keys()
         keys.update(self._get_set_of_sequence_keys())
-        for key in sorted(keys):
-            if pd.isna(data[key]).any():
-                errors.append(f"Key column '{key}' contains missing values.")
+        for key in sorted(keys, key=lambda key: key if isinstance(key, str) else key[0]):
+            key_list = [key] if isinstance(key, str) else list(key)
+            if pd.isna(data[key_list]).all(axis=1).any():
+                key = f"'{key}'" if isinstance(key, str) else f'{key}'
+                errors.append(f'Key column {key} contains missing values.')
 
         return errors
 
     def _validate_key_values_are_unique(self, data):
         errors = []
         keys = self._get_primary_and_alternate_keys()
-        for key in sorted(keys):
-            repeated_values = set(data[key][data[key].duplicated()])
-            if repeated_values:
-                repeated_values = _format_invalid_values_string(repeated_values, 3)
-                errors.append(f"Key column '{key}' contains repeating values: " + repeated_values)
+        for key in sorted(keys, key=lambda key: key if isinstance(key, str) else key[0]):
+            key_list = [key] if isinstance(key, str) else list(key)
+            repeated_values = data[key_list][data[key_list].duplicated()]
+            if not repeated_values.empty:
+                if len(repeated_values.columns) == 1:
+                    repeated_values = ' ' + _format_invalid_values_string(
+                        set(repeated_values[key]), 3
+                    )
+                else:
+                    repeated_values = '\n' + _format_invalid_values_string(
+                        repeated_values.drop_duplicates(), 3
+                    )
+
+                key = f"'{key}'" if isinstance(key, str) else f'{key}'
+                errors.append(f'Key column {key} contains repeating values:' + repeated_values)
 
         return errors
 
     def _validate_primary_key(self, data):
         error = []
-        is_int = self.primary_key and pd.api.types.is_integer_dtype(data[self.primary_key])
-        regex = self.columns.get(self.primary_key, {}).get('regex_format')
-        if is_int and regex:
-            possible_characters = get_possible_chars(regex, 1)
-            if '0' in possible_characters:
-                error.append(f'Primary key "{self.primary_key}" {INT_REGEX_ZERO_ERROR_MESSAGE}')
+        primary_key_list = _cast_to_iterable(self.primary_key) if self.primary_key else []
+        for key in primary_key_list:
+            is_int = pd.api.types.is_integer_dtype(data[key])
+            regex = self.columns.get(key, {}).get('regex_format')
+            if is_int and regex:
+                possible_characters = get_possible_chars(regex, 1)
+                if '0' in possible_characters:
+                    error.append(f'Primary key column "{key}" {INT_REGEX_ZERO_ERROR_MESSAGE}')
 
         return error
 

sdv/multi_table/base.py

@@ -114,6 +114,20 @@ def _check_metadata_updated(self):
                 ' in future SDV versions.'
             )
 
+    def _handle_composite_keys(self):
+        """Validates that composite keys are not used in Public SDV."""
+        composite_key_tables = []
+        for table, table_metadata in self.metadata.tables.items():
+            if table_metadata._primary_key_is_composite:
+                composite_key_tables.append(table)
+
+        if composite_key_tables:
+            raise SynthesizerInputError(
+                'Your metadata contains composite keys (primary key of tables '
+                f'{composite_key_tables} have multiple columns). Composite keys are '
+                'not supported in SDV Community.'
+            )
+
     def __init__(self, metadata, locales=['en_US'], synthesizer_kwargs=None):
         self.metadata = metadata
         if type(metadata) is MultiTableMetadata:
@@ -123,6 +137,7 @@ def __init__(self, metadata, locales=['en_US'], synthesizer_kwargs=None):
             self.metadata.validate()
 
         self._check_metadata_updated()
+        self._handle_composite_keys()
         self.locales = locales
         self.verbose = False
         self.extended_columns = defaultdict(dict)

sdv/single_table/base.py

@@ -150,6 +150,15 @@ def _validate_regex_format(self):
                 )
                 _check_regex_format(self._table_name, column_name, regex)
 
+    def _handle_composite_keys(self, single_table_metadata):
+        """Validates that composite keys are not used in Public SDV."""
+        if single_table_metadata._primary_key_is_composite:
+            raise SynthesizerInputError(
+                'Your metadata contains composite keys (primary key of table '
+                f"'{self._table_name}' has multiple columns). Composite keys are "
+                'not supported in SDV Community.'
+            )
+
     def __init__(
         self, metadata, enforce_min_max_values=True, enforce_rounding=True, locales=['en_US']
     ):
@@ -172,6 +181,8 @@ def __init__(
 
         self.metadata.validate()
         self._check_metadata_updated()
+        single_table_metadata = self.metadata._convert_to_single_table()
+        self._handle_composite_keys(single_table_metadata)
 
         # Points to a metadata object that conserves the initialized status of the synthesizer
         self._original_metadata = deepcopy(self.metadata)
@@ -180,7 +191,7 @@ def __init__(
         self.enforce_rounding = enforce_rounding
         self.locales = locales
         self._data_processor = DataProcessor(
-            metadata=self.metadata._convert_to_single_table(),
+            metadata=single_table_metadata,
             enforce_rounding=self.enforce_rounding,
             enforce_min_max_values=self.enforce_min_max_values,
             locales=self.locales,

tests/integration/metadata/test_metadata.py

@@ -993,7 +993,7 @@ def test_validate_metadata_with_reused_foreign_keys():
     # Run and Assert
     error_msg = re.escape(
         'Relationships:\n'
-        'Relationship between tables (A2, A3) uses a foreign key column '
+        'Relationship between tables (A2, A3) uses a foreign key '
         "('fk3_A1_A2') that is already used in another relationship."
     )
     with pytest.raises(InvalidMetadataError, match=error_msg):

tests/integration/metadata/test_multi_table.py

@@ -22,6 +22,47 @@ def test_multi_table_metadata():
     assert instance.relationships == []
 
 
+def test_multi_table_metadata_composite_keys():
+    """Test ``MultiTableMetadata`` with composite keys."""
+    # Setup
+    metadata_dict = {
+        'tables': {
+            'table1': {
+                'columns': {
+                    'table1_id': {'sdtype': 'id'},
+                    'cat_col': {'sdtype': 'categorical'},
+                },
+                'primary_key': ['table1_id', 'cat_col'],
+            },
+            'table2': {
+                'columns': {
+                    'pk': {'sdtype': 'id'},
+                    'fk1': {'sdtype': 'id'},
+                    'fk2': {'sdtype': 'categorical'},
+                },
+                'primary_key': 'pk',
+            },
+        },
+        'relationships': [
+            {
+                'parent_table_name': 'table1',
+                'parent_primary_key': ['table1_id', 'cat_col'],
+                'child_table_name': 'table2',
+                'child_foreign_key': ['fk1', 'fk2'],
+            },
+        ],
+    }
+
+    # Run
+    instance = MultiTableMetadata.load_from_dict(metadata_dict)
+    result = instance.to_dict()
+
+    # Assert
+    instance.validate()
+    assert result == {**metadata_dict, 'METADATA_SPEC_VERSION': 'MULTI_TABLE_V1'}
+    assert instance.relationships == metadata_dict['relationships']
+
+
 @patch('rdt.transformers')
 def test_add_column_relationship(mock_rdt_transformers):
     """Test ``add_column_relationship`` method."""