Add support for using data with `pyarrow`-backed `pandas` extension dtypes

[jrbourbeau]

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:

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 <module>
    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 #7760, #8480) and it would be great if pyarrow-backed extension dtypes were also supported.

[trivialfis]

Hi @jrbourbeau , thank you for raising the issue. A quick search didn't return any relevant document for enumerating the types. Could you please share some references I can use?

https://pandas.pydata.org/pandas-docs/stable/ecosystem.html#ecosystem-extensions
https://pandas.pydata.org/pandas-docs/stable/development/extending.html#compatibility-with-apache-arrow
https://pandas.pydata.org/pandas-docs/stable/user_guide/basics.html#dtypes

[jrbourbeau]

Thanks @trivialfis! Yeah, I'm under the impression that extensive documentation around the new pyarrow dtypes is still in development. IIUC I think any pyarrow data type (listed here https://arrow.apache.org/docs/python/api/datatypes.html) can now be converted to a pandas dtype through pd.ArrowDtype.

In [1]: import pandas as pd

In [2]: import pyarrow as pa

In [3]: pd.ArrowDtype(pa.uint8())
Out[3]: uint8[pyarrow]

Also cc @mroeschke who is the expert in on the pandas side

[mroeschke]

Sorry this is not nicely document as of the pandas 1.5 release. The public API (namely pd.ArrowDtype and pd.arrays.ArrowExtensionArray) can be found here: https://pandas.pydata.org/pandas-docs/stable/reference/arrays.html#pyarrow

Feel free to ping me with any further questions!

[trivialfis]

I gained some basic understanding of arrow recently due to work on CUDA IPC. I think XGBoost needs a new interface with pyarrow. Will try to build a prototype later.

[trivialfis]

Hey, out of curiosity, is this going to be used in dask in recent future? It would be really exciting!

[jrbourbeau]

I think XGBoost needs a new interface with pyarrow. Will try to build a prototype later.

Woo! Feel free to ping me into anything if you think it'd be useful

Hey, out of curiosity, is this going to be used in dask in recent future? It would be really exciting!

Yes! We're currently working on improving Dask's support for pyarrow-backed dtypes. FWIW I actually ran across this issue originally using xgboost.dask and a Dask DataFrame with int64[pyarrow], string[pyarrow], etc. data. If you have some examples of where pyarrow-backed dtypes have been / would be useful to you, @mroeschke and I would love to hear about them : )

[trivialfis]

Hi, I opened a PR #8653 . It's not completed yet since boolean and dictionary are not supported. Please help take a look when you have additional bandwidth (would be great if you have ways to improve the pandas handling code, it's becoming more and more complicated now).

[fingoldo]

Hi, I opened a PR #8653 . It's not completed yet since boolean and dictionary are not supported. Please help take a look when you have additional bandwidth (would be great if you have ways to improve the pandas handling code, it's becoming more and more complicated now).

Handling of caterorical dtype backed by string[pyarrow] and large_string[pyarrow] would be very desirable.

Currently (in 3.1.1) it fails with

---

TypeError Traceback (most recent call last)
Cell In[7], line 6
1 from xgboost import XGBRegressor
3 XGBRegressor(
4 verbosity=3,
5 enable_categorical=True,
----> 6 ).fit(
7 df,
8 df["f1"],
9 ).predict(df)

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:774, in require_keyword_args..throw_if..inner_f(*args, **kwargs)
772 for k, arg in zip(sig.parameters, args):
773 kwargs[k] = arg
--> 774 return func(**kwargs)

File /venv/main/lib/python3.12/site-packages/xgboost/sklearn.py:1340, in XGBModel.fit(self, X, y, sample_weight, base_margin, eval_set, verbose, xgb_model, sample_weight_eval_set, base_margin_eval_set, feature_weights)
1337 model, feature_types = get_model_categories(X, model, self.feature_types)
1339 evals_result: EvalsLog = {}
-> 1340 train_dmatrix, evals = _wrap_evaluation_matrices(
1341 missing=self.missing,
1342 X=X,
1343 y=y,
1344 group=None,
1345 qid=None,
1346 sample_weight=sample_weight,
1347 base_margin=base_margin,
1348 feature_weights=feature_weights,
1349 eval_set=eval_set,
1350 sample_weight_eval_set=sample_weight_eval_set,
1351 base_margin_eval_set=base_margin_eval_set,
1352 eval_group=None,
1353 eval_qid=None,
1354 create_dmatrix=self._create_dmatrix,
1355 enable_categorical=self.enable_categorical,
1356 feature_types=feature_types,
1357 )
1359 if callable(self.objective):
1360 obj: Optional[Objective] = _objective_decorator(self.objective)

File /venv/main/lib/python3.12/site-packages/xgboost/sklearn.py:701, in _wrap_evaluation_matrices(missing, X, y, group, qid, sample_weight, base_margin, feature_weights, eval_set, sample_weight_eval_set, base_margin_eval_set, eval_group, eval_qid, create_dmatrix, enable_categorical, feature_types)
696 """Convert array_like evaluation matrices into DMatrix. Perform sanity checks on the
697 way.
698
699 """
700 # Feature_types contains the optional reference categories from the booster object.
--> 701 train_dmatrix = create_dmatrix(
702 data=X,
703 label=y,
704 group=group,
705 qid=qid,
706 weight=sample_weight,
707 base_margin=base_margin,
708 feature_weights=feature_weights,
709 missing=missing,
710 enable_categorical=enable_categorical,
711 feature_types=feature_types,
712 ref=None,
713 )
715 n_validation = 0 if eval_set is None else len(eval_set)
716 if hasattr(train_dmatrix, "get_categories"):

File /venv/main/lib/python3.12/site-packages/xgboost/sklearn.py:1259, in XGBModel._create_dmatrix(self, ref, **kwargs)
1257 except TypeError: # QuantileDMatrix supports lesser types than DMatrix
1258 pass
-> 1259 return DMatrix(**kwargs, nthread=self.n_jobs)

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:774, in require_keyword_args..throw_if..inner_f(*args, **kwargs)
772 for k, arg in zip(sig.parameters, args):
773 kwargs[k] = arg
--> 774 return func(**kwargs)

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:999, in DMatrix.init(self, data, label, weight, base_margin, missing, silent, feature_names, feature_types, nthread, group, qid, label_lower_bound, label_upper_bound, feature_weights, enable_categorical, data_split_mode)
996 assert self.handle is not None
997 return
--> 999 handle, feature_names, feature_types = dispatch_data_backend(
1000 data=data,
1001 missing=self.missing,
1002 threads=self.nthread,
1003 feature_names=feature_names,
1004 feature_types=feature_types,
1005 enable_categorical=enable_categorical,
1006 data_split_mode=data_split_mode,
1007 )
1008 assert handle is not None
1009 self.handle = handle

File /venv/main/lib/python3.12/site-packages/xgboost/data.py:1419, in dispatch_data_backend(data, missing, threads, feature_names, feature_types, enable_categorical, data_split_mode)
1417 data = pd.DataFrame(data)
1418 if _is_pandas_df(data):
-> 1419 return _from_pandas_df(
1420 data=data,
1421 enable_categorical=enable_categorical,
1422 missing=missing,
1423 nthread=threads,
1424 feature_names=feature_names,
1425 feature_types=feature_types,
1426 data_split_mode=data_split_mode,
1427 )
1428 if _is_cudf_df(data) or _is_cudf_ser(data):
1429 return _from_cudf_df(
1430 data=data,
1431 missing=missing,
(...) 1435 enable_categorical=enable_categorical,
1436 )

File /venv/main/lib/python3.12/site-packages/xgboost/data.py:700, in _from_pandas_df(data, enable_categorical, missing, nthread, feature_names, feature_types, data_split_mode)
690 def _from_pandas_df(
691 *,
692 data: "PdDataFrame",
(...) 698 data_split_mode: DataSplitMode = DataSplitMode.ROW,
699 ) -> DispatchedDataBackendReturnType:
--> 700 df, feature_names, feature_types = _transform_pandas_df(
701 data, enable_categorical, feature_names, feature_types
702 )
704 handle = ctypes.c_void_p()
705 _check_call(
706 _LIB.XGDMatrixCreateFromColumnar(
707 df.array_interface(),
(...) 712 )
713 )

File /venv/main/lib/python3.12/site-packages/xgboost/data.py:668, in _transform_pandas_df(data, enable_categorical, feature_names, feature_types, meta)
662 feature_names, feature_types = pandas_feature_info(
663 data, meta, feature_names, feature_types, enable_categorical
664 )
666 arrays = pandas_transform_data(data)
667 return (
--> 668 PandasTransformed(arrays, ref_categories=ref_categories),
669 feature_names,
670 feature_types,
671 )

File /venv/main/lib/python3.12/site-packages/xgboost/data.py:625, in PandasTransformed.init(self, columns, ref_categories)
622 for col in self.columns:
623 if _is_df_cat(col):
624 # Categorical column
--> 625 jnames, jcodes, buf = pd_cat_inf(col.categories, col.codes)
626 self.temporary_buffers.append(buf)
627 aitfs.append((jnames, jcodes))

File /venv/main/lib/python3.12/site-packages/xgboost/_data_utils.py:407, in pd_cat_inf(cats, codes)
404 # pandas uses -1 to represent missing values for categorical features
405 codes = codes.replace(-1, np.nan)
--> 407 if np.issubdtype(cats.dtype, np.floating) or np.issubdtype(cats.dtype, np.integer):
408 # Numeric index type
409 name_values_num = cats.values
410 jarr_values = array_interface_dict(name_values_num)

File /venv/main/lib/python3.12/site-packages/numpy/core/numerictypes.py:417, in issubdtype(arg1, arg2)
359 r"""
360 Returns True if first argument is a typecode lower/equal in type hierarchy.
361
(...) 414
415 """
416 if not issubclass_(arg1, generic):
--> 417 arg1 = dtype(arg1).type
418 if not issubclass_(arg2, generic):
419 arg2 = dtype(arg2).type

TypeError: Cannot interpret 'large_string[pyarrow]' as a data type

reproducer:

import pandas as pd
import pyarrow as pa

# Arrow float32 arrays
a = pa.array([1.1, 2.2, 3.3], type=pa.float32())
b = pa.array([4.4, 5.5, 6.6], type=pa.float32())
c = pa.array([7.7, 8.8, 9.9], type=pa.float32())
d = pa.array([10.1, 11.2, 12.3], type=pa.float32())

# Arrow dictionary (categorical) array: dictionary values ["cat", "dog"]
indices = pa.array([0, 1, 0], type=pa.int8())
dictionary = pa.array(["cat", "dog"], type=pa.large_utf8())
cat = pa.DictionaryArray.from_arrays(indices, dictionary)

# Build pandas Series using Arrow-backed dtypes
s_a = pd.Series(a, dtype=pd.ArrowDtype(pa.float32()))
s_b = pd.Series(b, dtype=pd.ArrowDtype(pa.float32()))
s_c = pd.Series(c, dtype=pd.ArrowDtype(pa.float32()))
s_d = pd.Series(d, dtype=pd.ArrowDtype(pa.float32()))
s_cat = pd.Series(cat, dtype=pd.ArrowDtype(cat.type)).astype('category')  # Arrow dictionary type

df = pd.DataFrame({
    "f1": s_a,
    "f2": s_b,
    "f3": s_c,
    "f4": s_d,
    "animal": s_cat
})

print(df)
print("\nDtypes:")
print(df.dtypes)

from xgboost import  XGBRegressor

XGBRegressor(
    verbosity=3,
    enable_categorical=True,
).fit(
    df,
    df["f1"],
).predict(df)

[trivialfis]

It's not too difficult to create basic support: #11801

However, it seems quite challenging to create a non-trivial dictionary column for pandas using arrow at the moment. string[pyarrow] is not a categorical column, I don't know how to apply an arrow-backed pandas column for a categorical feature in a practical use case.

[fingoldo]

It's not too difficult to create basic support: #11801

However, it seems quite challenging to create a non-trivial dictionary column for pandas using arrow at the moment. string[pyarrow] is not a categorical column, I don't know how to apply an arrow-backed pandas column for a categorical feature in a practical use case.

Thank you for this amazingly fast turnaround! You need more time to decide if your PR is safe to merge, did I get it correct?

[trivialfis]

I need to know how to use it for a non trivial case so that I can do some real world tests. Do you know how to create a pandas dataframe backed by arrow dictionary for a real world dataset (without manually specifying the encoding)

[fingoldo]

I need to know how to use it for a non trivial case so that I can do some real world tests. Do you know how to create a pandas dataframe backed by arrow dictionary for a real world dataset (without manually specifying the encoding)

Oh. That's simple. I came across this problem when trying to apply xgboost to a pandas view of a polars dataframe.
View was constructed using zero-copy conversion from polars to arrow, then from arrow to pandas.
As you can see in this example, converting such str pandas columns to categoricals result in above mentioned dtype, category backed by pyarrow:

import polars as pl, pandas as pd
import numpy as np

df = pl.DataFrame(
    {
        "price": np.random.uniform(10, 100, size=5),  # float column
        "weight": np.random.normal(70, 10, size=5),  # float column
        "color": ["red", "blue", "green", "blue", "red"],  # categorical (string)
    }
)

converted_pd_df = df.to_arrow().to_pandas(
    types_mapper=pd.ArrowDtype
)  # keep Arrow-backed columns for zero-copy)

converted_pd_df["color"] = converted_pd_df["color"].astype("category")
print(converted_pd_df["color"])

0 red
1 blue
2 green
3 blue
4 red
Name: color, dtype: category
Categories (3, large_string[pyarrow]): [blue, green, red]

While there are workarounds for my particular use case, I reported this inconsistency so that other users don't have to meet it (in this exact and related forms ) in future at all.

[trivialfis]

@fingoldo Thank you for sharing. I misunderstood something there. You want to have a normal, numpy-backed pandas categorical column, except that the index is backed by arrow strings instead of numpy strings. But you don't want the arrow dictionary to back the pandas categorical column:

indices = pa.array([0, 1, 0], type=pa.int8())
dictionary = pa.array(["cat", "dog"], type=pa.large_utf8())
cat = pa.DictionaryArray.from_arrays(indices, dictionary)

a = pd.Series(cat, dtype=pd.ArrowDtype(cat.type)) # We can support this in the future
b = pd.Series(cat, dtype=pd.ArrowDtype(cat.type)).astype("category") # But not this

print(a.dtype)
>>> dictionary<values=large_string, indices=int8, ordered=0>[pyarrow]  # backed by arrow dictionary
print(b.dtype)
>>> category # backed by numpy, with arrow strings as cat values. This is not a well-defined type.

I don't think this is the best solution, as it combines too many elements. We will try to support dictionary-backed columns, but not numpy-backed columns with arrow strings (which is object, btw).

If you want to process categorical features and are using Polars, I suggest using Polars directly with XGBoost by specifying either categorical or enum: https://docs.pola.rs/user-guide/expressions/categorical-data-and-enums/. There's no need for these conversions.