MAINT Fix pandas nightly (#728)

* replace set np.NaN by pd.NA * Update skrub/_fuzzy_join.py --------- Co-authored-by: Gael Varoquaux <[email protected]>
skrub-data · Sep 1, 2023 · dcf6061 · dcf6061
1 parent 1cd57e3
commit dcf6061
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/skrub/_fuzzy_join.py b/skrub/_fuzzy_join.py
@@ -506,9 +506,11 @@ def fuzzy_join(
     if drop_unmatched:
         df_joined.drop(columns=["fj_idx"], inplace=True)
     else:
-        idx = df_joined.index[df_joined["fj_nan"] == 1]
-        if len(idx) != 0:
-            df_joined.iloc[idx, df_joined.columns.get_loc("fj_idx") :] = np.NaN
+        mask_na = df_joined["fj_nan"] == 1
+        if mask_na.any():
+            right_cols = df_joined.columns[df_joined.columns.get_loc("fj_idx") :]
+            df_joined[right_cols] = pd.DataFrame.convert_dtypes(df_joined[right_cols])
+            df_joined.loc[mask_na, right_cols] = pd.NA
         df_joined.drop(columns=["fj_idx", "fj_nan"], inplace=True)
 
     if return_score: