wip

Author: gsheni

sdv/multi_table/hma.py

@@ -779,8 +779,19 @@ def _find_parent_ids(self, child_table, parent_table, child_name, parent_name, f
         return likelihoods.apply(self._find_parent_id, axis=1, num_rows=num_rows)
 
     def _add_foreign_key_columns(self, child_table, parent_table, child_name, parent_name):
+        parent_primary_key = self.metadata.tables[parent_name].primary_key
+        parent_id_values = parent_table[parent_primary_key].dropna().unique()
         for foreign_key in self.metadata._get_foreign_keys(parent_name, child_name):
-            if foreign_key not in child_table:
+            needs_assignment = True
+            if foreign_key in child_table:
+                child_column = child_table[foreign_key]
+                if not child_column.dropna().empty:
+                    # check if child column (FK) contains IDs in parent table
+                    is_valid = child_column.dropna().isin(parent_id_values).all()
+                    if is_valid:
+                        needs_assignment = False
+
+            if needs_assignment:
                 parent_ids = self._find_parent_ids(
                     child_table=child_table,
                     parent_table=parent_table,

tests/integration/multi_table/conftest.py

@@ -234,8 +234,6 @@ def data_metadata_multiple_foreign_keys():
     })
     assert data['child']['parent_1_id'].equals(data['parent']['parent_id'])
     assert data['child']['parent_2_id'].equals(data['second_parent']['parent_id'])
-    metadata.validate()
-    metadata.validate_data(data)
     return data, metadata
 
 
@@ -259,6 +257,4 @@ def data_metadata_multiple_foreign_keys_subset(data_metadata_multiple_foreign_ke
     }
     assert set(data['child']['parent_1_id']).issubset(set(data['parent']['parent_id']))
     assert set(data['child']['parent_2_id']).issubset(set(data['second_parent']['parent_id']))
-    metadata.validate()
-    metadata.validate_data(data)
     return data, metadata

tests/integration/multi_table/test_hma.py

@@ -2961,14 +2961,16 @@ def test_multiple_fks(self, data_metadata_multiple_foreign_keys):
         synthetic_data = synthesizer.sample(scale=1.0)
 
         # Assert
-        for each_parent_id in synthetic_data['child']['parent_1_id'].tolist():
-            assert each_parent_id in set(synthetic_data['parent']['parent_id'])
-        for each_parent_id in synthetic_data['child']['parent_2_id'].tolist():
-            assert each_parent_id in set(synthetic_data['second_parent']['parent_id'])
+        assert set(synthetic_data['child']['parent_1_id']).issubset(
+            synthetic_data['parent']['parent_id']
+        )
+        assert set(synthetic_data['child']['parent_2_id']).issubset(
+            synthetic_data['second_parent']['parent_id']
+        )
         synthesizer.validate(synthetic_data)
 
-    def test_multiple_fks_mismatched(self, data_metadata_multiple_foreign_keys_subset):
-        """Test support for parent and child with multiple foreign keys."""
+    def test_multiple_fks_subset(self, data_metadata_multiple_foreign_keys_subset):
+        """Test support for parent and child with multiple foreign keys (subset in child)."""
         # Setup
         data, metadata = data_metadata_multiple_foreign_keys_subset
 
@@ -2978,8 +2980,10 @@ def test_multiple_fks_mismatched(self, data_metadata_multiple_foreign_keys_subse
         synthetic_data = synthesizer.sample(scale=1.0)
 
         # Assert
-        for each_parent_id in synthetic_data['child']['parent_1_id'].tolist():
-            assert each_parent_id in set(synthetic_data['parent']['parent_id'])
-        for each_parent_id in synthetic_data['child']['parent_2_id'].tolist():
-            assert each_parent_id in set(synthetic_data['second_parent']['parent_id'])
+        assert set(synthetic_data['child']['parent_1_id']).issubset(
+            synthetic_data['parent']['parent_id']
+        )
+        assert set(synthetic_data['child']['parent_2_id']).issubset(
+            synthetic_data['second_parent']['parent_id']
+        )
         synthesizer.validate(synthetic_data)

tests/unit/multi_table/test_hma.py

@@ -869,9 +869,17 @@ def test__add_foreign_key_columns(self):
         """Test that the ``_add_foreign_key_columns`` method adds foreign keys."""
         # Setup
         instance = Mock()
-        metadata = Mock()
-        metadata._get_foreign_keys.return_value = ['primary_user_id', 'secondary_user_id']
-        instance.metadata = metadata
+        mock_users_metadata = Mock()
+        mock_users_metadata.primary_key = 'user_id'
+        mock_transactions_metadata = Mock()
+        mock_transactions_metadata.primary_key = 'transaction_id'
+        mock_metadata = Mock()
+        mock_metadata.tables = {
+            'users': mock_users_metadata,
+            'transactions': mock_transactions_metadata,
+        }
+        mock_metadata._get_foreign_keys.return_value = ['primary_user_id', 'secondary_user_id']
+        instance.metadata = mock_metadata
 
         instance._find_parent_ids.return_value = pd.Series([2, 1, 2], name='secondary_user_id')