use existing logic

Author: gsheni

sdv/multi_table/hma.py

@@ -779,18 +779,38 @@ 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):
+        """Add foreign key columns in the child table.
+
+        This function adds foreign key columns to a child table.
+        If the foreign key column does not exist in the child table, it adds the column.
+        If the foreign key column already exists in the child table (e.g., when it is also a PK)
+        and it contains invalid references (FKs not found in parent table), it overwrites the
+        foreign key values (from the parent table's PK).
+
+        Args:
+            child_table (pd.DataFrame): The child table which may or may not contain the FK columns.
+            parent_table (pd.DataFrame): The parent table.
+            child_name (str): The name of the child table in the metadata.
+            parent_name (str): The name of the parent table in the metadata.
+
+        Returns:
+            None: The child_table is modified in-place.
+        """
         parent_primary_key = self.metadata.tables[parent_name].primary_key
         parent_id_values = None
         for foreign_key in self.metadata._get_foreign_keys(parent_name, child_name):
-            needs_assignment = True
-            if foreign_key in child_table:
-                child_column = child_table[foreign_key]
-                if not child_column.dropna().empty:
+            needs_assignment = foreign_key not in child_table
+
+            if not needs_assignment:
+                child_column = child_table[foreign_key].dropna()
+                if child_column.empty:
+                    needs_assignment = True
+                else:
                     if parent_id_values is None:
                         parent_id_values = parent_table[parent_primary_key].dropna().unique()
-                    is_valid = child_column.dropna().isin(parent_id_values).all()
-                    if is_valid:
-                        needs_assignment = False
+
+                    if not child_column.isin(parent_id_values).all():
+                        needs_assignment = True
 
             if needs_assignment:
                 parent_ids = self._find_parent_ids(

tests/unit/multi_table/test_hma.py

@@ -912,6 +912,41 @@ def test__add_foreign_key_columns(self):
         pd.testing.assert_frame_equal(expected_parent_table, parent_table)
         pd.testing.assert_frame_equal(expected_child_table, child_table)
 
+    def test__add_foreign_key_columns_fk_already_in_child(self):
+        """Test that ``_add_foreign_key_columns`` does not add a FK already in child table."""
+        # Setup
+        instance = Mock()
+        mock_child_metadata = Mock()
+        mock_child_metadata.primary_key = 'parent_id'
+        mock_parent_metadata = Mock()
+        mock_parent_metadata.primary_key = 'parent_id_pk'
+        mock_metadata = Mock()
+        mock_metadata.tables = {
+            'child': mock_child_metadata,
+            'parent': mock_parent_metadata,
+        }
+        mock_metadata._get_foreign_keys.return_value = ['parent_id']
+        instance.metadata = mock_metadata
+        child_table = pd.DataFrame({'parent_id': [1, 2, 3], 'col_num': [10.1, 11.2, 12.3]})
+        parent_table = pd.DataFrame({'parent_id_pk': [1, 2, 3, 4], 'col_cat': ['A', 'B', 'B', 'C']})
+
+        # Run
+        HMASynthesizer._add_foreign_key_columns(
+            instance, child_table, parent_table, child_name='child', parent_name='parent'
+        )
+
+        # Assert
+        expected_parent_table = pd.DataFrame({
+            'parent_id_pk': pd.Series([1, 2, 3, 4], dtype=np.int64),
+            'col_cat': pd.Series(['A', 'B', 'B', 'C'], dtype=object),
+        })
+        expected_child_table = pd.DataFrame({
+            'parent_id': pd.Series([1, 2, 3], dtype=np.int64),
+            'col_num': pd.Series([10.1, 11.2, 12.3], dtype='float64'),
+        })
+        pd.testing.assert_frame_equal(expected_parent_table, parent_table)
+        pd.testing.assert_frame_equal(expected_child_table, child_table)
+
     def test__estimate_num_columns_to_be_modeled_multiple_foreign_keys(self):
         """Test it when there are two relationships between a parent and a child tables.