From b17ae25a7dc95bfac0a11758f9180b49a96e6dd6 Mon Sep 17 00:00:00 2001
From: wiedld <dlw405@gmail.com>
Date: Tue, 27 Aug 2024 10:20:56 -0700
Subject: [PATCH] feat(12118): physical plan support for Utf8View

---
 .../substrait/src/physical_plan/consumer.rs   | 15 ++++++++++++++-
 .../substrait/src/physical_plan/producer.rs   | 19 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/datafusion/substrait/src/physical_plan/consumer.rs b/datafusion/substrait/src/physical_plan/consumer.rs
index 5a8b888ef1cc..4bdf443cabe3 100644
--- 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
index 57fe68c4a780..54a4079ec3f1 100644
--- 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,
             })),
         }),