diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py index c5ac0bb89..aa01240fc 100644 --- a/python/datafusion/dataframe.py +++ b/python/datafusion/dataframe.py @@ -175,7 +175,23 @@ def with_column_renamed(self, old_name: str, new_name: str) -> DataFrame: Returns: DataFrame with the column renamed. """ - return DataFrame(self.df.with_column_renamed(old_name, new_name)) + return DataFrame(self.df.rename({old_name: new_name})) + + def rename(self, mapping: dict[str, str]) -> DataFrame: + r"""Rename one or multiple columns by applying a new projection. + + This is a no-op if the column to be renamed does not exist. + + The method supports case sensitive rename with wrapping column name + into one the following symbols (" or ' or \`). + + Args: + mapping: mapping of old (key) to new (value) names + + Returns: + DataFrame with one or multiple columns renamed. + """ + return DataFrame(self.df.rename(mapping)) def aggregate( self, group_by: list[Expr] | Expr, aggs: list[Expr] | Expr diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py index e89c57159..71d8764bc 100644 --- a/python/tests/test_dataframe.py +++ b/python/tests/test_dataframe.py @@ -205,8 +205,8 @@ def test_with_column(df): assert result.column(2) == pa.array([5, 7, 9]) -def test_with_column_renamed(df): - df = df.with_column("c", column("a") + column("b")).with_column_renamed("c", "sum") +def test_rename(df): + df = df.with_column("c", column("a") + column("b")).rename({"c": "sum"}) result = df.collect()[0] diff --git a/src/dataframe.rs b/src/dataframe.rs index e77ca8425..c5a76e6db 100644 --- a/src/dataframe.rs +++ b/src/dataframe.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::ffi::CString; use std::sync::Arc; @@ -180,14 +181,13 @@ impl PyDataFrame { Ok(Self::new(df)) } - /// Rename one column by applying a new projection. This is a no-op if the column to be + /// Rename single or multiple columns by applying a new projection. This is a no-op if the column to be /// renamed does not exist. - fn with_column_renamed(&self, old_name: &str, new_name: &str) -> PyResult { - let df = self - .df - .as_ref() - .clone() - .with_column_renamed(old_name, new_name)?; + fn rename(&self, mapping: HashMap) -> PyResult { + let mut df = self.df.as_ref().clone(); + for (old_name, new_name) in mapping.iter() { + df = df.with_column_renamed(old_name, new_name)? + } Ok(Self::new(df)) }