Add alternate keys to dropped columns

sdv-dev · Nov 22, 2023 · 33ee1c3 · 33ee1c3
1 parent 150f3a5
commit 33ee1c3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/sdmetrics/single_table/detection/base.py b/sdmetrics/single_table/detection/base.py
@@ -9,7 +9,7 @@
 from sdmetrics.errors import IncomputableMetricError
 from sdmetrics.goal import Goal
 from sdmetrics.single_table.base import SingleTableMetric
-from sdmetrics.utils import HyperTransformer
+from sdmetrics.utils import HyperTransformer, get_alternate_keys
 
 LOGGER = logging.getLogger(__name__)
 
@@ -73,6 +73,7 @@ def compute(cls, real_data, synthetic_data, metadata=None):
 
         if metadata is not None:
             drop_columns = []
+            drop_columns.extend(get_alternate_keys(metadata))
             if 'columns' in metadata:
                 for column in metadata['columns']:
                     if ('primary_key' in metadata and

diff --git a/tests/unit/single_table/detection/test_detection.py b/tests/unit/single_table/detection/test_detection.py
@@ -84,7 +84,8 @@ def test_ignore_keys_detection_metrics(self, fit_transform_mock, transform_mock)
             'col3': [5, 6, 7, 8, 9],
             'ID_3': ['a', 'b', 'c', 'd', 'e'],
             'blob': ['Hello world!', 'Hello world!', 'This is SDV', 'This is SDV', 'Hello world!'],
-            'col4': [1, 3, 9, 2, 1]
+            'col4': [1, 3, 9, 2, 1],
+            'col5': [10, 20, 30, 40, 50]
         })
         synthetic_data = pd.DataFrame({
             'ID_1': [1, 3, 4, 2, 2],
@@ -94,7 +95,8 @@ def test_ignore_keys_detection_metrics(self, fit_transform_mock, transform_mock)
             'col3': [55, 66, 77, 88, 99],
             'ID_3': ['a', 'b', 'e', 'd', 'c'],
             'blob': ['Hello world!', 'Hello world!', 'This is SDV', 'This is SDV', 'Hello world!'],
-            'col4': [4, 1, 3, 1, 9]
+            'col4': [4, 1, 3, 1, 9],
+            'col5': [10, 20, 30, 40, 50]
         })
         metadata = {
             'columns': {
@@ -105,9 +107,11 @@ def test_ignore_keys_detection_metrics(self, fit_transform_mock, transform_mock)
                 'col3': {'sdtype': 'numerical'},
                 'ID_3': {'sdtype': 'id'},
                 'blob': {'sdtype': 'text'},
-                'col4': {'sdtype': 'numerical', 'pii': False}
+                'col4': {'sdtype': 'numerical', 'pii': False},
+                'col5': {'sdtype': 'numerical'}
             },
-            'primary_key': {'ID_1', 'ID_2'}
+            'primary_key': {'ID_1', 'ID_2'},
+            'alternate_keys': ['col5']
         }
 
         expected_real_dataframe = pd.DataFrame({