feat: alias with metadata (#1111)

* feat: alias with metadata

* fmt

Author: chenkovsky

python/datafusion/expr.py

@@ -406,9 +406,17 @@ def column(value: str) -> 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`.

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:

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"}

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"}

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.

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),
+        ),
     })
 }