From 96907feab9dea1eb0bfa800058fc290ca08e4d88 Mon Sep 17 00:00:00 2001 From: francis-du Date: Wed, 7 Sep 2022 15:34:21 +0800 Subject: [PATCH] feat: add with_column_renamed funcation for dataframe --- 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 bed0a91a6..f4d1a5fc5 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 80963f7f0..b07e6d77b 100644 --- a/src/dataframe.rs +++ b/src/dataframe.rs @@ -92,6 +92,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 { + let df = self.df.with_column_renamed(old_name, new_name)?; + Ok(Self::new(df)) + } + fn aggregate(&self, group_by: Vec, aggs: Vec) -> PyResult { let group_by = group_by.into_iter().map(|e| e.into()).collect(); let aggs = aggs.into_iter().map(|e| e.into()).collect();