Remove unnecessary OneHotEncoder when there are nan values

rdt/transformers/categorical.py

@@ -634,7 +634,7 @@ def _transform(self, data):
         """
         data = self._prepare_data(data)
         unique_data = {np.nan if pd.isna(x) else x for x in pd.unique(data)}
-        unseen_categories = unique_data - set(self.dummies)
+        unseen_categories = unique_data - {np.nan if pd.isna(x) else x for x in self.dummies}
         if unseen_categories:
             # Select only the first 5 unseen categories to avoid flooding the console.
             examples_unseen_categories = set(list(unseen_categories)[:5])

tests/integration/transformers/test_categorical.py

@@ -1,4 +1,5 @@
 import pickle
+import warnings
 from io import BytesIO
 
 import numpy as np
@@ -424,6 +425,26 @@ def test_one_hot_numerical_nans():
     pd.testing.assert_frame_equal(reverse, data)
 
 
+def test_one_hot_doesnt_warn(tmp_path):
+    """Ensure OneHotEncoder doesn't warn when saving and loading GH#616."""
+    # Setup
+    data = pd.DataFrame({'column_name': [1.0, 2.0, np.nan, 2.0, 3.0, np.nan, 3.0]})
+    ohe = OneHotEncoder()
+
+    # Run
+    ohe.fit(data, 'column_name')
+    tmp = tmp_path / 'ohe.pkl'
+    with open(tmp.name, 'wb') as f:
+        pickle.dump(ohe, f)
+    with open(tmp.name, 'rb') as f:
+        ohe_loaded = pickle.load(f)
+
+    # Assert
+    with warnings.catch_warnings():
+        warnings.simplefilter('error')
+        ohe_loaded.transform(data)
+
+
 def test_label_numerical_2d_array():
     """Ensure LabelEncoder works on numerical + nan only columns."""
 
@@ -432,6 +453,7 @@ def test_label_numerical_2d_array():
 
     transformer = LabelEncoder()
     transformer.fit(data, column)
+
     transformed = pd.DataFrame([0., 1., 2., 3.], columns=['column_name'])
     reverse = transformer.reverse_transform(transformed)