fix dataset parsing for categories

Author: mfeurer

openml/datasets/dataset.py

@@ -408,8 +408,8 @@ def _unpack_categories(series, categories):
                 col.append(categories[int(x)])
             except (TypeError, ValueError):
                 col.append(np.nan)
-        return pd.Series(col, index=series.index, dtype='category',
-                         name=series.name)
+        raw_cat = pd.Categorical(col, ordered=True, categories=categories)
+        return pd.Series(raw_cat, index=series.index, name=series.name)
 
     def _download_data(self) -> None:
         """ Download ARFF data file to standard cache directory. Set `self.data_file`. """

tests/test_datasets/test_dataset.py

@@ -192,6 +192,18 @@ def test_dataset_format_constructor(self):
                 format='arff'
             )
 
+    def test_get_data_with_nonexisting_class(self):
+        # This class is using the anneal dataset with labels [1, 2, 3, 4, 5, 'U']. However,
+        # label 4 does not exist and we test that the features 5 and 'U' are correctly mapped to
+        # indices 4 and 5, and that nothing is mapped to index 3.
+        _, y = self.dataset.get_data('class', dataset_format='dataframe')
+        self.assertEqual(list(y.dtype.categories), ['1', '2', '3', '4', '5', 'U'])
+        _, y = self.dataset.get_data('class', dataset_format='array')
+        self.assertEqual(np.min(y), 0)
+        self.assertEqual(np.max(y), 5)
+        # Check that the
+        self.assertEqual(np.sum(y == 3), 0)
+
 
 class OpenMLDatasetTestOnTestServer(TestBase):
     def setUp(self):