linkrs_query/executor/expression/functions/builtin/
string.rs

1//! Implementation of string functions
2
3use crate::executor::expression::ExpressionError;
4use linkrs_core::value::list::List;
5use linkrs_core::value::NullType;
6use linkrs_core::Value;
7
8define_function_enum! {
9    /// String function enumeration
10    pub enum StringFunction {
11        Upper => {
12            name: "upper",
13            arity: 1,
14            variadic: false,
15            description: "convert to uppercase",
16            handler: execute_upper
17        },
18        Lower => {
19            name: "lower",
20            arity: 1,
21            variadic: false,
22            description: "Convert to lowercase",
23            handler: execute_lower
24        },
25        Trim => {
26            name: "trim",
27            arity: 1,
28            variadic: false,
29            description: "Remove header and footer",
30            handler: execute_trim
31        },
32        Substring => {
33            name: "substring",
34            arity: 3,
35            variadic: false,
36            description: "Get substring",
37            handler: execute_substring
38        },
39        Concat => {
40            name: "concat",
41            arity: 1,
42            variadic: true,
43            description: "connection string",
44            handler: execute_concat
45        },
46        Replace => {
47            name: "replace",
48            arity: 2,
49            variadic: false,
50            description: "Replacement String",
51            handler: execute_replace
52        },
53        Contains => {
54            name: "contains",
55            arity: 2,
56            variadic: false,
57            description: "Checks if it contains substrings",
58            handler: execute_contains
59        },
60        StartsWith => {
61            name: "starts_with",
62            arity: 2,
63            variadic: false,
64            description: "Checks if it starts with the specified string",
65            handler: execute_starts_with
66        },
67        EndsWith => {
68            name: "ends_with",
69            arity: 2,
70            variadic: false,
71            description: "Checks if the specified string ends",
72            handler: execute_ends_with
73        },
74        Split => {
75            name: "split",
76            arity: 2,
77            variadic: false,
78            description: "Split String",
79            handler: execute_split
80        },
81        Lpad => {
82            name: "lpad",
83            arity: 3,
84            variadic: false,
85            description: "Left padding string",
86            handler: execute_lpad
87        },
88        Rpad => {
89            name: "rpad",
90            arity: 3,
91            variadic: false,
92            description: "Right Fill String",
93            handler: execute_rpad
94        },
95        ConcatWs => {
96            name: "concat_ws",
97            arity: 2,
98            variadic: true,
99            description: "Concatenate strings using delimiters",
100            handler: execute_concat_ws
101        },
102        Strcasecmp => {
103            name: "strcasecmp",
104            arity: 2,
105            variadic: false,
106            description: "Compare strings case-insensitively",
107            handler: execute_strcasecmp
108        },
109        Levenshtein => {
110            name: "levenshtein",
111            arity: 2,
112            variadic: false,
113            description: "Calculate Levenshtein edit distance between two strings",
114            handler: execute_levenshtein
115        },
116        SplitPart => {
117            name: "split_part",
118            arity: 3,
119            variadic: false,
120            description: "Split string by delimiter and return Nth part",
121            handler: execute_split_part
122        },
123        Initcap => {
124            name: "initcap",
125            arity: 1,
126            variadic: false,
127            description: "Capitalize first letter of each word",
128            handler: execute_initcap
129        },
130        Repeat => {
131            name: "repeat",
132            arity: 2,
133            variadic: false,
134            description: "Repeat string N times",
135            handler: execute_repeat
136        },
137        Position => {
138            name: "position",
139            arity: 2,
140            variadic: false,
141            description: "Find position of substring",
142            handler: execute_position
143        },
144        Left => {
145            name: "left",
146            arity: 2,
147            variadic: false,
148            description: "Get first N characters of string",
149            handler: execute_left
150        },
151        Right => {
152            name: "right",
153            arity: 2,
154            variadic: false,
155            description: "Get last N characters of string",
156            handler: execute_right
157        },
158        StringInsert => {
159            name: "insert",
160            arity: 4,
161            variadic: false,
162            description: "Insert substring at specified position",
163            handler: execute_string_insert
164        },
165        Translate => {
166            name: "translate",
167            arity: 3,
168            variadic: false,
169            description: "Replace characters in string using mapping",
170            handler: execute_translate
171        },
172        Format => {
173            name: "format",
174            arity: 2,
175            variadic: true,
176            description: "Format string with placeholders",
177            handler: execute_format
178        },
179        StringSplit => {
180            name: "string_split",
181            arity: 2,
182            variadic: false,
183            description: "Split string into substrings by delimiter",
184            handler: execute_string_split
185        },
186        Reverse => {
187            name: "reverse_string",
188            arity: 1,
189            variadic: false,
190            description: "Reverse a string",
191            handler: execute_reverse
192        },
193    }
194}
195
196define_unary_string_fn!(execute_upper, |s: &str| s.to_uppercase(), "upper");
197define_unary_string_fn!(execute_lower, |s: &str| s.to_lowercase(), "lower");
198define_unary_string_fn!(execute_trim, |s: &str| s.trim().to_string(), "trim");
199
200fn execute_substring(args: &[Value]) -> Result<Value, ExpressionError> {
201    match (&args[0], &args[1], &args[2]) {
202        (Value::String(s), Value::Int(start), Value::Int(len)) => {
203            if *start < 0 || *len <= 0 {
204                return Ok(Value::string(String::new()));
205            }
206            let result: String = s
207                .chars()
208                .skip(*start as usize)
209                .take(*len as usize)
210                .collect();
211            Ok(Value::string(result))
212        }
213        (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
214            Ok(Value::Null(NullType::Null))
215        }
216        _ => Err(ExpressionError::type_error(
217            "The substring function takes a string and two integers.",
218        )),
219    }
220}
221
222fn execute_concat(args: &[Value]) -> Result<Value, ExpressionError> {
223    let mut result = String::new();
224    for arg in args {
225        match arg {
226            Value::String(s) => result.push_str(s),
227            Value::Null(_) => return Ok(Value::Null(NullType::Null)),
228            _ => {
229                return Err(ExpressionError::type_error(
230                    "The concat function requires a string type",
231                ))
232            }
233        }
234    }
235    Ok(Value::string(result))
236}
237
238fn execute_replace(args: &[Value]) -> Result<Value, ExpressionError> {
239    match (&args[0], &args[1]) {
240        (Value::String(s), Value::String(from)) => Ok(Value::string(s.replace(from.as_str(), ""))),
241        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
242        _ => Err(ExpressionError::type_error(
243            "The replace function requires a string type",
244        )),
245    }
246}
247
248fn execute_contains(args: &[Value]) -> Result<Value, ExpressionError> {
249    match (&args[0], &args[1]) {
250        (Value::String(s), Value::String(sub)) => Ok(Value::Bool(s.contains(sub.as_str()))),
251        (Value::List(list), Value::String(target)) => Ok(Value::Bool(
252            list.values
253                .iter()
254                .any(|v| matches!(v, Value::String(s) if s == target)),
255        )),
256        (Value::List(list), Value::Int(target)) => Ok(Value::Bool(
257            list.values
258                .iter()
259                .any(|v| matches!(v, Value::Int(i) if *i == *target)),
260        )),
261        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
262        _ => Err(ExpressionError::type_error(
263            "The contains function requires a string or list type.",
264        )),
265    }
266}
267
268define_binary_string_bool_fn!(
269    execute_starts_with,
270    |s: &str, prefix: &str| s.starts_with(prefix),
271    "starts_with"
272);
273define_binary_string_bool_fn!(
274    execute_ends_with,
275    |s: &str, suffix: &str| s.ends_with(suffix),
276    "ends_with"
277);
278
279fn execute_split(args: &[Value]) -> Result<Value, ExpressionError> {
280    use linkrs_core::value::list::List;
281    match (&args[0], &args[1]) {
282        (Value::String(s), Value::String(delimiter)) => {
283            let parts: Vec<Value> = s.split(delimiter.as_str()).map(Value::string).collect();
284            Ok(Value::list(List { values: parts }))
285        }
286        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
287        _ => Err(ExpressionError::type_error("split requires string type")),
288    }
289}
290
291fn execute_lpad(args: &[Value]) -> Result<Value, ExpressionError> {
292    match (&args[0], &args[1], &args[2]) {
293        (Value::String(s), Value::Int(len), Value::String(pad)) => {
294            if *len < 0 {
295                return Err(ExpressionError::type_error(
296                    "The lpad function requires a non-negative length",
297                ));
298            }
299            let len = *len as usize;
300            let char_count = s.chars().count();
301            if char_count >= len {
302                Ok(Value::string(s.chars().take(len).collect::<String>()))
303            } else {
304                let pad_chars: Vec<char> = pad.chars().collect();
305                if pad_chars.is_empty() {
306                    return Ok(Value::string(s.clone()));
307                }
308                let pad_str: String = pad_chars
309                    .into_iter()
310                    .cycle()
311                    .take(len - char_count)
312                    .collect();
313                Ok(Value::string(format!("{}{}", pad_str, s)))
314            }
315        }
316        (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
317            Ok(Value::Null(NullType::Null))
318        }
319        _ => Err(ExpressionError::type_error(
320            "The lpad function takes 3 arguments",
321        )),
322    }
323}
324
325fn execute_rpad(args: &[Value]) -> Result<Value, ExpressionError> {
326    match (&args[0], &args[1], &args[2]) {
327        (Value::String(s), Value::Int(len), Value::String(pad)) => {
328            if *len < 0 {
329                return Err(ExpressionError::type_error(
330                    "The rpad function requires a non-negative length",
331                ));
332            }
333            let len = *len as usize;
334            let char_count = s.chars().count();
335            if char_count >= len {
336                Ok(Value::string(s.chars().take(len).collect::<String>()))
337            } else {
338                let pad_chars: Vec<char> = pad.chars().collect();
339                if pad_chars.is_empty() {
340                    return Ok(Value::string(s.clone()));
341                }
342                let pad_str: String = pad_chars
343                    .into_iter()
344                    .cycle()
345                    .take(len - char_count)
346                    .collect();
347                Ok(Value::string(format!("{}{}", s, pad_str)))
348            }
349        }
350        (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
351            Ok(Value::Null(NullType::Null))
352        }
353        _ => Err(ExpressionError::type_error(
354            "The rpad function takes string, integer, and string arguments",
355        )),
356    }
357}
358
359fn execute_concat_ws(args: &[Value]) -> Result<Value, ExpressionError> {
360    if args.len() < 2 {
361        return Err(ExpressionError::type_error(
362            "The concat_ws function takes at least 2 arguments",
363        ));
364    }
365    let separator = match &args[0] {
366        Value::String(s) => s.to_string(),
367        Value::Null(_) => return Ok(Value::Null(NullType::Null)),
368        _ => {
369            return Err(ExpressionError::type_error(
370                "The concat_ws function requires a string type for the first argument",
371            ))
372        }
373    };
374    let mut result = String::new();
375    for (i, arg) in args[1..].iter().enumerate() {
376        match arg {
377            Value::String(s) => {
378                if i > 0 {
379                    result.push_str(&separator);
380                }
381                result.push_str(s);
382            }
383            Value::Null(_) => return Ok(Value::Null(NullType::Null)),
384            _ => {
385                return Err(ExpressionError::type_error(
386                    "The concat_ws function requires the string type",
387                ))
388            }
389        }
390    }
391    Ok(Value::string(result))
392}
393
394fn execute_strcasecmp(args: &[Value]) -> Result<Value, ExpressionError> {
395    match (&args[0], &args[1]) {
396        (Value::String(a), Value::String(b)) => {
397            let cmp = a.to_lowercase().cmp(&b.to_lowercase());
398            Ok(Value::Int(match cmp {
399                std::cmp::Ordering::Less => -1,
400                std::cmp::Ordering::Equal => 0,
401                std::cmp::Ordering::Greater => 1,
402            }))
403        }
404        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
405        _ => Err(ExpressionError::type_error(
406            "The strcasecmp function requires the string type",
407        )),
408    }
409}
410
411fn execute_levenshtein(args: &[Value]) -> Result<Value, ExpressionError> {
412    match (&args[0], &args[1]) {
413        (Value::String(s1), Value::String(s2)) => {
414            let dist = levenshtein_distance(s1, s2);
415            Ok(Value::Int(dist as i32))
416        }
417        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
418        _ => Err(ExpressionError::type_error(
419            "levenshtein requires string arguments",
420        )),
421    }
422}
423
424fn execute_split_part(args: &[Value]) -> Result<Value, ExpressionError> {
425    match (&args[0], &args[1], &args[2]) {
426        (Value::String(s), Value::String(delimiter), Value::Int(n)) => {
427            if *n <= 0 {
428                return Err(ExpressionError::type_error(
429                    "split_part index must be positive",
430                ));
431            }
432            let parts: Vec<&str> = s.split(delimiter.as_str()).collect();
433            let idx = (*n - 1) as usize;
434            if idx < parts.len() {
435                Ok(Value::string(parts[idx]))
436            } else {
437                Ok(Value::string(String::new()))
438            }
439        }
440        (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
441            Ok(Value::Null(NullType::Null))
442        }
443        _ => Err(ExpressionError::type_error(
444            "split_part requires string, string, and integer arguments",
445        )),
446    }
447}
448
449fn execute_initcap(args: &[Value]) -> Result<Value, ExpressionError> {
450    match &args[0] {
451        Value::String(s) => {
452            let result: String = s
453                .split_whitespace()
454                .map(|word| {
455                    let mut chars = word.chars();
456                    match chars.next() {
457                        None => String::new(),
458                        Some(c) => c.to_uppercase().to_string() + &chars.as_str().to_lowercase(),
459                    }
460                })
461                .collect::<Vec<_>>()
462                .join(" ");
463            Ok(Value::string(result))
464        }
465        Value::Null(_) => Ok(Value::Null(NullType::Null)),
466        _ => Err(ExpressionError::type_error(
467            "initcap requires a string type",
468        )),
469    }
470}
471
472fn execute_repeat(args: &[Value]) -> Result<Value, ExpressionError> {
473    match (&args[0], &args[1]) {
474        (Value::String(s), Value::Int(n)) => {
475            if *n < 0 {
476                return Err(ExpressionError::type_error(
477                    "repeat count must be non-negative",
478                ));
479            }
480            Ok(Value::string(s.repeat(*n as usize)))
481        }
482        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
483        _ => Err(ExpressionError::type_error(
484            "repeat requires string and integer arguments",
485        )),
486    }
487}
488
489fn execute_position(args: &[Value]) -> Result<Value, ExpressionError> {
490    match (&args[0], &args[1]) {
491        (Value::String(s), Value::String(sub)) => {
492            if let Some(byte_idx) = s.find(sub.as_str()) {
493                let char_pos = s[..byte_idx].chars().count() + 1;
494                Ok(Value::Int(char_pos as i32))
495            } else {
496                Ok(Value::Int(0))
497            }
498        }
499        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
500        _ => Err(ExpressionError::type_error(
501            "position requires string arguments",
502        )),
503    }
504}
505
506fn execute_left(args: &[Value]) -> Result<Value, ExpressionError> {
507    match (&args[0], &args[1]) {
508        (Value::String(s), Value::Int(n)) => {
509            let n = *n;
510            if n >= s.chars().count() as i32 {
511                Ok(Value::string(s.clone()))
512            } else if n <= 0 {
513                Ok(Value::string(String::new()))
514            } else {
515                let result: String = s.chars().take(n as usize).collect();
516                Ok(Value::string(result))
517            }
518        }
519        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
520        _ => Err(ExpressionError::type_error(
521            "left requires string and integer arguments",
522        )),
523    }
524}
525
526fn execute_right(args: &[Value]) -> Result<Value, ExpressionError> {
527    match (&args[0], &args[1]) {
528        (Value::String(s), Value::Int(n)) => {
529            let n = *n;
530            let char_count = s.chars().count();
531            if n >= char_count as i32 {
532                Ok(Value::string(s.clone()))
533            } else if n <= 0 {
534                Ok(Value::string(String::new()))
535            } else {
536                let result: String = s.chars().skip(char_count - n as usize).collect();
537                Ok(Value::string(result))
538            }
539        }
540        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
541        _ => Err(ExpressionError::type_error(
542            "right requires string and integer arguments",
543        )),
544    }
545}
546
547fn execute_string_insert(args: &[Value]) -> Result<Value, ExpressionError> {
548    match (&args[0], &args[1], &args[2], &args[3]) {
549        (Value::String(s), Value::Int(pos), Value::Int(len), Value::String(newsub)) => {
550            let char_count = s.chars().count();
551            let pos = (*pos).max(0) as usize;
552            let del = (*len).max(0) as usize;
553            if pos > char_count {
554                return Ok(Value::string(s.clone()));
555            }
556            let end = (pos + del).min(char_count);
557            let prefix: String = s.chars().take(pos).collect();
558            let suffix: String = s.chars().skip(end).collect();
559            Ok(Value::string(format!("{}{}{}", prefix, newsub, suffix)))
560        }
561        (Value::Null(_), _, _, _)
562        | (_, Value::Null(_), _, _)
563        | (_, _, Value::Null(_), _)
564        | (_, _, _, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
565        _ => Err(ExpressionError::type_error(
566            "insert requires string, int, int, string arguments",
567        )),
568    }
569}
570
571fn execute_translate(args: &[Value]) -> Result<Value, ExpressionError> {
572    match (&args[0], &args[1], &args[2]) {
573        (Value::String(s), Value::String(from), Value::String(to)) => {
574            let result: String = s
575                .chars()
576                .map(|c| {
577                    if let Some(pos) = from.chars().position(|fc| fc == c) {
578                        to.chars().nth(pos).unwrap_or(c)
579                    } else {
580                        c
581                    }
582                })
583                .collect();
584            Ok(Value::string(result))
585        }
586        (Value::Null(_), _, _) | (_, Value::Null(_), _) | (_, _, Value::Null(_)) => {
587            Ok(Value::Null(NullType::Null))
588        }
589        _ => Err(ExpressionError::type_error(
590            "translate requires string, string, string arguments",
591        )),
592    }
593}
594
595fn execute_format(args: &[Value]) -> Result<Value, ExpressionError> {
596    if args.len() < 2 {
597        return Err(ExpressionError::type_error(
598            "format requires at least 2 arguments (format string + values)",
599        ));
600    }
601    let format_str = match &args[0] {
602        Value::String(s) => s.to_string(),
603        Value::Null(_) => return Ok(Value::Null(NullType::Null)),
604        _ => {
605            return Err(ExpressionError::type_error(
606                "format first argument must be a string",
607            ))
608        }
609    };
610    let mut result = format_str;
611    for (i, arg) in args[1..].iter().enumerate() {
612        let placeholder = format!("{{{}}}", i);
613        let replacement = match arg {
614            Value::Null(_) => "NULL".to_string(),
615            Value::String(s) => s.to_string(),
616            other => format!("{}", other),
617        };
618        result = result.replace(&placeholder, &replacement);
619    }
620    Ok(Value::string(result))
621}
622
623fn execute_string_split(args: &[Value]) -> Result<Value, ExpressionError> {
624    match (&args[0], &args[1]) {
625        (Value::String(s), Value::String(delimiter)) => {
626            let parts: Vec<Value> = s.split(delimiter.as_str()).map(Value::string).collect();
627            Ok(Value::list(List { values: parts }))
628        }
629        (Value::Null(_), _) | (_, Value::Null(_)) => Ok(Value::Null(NullType::Null)),
630        _ => Err(ExpressionError::type_error(
631            "string_split requires string and string arguments",
632        )),
633    }
634}
635
636fn execute_reverse(args: &[Value]) -> Result<Value, ExpressionError> {
637    match &args[0] {
638        Value::String(s) => Ok(Value::string(s.chars().rev().collect::<String>())),
639        Value::Null(_) => Ok(Value::Null(NullType::Null)),
640        _ => Err(ExpressionError::type_error("reverse requires a string")),
641    }
642}
643
644fn levenshtein_distance(s1: &str, s2: &str) -> usize {
645    let len1 = s1.chars().count();
646    let len2 = s2.chars().count();
647    if len1 == 0 {
648        return len2;
649    }
650    if len2 == 0 {
651        return len1;
652    }
653
654    let mut prev_row: Vec<usize> = (0..=len2).collect();
655    let mut curr_row = vec![0usize; len2 + 1];
656
657    for (i, c1) in s1.chars().enumerate() {
658        curr_row[0] = i + 1;
659        for (j, c2) in s2.chars().enumerate() {
660            let cost = if c1 == c2 { 0 } else { 1 };
661            curr_row[j + 1] = (prev_row[j + 1] + 1)
662                .min(curr_row[j] + 1)
663                .min(prev_row[j] + cost);
664        }
665        std::mem::swap(&mut prev_row, &mut curr_row);
666    }
667
668    prev_row[len2]
669}
670
671#[cfg(test)]
672mod tests;
linkrs_query/executor/expression/functions/builtin/string.rs

linkrs_query/executor/expression/functions/builtin/
string.rs
