fix: substr - correct behaivour with negative start pos (#1660)

apache · Jan 31, 2022 · b6ace16 · b6ace16
1 parent 940d4eb
commit b6ace16
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 10 deletions.
diff --git a/datafusion/src/physical_plan/functions.rs b/datafusion/src/physical_plan/functions.rs
@@ -3582,6 +3582,18 @@ mod tests {
             StringArray
         );
         #[cfg(feature = "unicode_expressions")]
+        test_function!(
+            Substr,
+            &[
+                lit(ScalarValue::Utf8(Some("joséésoj".to_string()))),
+                lit(ScalarValue::Int64(Some(-5))),
+            ],
+            Ok(Some("joséésoj")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        #[cfg(feature = "unicode_expressions")]
         test_function!(
             Substr,
             &[
@@ -3680,6 +3692,61 @@ mod tests {
             StringArray
         );
         #[cfg(feature = "unicode_expressions")]
+        test_function!(
+            Substr,
+            &[
+                lit(ScalarValue::Utf8(Some("alphabet".to_string()))),
+                lit(ScalarValue::Int64(Some(0))),
+                lit(ScalarValue::Int64(Some(5))),
+            ],
+            Ok(Some("alph")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        // starting from 5 (10 + -5)
+        #[cfg(feature = "unicode_expressions")]
+        test_function!(
+            Substr,
+            &[
+                lit(ScalarValue::Utf8(Some("alphabet".to_string()))),
+                lit(ScalarValue::Int64(Some(-5))),
+                lit(ScalarValue::Int64(Some(10))),
+            ],
+            Ok(Some("alph")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        // starting from -1 (4 + -5)
+        #[cfg(feature = "unicode_expressions")]
+        test_function!(
+            Substr,
+            &[
+                lit(ScalarValue::Utf8(Some("alphabet".to_string()))),
+                lit(ScalarValue::Int64(Some(-5))),
+                lit(ScalarValue::Int64(Some(4))),
+            ],
+            Ok(Some("")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        // starting from 0 (5 + -5)
+        #[cfg(feature = "unicode_expressions")]
+        test_function!(
+            Substr,
+            &[
+                lit(ScalarValue::Utf8(Some("alphabet".to_string()))),
+                lit(ScalarValue::Int64(Some(-5))),
+                lit(ScalarValue::Int64(Some(5))),
+            ],
+            Ok(Some("")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        #[cfg(feature = "unicode_expressions")]
         test_function!(
             Substr,
             &[

diff --git a/datafusion/src/physical_plan/unicode_expressions.rs b/datafusion/src/physical_plan/unicode_expressions.rs
@@ -457,21 +457,29 @@ pub fn substr<T: StringOffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
                                 start,
                                 count
                             )))
-                        } else if start <= 0 {
-                            Ok(Some(string.to_string()))
                         } else {
                             let graphemes = string.graphemes(true).collect::<Vec<&str>>();
-                            let start_pos = start as usize - 1;
-                            let count_usize = count as usize;
-                            if graphemes.len() < start_pos {
+                            let (start_pos, end_pos) = if start <= 0 {
+                                let end_pos = start + count - 1;
+                                (
+                                    0_usize,
+                                    if end_pos < 0 {
+                                        // we use 0 as workaround for usize to return empty string
+                                        0
+                                    } else {
+                                        end_pos as usize
+                                    },
+                                )
+                            } else {
+                                ((start - 1) as usize, (start + count - 1) as usize)
+                            };
+
+                            if end_pos == 0 || graphemes.len() < start_pos {
                                 Ok(Some("".to_string()))
-                            } else if graphemes.len() < start_pos + count_usize {
+                            } else if graphemes.len() < end_pos {
                                 Ok(Some(graphemes[start_pos..].concat()))
                             } else {
-                                Ok(Some(
-                                    graphemes[start_pos..start_pos + count_usize]
-                                        .concat(),
-                                ))
+                                Ok(Some(graphemes[start_pos..end_pos].concat()))
                             }
                         }
                     }