Minor Fixes: Dask Extra, OneHotEncode example, categorize_pandas tests

ibisml/transforms/encode.py

@@ -11,6 +11,29 @@
 
 
 class OneHotEncode(Transform):
+    """
+    A transformation class that applies one-hot encoding to specified categorical columns.
+
+    Examples
+    ----------
+    >>> from ibis.interactive import *
+    >>> import ibisml as ml
+    >>> penguins = ex.penguins.fetch()
+    >>> recipe = ml.Recipe(ml.OneHotEncode("species"))
+    >>> recipe.fit(penguins).transform(penguins).table.select(s.startswith("species"))
+    ┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
+    ┃ species_Adelie ┃ species_Chinstrap ┃ species_Gentoo ┃
+    ┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
+    │ int8           │ int8              │ int8           │
+    ├────────────────┼───────────────────┼────────────────┤
+    │              1 │                 0 │              0 │
+    │              1 │                 0 │              0 │
+    │              1 │                 0 │              0 │
+    │              1 │                 0 │              0 │
+    │              … │                 … │              … │
+    └────────────────┴───────────────────┴────────────────┘
+    """
+
     def __init__(self, categories: dict[str, list[Any]]):
         self.categories = categories
 

pyproject.toml

@@ -13,6 +13,9 @@ dependencies = [
     "ibis-framework"
 ]
 
+[project.optional-dependencies]
+dask = ["dask[dataframe]"]
+
 [tool.setuptools]
 include-package-data = false
 

tests/test_core.py

@@ -1,6 +1,7 @@
 from textwrap import dedent
 
 import ibis
+import pandas as pd
 import pyarrow as pa
 
 import ibisml as ml
@@ -50,3 +51,41 @@ def test_transform_result_repr():
             d  int64
         }"""
     )
+
+
+def test_categorize_pandas():
+    categories = {
+        "num_col": ml.core.Categories(pa.array([1, 2, 3, 4, 5]), ordered=True),
+        "str_col": ml.core.Categories(pa.array(["a", "b", "c"]), ordered=False),
+    }
+
+    df = pd.DataFrame(
+        {
+            "num_col": [
+                0,
+                1,
+                2,
+                -1,
+                4,
+            ],
+            "str_col": [
+                0,
+                1,
+                2,
+                0,
+                -1,
+            ],
+        }
+    )
+
+    t = ibis.table({"num_col": "int64", "str_col": "string"})
+    res = ml.TransformResult(t, categories=categories)
+
+    transformed_df = res._categorize_pandas(df)
+
+    assert isinstance(transformed_df["num_col"].dtype, pd.CategoricalDtype)
+    assert isinstance(transformed_df["str_col"].dtype, pd.CategoricalDtype)
+    assert transformed_df["num_col"].cat.ordered is True
+    assert transformed_df["str_col"].cat.ordered is False
+    assert transformed_df["num_col"].isna().sum() == 1
+    assert transformed_df["str_col"].isna().sum() == 1