diff --git a/sdv/multi_table/utils.py b/sdv/multi_table/utils.py index 36069d83a..7876b206d 100644 --- a/sdv/multi_table/utils.py +++ b/sdv/multi_table/utils.py @@ -499,6 +499,32 @@ def _subsample_table_and_descendants(data, metadata, table, num_rows, drop_missi _drop_rows(data, metadata, drop_missing_values) +def _get_key_values(table, columns): + """Get a Series of key values for the given columns. + + A single-column key is returned as-is. A composite key is returned as a Series of + tuples so it can be compared with ``isin`` and ``unique``. + + Args: + table (pandas.DataFrame): + Table containing the key columns. + columns (str or list[str] or tuple[str]): + Name of the key column, or list/tuple of names for a composite key. + + Returns: + pandas.Series: + Key values aligned to ``table.index``. + """ + columns = _cast_to_iterable(columns) + if len(columns) == 1: + return table[columns[0]] + + return pd.Series( + list(table[list(columns)].itertuples(index=False, name=None)), + index=table.index, + ) + + def _get_primary_keys_referenced(data, metadata): """Get the primary keys referenced by the relationships. @@ -519,7 +545,9 @@ def _get_primary_keys_referenced(data, metadata): parent_table = relationship['parent_table_name'] child_table = relationship['child_table_name'] foreign_key = relationship['child_foreign_key'] - primary_keys_referenced[parent_table].update(set(data[child_table][foreign_key].unique())) + primary_keys_referenced[parent_table].update( + set(_get_key_values(data[child_table], foreign_key).unique()) + ) return primary_keys_referenced @@ -536,7 +564,7 @@ def _subsample_parent( Args: parent_table (pandas.DataFrame): Parent table to subsample. - parent_primary_key (str): + parent_primary_key (str or list[str]): Name of the primary key of the parent table. parent_pk_referenced_before (set): Set of the primary keys referenced before any subsampling. @@ -551,9 +579,10 @@ def _subsample_parent( total_dropped = len(dereferenced_pk_parent) drop_proportion = total_dropped / total_referenced - parent_table = parent_table[~parent_table[parent_primary_key].isin(dereferenced_pk_parent)] + parent_keys = _get_key_values(parent_table, parent_primary_key) + parent_table = parent_table[~parent_keys.isin(dereferenced_pk_parent)] unreferenced_data = parent_table[ - ~parent_table[parent_primary_key].isin(parent_pk_referenced_before) + ~_get_key_values(parent_table, parent_primary_key).isin(parent_pk_referenced_before) ] # Randomly drop a proportional amount of never-referenced rows diff --git a/tests/integration/utils/test_poc.py b/tests/integration/utils/test_poc.py index cf90cabd0..b6d74c790 100644 --- a/tests/integration/utils/test_poc.py +++ b/tests/integration/utils/test_poc.py @@ -295,3 +295,72 @@ def test_get_random_subset_with_missing_values(metadata, data): # Assert assert len(result['child']) == 3 assert result['child']['parent_id'].isna().sum() > 0 + + +def test_get_random_subset_composite_keys(): + """Test ``get_random_subset`` when a relationship uses composite keys.""" + # Setup + parent = pd.DataFrame({ + 'id_1': list(range(10)), + 'id_2': list('ABCDEFGHIJ'), + 'col': range(10), + }) + child = pd.DataFrame({ + 'child_id': list(range(20)), + 'fk_1': [i % 10 for i in range(20)], + 'fk_2': list('ABCDEFGHIJ' * 2), + 'col': range(20), + }) + data = {'parent': parent, 'child': child} + metadata = Metadata.load_from_dict({ + 'tables': { + 'parent': { + 'columns': { + 'id_1': {'sdtype': 'id'}, + 'id_2': {'sdtype': 'id'}, + 'col': {'sdtype': 'numerical'}, + }, + 'primary_key': ['id_1', 'id_2'], + }, + 'child': { + 'columns': { + 'child_id': {'sdtype': 'id'}, + 'fk_1': {'sdtype': 'id'}, + 'fk_2': {'sdtype': 'id'}, + 'col': {'sdtype': 'numerical'}, + }, + 'primary_key': 'child_id', + }, + }, + 'relationships': [ + { + 'parent_table_name': 'parent', + 'parent_primary_key': ['id_1', 'id_2'], + 'child_table_name': 'child', + 'child_foreign_key': ['fk_1', 'fk_2'], + } + ], + }) + + # Run + result_from_child = get_random_subset(data, metadata, 'child', 8, verbose=False) + result_from_parent = get_random_subset(data, metadata, 'parent', 5, verbose=False) + + # Assert + assert len(result_from_child['child']) == 8 + parent_keys = set( + result_from_child['parent'][['id_1', 'id_2']].itertuples(index=False, name=None) + ) + child_keys = set( + result_from_child['child'][['fk_1', 'fk_2']].itertuples(index=False, name=None) + ) + assert child_keys.issubset(parent_keys) + + assert len(result_from_parent['parent']) == 5 + parent_keys = set( + result_from_parent['parent'][['id_1', 'id_2']].itertuples(index=False, name=None) + ) + child_keys = set( + result_from_parent['child'][['fk_1', 'fk_2']].itertuples(index=False, name=None) + ) + assert child_keys.issubset(parent_keys) diff --git a/tests/unit/multi_table/test_utils.py b/tests/unit/multi_table/test_utils.py index 95dceb15b..51282fb0e 100644 --- a/tests/unit/multi_table/test_utils.py +++ b/tests/unit/multi_table/test_utils.py @@ -1518,6 +1518,59 @@ def test__get_primary_keys_referenced(): assert result == expected_result +def test__get_primary_keys_referenced_composite_keys(): + """Test ``_get_primary_keys_referenced`` when a relationship uses composite keys.""" + data = { + 'parent': pd.DataFrame({ + 'id_1': [0, 1, 2, 3], + 'id_2': ['A', 'B', 'A', 'B'], + 'col': [1, 2, 3, 4], + }), + 'child': pd.DataFrame({ + 'fk_1': [0, 0, 1, 2, 1], + 'fk_2': ['A', 'A', 'B', 'A', 'B'], + 'pk_c': [10, 11, 12, 13, 14], + }), + } + metadata = Metadata().load_from_dict({ + 'tables': { + 'parent': { + 'columns': { + 'id_1': {'sdtype': 'id'}, + 'id_2': {'sdtype': 'id'}, + 'col': {'sdtype': 'numerical'}, + }, + 'primary_key': ['id_1', 'id_2'], + }, + 'child': { + 'columns': { + 'fk_1': {'sdtype': 'id'}, + 'fk_2': {'sdtype': 'id'}, + 'pk_c': {'sdtype': 'id'}, + }, + 'primary_key': 'pk_c', + }, + }, + 'relationships': [ + { + 'parent_table_name': 'parent', + 'child_table_name': 'child', + 'parent_primary_key': ['id_1', 'id_2'], + 'child_foreign_key': ['fk_1', 'fk_2'], + } + ], + }) + + # Run + result = _get_primary_keys_referenced(data, metadata) + + # Assert + expected_result = { + 'parent': {(0, 'A'), (1, 'B'), (2, 'A')}, + } + assert result == expected_result + + def test__subsample_parent_all_reeferenced_before(): """Test the ``_subsample_parent`` when all primary key were referenced before. @@ -1590,6 +1643,33 @@ def test__subsample_parent_not_all_referenced_before(): assert set(data['parent']['pk_p']).issubset({1, 2, 3, 4, 6, 7, 8}) +def test__subsample_parent_composite_keys(): + """Test ``_subsample_parent`` when the parent primary key is composite. + + All composite keys were referenced before. Keys ``(2, 'B')`` and ``(3, 'A')`` are no + longer referenced and should be dropped. + """ + # Setup + parent_table = pd.DataFrame({ + 'id_1': [1, 1, 2, 2, 3], + 'id_2': ['A', 'B', 'A', 'B', 'A'], + 'col': [10, 11, 12, 13, 14], + }) + referenced_before = {(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B'), (3, 'A')} + dereferenced = {(2, 'B'), (3, 'A')} + + # Run + result = _subsample_parent(parent_table, ['id_1', 'id_2'], referenced_before, dereferenced) + + # Assert + expected = pd.DataFrame({ + 'id_1': [1, 1, 2], + 'id_2': ['A', 'B', 'A'], + 'col': [10, 11, 12], + }) + pd.testing.assert_frame_equal(result.reset_index(drop=True), expected) + + def test__subsample_ancestors(): """Test the ``_subsample_ancestors`` method.""" # Setup