Skip to content

Commit f90c8ac

Browse files
committed
Fix pandas FutureWarning when likelihoods is entirely nan
1 parent 46989cb commit f90c8ac

File tree

3 files changed

+3
-1
lines changed

3 files changed

+3
-1
lines changed

sdv/multi_table/hma.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -635,7 +635,7 @@ def _find_parent_id(likelihoods, num_rows):
635635
# Some rows got singular matrix error and the rest were 0
636636
# Fallback to num_rows on the singular matrix rows and
637637
# keep 0s on the rest.
638-
likelihoods = likelihoods.fillna(num_rows)
638+
likelihoods = likelihoods.astype(float).fillna(num_rows)
639639
else:
640640
# at least one row got a valid likelihood, so fill the
641641
# rows that got a singular matrix error with the mean

tests/integration/multi_table/test_hma.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,7 @@ def test_save_and_load(self, tmp_path):
388388
assert loaded_synthesizer.get_info() == synthesizer.get_info()
389389
assert loaded_synthesizer.metadata.to_dict() == metadata.to_dict()
390390

391+
@pytest.mark.filterwarnings('error::FutureWarning')
391392
def test_hma_primary_key_and_foreign_key_only(self):
392393
"""Test that ``HMASynthesizer`` can handle tables with primary and foreign keys only."""
393394
# Setup

tests/unit/multi_table/test_hma.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1038,6 +1038,7 @@ def test__estimate_num_columns_to_be_modeled_different_distributions(self):
10381038

10391039
assert num_table_cols == estimated_num_columns[table_name]
10401040

1041+
@pytest.mark.filterwarnings('error::FutureWarning')
10411042
def test__estimate_num_columns_to_be_modeled(self):
10421043
"""Test the estimated number of columns is exactly the number of columns to be modeled.
10431044

0 commit comments

Comments
 (0)