feat(12118): physical plan support for Utf8View

apache · Aug 27, 2024 · b17ae25 · b17ae25
1 parent d7be771
commit b17ae25
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/datafusion/substrait/src/physical_plan/consumer.rs b/datafusion/substrait/src/physical_plan/consumer.rs
@@ -37,6 +37,11 @@ use substrait::proto::{
     expression::MaskExpression, read_rel::ReadType, rel::RelType, Rel,
 };
 
+use crate::variation_const::{
+    DEFAULT_CONTAINER_TYPE_VARIATION_REF, LARGE_CONTAINER_TYPE_VARIATION_REF,
+    VIEW_CONTAINER_TYPE_VARIATION_REF,
+};
+
 /// Convert Substrait Rel to DataFusion ExecutionPlan
 #[async_recursion]
 pub async fn from_substrait_rel(
@@ -177,7 +182,15 @@ fn to_field(name: &String, r#type: &Type) -> Result<Field> {
         }
         Kind::String(string) => {
             nullable = is_nullable(string.nullability);
-            Ok(DataType::Utf8)
+            match string.type_variation_reference {
+                DEFAULT_CONTAINER_TYPE_VARIATION_REF => Ok(DataType::Utf8),
+                LARGE_CONTAINER_TYPE_VARIATION_REF => Ok(DataType::LargeUtf8),
+                VIEW_CONTAINER_TYPE_VARIATION_REF => Ok(DataType::Utf8View),
+                _ => substrait_err!(
+                    "Invalid type variation found for substrait string type class: {}",
+                    string.type_variation_reference
+                ),
+            }
         }
         _ => substrait_err!(
             "Unsupported kind: {:?} in the type with name {}",

diff --git a/datafusion/substrait/src/physical_plan/producer.rs b/datafusion/substrait/src/physical_plan/producer.rs
@@ -35,6 +35,11 @@ use substrait::proto::ReadRel;
 use substrait::proto::Rel;
 use substrait::proto::{extensions, NamedStruct, Type};
 
+use crate::variation_const::{
+    DEFAULT_CONTAINER_TYPE_VARIATION_REF, LARGE_CONTAINER_TYPE_VARIATION_REF,
+    VIEW_CONTAINER_TYPE_VARIATION_REF,
+};
+
 /// Convert DataFusion ExecutionPlan to Substrait Rel
 pub fn to_substrait_rel(
     plan: &dyn ExecutionPlan,
@@ -155,7 +160,19 @@ fn to_substrait_type(data_type: &DataType, nullable: bool) -> Result<Type> {
         }),
         DataType::Utf8 => Ok(Type {
             kind: Some(Kind::String(SubstraitString {
-                type_variation_reference: 0,
+                type_variation_reference: DEFAULT_CONTAINER_TYPE_VARIATION_REF,
+                nullability,
+            })),
+        }),
+        DataType::LargeUtf8 => Ok(Type {
+            kind: Some(Kind::String(SubstraitString {
+                type_variation_reference: LARGE_CONTAINER_TYPE_VARIATION_REF,
+                nullability,
+            })),
+        }),
+        DataType::Utf8View => Ok(Type {
+            kind: Some(Kind::String(SubstraitString {
+                type_variation_reference: VIEW_CONTAINER_TYPE_VARIATION_REF,
                 nullability,
             })),
         }),