wip

Author: gsheni

sdv/multi_table/base.py

@@ -466,18 +466,17 @@ def _validate_table_name(self, table_name):
             )
 
     def _assign_table_transformers(self, synthesizer, table_name, table_data):
-        """Update the ``synthesizer`` to ignore the foreign keys while preprocessing the data."""
-        synthesizer.auto_assign_transformers(table_data)
+        """Update the ``synthesizer`` to ignore the foreign keys while preprocessing the data.
 
-        # foreign_key_columns = self.metadata._get_all_foreign_keys(table_name)
-        # column_name_to_transformers = {column_name: None for column_name in foreign_key_columns}
-        # keep it in it's raw form, HMA can use it group later on,
-        # prevent pre-process FK columns to numerical
+        If the foreign key is also the primary key for the table, then no transformer assigned
+        for that column.
 
+        """
+        synthesizer.auto_assign_transformers(table_data)
         column_name_to_transformers = {}
         primary_key = self.metadata.tables[table_name].primary_key
         for relation in self.metadata.relationships:
-            column_name = deepcopy(relation['child_foreign_key'])
+            column_name = relation['child_foreign_key']
             if (
                 relation['child_table_name'] == table_name
                 and relation['child_foreign_key'] != primary_key

sdv/multi_table/hma.py

@@ -349,9 +349,6 @@ def _get_extension(self, child_name, child_table, foreign_key, progress_bar_desc
 
         extension_rows = []
         foreign_key_columns = self.metadata._get_all_foreign_keys(child_name)
-
-        # foreign_key_values = child_table.index.unique()
-        # only do if FK not the primary key
         foreign_key_values = child_table[foreign_key].unique()
         child_table = child_table.set_index(foreign_key)
 
@@ -451,10 +448,6 @@ def _augment_table(self, table, tables, table_name):
 
             foreign_keys = self.metadata._get_foreign_keys(table_name, child_name)
 
-            primary_key = self.metadata.tables[child_name].primary_key
-            if primary_key in foreign_keys:
-                child_table = child_table.reset_index(drop=False)
-            # check here
             for foreign_key in foreign_keys:
                 progress_bar_desc = (
                     f'({self._learned_relationships + 1}/{len(self.metadata.relationships)}) '
@@ -500,15 +493,20 @@ def _augment_tables(self, processed_data):
             processed_data (dict):
                 Dictionary mapping each table name to a preprocessed ``pandas.DataFrame``.
         """
-        # data processor sets index
         augmented_data = deepcopy(processed_data)
         self._augmented_tables = []
         self._learned_relationships = 0
         parent_map = self.metadata._get_parent_map()
         self._print(text='Learning relationships:')
+
+        for table_name in processed_data:
+            foreign_keys = self.metadata._get_all_foreign_keys(table_name)
+            primary_key = self.metadata.tables[table_name].primary_key
+            if primary_key in foreign_keys:
+                augmented_data[table_name] = augmented_data[table_name].reset_index(drop=False)
+
         for table_name in processed_data:
             if not parent_map.get(table_name):
-                # only changing the child tables
                 self._augment_table(augmented_data[table_name], augmented_data, table_name)
 
         LOGGER.info('Augmentation Complete')
@@ -528,10 +526,6 @@ def _pop_foreign_keys(self, table_data, table_name):
                 A dictionary mapping with the foreign key and it's values within the table.
         """
         foreign_keys = self.metadata._get_all_foreign_keys(table_name)
-        primary_key = self.metadata.tables[table_name].primary_key
-        if primary_key in foreign_keys:
-            table_data = table_data.reset_index(drop=False)
-
         keys = {}
         for fk in foreign_keys:
             keys[fk] = table_data.pop(fk).to_numpy()

sdv/sampling/hierarchical_sampler.py

@@ -102,9 +102,9 @@ def _add_child_rows(self, child_name, parent_name, parent_row, sampled_data, num
         if len(sampled_rows):
             parent_key = self.metadata.tables[parent_name].primary_key
             if foreign_key in sampled_rows:
-                # If foreign key is in sampeld rows raises `SettingWithCopyWarning`
+                # If foreign key is in sampled rows raises `SettingWithCopyWarning`
                 row_indices = sampled_rows.index
-                sampled_rows[foreign_key].iloc[row_indices] = parent_row[parent_key]
+                sampled_rows.loc[row_indices, foreign_key] = parent_row[parent_key]
             else:
                 sampled_rows[foreign_key] = (
                     parent_row[parent_key] if parent_row is not None else np.nan