From 32bb06500dd2ad12103b42527350e775d9255aa1 Mon Sep 17 00:00:00 2001
From: Jiayu Liu <jiayu.liu@airbnb.com>
Date: Mon, 31 Oct 2022 17:32:30 +0800
Subject: [PATCH] update uuid function

---
 datafusion/expr/src/function.rs                    |  2 +-
 datafusion/physical-expr/src/string_expressions.rs | 14 +++++++++++++-
 datafusion/proto/src/from_proto.rs                 |  1 +
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
index fa6fe81529c17..fb6858fd963b1 100644
--- a/datafusion/expr/src/function.rs
+++ b/datafusion/expr/src/function.rs
@@ -157,6 +157,7 @@ pub fn return_type(
             utf8_to_int_type(&input_expr_types[0], "octet_length")
         }
         BuiltinScalarFunction::Random => Ok(DataType::Float64),
+        BuiltinScalarFunction::Uuid => Ok(DataType::Utf8),
         BuiltinScalarFunction::RegexpReplace => {
             utf8_to_str_type(&input_expr_types[0], "regex_replace")
         }
@@ -191,7 +192,6 @@ pub fn return_type(
         BuiltinScalarFunction::StartsWith => Ok(DataType::Boolean),
         BuiltinScalarFunction::Strpos => utf8_to_int_type(&input_expr_types[0], "strpos"),
         BuiltinScalarFunction::Substr => utf8_to_str_type(&input_expr_types[0], "substr"),
-        BuiltinScalarFunction::Uuid => Ok(DataType::Utf8),
         BuiltinScalarFunction::ToHex => Ok(match input_expr_types[0] {
             DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
                 DataType::Utf8
diff --git a/datafusion/physical-expr/src/string_expressions.rs b/datafusion/physical-expr/src/string_expressions.rs
index ebb0ad8ce2d4a..5733234f870b0 100644
--- a/datafusion/physical-expr/src/string_expressions.rs
+++ b/datafusion/physical-expr/src/string_expressions.rs
@@ -32,6 +32,7 @@ use datafusion_common::ScalarValue;
 use datafusion_common::{DataFusionError, Result};
 use datafusion_expr::ColumnarValue;
 use std::any::type_name;
+use std::iter;
 use std::sync::Arc;
 use uuid::Uuid;
 
@@ -591,5 +592,16 @@ pub fn upper(args: &[ColumnarValue]) -> Result<ColumnarValue> {
 /// Prints random (v4) uuid values per row
 /// uuid() = 'a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11'
 pub fn uuid(args: &[ColumnarValue]) -> Result<ColumnarValue> {
-    handle(args, |_| Uuid::new_v4().to_string(), "uuid")
+    let len: usize = match &args[0] {
+        ColumnarValue::Array(array) => array.len(),
+        _ => {
+            return Err(DataFusionError::Internal(
+                "Expect uuid function to take no param".to_string(),
+            ))
+        }
+    };
+
+    let values = iter::repeat_with(|| Uuid::new_v4().to_string()).take(len);
+    let array = GenericStringArray::<i32>::from_iter_values(values);
+    Ok(ColumnarValue::Array(Arc::new(array)))
 }
diff --git a/datafusion/proto/src/from_proto.rs b/datafusion/proto/src/from_proto.rs
index cd54929b8eeb7..ed6fc006f2620 100644
--- a/datafusion/proto/src/from_proto.rs
+++ b/datafusion/proto/src/from_proto.rs
@@ -969,6 +969,7 @@ pub fn parse_expr(
                     parse_expr(&args[1], registry)?,
                 )),
                 ScalarFunction::Random => Ok(random()),
+                ScalarFunction::Uuid => Ok(uuid()),
                 ScalarFunction::Repeat => Ok(repeat(
                     parse_expr(&args[0], registry)?,
                     parse_expr(&args[1], registry)?,