Union mode

apache · Jul 8, 2024 · 8b9715d · 8b9715d
1 parent 9fd793a
commit 8b9715d
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 7 deletions.
diff --git a/datafusion-examples/examples/logical_type.rs b/datafusion-examples/examples/logical_type.rs
@@ -8,7 +8,7 @@ use datafusion::prelude::SessionContext;
 use datafusion_common::logical_type::field::LogicalField;
 use datafusion_common::logical_type::schema::{LogicalSchema, LogicalSchemaRef};
 use datafusion_common::logical_type::signature::LogicalType;
-use datafusion_common::logical_type::ExtensionType;
+use datafusion_common::logical_type::{ExtensionType, ExtensionTypeRef};
 use datafusion_expr::{Expr, TableType};
 use std::any::Any;
 use std::sync::Arc;
@@ -64,7 +64,7 @@ impl TableProvider for ExampleTableSource {
 
     fn schema(&self) -> LogicalSchemaRef {
         // TODO: ugly?
-        let custom_magical_type: Arc<dyn ExtensionType + Send + Sync> =
+        let custom_magical_type: ExtensionTypeRef =
             Arc::new(CustomMagicalType::default());
 
         // This schema will be equivalent to:

diff --git a/datafusion/common/src/logical_type/fields.rs b/datafusion/common/src/logical_type/fields.rs
@@ -18,7 +18,7 @@
 use std::ops::Deref;
 use std::sync::Arc;
 
-use arrow_schema::{Field, Fields, UnionFields};
+use arrow_schema::{Field, FieldRef, Fields, UnionFields};
 
 use super::field::{LogicalField, LogicalFieldRef};
 
@@ -153,3 +153,13 @@ impl From<UnionFields> for LogicalUnionFields {
         Self::from(&value)
     }
 }
+
+impl Into<UnionFields> for LogicalUnionFields {
+    fn into(self) -> UnionFields {
+        UnionFields::from_iter(
+            self.0
+                .into_iter()
+                .map(|(i, f)| (*i, FieldRef::new(f.as_ref().clone().into()))),
+        )
+    }
+}
diff --git a/datafusion/common/src/logical_type/mod.rs b/datafusion/common/src/logical_type/mod.rs
@@ -196,7 +196,7 @@ impl From<&DataType> for LogicalType {
                 sorted.clone(),
             ),
             DataType::RunEndEncoded(_, f) => f.data_type().into(),
-            DataType::Union(f, _) => LogicalType::Union(f.into()),
+            DataType::Union(f, mode) => LogicalType::Union(f.into(), mode.clone()),
         }
     }
 }
diff --git a/datafusion/common/src/logical_type/signature.rs b/datafusion/common/src/logical_type/signature.rs
@@ -18,7 +18,7 @@
 use core::fmt;
 use std::sync::Arc;
 
-use arrow_schema::{IntervalUnit, TimeUnit};
+use arrow_schema::{DataType, FieldRef, IntervalUnit, TimeUnit, UnionMode};
 
 use super::{
     field::LogicalFieldRef,
@@ -53,7 +53,7 @@ pub enum LogicalType {
     Map(LogicalFieldRef, bool),
     Decimal128(u8, i8),
     Decimal256(u8, i8),
-    Union(LogicalUnionFields), // TODO: extension signatures?
+    Union(LogicalUnionFields, UnionMode), // TODO: extension signatures?
 }
 
 impl fmt::Display for LogicalType {
@@ -135,3 +135,41 @@ impl LogicalType {
         matches!(self, Null)
     }
 }
+
+impl Into<DataType> for LogicalType {
+    fn into(self) -> DataType {
+        match self {
+            LogicalType::Null => DataType::Null,
+            LogicalType::Int8 => DataType::Int8,
+            LogicalType::Int16 => DataType::Int16,
+            LogicalType::Int32 => DataType::Int32,
+            LogicalType::Int64 => DataType::Int64,
+            LogicalType::UInt8 => DataType::UInt8,
+            LogicalType::UInt16 => DataType::UInt16,
+            LogicalType::UInt32 => DataType::UInt32,
+            LogicalType::UInt64 => DataType::UInt64,
+            LogicalType::Boolean => DataType::Boolean,
+            LogicalType::Float16 => DataType::Float16,
+            LogicalType::Float32 => DataType::Float32,
+            LogicalType::Float64 => DataType::Float64,
+            LogicalType::Utf8 => DataType::Utf8,
+            LogicalType::Binary => DataType::Binary,
+            LogicalType::Date => DataType::Date32,
+            LogicalType::Time32(tu) => DataType::Time32(tu),
+            LogicalType::Time64(tu) => DataType::Time64(tu),
+            LogicalType::Timestamp(tu, tz) => DataType::Timestamp(tu, tz),
+            LogicalType::Duration(tu) => DataType::Duration(tu),
+            LogicalType::Interval(iu) => DataType::Interval(iu),
+            LogicalType::List(field) => {
+                DataType::List(FieldRef::new(field.as_ref().clone().into()))
+            }
+            LogicalType::Struct(fields) => DataType::Struct(fields.into()),
+            LogicalType::Map(field, v) => {
+                DataType::Map(FieldRef::new(field.as_ref().clone().into()), v)
+            }
+            LogicalType::Decimal128(a, b) => DataType::Decimal128(a, b),
+            LogicalType::Decimal256(a, b) => DataType::Decimal256(a, b),
+            LogicalType::Union(union, mode) => DataType::Union(union.into(), mode),
+        }
+    }
+}
diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs
@@ -1012,7 +1012,7 @@ impl Unparser<'_> {
             LogicalType::Struct(_) => {
                 not_impl_err!("Unsupported DataType: conversion: {data_type:?}")
             }
-            LogicalType::Union(_) => {
+            LogicalType::Union(_, _) => {
                 not_impl_err!("Unsupported DataType: conversion: {data_type:?}")
             }
             LogicalType::Decimal128(precision, scale)