refactor: Remove dead code & consolidate with_column (#4787)

py-polars/polars/internals/dataframe/frame.py

@@ -4016,15 +4016,11 @@ def with_column(self, column: pli.Series | pli.Expr) -> DataFrame:
         └──────┴─────┘
 
         """
-        if isinstance(column, list):
-            raise ValueError(
-                "`with_column` expects a single expression, not a list. Consider using"
-                " `with_columns`"
-            )
-        if isinstance(column, pli.Expr):
-            return self.with_columns([column])
-        else:
-            return self._from_pydf(self._df.with_column(column._s))
+        return (
+            self.lazy()
+            .with_column(column)
+            .collect(no_optimization=True, string_cache=False)
+        )
 
     def hstack(
         self: DF,

py-polars/polars/internals/lazyframe/frame.py

@@ -1678,14 +1678,14 @@ def with_context(self, other: LDF | list[LDF]) -> LDF:
 
         return self._from_pyldf(self._ldf.with_context([lf._ldf for lf in other]))
 
-    def with_column(self: LDF, expr: pli.Expr) -> LDF:
+    def with_column(self: LDF, column: pli.Series | pli.Expr) -> LDF:
         """
         Add or overwrite column in a DataFrame.
 
         Parameters
         ----------
-        expr
-            Expression that evaluates to column.
+        column
+            Expression that evaluates to column or a Series to use.
 
         Examples
         --------
@@ -1723,7 +1723,12 @@ def with_column(self: LDF, expr: pli.Expr) -> LDF:
         └──────┴─────┘
 
         """
-        return self.with_columns([expr])
+        if not isinstance(column, (pli.Expr, pli.Series)):
+            raise TypeError(
+                "`with_column` expects a single Expr or Series. "
+                "Consider using `with_columns` if you need multiple columns."
+            )
+        return self.with_columns([column])
 
     def drop(self: LDF, columns: str | list[str]) -> LDF:
         """

py-polars/src/dataframe.rs

@@ -18,7 +18,6 @@ use polars_core::prelude::QuantileInterpolOptions;
 use polars_core::utils::arrow::compute::cast::CastOptions;
 use polars_core::utils::try_get_supertype;
 use polars_lazy::frame::pivot::{pivot, pivot_stable};
-use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyTuple};
 
@@ -819,12 +818,6 @@ impl PyDataFrame {
         Ok(())
     }
 
-    pub fn with_column(&mut self, s: PySeries) -> PyResult<Self> {
-        let mut df = self.df.clone();
-        df.with_column(s.series).map_err(PyPolarsErr::from)?;
-        Ok(df.into())
-    }
-
     /// Get datatypes
     pub fn dtypes(&self, py: Python) -> PyObject {
         let iter = self
@@ -918,16 +911,6 @@ impl PyDataFrame {
         Ok(PyDataFrame::new(df))
     }
 
-    pub fn filter(&self, mask: &PySeries) -> PyResult<Self> {
-        let filter_series = &mask.series;
-        if let Ok(ca) = filter_series.bool() {
-            let df = self.df.filter(ca).map_err(PyPolarsErr::from)?;
-            Ok(PyDataFrame::new(df))
-        } else {
-            Err(PyRuntimeError::new_err("Expected a boolean mask"))
-        }
-    }
-
     pub fn take(&self, indices: Wrap<Vec<IdxSize>>) -> PyResult<Self> {
         let indices = indices.0;
         let indices = IdxCa::from_vec("", indices);
@@ -962,11 +945,6 @@ impl PyDataFrame {
         Ok(())
     }
 
-    pub fn rename(&mut self, column: &str, new_col: &str) -> PyResult<()> {
-        self.df.rename(column, new_col).map_err(PyPolarsErr::from)?;
-        Ok(())
-    }
-
     pub fn replace_at_idx(&mut self, index: usize, new_col: PySeries) -> PyResult<()> {
         self.df
             .replace_at_idx(index, new_col.series)

py-polars/tests/unit/test_df.py

@@ -961,7 +961,7 @@ def test_literal_series() -> None:
     )
     out = (
         df.lazy()
-        .with_column(pl.Series("e", [2, 1, 3], pl.Int32))  # type: ignore[arg-type]
+        .with_column(pl.Series("e", [2, 1, 3], pl.Int32))
         .with_column(pl.col("e").cast(pl.Float32))
         .collect()
     )