TST Fix openml parser implementation for pandas-dev (scikit-learn#26386)

thomasjpfan · REDVM · commit 6d726f305b5b · 2023-11-16T16:26:10.000+01:00
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -221,6 +221,9 @@ Changelog
   is deprecated and will be removed in v1.5.
   :pr:`25784` by :user:`Jérémie du Boisberranger`.
 
+- |Fix| :func:`datasets.fetch_openml` returns improved data types when
+  `as_frame=True` and `parser="liac-arff"`. :pr:`26386` by `Thomas Fan`_.
+
 :mod:`sklearn.decomposition`
 ............................
 
diff --git a/sklearn/datasets/_arff_parser.py b/sklearn/datasets/_arff_parser.py
@@ -199,6 +199,11 @@ def _io_to_generator(gzip_file):
             dfs.append(
                 pd.DataFrame(data, columns=columns_names, copy=False)[columns_to_keep]
             )
+        # dfs[0] contains only one row, which may not have enough data to infer to
+        # column's dtype. Here we use `dfs[1]` to configure the dtype in dfs[0]
+        if len(dfs) >= 2:
+            dfs[0] = dfs[0].astype(dfs[1].dtypes)
+
         frame = pd.concat(dfs, ignore_index=True)
         del dfs, first_df
 
diff --git a/sklearn/datasets/tests/test_openml.py b/sklearn/datasets/tests/test_openml.py
@@ -925,7 +925,7 @@ def datasets_missing_values():
         # with casting it will be transformed to either float or Int64
         (40966, "pandas", 1, 77, 0),
         # titanic
-        (40945, "liac-arff", 3, 5, 0),
+        (40945, "liac-arff", 3, 6, 0),
         (40945, "pandas", 3, 3, 3),
     ],
 )

Original file line number	Diff line number	Diff line change
`@@ -925,7 +925,7 @@ def datasets_missing_values():`
`925`	`925`	`# with casting it will be transformed to either float or Int64`
`926`	`926`	`(40966, "pandas", 1, 77, 0),`
`927`	`927`	`# titanic`
`928`		`- (40945, "liac-arff", 3, 5, 0),`
	`928`	`+ (40945, "liac-arff", 3, 6, 0),`
`929`	`929`	`(40945, "pandas", 3, 3, 3),`
`930`	`930`	`],`
`931`	`931`	`)`