Skip to content

Fix get_random_subset crash with composite keys #2982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cpruijsen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix get_random_subset crash with composite keys #2982

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 33 additions & 4 deletions sdv/multi_table/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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

Expand All @@ -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.
Expand All @@ -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
Expand Down
69 changes: 69 additions & 0 deletions tests/integration/utils/test_poc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
80 changes: 80 additions & 0 deletions tests/unit/multi_table/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down