Update msg

Author: fealho

sdv/multi_table/hma.py

@@ -18,6 +18,11 @@
 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
 
 
 class HMASynthesizer(BaseHierarchicalSampler, BaseMultiTableSynthesizer):
@@ -102,7 +107,7 @@ def _get_num_extended_columns(
 
     @classmethod
     def _estimate_columns_traversal(
-        cls, metadata, table_name, columns_per_table, visited, distributions=None
+        cls, metadata, table_name, columns_per_table, visited, distributions=None, max_total=None
     ):
         """Given a table, estimate how many columns each parent will model.
 
@@ -118,16 +123,21 @@ 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)
+                cls._estimate_columns_traversal(
+                    metadata, child_name, columns_per_table, visited, distributions, max_total
+                )
 
             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
+
         visited.add(table_name)
 
     @classmethod
-    def _estimate_num_columns(cls, metadata, distributions=None):
+    def _estimate_num_columns(cls, metadata, distributions=None, max_total=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
@@ -153,9 +163,14 @@ def _estimate_num_columns(cls, metadata, distributions=None):
         # each table will model
         visited = set()
         for table_name in _get_root_tables(metadata.relationships):
-            cls._estimate_columns_traversal(
-                metadata, table_name, columns_per_table, visited, distributions
-            )
+            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:
+                break
 
         return columns_per_table
 
@@ -235,19 +250,29 @@ def _print_estimate_warning(self):
         metadata_columns = self._get_num_data_columns(self.metadata)
         print_table = []
         distributions = self._get_distributions()
-        for table, est_cols in self._estimate_num_columns(self.metadata, distributions).items():
+        estimated_columns = self._estimate_num_columns(
+            self.metadata, distributions, max_total=PERFORMANCE_ALERT_DISPLAY_CAP
+        )
+        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 = (
+                f'{PERFORMANCE_ALERT_DISPLAY_CAP}+'
+                if total_est_cols > PERFORMANCE_ALERT_DISPLAY_CAP
+                else f'{total_est_cols}'
+            )
             self._print(
                 'PerformanceAlert: Using the HMASynthesizer on this metadata '
                 'schema is not recommended. To model this data, HMA will '
-                f'generate a large number of columns. ({total_est_cols} columns)\n\n'
+                f'generate a large number of columns. ({display_total} columns)\n\n'
             )
             self._print(
                 pd.DataFrame(

tests/unit/multi_table/test_hma.py

@@ -7,7 +7,11 @@
 
 from sdv.errors import SynthesizerInputError
 from sdv.metadata.metadata import Metadata
-from sdv.multi_table.hma import HMASynthesizer
+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
 
@@ -1285,3 +1289,43 @@ 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)