correct dtype for power (#4246)

pola-rs · Aug 4, 2022 · 52120c6 · 52120c6
1 parent 99ed0d8
commit 52120c6
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/dsl/function_expr/mod.rs b/polars/polars-lazy/src/dsl/function_expr/mod.rs
@@ -105,11 +105,20 @@ impl FunctionExpr {
         };
 
         let same_type = || map_dtype(&|dtype| dtype.clone());
+        let super_type = || {
+            let mut first = fields[0].clone();
+            let mut st = first.data_type().clone();
+            for field in &fields[1..] {
+                st = get_supertype(&st, field.data_type())?
+            }
+            first.coerce(st);
+            Ok(first)
+        };
 
         use FunctionExpr::*;
         match self {
             NullCount => with_dtype(IDX_DTYPE),
-            Pow => float_dtype(),
+            Pow => super_type(),
             #[cfg(feature = "row_hash")]
             Hash(..) => with_dtype(DataType::UInt64),
             #[cfg(feature = "is_in")]

diff --git a/py-polars/tests/test_schema.py b/py-polars/tests/test_schema.py
@@ -36,3 +36,19 @@ def test_fill_null_minimal_upcast_4056() -> None:
 def test_with_column_duplicates() -> None:
     df = pl.DataFrame({"a": [0, None, 2, 3, None], "b": [None, 1, 2, 3, None]})
     assert df.with_columns([pl.all().alias("same")]).columns == ["a", "b", "same"]
+
+
+def test_pow_dtype() -> None:
+    df = pl.DataFrame(
+        {
+            "foo": [1, 2, 3, 4, 5],
+        }
+    ).lazy()
+
+    df = df.with_columns([pl.col("foo").cast(pl.UInt32)]).with_columns(
+        [
+            (pl.col("foo") * 2**2).alias("scaled_foo"),
+            (pl.col("foo") * 2**2.1).alias("scaled_foo2"),
+        ]
+    )
+    assert df.collect().dtypes == [pl.UInt32, pl.UInt32, pl.Float64]