fix add string name (#2222)

pola-rs · Dec 29, 2021 · 27008e3 · 27008e3
1 parent 315e304
commit 27008e3
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/polars/polars-core/src/chunked_array/arithmetic.rs b/polars/polars-core/src/chunked_array/arithmetic.rs
@@ -18,7 +18,7 @@ macro_rules! apply_operand_on_chunkedarray_by_iter {
                         .into_no_null_iter()
                         .zip($rhs.into_no_null_iter())
                         .map(|(left, right)| left $operand right)
-                        .collect();
+                        .collect_trusted();
                         a.into_inner()
                     },
                     (false, _) => {
@@ -33,7 +33,7 @@ macro_rules! apply_operand_on_chunkedarray_by_iter {
                         .into_iter()
                         .zip($rhs.into_no_null_iter())
                         .map(|(opt_left, right)| opt_left.map(|left| left $operand right))
-                        .collect()
+                        .collect_trusted()
                     },
                     (_, _) => {
                     $self.into_iter()
@@ -363,13 +363,16 @@ impl Add for &Utf8Chunked {
         }
 
         // todo! add no_null variants. Need 4 paths.
-        self.into_iter()
+        let mut ca: Self::Output = self
+            .into_iter()
             .zip(rhs.into_iter())
             .map(|(opt_l, opt_r)| match (opt_l, opt_r) {
                 (Some(l), Some(r)) => Some(concat_strings(l, r)),
                 _ => None,
             })
-            .collect()
+            .collect_trusted();
+        ca.rename(self.name());
+        ca
     }
 }
 
@@ -385,16 +388,18 @@ impl Add<&str> for &Utf8Chunked {
     type Output = Utf8Chunked;
 
     fn add(self, rhs: &str) -> Self::Output {
-        match self.has_validity() {
+        let mut ca: Self::Output = match self.has_validity() {
             false => self
                 .into_no_null_iter()
                 .map(|l| concat_strings(l, rhs))
-                .collect(),
+                .collect_trusted(),
             _ => self
                 .into_iter()
                 .map(|opt_l| opt_l.map(|l| concat_strings(l, rhs)))
-                .collect(),
-        }
+                .collect_trusted(),
+        };
+        ca.rename(self.name());
+        ca
     }
 }
 

diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -3841,7 +3841,7 @@ def fold(
         ... )
         >>> df.fold(lambda s1, s2: s1 + s2)
         shape: (3,)
-        Series: '' [str]
+        Series: 'a' [str]
         [
             "foo11.0"
             "bar22.0"

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1648,7 +1648,6 @@ def test_arithmetic() -> None:
         _ = df + [1]  # type: ignore
 
 
-@pytest.mark.skip("column names of result are empty, see issue #2217")
 def test_add_string() -> None:
     df = pl.DataFrame({"a": ["hi", "there"], "b": ["hello", "world"]})
     result = df + " hello"