fix(expr): correctly handle unicode for substr

src/expr/src/vector_op/substr.rs

@@ -12,40 +12,52 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::cmp::{max, min};
+use std::cmp::min;
 use std::fmt::Write;
 
 use risingwave_expr_macro::function;
 
-use crate::{bail, Result};
+use crate::{ExprError, Result};
 
 #[function("substr(varchar, int32) -> varchar")]
 pub fn substr_start(s: &str, start: i32, writer: &mut dyn Write) -> Result<()> {
-    let start = (start.saturating_sub(1).max(0) as usize).min(s.len());
-    writer.write_str(&s[start..]).unwrap();
+    let skip = start.saturating_sub(1).max(0) as usize;
+
+    let substr = s.chars().skip(skip);
+    for char in substr {
+        writer.write_char(char).unwrap();
+    }
+
     Ok(())
 }
 
-// #[function("substr(varchar, 0, int32) -> varchar")]
+// #[function("substr(varchar, 1, int32) -> varchar")]
+// TODO: when is this function called?
 pub fn substr_for(s: &str, count: i32, writer: &mut dyn Write) -> Result<()> {
-    let end = min(count as usize, s.len());
-    writer.write_str(&s[..end]).unwrap();
-    Ok(())
+    substr_start_for(s, 1, count, writer)
 }
 
 #[function("substr(varchar, int32, int32) -> varchar")]
 pub fn substr_start_for(s: &str, start: i32, count: i32, writer: &mut dyn Write) -> Result<()> {
     if count < 0 {
-        bail!("length in substr should be non-negative: {}", count);
+        return Err(ExprError::InvalidParam {
+            name: "length",
+            reason: "negative substring length not allowed".to_string(),
+        });
     }
-    let start = start.saturating_sub(1);
-    // NOTE: we use `s.len()` here as an upper bound.
-    // This is so it will return an empty slice if it exceeds
-    // the length of `s`.
-    // 0 <= begin <= s.len()
-    let begin = min(max(start, 0) as usize, s.len());
-    let end = (start.saturating_add(count).max(0) as usize).min(s.len());
-    writer.write_str(&s[begin..end]).unwrap();
+
+    let skip = start.saturating_sub(1).max(0) as usize;
+    let take = if start >= 1 {
+        count as usize
+    } else {
+        count.saturating_add(start.saturating_sub(1)) as usize
+    };
+
+    let substr = s.chars().skip(skip).take(take);
+    for char in substr {
+        writer.write_char(char).unwrap();
+    }
+
     Ok(())
 }
 
@@ -56,6 +68,7 @@ mod tests {
     #[test]
     fn test_substr() -> Result<()> {
         let s = "cxscgccdd";
+        let us = "上海自来水来自海上";
 
         let cases = [
             (s, Some(4), None, "cgccdd"),
@@ -64,6 +77,14 @@ mod tests {
             (s, Some(4), Some(2), "cg"),
             (s, Some(-1), Some(-5), "[unused result]"),
             (s, Some(-1), Some(5), "cxs"),
+            // Unicode test
+            (us, Some(1), Some(3), "上海自"),
+            (us, Some(3), Some(3), "自来水"),
+            (us, None, Some(5), "上海自来水"),
+            (us, Some(6), Some(2), "来自"),
+            (us, Some(6), Some(100), "来自海上"),
+            (us, Some(6), None, "来自海上"),
+            ("Mér", Some(1), Some(2), "Mé"),
         ];
 
         for (s, off, len, expected) in cases {