convert pd.NA to np.nan

Author: tvdboom

category_encoders/ordinal.py

@@ -193,9 +193,9 @@ def ordinal_encoding(X_in, mapping=None, cols=None, handle_unknown='value', hand
                 column = switch.get('col')
                 col_mapping = switch['mapping']
 
-                # Treat None as np.nan
-                X[column] = pd.Series([el if el is not None else np.NaN for el in X[column]], index=X[column].index)
-                X[column] = X[column].map(col_mapping)
+                # Convert to object to accept np.nan (dtype string doesn't)
+                # fillna changes None and pd.NA to np.nan
+                X[column] = X[column].astype("object").fillna(np.nan).map(col_mapping)
                 if util.is_category(X[column].dtype):
                     nan_identity = col_mapping.loc[col_mapping.index.isna()].values[0]
                     X[column] = X[column].cat.add_categories(nan_identity)

tests/test_encoders.py

@@ -199,33 +199,37 @@ def test_handle_unknown_return_nan(self):
                     self.assertTrue(result[1:].isnull().all())
 
     def test_handle_missing_return_nan_train(self):
-        X = pd.DataFrame({'city': ['chicago', 'los angeles', np.NaN]})
+        X_np = pd.DataFrame({'city': ['chicago', 'los angeles', np.NaN]})
+        X_pd = pd.DataFrame({'city': ['chicago', 'los angeles', pd.NA]}, dtype="string")
         y = pd.Series([1, 0, 1])
 
         for encoder_name in (set(encoders.__all__) - {'HashingEncoder'}):  # HashingEncoder supports new values by design -> excluded
-            with self.subTest(encoder_name=encoder_name):
-                enc = getattr(encoders, encoder_name)(handle_missing='return_nan')
-                result = enc.fit_transform(X, y).iloc[2, :]
+            for X in (X_np, X_pd):
+                with self.subTest(encoder_name=encoder_name):
+                    enc = getattr(encoders, encoder_name)(handle_missing='return_nan')
+                    result = enc.fit_transform(X, y).iloc[2, :]
 
-                if len(result) == 1:
-                    self.assertTrue(result.isnull().all())
-                else:
-                    self.assertTrue(result[1:].isnull().all())
+                    if len(result) == 1:
+                        self.assertTrue(result.isnull().all())
+                    else:
+                        self.assertTrue(result[1:].isnull().all())
 
     def test_handle_missing_return_nan_test(self):
         X = pd.DataFrame({'city': ['chicago', 'los angeles', 'chicago']})
-        X_t = pd.DataFrame({'city': ['chicago', 'los angeles', np.NaN]})
+        X_np = pd.DataFrame({'city': ['chicago', 'los angeles', np.NaN]})
+        X_pd = pd.DataFrame({'city': ['chicago', 'los angeles', pd.NA]}, dtype="string")
         y = pd.Series([1, 0, 1])
 
         for encoder_name in (set(encoders.__all__) - {'HashingEncoder'}):  # HashingEncoder supports new values by design -> excluded
-            with self.subTest(encoder_name=encoder_name):
-                enc = getattr(encoders, encoder_name)(handle_missing='return_nan')
-                result = enc.fit(X, y).transform(X_t).iloc[2, :]
-
-                if len(result) == 1:
-                    self.assertTrue(result.isnull().all())
-                else:
-                    self.assertTrue(result[1:].isnull().all())
+            for X_na in (X_np, X_pd):
+                with self.subTest(encoder_name=encoder_name):
+                    enc = getattr(encoders, encoder_name)(handle_missing='return_nan')
+                    result = enc.fit(X, y).transform(X_na).iloc[2, :]
+
+                    if len(result) == 1:
+                        self.assertTrue(result.isnull().all())
+                    else:
+                        self.assertTrue(result[1:].isnull().all())
 
     def test_handle_unknown_value(self):
         train = pd.DataFrame({'city': ['chicago', 'los angeles']})