diff --git a/openml/datasets/functions.py b/openml/datasets/functions.py index 970331cf1..d32b1210b 100644 --- a/openml/datasets/functions.py +++ b/openml/datasets/functions.py @@ -577,13 +577,18 @@ def attributes_arff_from_df(df: pd.DataFrame) -> list[tuple[str, list[str] | str categories = df[column_name].cat.categories categories_dtype = pd.api.types.infer_dtype(categories) if categories_dtype not in ("string", "unicode"): - raise ValueError( - f"The column '{column_name}' of the dataframe is of " - "'category' dtype. Therefore, all values in " - "this columns should be string. Please " - "convert the entries which are not string. " - f"Got {categories_dtype} dtype in this column.", - ) + try: + # [ENH] Attempt to convert categories to string automatically + new_categories = categories.astype(str) + df[column_name] = df[column_name].cat.rename_categories(new_categories) + categories = df[column_name].cat.categories + except Exception: + raise ValueError( + f"The column '{column_name}' of the dataframe is of " + "'category' dtype. Therefore, all values in " + "this columns should be string. Automatic conversion failed. " + f"Got {categories_dtype} dtype in this column.", + ) attributes_arff.append((column_name, categories.tolist())) elif column_dtype == "boolean": # boolean are encoded as categorical. diff --git a/tests/test_datasets/test_dataset_functions.py b/tests/test_datasets/test_dataset_functions.py index fae4a1c4e..ed74d77ca 100644 --- a/tests/test_datasets/test_dataset_functions.py +++ b/tests/test_datasets/test_dataset_functions.py @@ -663,7 +663,7 @@ def test_attributes_arff_from_df_numeric_column(self): ] def test_attributes_arff_from_df_mixed_dtype_categories(self): - # liac-arff imposed categorical attributes to be of sting dtype. We + # liac-arff imposed categorical attributes to be of string dtype. We # raise an error if this is not the case. df = pd.DataFrame([[1], ["2"], [3.0]]) df[0] = df[0].astype("category") @@ -671,6 +671,22 @@ def test_attributes_arff_from_df_mixed_dtype_categories(self): with pytest.raises(ValueError, match=err_msg): attributes_arff_from_df(df) + def test_attributes_arff_from_df_integer_categories(self): + df = pd.DataFrame({"category": [1, 2, 1]}) + df["category"] = df["category"].astype("category") + + attributes = attributes_arff_from_df(df) + assert attributes == [("category", ["1", "2"])] + + def test_attributes_arff_from_df_missing_values(self): + df = pd.DataFrame({"category": [1, None, 2]}) + df["category"] = df["category"].astype("category") + + attributes = attributes_arff_from_df(df) + + assert attributes == [("category", ["1", "2"])] + assert df["category"].isna().sum() == 1 + def test_attributes_arff_from_df_unknown_dtype(self): # check that an error is raised when the dtype is not supptagorted by # liac-arff