risingwave_expr_impl/scalar/
overlay.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Write;
16
17use risingwave_expr::{ExprError, Result, function};
18
19/// Replaces a substring of the given string with a new substring.
20///
21/// ```slt
22/// query T
23/// select overlay('αβγδεζ' placing '💯' from 3);
24/// ----
25/// αβ💯δεζ
26/// ```
27#[function("overlay(varchar, varchar, int4) -> varchar")]
28pub fn overlay(s: &str, new_sub_str: &str, start: i32, writer: &mut impl Write) -> Result<()> {
29    let sub_len = new_sub_str
30        .chars()
31        .count()
32        .try_into()
33        .map_err(|_| ExprError::NumericOutOfRange)?;
34    overlay_for(s, new_sub_str, start, sub_len, writer)
35}
36
37/// Replaces a substring of the given string with a new substring.
38///
39/// ```slt
40/// statement error not positive
41/// select overlay('αβγδεζ' placing '①②③' from 0);
42///
43/// query T
44/// select overlay('αβγδεζ' placing '①②③' from 10);
45/// ----
46/// αβγδεζ①②③
47///
48/// query T
49/// select overlay('αβγδεζ' placing '①②③' from 4 for 2);
50/// ----
51/// αβγ①②③ζ
52///
53/// query T
54/// select overlay('αβγδεζ' placing '①②③' from 4);
55/// ----
56/// αβγ①②③
57///
58/// query T
59/// select overlay('αβγδεζ' placing '①②③' from 2 for 4);
60/// ----
61/// α①②③ζ
62///
63/// query T
64/// select overlay('αβγδεζ' placing '①②③' from 2 for 7);
65/// ----
66/// α①②③
67///
68/// query T
69/// select overlay('αβγδεζ' placing '①②③' from 4 for 0);
70/// ----
71/// αβγ①②③δεζ
72///
73/// query T
74/// select overlay('αβγδεζ' placing '①②③' from 4 for -2);
75/// ----
76/// αβγ①②③βγδεζ
77///
78/// query T
79/// select overlay('αβγδεζ' placing '①②③' from 4 for -1000);
80/// ----
81/// αβγ①②③αβγδεζ
82/// ```
83#[function("overlay(varchar, varchar, int4, int4) -> varchar")]
84pub fn overlay_for(
85    s: &str,
86    new_sub_str: &str,
87    start: i32,
88    count: i32,
89    writer: &mut impl Write,
90) -> Result<()> {
91    if start <= 0 {
92        return Err(ExprError::InvalidParam {
93            name: "start",
94            reason: format!("{start} is not positive").into(),
95        });
96    }
97
98    let mut chars = s.char_indices().skip(start as usize - 1).peekable();
99
100    // write the substring before the overlay.
101    let leading = match chars.peek() {
102        Some((i, _)) => &s[..*i],
103        None => s,
104    };
105    writer.write_str(leading).unwrap();
106
107    // write the new substring.
108    writer.write_str(new_sub_str).unwrap();
109
110    let Ok(count) = count.try_into() else {
111        // For negative `count`, which is rare in practice, we hand over to `substr`
112        let start_right = start
113            .checked_add(count)
114            .ok_or(ExprError::NumericOutOfRange)?;
115        return super::substr::substr_start(s, start_right, writer);
116    };
117
118    // write the substring after the overlay.
119    if let Some((i, _)) = chars.nth(count) {
120        writer.write_str(&s[i..]).unwrap();
121    }
122
123    Ok(())
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn test_overlay() {
132        case("aaa__aaa", "XY", 4, None, "aaaXYaaa");
133        // Place at end.
134        case("aaa", "XY", 4, None, "aaaXY");
135        // Place at start.
136        case("aaa", "XY", 1, Some(0), "XYaaa");
137        // Replace shorter string.
138        case("aaa_aaa", "XYZ", 4, Some(1), "aaaXYZaaa");
139        case("aaaaaa", "XYZ", 4, Some(0), "aaaXYZaaa");
140        // Replace longer string.
141        case("aaa___aaa", "X", 4, Some(3), "aaaXaaa");
142        // start too large.
143        case("aaa", "XY", 123, None, "aaaXY");
144        // count too small or large.
145        case("aaa", "X", 4, Some(-123), "aaaXaaa");
146        case("aaa_", "X", 4, Some(123), "aaaX");
147        // very large start and count
148        case("aaa", "X", i32::MAX, Some(i32::MAX), "aaaX");
149
150        #[track_caller]
151        fn case(s: &str, new_sub_str: &str, start: i32, count: Option<i32>, expected: &str) {
152            let mut writer = String::new();
153            match count {
154                None => overlay(s, new_sub_str, start, &mut writer),
155                Some(count) => overlay_for(s, new_sub_str, start, count, &mut writer),
156            }
157            .unwrap();
158            assert_eq!(writer, expected);
159        }
160    }
161}
risingwave_expr_impl/scalar/overlay.rs

risingwave_expr_impl/scalar/
overlay.rs
