Merge pull request #12 from adrinjalali/SelectCol

MNT remove skrub dep and have our own SelectCols

Author: koaning

.github/workflows/unittest.yml

@@ -29,5 +29,5 @@ jobs:
         source .venv/bin/activate
         which python
         python --version
-        uv pip install -e . pytest
+        uv pip install -e .[test]
         uv run pytest

playtime/__init__.py

@@ -3,7 +3,7 @@
 from sklearn.preprocessing import FunctionTransformer, OneHotEncoder, SplineTransformer
 from sklearn.compose import make_column_transformer
 from sklearn.feature_extraction.text import CountVectorizer
-from skrub import SelectCols
+from .estimators import SelectCols
 from .transformer_functions import datetime_feats
 from .formula import PlaytimePipeline
 from typing import Any
@@ -41,9 +41,7 @@ def select(*colnames: str) -> PlaytimePipeline:
     pipeline = select("col_a", "col_b")
     ```
     """
-    return PlaytimePipeline(
-        pipeline=make_pipeline(SelectCols([col for col in colnames]))
-    )
+    return PlaytimePipeline(pipeline=make_pipeline(SelectCols(colnames)))
 
 
 def onehot(*colnames: str, **kwargs) -> PlaytimePipeline:

playtime/estimators.py

@@ -0,0 +1,74 @@
+import narwhals as nw
+from sklearn.base import BaseEstimator, TransformerMixin
+
+
+class SelectCols(TransformerMixin, BaseEstimator):
+    """Select a subset of a DataFrame's columns.
+
+    A ``ValueError`` is raised if any of the provided column names are not in
+    the dataframe.
+
+    Accepts anything accepted by :func:`narwhals.from_native(..., eager_only=True)`.
+
+    Arguments
+        cols : list of str or str
+            The columns to select. A single column name can be passed as a ``str``:
+            ``"col_name"`` is the same as ``["col_name"]``.
+
+    **Usage**
+    ```python
+    >>> import polars as pl
+    >>> from playtime.estimators import SelectCols
+    >>> df = pl.DataFrame({"A": [1, 2], "B": [10, 20], "C": ["x", "y"]})
+    >>> df
+       A   B  C
+    0  1  10  x
+    1  2  20  y
+    >>> SelectCols(["C", "A"]).fit_transform(df)
+       C  A
+    0  x  1
+    1  y  2
+    >>> SelectCols(["X", "A"]).fit_transform(df)
+    Traceback (most recent call last):
+        ...
+    ValueError: The following columns are requested for selection but missing from dataframe: ['X']
+    ```
+    """
+
+    def __init__(self, cols=None):
+        self.cols = cols
+
+    def fit(self, X, y=None):
+        """Fit the transformer.
+
+        Arguments
+            X : DataFrame or None
+                If `X` is a DataFrame, the transformer checks that all the column
+                names provided in ``self.cols`` can be found in `X`.
+
+            y : None
+                Unused.
+
+        Returns
+        -------
+        SelectCols
+            The transformer itself.
+        """
+        nw.from_native(X, eager_only=True).select(self.cols)
+        return self
+
+    def transform(self, X):
+        """Transform a dataframe by selecting columns.
+
+        Parameters
+        ----------
+        X : DataFrame
+            The DataFrame on which to apply the selection.
+
+        Returns
+        -------
+        DataFrame
+            The input DataFrame ``X`` after selecting only the columns listed
+            in ``self.cols`` (in the provided order).
+        """
+        return nw.from_native(X, eager_only=True).select(self.cols)

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
 
 [project.optional-dependencies]
 lint = ["pre-commit"]
+test = ["pytest", "pandas"]
 
 [project.urls]
 repository = "https://github.com/koaning/playtime"

tests/test_estimators.py

@@ -0,0 +1,19 @@
+import pytest
+
+from playtime.estimators import SelectCols
+
+
+@pytest.mark.parametrize("lib", ["pandas", "polars"])
+def test_select_cols(lib):
+    lib = pytest.importorskip(lib)
+
+    df = lib.DataFrame({"a": [1, 2], "b": [10, 20], "c": ["x", "y"]})
+    tfm = SelectCols(["a", "b"])
+    # This should work w/o calling `fit`
+    assert tfm.transform(df).columns == ["a", "b"]
+    # This should also pass
+    tfm.fit(df)
+
+    # In reality it's either a KeyError or a pl.exceptions.ColumnNotFoundError
+    with pytest.raises(Exception):
+        SelectCols(["a", "b", "d"]).fit(df)