fix(rust, python): fix owned arithmetic schema (#5685)

pola-rs · Nov 30, 2022 · 62e8b0f · 62e8b0f
1 parent 02463f6
commit 62e8b0f
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 4 deletions.
diff --git a/polars/polars-core/src/chunked_array/arithmetic.rs b/polars/polars-core/src/chunked_array/arithmetic.rs
@@ -136,8 +136,9 @@ where
             let opt_lhs = lhs.get(0);
             match opt_lhs {
                 None => ChunkedArray::full_null(lhs.name(), rhs.len()),
-                Some(lhs) => {
-                    rhs.apply_mut(|rhs| operation(lhs, rhs));
+                Some(lhs_val) => {
+                    rhs.apply_mut(|rhs| operation(lhs_val, rhs));
+                    rhs.rename(lhs.name());
                     rhs
                 }
             }

diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
@@ -425,7 +425,7 @@ fn test_lazy_query_9() -> PolarsResult<()> {
 ))]
 fn test_lazy_query_10() {
     use polars_core::export::chrono::Duration as ChronoDuration;
-    let date = NaiveDate::from_ymd(2021, 3, 5);
+    let date = NaiveDate::from_ymd_opt(2021, 3, 5).unwrap();
     let x: Series = DatetimeChunked::from_naive_datetime(
         "x",
         [
@@ -502,7 +502,7 @@ fn test_lazy_query_10() {
     feature = "dtype-datetime"
 ))]
 fn test_lazy_query_7() {
-    let date = NaiveDate::from_ymd(2021, 3, 5);
+    let date = NaiveDate::from_ymd_opt(2021, 3, 5).unwrap();
     let dates = [
         NaiveDateTime::new(date, NaiveTime::from_hms_opt(12, 0, 0).unwrap()),
         NaiveDateTime::new(date, NaiveTime::from_hms_opt(12, 1, 0).unwrap()),

diff --git a/py-polars/tests/unit/test_schema.py b/py-polars/tests/unit/test_schema.py
@@ -267,3 +267,14 @@ def test_boolean_agg_schema() -> None:
             == agg_df.schema
             == {"x": pl.Int64, "max_y": pl.Boolean}
         )
+
+
+def test_schema_owned_arithmetic_5669() -> None:
+    df = (
+        pl.DataFrame({"A": [1, 2, 3]})
+        .lazy()
+        .filter(pl.col("A") >= 3)
+        .with_column(-pl.col("A").alias("B"))
+        .collect()
+    )
+    assert df.columns == ["A", "literal"], df.columns