Autodetect primary keys for pii columns

sdv-dev · Jan 5, 2024 · 20f6195 · 20f6195
1 parent 64e8df2
commit 20f6195
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/sdv/metadata/single_table.py b/sdv/metadata/single_table.py
@@ -378,6 +378,7 @@ def _detect_columns(self, data):
             data (pandas.DataFrame):
                 The data to be analyzed.
         """
+        first_pii_field = None
         for field in data:
             column_data = data[field]
             clean_data = column_data.dropna()
@@ -410,12 +411,19 @@ def _detect_columns(self, data):
 
             if sdtype in self._REFERENCE_TO_SDTYPE.values() or sdtype == 'unknown':
                 column_dict['pii'] = True
+                if first_pii_field is None:
+                    first_pii_field = field
+
             elif sdtype == 'datetime' and dtype == 'O':
                 datetime_format = get_datetime_format(column_data.iloc[:100])
                 column_dict['datetime_format'] = datetime_format
 
             self.columns[field] = deepcopy(column_dict)
 
+        # When no primary key column was set, choose the first pii field
+        if self.primary_key is None and first_pii_field:
+            self.primary_key = first_pii_field
+
     def detect_from_dataframe(self, data):
         """Detect the metadata from a ``pd.DataFrame`` object.
 

diff --git a/tests/integration/metadata/test_multi_table.py b/tests/integration/metadata/test_multi_table.py
@@ -197,6 +197,7 @@ def test_detect_from_dataframes():
                     'billing_address': {'sdtype': 'unknown', 'pii': True},
                     'credit_card_number': {'sdtype': 'credit_card_number', 'pii': True}
                 },
+                'primary_key': 'guest_email'
             }
         },
         'relationships': [
@@ -261,6 +262,7 @@ def test_detect_from_csvs(tmp_path):
                     'billing_address': {'sdtype': 'unknown', 'pii': True},
                     'credit_card_number': {'sdtype': 'credit_card_number', 'pii': True}
                 },
+                'primary_key': 'guest_email'
             }
         },
         'relationships': [

diff --git a/tests/integration/metadata/test_single_table.py b/tests/integration/metadata/test_single_table.py
@@ -293,6 +293,7 @@ def test_detect_from_dataframe_with_pii_names():
     # Assert
     expected_metadata = {
         'METADATA_SPEC_VERSION': 'SINGLE_TABLE_V1',
+        'primary_key': 'USER PHONE NUMBER',
         'columns': {
             'USER PHONE NUMBER': {'sdtype': 'phone_number', 'pii': True},
             'addr_line_1': {'sdtype': 'street_address', 'pii': True},