wip

Author: gsheni

sdv/multi_table/base.py

@@ -471,6 +471,8 @@ def _assign_table_transformers(self, synthesizer, table_name, table_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
 
         column_name_to_transformers = {}
         primary_key = self.metadata.tables[table_name].primary_key

sdv/multi_table/hma.py

@@ -349,6 +349,9 @@ 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)
 
@@ -447,6 +450,11 @@ def _augment_table(self, table, tables, table_name):
                 child_table = tables[child_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)}) '
@@ -492,13 +500,15 @@ 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:
             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')
@@ -518,6 +528,10 @@ 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()

tests/integration/multi_table/test_hma.py

@@ -2903,7 +2903,7 @@ def data_metadata_1_to_1():
     metadata = Metadata.load_from_dict(metadata_dict)
     metadata.validate()
     metadata.validate_data(data)
-    metadata.remove_primary_key('rooms')
+    # metadata.remove_primary_key('rooms')
     return data, metadata