## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/functions/src/core/overlay.rs

    Line
    Count
    Source
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    use std::any::Any;
    19
    use std::sync::Arc;
    20
    21
    use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
    22
    use arrow::datatypes::DataType;
    23
    24
    use crate::utils::{make_scalar_function, utf8_to_str_type};
    25
    use datafusion_common::cast::{
    26
        as_generic_string_array, as_int64_array, as_string_view_array,
    27
    };
    28
    use datafusion_common::{exec_err, Result};
    29
    use datafusion_expr::{ColumnarValue, Documentation, TypeSignature, Volatility};
    30
    use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
    31
    use datafusion_macros::user_doc;
    32
    33
    #[user_doc(
    34
        doc_section(label = "String Functions"),
    35
        description = "Returns the string which is replaced by another string from the specified position and specified count length.",
    36
        syntax_example = "overlay(str PLACING substr FROM pos [FOR count])",
    37
        sql_example = r#"```sql
    38
    > select overlay('Txxxxas' placing 'hom' from 2 for 4);
    39
    +--------------------------------------------------------+
    40
    | overlay(Utf8("Txxxxas"),Utf8("hom"),Int64(2),Int64(4)) |
    41
    +--------------------------------------------------------+
    42
    | Thomas                                                 |
    43
    +--------------------------------------------------------+
    44
    ```"#,
    45
        standard_argument(name = "str", prefix = "String"),
    46
        argument(name = "substr", description = "Substring to replace in str."),
    47
        argument(
    48
            name = "pos",
    49
            description = "The start position to start the replace in str."
    50
        ),
    51
        argument(
    52
            name = "count",
    53
            description = "The count of characters to be replaced from start position of str. If not specified, will use substr length instead."
    54
        )
    55
    )]
    56
    #[derive(Debug, PartialEq, Eq, Hash)]
    57
    pub struct OverlayFunc {
    58
        signature: Signature,
    59
    }
    60
    61
    impl Default for OverlayFunc {
    62
    0
        fn default() -> Self {
    63
    0
            Self::new()
    64
    0
        }
    65
    }
    66
    67
    impl OverlayFunc {
    68
    0
        pub fn new() -> Self {
    69
            use DataType::*;
    70
    0
            Self {
    71
    0
                signature: Signature::one_of(
    72
    0
                    vec![
    73
    0
                        TypeSignature::Exact(vec![Utf8View, Utf8View, Int64, Int64]),
    74
    0
                        TypeSignature::Exact(vec![Utf8, Utf8, Int64, Int64]),
    75
    0
                        TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, Int64, Int64]),
    76
    0
                        TypeSignature::Exact(vec![Utf8View, Utf8View, Int64]),
    77
    0
                        TypeSignature::Exact(vec![Utf8, Utf8, Int64]),
    78
    0
                        TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, Int64]),
    79
    0
                    ],
    80
    0
                    Volatility::Immutable,
    81
    0
                ),
    82
    0
            }
    83
    0
        }
    84
    }
    85
    86
    impl ScalarUDFImpl for OverlayFunc {
    87
    0
        fn as_any(&self) -> &dyn Any {
    88
    0
            self
    89
    0
        }
    90
    91
    0
        fn name(&self) -> &str {
    92
    0
            "overlay"
    93
    0
        }
    94
    95
    0
        fn signature(&self) -> &Signature {
    96
    0
            &self.signature
    97
    0
        }
    98
    99
    0
        fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    100
    0
            utf8_to_str_type(&arg_types[0], "overlay")
    101
    0
        }
    102
    103
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    104
    0
            match args.args[0].data_type() {
    105
                DataType::Utf8View | DataType::Utf8 => {
    106
    0
                    make_scalar_function(overlay::<i32>, vec![])(&args.args)
    107
                }
    108
                DataType::LargeUtf8 => {
    109
    0
                    make_scalar_function(overlay::<i64>, vec![])(&args.args)
    110
                }
    111
    0
                other => exec_err!("Unsupported data type {other:?} for function overlay"),
    112
            }
    113
    0
        }
    114
    115
    0
        fn documentation(&self) -> Option<&Documentation> {
    116
    0
            self.doc()
    117
    0
        }
    118
    }
    119
    120
    macro_rules! process_overlay {
    121
        // For the three-argument case
    122
        ($string_array:expr, $characters_array:expr, $pos_num:expr) => {{
    123
            $string_array
    124
            .iter()
    125
            .zip($characters_array.iter())
    126
            .zip($pos_num.iter())
    127
    0
            .map(|((string, characters), start_pos)| {
    128
    0
                match (string, characters, start_pos) {
    129
    0
                    (Some(string), Some(characters), Some(start_pos)) => {
    130
    0
                        let string_len = string.chars().count();
    131
    0
                        let characters_len = characters.chars().count();
    132
    0
                        let replace_len = characters_len as i64;
    133
    0
                        let mut res =
    134
    0
                            String::with_capacity(string_len.max(characters_len));
    135
    136
                        //as sql replace index start from 1 while string index start from 0
    137
    0
                        if start_pos > 1 && start_pos - 1 < string_len as i64 {
    138
    0
                            let start = (start_pos - 1) as usize;
    139
    0
                            res.push_str(&string[..start]);
    140
    0
                        }
    141
    0
                        res.push_str(characters);
    142
                        // if start + replace_len - 1 >= string_length, just to string end
    143
    0
                        if start_pos + replace_len - 1 < string_len as i64 {
    144
    0
                            let end = (start_pos + replace_len - 1) as usize;
    145
    0
                            res.push_str(&string[end..]);
    146
    0
                        }
    147
    0
                        Ok(Some(res))
    148
                    }
    149
    0
                    _ => Ok(None),
    150
                }
    151
    0
            })
    152
            .collect::<Result<GenericStringArray<T>>>()
    153
        }};
    154
    155
        // For the four-argument case
    156
        ($string_array:expr, $characters_array:expr, $pos_num:expr, $len_num:expr) => {{
    157
            $string_array
    158
            .iter()
    159
            .zip($characters_array.iter())
    160
            .zip($pos_num.iter())
    161
            .zip($len_num.iter())
    162
    0
            .map(|(((string, characters), start_pos), len)| {
    163
    0
                match (string, characters, start_pos, len) {
    164
    0
                    (Some(string), Some(characters), Some(start_pos), Some(len)) => {
    165
    0
                        let string_len = string.chars().count();
    166
    0
                        let characters_len = characters.chars().count();
    167
    0
                        let replace_len = len.min(string_len as i64);
    168
    0
                        let mut res =
    169
    0
                            String::with_capacity(string_len.max(characters_len));
    170
    171
                        //as sql replace index start from 1 while string index start from 0
    172
    0
                        if start_pos > 1 && start_pos - 1 < string_len as i64 {
    173
    0
                            let start = (start_pos - 1) as usize;
    174
    0
                            res.push_str(&string[..start]);
    175
    0
                        }
    176
    0
                        res.push_str(characters);
    177
                        // if start + replace_len - 1 >= string_length, just to string end
    178
    0
                        if start_pos + replace_len - 1 < string_len as i64 {
    179
    0
                            let end = (start_pos + replace_len - 1) as usize;
    180
    0
                            res.push_str(&string[end..]);
    181
    0
                        }
    182
    0
                        Ok(Some(res))
    183
                    }
    184
    0
                    _ => Ok(None),
    185
                }
    186
    0
            })
    187
            .collect::<Result<GenericStringArray<T>>>()
    188
        }};
    189
    }
    190
    191
    /// OVERLAY(string1 PLACING string2 FROM integer FOR integer2)
    192
    /// Replaces a substring of string1 with string2 starting at the integer bit
    193
    /// pgsql overlay('Txxxxas' placing 'hom' from 2 for 4) → Thomas
    194
    /// overlay('Txxxxas' placing 'hom' from 2) -> Thomxas, without for option, str2's len is instead
    195
    0
    fn overlay<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
    196
    0
        let use_string_view = args[0].data_type() == &DataType::Utf8View;
    197
    0
        if use_string_view {
    198
    0
            string_view_overlay::<T>(args)
    199
        } else {
    200
    0
            string_overlay::<T>(args)
    201
        }
    202
    0
    }
    203
    204
    0
    pub fn string_overlay<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
    205
    0
        match args.len() {
    206
            3 => {
    207
    0
                let string_array = as_generic_string_array::<T>(&args[0])?;
    208
    0
                let characters_array = as_generic_string_array::<T>(&args[1])?;
    209
    0
                let pos_num = as_int64_array(&args[2])?;
    210
    211
    0
                let result = process_overlay!(string_array, characters_array, pos_num)?;
    212
    0
                Ok(Arc::new(result) as ArrayRef)
    213
            }
    214
            4 => {
    215
    0
                let string_array = as_generic_string_array::<T>(&args[0])?;
    216
    0
                let characters_array = as_generic_string_array::<T>(&args[1])?;
    217
    0
                let pos_num = as_int64_array(&args[2])?;
    218
    0
                let len_num = as_int64_array(&args[3])?;
    219
    220
    0
                let result =
    221
    0
                    process_overlay!(string_array, characters_array, pos_num, len_num)?;
    222
    0
                Ok(Arc::new(result) as ArrayRef)
    223
            }
    224
    0
            other => {
    225
    0
                exec_err!("overlay was called with {other} arguments. It requires 3 or 4.")
    226
            }
    227
        }
    228
    0
    }
    229
    230
    0
    pub fn string_view_overlay<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
    231
    0
        match args.len() {
    232
            3 => {
    233
    0
                let string_array = as_string_view_array(&args[0])?;
    234
    0
                let characters_array = as_string_view_array(&args[1])?;
    235
    0
                let pos_num = as_int64_array(&args[2])?;
    236
    237
    0
                let result = process_overlay!(string_array, characters_array, pos_num)?;
    238
    0
                Ok(Arc::new(result) as ArrayRef)
    239
            }
    240
            4 => {
    241
    0
                let string_array = as_string_view_array(&args[0])?;
    242
    0
                let characters_array = as_string_view_array(&args[1])?;
    243
    0
                let pos_num = as_int64_array(&args[2])?;
    244
    0
                let len_num = as_int64_array(&args[3])?;
    245
    246
    0
                let result =
    247
    0
                    process_overlay!(string_array, characters_array, pos_num, len_num)?;
    248
    0
                Ok(Arc::new(result) as ArrayRef)
    249
            }
    250
    0
            other => {
    251
    0
                exec_err!("overlay was called with {other} arguments. It requires 3 or 4.")
    252
            }
    253
        }
    254
    0
    }
    255
    256
    #[cfg(test)]
    257
    mod tests {
    258
        use arrow::array::{Int64Array, StringArray};
    259
    260
        use super::*;
    261
    262
        #[test]
    263
        fn to_overlay() -> Result<()> {
    264
            let string =
    265
                Arc::new(StringArray::from(vec!["123", "abcdefg", "xyz", "Txxxxas"]));
    266
            let replace_string =
    267
                Arc::new(StringArray::from(vec!["abc", "qwertyasdfg", "ijk", "hom"]));
    268
            let start = Arc::new(Int64Array::from(vec![4, 1, 1, 2])); // start
    269
            let end = Arc::new(Int64Array::from(vec![5, 7, 2, 4])); // replace len
    270
    271
            let res = overlay::<i32>(&[string, replace_string, start, end]).unwrap();
    272
            let result = as_generic_string_array::<i32>(&res).unwrap();
    273
            let expected = StringArray::from(vec!["abc", "qwertyasdfg", "ijkz", "Thomas"]);
    274
            assert_eq!(&expected, result);
    275
    276
            Ok(())
    277
        }
    278
    }