diff --git a/sdmetrics/multi_table/statistical/constraints/__init__.py b/sdmetrics/multi_table/statistical/constraints/__init__.py index 5b984c03..bb6ef3ec 100644 --- a/sdmetrics/multi_table/statistical/constraints/__init__.py +++ b/sdmetrics/multi_table/statistical/constraints/__init__.py @@ -1,7 +1,11 @@ """Constraints.""" from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.carry_over_columns import CarryOverColumns from sdmetrics.multi_table.statistical.constraints.column_formula import ColumnFormula +from sdmetrics.multi_table.statistical.constraints.denormalized_table import ( + DenormalizedTable, +) from sdmetrics.multi_table.statistical.constraints.chained_inequality import ChainedInequality from sdmetrics.multi_table.statistical.constraints.inequality import Inequality from sdmetrics.multi_table.statistical.constraints.fixed_combinations import FixedCombinations @@ -9,7 +13,20 @@ from sdmetrics.multi_table.statistical.constraints.fixed_null_combinations import ( FixedNullCombinations, ) +from sdmetrics.multi_table.statistical.constraints.foreign_to_foreign_key import ( + ForeignToForeignKey, +) +from sdmetrics.multi_table.statistical.constraints.foreign_to_primary_key_subset import ( + ForeignToPrimaryKeySubset, +) from sdmetrics.multi_table.statistical.constraints.mixed_scales import MixedScales +from sdmetrics.multi_table.statistical.constraints.polymorphic_relationship import ( + PolymorphicRelationship, +) +from sdmetrics.multi_table.statistical.constraints.primary_to_primary_key_subset import ( + PrimaryToPrimaryKeySubset, +) +from sdmetrics.multi_table.statistical.constraints.reference_table import ReferenceTable from sdmetrics.multi_table.statistical.constraints.one_hot_encoding import OneHotEncoding from sdmetrics.multi_table.statistical.constraints.range import Range from sdmetrics.multi_table.statistical.constraints.referential_hierarchy import ( @@ -18,13 +35,20 @@ __all__ = ( BaseConstraint, + CarryOverColumns, ColumnFormula, + DenormalizedTable, ChainedInequality, Inequality, FixedIncrements, FixedCombinations, FixedNullCombinations, + ForeignToForeignKey, + ForeignToPrimaryKeySubset, MixedScales, + PolymorphicRelationship, + PrimaryToPrimaryKeySubset, + ReferenceTable, OneHotEncoding, Range, SelfReferentialHierarchy, diff --git a/sdmetrics/multi_table/statistical/constraints/_utils.py b/sdmetrics/multi_table/statistical/constraints/_utils.py index dddfef4f..1baebc39 100644 --- a/sdmetrics/multi_table/statistical/constraints/_utils.py +++ b/sdmetrics/multi_table/statistical/constraints/_utils.py @@ -8,6 +8,10 @@ import pandas as pd from pandas.core.tools.datetimes import _guess_datetime_format_for_array +from sdmetrics.multi_table.statistical.constraints.error import ( + ConstraintNotApplicableError, +) + PRECISION_LEVELS = { '%Y': 1, # Year '%y': 1, # Year without century (same precision as %Y) @@ -95,6 +99,190 @@ def _is_list_of_type(values, type_to_check=str): return isinstance(values, list) and all(isinstance(value, type_to_check) for value in values) +def _get_key_values(table_data, key_columns): + """Return one hashable value per row for the given key columns. + + Every missing value is mapped to ``None`` so that two rows that are null in the + same columns produce equal values. + + Args: + table_data (pandas.DataFrame): + The data of the table. + key_columns (list[str]): + The names of the columns that make up the key. + + Returns: + pandas.Series: + A tuple with the values of ``key_columns`` for every row. + """ + return pd.Series( + [_tuple_from_columns(row, key_columns) for _, row in table_data[key_columns].iterrows()], + index=table_data.index, + dtype=object, + ) + + +def _validate_foreign_to_primary_key_subset_input( + parent_table_name, + child_table_name, + child_foreign_key, + conditional_column_name, + conditional_values, +): + """Validate the input for the ForeignToPrimaryKeySubset constraint.""" + if not isinstance(parent_table_name, str): + raise TypeError('`parent_table_name` must be a string.') + + if not isinstance(child_table_name, str): + raise TypeError('`child_table_name` must be a string.') + + if not isinstance(child_foreign_key, str) and not _is_list_of_type(child_foreign_key): + raise TypeError('`child_foreign_key` must be a string or a list of strings.') + + if not isinstance(conditional_column_name, str): + raise TypeError('`conditional_column_name` must be a string.') + + if not isinstance(conditional_values, list): + raise TypeError('`conditional_values` must be a list.') + + +def _validate_foreign_to_primary_key_subset( + data, + parent_primary_key, + parent_table_name, + child_table_name, + child_foreign_key, + conditional_column_name, + conditional_values, +): + """Validate the ForeignToPrimaryKeySubset constraint.""" + parent_primary_key = _cast_to_iterable(parent_primary_key) + child_foreign_key = _cast_to_iterable(child_foreign_key) + indicator_col = _create_unique_name('_merge', parent_primary_key + child_foreign_key) + parent_table = data[parent_table_name] + merged_parent = ( + parent_table[parent_primary_key] + .merge( + data[child_table_name][child_foreign_key].drop_duplicates(), + left_on=parent_primary_key, + right_on=child_foreign_key, + how='left', + indicator=indicator_col, + ) + .set_index(parent_table.index) + ) + filtered_parent = parent_table[merged_parent[indicator_col] == 'both'] + table_to_valid_rows = _get_table_to_valid_rows(data) + if not set(filtered_parent[conditional_column_name]).issubset(conditional_values): + good_parent_value_index = filtered_parent[conditional_column_name].isin(conditional_values) + good_parent_values = filtered_parent.loc[good_parent_value_index][parent_primary_key] + invalid_rows = ( + data[child_table_name][child_foreign_key] + .merge( + good_parent_values, + left_on=child_foreign_key, + right_on=parent_primary_key, + how='left', + indicator=indicator_col, + ) + .set_index(data[child_table_name].index)[indicator_col] + == 'left_only' + ) + table_to_valid_rows[child_table_name][invalid_rows] = False + + return table_to_valid_rows + + +def _validate_foreign_to_foreign_key_input(columns, foreign_key_generation): + """Validates a list of foreign key specifications. + + Args: + columns (list[dict]): + A list of dictionaries, each specifying a foreign key that are all connected. + Each dictionary should have the keys: + - 'table_name' (str): The name of the table. + - 'foreign_key' (str or tuple[str]): The foreign key column(s). + foreign_key_generation (str): + Method to use to generate new foreign key values. Must be one of ['new', 'reuse']. + + Raises: + ValueError: + If the ``columns`` value is not instance of list or dictionaries do not + contain the right inputs, or if ``foreign_key_generation`` value is not a string or + is an invalid option. + """ + expected_length = None + + if not isinstance(columns, list): + raise ValueError('columns must be a list of dictionaries') + + for entry in columns: + if not isinstance(entry, dict): + raise ValueError('Each entry in columns must be a dictionary') + + table_name = entry.get('table_name') + foreign_key = entry.get('foreign_key') + + if 'table_name' not in entry or not isinstance(table_name, str): + raise ValueError("Each dictionary must have a 'table_name' key with a string value") + + if 'foreign_key' not in entry: + raise ValueError("Each dictionary must have a 'foreign_key' key") + + if isinstance(foreign_key, str): + key_columns = [foreign_key] + elif isinstance(foreign_key, tuple) and all(isinstance(col, str) for col in foreign_key): + key_columns = list(foreign_key) + else: + raise ValueError("'foreign_key' must be a string or a tuple of strings") + + if expected_length is None: + expected_length = len(key_columns) + + elif len(key_columns) != expected_length: + raise ValueError( + 'All foreign key entries must have the same number of columns. ' + f"Entry for table '{table_name}' has {len(key_columns)} columns, " + f'expected {expected_length}.' + ) + + if not isinstance(foreign_key_generation, str): + raise ValueError('`foreign_key_generation` must be a string.') + + if foreign_key_generation not in ['new', 'reuse']: + raise ValueError( + f"Unrecognized `foreign_key_generation` value '{foreign_key_generation}'. " + "Must be one of ['new', 'reuse']." + ) + + +def _get_primary_key(metadata, table_name): + """Return the primary key of a table, as it is written in the metadata. + + Args: + metadata (dict): + The multi table metadata. + table_name (str): + The name of the table to get the primary key of. + + Returns: + str: + The name of the primary key column. + + Raises: + ConstraintNotApplicableError: + If the metadata does not give a primary key for the table. + """ + tables_metadata = (metadata or {}).get('tables', {}) + primary_key = tables_metadata.get(table_name, {}).get('primary_key') + if not isinstance(primary_key, str): + raise ConstraintNotApplicableError( + f"The table '{table_name}' does not have a primary key in the metadata." + ) + + return primary_key + + def _get_table_to_valid_rows(data): return {table: pd.Series(True, index=data[table].index) for table in data} diff --git a/sdmetrics/multi_table/statistical/constraints/carry_over_columns.py b/sdmetrics/multi_table/statistical/constraints/carry_over_columns.py new file mode 100644 index 00000000..25a5ba20 --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/carry_over_columns.py @@ -0,0 +1,148 @@ +"""Carry Over Columns Constraint.""" + +from copy import deepcopy + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _get_table_to_valid_rows, + _replace_nans_with_none, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +def _validate_carry_over_columns(data, common_column_info): + """Validate the CarryOverColumn constraint for the data. + + Validate that all carry over column rows share the same value for values across all + 'key_column_name' columns. Within a table, the carry over column's values should be + consistent for the same value of the key column. Between tables, the carry over column's + values should be consistent for the same value of the key column + + Args: + data (dict[pd.DataFrame]): + The data dictionary to validate. + common_column_info (list[dict]): + A list of dictionaries where each dictionary has the keys 'table_name', + 'carryover_column_name', and 'key_column_name'. + """ + keys_to_column_info = {} + key_value_pair = {} + table_to_valid_rows = _get_table_to_valid_rows(data) + for column_info in common_column_info: + carry_over_column = column_info['carryover_column_name'] + key_column = column_info['key_column_name'] + table_name = column_info['table_name'] + table = deepcopy(data[table_name]) + carry_over_grouped_by_key = table.groupby(key_column, dropna=False)[[carry_over_column]] + carry_over_unique_counts = carry_over_grouped_by_key.nunique(dropna=False) + keys_with_over_1_value = (carry_over_unique_counts > 1).any(axis=1) + inconsistent_vals = carry_over_unique_counts[keys_with_over_1_value].index + if len(inconsistent_vals) > 0: + table_to_valid_rows[table_name].loc[inconsistent_vals] = False + + table[key_column] = _replace_nans_with_none(table[key_column]) + table[carry_over_column] = _replace_nans_with_none(table[carry_over_column]) + mapping = table.set_index(key_column)[carry_over_column].to_dict() + for key, value in mapping.items(): + existing_value = key_value_pair.get(key) + if existing_value is None: + key_value_pair[key] = value + keys_to_column_info[key] = { + 'table_name': table_name, + 'carryover_column_name': carry_over_column, + 'key_column_name': key_column, + } + elif existing_value != value: + key_matches = table[key_column] == key + carry_over_matches = table[carry_over_column] == value + invalid_rows = table[key_matches & carry_over_matches] + table_to_valid_rows[table_name].loc[invalid_rows.index] = False + + return table_to_valid_rows + + +class CarryOverColumns(BaseConstraint): + """Constraint for a table that carries columns. + + The Carry Over Columns constraint that checks columns that were carried over + from a parent table to a child table. + + Args: + common_column_info (list[dict]): + A list of dictionaries containing the following keys: + - `table_name`: The name of the table. + - `carryover_column_name`: The name of the column to carry over. + - `key_column_name`: The name of the column to use as the shared key + for the carried over column. Must be a PII or ID sdtype. + """ + + _is_single_table = False + + def __init__(self, common_column_info): + super().__init__() + + expected_keys = {'table_name', 'carryover_column_name', 'key_column_name'} + if not isinstance(common_column_info, list): + raise TypeError('`common_column_info` must be a list.') + + for column_info in common_column_info: + if not isinstance(column_info, dict): + raise TypeError('Each element of `common_column_info` must be a dictionary.') + if not set(column_info.keys()) == expected_keys: + raise ValueError( + "Each element of `common_column_info` must have the keys 'table_name', " + "'carryover_column_name', and 'key_column_name'." + ) + + all_values_str = all(isinstance(column_info[key], str) for key in expected_keys) + if not all_values_str: + raise TypeError( + "The values of 'table_name', 'carryover_column_name', and 'key_column_name' " + 'in each element of `common_column_info` must be strings.' + ) + + self.common_column_info = common_column_info + self.table_name = None + table_names = set(column_info['table_name'] for column_info in common_column_info) + if len(table_names) == 1: + # required to work in single table synthesizer + self.table_name = table_names.pop() + self._is_single_table = True + + def _validate_data(self, data, metadata=None): + """Check that every table and all the referenced columns exist in the data.""" + for column_info in self.common_column_info: + table_name = column_info['table_name'] + if table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{table_name}' is missing from the data." + ) + + columns = data[table_name].columns + table_columns = [column_info['key_column_name'], column_info['carryover_column_name']] + missing_columns = [ + column_name for column_name in table_columns if column_name not in columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table '{table_name}'." + ) + + def _is_valid(self, data, metadata=None): + """Check that the data is valid. + + A row is considered invalid if the value of its carryover column does not match + the value that other rows with the same key have, in any of the tables. + + Args: + data (dict[str, pandas.DataFrame]): + The data dictionary. + metadata (dict): + Metadata as a dictionary. + + Returns: + dict[str, pandas.Series]: + Whether each row is valid. + """ + return _validate_carry_over_columns(data, self.common_column_info) diff --git a/sdmetrics/multi_table/statistical/constraints/denormalized_table.py b/sdmetrics/multi_table/statistical/constraints/denormalized_table.py new file mode 100644 index 00000000..52ac84a0 --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/denormalized_table.py @@ -0,0 +1,101 @@ +"""Denormalized Table Constraint.""" + +from sdmetrics.multi_table.statistical.constraints._utils import _get_table_to_valid_rows +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +class DenormalizedTable(BaseConstraint): + """Constraint for a table that contains denormalized columns. + + A denormalized table repeats the columns of a parent table on every row that + references the same parent. The data adheres to this constraint when every row + that shares a ``denormalized_primary_key`` value also shares the exact same + values for all the ``denormalized_column_names``. + + Args: + table_name (str): + The name of the denormalized table. + denormalized_primary_key (str): + The name of the column that contains the primary key of the parent table. + denormalized_column_names (list[str] or None): + The names of the columns that come from the parent table. If ``None`` or + empty, there is nothing to check and every row is considered valid. + """ + + def __init__(self, table_name, denormalized_primary_key, denormalized_column_names=None): + super().__init__() + if not isinstance(table_name, str): + raise ValueError("The 'table_name' parameter must be a string.") + + if not isinstance(denormalized_primary_key, str): + raise ValueError("The 'denormalized_primary_key' parameter must be a string.") + + if denormalized_column_names is None: + denormalized_column_names = [] + + is_list_of_strings = isinstance(denormalized_column_names, list) and all( + isinstance(column_name, str) for column_name in denormalized_column_names + ) + if not is_list_of_strings: + raise ValueError("The 'denormalized_column_names' parameter must be a list of strings.") + + if denormalized_primary_key in denormalized_column_names: + raise ValueError( + f"The column '{denormalized_primary_key}' cannot be both the " + "'denormalized_primary_key' and one of the 'denormalized_column_names'." + ) + + self.table_name = table_name + self.denormalized_primary_key = denormalized_primary_key + self.denormalized_column_names = list(denormalized_column_names) + + def _validate_data(self, data, metadata=None): + """Check that the table and all the referenced columns exist in the data.""" + if self.table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{self.table_name}' is missing from the data." + ) + + columns = data[self.table_name].columns + missing_columns = [ + column_name + for column_name in [self.denormalized_primary_key, *self.denormalized_column_names] + if column_name not in columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table '{self.table_name}'." + ) + + def _is_valid(self, data, metadata=None): + """Check that the data is valid. + + A row is considered invalid if the value in any column in denormalized_column_names + does not match the value for other instances of the same key. + + Args: + data (dict[str, pandas.DataFrame]): + The data dictionary. + metadata (dict): + Metadata as a dictionary. + + Returns: + dict[str, pandas.Series]: + Whether each row is valid. + """ + table = data[self.table_name] + table_to_valid_rows = _get_table_to_valid_rows(data) + if not self.denormalized_column_names or table.empty: + return table_to_valid_rows + + counts_per_row = table.groupby(self.denormalized_primary_key, dropna=False)[ + self.denormalized_column_names + ].transform(lambda col: col.nunique(dropna=False)) + + row_invalid = (counts_per_row > 1).any(axis=1) + if row_invalid.any(): + table_to_valid_rows[self.table_name].loc[row_invalid] = False + + return table_to_valid_rows diff --git a/sdmetrics/multi_table/statistical/constraints/foreign_to_foreign_key.py b/sdmetrics/multi_table/statistical/constraints/foreign_to_foreign_key.py new file mode 100644 index 00000000..136a425f --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/foreign_to_foreign_key.py @@ -0,0 +1,78 @@ +"""Foreign To Foreign Key Constraint.""" + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _get_table_to_valid_rows, + _validate_foreign_to_foreign_key_input, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +def _get_key_columns(foreign_key): + """Return the list of columns that make up a foreign key.""" + if isinstance(foreign_key, tuple): + return list(foreign_key) + + return [foreign_key] + + +class ForeignToForeignKey(BaseConstraint): + """Constraint to check many-to-many foreign key relationships. + + Args: + columns (list[dict]): + A list of dictionaries, each specifying a foreign key from a table + that is logically connected to others. Each dictionary should contain: + - `'table_name' (str)`: The name of the table containing the foreign key. + - `'foreign_key' (str | tuple[str])`: The name of the foreign key column, or a + tuple of column names if the foreign key is composite. + foreign_key_generation (str): + How to generate foreign key values. Must be on of `'new'` and `'reuse'`. If `'new'`, + the synthetic data will create entirely new foreign key values that will be shared + between the tables. If `'reuse'`, the same foreign key values will be reused from + the original data. Defaults to `'new'`. + """ + + _is_single_table = False + + def __init__(self, columns, foreign_key_generation='new'): + super().__init__() + + _validate_foreign_to_foreign_key_input(columns, foreign_key_generation) + self.columns = columns + self.foreign_key_generation = foreign_key_generation + + def _validate_data(self, data, metadata=None): + """Check that every table and all the referenced columns exist in the data.""" + for column_info in self.columns: + table_name = column_info['table_name'] + if table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{table_name}' is missing from the data." + ) + + columns = data[table_name].columns + missing_columns = [ + column_name + for column_name in _get_key_columns(column_info['foreign_key']) + if column_name not in columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table '{table_name}'." + ) + + def _is_valid(self, data, metadata=None): + """Check that the data is valid. + + Args: + data (dict[str, pandas.DataFrame]): + The data dictionary. + metadata (dict): + Metadata as a dictionary. + + Returns: + dict[str, pandas.Series]: + """ + return _get_table_to_valid_rows(data) diff --git a/sdmetrics/multi_table/statistical/constraints/foreign_to_primary_key_subset.py b/sdmetrics/multi_table/statistical/constraints/foreign_to_primary_key_subset.py new file mode 100644 index 00000000..2cbcd6df --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/foreign_to_primary_key_subset.py @@ -0,0 +1,111 @@ +"""Foreign To Primary Key Subset Constraint.""" + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _get_primary_key, + _validate_foreign_to_primary_key_subset, + _validate_foreign_to_primary_key_subset_input, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +class ForeignToPrimaryKeySubset(BaseConstraint): + """Constraint for a foreign key that may only reference a subset of the parent rows. + + Args: + parent_table_name (str): + Name of the parent table. + child_table_name (str): + Name of the child table. + child_foreign_key (str or list[str]): + Name of the column (or list of column names for composite keys) in the child table + that is a foreign key to the parent table. + conditional_column_name (str): + Name of the column in the parent table that defines the subset of valid primary key + values. + conditional_values (list): + List of values in the ``conditional_column_name`` column that define the subset of + valid primary key values. + """ + + _is_single_table = False + + def __init__( + self, + parent_table_name, + child_table_name, + child_foreign_key, + conditional_column_name, + conditional_values, + ): + super().__init__() + _validate_foreign_to_primary_key_subset_input( + parent_table_name, + child_table_name, + child_foreign_key, + conditional_column_name, + conditional_values, + ) + self.parent_table_name = parent_table_name + self.child_table_name = child_table_name + self.child_foreign_key = child_foreign_key + self.conditional_column_name = conditional_column_name + self.conditional_values = conditional_values + self._parent_primary_key = None + + def _validate_data(self, data, metadata=None): + """Check that both tables and all the referenced columns exist in the data.""" + table_to_columns = { + self.parent_table_name: [self.conditional_column_name], + self.child_table_name: [self.child_foreign_key], + } + for table_name, table_columns in table_to_columns.items(): + if table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{table_name}' is missing from the data." + ) + + columns = data[table_name].columns + missing_columns = [ + column_name for column_name in table_columns if column_name not in columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table '{table_name}'." + ) + + primary_key = _get_primary_key(metadata, self.parent_table_name) + if primary_key not in data[self.parent_table_name].columns: + raise ConstraintNotApplicableError( + f"The column(s) '{primary_key}' are missing from the table " + f"'{self.parent_table_name}'." + ) + + self._parent_primary_key = primary_key + + def _is_valid(self, data, metadata=None): + """Check that the data is valid. + + Args: + data (dict[str, pandas.DataFrame]): + The data dictionary. + metadata (dict): + Metadata as a dictionary. + + Returns: + dict[str, pandas.Series]: + """ + primary_key = self._parent_primary_key + if primary_key is None: + primary_key = _get_primary_key(metadata, self.parent_table_name) + + return _validate_foreign_to_primary_key_subset( + data, + primary_key, + self.parent_table_name, + self.child_table_name, + self.child_foreign_key, + self.conditional_column_name, + self.conditional_values, + ) diff --git a/sdmetrics/multi_table/statistical/constraints/polymorphic_relationship.py b/sdmetrics/multi_table/statistical/constraints/polymorphic_relationship.py new file mode 100644 index 00000000..cd5202e6 --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/polymorphic_relationship.py @@ -0,0 +1,295 @@ +"""Polymorphic Relationship Constraint.""" + +import itertools +from collections import defaultdict + +import pandas as pd + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _cast_to_iterable, + _create_unique_name, + _get_primary_key, + _get_table_to_valid_rows, + _is_list_of_type, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +class PolymorphicRelationship(BaseConstraint): + """Polymorphic relationship constraint. + + This constraint handles the case where a single column acts as a foreign key into + multiple possible tables, depending on either (1) the format of the foreign key + itself or (2) the value in another column. + + Args: + table_name (str) + The name of the table that contains the shared foreign key. + foreign_key (str or list[str]) + The name of the shared foreign key column that is present in the table. + parent_table_names (list[str]) + A list of table names that the foreign key values refer to. + type_column_name (str, optional) + The name of the categorical column in the table that sets which table the + foreign key references. If None, attempts to detect the parent table from the + foreign key value. Defaults to None. + type_value_to_table (dict, optional) + A map of category values in the type column to the parent table being referenced. + If None and `type_column_name` is passed, the table name is used as the type value + for each parent. Defaults to None. + """ + + _is_single_table = False + + def _validate_polymorphic_relationship_inputs( + self, + table_name, + foreign_key, + parent_table_names, + type_column_name, + type_value_to_table, + ): + if not isinstance(table_name, str): + raise TypeError('`table_name` must be a string.') + + if not isinstance(foreign_key, str) and not _is_list_of_type(foreign_key): + raise TypeError('`foreign_key` must be a string or a list of strings.') + + if not _is_list_of_type(parent_table_names, str): + raise TypeError('`parent_table_names` must be a list of strings.') + elif table_name in parent_table_names: + raise ConstraintNotApplicableError( + f"Table name '{table_name}' cannot also be in `parent_table_names`." + ) + + if type_column_name is not None and not isinstance(type_column_name, str): + raise TypeError('`type_column_name` must be a string or None.') + + if type_column_name in _cast_to_iterable(foreign_key): + raise ValueError('`foreign_key` and `type_column_name` must be different columns.') + + if type_value_to_table is not None: + if not isinstance(type_value_to_table, dict): + raise TypeError('`type_value_to_table` must be a dict or `None`.') + + extra_tables = set(type_value_to_table.values()) - set(parent_table_names) + missing_tables = set(parent_table_names) - set(type_value_to_table.values()) + if extra_tables: + extra = "', '".join(list(extra_tables)) + raise ValueError( + f"Table(s) '{extra}' in `type_values_to_table` not found " + 'in `parent_table_names` list.' + ) + if missing_tables: + missing = "', '".join(list(missing_tables)) + raise ValueError( + f"Table(s) '{missing}' in `parent_table_names` do not have any " + 'type value associated with them in `type_values_to_table`.' + ) + + def _get_parent_type_dicts(self): + type_value_to_parent = self.type_value_to_parent + parent_to_type_values = defaultdict(list) + if self.type_value_to_parent is not None: + for value, parent in self.type_value_to_parent.items(): + parent_to_type_values[parent].append(value) + else: + type_value_to_parent = {parent: parent for parent in self.parent_tables} + parent_to_type_values = {parent: [parent] for parent in self.parent_tables} + + return parent_to_type_values, type_value_to_parent + + def __init__( + self, + table_name, + foreign_key, + parent_table_names, + type_column_name=None, + type_value_to_table=None, + ): + super().__init__() + + self._validate_polymorphic_relationship_inputs( + table_name, + foreign_key, + parent_table_names, + type_column_name, + type_value_to_table, + ) + self.table_name = table_name + self.foreign_key = foreign_key + self._num_fk_cols = 1 if isinstance(foreign_key, str) else len(self.foreign_key) + self.parent_tables = parent_table_names + self.type_column = type_column_name + self.type_value_to_table = type_value_to_table + + self.type_value_to_parent = type_value_to_table if type_column_name else None + self._parent_to_type_values, self._type_value_to_parent = self._get_parent_type_dicts() + + def _validate_data(self, data, metadata=None): + """Check that every table and all the referenced columns exist in the data.""" + if self.table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{self.table_name}' is missing from the data." + ) + + columns = data[self.table_name].columns + table_columns = [self.foreign_key] + if self.type_column is not None: + table_columns.append(self.type_column) + + missing_columns = [ + column_name for column_name in table_columns if column_name not in columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table '{self.table_name}'." + ) + + for parent_table_name in self.parent_tables: + if parent_table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{parent_table_name}' is missing from the data." + ) + + primary_key = _get_primary_key(metadata, parent_table_name) + if primary_key not in data[parent_table_name].columns: + raise ConstraintNotApplicableError( + f"The column(s) '{primary_key}' are missing from the table " + f"'{parent_table_name}'." + ) + + def _get_foreign_key_groups(self, data, metadata): + table_data = data[self.table_name] + type_column = self.type_column + parent_to_types = self._parent_to_type_values + child_groups = {} + foreign_key = _cast_to_iterable(self.foreign_key) + if self.type_column is not None: + for parent, type_values in parent_to_types.items(): + child_rows_mask = table_data[type_column].isin(type_values) + child_groups[parent] = table_data[foreign_key][child_rows_mask] + + else: + child_dtypes = list(table_data[foreign_key].dtypes) + for parent in self.parent_tables: + parent_pk = _cast_to_iterable(_get_primary_key(metadata, parent)) + parent_pk_values = data[parent][parent_pk].astype({ + pk_col: dtype for pk_col, dtype in zip(parent_pk, child_dtypes) + }) + indicator_col = _create_unique_name('_merge', foreign_key + parent_pk) + merged_child = table_data[foreign_key].merge( + parent_pk_values, + left_on=foreign_key, + right_on=parent_pk, + how='left', + indicator=indicator_col, + ) + child_groups[parent] = table_data[merged_child[indicator_col] == 'both'][ + foreign_key + ] + + return child_groups + + def _validate_type_column(self, table_data): + type_column = table_data[self.type_column] + type_values = ( + self.parent_tables + if not self.type_value_to_parent + else self.type_value_to_parent.keys() + ) + bad_type_values = ~(type_column.isin(type_values) | type_column.isna()) + + return ~bad_type_values + + def _validate_parent_primary_keys(self, data, metadata): + is_valid_pk = { + parent: pd.Series(True, index=data[parent].index) for parent in self.parent_tables + } + primary_keys = {parent: _get_primary_key(metadata, parent) for parent in self.parent_tables} + overlapping_parents = {} + for parent1, parent2 in itertools.combinations(primary_keys, 2): + table1_pk = _cast_to_iterable(primary_keys[parent1]) + table2_pk = _cast_to_iterable(primary_keys[parent2]) + indicator_col = _create_unique_name('_merge', table1_pk + table2_pk) + invalid_parent2_mask = ( + data[parent2][table2_pk] + .astype('object') + .merge( + data[parent1][table1_pk].astype('object'), + left_on=table2_pk, + right_on=table1_pk, + how='left', + indicator=indicator_col, + ) + .set_index(data[parent2].index)[indicator_col] + == 'both' + ) + if any(invalid_parent2_mask): + overlapping_parents[(parent1, parent2)] = invalid_parent2_mask + + for (_, table2), invalid_pk_mask in overlapping_parents.items(): + is_valid_pk[table2] &= ~invalid_pk_mask + + return is_valid_pk + + def _validate_polymorphic_relationship_with_data(self, data, metadata): + table_to_valid_rows = _get_table_to_valid_rows(data) + table_data = data[self.table_name] + valid_rows = table_to_valid_rows[self.table_name] + referenced = pd.Series(False, index=table_data.index) + referenced[table_data[_cast_to_iterable(self.foreign_key)].isna().all(axis=1)] = True + + if self.type_column: + valid_rows[~self._validate_type_column(table_data)] = False + else: + table_to_valid_rows.update(self._validate_parent_primary_keys(data, metadata)) + valid_rows = table_to_valid_rows[self.table_name] + + child_groups = self._get_foreign_key_groups(data, metadata) + for parent, child_ids in child_groups.items(): + referenced[child_ids.index] = True + primary_key = _cast_to_iterable(_get_primary_key(metadata, parent)) + foreign_key = list(child_ids.columns) + indicator_col = _create_unique_name('_merge', primary_key + foreign_key) + unknown_keys_mask = ( + child_ids + .dropna(how='all') + .merge( + data[parent][primary_key], + left_on=foreign_key, + right_on=primary_key, + how='left', + indicator=indicator_col, + ) + .set_index(child_ids.dropna(how='all').index)[indicator_col] + == 'left_only' + ) + if any(unknown_keys_mask): + valid_rows[valid_rows & unknown_keys_mask] = False + + if not all(referenced): + valid_rows[~referenced] = False + + table_to_valid_rows[self.table_name] = valid_rows + return table_to_valid_rows + + def _is_valid(self, data, metadata): + """Return whether or not each row in the data dictionary is valid for the constraint. + + Args: + data (dict[str, pd.DataFrame]): + Table data. + metadata (sdv.Metadata): + Metadata for the data. + + Returns: + dict[str, pd.Series]: + A dictionary mapping the table name to a Series where each row is=True or False + depending on if it's valid. + """ + table_to_valid_rows = self._validate_polymorphic_relationship_with_data(data, metadata) + + return table_to_valid_rows diff --git a/sdmetrics/multi_table/statistical/constraints/primary_to_primary_key_subset.py b/sdmetrics/multi_table/statistical/constraints/primary_to_primary_key_subset.py new file mode 100644 index 00000000..287917c8 --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/primary_to_primary_key_subset.py @@ -0,0 +1,157 @@ +"""Primary To Primary Key Subset Constraint.""" + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _cast_to_iterable, + _create_unique_name, + _get_primary_key, + _get_table_to_valid_rows, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +class PrimaryToPrimaryKeySubset(BaseConstraint): + """Constraint for tables that hold a subset of the rows of a main table. + + Args: + main_table_name (str): + The name of the main table, the one that holds every possible row. + conditional_column_name (str): + The name of the column of the main table that controls whether a connection + is allowed. + relationships (dict): + A dictionary that maps the name of every connected table to the list of + conditional values that allow the connection. + """ + + _is_single_table = False + + @staticmethod + def _validate_inputs(main_table_name, conditional_column_name, relationships): + if not all(isinstance(value, str) for value in [main_table_name, conditional_column_name]): + raise ValueError('`main_table_name` and `conditional_column_name` must be strings') + + if not isinstance(relationships, dict) or not all( + isinstance(k, str) and isinstance(v, list) for k, v in relationships.items() + ): + raise ValueError( + '`relationships` must be a a dict that maps the name of the connected table to a ' + 'list of values that are acceptable for a connection to be made.' + ) + + def __init__(self, main_table_name, conditional_column_name, relationships): + super().__init__() + self._validate_inputs(main_table_name, conditional_column_name, relationships) + self.main_table_name = main_table_name + self.conditional_column_name = conditional_column_name + self.relationships = relationships + + def _validate_data(self, data, metadata=None): + """Check that every table and all the referenced columns exist in the data.""" + if self.main_table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{self.main_table_name}' is missing from the data." + ) + + main_columns = data[self.main_table_name].columns + main_primary_key = _get_primary_key(metadata, self.main_table_name) + missing_columns = [ + column_name + for column_name in [main_primary_key, self.conditional_column_name] + if column_name not in main_columns + ] + if missing_columns: + missing_columns = "', '".join(missing_columns) + raise ConstraintNotApplicableError( + f"The column(s) '{missing_columns}' are missing from the table " + f"'{self.main_table_name}'." + ) + + for table_name in self.relationships: + if table_name not in data: + raise ConstraintNotApplicableError( + f"The table '{table_name}' is missing from the data." + ) + + primary_key = _get_primary_key(metadata, table_name) + if primary_key not in data[table_name].columns: + raise ConstraintNotApplicableError( + f"The column(s) '{primary_key}' are missing from the table '{table_name}'." + ) + + def _get_metadata_parameters(self, metadata): + """Get the metadata parameters for the constraint. + + Return all the necessary metadata parameters to compute the updated metadata: + - table_to_pk: A dictionary that maps the table name to its primary key. + - main_table_columns: The columns of the main table. + - tables_to_column_names: A dictionary that maps the table name to the column names mapping + of its related table. The column names mapping is a dictionary that maps the original + column names to the new column names after merging the related table + into the main table. + + Args: + metadata (dict): + The input metadata for the constraint. + """ + table_to_pk = {} + table_to_pk[self.main_table_name] = _get_primary_key(metadata, self.main_table_name) + main_table_columns = metadata['tables'][self.main_table_name]['columns'].keys() + tables_to_column_names = {} + main_table_columns = list(main_table_columns) + existing_column_names = list(main_table_columns) + for table_name in self.relationships: + table_to_pk[table_name] = _get_primary_key(metadata, table_name) + table_columns = metadata['tables'][table_name]['columns'].keys() + column_names_to_merge = [f'{table_name}_{column_name}' for column_name in table_columns] + column_names_to_merge = [ + _create_unique_name(column_name, existing_column_names) + for column_name in column_names_to_merge + ] + conditional_value_to_column_name = dict(zip(table_columns, column_names_to_merge)) + for pk_col in _cast_to_iterable(table_to_pk[table_name]): + del conditional_value_to_column_name[pk_col] + + tables_to_column_names[table_name] = conditional_value_to_column_name + existing_column_names.extend(column_names_to_merge) + + return table_to_pk, main_table_columns, tables_to_column_names + + def _is_valid(self, data, metadata=None): + """Get all valid rows. + + A valid row is a row in the related table that has primary key value that is a subset + of the main table's primary key values and the condition matches. + + Args: + data (dict ): + Table data. + + Returns: + dict : + A dictionary mapping the table name to a Series where each row is=True or False + depending on if it's valid. + """ + table_to_pks, _, _ = self._get_metadata_parameters(metadata) + table_to_valid_rows = _get_table_to_valid_rows(data) + main_table_pk = _cast_to_iterable(table_to_pks[self.main_table_name]) + main_table_keys = data[self.main_table_name][main_table_pk] + conditional_col = data[self.main_table_name][self.conditional_column_name] + for table_name, conditional_values in self.relationships.items(): + valid_primary_key_values = main_table_keys[conditional_col.isin(conditional_values)] + table_pk = _cast_to_iterable(table_to_pks[table_name]) + indicator_col = _create_unique_name('_merge', table_pk + main_table_pk) + valid_pk_mask = ( + data[table_name][table_pk] + .astype('object') + .merge( + valid_primary_key_values.astype('object'), + left_on=table_pk, + right_on=main_table_pk, + how='left', + indicator=indicator_col, + ) + ) + table_to_valid_rows[table_name][valid_pk_mask[indicator_col] == 'left_only'] = False + + return table_to_valid_rows diff --git a/sdmetrics/multi_table/statistical/constraints/reference_table.py b/sdmetrics/multi_table/statistical/constraints/reference_table.py new file mode 100644 index 00000000..4529ee1e --- /dev/null +++ b/sdmetrics/multi_table/statistical/constraints/reference_table.py @@ -0,0 +1,87 @@ +"""Reference Table Constraint.""" + +from sdmetrics.multi_table.statistical.constraints._utils import ( + _get_table_to_valid_rows, +) +from sdmetrics.multi_table.statistical.constraints.base import BaseConstraint +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +class ReferenceTable(BaseConstraint): + """Constraint for tables whose rows connect to a reference table. + + Args: + reference_table_names (list[str]): + Names of the reference tables. + """ + + _is_single_table = False + + def __init__(self, reference_table_names): + super().__init__() + + if not isinstance(reference_table_names, list) or not all( + isinstance(name, str) for name in reference_table_names + ): + raise ValueError("'reference_table_names' must be a list of strings.") + + self.reference_table_names = reference_table_names + + def _validate_constraint_with_metadata(self, metadata): + """Validate the metadata for the constraint. + + This method: + - Validates that each reference table exists in the metadata. + - Validates that no reference table is a child of another table. + A reference table can be the child of another reference table. + + Args: + metadata (dict): + The metadata for the dataset. + + Raises: + ConstraintNotMetError: + If any reference table is missing from metadata + or is a child of a non-reference table. + """ + if any(table not in metadata['tables'] for table in self.reference_table_names): + missing = set(self.reference_table_names) - set(metadata['tables']) + raise ConstraintNotApplicableError( + f"Reference table(s) '{sorted(missing)}' missing from metadata." + ) + + invalid_pairs = set() + for relationship in metadata['relationships']: + parent = relationship['parent_table_name'] + child = relationship['child_table_name'] + if child in self.reference_table_names and parent not in self.reference_table_names: + invalid_pairs.add((child, parent)) + + if invalid_pairs: + raise ConstraintNotApplicableError( + 'Reference tables cannot be children of non-reference tables. ' + f"The following child-parent pairs are invalid: '{sorted(invalid_pairs)}'" + ) + + def _validate_data(self, data, metadata=None): + """No data validation needed for reference tables.""" + pass + + def _is_valid(self, data, metadata=None): + """Get valid rows. + + All rows are valid. + + Args: + data (dict[str, pd.DataFrame]): + Table data. + + Returns: + dict[str, pd.Series]: + A dictionary mapping the table name to a Series where each row is=True or False + depending on if it's valid. + """ + if metadata is not None: + self._validate_constraint_with_metadata(metadata) + + return _get_table_to_valid_rows(data) diff --git a/tests/unit/multi_table/statistical/constraints/test_carry_over_columns.py b/tests/unit/multi_table/statistical/constraints/test_carry_over_columns.py new file mode 100644 index 00000000..ba7d10ea --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_carry_over_columns.py @@ -0,0 +1,311 @@ +import re + +import numpy as np +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import CarryOverColumns +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def metadata(): + return { + 'tables': { + 'main_table': { + 'columns': { + 'primary_key': {'sdtype': 'id'}, + 'parent_1': {'sdtype': 'categorical'}, + }, + 'primary_key': 'primary_key', + }, + 'carry_over_1': { + 'columns': { + 'child_1': {'sdtype': 'categorical'}, + 'child_2': {'sdtype': 'categorical'}, + 'key_column_1': {'sdtype': 'id'}, + 'key_column_2': {'sdtype': 'id'}, + }, + }, + 'carry_over_2': { + 'columns': {'child_3': {'sdtype': 'categorical'}, 'foreign_key': {'sdtype': 'id'}}, + }, + }, + } + + +@pytest.fixture +def data(): + return { + 'main_table': pd.DataFrame({ + 'primary_key': [1, 2, 3, 4, 5, 6], + 'parent_1': ['a', 'b', 'c', 'a', 'b', 'c'], + }), + 'carry_over_1': pd.DataFrame({ + 'child_1': ['a', 'a', 'c', 'c', 'd', 'e', 'f'], + 'child_2': ['b', 'b', 'a', 'a', 'd', 'e', 'f'], + 'key_column_1': [1, 1, 3, 3, 7, 8, 9], + 'key_column_2': [2, 2, 4, 4, 7, 8, 9], + }), + 'carry_over_2': pd.DataFrame({ + 'child_3': ['a', 'b', 'c', 'b', 'c', 'd', 'e'], + 'foreign_key': [1, 2, 3, 5, 6, 7, 8], + }), + } + + +@pytest.fixture +def common_column_info(): + return [ + { + 'table_name': 'main_table', + 'key_column_name': 'primary_key', + 'carryover_column_name': 'parent_1', + }, + { + 'table_name': 'carry_over_1', + 'key_column_name': 'key_column_1', + 'carryover_column_name': 'child_1', + }, + { + 'table_name': 'carry_over_1', + 'key_column_name': 'key_column_2', + 'carryover_column_name': 'child_2', + }, + { + 'table_name': 'carry_over_2', + 'key_column_name': 'foreign_key', + 'carryover_column_name': 'child_3', + }, + ] + + +@pytest.fixture +def constraint(common_column_info): + return CarryOverColumns(common_column_info=common_column_info) + + +class TestCarryOverColumns: + def test___init__invalid_parameters(self, common_column_info): + """Test the ``__init__`` method errors with invalid arguments.""" + # Setup + err_msg = re.escape('`common_column_info` must be a list.') + + # Run and Assert + with pytest.raises(TypeError, match=err_msg): + CarryOverColumns(common_column_info=common_column_info[0]) + + def test___init__invalid_keys(self, common_column_info): + """Test the ``__init__`` method errors if an entry has the wrong keys.""" + # Setup + err_msg = re.escape( + "Each element of `common_column_info` must have the keys 'table_name', " + "'carryover_column_name', and 'key_column_name'." + ) + not_a_dict_msg = re.escape('Each element of `common_column_info` must be a dictionary.') + + # Run and Assert + with pytest.raises(ValueError, match=err_msg): + CarryOverColumns(common_column_info=[common_column_info[0], {'table_name': 'tableA'}]) + + with pytest.raises(TypeError, match=not_a_dict_msg): + CarryOverColumns(common_column_info=[common_column_info[0], 'tableA']) + + def test___init__invalid_values(self, common_column_info): + """Test the ``__init__`` method errors if an entry has non string values.""" + # Setup + invalid_info = {**common_column_info[1], 'key_column_name': 1} + err_msg = re.escape( + "The values of 'table_name', 'carryover_column_name', and 'key_column_name' " + 'in each element of `common_column_info` must be strings.' + ) + + # Run and Assert + with pytest.raises(TypeError, match=err_msg): + CarryOverColumns(common_column_info=[common_column_info[0], invalid_info]) + + def test__validate_data_missing_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if one of the tables is not in the data.""" + # Setup + del data['main_table'] + expected_error = re.escape("The table 'main_table' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_key_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the key column is not in the table.""" + # Setup + del data['main_table']['primary_key'] + expected_error = re.escape( + "The column(s) 'primary_key' are missing from the table 'main_table'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_carryover_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the carryover column is not in the table.""" + # Setup + del data['carry_over_1']['child_1'] + expected_error = re.escape( + "The column(s) 'child_1' are missing from the table 'carry_over_1'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__is_valid(self, data, metadata, constraint): + """Test ``_is_valid`` considers every row valid when the values match up.""" + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 6)) + pd.testing.assert_series_equal(is_valid['carry_over_1'], pd.Series([True] * 7)) + pd.testing.assert_series_equal(is_valid['carry_over_2'], pd.Series([True] * 7)) + + def test__is_valid_with_inconsistent_carry_over_table(self, data, metadata, constraint): + """Test ``_is_valid`` flags a key that is inconsistent within a single table.""" + # Setup + data['carry_over_1'].loc[3, 'child_1'] = 'd' + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + expected_carry_over_1 = pd.Series([True, True, True, False, True, True, True]) + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 6)) + pd.testing.assert_series_equal(is_valid['carry_over_1'], expected_carry_over_1) + pd.testing.assert_series_equal(is_valid['carry_over_2'], pd.Series([True] * 7)) + + def test__is_valid_with_invalid_values(self, data, metadata, constraint): + """Test ``_is_valid`` flags every row of a key that does not match up.""" + # Setup + data['carry_over_2'].loc[0, 'child_3'] = 'b' + data['carry_over_2'].loc[3, 'child_3'] = 'z' + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + expected_carry_over_2 = pd.Series([False, True, True, False, True, True, True]) + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 6)) + pd.testing.assert_series_equal(is_valid['carry_over_1'], pd.Series([True] * 7)) + pd.testing.assert_series_equal(is_valid['carry_over_2'], expected_carry_over_2) + + def test__is_valid_with_nans(self, data, metadata, constraint): + """Test ``_is_valid`` treats a missing carryover value as its own value.""" + # Setup + data['carry_over_1'].loc[4, 'child_1'] = np.nan + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 6)) + pd.testing.assert_series_equal(is_valid['carry_over_1'], pd.Series([True] * 7)) + pd.testing.assert_series_equal(is_valid['carry_over_2'], pd.Series([True] * 7)) + + def test__is_valid_with_unmatched_key(self, data, metadata, constraint): + """Test ``_is_valid`` considers a key that is only in one table valid.""" + # Setup + data['carry_over_2'].loc[6, 'foreign_key'] = 99 + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 6)) + pd.testing.assert_series_equal(is_valid['carry_over_1'], pd.Series([True] * 7)) + pd.testing.assert_series_equal(is_valid['carry_over_2'], pd.Series([True] * 7)) + + def test__is_valid_other_tables(self, data, metadata, constraint): + """Test ``_is_valid`` considers every row of the other tables valid.""" + # Setup + data['other_table'] = pd.DataFrame({'a': ['a', 'b']}) + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['other_table'], pd.Series([True, True])) + + def test__is_valid_mismatch(self, common_column_info): + """Test that rows are marked invalid if there are mismatching carry over columns. + + This test will check rows where: + - The carry over column is mismatched for the same key value. + - The same key value has multiple carry over column values in one table. + """ + # Setup + data = { + 'main_table': pd.DataFrame({ + 'primary_key': [1, 2, 3, 4, 5, 6], + 'parent_1': ['a', 'b', 'c', 'a', 'b', 'c'], + }), + 'carry_over_1': pd.DataFrame({ + 'child_1': ['a', 'b', 'c', 'c', 'd', 'e', 'f'], + 'child_2': ['b', 'b', 'a', 'a', 'd', 'e', 'f'], + 'key_column_1': [1, 1, 3, 3, 7, 8, 9], + 'key_column_2': [2, 2, 4, 4, 7, 8, 9], + }), + 'carry_over_2': pd.DataFrame({ + 'child_3': ['d', 'b', 'c', 'b', 'c', 'd', 'e'], + 'foreign_key': [1, 2, 3, 5, 6, 7, 8], + }), + } + common_column_info = [ + { + 'table_name': 'main_table', + 'key_column_name': 'primary_key', + 'carryover_column_name': 'parent_1', + }, + { + 'table_name': 'carry_over_1', + 'key_column_name': 'key_column_1', + 'carryover_column_name': 'child_1', + }, + { + 'table_name': 'carry_over_2', + 'key_column_name': 'foreign_key', + 'carryover_column_name': 'child_3', + }, + ] + constraint = CarryOverColumns(common_column_info) + + # Run + valid_rows = constraint._is_valid(data) + + # Assert + expected = { + 'main_table': pd.Series([True] * 6), + 'carry_over_1': pd.Series([True, False, True, True, True, True, True]), + 'carry_over_2': pd.Series([False, True, True, True, True, True, True]), + } + for key in data.keys(): + pd.testing.assert_series_equal(expected[key], valid_rows[key]) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_invalid_data(self, data, metadata, constraint): + """Test ``get_score`` counts the rows of every table involved.""" + # Setup + data['carry_over_1'].loc[1, 'child_1'] = 'b' + + # Run & Assert + assert constraint.get_score(data, metadata) == 0.95 + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata)) diff --git a/tests/unit/multi_table/statistical/constraints/test_denormalized_table.py b/tests/unit/multi_table/statistical/constraints/test_denormalized_table.py new file mode 100644 index 00000000..40b9aa77 --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_denormalized_table.py @@ -0,0 +1,154 @@ +import re + +import numpy as np +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import DenormalizedTable +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def data(): + return { + 'tableA': pd.DataFrame({ + 'id': [1, 1, 2, 2, 3], + 'dob': ['1990-01-01', '1990-01-01', '1985-05-05', '1985-05-05', '1970-02-02'], + 'name': ['Ann', 'Ann', 'Bob', 'Bob', 'Cam'], + 'last_name': ['A', 'A', 'B', 'B', 'C'], + 'amount': [10, 20, 30, 40, 50], + }) + } + + +@pytest.fixture +def constraint(): + return DenormalizedTable( + table_name='tableA', + denormalized_primary_key='id', + denormalized_column_names=['dob', 'name', 'last_name'], + ) + + +class TestDenormalizedTable: + def test___init__invalid_parameters(self): + """Test ``__init__`` validates the parameter types.""" + # Run & Assert 1 + with pytest.raises(ValueError, match="The 'table_name' parameter must be a string."): + DenormalizedTable(table_name=1, denormalized_primary_key='id') + + # Run & Assert 2 + error_message = "The 'denormalized_column_names' parameter must be a list of strings." + with pytest.raises(ValueError, match=error_message): + DenormalizedTable( + table_name='tableA', + denormalized_primary_key='id', + denormalized_column_names='name', + ) + + def test___init__primary_key_in_column_names(self): + """Test ``__init__`` errors if the primary key is also a denormalized column.""" + # Run & Assert + with pytest.raises(ValueError, match='cannot be both'): + DenormalizedTable( + table_name='tableA', + denormalized_primary_key='id', + denormalized_column_names=['id'], + ) + + def test__validate_data_missing_table(self, constraint): + """Test ``_validate_data`` errors if the table is not in the data.""" + # Setup + expected_error = re.escape("The table 'tableA' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data({'OtherTable': pd.DataFrame()}) + + def test__validate_data_missing_columns(self, data, constraint): + """Test ``_validate_data`` errors if a referenced column is not in the table.""" + # Setup + del data['tableA']['name'] + expected_error = re.escape("The column(s) 'name' are missing from the table 'tableA'.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data) + + def test__is_valid(self, data, constraint): + """Test ``_is_valid`` method.""" + # Run + is_valid = constraint._is_valid(data) + + # Assert + expected = pd.Series([True, True, True, True, True]) + pd.testing.assert_series_equal(is_valid['tableA'], expected) + + def test__is_valid_no_denormalized_columns(self, data): + """Test ``_is_valid`` return every row is valid when there is nothing to check.""" + # Setup + instance = DenormalizedTable(table_name='tableA', denormalized_primary_key='id') + + # Run + is_valid = instance._is_valid(data) + + # Assert + pd.testing.assert_series_equal(is_valid['tableA'], pd.Series([True] * 5)) + + def test__is_valid_with_nans(self, constraint): + """Test ``_is_valid`` missing values are treated as equal to each other.""" + # Setup + data = { + 'tableA': pd.DataFrame({ + 'id': [1, 1, None, None], + 'dob': [np.nan, np.nan, '1970-02-02', '1970-02-02'], + 'name': ['Ann', 'Ann', 'Cam', 'Cam'], + 'last_name': ['A', 'A', 'C', 'C'], + }) + } + + # Run + is_valid = constraint._is_valid(data) + + # Assert + pd.testing.assert_series_equal(is_valid['tableA'], pd.Series([True] * 4)) + + def test__is_valid_empty_table(self, data, constraint): + """Test it returns all true when the table is empty.""" + # Setup + empty = data['tableA'].iloc[0:0].copy() + data['tableA'] = empty + + # Run + valid_rows = constraint._is_valid(data) + + # Assert + assert valid_rows['tableA'].empty + + def test__is_valid_inconsistent_other_denorm_column(self, data, constraint): + """Variation in any denormalized column marks all rows for that key invalid.""" + # Setup + data['tableA'].loc[3, 'name'] = 'Cam' + + # Run + valid_rows = constraint._is_valid(data) + + # Assert + expected = pd.Series([True, True, False, False, True]) + pd.testing.assert_series_equal(valid_rows['tableA'], expected) + + def test_get_score(self, data, constraint): + """Test get_score returns the proportion of valid rows.""" + # Setup + data['tableA'].loc[3, 'name'] = 'Cam' + + # Run & Assert + assert constraint.get_score(data) == 0.6 + + def test_get_score_empty_table(self, data, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data['tableA'] = data['tableA'].iloc[:0] + + # Run & Assert + assert pd.isna(constraint.get_score(data)) diff --git a/tests/unit/multi_table/statistical/constraints/test_foreign_to_foreign_key.py b/tests/unit/multi_table/statistical/constraints/test_foreign_to_foreign_key.py new file mode 100644 index 00000000..9cb2d400 --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_foreign_to_foreign_key.py @@ -0,0 +1,270 @@ +import re + +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import ForeignToForeignKey +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def metadata(): + return { + 'tables': { + 'users': { + 'primary_key': 'user_id', + 'columns': { + 'user_id': {'sdtype': 'id'}, + 'product_id': {'sdtype': 'id'}, + 'company_name': {'sdtype': 'company', 'pii': True}, + 'user_name': {'sdtype': 'name', 'pii': True}, + }, + }, + 'transactions': { + 'primary_key': 'transaction_id', + 'columns': { + 'transaction_id': {'sdtype': 'id'}, + 'product_id': {'sdtype': 'id'}, + 'company_name': {'sdtype': 'company', 'pii': True}, + 'amount': {'sdtype': 'unknown', 'pii': True}, + }, + }, + } + } + + +@pytest.fixture +def data(): + users_data = pd.DataFrame({ + 'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + 'product_id': [101, 102, 101, 103, 104, 105, 106, 101, 102, 103], + 'company_name': [ + 'TechCorp', + 'MobileInc', + 'TechCorp', + 'GadgetWorks', + 'ScreenMasters', + 'KeySolutions', + 'MouseMakers', + 'TechCorp', + 'MobileInc', + 'GadgetWorks', + ], + 'user_name': [ + 'Alice', + 'Bob', + 'Charlie', + 'David', + 'Eve', + 'Frank', + 'Grace', + 'Hank', + 'Ivy', + 'Jack', + ], + }) + + transactions_data = pd.DataFrame({ + 'transaction_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010], + 'product_id': [101, 102, 101, 103, 104, 105, 106, 101, 102, 103], + 'company_name': [ + 'TechCorp', + 'MobileInc', + 'TechCorp', + 'GadgetWorks', + 'ScreenMasters', + 'KeySolutions', + 'MouseMakers', + 'TechCorp', + 'MobileInc', + 'GadgetWorks', + ], + 'amount': [50.00, 75.00, 20.00, 60.00, 90.00, 100.00, 150.00, 55.00, 80.00, 30.00], + }) + return {'users': users_data, 'transactions': transactions_data} + + +@pytest.fixture +def columns(): + return [ + {'table_name': 'users', 'foreign_key': 'product_id'}, + {'table_name': 'transactions', 'foreign_key': 'product_id'}, + ] + + +@pytest.fixture +def constraint(columns): + return ForeignToForeignKey(columns=columns) + + +class TestForeignToForeignKey: + def test___init__(self, columns): + """Test the ``__init__`` method sets the parameters.""" + # Run + instance = ForeignToForeignKey(columns=columns) + + # Assert + assert instance.columns == columns + assert instance.foreign_key_generation == 'new' + + def test___init__with_composite_keys(self): + """Test the ``__init__`` method accepts composite foreign keys.""" + # Setup + columns = [ + {'table_name': 'users', 'foreign_key': ('product_id', 'company_name')}, + {'table_name': 'transactions', 'foreign_key': ('product_id', 'company_name')}, + ] + + # Run + instance = ForeignToForeignKey(columns=columns, foreign_key_generation='reuse') + + # Assert + assert instance.columns == columns + assert instance.foreign_key_generation == 'reuse' + + def test___init__invalid_columns(self, columns): + """Test the ``__init__`` method errors if ``columns`` is malformed.""" + # Run and Assert + err_msg = re.escape('columns must be a list of dictionaries') + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=columns[0]) + + err_msg = re.escape('Each entry in columns must be a dictionary') + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], 'transactions']) + + def test___init__invalid_entry_keys(self, columns): + """Test the ``__init__`` method errors if an entry has the wrong keys.""" + # Run and Assert + err_msg = re.escape("Each dictionary must have a 'table_name' key with a string value") + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], {'foreign_key': 'product_id'}]) + + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], {**columns[1], 'table_name': 1}]) + + err_msg = re.escape("Each dictionary must have a 'foreign_key' key") + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], {'table_name': 'transactions'}]) + + err_msg = re.escape("'foreign_key' must be a string or a tuple of strings") + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], {**columns[1], 'foreign_key': ['a']}]) + + def test___init__mismatched_composite_keys(self, columns): + """Test the ``__init__`` method errors if the keys do not have the same size.""" + # Setup + composite_info = {**columns[1], 'foreign_key': ('product_id', 'company_name')} + err_msg = re.escape( + 'All foreign key entries must have the same number of columns. Entry for table ' + "'transactions' has 2 columns, expected 1." + ) + + # Run and Assert + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=[columns[0], composite_info]) + + def test___init__invalid_foreign_key_generation(self, columns): + """Test the ``__init__`` method errors with an unknown foreign key generation.""" + # Run and Assert + err_msg = re.escape('`foreign_key_generation` must be a string.') + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=columns, foreign_key_generation=1) + + err_msg = re.escape( + "Unrecognized `foreign_key_generation` value 'copy'. Must be one of ['new', 'reuse']." + ) + with pytest.raises(ValueError, match=err_msg): + ForeignToForeignKey(columns=columns, foreign_key_generation='copy') + + def test__validate_data_missing_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if one of the tables is not in the data.""" + # Setup + del data['transactions'] + expected_error = re.escape("The table 'transactions' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a foreign key column is not in the table.""" + # Setup + del data['users']['product_id'] + expected_error = re.escape("The column(s) 'product_id' are missing from the table 'users'.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_composite_column(self, data, metadata): + """Test ``_validate_data`` checks every column of a composite foreign key.""" + # Setup + instance = ForeignToForeignKey( + columns=[ + {'table_name': 'users', 'foreign_key': ('product_id', 'company_name')}, + {'table_name': 'transactions', 'foreign_key': ('product_id', 'company_name')}, + ] + ) + del data['transactions']['company_name'] + expected_error = re.escape( + "The column(s) 'company_name' are missing from the table 'transactions'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + instance._validate_data(data, metadata) + + def test__is_valid(self, data, metadata, constraint): + """Test the ``_is_valid`` method returns True for all rows.""" + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + assert set(is_valid) == {'users', 'transactions'} + pd.testing.assert_series_equal(is_valid['users'], pd.Series([True] * 10)) + pd.testing.assert_series_equal(is_valid['transactions'], pd.Series([True] * 10)) + + def test__is_valid_with_unshared_values(self, data, metadata, constraint): + """Test ``_is_valid`` also accepts a value that only one of the tables holds.""" + # Setup + data['users'].loc[2, 'product_id'] = 999 + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['users'], pd.Series([True] * 10)) + pd.testing.assert_series_equal(is_valid['transactions'], pd.Series([True] * 10)) + + def test__is_valid_other_tables(self, data, metadata, constraint): + """Test ``_is_valid`` considers every row of the other tables valid.""" + # Setup + data['other_table'] = pd.DataFrame({'a': ['a', 'b']}) + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['other_table'], pd.Series([True, True])) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_with_unshared_values(self, data, metadata, constraint): + """Test ``get_score`` stays at one when the tables do not share a value.""" + # Setup + data['users'].loc[2, 'product_id'] = 999 + + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata)) diff --git a/tests/unit/multi_table/statistical/constraints/test_foreign_to_primary_key_subset.py b/tests/unit/multi_table/statistical/constraints/test_foreign_to_primary_key_subset.py new file mode 100644 index 00000000..9f94effa --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_foreign_to_primary_key_subset.py @@ -0,0 +1,285 @@ +import re + +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import ForeignToPrimaryKeySubset +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def metadata(): + return { + 'tables': { + 'users': { + 'primary_key': 'user_id', + 'columns': { + 'user_id': {'sdtype': 'id'}, + 'company_name': {'sdtype': 'company', 'pii': True}, + 'user_name': {'sdtype': 'name', 'pii': True}, + }, + }, + 'transactions': { + 'primary_key': 'transaction_id', + 'columns': { + 'transaction_id': {'sdtype': 'id'}, + 'user_id': {'sdtype': 'id'}, + 'amount': {'sdtype': 'unknown', 'pii': True}, + }, + }, + } + } + + +@pytest.fixture +def data(): + users_data = pd.DataFrame({ + 'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + 'company_name': [ + 'TechCorp', + 'MobileInc', + 'TechCorp', + 'GadgetWorks', + 'ScreenMasters', + 'KeySolutions', + 'MouseMakers', + 'TechCorp', + 'MobileInc', + 'GadgetWorks', + ], + 'user_name': [ + 'Alice', + 'Bob', + 'Charlie', + 'David', + 'Eve', + 'Frank', + 'Grace', + 'Hank', + 'Ivy', + 'Jack', + ], + }) + + transactions_data = pd.DataFrame({ + 'transaction_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010], + 'user_id': [1, 3, 8, 1, 3, 8, 1, 3, 8, 1], + 'amount': [50.00, 75.00, 20.00, 60.00, 90.00, 100.00, 150.00, 55.00, 80.00, 30.00], + }) + return {'users': users_data, 'transactions': transactions_data} + + +@pytest.fixture +def constraint(): + return ForeignToPrimaryKeySubset( + parent_table_name='users', + child_table_name='transactions', + child_foreign_key='user_id', + conditional_column_name='company_name', + conditional_values=['TechCorp'], + ) + + +class TestForeignToPrimaryKeySubset: + def test___init__(self, constraint): + """Test the ``__init__`` method sets the parameters.""" + # Assert + assert constraint.parent_table_name == 'users' + assert constraint.child_table_name == 'transactions' + assert constraint.child_foreign_key == 'user_id' + assert constraint.conditional_column_name == 'company_name' + assert constraint.conditional_values == ['TechCorp'] + assert constraint._parent_primary_key is None + + def test___init__invalid_parameters(self): + """Test the ``__init__`` method errors with invalid arguments.""" + # Setup + parameters = { + 'parent_table_name': 'users', + 'child_table_name': 'transactions', + 'child_foreign_key': 'user_id', + 'conditional_column_name': 'company_name', + 'conditional_values': ['TechCorp'], + } + + # Run and Assert + for parameter_name in [ + 'parent_table_name', + 'child_table_name', + 'conditional_column_name', + ]: + err_msg = re.escape(f'`{parameter_name}` must be a string.') + with pytest.raises(TypeError, match=err_msg): + ForeignToPrimaryKeySubset(**{**parameters, parameter_name: 1}) + + err_msg = re.escape('`child_foreign_key` must be a string or a list of strings.') + with pytest.raises(TypeError, match=err_msg): + ForeignToPrimaryKeySubset(**{**parameters, 'child_foreign_key': 1}) + + err_msg = re.escape('`conditional_values` must be a list.') + with pytest.raises(TypeError, match=err_msg): + ForeignToPrimaryKeySubset(**{**parameters, 'conditional_values': 'TechCorp'}) + + def test__validate_data_missing_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if one of the tables is not in the data.""" + # Setup + del data['users'] + expected_error = re.escape("The table 'users' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the conditional column is not in the parent.""" + # Setup + del data['users']['company_name'] + expected_error = re.escape( + "The column(s) 'company_name' are missing from the table 'users'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_primary_key(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the parent has no primary key in the metadata.""" + # Setup + del metadata['tables']['users']['primary_key'] + expected_error = re.escape("The table 'users' does not have a primary key in the metadata.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_stores_the_primary_key(self, data, metadata, constraint): + """Test ``_validate_data`` remembers the primary key that it resolved.""" + # Run + constraint._validate_data(data, metadata) + + # Assert + assert constraint._parent_primary_key == 'user_id' + + def test__is_valid(self, data, metadata, constraint): + """Test ``_is_valid`` considers valid every child row that references a subset row.""" + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['transactions'], pd.Series([True] * 10)) + pd.testing.assert_series_equal(is_valid['users'], pd.Series([True] * 10)) + + def test__is_valid_with_invalid_values(self, data, metadata, constraint): + """Test ``_is_valid`` flags a child row that references a row out of the subset.""" + # Setup + data['transactions']['user_id'] = [1, 2, 3, 8, 1, 3, 8, 1, 3, 8] + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + expected = pd.Series([True, False, True, True, True, True, True, True, True, True]) + pd.testing.assert_series_equal(is_valid['transactions'], expected) + + def test__is_valid_with_several_conditional_values(self, data, metadata): + """Test ``_is_valid`` accepts every conditional value that is allowed.""" + # Setup + data['transactions']['user_id'] = [1, 2, 3, 8, 1, 3, 8, 1, 3, 8] + instance = ForeignToPrimaryKeySubset( + parent_table_name='users', + child_table_name='transactions', + child_foreign_key='user_id', + conditional_column_name='company_name', + conditional_values=['TechCorp', 'MobileInc'], + ) + + # Run + is_valid = instance._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['transactions'], pd.Series([True] * 10)) + + def test__is_valid_with_unknown_parent(self, data, metadata, constraint): + """Test ``_is_valid`` accepts an unknown parent while every known one is allowed.""" + # Setup + data['transactions']['user_id'] = [1, 3, 8, 1, 3, 8, 1, 3, 8, 999] + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['transactions'], pd.Series([True] * 10)) + + def test__is_valid_with_unknown_parent_and_invalid_values(self, data, metadata, constraint): + """Test ``_is_valid`` flags an unknown parent once a known one is out of the subset.""" + # Setup + data['transactions']['user_id'] = [1, 2, 8, 1, 3, 8, 1, 3, 8, 999] + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + expected = pd.Series([True, False, True, True, True, True, True, True, True, False]) + pd.testing.assert_series_equal(is_valid['transactions'], expected) + + def test__is_valid_other_tables(self, data, metadata, constraint): + """Test ``_is_valid`` considers every row of the other tables valid.""" + # Setup + data['other_table'] = pd.DataFrame({'a': ['a', 'b']}) + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['other_table'], pd.Series([True, True])) + + def test__is_valid_with_another_data(self): + """Test that all child foreign keys match a primary key that matches the condition.""" + # Setup + data = { + 'parent_table': pd.DataFrame({ + 'parent_pk': [1, 2, 3, 4], + 'conditional_column': ['value_1', 'value_2', 'value_3', 'value_4'], + }), + 'child_table': pd.DataFrame({'child_fk': [1, 2, 1, 2, 3, 4, 5]}), + } + constraint = ForeignToPrimaryKeySubset( + 'parent_table', + 'child_table', + 'child_fk', + 'conditional_column', + ['value_1', 'value_2'], + ) + constraint._parent_primary_key = 'parent_pk' + + # Run + valid_rows = constraint._is_valid(data) + + # Assert + expected = { + 'parent_table': pd.Series([True, True, True, True]), + 'child_table': pd.Series([True, True, True, True, False, False, False]), + } + for table_name, data in expected.items(): + pd.testing.assert_series_equal(data, valid_rows[table_name]) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_invalid_data(self, data, metadata, constraint): + """Test ``get_score`` counts the rows of every table involved.""" + # Setup + data['transactions']['user_id'] = [1, 2, 3, 8, 1, 3, 8, 1, 3, 8] + + # Run & Assert + assert constraint.get_score(data, metadata) == 0.95 + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata)) diff --git a/tests/unit/multi_table/statistical/constraints/test_polymorphic_relationship.py b/tests/unit/multi_table/statistical/constraints/test_polymorphic_relationship.py new file mode 100644 index 00000000..4b70e812 --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_polymorphic_relationship.py @@ -0,0 +1,397 @@ +import re +from unittest.mock import Mock + +import numpy as np +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import PolymorphicRelationship +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def data(): + return { + 'parent1': pd.DataFrame({'primary_key': [1, 2]}), + 'parent2': pd.DataFrame({'primary_key': [10, 20]}), + 'table': pd.DataFrame({ + 'foreign_key': [1, 2, 10, 20, np.nan], + 'type': ['DEBIT', 'DEBIT', 'CREDIT', 'CREDIT', 'DEBIT'], + }), + } + + +@pytest.fixture +def metadata(): + return { + 'tables': { + 'table': { + 'columns': {'foreign_key': {'sdtype': 'id'}, 'type': {'sdtype': 'categorical'}} + }, + 'parent1': {'columns': {'primary_key': {'sdtype': 'id'}}, 'primary_key': 'primary_key'}, + 'parent2': {'columns': {'primary_key': {'sdtype': 'id'}}, 'primary_key': 'primary_key'}, + } + } + + +@pytest.fixture +def constraint(): + return PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1', 'parent2'], + type_column_name='type', + type_value_to_table={'CREDIT': 'parent2', 'DEBIT': 'parent1'}, + ) + + +class TestPolymorphicRelationship: + def test___init__(self, constraint): + """Test the ``__init__`` method sets the parameters.""" + # Assert + assert constraint.table_name == 'table' + assert constraint.foreign_key == 'foreign_key' + assert constraint.parent_tables == ['parent1', 'parent2'] + assert constraint.type_column == 'type' + assert constraint.type_value_to_table == { + 'CREDIT': 'parent2', + 'DEBIT': 'parent1', + } + + def test___init__without_type_column(self): + """Test the ``__init__`` method accepts a constraint without a type column.""" + # Run + instance = PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1'], + ) + + # Assert + assert instance.type_column is None + assert instance.type_value_to_table is None + + def test___init__invalid_parameters(self): + """Test the ``__init__`` method errors with invalid arguments.""" + # Setup + parameters = { + 'table_name': 'table', + 'foreign_key': 'foreign_key', + 'parent_table_names': ['parent1'], + } + + # Run and Assert + err_msg = re.escape('`table_name` must be a string.') + with pytest.raises(TypeError, match=err_msg): + PolymorphicRelationship(**{**parameters, 'table_name': 1}) + + err_msg = re.escape('`foreign_key` must be a string or a list of strings.') + with pytest.raises(TypeError, match=err_msg): + PolymorphicRelationship(**{**parameters, 'foreign_key': 1}) + + err_msg = re.escape('`parent_table_names` must be a list of strings.') + with pytest.raises(TypeError, match=err_msg): + PolymorphicRelationship(**{**parameters, 'parent_table_names': 'parent1'}) + + err_msg = re.escape('`type_column_name` must be a string or None.') + with pytest.raises(TypeError, match=err_msg): + PolymorphicRelationship(**{**parameters, 'type_column_name': 1}) + + def test___init__table_name_in_parent_table_names(self): + """Test the ``__init__`` method errors if the table is one of its own parents.""" + # Setup + err_msg = re.escape("Table name 'parent1' cannot also be in `parent_table_names`.") + + # Run and Assert + with pytest.raises(ConstraintNotApplicableError, match=err_msg): + PolymorphicRelationship( + table_name='parent1', + foreign_key='foreign_key', + parent_table_names=['parent1'], + ) + + def test___init__type_column_is_the_foreign_key(self): + """Test the ``__init__`` method errors if the type column is the foreign key.""" + # Setup + err_msg = re.escape('`foreign_key` and `type_column_name` must be different columns.') + + # Run and Assert + with pytest.raises(ValueError, match=err_msg): + PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1'], + type_column_name='foreign_key', + ) + + def test___init__invalid_type_value_to_table(self): + """Test the ``__init__`` method errors if the type mapping is inconsistent.""" + # Setup + parameters = { + 'table_name': 'table', + 'foreign_key': 'foreign_key', + 'parent_table_names': ['parent1'], + } + + # Run and Assert + err_msg = re.escape('`type_value_to_table` must be a dict or `None`.') + with pytest.raises(TypeError, match=err_msg): + PolymorphicRelationship(**{ + **parameters, + 'type_column_name': 'type', + 'type_value_to_table': 'debit', + }) + + err_msg = re.escape( + "Table(s) 'parent2' in `type_values_to_table` not found in `parent_table_names` list." + ) + with pytest.raises(ValueError, match=err_msg): + PolymorphicRelationship(**{ + **parameters, + 'type_column_name': 'type', + 'type_value_to_table': {'CREDIT': 'parent2'}, + }) + + err_msg = re.escape( + "Table(s) 'parent2' in `parent_table_names` do not have any type value " + 'associated with them in `type_values_to_table`.' + ) + with pytest.raises(ValueError, match=err_msg): + PolymorphicRelationship(**{ + **parameters, + 'parent_table_names': ['parent1', 'parent2'], + 'type_column_name': 'type', + 'type_value_to_table': {'DEBIT': 'parent1'}, + }) + + def test__validate_data_missing_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a parent table is not in the data.""" + # Setup + del data['parent1'] + expected_error = re.escape("The table 'parent1' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the type column is not in the table.""" + # Setup + del data['table']['type'] + expected_error = re.escape("The column(s) 'type' are missing from the table 'table'.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_primary_key_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the parent does not hold its primary key.""" + # Setup + data['parent1']['id'] = data['parent1']['primary_key'] + del data['parent1']['primary_key'] + expected_error = re.escape( + "The column(s) 'primary_key' are missing from the table 'parent1'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_primary_key(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a parent has no primary key in the metadata.""" + # Setup + del metadata['tables']['parent1']['primary_key'] + expected_error = re.escape( + "The table 'parent1' does not have a primary key in the metadata." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__get_foreign_key_groups(self, metadata): + """Test helper to group foreign keys by parent with a type column.""" + # Setup + instance = PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1', 'parent2'], + ) + instance_type_col = PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1', 'parent2'], + type_column_name='type', + ) + data = { + 'table': pd.DataFrame({ + 'foreign_key': ['id0', 'id0', 'id0', 'id1', 'id2', 0, 0, 0, 2, 2], + 'type': ['parent1'] * 5 + ['parent2'] * 5, + }), + 'parent1': pd.DataFrame({'primary_key': ['id0', 'id1', 'id2']}), + 'parent2': pd.DataFrame({'primary_key': [0, 1, 2]}), + } + + # Run + child_groups = instance._get_foreign_key_groups(data, metadata) + child_groups_type_col = instance_type_col._get_foreign_key_groups(data, metadata) + + # Assert + for result in (child_groups, child_groups_type_col): + assert set(result.keys()) == {'parent1', 'parent2'} + expected_parent1_group = pd.DataFrame({ + 'foreign_key': ['id0', 'id0', 'id0', 'id1', 'id2'] + }) + pd.testing.assert_frame_equal( + result['parent1'], expected_parent1_group, check_names=False + ) + expected_parent2_group = pd.DataFrame( + {'foreign_key': [0, 0, 0, 2, 2]}, index=[5, 6, 7, 8, 9], dtype='object' + ) + + pd.testing.assert_frame_equal( + result['parent2'], expected_parent2_group, check_names=False + ) + + def test__validate_type_column(self): + """Test type column validation.""" + instance = PolymorphicRelationship( + table_name='table', + foreign_key='fk', + parent_table_names=['parent1', 'parent2'], + type_column_name='type', + ) + table_data = pd.DataFrame({ + 'fk': ['id0', 'id0', 'id0', 'id1', 'id2', 0, 0, 0, 2, 2], + 'type': ['parent1'] * 5 + ['parent2'] * 4 + ['unknown'], + }) + valid_type = ['parent1'] * 5 + ['parent2'] * 5 + + # Run and Assert + unknown_is_valid = instance._validate_type_column(table_data) + + instance.type_value_to_parent = {'parent1': 'parent1', 'parent2': 'parent2'} + extra_is_valid = instance._validate_type_column(table_data) + + table_data['type'] = valid_type + valid = instance._validate_type_column(table_data) + + # Assert + expected_is_valid = pd.Series([True] * 9 + [False], name='type') + pd.testing.assert_series_equal(extra_is_valid, expected_is_valid) + pd.testing.assert_series_equal(unknown_is_valid, expected_is_valid) + pd.testing.assert_series_equal(valid, pd.Series([True] * 10, name='type')) + + def test__validate_parent_primary_keys(self, metadata): + """Test validating that primary key values do not overlap.""" + instance = PolymorphicRelationship( + table_name='table', + foreign_key='fk', + parent_table_names=['parent1', 'parent2'], + ) + data = { + 'table': pd.DataFrame({ + 'fk': ['id0', 'id0', 'id0', 'id1', 'id2', 0, 0, 0, 2, 2], + 'type': ['parent1'] * 5 + ['parent2'] * 4 + ['unknown'], + }), + 'parent1': pd.DataFrame({'primary_key': ['id0', 'id1', 'id2']}), + 'parent2': pd.DataFrame({'primary_key': ['id2', 'id3', 'id4']}), + } + + # Run and Assert + overlap_is_valid_dict = instance._validate_parent_primary_keys(data, metadata) + + # Assert + assert set(overlap_is_valid_dict.keys()) == {'parent1', 'parent2'} + assert all(overlap_is_valid_dict['parent1']) + expected_is_valid_parent2 = pd.Series([False, True, True]) + pd.testing.assert_series_equal( + overlap_is_valid_dict['parent2'], expected_is_valid_parent2, check_names=False + ) + + def test__validate_polymorphic_relationship_with_data(self, metadata): + """Test validation with data without erroring.""" + instance = PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1', 'parent2'], + type_column_name='type', + ) + instance._get_foreign_key_groups = Mock( + return_value={ + 'parent1': pd.DataFrame( + {'foreign_key': ['bad_key', 'id0', 'id1', 'id2']}, index=range(1, 5) + ), + 'parent2': pd.DataFrame( + {'foreign_key': [0, 0, 0, 2, 'unknown']}, index=range(5, 10) + ), + } + ) + data = { + 'table': pd.DataFrame({ + 'foreign_key': ['id0', 'bad_key', 'id0', 'id1', 'id2', 0, 0, 0, 2, 'unknown'], + 'type': ['extra'] + ['parent1'] * 4 + ['parent2'] * 5, + }), + 'parent1': pd.DataFrame({'primary_key': ['id0', 'id1', 'id2']}), + 'parent2': pd.DataFrame({'primary_key': [0, 1, 2]}), + } + + # Run + is_valid_dict = instance._validate_polymorphic_relationship_with_data(data, metadata) + + # Assert + instance._get_foreign_key_groups.assert_called_once_with(data, metadata) + expected_is_valid_table = pd.Series([False, False] + [True] * 7 + [False]) + for table, is_valid in is_valid_dict.items(): + if table == 'table': + pd.testing.assert_series_equal(is_valid, expected_is_valid_table) + else: + assert all(is_valid) + + def test__is_valid(self, metadata): + """Test ``_is_valid`` method.""" + # Setup + instance = PolymorphicRelationship( + table_name='table', + foreign_key='foreign_key', + parent_table_names=['parent1', 'parent2'], + ) + data = { + 'table': pd.DataFrame({'foreign_key': list(range(5)) * 2}), + 'parent1': pd.DataFrame({ + 'primary_key': [0, 1, 2], + }), + 'parent2': pd.DataFrame({ + 'primary_key': [3, 4], + }), + } + + # Run + is_valid = instance._is_valid(data, metadata) + + # Assert + assert set(is_valid.keys()) == {'table', 'parent1', 'parent2'} + pd.testing.assert_series_equal(is_valid['table'], pd.Series([True] * 10)) + pd.testing.assert_series_equal(is_valid['parent1'], pd.Series([True] * 3)) + pd.testing.assert_series_equal(is_valid['parent2'], pd.Series([True] * 2)) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_invalid_data(self, data, metadata, constraint): + """Test ``get_score`` counts the rows of every table involved.""" + # Setup + data['table']['type'] = ['CREDIT', 'DEBIT', 'DEBIT', 'CREDIT', 'DEBIT'] + + # Run & Assert + assert constraint.get_score(data, metadata) == pytest.approx(7 / 9) + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata)) diff --git a/tests/unit/multi_table/statistical/constraints/test_primary_to_primary_key_subset.py b/tests/unit/multi_table/statistical/constraints/test_primary_to_primary_key_subset.py new file mode 100644 index 00000000..f915f335 --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_primary_to_primary_key_subset.py @@ -0,0 +1,336 @@ +import re +from unittest.mock import Mock + +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import PrimaryToPrimaryKeySubset +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def data(): + return { + 'main_table': pd.DataFrame({ + 'primary_key': [1, 2, 3, 4, 5], + 'condition_column': [ + 'conditional_value_1', + 'conditional_value_2', + 'conditional_value_1', + 'conditional_value_2', + 'conditional_value_1', + ], + }), + 'table_1': pd.DataFrame({ + 'col_1': [1, 3, 5], + 'col_2': [7, 8, 10], + 'col_3': ['A', 'A', 'B'], + }), + 'table_2': pd.DataFrame({'col_4': ['2', '4'], 'col_5': [14.5, 16.7], 'col_6': ['D', 'E']}), + } + + +@pytest.fixture +def metadata(): + return { + 'tables': { + 'main_table': { + 'columns': { + 'primary_key': {'sdtype': 'id'}, + 'condition_column': {'sdtype': 'categorical'}, + }, + 'primary_key': 'primary_key', + }, + 'table_1': { + 'columns': { + 'col_1': {'sdtype': 'id'}, + 'col_2': {'sdtype': 'numerical'}, + 'col_3': {'sdtype': 'categorical'}, + }, + 'primary_key': 'col_1', + }, + 'table_2': { + 'columns': { + 'col_4': {'sdtype': 'id'}, + 'col_5': {'sdtype': 'numerical'}, + 'col_6': {'sdtype': 'categorical'}, + }, + 'primary_key': 'col_4', + }, + }, + 'relationships': [{'table_1': ['conditional_value_1'], 'table_2': ['conditional_value_2']}], + } + + +@pytest.fixture +def constraint(): + return PrimaryToPrimaryKeySubset( + main_table_name='main_table', + conditional_column_name='condition_column', + relationships={'table_1': ['conditional_value_1'], 'table_2': ['conditional_value_2']}, + ) + + +class TestPrimaryToPrimaryKeySubset: + def test___init__(self, constraint): + """Test the ``__init__`` method sets the parameters.""" + # Assert + assert constraint.main_table_name == 'main_table' + assert constraint.conditional_column_name == 'condition_column' + assert constraint.relationships == { + 'table_1': ['conditional_value_1'], + 'table_2': ['conditional_value_2'], + } + + def test___init__invalid_parameters(self): + """Test the ``__init__`` method errors with invalid arguments.""" + # Setup + parameters = { + 'main_table_name': 'main_table', + 'conditional_column_name': 'condition_column', + 'relationships': { + 'table_1': ['conditional_value_1'], + 'table_2': ['conditional_value_2'], + }, + } + err_msg = re.escape('`main_table_name` and `conditional_column_name` must be strings') + + # Run and Assert + with pytest.raises(ValueError, match=err_msg): + PrimaryToPrimaryKeySubset(**{**parameters, 'main_table_name': 1}) + + with pytest.raises(ValueError, match=err_msg): + PrimaryToPrimaryKeySubset(**{**parameters, 'conditional_column_name': 1}) + + def test___init__invalid_relationships(self): + """Test the ``__init__`` method errors if a relationship is malformed.""" + # Setup + parameters = { + 'main_table_name': 'main_table', + 'conditional_column_name': 'condition_column', + 'relationships': { + 'table_1': ['conditional_value_1'], + 'table_2': ['conditional_value_2'], + }, + } + err_msg = re.escape( + '`relationships` must be a a dict that maps the name of the connected table to a ' + 'list of values that are acceptable for a connection to be made.' + ) + + # Run and Assert + with pytest.raises(ValueError, match=err_msg): + PrimaryToPrimaryKeySubset(**{**parameters, 'relationships': ['table_1']}) + + with pytest.raises(ValueError, match=err_msg): + PrimaryToPrimaryKeySubset(**{ + **parameters, + 'relationships': {'table_1': 'conditional_value_1'}, + }) + + with pytest.raises(ValueError, match=err_msg): + PrimaryToPrimaryKeySubset(**{**parameters, 'relationships': {1: ['a']}}) + + def test__validate_data_missing_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a connected table is not in the data.""" + # Setup + del data['table_2'] + expected_error = re.escape("The table 'table_2' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_main_table(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the main table is not in the data.""" + # Setup + del data['main_table'] + expected_error = re.escape("The table 'main_table' is missing from the data.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if the conditional column is not in the main table.""" + # Setup + del data['main_table']['condition_column'] + expected_error = re.escape( + "The column(s) 'condition_column' are missing from the table 'main_table'." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_connected_primary_key_column(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a connected table does not hold its primary key.""" + # Setup + del data['table_1']['col_1'] + expected_error = re.escape("The column(s) 'col_1' are missing from the table 'table_1'.") + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__validate_data_missing_primary_key(self, data, metadata, constraint): + """Test ``_validate_data`` errors if a table has no primary key in the metadata.""" + # Setup + del metadata['tables']['table_1']['primary_key'] + expected_error = re.escape( + "The table 'table_1' does not have a primary key in the metadata." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + constraint._validate_data(data, metadata) + + def test__get_metadata_parameters(self, metadata): + """Test the ``_get_metadata_parameters`` method.""" + # Setup + relationships = {'table_1': ['conditional_value_1'], 'table_2': ['conditional_value_2']} + constraint = PrimaryToPrimaryKeySubset('main_table', 'condition_column', relationships) + + # Run + results = constraint._get_metadata_parameters(metadata) + + # Assert + assert results[0] == {'main_table': 'primary_key', 'table_1': 'col_1', 'table_2': 'col_4'} + assert results[1] == ['primary_key', 'condition_column'] + assert results[2] == { + 'table_1': {'col_2': 'table_1_col_2', 'col_3': 'table_1_col_3'}, + 'table_2': {'col_5': 'table_2_col_5', 'col_6': 'table_2_col_6'}, + } + + def test__is_valid(self): + """Test that rows with mismatched keys are marked invalid.""" + # Setup + instance = PrimaryToPrimaryKeySubset( + 'main', 'color', relationships={'ref1': ['blue'], 'ref2': ['green', 'red']} + ) + instance._get_metadata_parameters = Mock( + return_value=( + { + 'main': 'pk1', + 'ref1': 'pk2', + 'ref2': 'pk3', + }, + '', + '', + ) + ) + data = { + 'main': pd.DataFrame({ + 'pk1': range(20), + 'color': (['blue'] * 5) + (['green'] * 5) + (['red'] * 5) + (['orange'] * 5), + }), + 'ref1': pd.DataFrame({'pk2': [9, 1, 4, 2, 6, 21]}), + 'ref2': pd.DataFrame({'pk3': [6, 7, 8, 1, 16, 22]}), + } + + # Run + valid_rows = instance._is_valid(data) + + # Assert + expected = { + 'main': pd.Series([True] * 20), + 'ref1': pd.Series([False, True, True, True, False, False]), + 'ref2': pd.Series([True, True, True, False, False, False]), + } + for table, data in expected.items(): + pd.testing.assert_series_equal(data, valid_rows[table]) + + def test__is_valid_default(self, data, metadata): + """Test ``_is_valid`` considers valid every connected row that is allowed.""" + # Setup + instance = PrimaryToPrimaryKeySubset( + main_table_name='main_table', + conditional_column_name='condition_column', + relationships={'table_1': ['conditional_value_1']}, + ) + + # Run + is_valid = instance._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 5)) + pd.testing.assert_series_equal(is_valid['table_1'], pd.Series([True] * 3)) + + def test__is_valid_with_invalid_values(self, data, metadata): + """Test ``_is_valid`` flags a connected row whose main row is not allowed.""" + # Setup + data['table_1']['col_1'] = [1, 2, 5] + instance = PrimaryToPrimaryKeySubset( + main_table_name='main_table', + conditional_column_name='condition_column', + relationships={'table_1': ['conditional_value_1']}, + ) + + # Run + is_valid = instance._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['table_1'], pd.Series([True, False, True])) + + def test__is_valid_with_unknown_key(self, data, metadata): + """Test ``_is_valid`` flags a connected row that the main table does not hold.""" + # Setup + data['table_1']['col_1'] = [1, 3, 99] + instance = PrimaryToPrimaryKeySubset( + main_table_name='main_table', + conditional_column_name='condition_column', + relationships={'table_1': ['conditional_value_1']}, + ) + + # Run + is_valid = instance._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['table_1'], pd.Series([True, True, False])) + + def test__is_valid_several_relationships(self, data, metadata, constraint): + """Test ``_is_valid`` checks every connected table against its own values.""" + # Setup + data['table_2']['col_4'] = [2, 4] + + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['main_table'], pd.Series([True] * 5)) + pd.testing.assert_series_equal(is_valid['table_1'], pd.Series([True] * 3)) + pd.testing.assert_series_equal(is_valid['table_2'], pd.Series([True] * 2)) + + def test__is_valid_with_mismatched_key_types(self, data, metadata, constraint): + """Test ``_is_valid`` flags a key that does not have the type of the main key.""" + # Run + is_valid = constraint._is_valid(data, metadata) + + # Assert + pd.testing.assert_series_equal(is_valid['table_1'], pd.Series([True] * 3)) + pd.testing.assert_series_equal(is_valid['table_2'], pd.Series([False] * 2)) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Setup + data['table_2']['col_4'] = [2, 4] + + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_invalid_data(self, data, metadata, constraint): + """Test ``get_score`` counts the rows of every table involved.""" + # Setup + data['table_2']['col_4'] = [2, 4] + data['table_1']['col_1'] = [1, 2, 5] + + # Run & Assert + assert constraint.get_score(data, metadata) == 0.9 + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata)) diff --git a/tests/unit/multi_table/statistical/constraints/test_reference_table.py b/tests/unit/multi_table/statistical/constraints/test_reference_table.py new file mode 100644 index 00000000..6c033171 --- /dev/null +++ b/tests/unit/multi_table/statistical/constraints/test_reference_table.py @@ -0,0 +1,186 @@ +import re + +import pandas as pd +import pytest + +from sdmetrics.multi_table.statistical.constraints import ReferenceTable +from sdmetrics.multi_table.statistical.constraints.error import ConstraintNotApplicableError + + +@pytest.fixture +def metadata(): + """Metadata for the test. + + It has the following relationships: + - grandparent -> parent + - parent -> child + - grandparent -> child + """ + return { + 'tables': { + 'grandparent': { + 'columns': {'pk': {'sdtype': 'id'}, 'col': {'sdtype': 'categorical'}}, + 'primary_key': 'pk', + }, + 'parent': { + 'columns': { + 'pk': {'sdtype': 'id'}, + 'fk': {'sdtype': 'id'}, + 'col': {'sdtype': 'categorical'}, + }, + 'primary_key': 'pk', + }, + 'child': { + 'columns': { + 'pk': {'sdtype': 'id'}, + 'fk_parent': {'sdtype': 'id'}, + 'fk_grandparent': {'sdtype': 'id'}, + 'col': {'sdtype': 'categorical'}, + }, + 'primary_key': 'pk', + }, + }, + 'relationships': [ + { + 'parent_table_name': 'grandparent', + 'child_table_name': 'parent', + 'parent_primary_key': 'pk', + 'child_foreign_key': 'fk', + }, + { + 'parent_table_name': 'parent', + 'child_table_name': 'child', + 'parent_primary_key': 'pk', + 'child_foreign_key': 'fk_parent', + }, + { + 'parent_table_name': 'grandparent', + 'child_table_name': 'child', + 'parent_primary_key': 'pk', + 'child_foreign_key': 'fk_grandparent', + }, + ], + } + + +@pytest.fixture +def data(): + return { + 'grandparent': pd.DataFrame({'pk': range(5), 'col': ['A', 'B', 'C', 'D', 'E']}), + 'parent': pd.DataFrame({ + 'pk': range(5), + 'fk': [0, 1, 1, 2, 4], + 'col': ['A', 'B', 'C', 'D', 'E'], + }), + 'child': pd.DataFrame({ + 'pk': range(5), + 'fk_parent': [0, 1, 2, 3, 4], + 'fk_grandparent': [0, 1, 1, 2, 4], + 'col': ['X', 'Y', 'Z', 'X', 'Y'], + }), + } + + +@pytest.fixture +def constraint(): + return ReferenceTable(reference_table_names=['grandparent']) + + +class TestReferenceTable: + def test___init__(self, constraint): + """Test the ``__init__`` method sets the parameters.""" + # Assert + assert constraint.reference_table_names == ['grandparent'] + + def test___init___invalid_reference_table_type(self): + """Test the ``__init__`` method when reference_table_names is not a list.""" + # Run and Assert + with pytest.raises(ValueError, match="'reference_table_names' must be a list of strings."): + ReferenceTable('not_a_list') + + def test___init___invalid_reference_table_names(self): + """Test the ``__init__`` method when reference_table_names is not a list of strings.""" + # Run and Assert + with pytest.raises(ValueError, match="'reference_table_names' must be a list of strings."): + ReferenceTable(['string', 10]) + + def test__validate_constraint_with_metadata(self, metadata, constraint): + """Test ``_validate_constraint_with_metadata`` passes for a table with no parent.""" + # Run & Assert + constraint._validate_constraint_with_metadata(metadata) + + def test__validate_constraint_with_metadata_reference_parent(self, metadata): + """Test a reference table may be the child of another reference table.""" + # Setup + instance = ReferenceTable(reference_table_names=['grandparent', 'parent']) + + # Run & Assert + instance._validate_constraint_with_metadata(metadata) + + def test__validate_constraint_with_metadata_every_table(self, metadata): + """Test every table of the dataset may be a reference table.""" + # Setup + instance = ReferenceTable(reference_table_names=['grandparent', 'parent', 'child']) + + # Run & Assert + instance._validate_constraint_with_metadata(metadata) + + def test__validate_constraint_with_metadata_missing_table(self, metadata): + """Test ``_validate_constraint_with_metadata`` errors if a table is not in the metadata.""" + # Setup + instance = ReferenceTable(reference_table_names=['City', 'Country']) + expected_error = re.escape( + "Reference table(s) '['City', 'Country']' missing from metadata." + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + instance._validate_constraint_with_metadata(metadata) + + def test__validate_constraint_with_metadata_non_reference_parent(self, metadata): + """Test ``_validate_constraint_with_metadata`` errors on a non reference parent.""" + # Setup + instance = ReferenceTable(reference_table_names=['parent']) + expected_error = re.escape( + 'Reference tables cannot be children of non-reference tables. The following ' + "child-parent pairs are invalid: '[('parent', 'grandparent')]'" + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + instance._validate_constraint_with_metadata(metadata) + + def test__validate_constraint_with_metadata_several_non_reference_parents(self, metadata): + """Test ``_validate_constraint_with_metadata`` reports every invalid pair.""" + # Setup + instance = ReferenceTable(reference_table_names=['child']) + expected_error = re.escape( + 'Reference tables cannot be children of non-reference tables. The following ' + "child-parent pairs are invalid: '[('child', 'grandparent'), ('child', 'parent')]'" + ) + + # Run & Assert + with pytest.raises(ConstraintNotApplicableError, match=expected_error): + instance._validate_constraint_with_metadata(metadata) + + def test__is_valid(self, data, constraint): + """Test that all rows are valid.""" + # Run + valid_rows = constraint._is_valid(data) + + # Assert + for column in valid_rows.values(): + assert all(column) + + def test_get_score(self, data, metadata, constraint): + """Test ``get_score`` returns the proportion of valid rows.""" + # Run & Assert + assert constraint.get_score(data, metadata) == 1.0 + + def test_get_score_empty_tables(self, data, metadata, constraint): + """Test ``get_score`` returns NaN when there are no rows to check.""" + # Setup + data = {table: table_data.iloc[:0] for table, table_data in data.items()} + + # Run & Assert + assert pd.isna(constraint.get_score(data, metadata))