feat: support customize metadata in alias for dataframe api (#15120)

* feat: support customize metadata in alias for dataframe api

* update doc

* remove clone

Author: chenkovsky

datafusion/core/src/dataframe/mod.rs

@@ -2039,6 +2039,7 @@ impl DataFrame {
                             })),
                             relation: None,
                             name: field.name().to_string(),
+                            metadata: None,
                         }),
                         Err(_) => col(field.name()),
                     }

datafusion/core/tests/dataframe/mod.rs

@@ -5374,6 +5374,21 @@ async fn test_alias() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_alias_with_metadata() -> Result<()> {
+    let mut metadata = HashMap::new();
+    metadata.insert(String::from("k"), String::from("v"));
+    let df = create_test_table("test")
+        .await?
+        .select(vec![col("a").alias_with_metadata("b", Some(metadata))])?
+        .alias("table_alias")?;
+    let df = df.select(vec![col("table_alias.b")])?;
+    let schema = df.schema();
+    let metadata = schema.field(0).metadata();
+    assert_eq!(metadata.get("k"), Some(&String::from("v")));
+    Ok(())
+}
+
 // Use alias to perform a self-join
 // Issue: https://github.com/apache/datafusion/issues/14112
 #[tokio::test]

datafusion/expr/src/expr.rs

@@ -391,11 +391,34 @@ impl Unnest {
 }
 
 /// Alias expression
-#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
+#[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Alias {
     pub expr: Box<Expr>,
     pub relation: Option<TableReference>,
     pub name: String,
+    pub metadata: Option<std::collections::HashMap<String, String>>,
+}
+
+impl Hash for Alias {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.expr.hash(state);
+        self.relation.hash(state);
+        self.name.hash(state);
+    }
+}
+
+impl PartialOrd for Alias {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        let cmp = self.expr.partial_cmp(&other.expr);
+        let Some(std::cmp::Ordering::Equal) = cmp else {
+            return cmp;
+        };
+        let cmp = self.relation.partial_cmp(&other.relation);
+        let Some(std::cmp::Ordering::Equal) = cmp else {
+            return cmp;
+        };
+        self.name.partial_cmp(&other.name)
+    }
 }
 
 impl Alias {
@@ -409,8 +432,17 @@ impl Alias {
             expr: Box::new(expr),
             relation: relation.map(|r| r.into()),
             name: name.into(),
+            metadata: None,
         }
     }
+
+    pub fn with_metadata(
+        mut self,
+        metadata: Option<std::collections::HashMap<String, String>>,
+    ) -> Self {
+        self.metadata = metadata;
+        self
+    }
 }
 
 /// Binary expression
@@ -1278,6 +1310,27 @@ impl Expr {
         Expr::Alias(Alias::new(self, None::<&str>, name.into()))
     }
 
+    /// Return `self AS name` alias expression with metadata
+    ///
+    /// The metadata will be attached to the Arrow Schema field when the expression
+    /// is converted to a field via `Expr.to_field()`.
+    ///
+    /// # Example
+    /// ```
+    /// # use datafusion_expr::col;
+    /// use std::collections::HashMap;
+    /// let metadata = HashMap::from([("key".to_string(), "value".to_string())]);
+    /// let expr = col("foo").alias_with_metadata("bar", Some(metadata));
+    /// ```
+    ///
+    pub fn alias_with_metadata(
+        self,
+        name: impl Into<String>,
+        metadata: Option<std::collections::HashMap<String, String>>,
+    ) -> Expr {
+        Expr::Alias(Alias::new(self, None::<&str>, name.into()).with_metadata(metadata))
+    }
+
     /// Return `self AS name` alias expression with a specific qualifier
     pub fn alias_qualified(
         self,
@@ -1287,6 +1340,28 @@ impl Expr {
         Expr::Alias(Alias::new(self, relation, name.into()))
     }
 
+    /// Return `self AS name` alias expression with a specific qualifier and metadata
+    ///
+    /// The metadata will be attached to the Arrow Schema field when the expression
+    /// is converted to a field via `Expr.to_field()`.
+    ///
+    /// # Example
+    /// ```
+    /// # use datafusion_expr::col;
+    /// use std::collections::HashMap;
+    /// let metadata = HashMap::from([("key".to_string(), "value".to_string())]);
+    /// let expr = col("foo").alias_qualified_with_metadata(Some("tbl"), "bar", Some(metadata));
+    /// ```
+    ///
+    pub fn alias_qualified_with_metadata(
+        self,
+        relation: Option<impl Into<TableReference>>,
+        name: impl Into<String>,
+        metadata: Option<std::collections::HashMap<String, String>>,
+    ) -> Expr {
+        Expr::Alias(Alias::new(self, relation, name.into()).with_metadata(metadata))
+    }
+
     /// Remove an alias from an expression if one exists.
     ///
     /// If the expression is not an alias, the expression is returned unchanged.
@@ -1738,11 +1813,13 @@ impl NormalizeEq for Expr {
                     expr: self_expr,
                     relation: self_relation,
                     name: self_name,
+                    ..
                 }),
                 Expr::Alias(Alias {
                     expr: other_expr,
                     relation: other_relation,
                     name: other_name,
+                    ..
                 }),
             ) => {
                 self_name == other_name
@@ -2088,6 +2165,7 @@ impl HashNode for Expr {
                 expr: _expr,
                 relation,
                 name,
+                ..
             }) => {
                 relation.hash(state);
                 name.hash(state);

datafusion/expr/src/expr_schema.rs

@@ -343,7 +343,16 @@ impl ExprSchemable for Expr {
     fn metadata(&self, schema: &dyn ExprSchema) -> Result<HashMap<String, String>> {
         match self {
             Expr::Column(c) => Ok(schema.metadata(c)?.clone()),
-            Expr::Alias(Alias { expr, .. }) => expr.metadata(schema),
+            Expr::Alias(Alias { expr, metadata, .. }) => {
+                let mut ret = expr.metadata(schema)?;
+                if let Some(metadata) = metadata {
+                    if !metadata.is_empty() {
+                        ret.extend(metadata.clone());
+                        return Ok(ret);
+                    }
+                }
+                Ok(ret)
+            }
             Expr::Cast(Cast { expr, .. }) => expr.metadata(schema),
             _ => Ok(HashMap::new()),
         }

datafusion/expr/src/tree_node.rs

@@ -132,7 +132,10 @@ impl TreeNode for Expr {
                 expr,
                 relation,
                 name,
-            }) => f(*expr)?.update_data(|e| e.alias_qualified(relation, name)),
+                metadata,
+            }) => f(*expr)?.update_data(|e| {
+                e.alias_qualified_with_metadata(relation, name, metadata)
+            }),
             Expr::InSubquery(InSubquery {
                 expr,
                 subquery,

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -492,8 +492,11 @@ fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Project
                 expr,
                 relation,
                 name,
+                metadata,
             }) => rewrite_expr(*expr, &prev_projection).map(|result| {
-                result.update_data(|expr| Expr::Alias(Alias::new(expr, relation, name)))
+                result.update_data(|expr| {
+                    Expr::Alias(Alias::new(expr, relation, name).with_metadata(metadata))
+                })
             }),
             e => rewrite_expr(e, &prev_projection),
         }

datafusion/proto/proto/datafusion.proto

@@ -484,6 +484,7 @@ message AliasNode {
   LogicalExprNode expr = 1;
   string alias = 2;
   repeated TableReference relation = 3;
+  map<string, string> metadata = 4;
 }
 
 message BinaryExprNode {

datafusion/proto/src/generated/pbjson.rs

@@ -19,6 +19,8 @@
 //! DataFusion logical plans to be serialized and transmitted between
 //! processes.
 
+use std::collections::HashMap;
+
 use datafusion_common::{TableReference, UnnestOptions};
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::expr::{
@@ -200,6 +202,7 @@ pub fn serialize_expr(
             expr,
             relation,
             name,
+            metadata,
         }) => {
             let alias = Box::new(protobuf::AliasNode {
                 expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)),
@@ -208,6 +211,7 @@ pub fn serialize_expr(
                     .map(|r| vec![r.into()])
                     .unwrap_or(vec![]),
                 alias: name.to_owned(),
+                metadata: metadata.to_owned().unwrap_or(HashMap::new()),
             });
             protobuf::LogicalExprNode {
                 expr_type: Some(ExprType::Alias(alias)),

datafusion/proto/src/generated/prost.rs

@@ -363,6 +363,7 @@ pub(super) fn inject_column_aliases(
                 expr: Box::new(expr.clone()),
                 relation,
                 name: col_alias.value,
+                metadata: None,
             })
         })
         .collect::<Vec<_>>();