HMA sampling raises a pandas FutureWarning: Downcasting object dtype arrays

sdv/multi_table/hma.py

@@ -635,7 +635,7 @@ def _find_parent_id(likelihoods, num_rows):
             # Some rows got singular matrix error and the rest were 0
             # Fallback to num_rows on the singular matrix rows and
             # keep 0s on the rest.
-            likelihoods = likelihoods.fillna(num_rows)
+            likelihoods = likelihoods.astype(float).fillna(num_rows)
         else:
             # at least one row got a valid likelihood, so fill the
             # rows that got a singular matrix error with the mean

tests/integration/multi_table/test_hma.py

@@ -388,6 +388,7 @@ def test_save_and_load(self, tmp_path):
         assert loaded_synthesizer.get_info() == synthesizer.get_info()
         assert loaded_synthesizer.metadata.to_dict() == metadata.to_dict()
 
+    @pytest.mark.filterwarnings('error::FutureWarning')
     def test_hma_primary_key_and_foreign_key_only(self):
         """Test that ``HMASynthesizer`` can handle tables with primary and foreign keys only."""
         # Setup

tests/unit/multi_table/test_hma.py

@@ -1038,6 +1038,7 @@ def test__estimate_num_columns_to_be_modeled_different_distributions(self):
 
             assert num_table_cols == estimated_num_columns[table_name]
 
+    @pytest.mark.filterwarnings('error::FutureWarning')
     def test__estimate_num_columns_to_be_modeled(self):
         """Test the estimated number of columns is exactly the number of columns to be modeled.