fix: encode function should work with strings and binary

closes apache#14055

Author: mesejo

datafusion/functions/src/encoding/inner.rs

@@ -87,7 +87,18 @@ impl ScalarUDFImpl for EncodeFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        Ok(arg_types[0].to_owned())
+        use DataType::*;
+
+        Ok(match arg_types[0] {
+            Utf8 => Utf8,
+            LargeUtf8 => LargeUtf8,
+            Binary => Utf8,
+            LargeBinary => LargeUtf8,
+            Null => Null,
+            _ => {
+                return plan_err!("The encode function can only accept Utf8 or Binary or Null.");
+            }
+        })
     }
 
     fn invoke_batch(
@@ -112,12 +123,18 @@ impl ScalarUDFImpl for EncodeFunc {
         }
 
         match arg_types[0] {
-            DataType::Utf8 | DataType::Utf8View | DataType::Binary | DataType::Null => {
+            DataType::Utf8 | DataType::Utf8View | DataType::Null => {
                 Ok(vec![DataType::Utf8; 2])
             }
-            DataType::LargeUtf8 | DataType::LargeBinary => {
+            DataType::LargeUtf8 => {
                 Ok(vec![DataType::LargeUtf8, DataType::Utf8])
             }
+            DataType::Binary => {
+                Ok(vec![DataType::Binary, DataType::Utf8])
+            }
+            DataType::LargeBinary => {
+                Ok(vec![DataType::LargeBinary, DataType::Utf8])
+            }
             _ => plan_err!(
                 "1st argument should be Utf8 or Binary or Null, got {:?}",
                 arg_types[0]

datafusion/sqllogictest/test_files/encoding.slt

@@ -23,8 +23,9 @@ CREATE TABLE test(
   hex_field TEXT
 ) as VALUES
   (0, 'abc',  encode('abc', 'base64'), encode('abc', 'hex')),
-  (1, 'qweqwe',    encode('qweqwe', 'base64'), encode('qweqwe', 'hex')),
-  (2, NULL, NULL, NULL)
+  (1, 'qweqwe', encode('qweqwe', 'base64'), encode('qweqwe', 'hex')),
+  (2, NULL, NULL, NULL),
+  (3, X'8f50d3f60eae370ddbf85c86219c55108a350165', encode('abc', 'base64'), encode('abc', 'hex'))
 ;
 
 # errors
@@ -50,27 +51,31 @@ SELECT encode(bin_field, 'hex') FROM test ORDER BY num;
 616263
 717765717765
 NULL
+8f50d3f60eae370ddbf85c86219c55108a350165
 
 query T
 SELECT arrow_cast(decode(base64_field, 'base64'), 'Utf8') FROM test ORDER BY num;
 ----
 abc
 qweqwe
 NULL
+abc
 
 query T
 SELECT arrow_cast(decode(hex_field, 'hex'), 'Utf8') FROM test ORDER BY num;
 ----
 abc
 qweqwe
 NULL
+abc
 
 query T
 select to_hex(num) from test ORDER BY num;
 ----
 0
 1
 2
+3
 
 # test for Utf8View support for encode
 statement ok
@@ -101,6 +106,4 @@ FROM test_utf8view;
 Andrew QW5kcmV3 416e64726577 X WA 58
 Xiangpeng WGlhbmdwZW5n 5869616e6770656e67 Xiangpeng WGlhbmdwZW5n 5869616e6770656e67
 Raphael UmFwaGFlbA 5261706861656c R Ug 52
-NULL NULL NULL R Ug 52
-
-select encode(X'8f50d3f60eae370ddbf85c86219c55108a350165', 'hex');
+NULL NULL NULL R Ug 52