This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new c9f15547 feat: alias with metadata (#1111)
c9f15547 is described below
commit c9f15547cb8019068bbf2dc8eaf148d6eb42bd48
Author: Chen Chongchen <[email protected]>
AuthorDate: Fri Apr 25 21:01:54 2025 +0800
feat: alias with metadata (#1111)
* feat: alias with metadata
* fmt
---
python/datafusion/expr.py | 14 +++++++++++---
python/datafusion/functions.py | 15 ++++++++++++---
python/tests/test_expr.py | 5 +++++
python/tests/test_functions.py | 5 +++++
src/expr.rs | 6 ++++--
src/functions.rs | 9 +++++++--
6 files changed, 44 insertions(+), 10 deletions(-)
diff --git a/python/datafusion/expr.py b/python/datafusion/expr.py
index 2697d814..01e1f3de 100644
--- a/python/datafusion/expr.py
+++ b/python/datafusion/expr.py
@@ -406,9 +406,17 @@ class Expr:
"""Creates a new expression representing a column."""
return Expr(expr_internal.RawExpr.column(value))
- def alias(self, name: str) -> Expr:
- """Assign a name to the expression."""
- return Expr(self.expr.alias(name))
+ def alias(self, name: str, metadata: Optional[dict[str, str]] = None) ->
Expr:
+ """Assign a name to the expression.
+
+ Args:
+ name: The name to assign to the expression.
+ metadata: Optional metadata to attach to the expression.
+
+ Returns:
+ A new expression with the assigned name.
+ """
+ return Expr(self.expr.alias(name, metadata))
def sort(self, ascending: bool = True, nulls_first: bool = True) ->
SortExpr:
"""Creates a sort :py:class:`Expr` from an existing :py:class:`Expr`.
diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py
index 5cf914e1..f430cdf4 100644
--- a/python/datafusion/functions.py
+++ b/python/datafusion/functions.py
@@ -372,9 +372,18 @@ def order_by(expr: Expr, ascending: bool = True,
nulls_first: bool = True) -> So
return SortExpr(expr, ascending=ascending, nulls_first=nulls_first)
-def alias(expr: Expr, name: str) -> Expr:
- """Creates an alias expression."""
- return Expr(f.alias(expr.expr, name))
+def alias(expr: Expr, name: str, metadata: Optional[dict[str, str]] = None) ->
Expr:
+ """Creates an alias expression with an optional metadata dictionary.
+
+ Args:
+ expr: The expression to alias
+ name: The alias name
+ metadata: Optional metadata to attach to the column
+
+ Returns:
+ An expression with the given alias
+ """
+ return Expr(f.alias(expr.expr, name, metadata))
def col(name: str) -> Expr:
diff --git a/python/tests/test_expr.py b/python/tests/test_expr.py
index 926e6984..dcf75f02 100644
--- a/python/tests/test_expr.py
+++ b/python/tests/test_expr.py
@@ -247,3 +247,8 @@ def test_fill_null(df):
assert result.column(0) == pa.array([1, 2, 100])
assert result.column(1) == pa.array([4, 25, 6])
assert result.column(2) == pa.array([1234, 1234, 8])
+
+
+def test_alias_with_metadata(df):
+ df = df.select(col("a").alias("b", {"key": "value"}))
+ assert df.schema().field("b").metadata == {b"key": b"value"}
diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py
index 37f2075f..90cf01f7 100644
--- a/python/tests/test_functions.py
+++ b/python/tests/test_functions.py
@@ -1231,3 +1231,8 @@ def test_between_default(df):
actual = df.collect()[0].to_pydict()
assert actual == expected
+
+
+def test_alias_with_metadata(df):
+ df = df.select(f.alias(f.col("a"), "b", {"key": "value"}))
+ assert df.schema().field("b").metadata == {b"key": b"value"}
diff --git a/src/expr.rs b/src/expr.rs
index fe0e76da..7d4aa879 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -22,6 +22,7 @@ use datafusion::logical_expr::{
};
use pyo3::IntoPyObjectExt;
use pyo3::{basic::CompareOp, prelude::*};
+use std::collections::HashMap;
use std::convert::{From, Into};
use std::sync::Arc;
use window::PyWindowFrame;
@@ -275,8 +276,9 @@ impl PyExpr {
}
/// assign a name to the PyExpr
- pub fn alias(&self, name: &str) -> PyExpr {
- self.expr.clone().alias(name).into()
+ #[pyo3(signature = (name, metadata=None))]
+ pub fn alias(&self, name: &str, metadata: Option<HashMap<String, String>>)
-> PyExpr {
+ self.expr.clone().alias_with_metadata(name, metadata).into()
}
/// Create a sort PyExpr from an existing PyExpr.
diff --git a/src/functions.rs b/src/functions.rs
index 476c2b80..caa79b8a 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.
+use std::collections::HashMap;
+
use datafusion::functions_aggregate::all_default_aggregate_functions;
use datafusion::functions_window::all_default_window_functions;
use datafusion::logical_expr::expr::WindowFunctionParams;
@@ -205,10 +207,13 @@ fn order_by(expr: PyExpr, asc: bool, nulls_first: bool)
-> PyResult<PySortExpr>
/// Creates a new Alias Expr
#[pyfunction]
-fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
+#[pyo3(signature = (expr, name, metadata=None))]
+fn alias(expr: PyExpr, name: &str, metadata: Option<HashMap<String, String>>)
-> PyResult<PyExpr> {
let relation: Option<TableReference> = None;
Ok(PyExpr {
- expr: datafusion::logical_expr::Expr::Alias(Alias::new(expr.expr,
relation, name)),
+ expr: datafusion::logical_expr::Expr::Alias(
+ Alias::new(expr.expr, relation, name).with_metadata(metadata),
+ ),
})
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
