coerce substr

jonathanc-n · Oct 29, 2024 · 5781997 · 5781997
1 parent 1fd6116
commit 5781997
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 11 deletions.
diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs
@@ -94,6 +94,7 @@ impl ScalarUDFImpl for SubstrFunc {
         let first_data_type = match &arg_types[0] {
             DataType::Null => Ok(DataType::Utf8),
             DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8 => Ok(arg_types[0].clone()),
+            DataType::Int64 | DataType::Int32 => Ok(DataType::Utf8), // Allows for first arg to be int
             DataType::Dictionary(key_type, value_type) => {
                 if key_type.is_integer() {
                     match value_type.as_ref() {

diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt
@@ -193,11 +193,10 @@ SELECT substr('alphabet', 3, CAST(NULL AS int))
 ----
 NULL
 
-statement error The first argument of the substr function can only be a string, but got Int64
-SELECT substr(1, 3)
-
-statement error The first argument of the substr function can only be a string, but got Int64
-SELECT substr(1, 3, 4)
+query T
+select substr(123456, 3);
+----
+3456
 
 query T
 SELECT translate('12345', '143', 'ax')

diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt
@@ -123,12 +123,6 @@ SELECT substr('Hello🌏世界', 5, 3)
 ----
 o🌏世
 
-statement error The first argument of the substr function can only be a string, but got Int64
-SELECT substr(1, 3)
-
-statement error The first argument of the substr function can only be a string, but got Int64
-SELECT substr(1, 3, 4)
-
 statement error Execution error: negative substring length not allowed
 select substr(arrow_cast('foo', 'Utf8View'), 1, -1);