[python] zip with and some ergonomic changes

py-polars/pypolars/frame.py

@@ -265,6 +265,10 @@ def __getitem__(self, item):
     def __setitem__(self, key, value):
         # df["foo"] = series
         if isinstance(key, str):
+            try:
+                self.drop_in_place(key)
+            except:
+                pass
             self.hstack([Series(key, value)])
         # df[idx] = series
         elif isinstance(key, int):

py-polars/pypolars/lazy.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import os
-from typing import Union, List
+from typing import Union, List, Callable
 from pypolars.frame import DataFrame, wrap_df
 from .pypolars import PyLazyFrame, col, lit, binary_expr, PyExpr, PyLazyGroupBy, when
 
@@ -34,6 +34,21 @@ def from_pyldf(ldf: "PyLazyFrame") -> "LazyFrame":
         self._ldf = ldf
         return self
 
+    def pipe(self, func: Callable, *args, **kwargs):
+        """
+        Apply a function on Self
+
+        Parameters
+        ----------
+        func
+            Callable
+        args
+            Arguments
+        kwargs
+            Keyword arguments
+        """
+        return func(self, *args, **kwargs)
+
     def describe_plan(self) -> str:
         return self._ldf.describe_plan()
 
@@ -43,7 +58,15 @@ def describe_optimized_plan(self) -> str:
     def sort(self, by_column: str) -> LazyFrame:
         return wrap_ldf(self._ldf.sort(by_column))
 
-    def collect(self) -> DataFrame:
+    def collect(
+        self,
+        type_coercion: bool = True,
+        predicate_pushdown: bool = True,
+        projection_pushdown: bool = True,
+    ) -> DataFrame:
+        self._ldf = self._ldf.optimization_toggle(
+            type_coercion, predicate_pushdown, projection_pushdown
+        )
         return wrap_df(self._ldf.collect())
 
     def filter(self, predicate: PyExpr) -> LazyFrame:

py-polars/pypolars/series.py

@@ -542,6 +542,12 @@ def len(self) -> int:
     def __len__(self):
         return self.len()
 
+    def cast_utf8(self):
+        """
+        Cast this Series
+        """
+        return wrap_s(self._s.cast_utf8())
+
     def cast_u8(self):
         """
         Cast this Series
@@ -813,3 +819,20 @@ def shift(self, periods: int) -> Series:
             Number of places to shift (may be negative).
         """
         return wrap_s(self._s.shift(periods))
+
+    def zip_with(self, mask: Series, other: Series) -> Series:
+        """
+        Where mask evaluates true take values from self. Where mask evaluates false, take values from other.
+
+        Parameters
+        ----------
+        mask
+            Boolean Series
+        other
+            Series of same type
+
+        Returns
+        -------
+        New Series
+        """
+        return wrap_s(self._s.zip_with(mask._s, other._s))

py-polars/src/series.rs

@@ -478,6 +478,15 @@ impl PySeries {
         let s = self.series.shift(periods).map_err(PyPolarsEr::from)?;
         Ok(PySeries::new(s))
     }
+
+    pub fn zip_with(&self, mask: &PySeries, other: &PySeries) -> PyResult<Self> {
+        let mask = mask.series.bool().map_err(PyPolarsEr::from)?;
+        let s = self
+            .series
+            .zip_with(mask, &other.series)
+            .map_err(PyPolarsEr::from)?;
+        Ok(PySeries::new(s))
+    }
 }
 
 macro_rules! impl_ufuncs {
@@ -677,6 +686,7 @@ impl_cast!(cast_date32, Date32Type);
 impl_cast!(cast_date64, Date64Type);
 impl_cast!(cast_time64ns, Time64NanosecondType);
 impl_cast!(cast_duration_ns, DurationNanosecondType);
+impl_cast!(cast_utf8, Utf8Type);
 
 macro_rules! impl_arithmetic {
     ($name:ident, $type:ty, $operand:tt) => {