python make division consistent

polars/Cargo.toml

@@ -93,6 +93,7 @@ rank = ["polars-core/rank", "polars-lazy/rank"]
 diff = ["polars-core/diff", "polars-lazy/diff"]
 moment = ["polars-core/moment", "polars-lazy/moment"]
 arange = ["polars-lazy/arange"]
+true_div = ["polars-lazy/true_div"]
 
 # don't use this
 private = []

polars/polars-lazy/Cargo.toml

@@ -30,6 +30,7 @@ dtype-date64 = ["polars-core/dtype-date64"]
 # uncomment to have datafusion integration
 # when uncommenting we both need to point to the same arrow version
 #ooc = ["datafusion", "tokio"]
+true_div = []
 
 # operations
 is_in = ["polars-core/is_in"]

polars/polars-lazy/src/dsl.rs

@@ -458,6 +458,8 @@ pub enum Operator {
     Minus,
     Multiply,
     Divide,
+    #[cfg(feature = "true_div")]
+    TrueDivide,
     Modulus,
     And,
     Or,

polars/polars-lazy/src/logical_plan/optimizer/simplify_expr.rs

@@ -287,7 +287,7 @@ impl OptimizationRule for SimplifyExprRule {
                 Some(AExpr::Literal(LiteralValue::Null))
             }
 
-            // lit(left) + lit(right) => lit(left = right)
+            // lit(left) + lit(right) => lit(left + right)
             AExpr::BinaryExpr { left, op, right } => {
                 let left = expr_arena.get(*left);
                 let right = expr_arena.get(*right);
@@ -297,6 +297,51 @@ impl OptimizationRule for SimplifyExprRule {
                     Operator::Minus => eval_binary_same_type!(left, -, right),
                     Operator::Multiply => eval_binary_same_type!(left, *, right),
                     Operator::Divide => eval_binary_same_type!(left, /, right),
+                    #[cfg(feature = "true_div")]
+                    Operator::TrueDivide => {
+                        if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = (left, right)
+                        {
+                            return match (lit_left, lit_right) {
+                                (LiteralValue::Float32(x), LiteralValue::Float32(y)) => {
+                                    Some(AExpr::Literal(LiteralValue::Float32(x / y)))
+                                }
+                                (LiteralValue::Float64(x), LiteralValue::Float64(y)) => {
+                                    Some(AExpr::Literal(LiteralValue::Float64(x / y)))
+                                }
+                                #[cfg(feature = "dtype-i8")]
+                                (LiteralValue::Int8(x), LiteralValue::Int8(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                #[cfg(feature = "dtype-i16")]
+                                (LiteralValue::Int16(x), LiteralValue::Int16(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                (LiteralValue::Int32(x), LiteralValue::Int32(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                (LiteralValue::Int64(x), LiteralValue::Int64(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                #[cfg(feature = "dtype-u8")]
+                                (LiteralValue::UInt8(x), LiteralValue::UInt8(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                #[cfg(feature = "dtype-u16")]
+                                (LiteralValue::UInt16(x), LiteralValue::UInt16(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                (LiteralValue::UInt32(x), LiteralValue::UInt32(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                #[cfg(feature = "dtype-u64")]
+                                (LiteralValue::UInt64(x), LiteralValue::UInt64(y)) => Some(
+                                    AExpr::Literal(LiteralValue::Float64(*x as f64 / *y as f64)),
+                                ),
+                                _ => None,
+                            };
+                        }
+                        None
+                    }
                     Operator::Modulus => eval_binary_same_type!(left, %, right),
                     Operator::Lt => eval_binary_bool_type!(left, <, right),
                     Operator::Gt => eval_binary_bool_type!(left, >, right),

polars/polars-lazy/src/physical_plan/expressions/binary.rs

@@ -40,6 +40,14 @@ pub(crate) fn apply_operator(left: &Series, right: &Series, op: Operator) -> Res
         Operator::Minus => Ok(left - right),
         Operator::Multiply => Ok(left * right),
         Operator::Divide => Ok(left / right),
+        #[cfg(feature = "true_div")]
+        Operator::TrueDivide => {
+            use DataType::*;
+            match left.dtype() {
+                Date32 | Date64 | Time64(_) | Float32 | Float64 => Ok(left / right),
+                _ => Ok(&left.cast_with_dtype(&Float64)? / &right.cast_with_dtype(&Float64)?),
+            }
+        }
         Operator::And => Ok((left.bool()? & right.bool()?).into_series()),
         Operator::Or => Ok((left.bool()? | right.bool()?).into_series()),
         Operator::Modulus => Ok(left % right),

py-polars/Cargo.toml

@@ -72,7 +72,8 @@ features = [
     "rank",
     "diff",
     "moment",
-    "arange"
+    "arange",
+    "true_div"
 ]
 
 #[patch.crates-io]

py-polars/polars/lazy/expr.py

@@ -116,6 +116,12 @@ def __truediv__(self, other: Any) -> "Expr":
     def __rtruediv__(self, other: Any) -> "Expr":
         return wrap_expr(self.__to_pyexpr(other) / self._pyexpr)
 
+    def __floordiv__(self, other: Any) -> "Expr":
+        return wrap_expr(self._pyexpr // self.__to_pyexpr(other))
+
+    def __rfloordiv__(self, other: Any) -> "Expr":
+        return wrap_expr(self.__to_pyexpr(other) // self._pyexpr)
+
     def __mod__(self, other: Any) -> "Expr":
         return wrap_expr(self._pyexpr % self.__to_pyexpr(other))
 

py-polars/src/lazy/dsl.rs

@@ -52,11 +52,14 @@ impl PyNumberProtocol for PyExpr {
         Ok(dsl::binary_expr(lhs.inner, Operator::Multiply, rhs.inner).into())
     }
     fn __truediv__(lhs: Self, rhs: Self) -> PyResult<PyExpr> {
-        Ok(dsl::binary_expr(lhs.inner, Operator::Divide, rhs.inner).into())
+        Ok(dsl::binary_expr(lhs.inner, Operator::TrueDivide, rhs.inner).into())
     }
     fn __mod__(lhs: Self, rhs: Self) -> PyResult<PyExpr> {
         Ok(dsl::binary_expr(lhs.inner, Operator::Modulus, rhs.inner).into())
     }
+    fn __floordiv__(lhs: Self, rhs: Self) -> PyResult<PyExpr> {
+        Ok(dsl::binary_expr(lhs.inner, Operator::Divide, rhs.inner).into())
+    }
 }
 
 #[pyproto]

py-polars/tests/test_lazy.py

@@ -491,11 +491,11 @@ def test_arithmetic():
         [
             (col("a") % 2).alias("1"),
             (2 % col("a")).alias("2"),
-            (1 / col("a")).alias("3"),
+            (1 // col("a")).alias("3"),
             (1 * col("a")).alias("4"),
             (1 + col("a")).alias("5"),
             (1 - col("a")).alias("6"),
-            (col("a") / 2).alias("7"),
+            (col("a") // 2).alias("7"),
             (col("a") * 2).alias("8"),
             (col("a") + 2).alias("9"),
             (col("a") - 2).alias("10"),

py-polars/tests/test_series.py

@@ -567,3 +567,15 @@ def test_list_concat_dispatch():
     assert df.select(pl.col("a").arr.concat(["b"]).alias("concat"))[
         "concat"
     ].series_equal(expected)
+
+
+def test_floor_divide():
+    s = pl.Series("a", [1, 2, 3])
+    assert (s // 2).to_list() == [0, 1, 1]
+    assert pl.DataFrame([s]).select(pl.col("a") // 2)["a"].to_list() == [0, 1, 1]
+
+
+def test_true_divide():
+    s = pl.Series("a", [1, 2])
+    assert (s / 2).to_list() == [0.5, 1.0]
+    assert pl.DataFrame([s]).select(pl.col("a") / 2)["a"].to_list() == [0.5, 1.0]