feat: implement gated validation for extension types to prevent metadata propagation

Author: cyberbeam524

datafusion/expr/src/logical_plan/builder.rs

@@ -22,6 +22,7 @@ use std::cmp::Ordering;
 use std::collections::{HashMap, HashSet};
 use std::iter::once;
 use std::sync::Arc;
+use datafusion_common::metadata::check_metadata_with_storage_equal;
 
 use crate::dml::CopyTo;
 use crate::expr::{Alias, PlannedReplaceSelectItem, Sort as SortExpr};
@@ -274,28 +275,42 @@ impl LogicalPlanBuilder {
         for j in 0..n_cols {
             let field = schema.field(j);
             let field_type = field.data_type();
-            let field_nullable = field.is_nullable();
-            let field_metadata = FieldMetadata::new_from_field(field);
+            let field_metadata = field.metadata();
+            let is_target_ext = field_metadata.contains_key(arrow_schema::extension::EXTENSION_TYPE_NAME_KEY);
+            let column_name = format!("column {}", j);
             for row in values.iter() {
                 let value = &row[j];
-                let data_type = value.get_type(schema)?;
-
-                if !data_type.equals_datatype(field_type)
-                    && !can_cast_types(&data_type, field_type)
-                {
-                    return exec_err!(
-                        "type mismatch and can't cast to got {} and {}",
-                        data_type,
-                        field_type
-                    );
+                let value_type = value.get_type(schema)?;
+                if value_type == DataType::Null { continue; }
+
+                if is_target_ext {
+                    let value_meta = value.metadata(schema)?;
+                    let value_meta_map = value_meta.to_hashmap();
+                    let is_value_ext = value_meta_map.contains_key(arrow_schema::extension::EXTENSION_TYPE_NAME_KEY);
+
+                    if is_value_ext && value_type == *field_type {
+                        check_metadata_with_storage_equal(
+                            (&value_type, Some(&value_meta_map)),
+                            (field_type, Some(field_metadata)),
+                            &column_name,
+                            " in VALUES list",
+                        )?;
+                    } else if !can_cast_types(&value_type, field_type) {
+                        return plan_err!("Cannot cast {} to extension type at column {}", value_type, column_name);
+                    }
+                } else {
+                    // Optimized path for standard types
+                    if !value_type.equals_datatype(field_type) && !can_cast_types(&value_type, field_type) {
+                        return plan_err!("Cannot cast {} to {} for column {}", value_type, field_type, column_name);
+                    }
                 }
             }
-            let metadata = if field_metadata.is_empty() {
-                None
-            } else {
-                Some(field_metadata)
-            };
-            fields.push_with_metadata(field_type.to_owned(), field_nullable, metadata);
+
+            fields.push_with_metadata(
+                field_type.clone(), 
+                field.is_nullable(), 
+                Some(FieldMetadata::new_from_field(field))
+            );
         }
 
         Self::infer_inner(values, fields, schema)

datafusion/sql/src/statement.rs

@@ -43,6 +43,7 @@ use datafusion_expr::dml::{CopyTo, InsertOp};
 use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check;
 use datafusion_expr::logical_plan::DdlStatement;
 use datafusion_expr::logical_plan::builder::project;
+use datafusion_common::metadata::check_metadata_with_storage_equal;
 use datafusion_expr::utils::expr_to_columns;
 use datafusion_expr::{
     Analyze, Cast, CreateCatalog, CreateCatalogSchema,
@@ -530,23 +531,40 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                                 .fields()
                                 .iter()
                                 .zip(input_fields)
-                                .map(|(field, input_field)| {
-                                    let metadata =
-                                        FieldMetadata::new_from_field(field.as_ref());
-                                    let alias_metadata = if metadata.is_empty() {
-                                        None
+                                .map(|(field, input_field)| -> Result<Expr> {
+                                    let input_meta = input_field.metadata();
+                                    let target_meta = field.metadata();
+                                    
+                                    let is_target_ext = target_meta.contains_key(arrow_schema::extension::EXTENSION_TYPE_NAME_KEY);
+                                    let is_input_ext = input_meta.contains_key(arrow_schema::extension::EXTENSION_TYPE_NAME_KEY);
+
+                                    if !is_target_ext && !is_input_ext {
+                                        return if input_field.data_type() == field.data_type() {
+                                            Ok(col(input_field.name()).alias(field.name()))
+                                        } else {
+                                            Ok(col(input_field.name())
+                                                .cast_to(field.data_type(), input_schema)?
+                                                .alias(field.name()))
+                                        };
+                                    }
+
+                                    let metadata_matches = check_metadata_with_storage_equal(
+                                        (input_field.data_type(), Some(input_meta)),
+                                        (field.data_type(), Some(target_meta)),
+                                        "input", "target"
+                                    ).is_ok();
+
+                                    if metadata_matches {
+                                        Ok(col(input_field.name()).alias(field.name()))
                                     } else {
-                                        Some(metadata)
-                                    };
-                                    Expr::Cast(Cast::new_from_field(
-                                        Box::new(col(input_field.name())),
-                                        Arc::clone(field),
-                                    ))
-                                    .alias_with_metadata(field.name(), alias_metadata)
-                                })
-                                .collect::<Vec<_>>();
-
-                            LogicalPlanBuilder::from(plan.clone())
+                                        Ok(Expr::Cast(datafusion_expr::Cast::new_from_field(
+                                            Box::new(col(input_field.name())),
+                                            Arc::clone(field),
+                                        )).alias(field.name()))
+                                    }
+                                }).collect::<Result<Vec<_>>>()?;
+
+                                LogicalPlanBuilder::from(plan.clone())
                                 .project(project_exprs)?
                                 .build()?
                         } else {