Encapsulate metadata for literals on to a FieldMetadata structure (#16317)

* Encapsulate FieldMetadata

* Add examples

* refactor

* impl `Default` for `FieldMetadata`

Author: alamb

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -40,6 +40,7 @@ use datafusion_common::{
     assert_batches_eq, assert_batches_sorted_eq, assert_contains, exec_err, not_impl_err,
     plan_err, DFSchema, DataFusionError, Result, ScalarValue,
 };
+use datafusion_expr::expr::FieldMetadata;
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion_expr::{
     lit_with_metadata, Accumulator, ColumnarValue, CreateFunction, CreateFunctionBody,
@@ -1535,7 +1536,7 @@ async fn test_metadata_based_udf_with_literal() -> Result<()> {
     let df = ctx.sql("select 0;").await?.select(vec![
         lit(5u64).alias_with_metadata("lit_with_doubling", Some(input_metadata.clone())),
         lit(5u64).alias("lit_no_doubling"),
-        lit_with_metadata(5u64, Some(input_metadata))
+        lit_with_metadata(5u64, Some(FieldMetadata::from(input_metadata)))
             .alias("lit_with_double_no_alias_metadata"),
     ])?;
 

datafusion/expr/src/expr.rs

@@ -30,7 +30,7 @@ use crate::logical_plan::Subquery;
 use crate::Volatility;
 use crate::{udaf, ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
 
-use arrow::datatypes::{DataType, FieldRef};
+use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::cse::{HashNode, NormalizeEq, Normalizeable};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion,
@@ -284,8 +284,8 @@ pub enum Expr {
     Column(Column),
     /// A named reference to a variable in a registry.
     ScalarVariable(DataType, Vec<String>),
-    /// A constant value along with associated metadata
-    Literal(ScalarValue, Option<BTreeMap<String, String>>),
+    /// A constant value along with associated [`FieldMetadata`].
+    Literal(ScalarValue, Option<FieldMetadata>),
     /// A binary expression such as "age > 21"
     BinaryExpr(BinaryExpr),
     /// LIKE expression
@@ -413,6 +413,168 @@ impl<'a> TreeNodeContainer<'a, Self> for Expr {
     }
 }
 
+/// Literal metadata
+///
+/// Stores metadata associated with a literal expressions
+/// and is designed to be fast to `clone`.
+///
+/// This structure is used to store metadata associated with a literal expression, and it
+/// corresponds to the `metadata` field on [`Field`].
+///
+/// # Example: Create [`FieldMetadata`] from a [`Field`]
+/// ```
+/// # use std::collections::HashMap;
+/// # use datafusion_expr::expr::FieldMetadata;
+/// # use arrow::datatypes::{Field, DataType};
+/// # let field = Field::new("c1", DataType::Int32, true)
+/// #  .with_metadata(HashMap::from([("foo".to_string(), "bar".to_string())]));
+/// // Create a new `FieldMetadata` instance from a `Field`
+/// let metadata = FieldMetadata::new_from_field(&field);
+/// // There is also a `From` impl:
+/// let metadata = FieldMetadata::from(&field);
+/// ```
+///
+/// # Example: Update a [`Field`] with [`FieldMetadata`]
+/// ```
+/// # use datafusion_expr::expr::FieldMetadata;
+/// # use arrow::datatypes::{Field, DataType};
+/// # let field = Field::new("c1", DataType::Int32, true);
+/// # let metadata = FieldMetadata::new_from_field(&field);
+/// // Add any metadata from `FieldMetadata` to `Field`
+/// let updated_field = metadata.add_to_field(field);
+/// ```
+///
+#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct FieldMetadata {
+    /// The inner metadata of a literal expression, which is a map of string
+    /// keys to string values.
+    ///
+    /// Note this is not a `HashMap because `HashMap` does not provide
+    /// implementations for traits like `Debug` and `Hash`.
+    inner: Arc<BTreeMap<String, String>>,
+}
+
+impl Default for FieldMetadata {
+    fn default() -> Self {
+        Self::new_empty()
+    }
+}
+
+impl FieldMetadata {
+    /// Create a new empty metadata instance.
+    pub fn new_empty() -> Self {
+        Self {
+            inner: Arc::new(BTreeMap::new()),
+        }
+    }
+
+    /// Merges two optional `FieldMetadata` instances, overwriting any existing
+    /// keys in `m` with keys from `n` if present
+    pub fn merge_options(
+        m: Option<&FieldMetadata>,
+        n: Option<&FieldMetadata>,
+    ) -> Option<FieldMetadata> {
+        match (m, n) {
+            (Some(m), Some(n)) => {
+                let mut merged = m.clone();
+                merged.extend(n.clone());
+                Some(merged)
+            }
+            (Some(m), None) => Some(m.clone()),
+            (None, Some(n)) => Some(n.clone()),
+            (None, None) => None,
+        }
+    }
+
+    /// Create a new metadata instance from a `Field`'s metadata.
+    pub fn new_from_field(field: &Field) -> Self {
+        let inner = field
+            .metadata()
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect();
+        Self {
+            inner: Arc::new(inner),
+        }
+    }
+
+    /// Create a new metadata instance from a map of string keys to string values.
+    pub fn new(inner: BTreeMap<String, String>) -> Self {
+        Self {
+            inner: Arc::new(inner),
+        }
+    }
+
+    /// Get the inner metadata as a reference to a `BTreeMap`.
+    pub fn inner(&self) -> &BTreeMap<String, String> {
+        &self.inner
+    }
+
+    /// Return the inner metadata
+    pub fn into_inner(self) -> Arc<BTreeMap<String, String>> {
+        self.inner
+    }
+
+    /// Adds metadata from `other` into `self`, overwriting any existing keys.
+    pub fn extend(&mut self, other: Self) {
+        let other = Arc::unwrap_or_clone(other.into_inner());
+        Arc::make_mut(&mut self.inner).extend(other);
+    }
+
+    /// Returns true if the metadata is empty.
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    /// Returns the number of key-value pairs in the metadata.
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    /// Updates the metadata on the Field with this metadata, if it is not empty.
+    pub fn add_to_field(&self, field: Field) -> Field {
+        if self.inner.is_empty() {
+            return field;
+        }
+
+        field.with_metadata(
+            self.inner
+                .iter()
+                .map(|(k, v)| (k.clone(), v.clone()))
+                .collect(),
+        )
+    }
+}
+
+impl From<&Field> for FieldMetadata {
+    fn from(field: &Field) -> Self {
+        Self::new_from_field(field)
+    }
+}
+
+impl From<BTreeMap<String, String>> for FieldMetadata {
+    fn from(inner: BTreeMap<String, String>) -> Self {
+        Self::new(inner)
+    }
+}
+
+impl From<std::collections::HashMap<String, String>> for FieldMetadata {
+    fn from(map: std::collections::HashMap<String, String>) -> Self {
+        Self::new(map.into_iter().collect())
+    }
+}
+
+/// From reference
+impl From<&std::collections::HashMap<String, String>> for FieldMetadata {
+    fn from(map: &std::collections::HashMap<String, String>) -> Self {
+        let inner = map
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect();
+        Self::new(inner)
+    }
+}
+
 /// UNNEST expression.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
 pub struct Unnest {

datafusion/expr/src/expr_rewriter/mod.rs

@@ -390,11 +390,7 @@ mod test {
                     } else {
                         utf8_val
                     };
-                    Ok(Transformed::yes(lit_with_metadata(
-                        utf8_val,
-                        metadata
-                            .map(|m| m.into_iter().collect::<HashMap<String, String>>()),
-                    )))
+                    Ok(Transformed::yes(lit_with_metadata(utf8_val, metadata)))
                 }
                 // otherwise, return None
                 _ => Ok(Transformed::no(expr)),

datafusion/expr/src/expr_schema.rs

@@ -423,12 +423,7 @@ impl ExprSchemable for Expr {
             Expr::Literal(l, metadata) => {
                 let mut field = Field::new(&schema_name, l.data_type(), l.is_null());
                 if let Some(metadata) = metadata {
-                    field = field.with_metadata(
-                        metadata
-                            .iter()
-                            .map(|(k, v)| (k.clone(), v.clone()))
-                            .collect(),
-                    );
+                    field = metadata.add_to_field(field);
                 }
                 Ok(Arc::new(field))
             }

datafusion/expr/src/literal.rs

@@ -17,34 +17,29 @@
 
 //! Literal module contains foundational types that are used to represent literals in DataFusion.
 
+use crate::expr::FieldMetadata;
 use crate::Expr;
 use datafusion_common::ScalarValue;
-use std::collections::HashMap;
 
 /// Create a literal expression
 pub fn lit<T: Literal>(n: T) -> Expr {
     n.lit()
 }
 
-pub fn lit_with_metadata<T: Literal>(
-    n: T,
-    metadata: impl Into<Option<HashMap<String, String>>>,
-) -> Expr {
-    let metadata = metadata.into();
+pub fn lit_with_metadata<T: Literal>(n: T, metadata: Option<FieldMetadata>) -> Expr {
     let Some(metadata) = metadata else {
         return n.lit();
     };
 
     let Expr::Literal(sv, prior_metadata) = n.lit() else {
         unreachable!();
     };
-
     let new_metadata = match prior_metadata {
         Some(mut prior) => {
             prior.extend(metadata);
             prior
         }
-        None => metadata.into_iter().collect(),
+        None => metadata,
     };
 
     Expr::Literal(sv, Some(new_metadata))

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -18,7 +18,7 @@
 //! Expression simplification API
 
 use std::borrow::Cow;
-use std::collections::{BTreeMap, HashSet};
+use std::collections::HashSet;
 use std::ops::Not;
 
 use arrow::{
@@ -58,6 +58,7 @@ use crate::{
     analyzer::type_coercion::TypeCoercionRewriter,
     simplify_expressions::unwrap_cast::try_cast_literal_to_type,
 };
+use datafusion_expr::expr::FieldMetadata;
 use indexmap::IndexSet;
 use regex::Regex;
 
@@ -523,9 +524,9 @@ struct ConstEvaluator<'a> {
 #[allow(clippy::large_enum_variant)]
 enum ConstSimplifyResult {
     // Expr was simplified and contains the new expression
-    Simplified(ScalarValue, Option<BTreeMap<String, String>>),
+    Simplified(ScalarValue, Option<FieldMetadata>),
     // Expr was not simplified and original value is returned
-    NotSimplified(ScalarValue, Option<BTreeMap<String, String>>),
+    NotSimplified(ScalarValue, Option<FieldMetadata>),
     // Evaluation encountered an error, contains the original expression
     SimplifyRuntimeError(DataFusionError, Expr),
 }
@@ -682,9 +683,7 @@ impl<'a> ConstEvaluator<'a> {
                 let m = f.metadata();
                 match m.is_empty() {
                     true => None,
-                    false => {
-                        Some(m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
-                    }
+                    false => Some(FieldMetadata::from(m)),
                 }
             });
         let col_val = match phys_expr.evaluate(&self.input_batch) {

datafusion/physical-expr/src/expressions/literal.rs

@@ -18,7 +18,6 @@
 //! Literal expressions for physical operations
 
 use std::any::Any;
-use std::collections::HashMap;
 use std::hash::Hash;
 use std::sync::Arc;
 
@@ -30,6 +29,7 @@ use arrow::{
     record_batch::RecordBatch,
 };
 use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::expr::FieldMetadata;
 use datafusion_expr::Expr;
 use datafusion_expr_common::columnar_value::ColumnarValue;
 use datafusion_expr_common::interval_arithmetic::Interval;
@@ -64,14 +64,13 @@ impl Literal {
     /// Create a literal value expression
     pub fn new_with_metadata(
         value: ScalarValue,
-        metadata: impl Into<Option<HashMap<String, String>>>,
+        metadata: Option<FieldMetadata>,
     ) -> Self {
-        let metadata = metadata.into();
         let mut field =
             Field::new(format!("{value}"), value.data_type(), value.is_null());
 
         if let Some(metadata) = metadata {
-            field = field.with_metadata(metadata);
+            field = metadata.add_to_field(field);
         }
 
         Self {