Does (or will) scikit learn support Arrow/pyarrow data? #25450

oliviermeslin · 2023-01-21T19:08:42Z

oliviermeslin
Jan 21, 2023

TL;DR

I have a general question regarding the compatibility of sklearn and pyarrow (and generally arrow). My impression is that sklearn does not support pyarrow Tables as inputs. Although it would imply a significant workload, I really think that adding arrow support would be valuable to sklearn users because it would make it possible to build faster, more efficient end-to-end pipelines based entirely on Arrow (getting rid of pandas).

Context

I'm currently developing a standard ML project for my employer. This project has two main steps:

1/ building the dataset of interest (implying a lot of data manipulations, joins, reshapes...);
2/ a standard sklearn pipeline 10000 with some preprocessing (StandardScaler() and OneHotEncoder()) and an XGBoost Regressor.

My pipeline looks like this:

preprocessor = ColumnTransformer(
    transformers=[
        ("num", StandardScaler(), make_column_selector(dtype_exclude=["category","object"])),
        ("cat", OneHotEncoder(), make_column_selector(dtype_include=["category","object"])),
    ]
)
my_pipeline= Pipeline(
    steps=[("preprocessor", preprocessor), ("classifier", xgb.XGBRegressor())]
)

Until now, I used pandas for step 1 and an sklearn for step 2. Recently, I replaced pandas with polars in step 1 for two reasons: 1/ switching to polars results in a massive speed-up of the data manipulation step (~10-20x faster), and 2/ XGBoost now supports Arrow tables (and a polars table can be easily converted to an Arrow table). Similar feature requests have been expressed for the two main competitors of XGBoost (see here for lightGBM and here for catboost).

Problem

My problem is the following: I tried to fit my pipeline using my_pipeline.fit(X_train, y) (where X_train is an Pyarrow Table and y is an Pyarrow FloatArray), I get the following error: ValueError: make_column_selector can only be applied to pandas dataframes.

I tried to adapt my pipeline in multiple ways, but all my tests resulted in various errors all related to the fact that sklearn could not use the data I was using in the fit function. I looked for answers online, but all I could find is this sentence in sklearn's FAQ: "Most of scikit-learn assumes data is in NumPy arrays or SciPy sparse matrices of a single numeric dtype".

All in all, it seems that my problem is that sklearn does not seem to support Arrow Tables and Arrays as input.

Questions

I have three questions:

Is it the case that sklearn does not support Arrow yet? That would not be surprising given that this framework is quite recent.
If arrow is not currently supported, did someone express a feature request regarding arrow support in sklearn?
If there is no such feature request yet, should I open one, given that I don't have the skills to suggest an implementation?

Answered by glemaitre

Jan 21, 2023

We don't support arrow for the moment. We have some ongoing work looking at the array API: #22352.

If the dataframe API (https://data-apis.org/dataframe-api/draft/API_specification/index.html) also become mature and packages are implementing it, we might be able to have a way forward to interoperability. But this is a bit early for the moment.

View full answer

glemaitre · 2023-01-21T20:04:05Z

glemaitre
Jan 21, 2023
Maintainer

We don't support arrow for the moment. We have some ongoing work looking at the array API: #22352.

If the dataframe API (https://data-apis.org/dataframe-api/draft/API_specification/index.html) also become mature and packages are implementing it, we might be able to have a way forward to interoperability. But this is a bit early for the moment.

2 replies

braaannigan Jan 25, 2023

Hi @glemaitre

Pandas has implemented the dataframe API (https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.__dataframe__.html) and Polars is implementing it.

I'd love to see Arrow support in sklearn. I think this would be justified as there is rapid growth in usage of Arrow-based or Arrow-friendly data processing tools (e.g. Polars, DuckDB, DataFusion)

MarcoGorelli Feb 10, 2023

Hey @braaannigan

As far as I understand, it's only the interchange protocol that's been implemented - the DataFrame API Standard is still quite far away

UPDATE as of March 2024:

Scikit-learn is already using the dataframe interchange protocol and the NumPy array interface to support Polars (it converts to numpy, it doesn't operate on arrow tables directly)

I think the dataframe API standard project is dead. Or at least, there's been no activity on it since I quit it

oliviermeslin · 2023-01-25T14:07:14Z

oliviermeslin
Jan 25, 2023
Author

Thanks a lot @glemaitre for your answer!

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Does (or will) scikit learn support Arrow/pyarrow data? #25450

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Replies: 2 comments 2 replies

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

Select a reply

Does (or will) scikit learn support Arrow/pyarrow data? #25450

oliviermeslin Jan 21, 2023

TL;DR

Context

Problem

Questions

Replies: 2 comments · 2 replies

glemaitre Jan 21, 2023 Maintainer

braaannigan Jan 25, 2023

MarcoGorelli Feb 10, 2023

oliviermeslin Jan 25, 2023 Author

oliviermeslin
Jan 21, 2023

Replies: 2 comments 2 replies

glemaitre
Jan 21, 2023
Maintainer

oliviermeslin
Jan 25, 2023
Author