risingwave_expr_impl/scalar/
string_to_array.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 auto_enums::auto_enum;
16use risingwave_common::types::ScalarRefImpl;
17use risingwave_expr::function;
18
19#[auto_enum(Iterator)]
20fn string_to_array_inner<'a>(s: &'a str, sep: Option<&'a str>) -> impl Iterator<Item = &'a str> {
21    match s.is_empty() {
22        true => std::iter::empty(),
23        #[nested]
24        _ => match sep {
25            Some(sep) if sep.is_empty() => std::iter::once(s),
26            Some(sep) => s.split(sep),
27            None => s.char_indices().map(move |(index, ch)| {
28                let len = ch.len_utf8();
29                &s[index..index + len]
30            }),
31        },
32    }
33}
34
35// Use cases shown in `e2e_test/batch/functions/string_to_array.slt.part`
36#[function("string_to_array(varchar, varchar) -> varchar[]")]
37pub fn string_to_array2(
38    s: &str,
39    sep: Option<&str>,
40    writer: &mut impl risingwave_common::array::ListWrite,
41) {
42    writer.write_iter(string_to_array_inner(s, sep).map(|s| Some(ScalarRefImpl::Utf8(s))));
43}
44
45#[function("string_to_array(varchar, varchar, varchar) -> varchar[]")]
46pub fn string_to_array3(
47    s: &str,
48    sep: Option<&str>,
49    null: Option<&str>,
50    writer: &mut impl risingwave_common::array::ListWrite,
51) {
52    match null {
53        Some(null) => {
54            let iter = string_to_array_inner(s, sep)
55                .map(|x| (x != null).then_some(ScalarRefImpl::Utf8(x)));
56            writer.write_iter(iter);
57        }
58        None => {
59            string_to_array2(s, sep, writer);
60        }
61    }
62}