[WIP] Interpolationjoiner dataframe api

.github/workflows/testing.yml

@@ -75,6 +75,7 @@ jobs:
       - shell: bash {0}
         run: |
           cp $GITHUB_WORKSPACE/pyproject.toml .
+          cp $GITHUB_WORKSPACE/conftest.py .
           $GITHUB_WORKSPACE/build_tools/github/test.sh
         working-directory: ${{ runner.temp }}
         name: 'Run tests'

CHANGES.rst

@@ -17,7 +17,7 @@ Major changes
 
 * :class:`InterpolationJoiner` was added to join two tables by using
   machine-learning to infer the matching rows from the second table.
-  :pr:`742` by :user:`Jérôme Dockès <jeromedockes>`.
+  :pr:`742` and :pr:`827` by :user:`Jérôme Dockès <jeromedockes>`.
 
 * Pipelines including :class:`TableVectorizer` can now be grid-searched, since
   we can now call `set_params` on the default transformers of :class:`TableVectorizer`.

conftest.py

@@ -0,0 +1,15 @@
+import pandas
+import pytest
+
+DATAFRAME_MODULES = [pandas]
+try:
+    import polars
+
+    DATAFRAME_MODULES.append(polars)
+except ImportError:
+    pass
+
+
+@pytest.fixture(params=DATAFRAME_MODULES)
+def px(request):
+    return request.param

setup.cfg

@@ -25,11 +25,12 @@ project_urls =
 include_package_data = True
 packages = find:
 install_requires =
-    scikit-learn>=1.2.1
+    scikit-learn>=1.3.0
     numpy>=1.23.5
     scipy>=1.9.3
-    pandas>=1.5.3
+    pandas>=2.1.0
     packaging>=23.1
+    dataframe-api-compat>=0.1.28
 python_requires = >=3.10
 
 [options.extras_require]
@@ -76,10 +77,10 @@ benchmarks =
 # Overwrite the previous install_requires for CI testing purposes
 # as defined in testing.yml.
 min-py310 =
-    scikit-learn==1.2.1
+    scikit-learn==1.3.0
     numpy==1.23.5
     scipy==1.9.3
-    pandas==1.5.3
+    pandas==2.1.0
 
 [flake8]
 # max line length for black

skrub/_dataframe/__init__.py

@@ -0,0 +1,4 @@
+from ._common import Selector, std, stdns
+from ._namespace import get_df_namespace, skrubns
+
+__all__ = ["get_df_namespace", "skrubns", "std", "stdns", "Selector"]

skrub/_dataframe/_common.py

@@ -0,0 +1,22 @@
+import enum
+
+
+class Selector(enum.Enum):
+    ALL = enum.auto()
+    NONE = enum.auto()
+    NUMERIC = enum.auto()
+    CATEGORICAL = enum.auto()
+
+
+def std(obj):
+    try:
+        return obj.__dataframe_consortium_standard__()
+    except AttributeError:
+        return obj.__column_consortium_standard__()
+
+
+def stdns(obj):
+    try:
+        return obj.__dataframe_consortium_standard__().__dataframe_namespace__()
+    except AttributeError:
+        return obj.__column_consortium_standard__().__column_namespace__()

skrub/_dataframe/_namespace.py

@@ -97,3 +97,8 @@ def get_df_namespace(*dfs):
                 "Only Pandas or Polars dataframes are currently supported, "
                 f"got {modules=!r}."
             )
+
+
+def skrubns(*dataframes):
+    ns, _ = get_df_namespace(*dataframes)
+    return ns

skrub/_dataframe/_pandas.py

@@ -9,6 +9,22 @@
 
 from skrub._utils import atleast_1d_or_none
 
+from ._common import Selector
+
+__all__ = [
+    "make_dataframe",
+    "make_series",
+    "aggregate",
+    "join",
+    "split_num_categ_cols",
+    "select",
+    "drop",
+    "Selector",
+    "concat_horizontal",
+    "any_rowwise",
+    "to_pandas",
+]
+
 
 def make_dataframe(X, index=None):
     """Convert an dictionary of columns into a Pandas dataframe.
@@ -329,4 +345,34 @@
 
 
 def select(dataframe, columns):
-    return dataframe[columns]
+    if not isinstance(columns, Selector):
+        return dataframe[columns]
+    if columns is Selector.ALL:
+        return dataframe
+    elif columns is Selector.NONE:
+        return dataframe[[]]
+    elif columns is Selector.NUMERIC:
+        return dataframe.select_dtypes("number")
+    elif columns is Selector.CATEGORICAL:
+        return dataframe.select_dtypes(["object", "string", "category"])
+    # we have covered all items in the enumeration
+    assert False
+
+
+def drop(dataframe, columns):
+    return dataframe.drop(select(dataframe, columns).columns.values, axis=1)
+
+
+def any_rowwise(dataframe):
+    return dataframe.any(axis=1)
+
+
+def concat_horizontal(dataframe, *other_dataframes):
+    other_dataframes = [
+        df.set_axis(dataframe.index, axis="index") for df in other_dataframes
+    ]
+    return pd.concat([dataframe] + list(other_dataframes), axis=1)
+
+
+def to_pandas(dataframe):
+    return dataframe

skrub/_dataframe/_polars.py

@@ -13,6 +13,23 @@
 
 from skrub._utils import atleast_1d_or_none
 
+from ._common import Selector
+
+__all__ = [
+    "POLARS_SETUP",
+    "make_dataframe",
+    "make_series",
+    "aggregate",
+    "join",
+    "split_num_categ_cols",
+    "select",
+    "drop",
+    "Selector",
+    "concat_horizontal",
+    "any_rowwise",
+    "to_pandas",
+]
+
 
 def make_dataframe(X, index=None):
     """Convert an dictionary of columns into a Polars dataframe.
@@ -263,5 +280,47 @@
     return num_cols, categ_cols
 
 
+def _check_selector(columns):
+    if not isinstance(columns, Selector):
+        return columns
+    if columns is Selector.ALL:
+        return cs.all()
+    elif columns is Selector.NONE:
+        return []
+    elif columns is Selector.NUMERIC:
+        return cs.numeric()
+    elif columns is Selector.CATEGORICAL:
+        return cs.string(include_categorical=True)
+    # we have covered all items in the enumeration
+    assert False
+
+
 def select(dataframe, columns):
-    return dataframe.select(columns)
+    return dataframe.select(_check_selector(columns))
+
+
+def drop(dataframe, columns):
+    return dataframe.drop(_check_selector(columns))
+
+
+def any_rowwise(dataframe):
+    return _collect(dataframe.select(pl.any_horizontal(pl.all()))).get_column("any")
+
+
+def concat_horizontal(dataframe, *other_dataframes):
+    return pl.concat(
+        [_collect(dataframe)] + [_collect(df) for df in other_dataframes],
+        how="horizontal",
+    )
+
+
+def _collect(dataframe):
+    if hasattr(dataframe, "collect"):
+        dataframe = dataframe.collect()
+    return dataframe
+
+
+def to_pandas(dataframe):
+    if hasattr(dataframe, "collect"):
+        dataframe = dataframe.collect()
+    return dataframe.to_pandas()

skrub/_dataframe/tests/test_all.py

@@ -0,0 +1,43 @@
+import pytest
+
+from skrub._dataframe import Selector, skrubns
+
+
+@pytest.fixture
+def df(px):
+    return px.DataFrame(
+        {"ID": [2, 3, 7], "name": ["ab", "cd", "01"], "temp": [20.3, 40.9, 11.5]}
+    )
+
+
+def test_select(df):
+    ns = skrubns(df)
+    assert list(ns.select(df, []).columns) == []
+    assert list(ns.select(df, ["name"]).columns) == ["name"]
+    assert list(ns.select(df, Selector.ALL).columns) == list(df.columns)
+    assert list(ns.select(df, Selector.NONE).columns) == []
+    assert list(ns.select(df, Selector.NUMERIC).columns) == ["ID", "temp"]
+    assert list(ns.select(df, Selector.CATEGORICAL).columns) == ["name"]
+
+
+def test_drop(df):
+    ns = skrubns(df)
+    assert list(ns.drop(df, []).columns) == list(df.columns)
+    assert list(ns.drop(df, ["name"]).columns) == ["ID", "temp"]
+    assert list(ns.drop(df, Selector.ALL).columns) == []
+    assert list(ns.drop(df, Selector.NONE).columns) == list(df.columns)
+    assert list(ns.drop(df, Selector.NUMERIC).columns) == ["name"]
+    assert list(ns.drop(df, Selector.CATEGORICAL).columns) == ["ID", "temp"]
+
+
+def test_concat_horizontal(df):
+    ns = skrubns(df)
+    df1 = (
+        df.__dataframe_consortium_standard__()
+        .rename_columns({c: f"{c}_1" for c in df.columns})
+        .dataframe
+    )
+    out = ns.concat_horizontal(df)
+    assert list(out.columns) == list(df.columns)
+    out = ns.concat_horizontal(df, df1)
+    assert list(out.columns) == list(df.columns) + list(df1.columns)

skrub/_dataframe/tests/test_common.py

@@ -0,0 +1,11 @@
+from skrub._dataframe import skrubns, std, stdns
+
+
+def test_std(px):
+    df = px.DataFrame({"A": [1, 2]})
+    assert hasattr(std(df), "dataframe")
+    assert hasattr(stdns(df), "dataframe_from_columns")
+    ns = skrubns(df)
+    s = ns.make_series([1, 2], name="A")
+    assert hasattr(std(s), "column")
+    assert hasattr(stdns(s), "dataframe_from_columns")