def + test multi-table

Author: R-Palazzo

sdv/multi_table/base.py

@@ -18,7 +18,6 @@
     check_synthesizer_version,
     generate_synthesizer_id,
 )
-from sdv.cag._errors import PatternNotMetError
 from sdv.errors import (
     ConstraintsNotMetError,
     InvalidDataError,
@@ -118,6 +117,7 @@ def __init__(self, metadata, locales=['en_US'], synthesizer_kwargs=None):
         self._table_synthesizers = {}
         self._table_parameters = defaultdict(dict)
         self._original_table_columns = {}
+        self._original_metadata = deepcopy(self.metadata)
         if synthesizer_kwargs is not None:
             warn_message = (
                 'The `synthesizer_kwargs` parameter is deprecated as of SDV 1.2.0 and does not '
@@ -323,6 +323,22 @@ def _validate_all_tables(self, data):
 
         return errors
 
+    def _validate_cags(self, data):
+        """Validate the data against the CAG patterns.
+
+        Args:
+            data (pandas.DataFrame):
+                The data to validate.
+        """
+        metadata = self.metadata
+        if hasattr(self, '_original_metadata'):
+            metadata = self._original_metadata
+
+        if hasattr(self, 'patterns'):
+            for pattern in self.patterns:
+                pattern.validate(data=data, metadata=metadata)
+                metadata = pattern.get_updated_metadata(metadata)
+
     def validate(self, data):
         """Validate the data.
 
@@ -334,8 +350,11 @@ def validate(self, data):
         """
         errors = []
         constraints_errors = []
-        cags_errors = []
-        self.metadata.validate_data(data)
+        metadata = self.metadata
+        if hasattr(self, '_original_metadata'):
+            metadata = self._original_metadata
+
+        metadata.validate_data(data)
         for table_name in data:
             if table_name in self._table_synthesizers:
                 try:
@@ -346,21 +365,13 @@ def validate(self, data):
                 # Validate rules specific to each synthesizer
                 errors += self._table_synthesizers[table_name]._validate(data[table_name])
 
-                # Validate single-table cags
-                if hasattr(self._table_synthesizers[table_name], '_validate_cags'):
-                    try:
-                        self._table_synthesizers[table_name]._validate_cags(data[table_name])
-                    except PatternNotMetError as error:
-                        cags_errors.append(error)
-
         if constraints_errors:
             raise ConstraintsNotMetError(constraints_errors)
 
         elif errors:
             raise InvalidDataError(errors)
 
-        elif cags_errors:
-            raise PatternNotMetError(cags_errors)
+        self._validate_cags(data)
 
     def _validate_table_name(self, table_name):
         if table_name not in self._table_synthesizers:
@@ -462,8 +473,8 @@ def preprocess(self, data):
         """
         list_of_changed_tables = self._store_and_convert_original_cols(data)
 
-        data = self._transform_helper(data)
         self.validate(data)
+        data = self._transform_helper(data)
         if self._fitted:
             warnings.warn(
                 'This model has already been fitted. To use the new preprocessed data, '

tests/integration/multi_table/test_hma.py

@@ -15,6 +15,8 @@
 from sdmetrics.reports.multi_table import DiagnosticReport
 
 from sdv import version
+from sdv.cag import Inequality
+from sdv.cag._errors import PatternNotMetError
 from sdv.datasets.demo import download_demo
 from sdv.datasets.local import load_csvs
 from sdv.errors import InvalidDataError, SamplingError, SynthesizerInputError, VersionError
@@ -2686,3 +2688,35 @@ def test__unsupported_regex_format():
     # Run and Assert
     with pytest.raises(SynthesizerInputError, match=expected_error):
         HMASynthesizer(metadata)
+
+
+def test_end_to_end_with_cags():
+    """Test HMA with a single-table cag."""
+    # Setup
+    data, metadata = download_demo('multi_table', 'fake_hotels')
+    synthesizer = HMASynthesizer(metadata)
+    pattern = Inequality(
+        low_column_name='checkin_date',
+        high_column_name='checkout_date',
+        strict_boundaries=False,
+        table_name='guests',
+    )
+    synthesizer.add_cag(patterns=[pattern])
+    data_guests = data['guests']
+    clean_data = data_guests[~(data_guests[['checkin_date', 'checkout_date']].isna().any(axis=1))]
+    data_invalid = clean_data.copy()
+    data_invalid.loc[0, 'checkin_date'] = '31 Dec 2020'
+    data['guests'] = clean_data
+    invalid_data = data.copy()
+    invalid_data['guests'] = data_invalid
+    expected_error_msg = re.escape('The inequality requirement is not met for row indices: [0]')
+
+    # Run
+    synthesizer.fit(data)
+    synthetic_data = synthesizer.sample(scale=1.0)
+
+    with pytest.raises(PatternNotMetError, match=expected_error_msg):
+        synthesizer.fit(invalid_data)
+
+    # Assert
+    synthesizer.validate(synthetic_data)

tests/unit/multi_table/test_base.py

@@ -410,16 +410,44 @@ def test_get_metadata(self):
         assert type(result) is Metadata
         assert expected_metadata.to_dict() == result.to_dict()
 
+    def test__validate_cags(self):
+        """Test the ``_validate_cags`` method."""
+        # Setup
+        data = pd.DataFrame()
+        original_metadata = Metadata()
+        metadata_1 = Metadata()
+        metadata_2 = Metadata()
+        instance = BaseMultiTableSynthesizer(original_metadata)
+        cag_mock_1 = Mock()
+        cag_mock_1.get_updated_metadata = Mock(return_value=metadata_1)
+        cag_mock_2 = Mock()
+        cag_mock_2.get_updated_metadata = Mock(return_value=metadata_2)
+        cag_mock_3 = Mock()
+        instance.patterns = [cag_mock_1, cag_mock_2, cag_mock_3]
+
+        # Run
+        instance._validate_cags(data)
+
+        # Assert
+        cag_mock_1.get_updated_metadata.assert_called_once_with(instance._original_metadata)
+        cag_mock_1.validate.assert_called_once_with(data=data, metadata=instance._original_metadata)
+        cag_mock_2.validate.assert_called_once_with(data=data, metadata=metadata_1)
+        cag_mock_3.validate.assert_called_once_with(data=data, metadata=metadata_2)
+
     def test_validate(self):
         """Test that no error is being raised when the data is valid."""
         # Setup
         metadata = get_multi_table_metadata()
         data = get_multi_table_data()
         instance = BaseMultiTableSynthesizer(metadata)
+        instance._validate_cags = Mock()
 
-        # Run and Assert
+        # Run
         instance.validate(data)
 
+        # Assert
+        instance._validate_cags.assert_called_once_with(data)
+
     def test_validate_missing_table(self):
         """Test that an error is being raised when there is a missing table in the dictionary."""
         # Setup