[arrow-datafusion-python] branch master updated: feat: add with_column_renamed funcation for dataframe (#47)

[agrove]

This is an automated email from the ASF dual-hosted git repository.

agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-python.git


The following commit(s) were added to refs/heads/master by this push:
     new b390d08  feat: add with_column_renamed funcation for dataframe (#47)
b390d08 is described below

commit b390d08641c0ee64a3042b9401349a1f569876c2
Author: Francis Du <fran...@francisdu.com>
AuthorDate: Thu Sep 8 11:33:59 2022 +0800

    feat: add with_column_renamed funcation for dataframe (#47)
---
 datafusion/tests/test_dataframe.py | 12 ++++++++++++
 src/dataframe.rs                   |  7 +++++++
 2 files changed, 19 insertions(+)

diff --git a/datafusion/tests/test_dataframe.py 
b/datafusion/tests/test_dataframe.py
index e739465..c9544ab 100644
--- a/datafusion/tests/test_dataframe.py
+++ b/datafusion/tests/test_dataframe.py
@@ -118,6 +118,18 @@ 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"
+    )
+
+    result = df.collect()[0]
+
+    assert result.schema.field(0).name == "a"
+    assert result.schema.field(1).name == "b"
+    assert result.schema.field(2).name == "sum"
+
+
 def test_udf(df):
     # is_null is a pa function over arrays
     is_null = udf(
diff --git a/src/dataframe.rs b/src/dataframe.rs
index f9f2f66..f6cb4f1 100644
--- a/src/dataframe.rs
+++ b/src/dataframe.rs
@@ -93,6 +93,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
+    /// renamed does not exist.
+    fn with_column_renamed(&self, old_name: &str, new_name: &str) -> 
PyResult<Self> {
+        let df = self.df.with_column_renamed(old_name, new_name)?;
+        Ok(Self::new(df))
+    }
+
     fn aggregate(&self, group_by: Vec<PyExpr>, aggs: Vec<PyExpr>) -> 
PyResult<Self> {
         let group_by = group_by.into_iter().map(|e| e.into()).collect();
         let aggs = aggs.into_iter().map(|e| e.into()).collect();