fix[python] floor divide floats (#4420)

pola-rs · Aug 14, 2022 · 67f4fda · 67f4fda
1 parent 2a7d835
commit 67f4fda
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 9 deletions.
diff --git a/polars/polars-lazy/src/dsl/expr.rs b/polars/polars-lazy/src/dsl/expr.rs
@@ -408,6 +408,7 @@ pub enum Operator {
     Multiply,
     Divide,
     TrueDivide,
+    FloorDivide,
     Modulus,
     And,
     Or,

diff --git a/polars/polars-lazy/src/logical_plan/format.rs b/polars/polars-lazy/src/logical_plan/format.rs
@@ -317,8 +317,9 @@ impl Debug for Operator {
             Plus => "+",
             Minus => "-",
             Multiply => "*",
-            Divide => "//",
+            Divide => "/",
             TrueDivide => "/",
+            FloorDivide => "//",
             Modulus => "%",
             And => "&",
             Or => "|",

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/simplify_expr.rs b/polars/polars-lazy/src/logical_plan/optimizer/simplify_expr.rs
@@ -348,6 +348,7 @@ impl OptimizationRule for SimplifyExprRule {
                         }
                         None
                     }
+                    Operator::FloorDivide => None,
                     Operator::Modulus => eval_binary_same_type!(left_aexpr, %, right_aexpr),
                     Operator::Lt => eval_binary_bool_type!(left_aexpr, <, right_aexpr),
                     Operator::Gt => eval_binary_bool_type!(left_aexpr, >, right_aexpr),

diff --git a/polars/polars-lazy/src/physical_plan/expressions/binary.rs b/polars/polars-lazy/src/physical_plan/expressions/binary.rs
@@ -46,13 +46,8 @@ fn apply_operator_owned(left: Series, right: Series, op: Operator) -> Result<Ser
         Operator::Minus => Ok(left - right),
         Operator::Multiply => Ok(left * right),
         Operator::Divide => Ok(&left / &right),
-        Operator::TrueDivide => {
-            use DataType::*;
-            match left.dtype() {
-                Date | Datetime(_, _) | Float32 | Float64 => Ok(&left / &right),
-                _ => Ok(&left.cast(&Float64)? / &right.cast(&Float64)?),
-            }
-        }
+        Operator::TrueDivide => apply_operator(&left, &right, op),
+        Operator::FloorDivide => apply_operator(&left, &right, op),
         Operator::And => left.bitand(&right),
         Operator::Or => left.bitor(&right),
         Operator::Xor => left.bitxor(&right),
@@ -81,6 +76,14 @@ pub fn apply_operator(left: &Series, right: &Series, op: Operator) -> Result<Ser
                 _ => Ok(&left.cast(&Float64)? / &right.cast(&Float64)?),
             }
         }
+        Operator::FloorDivide => {
+            use DataType::*;
+            match left.dtype() {
+                #[cfg(feature = "round_series")]
+                Float32 | Float64 => (left / right).floor(),
+                _ => Ok(left / right),
+            }
+        }
         Operator::And => left.bitand(right),
         Operator::Or => left.bitor(right),
         Operator::Xor => left.bitxor(right),

diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -66,7 +66,7 @@ impl PyExpr {
         Ok(dsl::binary_expr(self.inner.clone(), Operator::Modulus, rhs.inner).into())
     }
     fn __floordiv__(&self, rhs: Self) -> PyResult<PyExpr> {
-        Ok(dsl::binary_expr(self.inner.clone(), Operator::Divide, rhs.inner).into())
+        Ok(dsl::binary_expr(self.inner.clone(), Operator::FloorDivide, rhs.inner).into())
     }
 
     pub fn to_str(&self) -> String {

diff --git a/py-polars/tests/test_lazy.py b/py-polars/tests/test_lazy.py
@@ -860,6 +860,12 @@ def test_arithmetic() -> None:
     )
     assert out.frame_equal(expected)
 
+    # floating point floor divide
+    x = 10.4
+    step = 0.5
+    df = pl.DataFrame({"x": [x]})
+    assert df.with_columns(pl.col("x") // step)[0, 0] == x // step
+
 
 def test_ufunc() -> None:
     df = pl.DataFrame({"a": [1, 2]})