Supporting different dataframe libraries

[Vincent-Maladiere]

Hey all,

As we progressively broaden the scope of skrub, some long-term questions need to be addressed.

Question 1

We initially intended to work with pandas dataframes and numpy arrays, as scikit-learn does. However, with the rise of polars, duckdb, and many more, we have to wonder which modules to support.

We can stay lean and only choose to add polars for now, but what do we want in the long term?

[Edit]
It's important to note that for the moment, all skrub transformers convert dataframes to numpy arrays —some transformers like MinHashEncoder or DatetimeEncoder don't use scikit-learn estimators and could use dataframes instead of numpy arrays though.

For transformers that require scikit-learn, like the GapEncoder, the TargetEncoder, and the Joiner, we need to make a decision. Should we:

• raise an error right away when the input is not a pandas dataframe?
• try to convert the input to pandas, forwarding an error to the user if the input is too large or not compatible?
• implement our own out-of-core strategy around scikit-learn when the input is too large for memory, like dumping the input to disk and loading pandas dataframes by batch?

Question 2

We shouldn't try to adapt the codebase to each dataframe module because it would quickly become unsustainable. Instead, we have two options:

Option 1: The dataframe API

As introduced during EuroScipy 2023 by @MarcoGorelli, the dataframe API is now available for both pandas and polars dataframes, plus polars lazyframes!

We can use it as follows:

# pip install dataframe_api_compat
import pandas as pd
import polars as pl

dfs = [
    pd.DataFrame(dict(a=[1, 3, 4])),
	pl.DataFrame(dict(a=[1, 3, 4])),
	pl.LazyFrame(dict(a=[1, 3, 4])),
]

for df in dfs:
	df_std = df.__dataframe_consortium_standard__()
	df_std.get_column_by_name("a").mean()

See the API specs.

However, the dataframe API is still in its early days, and important features such as joining are not there (yet?).

Option 2: Ibis

Ibis has been around for some time and supports a wide array of different modules.

As an SQL engine, it can also directly query databases like PostgreSQL or Snowflake, which is a significant feature in the long term if we want to operate on remote databases.

I also like the simpler syntax (except for the setup boilerplate) and features that are already there.

# pip install 'ibis-framework[duckdb, polars]'
import sys
import ibis


def _is_pandas(df):
    if "pandas" not in sys.modules:
        return False

    import pandas as pd
        
    return isinstance(df, pd.DataFrame)
        

def _is_polars(df):
    if not "polars" in sys.modules:
        return False

    import polars as pl
        
    return isinstance(df, (pl.DataFrame, pl.LazyFrame))


dfs = [
    pd.DataFrame(dict(a=[1, 3, 4])),
    pl.DataFrame(dict(a=[1, 3, 4])),
    pl.LazyFrame(dict(a=[1, 3, 4])),
]

for df in dfs:
    if _is_pandas(df):
        con = ibis.pandas.connect()  
    elif _is_polars(df):
        con = ibis.polars.connect()
    else:
        raise NotImplementedError()
            
    con.do_connect({"my_df": df})  
    t = con.table("my_df")
    print(t["a"].mean())

---

Both solutions introduce an additional dependency (dataframe_api_compat or ibis) plus the desired modules.

What do you all think of this? Do you see some other ways forward?

[MarcoGorelli]

Thanks for starting the discussion!

Let's look at an example - GapEncoder from skrub. It looks to me like you:

• take in a dataframe
• scikit-learn converts it to object ndarray - and it has to be object because it needs to store strings. And it needs to store strings because you're doing NLP stuff like CountVectorizer

So, I don't think that the Consortium's DataFrame Standard (in its current form at least) would be sufficient for this task, because:

• you would need to convert to object ndarray
• you can't use df.__dataframe_consortium_standard__().to_array_object, because that requires you to specify a dtype from the Array API, and the Array API doesn't have object nor string in its dtypes: https://data-apis.org/array-api/2022.12/API_specification/data_types.html

I'm tempted to suggest you just use the interchange protocol to interchange to pandas, using .select_columns_by_name to only select the necessary columns, like plotly do (see plotly/plotly.py#4286)

[MarcoGorelli]

cc @rgommers @kkraus14 @thomasjpfan in case there's something I'm missing in my assessment above

[rgommers]

Both solutions introduce an additional dependency (dataframe_api_compat or ibis) plus the desired modules.

dataframe_api_compat is designed so it can be easily vendored. That's also what both scikit-learn and scipy do with array-api-compat. No one likes new dependencies:)

scikit-learn converts it to object ndarray

I would not make any long-term decisions based on object arrays being used. Both because strings fit better with dataframes in general (especially if we get a nice connection between the dataframe and array API standards), and because the implementation of a proper variable-length string dtype for NumPy 2.0 is just done (see the just submitted first version of NEP 55).

[MarcoGorelli]

Thanks for weighing in - I'll add that dataframe_api_compat is currently just a few pure-Python files, it's extremely light. It's 17 kB on PyPI, and has zero required dependencies (not even pandas or polars - it just uses whatever the user started with). So it should be a low-cost dependency to add, regardless of whether you vendor or not

Regarding strings - is it plausible that the String dtype might make its way into the Array API spec one day?

I chatted with @Vincent-Maladiere earlier, and he brought up the MinHashEncoder and DatetimeEncoder transformers as examples of classes in skrub which don't defer to scikit-learn. Support for other dataframe libraries doesn't have to be "all or nothing" - perhaps lets try the compat package out with those, and see how far we get?

I know for a fact that we won't get far enough, but this exercise will help highlight what else we're missing the Standard - then, we can address them 😎

Thanks all for the collaboration 🤝

[rgommers]

Regarding strings - is it plausible that the String dtype might make its way into the Array API spec one day?

I'm 99.x% certain that that won't happen within the next five years or so, because other libraries like JAX, CuPy and PyTorch don't have and won't care about strings.

[MarcoGorelli]

Just to check then - the dataframe API isn't a solution to skub supporting GapEncoder across dataframe libraries?

If so, I'm tempted to suggest: just check if the library is pandas or polars, and use .to_numpy if so. Hardcode the list of supported libraries - this wouldn't be as good as full cross-dataframe support, but it would be better than only supporting pandas IMO

[GaelVaroquaux]

Coming back from travels. My thoughts:

Question 1 (what to do with an input that is not a pandas dataframe):

My answer: use the dataframe API to convert it to pandas (interesting suggestions on how to do it above, eg by @MarcoGorelli). But, maybe worry about this after release 0.1. Consider vendoring dataframe_api_compat.

Question 2 (how to support many dataframe implementation):

My answer: favor the dataframe API, but consider special-casing a few implementations (for now I have in mind pandas and polars) to have efficient codebases (eg lazy for polars, .query for pandas). Also, cover these in tests.