Make GaussianCopula report the correct distribution name in the case of a fallback (#2401)

Author: fealho

sdv/multi_table/utils.py

@@ -443,7 +443,7 @@ def _drop_rows(data, metadata, drop_missing_values):
 
         if data[table].empty:
             raise InvalidDataError([
-                f"All references in table '{table}' are unknown and must be dropped."
+                f"All references in table '{table}' are unknown and must be dropped. "
                 'Try providing different data for this table.'
             ])
 
@@ -558,7 +558,7 @@ def _subsample_parent(
     parent_table = parent_table.drop(unreferenced_data_to_drop.index)
     if parent_table.empty:
         raise InvalidDataError([
-            f"All references in table '{parent_primary_key}' are unknown and must be dropped."
+            f"All references in table '{parent_primary_key}' are unknown and must be dropped. "
             'Try providing different data for this table.'
         ])
 

sdv/single_table/copulas.py

@@ -224,10 +224,11 @@ def get_learned_distributions(self):
         univariates = deepcopy(parameters['univariates'])
         learned_distributions = {}
         valid_columns = self._get_valid_columns_from_metadata(columns)
+        distribution_names = {v.__name__: k for k, v in self._DISTRIBUTIONS.items()}
         for column, learned_params in zip(columns, univariates):
             if column in valid_columns:
-                distribution = self.numerical_distributions.get(column, self.default_distribution)
-                learned_params.pop('type')
+                distribution_name = learned_params.pop('type').split('.')[-1]
+                distribution = distribution_names[distribution_name]
                 learned_distributions[column] = {
                     'distribution': distribution,
                     'learned_parameters': learned_params,

tests/integration/single_table/test_copulas.py

@@ -522,3 +522,35 @@ def test_user_warning_for_unused_numerical_distribution():
     )
     with pytest.warns(UserWarning, match=message):
         synthesizer.fit(data)
+
+
+def test_get_learned_distributions_fallback_distribution():
+    """Test it when the fallback distribution is used GH#2394."""
+    # Setup
+    data = pd.DataFrame(data={'A': np.concatenate([np.zeros(29), np.ones(21)])})
+    metadata = Metadata.load_from_dict({
+        'tables': {
+            'table': {
+                'columns': {
+                    'A': {
+                        'sdtype': 'numerical',
+                    },
+                },
+            },
+        },
+    })
+
+    # Run
+    synthesizer = GaussianCopulaSynthesizer(metadata, default_distribution='beta')
+    synthesizer.fit(data)
+
+    # Assert
+    assert synthesizer.get_learned_distributions() == {
+        'A': {
+            'distribution': 'norm',
+            'learned_parameters': {
+                'loc': 0.42,
+                'scale': 0.4935585071701226,
+            },
+        },
+    }

tests/unit/multi_table/test_utils.py

@@ -469,7 +469,7 @@ def test_drop_unknown_references_drop_all_rows(mock_get_rows_to_drop):
     # Run and Assert
     expected_message = re.escape(
         'The provided data does not match the metadata:\n'
-        "All references in table 'child' are unknown and must be dropped."
+        "All references in table 'child' are unknown and must be dropped. "
         'Try providing different data for this table.'
     )
     with pytest.raises(InvalidDataError, match=expected_message):

tests/unit/utils/test_utils.py

@@ -351,7 +351,7 @@ def test_drop_unknown_references_drop_all_rows(mock_get_rows_to_drop):
     # Run and Assert
     expected_message = re.escape(
         'The provided data does not match the metadata:\n'
-        "All references in table 'child' are unknown and must be dropped."
+        "All references in table 'child' are unknown and must be dropped. "
         'Try providing different data for this table.'
     )
     with pytest.raises(InvalidDataError, match=expected_message):