datafusion_functions/core/
overlay.rs

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
18use arrow::array::{
19    Array, ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringArrayType,
20    StringViewArray,
21};
22use arrow::buffer::NullBuffer;
23use arrow::datatypes::DataType;
24
25use crate::strings::{
26    BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringWriter,
27};
28use crate::utils::{make_scalar_function, utf8_to_str_type};
29use datafusion_common::cast::{
30    as_generic_string_array, as_int64_array, as_string_view_array,
31};
32use datafusion_common::{Result, exec_err};
33use datafusion_expr::{ColumnarValue, Documentation, TypeSignature, Volatility};
34use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
35use datafusion_macros::user_doc;
36
37#[user_doc(
38    doc_section(label = "String Functions"),
39    description = "Returns the string which is replaced by another string from the specified position and specified count length.",
40    syntax_example = "overlay(str PLACING substr FROM pos [FOR count])",
41    sql_example = r#"```sql
42> select overlay('Txxxxas' placing 'hom' from 2 for 4);
43+--------------------------------------------------------+
44| overlay(Utf8("Txxxxas"),Utf8("hom"),Int64(2),Int64(4)) |
45+--------------------------------------------------------+
46| Thomas                                                 |
47+--------------------------------------------------------+
48```"#,
49    standard_argument(name = "str", prefix = "String"),
50    argument(name = "substr", description = "Substring to replace in str."),
51    argument(
52        name = "pos",
53        description = "The start position to start the replace in str."
54    ),
55    argument(
56        name = "count",
57        description = "The count of characters to be replaced from start position of str. If not specified, will use substr length instead."
58    )
59)]
60#[derive(Debug, PartialEq, Eq, Hash)]
61pub struct OverlayFunc {
62    signature: Signature,
63}
64
65impl Default for OverlayFunc {
66    fn default() -> Self {
67        Self::new()
68    }
69}
70
71impl OverlayFunc {
72    pub fn new() -> Self {
73        use DataType::*;
74        Self {
75            signature: Signature::one_of(
76                vec![
77                    TypeSignature::Exact(vec![Utf8View, Utf8View, Int64, Int64]),
78                    TypeSignature::Exact(vec![Utf8, Utf8, Int64, Int64]),
79                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, Int64, Int64]),
80                    TypeSignature::Exact(vec![Utf8View, Utf8View, Int64]),
81                    TypeSignature::Exact(vec![Utf8, Utf8, Int64]),
82                    TypeSignature::Exact(vec![LargeUtf8, LargeUtf8, Int64]),
83                ],
84                Volatility::Immutable,
85            ),
86        }
87    }
88}
89
90impl ScalarUDFImpl for OverlayFunc {
91    fn name(&self) -> &str {
92        "overlay"
93    }
94
95    fn signature(&self) -> &Signature {
96        &self.signature
97    }
98
99    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
100        utf8_to_str_type(&arg_types[0], "overlay")
101    }
102
103    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
104        match args.args[0].data_type() {
105            DataType::Utf8View | DataType::Utf8 => {
106                make_scalar_function(overlay::<i32>, vec![])(&args.args)
107            }
108            DataType::LargeUtf8 => {
109                make_scalar_function(overlay::<i64>, vec![])(&args.args)
110            }
111            other => exec_err!("Unsupported data type {other:?} for function overlay"),
112        }
113    }
114
115    fn documentation(&self) -> Option<&Documentation> {
116        self.doc()
117    }
118}
119
120/// Computes the byte ranges of `string` to keep around the replaced span: the
121/// prefix is `string[..prefix_end]` and the suffix is `string[suffix_start..]`.
122///
123/// `start_pos` is a 1-based character position; the caller must ensure it is
124/// `>= 1`. `replace_len` is the number of characters of `string` to replace,
125/// and may be negative (in which case `suffix_start <= prefix_end` and the
126/// result re-emits part of the original string).
127///
128/// Matches PostgreSQL semantics for codepoint indices past the end of
129/// `string`: `prefix_end` and `suffix_start` clamp to `string.len()`.
130fn overlay_bounds(string: &str, start_pos: i64, replace_len: i64) -> (usize, usize) {
131    let start_char_idx = start_pos - 1;
132    let end_char_idx = start_char_idx.saturating_add(replace_len);
133
134    if string.is_ascii() {
135        // ASCII fast path: byte index == codepoint index.
136        let len = string.len() as i64;
137        let prefix_end = start_char_idx.clamp(0, len) as usize;
138        let suffix_start = end_char_idx.clamp(0, len) as usize;
139        return (prefix_end, suffix_start);
140    }
141
142    let prefix_target = usize::try_from(start_char_idx).unwrap_or(usize::MAX);
143    let suffix_target = usize::try_from(end_char_idx.max(0)).unwrap_or(usize::MAX);
144    let target_max = prefix_target.max(suffix_target);
145
146    // Single forward pass over codepoint boundaries records both targets.
147    // Either target falls through to `string.len()` if past the codepoint
148    // count.
149    let mut prefix_byte = string.len();
150    let mut suffix_byte = string.len();
151    for (count, (byte_idx, _)) in string.char_indices().enumerate() {
152        if count == prefix_target {
153            prefix_byte = byte_idx;
154        }
155        if count == suffix_target {
156            suffix_byte = byte_idx;
157        }
158        if count == target_max {
159            break;
160        }
161    }
162    (prefix_byte, suffix_byte)
163}
164
165/// Appends the overlay result for one non-null row into `builder`.
166#[inline]
167fn apply_overlay<B: BulkNullStringArrayBuilder>(
168    string: &str,
169    characters: &str,
170    start_pos: i64,
171    replace_len: i64,
172    builder: &mut B,
173) -> Result<()> {
174    if start_pos < 1 {
175        return exec_err!("overlay start position must be at least 1: {start_pos}");
176    }
177    let (prefix_end, suffix_start) = overlay_bounds(string, start_pos, replace_len);
178    builder.append_with(|w| {
179        w.write_str(&string[..prefix_end]);
180        w.write_str(characters);
181        w.write_str(&string[suffix_start..]);
182    });
183    Ok(())
184}
185
186#[inline]
187fn char_count(characters: &str) -> i64 {
188    if characters.is_ascii() {
189        characters.len() as i64
190    } else {
191        characters.chars().count() as i64
192    }
193}
194
195/// `OVERLAY(string PLACING substring FROM start [FOR count])`
196///
197/// Replaces a region of `string` with `substring`, starting at the 1-based
198/// character position `start`. If `count` is supplied, that many characters
199/// of `string` are replaced; otherwise `count` defaults to the character
200/// length of `substring`.
201///
202/// ```text
203/// overlay('Txxxxas' placing 'hom' from 2 for 4) → 'Thomas'
204/// overlay('Txxxxas' placing 'hom' from 2)       → 'Thomxas'
205/// ```
206fn overlay<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
207    if !matches!(args.len(), 3 | 4) {
208        return exec_err!(
209            "overlay was called with {} arguments. It requires 3 or 4.",
210            args.len()
211        );
212    }
213    let pos_array = as_int64_array(&args[2])?;
214    let len_array = if args.len() == 4 {
215        Some(as_int64_array(&args[3])?)
216    } else {
217        None
218    };
219
220    if args[0].data_type() == &DataType::Utf8View {
221        let string_array = as_string_view_array(&args[0])?;
222        let characters_array = as_string_view_array(&args[1])?;
223        let data_capacity = visible_view_bytes(string_array)
224            .saturating_add(visible_view_bytes(characters_array));
225        let builder = GenericStringArrayBuilder::<i32>::with_capacity(
226            string_array.len(),
227            data_capacity,
228        );
229        overlay_inner(
230            string_array,
231            characters_array,
232            pos_array,
233            len_array,
234            builder,
235        )
236    } else {
237        let string_array = as_generic_string_array::<T>(&args[0])?;
238        let characters_array = as_generic_string_array::<T>(&args[1])?;
239        let data_capacity = visible_offset_bytes(string_array)
240            .saturating_add(visible_offset_bytes(characters_array));
241        let builder = GenericStringArrayBuilder::<T>::with_capacity(
242            string_array.len(),
243            data_capacity,
244        );
245        overlay_inner(
246            string_array,
247            characters_array,
248            pos_array,
249            len_array,
250            builder,
251        )
252    }
253}
254
255/// Drives the per-row OVERLAY computation. A null in any input array
256/// produces a null output.
257fn overlay_inner<'a, V, B>(
258    string_array: V,
259    characters_array: V,
260    pos_array: &Int64Array,
261    len_array: Option<&Int64Array>,
262    mut builder: B,
263) -> Result<ArrayRef>
264where
265    V: StringArrayType<'a, Item = &'a str> + Copy,
266    B: BulkNullStringArrayBuilder,
267{
268    let len = string_array.len();
269    let nulls = NullBuffer::union_many([
270        string_array.nulls(),
271        characters_array.nulls(),
272        pos_array.nulls(),
273        len_array.and_then(|a| a.nulls()),
274    ]);
275
276    if let Some(nulls_ref) = nulls.as_ref() {
277        for i in 0..len {
278            if nulls_ref.is_null(i) {
279                builder.append_placeholder();
280                continue;
281            }
282            // SAFETY: `i < len`, and null bitmap check implies not-null
283            let string = unsafe { string_array.value_unchecked(i) };
284            let characters = unsafe { characters_array.value_unchecked(i) };
285            let start_pos = unsafe { pos_array.value_unchecked(i) };
286            let replace_len = match len_array {
287                Some(arr) => unsafe { arr.value_unchecked(i) },
288                None => char_count(characters),
289            };
290            apply_overlay(string, characters, start_pos, replace_len, &mut builder)?;
291        }
292    } else {
293        for i in 0..len {
294            // SAFETY: `i < len`, and no null bitmap means no nulls
295            let string = unsafe { string_array.value_unchecked(i) };
296            let characters = unsafe { characters_array.value_unchecked(i) };
297            let start_pos = unsafe { pos_array.value_unchecked(i) };
298            let replace_len = match len_array {
299                Some(arr) => unsafe { arr.value_unchecked(i) },
300                None => char_count(characters),
301            };
302            apply_overlay(string, characters, start_pos, replace_len, &mut builder)?;
303        }
304    }
305    builder.finish(nulls)
306}
307
308/// Bytes referenced by the visible window of `array`, computed from the
309/// per-view lengths.
310fn visible_view_bytes(array: &StringViewArray) -> usize {
311    array.lengths().map(|l| l as usize).sum()
312}
313
314/// Bytes referenced by the visible window of `array`, derived from the offset
315/// buffer.
316fn visible_offset_bytes<T: OffsetSizeTrait>(array: &GenericStringArray<T>) -> usize {
317    let offsets = array.value_offsets();
318    // `value_offsets()` always has `array.len() + 1` entries (≥1).
319    let first = offsets.first().copied().unwrap_or_default();
320    let last = offsets.last().copied().unwrap_or_default();
321    last.as_usize() - first.as_usize()
322}
323
324#[cfg(test)]
325mod tests {
326    use std::sync::Arc;
327
328    use arrow::array::StringArray;
329
330    use super::*;
331
332    #[test]
333    fn to_overlay() -> Result<()> {
334        let string =
335            Arc::new(StringArray::from(vec!["123", "abcdefg", "xyz", "Txxxxas"]));
336        let replace_string =
337            Arc::new(StringArray::from(vec!["abc", "qwertyasdfg", "ijk", "hom"]));
338        let start = Arc::new(Int64Array::from(vec![4, 1, 1, 2])); // start
339        let end = Arc::new(Int64Array::from(vec![5, 7, 2, 4])); // replace len
340
341        let res = overlay::<i32>(&[string, replace_string, start, end]).unwrap();
342        let result = as_generic_string_array::<i32>(&res).unwrap();
343        // First row: start=4 is past the end of "123" (len 3). PostgreSQL
344        // takes the whole string as prefix and appends the replacement.
345        let expected = StringArray::from(vec!["123abc", "qwertyasdfg", "ijkz", "Thomas"]);
346        assert_eq!(&expected, result);
347
348        Ok(())
349    }
350}
datafusion_functions/core/overlay.rs

datafusion_functions/core/
overlay.rs
