Add utf8view support for user defined

Author: zhuqi-lucas

datafusion/core/tests/user_defined/user_defined_plan.rs

@@ -63,6 +63,7 @@ use std::hash::Hash;
 use std::task::{Context, Poll};
 use std::{any::Any, collections::BTreeMap, fmt, sync::Arc};
 
+use arrow::array::{Array, ArrayRef, StringViewArray};
 use arrow::{
     array::{Int64Array, StringArray},
     datatypes::SchemaRef,
@@ -100,6 +101,7 @@ use datafusion_optimizer::AnalyzerRule;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
 use async_trait::async_trait;
+use datafusion_common::cast::as_string_view_array;
 use futures::{Stream, StreamExt};
 
 /// Execute the specified sql and return the resulting record batches
@@ -796,22 +798,30 @@ fn accumulate_batch(
     k: &usize,
 ) -> BTreeMap<i64, String> {
     let num_rows = input_batch.num_rows();
+
     // Assuming the input columns are
-    // column[0]: customer_id / UTF8
+    // column[0]: customer_id / UTF8 or UTF8View
     // column[1]: revenue: Int64
-    let customer_id =
-        as_string_array(input_batch.column(0)).expect("Column 0 is not customer_id");
 
+    let customer_id_column = input_batch.column(0);
     let revenue = as_int64_array(input_batch.column(1)).unwrap();
 
     for row in 0..num_rows {
-        add_row(
-            &mut top_values,
-            customer_id.value(row),
-            revenue.value(row),
-            k,
-        );
+        let customer_id = match customer_id_column.data_type() {
+            arrow::datatypes::DataType::Utf8 => {
+                let array = as_string_array(customer_id_column).unwrap();
+                array.value(row)
+            }
+            arrow::datatypes::DataType::Utf8View => {
+                let array = as_string_view_array(customer_id_column).unwrap();
+                array.value(row)
+            }
+            _ => panic!("Unsupported customer_id type"),
+        };
+
+        add_row(&mut top_values, customer_id, revenue.value(row), k);
     }
+
     top_values
 }
 
@@ -843,11 +853,22 @@ impl Stream for TopKReader {
                     self.state.iter().rev().unzip();
 
                 let customer: Vec<&str> = customer.iter().map(|&s| &**s).collect();
+
+                let customer_array: ArrayRef = match schema.field(0).data_type() {
+                    arrow::datatypes::DataType::Utf8 => {
+                        Arc::new(StringArray::from(customer))
+                    }
+                    arrow::datatypes::DataType::Utf8View => {
+                        Arc::new(StringViewArray::from(customer))
+                    }
+                    other => panic!("Unsupported customer_id output type: {:?}", other),
+                };
+
                 Poll::Ready(Some(
                     RecordBatch::try_new(
                         schema,
                         vec![
-                            Arc::new(StringArray::from(customer)),
+                            Arc::new(customer_array),
                             Arc::new(Int64Array::from(revenue)),
                         ],
                     )

datafusion/sql/src/planner.rs

@@ -72,7 +72,7 @@ impl ParserOptions {
             parse_float_as_decimal: false,
             enable_ident_normalization: true,
             support_varchar_with_length: true,
-            map_varchar_to_utf8view: false,
+            map_varchar_to_utf8view: true,
             enable_options_value_normalization: false,
             collect_spans: false,
         }

datafusion/sql/tests/sql_integration.rs

@@ -3360,7 +3360,7 @@ fn parse_decimals_parser_options() -> ParserOptions {
         parse_float_as_decimal: true,
         enable_ident_normalization: false,
         support_varchar_with_length: false,
-        map_varchar_to_utf8view: false,
+        map_varchar_to_utf8view: true,
         enable_options_value_normalization: false,
         collect_spans: false,
     }
@@ -3371,7 +3371,7 @@ fn ident_normalization_parser_options_no_ident_normalization() -> ParserOptions
         parse_float_as_decimal: true,
         enable_ident_normalization: false,
         support_varchar_with_length: false,
-        map_varchar_to_utf8view: false,
+        map_varchar_to_utf8view: true,
         enable_options_value_normalization: false,
         collect_spans: false,
     }
@@ -3382,7 +3382,7 @@ fn ident_normalization_parser_options_ident_normalization() -> ParserOptions {
         parse_float_as_decimal: true,
         enable_ident_normalization: true,
         support_varchar_with_length: false,
-        map_varchar_to_utf8view: false,
+        map_varchar_to_utf8view: true,
         enable_options_value_normalization: false,
         collect_spans: false,
     }