Store metadata as Metadata for BaseSynthesizer (#2422)

Author: fealho

sdv/metadata/multi_table.py

@@ -776,7 +776,7 @@ def _validate_all_tables(self, data):
                     self.tables[table_name].validate_data(table_data, table_sdtype_warnings)
 
             except InvalidDataError as error:
-                error_msg = f"Table: '{table_name}'"
+                error_msg = f'Errors in {table_name}:'
                 for _error in error.errors:
                     error_msg += f'\nError: {_error}'
 

sdv/sampling/independent_sampler.py

@@ -106,8 +106,7 @@ def _finalize(self, sampled_data):
                 try:
                     table_rows[name] = table_rows[name].dropna().astype(dtype)
                 except ValueError as e:
-                    column_metadata = metadata.columns.get(name)
-                    sdtype = column_metadata.get('sdtype')
+                    sdtype = metadata.columns.get(name).get('sdtype')
                     if sdtype not in dtypes_to_sdtype.values():
                         LOGGER.info(
                             f"The real data in '{table_name}' and column '{name}' was stored as "

sdv/sequential/par.py

@@ -34,8 +34,8 @@ class PARSynthesizer(LossValuesMixin, BaseSynthesizer):
     to be passed into PAR.
 
     Args:
-        metadata (sdv.metadata.SingleTableMetadata):
-            Single table metadata representing the data that this synthesizer will be used for.
+        metadata (sdv.metadata.Metadata):
+            Metadata representing the data that this synthesizer will be used for.
         enforce_min_max_values (bool):
             Specify whether or not to clip the data returned by ``reverse_transform`` of
             the numerical transformer, ``FloatFormatter``, to the min and max values seen
@@ -78,8 +78,9 @@ def _get_context_metadata(self):
         for column, column_metadata in self._extra_context_columns.items():
             context_columns_dict[column] = column_metadata
 
+        table_metadata = self._get_table_metadata()
         for column in context_columns:
-            context_columns_dict[column] = self.metadata.columns[column]
+            context_columns_dict[column] = table_metadata.columns[column]
 
         context_columns_dict = self._update_context_column_dict(context_columns_dict)
         context_metadata_dict = {'columns': context_columns_dict}
@@ -97,8 +98,9 @@ def _update_context_column_dict(self, context_columns_dict):
                 Updated context column metadata.
         """
         default_transformers_by_sdtype = deepcopy(self._data_processor._transformers_by_sdtype)
+        table_metadata = self._get_table_metadata()
         for column in self.context_columns:
-            column_metadata = self.metadata.columns[column]
+            column_metadata = table_metadata.columns[column]
             if default_transformers_by_sdtype.get(column_metadata['sdtype']):
                 context_columns_dict[column] = {'sdtype': 'numerical'}
 
@@ -107,8 +109,9 @@ def _update_context_column_dict(self, context_columns_dict):
     def _get_context_columns_for_processing(self):
         columns_to_be_processed = []
         default_transformers_by_sdtype = deepcopy(self._data_processor._transformers_by_sdtype)
+        table_metadata = self._get_table_metadata()
         for column in self.context_columns:
-            if default_transformers_by_sdtype.get(self.metadata.columns[column]['sdtype']):
+            if default_transformers_by_sdtype.get(table_metadata.columns[column]['sdtype']):
                 columns_to_be_processed.append(column)
 
         return columns_to_be_processed
@@ -136,15 +139,15 @@ def __init__(
             locales=locales,
         )
 
-        sequence_key = self.metadata.sequence_key
+        sequence_key = self._get_table_metadata().sequence_key
         self._sequence_key = list(_cast_to_iterable(sequence_key)) if sequence_key else None
         if not self._sequence_key:
             raise SynthesizerInputError(
                 'The PARSythesizer is designed for multi-sequence data, identifiable through a '
                 'sequence key. Your metadata does not include a sequence key.'
             )
 
-        self._sequence_index = self.metadata.sequence_index
+        self._sequence_index = self._get_table_metadata().sequence_index
         self.context_columns = context_columns or []
         self._validate_sequence_key_and_context_columns()
         self._extra_context_columns = {}
@@ -415,13 +418,14 @@ def _fit_sequence_columns(self, timeseries_data):
         )
         data_types = []
         context_types = []
+        table_metadata = self._get_table_metadata()
         for field in self._output_columns:
             dtype = timeseries_data[field].dtype
             kind = dtype.kind
             if kind in ('i', 'f'):
                 data_type = 'continuous'
                 # Check if metadata overrides this data type
-                if self.metadata.columns.get(field, {}).get('sdtype', None) == 'categorical':
+                if table_metadata.columns.get(field, {}).get('sdtype', None) == 'categorical':
                     data_type = 'categorical'
             elif kind in ('O', 'b'):
                 data_type = 'categorical'

sdv/single_table/base.py

@@ -99,8 +99,8 @@ def _update_default_transformers(self):
                 self._data_processor._update_transformers_by_sdtypes(sdtype, transformer)
 
     def _check_metadata_updated(self):
-        if self.metadata._updated:
-            self.metadata._updated = False
+        if self.metadata._check_updated_flag():
+            self.metadata._reset_updated_flag()
             warnings.warn(
                 "We strongly recommend saving the metadata using 'save_to_json' for replicability"
                 ' in future SDV versions.'
@@ -110,22 +110,22 @@ def __init__(
         self, metadata, enforce_min_max_values=True, enforce_rounding=True, locales=['en_US']
     ):
         self._validate_inputs(enforce_min_max_values, enforce_rounding)
-        self.metadata = metadata
-        self._table_name = Metadata.DEFAULT_SINGLE_TABLE_NAME
         if isinstance(metadata, Metadata):
             self._table_name = metadata._get_single_table_name()
-            self.metadata = metadata._convert_to_single_table()
-        elif isinstance(metadata, SingleTableMetadata):
+            self.metadata = metadata
+        else:
             warnings.warn(DEPRECATION_MSG, FutureWarning)
+            self._table_name = Metadata.DEFAULT_SINGLE_TABLE_NAME
+            self.metadata = Metadata.load_from_dict(metadata.to_dict(), self._table_name)
+            self.metadata.tables[self._table_name]._updated = metadata._updated
 
-        self._validate_inputs(enforce_min_max_values, enforce_rounding)
         self.metadata.validate()
         self._check_metadata_updated()
         self.enforce_min_max_values = enforce_min_max_values
         self.enforce_rounding = enforce_rounding
         self.locales = locales
         self._data_processor = DataProcessor(
-            metadata=self.metadata,
+            metadata=self.metadata._convert_to_single_table(),
             enforce_rounding=self.enforce_rounding,
             enforce_min_max_values=self.enforce_min_max_values,
             locales=self.locales,
@@ -158,7 +158,7 @@ def _validate_metadata(self, data):
         """Validate that the data follows the metadata."""
         errors = []
         try:
-            self.metadata.validate_data(data)
+            self.metadata.validate_data({self._table_name: data})
         except InvalidDataError as error:
             errors += error.errors
 
@@ -183,10 +183,13 @@ def _validate(self, data):
         """
         return []
 
+    def _get_table_metadata(self):
+        return self.metadata.tables.get(self._table_name, SingleTableMetadata())
+
     def _validate_primary_key(self, data):
-        primary_key = self.metadata.primary_key
+        primary_key = self._get_table_metadata().primary_key
         is_int = primary_key and pd.api.types.is_integer_dtype(data[primary_key])
-        regex = self.metadata.columns.get(primary_key, {}).get('regex_format')
+        regex = self._get_table_metadata().columns.get(primary_key, {}).get('regex_format')
         if is_int and regex:
             possible_characters = get_possible_chars(regex, 1)
             if '0' in possible_characters:
@@ -225,8 +228,8 @@ def validate(self, data):
             raise InvalidDataError(synthesizer_errors)
 
     def _validate_transformers(self, column_name_to_transformer):
-        primary_and_alternate_keys = self.metadata._get_primary_and_alternate_keys()
-        sequence_keys = self.metadata._get_set_of_sequence_keys()
+        primary_and_alternate_keys = self._get_table_metadata()._get_primary_and_alternate_keys()
+        sequence_keys = self._get_table_metadata()._get_set_of_sequence_keys()
         keys = primary_and_alternate_keys | sequence_keys
         for column, transformer in column_name_to_transformer.items():
             if transformer is None:
@@ -251,8 +254,9 @@ def _warn_quality_and_performance(self, column_name_to_transformer):
             column_name_to_transformer (dict):
                 Dict mapping column names to transformers to be used for that column.
         """
+        table_metadata = self._get_table_metadata()
         for column in column_name_to_transformer:
-            sdtype = self.metadata.columns.get(column, {}).get('sdtype')
+            sdtype = table_metadata.columns.get(column, {}).get('sdtype')
             if sdtype in {'categorical', 'boolean'}:
                 warnings.warn(
                     f"Replacing the default transformer for column '{column}' "
@@ -304,8 +308,10 @@ def get_parameters(self):
 
     def get_metadata(self):
         """Return the ``Metadata`` for this synthesizer."""
-        table_name = getattr(self, '_table_name', None)
-        return Metadata.load_from_dict(self.metadata.to_dict(), table_name)
+        if isinstance(self.metadata, SingleTableMetadata):
+            table_name = getattr(self, '_table_name', None)
+            return Metadata.load_from_dict(self.metadata.to_dict(), table_name)
+        return self.metadata
 
     def load_custom_constraint_classes(self, filepath, class_names):
         """Load a custom constraint class for the current synthesizer.
@@ -387,9 +393,10 @@ def get_transformers(self):
             )
 
         # Order the output to match metadata
+        table_metadata = self._get_table_metadata()
         ordered_field_transformers = {
             column_name: field_transformers.get(column_name)
-            for column_name in self.metadata.columns
+            for column_name in table_metadata.columns
             if column_name in field_transformers
         }
 

sdv/single_table/copulagan.py

@@ -164,7 +164,10 @@ def __init__(
             cuda=cuda,
         )
 
-        validate_numerical_distributions(numerical_distributions, self.metadata.columns)
+        validate_numerical_distributions(
+            numerical_distributions,
+            self._get_table_metadata().columns,
+        )
         self.numerical_distributions = numerical_distributions or {}
         self.default_distribution = default_distribution or 'beta'
 
@@ -177,7 +180,7 @@ def __init__(
         }
 
     def _create_gaussian_normalizer_config(self, processed_data):
-        columns = self.metadata.columns
+        columns = self._get_table_metadata().columns
         transformers = {}
         sdtypes = {}
         for column in processed_data.columns:

sdv/single_table/copulas.py

@@ -110,7 +110,10 @@ def __init__(
             enforce_rounding=enforce_rounding,
             locales=locales,
         )
-        validate_numerical_distributions(numerical_distributions, self.metadata.columns)
+        validate_numerical_distributions(
+            numerical_distributions,
+            self._get_table_metadata().columns,
+        )
 
         self.default_distribution = default_distribution or 'beta'
         self._default_distribution = self.get_distribution_class(self.default_distribution)
@@ -192,8 +195,9 @@ def _sample(self, num_rows, conditions=None):
 
     def _get_valid_columns_from_metadata(self, columns):
         valid_columns = []
+        table_metadata = self._get_table_metadata()
         for column in columns:
-            for valid_column in self.metadata.columns:
+            for valid_column in table_metadata.columns:
                 if column.startswith(valid_column):
                     valid_columns.append(column)
                     break

sdv/single_table/utils.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from sdv.errors import SynthesizerInputError
+from sdv.metadata import Metadata
 
 DISABLE_TMP_FILE = 'disable'
 IGNORED_DICT_KEYS = ['fitted', 'distribution', 'type']
@@ -19,7 +20,7 @@ def detect_discrete_columns(metadata, data, transformers):
     discrete.
 
     Args:
-        metadata (sdv.metadata.SingleTableMetadata):
+        metadata (sdv.metadata.Metadata):
             Metadata that belongs to the given ``data``.
 
         data (pandas.DataFrame):
@@ -33,6 +34,9 @@ def detect_discrete_columns(metadata, data, transformers):
         discrete_columns (list):
             A list of discrete columns to be used with some of ``sdv`` synthesizers.
     """
+    if isinstance(metadata, Metadata):
+        metadata = metadata._convert_to_single_table()
+
     discrete_columns = []
     for column in data.columns:
         if column in metadata.columns:

tests/integration/metadata/test_metadata.py

@@ -570,7 +570,7 @@ def test_single_table_compatibility(tmp_path):
     loaded_synthesizer = GaussianCopulaSynthesizer.load(model_path)
     assert isinstance(synthesizer, GaussianCopulaSynthesizer)
     assert loaded_synthesizer.get_info() == synthesizer.get_info()
-    assert loaded_synthesizer.metadata.to_dict() == metadata.to_dict()
+    assert loaded_synthesizer.metadata._convert_to_single_table().to_dict() == metadata.to_dict()
     loaded_sample = loaded_synthesizer.sample(10)
     synthesizer.validate(loaded_sample)
 

tests/integration/sequential/test_par.py

@@ -124,7 +124,7 @@ def test_save_and_load(tmp_path):
 
     # Assert
     assert isinstance(loaded_instance, PARSynthesizer)
-    assert metadata._convert_to_single_table().to_dict() == instance.metadata.to_dict()
+    assert metadata.to_dict() == instance.metadata.to_dict()
 
 
 def test_synthesize_sequences(tmp_path):
@@ -193,7 +193,7 @@ def test_synthesize_sequences(tmp_path):
     assert model_path.exists()
     assert model_path.is_file()
     assert loaded_synthesizer.get_info() == synthesizer.get_info()
-    assert loaded_synthesizer.metadata.to_dict() == metadata._convert_to_single_table().to_dict()
+    assert loaded_synthesizer.metadata.to_dict() == metadata.to_dict()
     synthesizer.validate(synthetic_data)
     synthesizer.validate(custom_synthetic_data)
     synthesizer.validate(custom_synthetic_data_conditional)

tests/integration/single_table/test_base.py

@@ -526,12 +526,7 @@ def test_save_and_load(tmp_path):
 
     # Assert
     assert isinstance(loaded_instance, BaseSingleTableSynthesizer)
-    assert loaded_instance.metadata.columns == {}
-    assert loaded_instance.metadata.primary_key is None
-    assert loaded_instance.metadata.alternate_keys == []
-    assert loaded_instance.metadata.sequence_key is None
-    assert loaded_instance.metadata.sequence_index is None
-    assert loaded_instance.metadata._version == 'SINGLE_TABLE_V1'
+    assert loaded_instance.metadata.tables == {}
     assert instance._synthesizer_id == loaded_instance._synthesizer_id
 
 
@@ -550,12 +545,8 @@ def test_save_and_load_no_id(tmp_path):
 
     # Assert
     assert isinstance(loaded_instance, BaseSingleTableSynthesizer)
-    assert loaded_instance.metadata.columns == {}
-    assert loaded_instance.metadata.primary_key is None
-    assert loaded_instance.metadata.alternate_keys == []
-    assert loaded_instance.metadata.sequence_key is None
-    assert loaded_instance.metadata.sequence_index is None
-    assert loaded_instance.metadata._version == 'SINGLE_TABLE_V1'
+
+    assert loaded_instance.metadata.tables == {}
     assert hasattr(instance, '_synthesizer_id') is False
     assert hasattr(loaded_instance, '_synthesizer_id') is True
     assert isinstance(loaded_instance._synthesizer_id, str) is True
@@ -718,10 +709,10 @@ def test_metadata_updated_warning(method, kwargs):
     single_metadata = metadata._convert_to_single_table()
     single_metadata.__getattribute__(method)(**kwargs)
     with pytest.warns(UserWarning, match=expected_message):
-        BaseSingleTableSynthesizer(single_metadata)
+        instance = BaseSingleTableSynthesizer(single_metadata)
 
     # Assert
-    assert single_metadata._updated is False
+    assert instance.metadata.tables['table']._updated is False
 
 
 def test_fit_raises_version_error():