[Feature] Composite Keys (#863)

Author: pvk-developer

README.md

@@ -91,10 +91,7 @@ my_report = QualityReport.load(filepath='demo_data_quality_report.pkl')
 # set by the real data
 from sdmetrics.single_column import BoundaryAdherence
 
-BoundaryAdherence.compute(
-    real_data['start_date'],
-    synthetic_data['start_date']
-)
+BoundaryAdherence.compute(real_data['start_date'], synthetic_data['start_date'])
 ```
 ```
 0.8503937007874016
@@ -104,11 +101,7 @@ BoundaryAdherence.compute(
 # calculate whether the synthetic data is new or whether it's an exact copy of the real data
 from sdmetrics.single_table import NewRowSynthesis
 
-NewRowSynthesis.compute(
-    real_data,
-    synthetic_data,
-    metadata
-)
+NewRowSynthesis.compute(real_data, synthetic_data, metadata)
 ```
 ```
 1.0

sdmetrics/column_pairs/statistical/cardinality_boundary_adherence.py

@@ -33,10 +33,10 @@ def compute_breakdown(real_data, synthetic_data):
         """Calculate the percentage of synthetic parents with cardinality in the correct range.
 
         Args:
-            real_data (tuple(pd.Series, pd.Series)):
-                A tuple with the real primary key Series as the first element and real
-                foreign keys Series as the second element.
-            synthetic_data (tuple(pd.Series, pd.Series)):
+            real_data (tuple(pd.DataFrame, pd.DataFrame) or tuple(pd.Series, pd.Series)):
+                A tuple with the real primary key DataFrame as the first element and real
+                foreign keys DataFrame as the second element.
+            synthetic_data (tuple(pd.DataFrame, pd.DataFrame) or tuple(pd.Series, pd.Series)):
                 A tuple with the synthetic primary key as the first element and synthetic
                 foreign keys as the second element.
 
@@ -45,10 +45,10 @@ def compute_breakdown(real_data, synthetic_data):
                 Metric output.
         """
         real_cardinality = pd.DataFrame(index=real_data[0].copy())
-        real_cardinality['cardinality'] = real_data[1].value_counts()
+        real_cardinality['cardinality'] = real_data[1].value_counts(dropna=False)
         real_cardinality = real_cardinality.fillna(0)
         synthetic_cardinality = pd.DataFrame(index=synthetic_data[0].copy())
-        synthetic_cardinality['cardinality'] = synthetic_data[1].value_counts()
+        synthetic_cardinality['cardinality'] = synthetic_data[1].value_counts(dropna=False)
         synthetic_cardinality = synthetic_cardinality.fillna(0)
 
         min_cardinality = real_cardinality['cardinality'].min()
@@ -66,10 +66,10 @@ def compute(cls, real_data, synthetic_data):
         """Calculate the percentage of synthetic parents with cardinality in the correct range.
 
         Args:
-            real_data (tuple(pd.Series, pd.Series)):
-                A tuple with the real primary key Series as the first element and real
-                foreign keys Series as the second element.
-            synthetic_data (tuple(pd.Series, pd.Series)):
+            real_data (tuple(pd.DataFrame, pd.DataFrame) or tuple(pd.Series, pd.Series)):
+                A tuple with the real primary key DataFrame as the first element and real
+                foreign keys DataFrame as the second element.
+            synthetic_data (tuple(pd.DataFrame, pd.DataFrame) or tuple(pd.Series, pd.Series)):
                 A tuple with the synthetic primary key as the first element and synthetic
                 foreign keys as the second element.
 

sdmetrics/column_pairs/statistical/contingency_similarity.py

@@ -117,7 +117,9 @@ def compute_breakdown(
         contingency_synthetic = synthetic.groupby(list(columns), dropna=False).size() / len(
             synthetic
         )
-        combined_index = contingency_real.index.union(contingency_synthetic.index, sort=False)
+        combined_index = contingency_real.index.union(
+            contingency_synthetic.index, sort=False
+        ).drop_duplicates()
         contingency_synthetic = contingency_synthetic.reindex(combined_index, fill_value=0)
         contingency_real = contingency_real.reindex(combined_index, fill_value=0)
         diff = abs(contingency_real - contingency_synthetic).fillna(0)

sdmetrics/column_pairs/statistical/referential_integrity.py

@@ -31,41 +31,70 @@ class ReferentialIntegrity(ColumnPairsMetric):
     goal = Goal.MAXIMIZE
     min_value = 0.0
     max_value = 1.0
+    INDICATOR_NAME = '__ri_indicator__'
+
+    @staticmethod
+    def _create_unique_name(name, list_names):
+        """Modify the ``name`` parameter if it already exists in the list of names."""
+        result = name
+        while result in list_names:
+            result += '_'
+
+        return result
 
     @classmethod
     def compute_breakdown(cls, real_data, synthetic_data):
         """Compute the score breakdown of the referential integrity metric.
 
         Args:
-            real_data (tuple of 2 pandas.Series):
+            real_data (tuple of 2 pandas.DataFrame):
                 (primary_key, foreign_key) columns from the real data.
-            synthetic_data (tuple of 2 pandas.Series):
+            synthetic_data (tuple of 2 pandas.DataFrame):
                 (primary_key, foreign_key) columns from the synthetic data.
 
         Returns:
             dict:
                 The score breakdown of the key uniqueness metric.
         """
-        if pd.isna(real_data[1]).any():
-            synthetic_data = list(synthetic_data)
-            synthetic_data[1] = synthetic_data[1].dropna()
-
-        missing_parents = not real_data[1].isin(real_data[0]).all()
+        real_pk_df, real_fk_df = real_data
+        synth_pk_df, synth_fk_df = synthetic_data
+        pk_columns = list(real_pk_df.columns)
+        fk_columns = list(real_fk_df.columns)
+        indicator_name = cls._create_unique_name(cls.INDICATOR_NAME, pk_columns + fk_columns)
+
+        real_merged = real_fk_df.merge(
+            real_pk_df.drop_duplicates(),
+            how='left',
+            left_on=fk_columns,
+            right_on=pk_columns,
+            indicator=indicator_name,
+        )
+        missing_parents = (real_merged[indicator_name] == 'left_only').any()
         if missing_parents:
             LOGGER.info("The real data has foreign keys that don't reference any primary key.")
 
-        score = synthetic_data[1].isin(synthetic_data[0]).mean()
+        if len(fk_columns) == 1 and pd.isna(real_fk_df[fk_columns[0]]).any():
+            synth_fk_df = synth_fk_df.dropna()
+
+        synth_merged = synth_fk_df.merge(
+            synth_pk_df.drop_duplicates(),
+            how='left',
+            left_on=fk_columns,
+            right_on=pk_columns,
+            indicator=indicator_name,
+        )
 
+        score = (synth_merged[indicator_name] == 'both').mean()
         return {'score': score}
 
     @classmethod
     def compute(cls, real_data, synthetic_data):
         """Compute the referential integrity of two columns.
 
         Args:
-            real_data (tuple of 2 pandas.Series):
+            real_data (tuple of 2 pandas.DataFrame):
                 (primary_key, foreign_key) columns from the real data.
-            synthetic_data (tuple of 2 pandas.Series):
+            synthetic_data (tuple of 2 pandas.DataFrame):
                 (primary_key, foreign_key) columns from the synthetic data.
 
         Returns:

sdmetrics/multi_table/README.md

@@ -91,6 +91,6 @@ For example, we could execute the same metric as before by passing the `metadata
 of having to specify the individual `foreign_keys`:
 
 ```python
-In [10]: LogisticParentChildDetection.compute(real_data, synthetic_data, metadata)
+In[10]: LogisticParentChildDetection.compute(real_data, synthetic_data, metadata)
 Out[10]: 0.7569444444444444
 ```

sdmetrics/reports/multi_table/_properties/inter_table_trends.py

@@ -11,6 +11,7 @@
     ColumnPairTrends as SingleTableColumnPairTrends,
 )
 from sdmetrics.reports.utils import PlotConfig
+from sdmetrics.utils import _cast_to_iterable
 
 
 class InterTableTrends(BaseMultiTableProperty):
@@ -50,16 +51,16 @@ def _denormalize_tables(self, real_data, synthetic_data, relationship):
         """
         parent = relationship['parent_table_name']
         child = relationship['child_table_name']
-        foreign_key = relationship['child_foreign_key']
-        primary_key = relationship['parent_primary_key']
+        foreign_key = _cast_to_iterable(relationship['child_foreign_key'])
+        primary_key = _cast_to_iterable(relationship['parent_primary_key'])
 
         real_parent = real_data[parent].add_prefix(f'{parent}.')
         real_child = real_data[child].add_prefix(f'{child}.')
         synthetic_parent = synthetic_data[parent].add_prefix(f'{parent}.')
         synthetic_child = synthetic_data[child].add_prefix(f'{child}.')
 
-        child_index = f'{child}.{foreign_key}'
-        parent_index = f'{parent}.{primary_key}'
+        child_index = [f'{child}.{key_col}' for key_col in foreign_key]
+        parent_index = [f'{parent}.{key_col}' for key_col in primary_key]
 
         denormalized_real = real_child.merge(
             real_parent, left_on=child_index, right_on=parent_index
@@ -101,7 +102,12 @@ def _merge_metadata(self, metadata, parent_table, child_table):
         merged_metadata['columns'] = {**child_cols, **parent_cols}
         if 'primary_key' in merged_metadata:
             primary_key = merged_metadata['primary_key']
-            merged_metadata['primary_key'] = f'{child_table}.{primary_key}'
+            if isinstance(primary_key, list):
+                merged_metadata['primary_key'] = [
+                    f'{child_table}.{pk_col}' for pk_col in primary_key
+                ]
+            else:
+                merged_metadata['primary_key'] = f'{child_table}.{primary_key}'
 
         return merged_metadata, list(parent_cols.keys()), list(child_cols.keys())
 
@@ -123,6 +129,7 @@ def _generate_details(self, real_data, synthetic_data, metadata, progress_bar=No
             parent = relationship['parent_table_name']
             child = relationship['child_table_name']
             foreign_key = relationship['child_foreign_key']
+            fk_tuple = tuple(foreign_key) if isinstance(foreign_key, list) else foreign_key
 
             denormalized_real, denormalized_synthetic = self._denormalize_tables(
                 real_data, synthetic_data, relationship
@@ -132,14 +139,14 @@ def _generate_details(self, real_data, synthetic_data, metadata, progress_bar=No
 
             parent_child_pairs = itertools.product(parent_cols, child_cols)
 
-            self._properties[(parent, child, foreign_key)] = SingleTableColumnPairTrends()
+            self._properties[(parent, child, fk_tuple)] = SingleTableColumnPairTrends()
             self._properties[
-                (parent, child, foreign_key)
+                (parent, child, fk_tuple)
             ].real_correlation_threshold = self.real_correlation_threshold
             self._properties[
-                (parent, child, foreign_key)
+                (parent, child, fk_tuple)
             ].real_association_threshold = self.real_association_threshold
-            details = self._properties[(parent, child, foreign_key)]._generate_details(
+            details = self._properties[(parent, child, fk_tuple)]._generate_details(
                 denormalized_real,
                 denormalized_synthetic,
                 merged_metadata,
@@ -149,7 +156,7 @@ def _generate_details(self, real_data, synthetic_data, metadata, progress_bar=No
 
             details['Parent Table'] = parent
             details['Child Table'] = child
-            details['Foreign Key'] = foreign_key
+            details['Foreign Key'] = str(foreign_key)
             if not details.empty:
                 details['Column 1'] = details['Column 1'].str.replace(
                     f'{parent}.', '', n=1, regex=False
@@ -233,18 +240,15 @@ def _compute_average_score(self, to_plot):
     def get_visualization(self, table_name=None):
         """Create a plot to show the inter table trends data.
 
-        Returns:
-            plotly.graph_objects._figure.Figure
-
         Args:
             table_name (str, optional):
                 Table to plot. Defaults to None.
 
-        Raises:
-            - ``ValueError`` if property has not been computed.
-
         Returns:
             plotly.graph_objects._figure.Figure
+
+        Raises:
+            - ``ValueError`` if property has not been computed.
         """
         if not self.is_computed:
             raise ValueError(

sdmetrics/reports/multi_table/base_multi_table_report.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 from sdmetrics.reports.base_report import BaseReport
+from sdmetrics.utils import _cast_to_iterable
 from sdmetrics.visualization import set_plotly_config
 
 
@@ -43,22 +44,36 @@ def _validate_data_format(self, real_data, synthetic_data):
     def _validate_relationships(self, real_data, synthetic_data, metadata):
         """Validate that the relationships are valid."""
         for rel in metadata.get('relationships', []):
-            parent_dtype = real_data[rel['parent_table_name']][rel['parent_primary_key']].dtype
-            child_dtype = real_data[rel['child_table_name']][rel['child_foreign_key']].dtype
-            if (parent_dtype == 'object' and child_dtype != 'object') or (
-                parent_dtype != 'object' and child_dtype == 'object'
-            ):
-                parent = rel['parent_table_name']
-                parent_key = rel['parent_primary_key']
-                child = rel['child_table_name']
-                child_key = rel['child_foreign_key']
+            parent = rel['parent_table_name']
+            parent_key = rel['parent_primary_key']
+            child = rel['child_table_name']
+            child_key = rel['child_foreign_key']
+            parent_key_str = f"'{parent_key}'" if isinstance(parent_key, str) else str(parent_key)
+            child_key_str = f"'{child_key}'" if isinstance(child_key, str) else str(child_key)
+            parent_primary_key = _cast_to_iterable(parent_key)
+            child_foreign_key = _cast_to_iterable(child_key)
+
+            if len(parent_primary_key) != len(child_foreign_key):
                 error_msg = (
                     f"The '{parent}' table and '{child}' table cannot be merged "
-                    'for computing the cardinality. Please make sure the primary key'
-                    f" in '{parent}' ('{parent_key}') and the foreign key in '{child}'"
-                    f" ('{child_key}') have the same data type."
+                    'for computing the cardinality. Please make sure the number of columns '
+                    f'in the primary key ({parent_key_str}) matches the number of '
+                    f'columns in the foreign key ({child_key_str}).'
                 )
                 raise ValueError(error_msg)
+            parent_dtypes = real_data[rel['parent_table_name']][parent_primary_key].dtypes
+            child_dtypes = real_data[rel['child_table_name']][child_foreign_key].dtypes
+            for parent_dtype, child_dtype in zip(parent_dtypes, child_dtypes):
+                if (parent_dtype == 'object' and child_dtype != 'object') or (
+                    parent_dtype != 'object' and child_dtype == 'object'
+                ):
+                    error_msg = (
+                        f"The '{parent}' table and '{child}' table cannot be merged "
+                        'for computing the cardinality. Please make sure the primary key'
+                        f" in '{parent}' ({parent_key_str}) and the foreign key in '{child}'"
+                        f' ({child_key_str}) have the same data types.'
+                    )
+                    raise ValueError(error_msg)
 
     def _validate_metadata_matches_data(self, real_data, synthetic_data, metadata):
         """Validate that the metadata matches the data."""

sdmetrics/reports/single_table/_properties/data_validity.py

@@ -39,13 +39,19 @@ def _generate_details(self, real_data, synthetic_data, metadata, progress_bar=No
                 The progress bar to use. Defaults to None.
         """
         column_names, metric_names, scores = [], [], []
+        column_sdtypes = [(col, metadata['columns'][col]['sdtype']) for col in metadata['columns']]
         error_messages = []
         primary_key = metadata.get('primary_key')
+        if isinstance(primary_key, list):
+            if len(primary_key) > 1:
+                column_sdtypes = [(primary_key, None)] + column_sdtypes
+            else:
+                primary_key = primary_key[0]
+
         alternate_keys = metadata.get('alternate_keys', [])
         sequence_index = metadata.get('sequence_index')
 
-        for column_name in metadata['columns']:
-            sdtype = metadata['columns'][column_name]['sdtype']
+        for column_name, sdtype in column_sdtypes:
             primary_key_match = column_name == primary_key
             alternate_key_match = column_name in alternate_keys
             is_unique = primary_key_match or alternate_key_match