Fix test

Author: fealho

sdv/multi_table/hma.py

@@ -16,13 +16,9 @@
 from sdv.sampling import BaseHierarchicalSampler
 
 LOGGER = logging.getLogger(__name__)
-MAX_NUMBER_OF_COLUMNS = 1000
-DEFAULT_EXTENDED_COLUMNS_DISTRIBUTION = 'truncnorm'
 PERFORMANCE_ALERT_DISPLAY_CAP = 1_000_000
-
-
-class _EarlyStopEstimation(Exception):
-    pass
+DEFAULT_EXTENDED_COLUMNS_DISTRIBUTION = 'truncnorm'
+MAX_NUMBER_OF_COLUMNS = 1000
 
 
 class HMASynthesizer(BaseHierarchicalSampler, BaseMultiTableSynthesizer):
@@ -107,7 +103,7 @@ def _get_num_extended_columns(
 
     @classmethod
     def _estimate_columns_traversal(
-        cls, metadata, table_name, columns_per_table, visited, distributions=None, max_total=None
+        cls, metadata, table_name, columns_per_table, visited, distributions=None
     ):
         """Given a table, estimate how many columns each parent will model.
 
@@ -123,21 +119,19 @@ def _estimate_columns_traversal(
         """
         for child_name in metadata._get_child_map()[table_name]:
             if child_name not in visited:
-                cls._estimate_columns_traversal(
-                    metadata, child_name, columns_per_table, visited, distributions, max_total
-                )
+                cls._estimate_columns_traversal(metadata, child_name, columns_per_table, visited)
 
             columns_per_table[table_name] += cls._get_num_extended_columns(
                 metadata, child_name, table_name, columns_per_table, distributions
             )
 
-            if max_total is not None and sum(columns_per_table.values()) > max_total:
-                raise _EarlyStopEstimation
+            if sum(columns_per_table.values()) > PERFORMANCE_ALERT_DISPLAY_CAP:
+                return
 
         visited.add(table_name)
 
     @classmethod
-    def _estimate_num_columns(cls, metadata, distributions=None, max_total=None):
+    def _estimate_num_columns(cls, metadata, distributions=None):
         """Estimate the number of columns that will be modeled for each table.
 
         This method estimates how many extended columns will be generated during the
@@ -163,13 +157,10 @@ def _estimate_num_columns(cls, metadata, distributions=None, max_total=None):
         # each table will model
         visited = set()
         for table_name in _get_root_tables(metadata.relationships):
-            try:
-                cls._estimate_columns_traversal(
-                    metadata, table_name, columns_per_table, visited, distributions, max_total
-                )
-            except _EarlyStopEstimation:
-                break
-            if max_total is not None and sum(columns_per_table.values()) > max_total:
+            cls._estimate_columns_traversal(
+                metadata, table_name, columns_per_table, visited, distributions
+            )
+            if sum(columns_per_table.values()) > PERFORMANCE_ALERT_DISPLAY_CAP:
                 break
 
         return columns_per_table
@@ -250,18 +241,14 @@ def _print_estimate_warning(self):
         metadata_columns = self._get_num_data_columns(self.metadata)
         print_table = []
         distributions = self._get_distributions()
-        estimated_columns = self._estimate_num_columns(
-            self.metadata, distributions, max_total=PERFORMANCE_ALERT_DISPLAY_CAP
-        )
+        estimated_columns = self._estimate_num_columns(self.metadata, distributions)
         for table, est_cols in estimated_columns.items():
             entry = []
             entry.append(table)
             entry.append(metadata_columns[table])
             total_est_cols += est_cols
             entry.append(est_cols)
             print_table.append(entry)
-            if total_est_cols > PERFORMANCE_ALERT_DISPLAY_CAP:
-                break
 
         if total_est_cols > MAX_NUMBER_OF_COLUMNS:
             display_total = (

tests/unit/multi_table/test_hma.py

@@ -8,9 +8,7 @@
 from sdv.errors import SynthesizerInputError
 from sdv.metadata.metadata import Metadata
 from sdv.multi_table.hma import (
-    PERFORMANCE_ALERT_DISPLAY_CAP,
     HMASynthesizer,
-    _EarlyStopEstimation,
 )
 from sdv.single_table.copulas import GaussianCopulaSynthesizer
 from tests.utils import get_multi_table_data, get_multi_table_metadata
@@ -135,6 +133,30 @@ def test__print_estimate_warning(self, get_distributions_mock, estimate_mock, ca
             match = re.search(constraint, captured.out + captured.err)
             assert match is None
 
+    @patch('sdv.multi_table.hma.HMASynthesizer._estimate_num_columns')
+    @patch('sdv.multi_table.hma.HMASynthesizer._get_distributions')
+    def test__print_estimate_warning_many_cols(self, get_distributions_mock, estimate_mock, capsys):
+        """Test that a warning appears if there are more than 1_000_000 expected columns"""
+        # Setup
+        metadata = get_multi_table_metadata()
+        estimate_mock.side_effect = [{'nesreca': 1_000_010}, {'nesreca': 10}]
+
+        # Run
+        HMASynthesizer(metadata)
+        captured = capsys.readouterr()
+
+        # Assert
+        expected_output = (
+            'PerformanceAlert: Using the HMASynthesizer on this metadata schema is not recommended.'
+            ' To model this data, HMA will generate a large number of columns. (1000000+ columns)\n'
+            '\n\nTable Name  # Columns in Metadata  Est # Columns\n'
+            '   nesreca                      1        1000010\n\n'
+            "We recommend simplifying your metadata schema using 'sdv.utils.poc.simplify_schema'."
+            '\nIf this is not possible, please visit datacebo.com and reach out to us for '
+            'enterprise solutions.\n\n'
+        )
+        assert captured.out == expected_output
+
     def test__get_extension_foreign_key_only(self):
         """Test the ``_get_extension`` method.
 
@@ -1289,43 +1311,3 @@ def test__estimate_num_columns_to_be_modeled_various_sdtypes(self):
                 num_table_cols -= 1
 
             assert num_table_cols == estimated_num_columns[table_name]
-
-    @patch('sdv.multi_table.hma.HMASynthesizer._estimate_num_columns')
-    @patch('sdv.multi_table.hma.HMASynthesizer._get_distributions')
-    def test__print_estimate_warning_capped_display_and_break(
-        self, get_distributions_mock, estimate_mock, capsys
-    ):
-        """When exceeding the cap, display 1_000_000+ and stop listing rows."""
-        # Setup
-        metadata = get_multi_table_metadata()
-        estimate_mock.return_value = {
-            'nesreca': PERFORMANCE_ALERT_DISPLAY_CAP + 5,
-            'oseba': 10,
-        }
-
-        # Run
-        HMASynthesizer(metadata)
-        captured = capsys.readouterr()
-
-        # Assert
-        assert 'PerformanceAlert:' in captured.out
-        assert f'{PERFORMANCE_ALERT_DISPLAY_CAP}+' in captured.out
-        assert 'nesreca' in captured.out
-        assert 'oseba' not in captured.out
-
-    def test__estimate_num_columns_early_stop_exception_is_handled(self):
-        """_estimate_num_columns should handle internal early-stop and return partial results."""
-        # Setup
-        metadata = get_multi_table_metadata()
-
-        with patch('sdv.multi_table.hma.HMASynthesizer._estimate_columns_traversal') as tr_mock:
-            tr_mock.side_effect = _EarlyStopEstimation
-
-            # Run
-            result = HMASynthesizer._estimate_num_columns(metadata, max_total=1)
-
-        # Assert
-        assert isinstance(result, dict)
-        for table_name in ['nesreca', 'oseba', 'upravna_enota']:
-            assert table_name in result
-            assert isinstance(result[table_name], int)