fix[rust, python]: make cut work on integer types (#4945)

pola-rs · Sep 23, 2022 · bec22c4 · bec22c4
1 parent a3e772b
commit bec22c4
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git a/polars/polars-algo/src/algo.rs b/polars/polars-algo/src/algo.rs
@@ -56,14 +56,17 @@ pub fn cut(
         .with_columns([col(category_str).cast(DataType::Categorical(None))])
         .collect()?;
 
-    s.sort(false).into_frame().join_asof(
-        &cuts,
-        var_name,
-        breakpoint_str,
-        AsofStrategy::Forward,
-        None,
-        None,
-    )
+    s.cast(&DataType::Float64)?
+        .sort(false)
+        .into_frame()
+        .join_asof(
+            &cuts,
+            var_name,
+            breakpoint_str,
+            AsofStrategy::Forward,
+            None,
+            None,
+        )
 }
 
 #[test]

diff --git a/py-polars/polars/internals/functions.py b/py-polars/polars/internals/functions.py
@@ -369,7 +369,8 @@ def cut(
     cuts_df = cuts_df.with_column(pli.col(category_label).cast(Categorical))
 
     result = (
-        s.sort()
+        s.cast(Float64)
+        .sort()
         .to_frame()
         .join_asof(
             cuts_df,

diff --git a/py-polars/tests/unit/test_functions.py b/py-polars/tests/unit/test_functions.py
@@ -107,6 +107,11 @@ def test_cut() -> None:
         ],
     }
 
+    # test cut on integers #4939
+    df = pl.DataFrame({"a": list(range(5))})
+    ser = df.select("a").to_series()
+    assert pl.cut(ser, bins=[-1, 1]).shape == (5, 3)
+
 
 def test_null_handling_correlation() -> None:
     df = pl.DataFrame({"a": [1, 2, 3, None, 4], "b": [1, 2, 3, 10, 4]})