Add support for using data with `pyarrow`-backed `pandas` extension dtypes
- Dominant language
- C++
- Stars
- 28.8k
- Forks
- 8.9k
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 54
Description
In `pandas=1.5`, `pandas` added support for using `pyarrow`-backed extension data dtypes. Using these data types (in particular `string[pyarrow]`) can lead to large performance improvements in terms of memory usage and computation wall time.
I went to use these new dtypes with `xgboost` and got a (very informative) error about them not being supported. Here's a minimal reproducer:
```python
import pandas as pd
import xgboost as xgb
df = pd.DataFrame({"name": ["alice", "bob", "rick"], "x": range(3), "y": [1.3, 7.2, 0.6]})
df = df.astype(
{
"name": "string[pyarrow]",
"x": "int64[pyarrow]",
"y": "float64[pyarrow]",
}
)
X = df.drop(columns=["name"])
y = df.loc[:, "name"]
dtrain = xgb.DMatrix(X, y)
output = xgb.train(
{"verbosity": 2, "tree_method": "hist", "objective": "reg:squarederror"},
dtrain,
num_boost_round=4,
evals=[(dtrain, "train")],
)
print(f"{output = }")
```
which outputs
```
Traceback (most recent call last):
File "/Users/james/projects/dask/dask/xgboost-pyarrow.py", line 15, in
dtrain = xgb.DMatrix(X, y)
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/core.py", line 620, in inner_f
return func(**kwargs)
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/core.py", line 743, in __init__
handle, feature_names, feature_types = dispatch_data_backend(
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/data.py", line 957, in dispatch_data_backend
return _from_pandas_df(data, enable_categorical, missing, threads,
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/data.py", line 404, in _from_pandas_df
data, feature_names, feature_types = _transform_pandas_df(
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/data.py", line 378, in _transform_pandas_df
_invalid_dataframe_dtype(data)
File "/Users/james/mambaforge/envs/dask-py39/lib/python3.9/site-packages/xgboost/data.py", line 270, in _invalid_dataframe_dtype
raise ValueError(msg)
ValueError: DataFrame.dtypes for data must be int, float, bool or category. When categorical type is supplied, The experimental DMatrix parameter`enable_categorical` must be set to `True`. Invalid columns:x: int64[pyarrow], y: double[pyarrow]
```
It looks like support for `pandas` nullable extension dtypes (e.g. `Int64`, `Float64`, etc.) has already been added to `xgboost` (xref https://github.com/dmlc/xgboost/pull/7760, https://github.com/dmlc/xgboost/pull/8480) and it would be great if `pyarrow`-backed extension dtypes were also supported.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.