Fix categorical data check. (#11172)

dmlc · Jan 22, 2025 · eb980fc · eb980fc
1 parent 0d7821f
commit eb980fc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/python-package/xgboost/data.py b/python-package/xgboost/data.py
@@ -912,12 +912,13 @@ def _arrow_feature_info(data: DataType) -> Tuple[List[str], List]:
 
     def map_type(name: str) -> str:
         col = table.column(name)
-        if isinstance(col, pa.DictionaryType):
+        if isinstance(col.type, pa.DictionaryType):
             raise NotImplementedError(
                 "Categorical feature is not yet supported with the current input data "
                 "type."
             )
             return CAT_T  # pylint: disable=unreachable
+
         return _arrow_dtype()[col.type]
 
     types = list(map(map_type, names))

diff --git a/tests/python/test_with_polars.py b/tests/python/test_with_polars.py
@@ -133,3 +133,13 @@ def test_regressor() -> None:
     predt1 = reg1.predict(X)
 
     np.testing.assert_allclose(predt0, predt1)
+
+def test_categorical() ->  None:
+    import polars as pl
+
+    df = pl.DataFrame(
+        {"f0": [1, 2, 3], "b": ["a", "b", "c"]},
+        schema=[("a", pl.Int64()), ("b", pl.Categorical())]
+    )
+    with pytest.raises(NotImplementedError, match="Categorical feature"):
+        xgb.DMatrix(df, enable_categorical=True)