Add support for auto-detect when a primary key is also a foreign key in Metadata

sdv/metadata/multi_table.py

@@ -541,18 +541,27 @@ def _detect_foreign_keys_by_column_name(self, data):
                 Dictionary of table names to dataframes.
                 NOTE: this is only used in SDV-Enterprise.
         """
-        for parent_candidate in self.tables.keys():
-            primary_key = self.tables[parent_candidate].primary_key
-            for child_candidate in self.tables.keys() - {parent_candidate}:
+        sorted_tables = sorted(self.tables.keys())
+        for parent_candidate in sorted_tables:
+            parent_meta = self.tables[parent_candidate]
+            primary_key = parent_meta.primary_key
+            if primary_key is None:
+                continue
+            for child_candidate in sorted_tables:
+                if child_candidate == parent_candidate:
+                    continue
                 child_meta = self.tables[child_candidate]
-                if primary_key in child_meta.columns.keys():
+                if primary_key in child_meta.columns:
                     try:
-                        original_foreign_key_sdtype = child_meta.columns[primary_key]['sdtype']
+                        original_sdinfo = child_meta.columns[primary_key]
+                        original_foreign_key_sdtype = original_sdinfo['sdtype']
+
                         if original_foreign_key_sdtype != 'id':
                             self.update_column(
-                                table_name=child_candidate, column_name=primary_key, sdtype='id'
+                                table_name=child_candidate,
+                                column_name=primary_key,
+                                sdtype='id',
                             )
-
                         self.add_relationship(
                             parent_candidate, child_candidate, primary_key, primary_key
                         )

tests/integration/metadata/conftest.py

@@ -10,36 +10,36 @@ def primary_key_to_primary_key():
         'tables': {
             'tableA': {
                 'columns': {
-                    'table_A_primary_key': {'sdtype': 'id'},
-                    'column_1': {'sdtype': 'categorical'},
+                    'table_id': {'sdtype': 'id'},
+                    'col1': {'sdtype': 'categorical'},
                 },
-                'primary_key': 'table_A_primary_key',
+                'primary_key': 'table_id',
             },
             'tableB': {
                 'columns': {
-                    'table_B_primary_key': {'sdtype': 'id'},
-                    'column_2': {'sdtype': 'categorical'},
+                    'table_id': {'sdtype': 'id'},
+                    'col2': {'sdtype': 'categorical'},
                 },
-                'primary_key': 'table_B_primary_key',
+                'primary_key': 'table_id',
             },
         },
         'relationships': [
             {
                 'parent_table_name': 'tableA',
-                'parent_primary_key': 'table_A_primary_key',
+                'parent_primary_key': 'table_id',
                 'child_table_name': 'tableB',
-                'child_foreign_key': 'table_B_primary_key',
+                'child_foreign_key': 'table_id',
             }
         ],
     })
     data = {
         'tableA': pd.DataFrame({
-            'table_A_primary_key': range(5),
-            'column_1': ['A', 'B', 'B', 'C', 'C'],
+            'table_id': range(5),
+            'col1': ['A', 'B', 'B', 'C', 'C'],
         }),
         'tableB': pd.DataFrame({
-            'table_B_primary_key': range(5),
-            'column_2': ['A', 'B', 'B', 'C', 'C'],
+            'table_id': range(5),
+            'col2': ['A', 'B', 'B', 'C', 'C'],
         }),
     }
     return data, metadata

tests/integration/metadata/test_metadata.py

@@ -2,6 +2,7 @@
 import re
 from copy import deepcopy
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -879,6 +880,44 @@ def test_detect_from_dataframes_invalid_format():
         Metadata.detect_from_dataframes(data)
 
 
+def test_no_duplicated_foreign_key_relationships_are_generated():
+    # Setup
+    parent_a = pd.DataFrame({
+        'id': ['id-' + str(i) for i in range(100)],
+        'col1': [round(i, 2) for i in np.random.uniform(low=0, high=10, size=100)],
+    })
+    parent_b = pd.DataFrame({
+        'id': ['id-' + str(i) for i in range(100)],
+        'col2': [round(i, 2) for i in np.random.uniform(low=0, high=10, size=100)],
+    })
+
+    child_c = pd.DataFrame({
+        'id': ['id-' + str(i) for i in np.random.randint(0, 100, size=1000)],
+        'col3': [round(i, 2) for i in np.random.uniform(low=0, high=10, size=1000)],
+    })
+
+    data = {'parent_a': parent_a, 'parent_b': parent_b, 'child_c': child_c}
+
+    # Run
+    metadata = Metadata.detect_from_dataframes(data)
+
+    # Assert
+    assert metadata.relationships == [
+        {
+            'parent_table_name': 'parent_a',
+            'child_table_name': 'child_c',
+            'parent_primary_key': 'id',
+            'child_foreign_key': 'id',
+        },
+        {
+            'parent_table_name': 'parent_a',
+            'child_table_name': 'parent_b',
+            'parent_primary_key': 'id',
+            'child_foreign_key': 'id',
+        },
+    ]
+
+
 def test_validate_metadata_with_reused_foreign_keys():
     # Setup
     metadata_dict = {
@@ -1373,44 +1412,48 @@ def test_validate_empty_metadata():
         synthesizer.fit(pd.DataFrame())
 
 
-def test_validate_pk_to_pk(primary_key_to_primary_key):
-    """Test validation to indicate a PK to PK relationship."""
+def test_validate_primary_key_to_primary_key(primary_key_to_primary_key):
+    """Test validate methods with primary key to primary key dataset."""
     # Setup
-    data, metadata_instance = primary_key_to_primary_key
+    data, metadata = primary_key_to_primary_key
+
+    # Run and Assert
+    metadata.validate()
+    metadata.validate_data(data)
+
+
+def test_primary_key_to_primary_key(primary_key_to_primary_key):
+    """Test metadata can auto-detect a primary key which is also a foreign key."""
+    # Setup
+    data, _ = primary_key_to_primary_key
 
     # Run
-    metadata_instance.validate()
-    metadata_instance.validate_data(data)
+    metadata = Metadata.detect_from_dataframes(data)
 
     # Assert
-    expected_metadata = {
+    metadata.validate()
+    metadata.validate_data(data)
+    assert metadata.to_dict() == {
         'tables': {
             'tableA': {
-                'columns': {
-                    'table_A_primary_key': {'sdtype': 'id'},
-                    'column_1': {'sdtype': 'categorical'},
-                },
-                'primary_key': 'table_A_primary_key',
+                'columns': {'table_id': {'sdtype': 'id'}, 'col1': {'sdtype': 'categorical'}},
+                'primary_key': 'table_id',
             },
             'tableB': {
-                'columns': {
-                    'table_B_primary_key': {'sdtype': 'id'},
-                    'column_2': {'sdtype': 'categorical'},
-                },
-                'primary_key': 'table_B_primary_key',
+                'columns': {'table_id': {'sdtype': 'id'}, 'col2': {'sdtype': 'categorical'}},
+                'primary_key': 'table_id',
             },
         },
         'relationships': [
             {
                 'parent_table_name': 'tableA',
-                'parent_primary_key': 'table_A_primary_key',
                 'child_table_name': 'tableB',
-                'child_foreign_key': 'table_B_primary_key',
+                'parent_primary_key': 'table_id',
+                'child_foreign_key': 'table_id',
             }
         ],
         'METADATA_SPEC_VERSION': 'V1',
     }
-    assert metadata_instance.to_dict() == expected_metadata
 
 
 def test_validate_pk_to_pk_email():

tests/unit/metadata/test_metadata.py

@@ -775,6 +775,75 @@ def test_detect_from_dataframe_bad_input_infer_keys(self):
         with pytest.raises(ValueError, match=expected_message):
             Metadata.detect_from_dataframe(data, infer_keys=infer_keys)
 
+    def test_detect_from_dataframe_primary_key_to_primary_key(self):
+        """Test primary to primary key relationship is detected if column name match."""
+        # Setup
+        data = {
+            'table1': pd.DataFrame({
+                'id': [1, 2, 3],
+            }),
+            'table2': pd.DataFrame({
+                'id': [1, 2, 3],
+            }),
+        }
+        instance = Metadata()
+        instance.detect_table_from_dataframe('table1', data['table1'])
+        instance.detect_table_from_dataframe('table2', data['table2'])
+
+        # Run
+        instance._detect_foreign_keys_by_column_name(data)
+
+        # Assert
+        assert instance.to_dict()['relationships'] == [
+            {
+                'parent_table_name': 'table1',
+                'child_table_name': 'table2',
+                'parent_primary_key': 'id',
+                'child_foreign_key': 'id',
+            }
+        ]
+
+    def test_detect_foreign_keys_sorting(self):
+        # Setup
+        mock_metadata_1 = Mock()
+        mock_metadata_1.primary_key = 'id'
+        mock_metadata_1.columns = {
+            'id': {'sdtype': 'id'},
+            'col1': {'sdtype': 'categorical'},
+        }
+        mock_metadata_2 = Mock()
+        mock_metadata_2.primary_key = 'id'
+        mock_metadata_2.columns = {
+            'id': {'sdtype': 'id'},
+            'col2': {'sdtype': 'numerical'},
+        }
+        data = {
+            'table1': pd.DataFrame({
+                'id': [1, 2, 3],
+                'col1': ['a', 'b', 'c'],
+            }),
+            'table2': pd.DataFrame({
+                'id': [1, 2, 3],
+                'col2': [1.1, 2.2, 3.3],
+            }),
+        }
+        instance = Metadata()
+        instance.tables = {
+            'table1': mock_metadata_1,
+            'table2': mock_metadata_2,
+        }
+        instance.add_relationship = Mock()
+
+        # Run
+        instance._detect_foreign_keys_by_column_name(data=data)
+
+        # Assert
+        expected_calls = [
+            call('table1', 'table2', 'id', 'id'),
+            call('table2', 'table1', 'id', 'id'),
+        ]
+        instance.add_relationship.assert_has_calls(expected_calls, any_order=False)
+
     def test__handle_table_name(self):
         """Test the ``_handle_table_name`` method."""
         # Setup