Feat/map sql extension types

datafusion/expr/src/logical_plan/builder.rs

@@ -50,14 +50,13 @@ use crate::{
 };
 
 use super::dml::InsertOp;
-use arrow::compute::can_cast_types;
 use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::metadata::FieldMetadata;
 use datafusion_common::{
     Column, Constraints, DFSchema, DFSchemaRef, FunctionalDependencies, NullEquality,
-    Result, ScalarValue, TableReference, ToDFSchema, UnnestOptions, exec_err,
+    Result, ScalarValue, TableReference, ToDFSchema, UnnestOptions,
     get_target_functional_dependencies, internal_datafusion_err, plan_datafusion_err,
     plan_err,
 };
@@ -275,23 +274,14 @@ impl LogicalPlanBuilder {
         let n_cols = values[0].len();
         let mut fields = ValuesFields::new();
         for j in 0..n_cols {
-            let field_type = schema.field(j).data_type();
-            let field_nullable = schema.field(j).is_nullable();
-            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!(
-                        "Types don't match and no valid cast exists, received data of type {} for field of type {}",
-                        data_type,
-                        field_type
-                    );
-                }
-            }
-            fields.push(field_type.to_owned(), field_nullable);
+            let field = schema.field(j);
+            let field_type = field.data_type();
+
+            fields.push_with_metadata(
+                field_type.clone(),
+                field.is_nullable(),
+                Some(FieldMetadata::new_from_field(field)),
+            );
         }
 
         Self::infer_inner(values, fields, schema)
@@ -1682,10 +1672,6 @@ impl ValuesFields {
         Self::default()
     }
 
-    pub fn push(&mut self, data_type: DataType, nullable: bool) {
-        self.push_with_metadata(data_type, nullable, None);
-    }
-
     pub fn push_with_metadata(
         &mut self,
         data_type: DataType,
@@ -3370,8 +3356,7 @@ mod tests {
 
         assert_eq!(
             err.strip_backtrace(),
-            "Execution error: Types don't match and no valid cast exists, \
-         received data of type Boolean for field of type Date32"
+            "Error during planning: Cannot automatically convert Boolean to Date32"
         );
     }
 }

datafusion/sql/Cargo.toml

@@ -52,6 +52,7 @@ recursive_protection = ["dep:recursive", "dep:stacker"]
 # They are used for testing purposes only, so they are in the dev-dependencies section.
 [dependencies]
 arrow = { workspace = true }
+arrow-schema = { workspace = true }
 bigdecimal = { workspace = true }
 chrono = { workspace = true }
 datafusion-common = { workspace = true, features = ["sql"] }

datafusion/sql/src/planner.rs

@@ -23,10 +23,12 @@ use std::vec;
 
 use crate::utils::make_decimal_type;
 use arrow::datatypes::*;
+use arrow_schema::extension::{ExtensionType, Json, Uuid};
 use datafusion_common::TableReference;
 use datafusion_common::config::SqlParserOptions;
 use datafusion_common::datatype::{DataTypeExt, FieldExt};
 use datafusion_common::error::add_possible_columns_to_diag;
+use datafusion_common::types::{DFExtensionType, DFJson, DFUuid};
 use datafusion_common::{DFSchema, DataFusionError, Result, not_impl_err, plan_err};
 use datafusion_common::{
     DFSchemaRef, Diagnostic, SchemaError, field_not_found, internal_err,
@@ -677,6 +679,31 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
 
         // If no type_planner can handle this type, use the default conversion
         match sql_type {
+            // Canonical Arrow extension types
+            SQLDataType::Uuid => {
+                let data_type = DataType::FixedSizeBinary(16);
+                let df_uuid = DFUuid::try_new(&data_type, Default::default())?;
+
+                // Re-construct the Arrow Uuid using the validated data from df_uuid
+                let arrow_uuid = Uuid::try_new(&df_uuid.storage_type().clone(), ())
+                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+                Ok(Arc::new(
+                    Field::new("", data_type, true).with_extension_type(arrow_uuid),
+                ))
+            }
+            SQLDataType::JSON => {
+                let data_type = DataType::Utf8;
+                let df_json = DFJson::try_new(&data_type, Default::default())?;
+
+                let arrow_json =
+                    Json::try_new(&df_json.storage_type().clone(), Default::default())
+                        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+
+                Ok(Arc::new(
+                    Field::new("", data_type, true).with_extension_type(arrow_json),
+                ))
+            }
             SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) => {
                 // Arrays may be multi-dimensional.
                 Ok(self.convert_data_type_to_field(inner_sql_type)?.into_list())

datafusion/sql/tests/cases/params.rs

@@ -835,11 +835,11 @@ fn test_insert_infer_with_metadata() {
         @r#"
     ** Initial Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
+      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
         Values: ($1, $2, $3)
     ** Final Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
+      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
         Values: (FixedSizeBinary(16, "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16") FieldMetadata { inner: {"ARROW:extension:name": "arrow.uuid"} } AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
     "#
     );
@@ -857,11 +857,11 @@ fn test_insert_infer_with_metadata() {
     ** Initial Plan:
     Prepare: "my_plan" [FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>, Utf8, Utf8]
       Dml: op=[Insert Into] table=[person_with_uuid_extension]
-        Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
+        Projection: column1 AS id, column2 AS first_name, column3 AS last_name
           Values: ($1, $2, $3)
     ** Final Plan:
     Dml: op=[Insert Into] table=[person_with_uuid_extension]
-      Projection: CAST(column1 AS FixedSizeBinary(16)<{"ARROW:extension:name": "arrow.uuid"}>) AS id, column2 AS first_name, column3 AS last_name
+      Projection: column1 AS id, column2 AS first_name, column3 AS last_name
         Values: (FixedSizeBinary(16, "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16") FieldMetadata { inner: {"ARROW:extension:name": "arrow.uuid"} } AS $1, Utf8("Alan") AS $2, Utf8("Turing") AS $3)
     "#
     );

datafusion/sql/tests/common/mod.rs

@@ -400,11 +400,6 @@ impl TypePlanner for CustomTypePlanner {
         sql_type: &sqlparser::ast::DataType,
     ) -> Result<Option<FieldRef>> {
         match sql_type {
-            sqlparser::ast::DataType::Uuid => Ok(Some(Arc::new(
-                Field::new("", DataType::FixedSizeBinary(16), true).with_metadata(
-                    Metadata::new().with("ARROW:extension:name", "arrow.uuid"),
-                ),
-            ))),
             sqlparser::ast::DataType::Datetime(precision) => {
                 let precision = match precision {
                     Some(0) => TimeUnit::Second,

datafusion/sqllogictest/src/test_context.rs

@@ -29,8 +29,7 @@ use arrow::array::{
 };
 use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::{
-    DataType, Field, FieldRef, Fields, Metadata, Schema, SchemaRef, TimeUnit, UInt32Type,
-    UnionFields,
+    DataType, Field, Fields, Schema, SchemaRef, TimeUnit, UInt32Type, UnionFields,
 };
 use arrow::record_batch::RecordBatch;
 use datafusion::catalog::{
@@ -41,7 +40,6 @@ use datafusion::common::stats::Precision;
 use datafusion::common::{DataFusionError, Result, not_impl_err};
 use datafusion::functions::math::abs;
 use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
-use datafusion::logical_expr::planner::TypePlanner;
 use datafusion::logical_expr::{
     ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
     Volatility, create_udf,
@@ -83,22 +81,6 @@ pub struct TestContext {
     test_dir: Option<TempDir>,
 }
 
-#[derive(Debug)]
-struct SqlLogicTestTypePlanner;
-
-impl TypePlanner for SqlLogicTestTypePlanner {
-    fn plan_type_field(&self, sql_type: &ast::DataType) -> Result<Option<FieldRef>> {
-        match sql_type {
-            ast::DataType::Uuid => Ok(Some(Arc::new(
-                Field::new("", DataType::FixedSizeBinary(16), true).with_metadata(
-                    Metadata::new().with("ARROW:extension:name", "arrow.uuid"),
-                ),
-            ))),
-            _ => Ok(None),
-        }
-    }
-}
-
 impl TestContext {
     pub fn new(ctx: SessionContext) -> Self {
         Self {
@@ -137,13 +119,6 @@ impl TestContext {
             }
         }
 
-        if matches!(
-            relative_path.file_name().and_then(|name| name.to_str()),
-            Some("cast_extension_type_metadata.slt")
-        ) {
-            state_builder =
-                state_builder.with_type_planner(Arc::new(SqlLogicTestTypePlanner));
-        }
 
         if matches!(
             relative_path.file_name().and_then(|name| name.to_str()),
@@ -176,6 +151,7 @@ impl TestContext {
             state_builder = state_builder.with_statistics_registry(registry);
         }
 
+
         let state = state_builder.build();
 
         let mut test_ctx = TestContext::new(SessionContext::new_with_state(state));