Clean up display and comparison of types with metadata

datafusion/common/src/datatype.rs

@@ -17,8 +17,11 @@
 
 //! [DataTypeExt] extension trait for converting DataTypes to Fields
 
-use crate::arrow::datatypes::{DataType, Field, FieldRef};
-use std::sync::Arc;
+use crate::{
+    arrow::datatypes::{DataType, Field, FieldRef},
+    metadata::FieldMetadata,
+};
+use std::{fmt::Display, sync::Arc};
 
 /// DataFusion extension methods for Arrow [`DataType`]
 pub trait DataTypeExt {
@@ -105,3 +108,150 @@ impl FieldExt for Arc<Field> {
         }
     }
 }
+
+#[derive(Debug, Clone, Copy, Eq)]
+pub struct SerializedTypeView<'a, 'b, 'c> {
+    arrow_type: &'a DataType,
+    extension_name: Option<&'b str>,
+    extension_metadata: Option<&'c str>,
+}
+
+impl<'a, 'b, 'c> SerializedTypeView<'a, 'b, 'c> {
+    pub fn new(
+        arrow_type: &'a DataType,
+        extension_name: Option<&'b str>,
+        extension_metadata: Option<&'c str>,
+    ) -> Self {
+        Self {
+            arrow_type,
+            extension_name,
+            extension_metadata,
+        }
+    }
+
+    pub fn from_type_and_metadata(
+        arrow_type: &'a DataType,
+        metadata: impl IntoIterator<Item = (&'b String, &'b String)>,
+    ) -> Self
+    where
+        'b: 'c,
+    {
+        let mut extension_name = None;
+        let mut extension_metadata = None;
+        for (k, v) in metadata {
+            match k.as_str() {
+                "ARROW:extension:name" => extension_name.replace(v.as_str()),
+                "ARROW:extension:metadata" => extension_metadata.replace(v.as_str()),
+                _ => None,
+            };
+        }
+
+        Self::new(arrow_type, extension_name, extension_metadata)
+    }
+
+    pub fn data_type(&self) -> Option<&DataType> {
+        if self.extension_name.is_some() {
+            None
+        } else {
+            Some(self.arrow_type)
+        }
+    }
+
+    pub fn arrow_type(&self) -> &DataType {
+        self.arrow_type
+    }
+
+    pub fn extension_name(&self) -> Option<&str> {
+        self.extension_name
+    }
+
+    pub fn extension_metadata(&self) -> Option<&str> {
+        if let Some(metadata) = self.extension_metadata {
+            if !metadata.is_empty() {
+                return Some(metadata);
+            }
+        }
+
+        None
+    }
+
+    pub fn to_field(&self) -> Field {
+        if let Some(extension_name) = self.extension_name() {
+            self.arrow_type.clone().into_nullable_field().with_metadata(
+                [
+                    (
+                        "ARROW:extension:name".to_string(),
+                        extension_name.to_string(),
+                    ),
+                    (
+                        "ARROW:extension:metadata".to_string(),
+                        self.extension_metadata().unwrap_or("").to_string(),
+                    ),
+                ]
+                .into(),
+            )
+        } else {
+            self.arrow_type.clone().into_nullable_field()
+        }
+    }
+
+    pub fn to_field_ref(&self) -> FieldRef {
+        self.to_field().into()
+    }
+}
+
+impl Display for SerializedTypeView<'_, '_, '_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match (self.extension_name(), self.extension_metadata()) {
+            (Some(name), None) => write!(f, "{}<{}>", name, self.arrow_type()),
+            (Some(name), Some(metadata)) => {
+                write!(f, "{}({})<{}>", name, metadata, self.arrow_type())
+            }
+            _ => write!(f, "{}", self.arrow_type()),
+        }
+    }
+}
+
+impl PartialEq<SerializedTypeView<'_, '_, '_>> for SerializedTypeView<'_, '_, '_> {
+    fn eq(&self, other: &SerializedTypeView) -> bool {
+        self.arrow_type() == other.arrow_type()
+            && self.extension_name() == other.extension_name()
+            && self.extension_metadata() == other.extension_metadata()
+    }
+}
+
+impl<'a> From<&'a DataType> for SerializedTypeView<'a, 'static, 'static> {
+    fn from(value: &'a DataType) -> Self {
+        Self::new(value, None, None)
+    }
+}
+
+impl<'a> From<&'a Field> for SerializedTypeView<'a, 'a, 'a> {
+    fn from(value: &'a Field) -> Self {
+        Self::from_type_and_metadata(value.data_type(), value.metadata())
+    }
+}
+
+impl<'a> From<&'a FieldRef> for SerializedTypeView<'a, 'a, 'a> {
+    fn from(value: &'a FieldRef) -> Self {
+        Self::from_type_and_metadata(value.data_type(), value.metadata())
+    }
+}
+
+impl<'a, 'b> From<(&'a DataType, &'b FieldMetadata)> for SerializedTypeView<'a, 'b, 'b> {
+    fn from(value: (&'a DataType, &'b FieldMetadata)) -> Self {
+        Self::from_type_and_metadata(value.0, value.1.inner())
+    }
+}
+
+impl<'a, 'b> From<(&'a DataType, Option<&'b FieldMetadata>)>
+    for SerializedTypeView<'a, 'b, 'b>
+{
+    fn from(value: (&'a DataType, Option<&'b FieldMetadata>)) -> Self {
+        if let Some(metadata) = value.1 {
+            Self::from_type_and_metadata(value.0, metadata.inner())
+        } else {
+            value.0.into()
+        }
+    }
+}

datafusion/common/src/metadata.rs

@@ -15,15 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::{collections::BTreeMap, fmt::Display, hash::Hash, sync::Arc};
 
 use arrow::datatypes::{DataType, Field};
 use hashbrown::HashMap;
 
-use crate::{error::_plan_err, DataFusionError, ScalarValue};
+use crate::{datatype::SerializedTypeView, DataFusionError, ScalarValue};
 
 /// A [`ScalarValue`] with optional [`FieldMetadata`]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Eq, PartialOrd)]
 pub struct ScalarAndMetadata {
     pub value: ScalarValue,
     pub metadata: Option<FieldMetadata>,
@@ -64,72 +64,74 @@ impl ScalarAndMetadata {
     }
 }
 
-/// Assert equality of data types where one or both sides may have field metadata
-///
-/// This currently compares absent metadata (e.g., one side was a DataType) and
-/// empty metadata (e.g., one side was a field where the field had no metadata)
-/// as equal and uses byte-for-byte comparison for the keys and values of the
-/// fields, even though this is potentially too strict for some cases (e.g.,
-/// extension types where extension metadata is represented by JSON, or cases
-/// where field metadata is orthogonal to the interpretation of the data type).
-///
-/// Returns a planning error with suitably formatted type representations if
-/// actual and expected do not compare to equal.
-pub fn check_metadata_with_storage_equal(
-    actual: (
-        &DataType,
-        Option<&std::collections::HashMap<String, String>>,
-    ),
-    expected: (
-        &DataType,
-        Option<&std::collections::HashMap<String, String>>,
-    ),
-    what: &str,
-    context: &str,
-) -> Result<(), DataFusionError> {
-    if actual.0 != expected.0 {
-        return _plan_err!(
-            "Expected {what} of type {}, got {}{context}",
-            format_type_and_metadata(expected.0, expected.1),
-            format_type_and_metadata(actual.0, actual.1)
-        );
-    }
+impl Display for ScalarAndMetadata {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let storage_type = self.value.data_type();
+        let serialized_type = SerializedTypeView::from((&storage_type, self.metadata()));
+
+        let metadata_without_extension_info = self
+            .metadata
+            .as_ref()
+            .map(|metadata| {
+                metadata
+                    .inner()
+                    .into_iter()
+                    .filter(|(k, _)| {
+                        *k != "ARROW:extension:name" && *k != "ARROW:extension:metadata"
+                    })
+                    .map(|(k, v)| (k.to_string(), v.to_string()))
+                    .collect::<BTreeMap<_, _>>()
+            })
+            .unwrap_or(BTreeMap::new());
+
+        match (
+            serialized_type.extension_name(),
+            serialized_type.extension_metadata(),
+        ) {
+            (Some(name), None) => write!(f, "{name}<{:?}>", self.value())?,
+            (Some(name), Some(metadata)) => {
+                write!(f, "{name}({metadata})<{:?}>", self.value())?
+            }
+            _ => write!(f, "{:?}", self.value())?,
+        }
 
-    let metadata_equal = match (actual.1, expected.1) {
-        (None, None) => true,
-        (None, Some(expected_metadata)) => expected_metadata.is_empty(),
-        (Some(actual_metadata), None) => actual_metadata.is_empty(),
-        (Some(actual_metadata), Some(expected_metadata)) => {
-            actual_metadata == expected_metadata
+        if !metadata_without_extension_info.is_empty() {
+            write!(
+                f,
+                " {:?}",
+                FieldMetadata::new(metadata_without_extension_info)
+            )
+        } else {
+            Ok(())
         }
-    };
-
-    if !metadata_equal {
-        return _plan_err!(
-            "Expected {what} of type {}, got {}{context}",
-            format_type_and_metadata(expected.0, expected.1),
-            format_type_and_metadata(actual.0, actual.1)
-        );
     }
+}
+
+impl From<ScalarValue> for ScalarAndMetadata {
+    fn from(value: ScalarValue) -> Self {
+        ScalarAndMetadata::new(value, None)
+    }
+}
 
-    Ok(())
+impl PartialEq<ScalarAndMetadata> for ScalarAndMetadata {
+    fn eq(&self, other: &ScalarAndMetadata) -> bool {
+        let metadata_equal = match (self.metadata(), other.metadata()) {
+            (None, None) => true,
+            (None, Some(other_meta)) => other_meta.is_empty(),
+            (Some(meta), None) => meta.is_empty(),
+            (Some(meta), Some(other_meta)) => meta == other_meta,
+        };
+
+        metadata_equal && self.value() == other.value()
+    }
 }
 
-/// Given a data type represented by storage and optional metadata, generate
-/// a user-facing string
-///
-/// This function exists to reduce the number of Field debug strings that are
-/// used to communicate type information in error messages and plan explain
-/// renderings.
-pub fn format_type_and_metadata(
-    data_type: &DataType,
-    metadata: Option<&std::collections::HashMap<String, String>>,
-) -> String {
-    match metadata {
-        Some(metadata) if !metadata.is_empty() => {
-            format!("{data_type}<{metadata:?}>")
+impl Hash for ScalarAndMetadata {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.value.hash(state);
+        if let Some(metadata) = self.metadata() {
+            metadata.hash(state);
         }
-        _ => data_type.to_string(),
     }
 }
 

datafusion/common/src/param_value.rs

@@ -15,8 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::datatype::SerializedTypeView;
 use crate::error::{_plan_datafusion_err, _plan_err};
-use crate::metadata::{check_metadata_with_storage_equal, ScalarAndMetadata};
+use crate::metadata::ScalarAndMetadata;
 use crate::{Result, ScalarValue};
 use arrow::datatypes::{DataType, Field, FieldRef};
 use std::collections::HashMap;
@@ -61,15 +62,15 @@ impl ParamValues {
                 // Verify if the types of the params matches the types of the values
                 let iter = expect.iter().zip(list.iter());
                 for (i, (param_type, lit)) in iter.enumerate() {
-                    check_metadata_with_storage_equal(
-                        (
-                            &lit.value.data_type(),
-                            lit.metadata.as_ref().map(|m| m.to_hashmap()).as_ref(),
-                        ),
-                        (param_type.data_type(), Some(param_type.metadata())),
-                        "parameter",
-                        &format!(" at index {i}"),
-                    )?;
+                    let actual_storage = lit.value().data_type();
+                    let actual_type =
+                        SerializedTypeView::from((&actual_storage, lit.metadata()));
+                    let expected_type = SerializedTypeView::from(param_type);
+                    if actual_type != expected_type {
+                        return _plan_err!(
+                            "Expected parameter of type {expected_type}, got {actual_type} at index {i}"
+                        );
+                    }
                 }
                 Ok(())
             }

datafusion/expr/src/expr.rs

@@ -32,6 +32,7 @@ use crate::{ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
 
 use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::cse::{HashNode, NormalizeEq, Normalizeable};
+use datafusion_common::metadata::ScalarAndMetadata;
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion,
 };
@@ -3256,10 +3257,7 @@ impl Display for Expr {
             }
             Expr::ScalarVariable(_, var_names) => write!(f, "{}", var_names.join(".")),
             Expr::Literal(v, metadata) => {
-                match metadata.as_ref().map(|m| m.is_empty()).unwrap_or(true) {
-                    false => write!(f, "{v:?} {:?}", metadata.as_ref().unwrap()),
-                    true => write!(f, "{v:?}"),
-                }
+                write!(f, "{}", ScalarAndMetadata::new(v.clone(), metadata.clone()))
             }
             Expr::Case(case) => {
                 write!(f, "CASE ")?;