test(core): categorize_pandas

Author: IndexSeek

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