Add primary key auto-detection for semantic sdtypes (#1731)

sdv-dev · Jan 10, 2024 · 15f18fd · 15f18fd
1 parent 6659835
commit 15f18fd
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 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()
@@ -407,15 +408,20 @@ def _detect_columns(self, data):
                         sdtype = 'unknown'
 
             column_dict = {'sdtype': sdtype}
-
             if sdtype in self._REFERENCE_TO_SDTYPE.values() or sdtype == 'unknown':
                 column_dict['pii'] = True
-            elif sdtype == 'datetime' and dtype == 'O':
+            if sdtype in self._REFERENCE_TO_SDTYPE.values() and first_pii_field is None:
+                first_pii_field = field
+            if 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
@@ -205,6 +205,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': [
@@ -269,6 +270,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
@@ -327,6 +327,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},