scikit-learn · glemaitre · Feb 1, 2022 · Feb 1, 2022 · Feb 1, 2022 · Feb 1, 2022
diff --git a/doc/whats_new/v1.1.rst b/doc/whats_new/v1.1.rst
@@ -145,6 +145,10 @@ Changelog
   3 times in case of a network failure with a delay between each try.
   :pr:`21901` by :user:`Rileran <rileran>`.
 
+- |Fix| :func:`datasets.fetch_openml` now downcasts whenever possible values
+  that are specified as nominal features instead of storing them as strings.
+  :pr:`22354` by :user:`Guillaume Lemaitre <glemaitre>`.
+
 :mod:`sklearn.decomposition`
 ............................
 

diff --git a/sklearn/datasets/tests/test_openml.py b/sklearn/datasets/tests/test_openml.py
@@ -32,7 +32,7 @@
 from sklearn.utils._testing import assert_allclose, assert_array_equal
 from urllib.error import HTTPError
 from sklearn.datasets.tests.test_common import check_return_X_y
-from sklearn.externals._arff import ArffContainerType
+from sklearn.externals._arff import ArffContainerType, _downcast
 from functools import partial
 from sklearn.utils._testing import fails_if_pypy
 
@@ -90,12 +90,16 @@ def decode_column(data_bunch, col_idx):
 
     data_downloaded = np.array(list(data_arff["data"]), dtype="O")
 
-    for i in range(len(data_bunch.feature_names)):
+    for column_idx, column_name in enumerate(data_bunch.feature_names):
         # XXX: Test per column, as this makes it easier to avoid problems with
         # missing values
 
+        column_downloaded = data_downloaded[:, column_idx]
+        if column_name in data_bunch.categories:
+            column_downloaded = np.array([_downcast(v) for v in column_downloaded])
+
         np.testing.assert_array_equal(
-            data_downloaded[:, i], decode_column(data_bunch, i)
+            column_downloaded, decode_column(data_bunch, column_idx)
         )
 
 
@@ -821,7 +825,7 @@ def test_fetch_openml_titanic_pandas(monkeypatch):
         "boat": object,
         "body": np.float64,
         "home.dest": object,
-        "survived": CategoricalDtype(["0", "1"]),
+        "survived": CategoricalDtype([0, 1]),
     }
 
     frame_columns = [

diff --git a/sklearn/externals/_arff.py b/sklearn/externals/_arff.py
@@ -280,21 +280,34 @@ def _unquote(v):
         return v
 
 
+def _downcast(value):
+    """Downcast a value to integral or float type if possible."""
+    if value is None:
+        return value
+    try:
+        return int(value)
+    except ValueError:
+        try:
+            return float(value)
+        except ValueError:
+            return value
+
+
 def _parse_values(s):
     '''(INTERNAL) Split a line into a list of values'''
     if not _RE_NONTRIVIAL_DATA.search(s):
         # Fast path for trivial cases (unfortunately we have to handle missing
         # values because of the empty string case :(.)
-        return [None if s in ('?', '') else s
+        return [None if s in ('?', '') else _downcast(s)
                 for s in next(csv.reader([s]))]
 
     # _RE_DENSE_VALUES tokenizes despite quoting, whitespace, etc.
     values, errors = zip(*_RE_DENSE_VALUES.findall(',' + s))
     if not any(errors):
-        return [_unquote(v) for v in values]
+        return [_downcast(_unquote(v)) for v in values]
     if _RE_SPARSE_LINE.match(s):
         try:
-            return {int(k): _unquote(v)
+            return {int(k): _downcast(_unquote(v))
                     for k, v in _RE_SPARSE_KEY_VALUES.findall(s)}
         except ValueError:
             # an ARFF syntax error in sparse data