Allow missing_values_proportion but validate it is zero for keys

Author: frances-h

sdv/single_table/_dayz_utils.py

@@ -35,7 +35,6 @@ def detect_column_parameters(data, metadata, table_name):
         dict: A dictionary containing the detected parameters.
     """
     table_metadata = metadata.tables[table_name]
-    table_keys = table_metadata._get_primary_and_alternate_keys()
     column_parameters = {}
     for column_name, column_metadata in table_metadata.columns.items():
         column_parameters[column_name] = {}
@@ -69,10 +68,9 @@ def detect_column_parameters(data, metadata, table_name):
                 'category_values': data[column_name].dropna().unique().tolist()
             }
 
-        if column_name not in table_keys:
-            column_parameters[column_name]['missing_values_proportion'] = (
-                data[column_name].isna().mean().item()
-            )
+        column_parameters[column_name]['missing_values_proportion'] = (
+            data[column_name].isna().mean()
+        )
 
     return {'columns': column_parameters}
 

sdv/single_table/dayz.py

@@ -57,15 +57,6 @@ def _validate_parameter_structure(dayz_parameters):
         _validate_table_parameter_dict_keys(table, table_parameters)
 
 
-def _validate_key_column(column_parameters, column_table_msg):
-    if 'missing_values_proportion' in column_parameters:
-        msg = (
-            f"Invalid 'missing_values_proportion' parameter for {column_table_msg}. Primary "
-            "and alternate keys can not have the 'missing_values_proportion' parameter set."
-        )
-        raise SynthesizerProcessingError(msg)
-
-
 def _validate_numerical_parameters(column_parameters, column_table_msg):
     for param in ['min_value', 'max_value']:
         if param in column_parameters and not _is_numerical(column_parameters[param]):
@@ -141,7 +132,7 @@ def _validate_categorical_parameters(column_parameters, column_table_msg):
         raise SynthesizerProcessingError(msg)
 
 
-def _validate_missing_value_parameters(column_parameters, column_table_msg):
+def _validate_missing_value_parameters(column_parameters, column_table_msg, is_key_column):
     missing_values_proportion = column_parameters['missing_values_proportion']
     if not _is_numerical(missing_values_proportion) or (
         missing_values_proportion < 0.0 or missing_values_proportion > 1.0
@@ -151,6 +142,12 @@ def _validate_missing_value_parameters(column_parameters, column_table_msg):
             'must be a float between 0.0 and 1.0.'
         )
         raise SynthesizerProcessingError(msg)
+    elif is_key_column and missing_values_proportion != 0:
+        msg = (
+            f"Invalid 'missing_values_proportion' parameter for {column_table_msg}. Primary "
+            "and alternate keys must have 'missing_values_proportion' parameter be zero."
+        )
+        raise SynthesizerProcessingError(msg)
 
 
 def _validate_column_parameters(table, column, column_metadata, column_parameters, is_key_column):
@@ -166,9 +163,6 @@ def _validate_column_parameters(table, column, column_metadata, column_parameter
         )
         raise SynthesizerProcessingError(msg)
 
-    if is_key_column:
-        _validate_key_column(column_parameters, column_table_msg)
-
     if sdtype == 'numerical':
         _validate_numerical_parameters(column_parameters, column_table_msg)
     elif sdtype == 'datetime':
@@ -177,7 +171,7 @@ def _validate_column_parameters(table, column, column_metadata, column_parameter
         _validate_categorical_parameters(column_parameters, column_table_msg)
 
     if 'missing_values_proportion' in column_parameters:
-        _validate_missing_value_parameters(column_parameters, column_table_msg)
+        _validate_missing_value_parameters(column_parameters, column_table_msg, is_key_column)
 
 
 def _validate_table_parameters(table, table_metadata, table_parameters):

tests/unit/single_table/test__dayz_utils.py

@@ -56,7 +56,7 @@ def test_detect_column_parameter():
     # Assert
     assert result == {
         'columns': {
-            'pk': {},
+            'pk': {'missing_values_proportion': 0.0},
             'num_col': {
                 'num_decimal_digits': 1,
                 'min_value': 1.0,
@@ -77,7 +77,7 @@ def test_detect_column_parameter():
                 'end_timestamp': '2020-01-03 00:00:00',
                 'missing_values_proportion': 0.25,
             },
-            'alt_key': {},
+            'alt_key': {'missing_values_proportion': 0.0},
         }
     }
 

tests/unit/single_table/test_dayz.py

@@ -98,6 +98,7 @@ def test__validate_column_parameter():
     column_metadata = {'sdtype': 'id'}
     bad_column_parameters = {'invalid_key': None}
     bad_missing_value = {'missing_values_proportion': 100}
+    bad_key_missing_value = {'missing_values_proportion': 0.5}
 
     # Run and Assert
     expected_bad_column_msg = re.escape(
@@ -118,10 +119,10 @@ def test__validate_column_parameter():
 
     expected_missing_values_with_key_msg = re.escape(
         "Invalid 'missing_values_proportion' parameter for column 'column' in table 'table'. "
-        "Primary and alternate keys can not have the 'missing_values_proportion' parameter set."
+        "Primary and alternate keys must have 'missing_values_proportion' parameter be zero."
     )
     with pytest.raises(SynthesizerProcessingError, match=expected_missing_values_with_key_msg):
-        _validate_column_parameters('table', 'column', column_metadata, bad_missing_value, True)
+        _validate_column_parameters('table', 'column', column_metadata, bad_key_missing_value, True)
 
 
 def test__validate_column_parameters_numerical():