datafusion_functions/unicode/
translate.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::{Array, ArrayRef, AsArray, GenericStringArray, StringArrayType};
19use arrow::buffer::NullBuffer;
20use arrow::datatypes::DataType;
21use datafusion_common::HashMap;
22
23use super::common::try_as_scalar_str;
24use crate::strings::{
25    BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder,
26    StringWriter,
27};
28use crate::utils::make_scalar_function;
29use datafusion_common::{Result, exec_err};
30use datafusion_expr::TypeSignature::Exact;
31use datafusion_expr::{
32    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
33    Volatility,
34};
35use datafusion_macros::user_doc;
36
37#[user_doc(
38    doc_section(label = "String Functions"),
39    description = "Performs character-wise substitution based on a mapping.",
40    syntax_example = "translate(str, from, to)",
41    sql_example = r#"```sql
42> select translate('twice', 'wic', 'her');
43+--------------------------------------------------+
44| translate(Utf8("twice"),Utf8("wic"),Utf8("her")) |
45+--------------------------------------------------+
46| there                                            |
47+--------------------------------------------------+
48```"#,
49    standard_argument(name = "str", prefix = "String"),
50    argument(name = "from", description = "The characters to be replaced."),
51    argument(
52        name = "to",
53        description = "The characters to replace them with. Each character in **from** that is found in **str** is replaced by the character at the same index in **to**. Any characters in **from** that don't have a corresponding character in **to** are removed. If a character appears more than once in **from**, the first occurrence determines the mapping."
54    )
55)]
56#[derive(Debug, PartialEq, Eq, Hash)]
57pub struct TranslateFunc {
58    signature: Signature,
59}
60
61impl Default for TranslateFunc {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl TranslateFunc {
68    pub fn new() -> Self {
69        use DataType::*;
70        Self {
71            signature: Signature::one_of(
72                vec![
73                    Exact(vec![Utf8View, Utf8, Utf8]),
74                    Exact(vec![Utf8, Utf8, Utf8]),
75                    Exact(vec![LargeUtf8, Utf8, Utf8]),
76                ],
77                Volatility::Immutable,
78            ),
79        }
80    }
81}
82
83impl ScalarUDFImpl for TranslateFunc {
84    fn name(&self) -> &str {
85        "translate"
86    }
87
88    fn signature(&self) -> &Signature {
89        &self.signature
90    }
91
92    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
93        Ok(arg_types[0].clone())
94    }
95
96    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
97        // When from and to are scalars, pre-build the translation map once
98        if let (Some(from_str), Some(to_str)) = (
99            try_as_scalar_str(&args.args[1]),
100            try_as_scalar_str(&args.args[2]),
101        ) {
102            let table = build_translate_table(from_str, to_str);
103
104            let string_array = args.args[0].to_array_of_size(args.number_rows)?;
105            let len = string_array.len();
106
107            let result = match string_array.data_type() {
108                DataType::Utf8View => {
109                    let arr = string_array.as_string_view();
110                    let builder = StringViewArrayBuilder::with_capacity(len);
111                    translate_with_table(&arr, &table, builder)
112                }
113                DataType::Utf8 => {
114                    let arr = string_array.as_string::<i32>();
115                    let builder = GenericStringArrayBuilder::<i32>::with_capacity(
116                        len,
117                        arr.value_data().len(),
118                    );
119                    translate_with_table(&arr, &table, builder)
120                }
121                DataType::LargeUtf8 => {
122                    let arr = string_array.as_string::<i64>();
123                    let builder = GenericStringArrayBuilder::<i64>::with_capacity(
124                        len,
125                        arr.value_data().len(),
126                    );
127                    translate_with_table(&arr, &table, builder)
128                }
129                other => {
130                    return exec_err!(
131                        "Unsupported data type {other:?} for function translate"
132                    );
133                }
134            }?;
135
136            return Ok(ColumnarValue::Array(result));
137        }
138
139        make_scalar_function(invoke_translate, vec![])(&args.args)
140    }
141
142    fn documentation(&self) -> Option<&Documentation> {
143        self.doc()
144    }
145}
146
147fn invoke_translate(args: &[ArrayRef]) -> Result<ArrayRef> {
148    let len = args[0].len();
149    match args[0].data_type() {
150        DataType::Utf8View => {
151            let string_array = args[0].as_string_view();
152            let from_array = args[1].as_string::<i32>();
153            let to_array = args[2].as_string::<i32>();
154            let builder = StringViewArrayBuilder::with_capacity(len);
155            translate(&string_array, from_array, to_array, builder)
156        }
157        DataType::Utf8 => {
158            let string_array = args[0].as_string::<i32>();
159            let from_array = args[1].as_string::<i32>();
160            let to_array = args[2].as_string::<i32>();
161            let builder = GenericStringArrayBuilder::<i32>::with_capacity(
162                len,
163                string_array.value_data().len(),
164            );
165            translate(&string_array, from_array, to_array, builder)
166        }
167        DataType::LargeUtf8 => {
168            let string_array = args[0].as_string::<i64>();
169            let from_array = args[1].as_string::<i32>();
170            let to_array = args[2].as_string::<i32>();
171            let builder = GenericStringArrayBuilder::<i64>::with_capacity(
172                len,
173                string_array.value_data().len(),
174            );
175            translate(&string_array, from_array, to_array, builder)
176        }
177        other => {
178            exec_err!("Unsupported data type {other:?} for function translate")
179        }
180    }
181}
182
183/// Replaces each character in string that matches a character in the from set
184/// with the corresponding character in the to set. If from is longer than to,
185/// occurrences of the extra characters in from are deleted.
186///
187/// translate('12345', '143', 'ax') = 'a2x5'
188fn translate<'a, S, O>(
189    string_array: &S,
190    from_array: &GenericStringArray<i32>,
191    to_array: &GenericStringArray<i32>,
192    mut builder: O,
193) -> Result<ArrayRef>
194where
195    S: StringArrayType<'a>,
196    O: BulkNullStringArrayBuilder,
197{
198    let mut from_map: HashMap<char, Option<char>> = HashMap::new();
199    let len = string_array.len();
200    let nulls = NullBuffer::union_many([
201        string_array.nulls(),
202        from_array.nulls(),
203        to_array.nulls(),
204    ]);
205
206    if let Some(nulls_ref) = nulls.as_ref() {
207        for i in 0..len {
208            if nulls_ref.is_null(i) {
209                builder.append_placeholder();
210                continue;
211            }
212
213            // SAFETY: union of input nulls is non-null at i, so each input is too.
214            let string = unsafe { string_array.value_unchecked(i) };
215            let from = unsafe { from_array.value_unchecked(i) };
216            let to = unsafe { to_array.value_unchecked(i) };
217            append_translated_row(&mut builder, string, from, to, &mut from_map);
218        }
219    } else {
220        for i in 0..len {
221            // SAFETY: i < len, and no input has a null buffer.
222            let string = unsafe { string_array.value_unchecked(i) };
223            let from = unsafe { from_array.value_unchecked(i) };
224            let to = unsafe { to_array.value_unchecked(i) };
225            append_translated_row(&mut builder, string, from, to, &mut from_map);
226        }
227    }
228
229    builder.finish(nulls)
230}
231
232#[inline]
233fn append_translated_row<B: BulkNullStringArrayBuilder>(
234    builder: &mut B,
235    string: &str,
236    from: &str,
237    to: &str,
238    from_map: &mut HashMap<char, Option<char>>,
239) {
240    if let Some(ascii_table) = build_ascii_translate_table(from, to) {
241        append_translated_ascii(builder, string, &ascii_table);
242        return;
243    }
244
245    from_map.clear();
246    let mut to_iter = to.chars();
247    for c in from.chars() {
248        let replacement = to_iter.next();
249        from_map.entry(c).or_insert(replacement);
250    }
251
252    builder.append_with(|w| write_translated_chars(w, string, from_map));
253}
254
255#[inline]
256fn write_translated_chars<W: StringWriter>(
257    w: &mut W,
258    input: &str,
259    from_map: &HashMap<char, Option<char>>,
260) {
261    for c in input.chars() {
262        match from_map.get(&c) {
263            Some(Some(r)) => w.write_char(*r),
264            Some(None) => {} // delete: `from` had no corresponding `to` char
265            None => w.write_char(c),
266        }
267    }
268}
269
270/// Sentinel value in the ASCII translate table indicating the character should
271/// be deleted (the `from` character has no corresponding `to` character).  Any
272/// value > 127 works since valid ASCII is 0–127.
273const ASCII_DELETE: u8 = 0xFF;
274
275/// Lookup table for ASCII-only translation. Entries 0..128 map input bytes to
276/// replacement bytes, or `ASCII_DELETE` if the character should be deleted.
277/// Entries 128..256 map to themselves so non-ASCII bytes pass through
278/// unchanged.
279#[derive(Debug)]
280struct AsciiTranslateTable {
281    map: [u8; 256],
282    has_delete: bool,
283}
284
285/// We use a byte-indexed table when both `from` and `to` strings are ASCII,
286/// otherwise a char-indexed map where `None` means delete.
287#[expect(
288    clippy::large_enum_variant,
289    reason = "one instance per call, passed by reference"
290)]
291enum TranslateTable {
292    Byte(AsciiTranslateTable),
293    Char(HashMap<char, Option<char>>),
294}
295
296#[inline]
297fn build_translate_table(from: &str, to: &str) -> TranslateTable {
298    if let Some(ascii) = build_ascii_translate_table(from, to) {
299        return TranslateTable::Byte(ascii);
300    }
301    let mut from_map: HashMap<char, Option<char>> = HashMap::with_capacity(from.len());
302    let mut to_iter = to.chars();
303    for c in from.chars() {
304        let replacement = to_iter.next();
305        from_map.entry(c).or_insert(replacement);
306    }
307    TranslateTable::Char(from_map)
308}
309
310/// Returns `None` if either string contains non-ASCII characters.
311fn build_ascii_translate_table(from: &str, to: &str) -> Option<AsciiTranslateTable> {
312    if !from.is_ascii() || !to.is_ascii() {
313        return None;
314    }
315
316    let to_bytes = to.as_bytes();
317    let mut map = std::array::from_fn::<u8, 256, _>(|i| i as u8);
318    let mut seen = [false; 128];
319    let mut has_delete = false;
320
321    for (i, from_byte) in from.bytes().enumerate() {
322        let idx = from_byte as usize;
323        if !seen[idx] {
324            seen[idx] = true;
325            if i < to_bytes.len() {
326                map[idx] = to_bytes[i];
327            } else {
328                map[idx] = ASCII_DELETE;
329                has_delete = true;
330            }
331        }
332    }
333
334    Some(AsciiTranslateTable { map, has_delete })
335}
336
337#[inline]
338fn append_translated_ascii<B: BulkNullStringArrayBuilder>(
339    builder: &mut B,
340    input: &str,
341    table: &AsciiTranslateTable,
342) {
343    // Fast path: equal-length byte-to-byte map when no deletions.
344    if !table.has_delete {
345        // SAFETY: ASCII source bytes map to ASCII replacements; non-ASCII
346        // bytes 128..256 map to themselves, so multi-byte UTF-8 sequences
347        // pass through unchanged. Output length equals input length and
348        // remains valid UTF-8.
349        unsafe {
350            builder.append_byte_map(input.as_bytes(), |b| table.map[b as usize]);
351        }
352    } else {
353        builder.append_with(|w| write_translated_ascii(w, input, table));
354    }
355}
356
357#[inline]
358fn write_translated_ascii<W: StringWriter>(
359    w: &mut W,
360    input: &str,
361    table: &AsciiTranslateTable,
362) {
363    let bytes = input.as_bytes();
364    let mut copy_start = 0;
365
366    for (i, &b) in bytes.iter().enumerate() {
367        let mapped = table.map[b as usize];
368        if mapped == b {
369            continue;
370        }
371
372        if copy_start < i {
373            w.write_str(&input[copy_start..i]);
374        }
375        if mapped != ASCII_DELETE {
376            w.write_char(mapped as char);
377        }
378        copy_start = i + 1;
379    }
380
381    if copy_start < input.len() {
382        w.write_str(&input[copy_start..]);
383    }
384}
385
386fn translate_with_table<'a, S, O>(
387    string_array: &S,
388    table: &TranslateTable,
389    mut builder: O,
390) -> Result<ArrayRef>
391where
392    S: StringArrayType<'a>,
393    O: BulkNullStringArrayBuilder,
394{
395    let len = string_array.len();
396    let nulls = string_array.nulls().cloned();
397
398    if let Some(nulls_ref) = nulls.as_ref() {
399        for i in 0..len {
400            if nulls_ref.is_null(i) {
401                builder.append_placeholder();
402                continue;
403            }
404
405            // SAFETY: input null buffer is non-null at i.
406            let s = unsafe { string_array.value_unchecked(i) };
407            apply_translate_table(&mut builder, s, table);
408        }
409    } else {
410        for i in 0..len {
411            // SAFETY: no null buffer means every index is valid.
412            let s = unsafe { string_array.value_unchecked(i) };
413            apply_translate_table(&mut builder, s, table);
414        }
415    }
416
417    builder.finish(nulls)
418}
419
420#[inline]
421fn apply_translate_table<B: BulkNullStringArrayBuilder>(
422    builder: &mut B,
423    input: &str,
424    table: &TranslateTable,
425) {
426    match table {
427        TranslateTable::Byte(t) => append_translated_ascii(builder, input, t),
428        TranslateTable::Char(m) => {
429            builder.append_with(|w| write_translated_chars(w, input, m))
430        }
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use std::sync::Arc;
437
438    use arrow::array::{Array, ArrayRef, StringArray, StringViewArray};
439    use arrow::datatypes::DataType::{Utf8, Utf8View};
440
441    use datafusion_common::{Result, ScalarValue};
442    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
443
444    use crate::unicode::translate::TranslateFunc;
445    use crate::utils::test::test_function;
446
447    #[test]
448    fn test_functions() -> Result<()> {
449        test_function!(
450            TranslateFunc::new(),
451            vec![
452                ColumnarValue::Scalar(ScalarValue::from("12345")),
453                ColumnarValue::Scalar(ScalarValue::from("143")),
454                ColumnarValue::Scalar(ScalarValue::from("ax"))
455            ],
456            Ok(Some("a2x5")),
457            &str,
458            Utf8,
459            StringArray
460        );
461        test_function!(
462            TranslateFunc::new(),
463            vec![
464                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
465                ColumnarValue::Scalar(ScalarValue::from("143")),
466                ColumnarValue::Scalar(ScalarValue::from("ax"))
467            ],
468            Ok(None),
469            &str,
470            Utf8,
471            StringArray
472        );
473        test_function!(
474            TranslateFunc::new(),
475            vec![
476                ColumnarValue::Scalar(ScalarValue::from("12345")),
477                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
478                ColumnarValue::Scalar(ScalarValue::from("ax"))
479            ],
480            Ok(None),
481            &str,
482            Utf8,
483            StringArray
484        );
485        test_function!(
486            TranslateFunc::new(),
487            vec![
488                ColumnarValue::Scalar(ScalarValue::from("12345")),
489                ColumnarValue::Scalar(ScalarValue::from("143")),
490                ColumnarValue::Scalar(ScalarValue::Utf8(None))
491            ],
492            Ok(None),
493            &str,
494            Utf8,
495            StringArray
496        );
497        test_function!(
498            TranslateFunc::new(),
499            vec![
500                ColumnarValue::Scalar(ScalarValue::from("abcabc")),
501                ColumnarValue::Scalar(ScalarValue::from("aa")),
502                ColumnarValue::Scalar(ScalarValue::from("de"))
503            ],
504            Ok(Some("dbcdbc")),
505            &str,
506            Utf8,
507            StringArray
508        );
509        test_function!(
510            TranslateFunc::new(),
511            vec![
512                ColumnarValue::Scalar(ScalarValue::from("é2íñ5")),
513                ColumnarValue::Scalar(ScalarValue::from("éñí")),
514                ColumnarValue::Scalar(ScalarValue::from("óü")),
515            ],
516            Ok(Some("ó2ü5")),
517            &str,
518            Utf8,
519            StringArray
520        );
521        // Non-ASCII input with ASCII scalar from/to.
522        test_function!(
523            TranslateFunc::new(),
524            vec![
525                ColumnarValue::Scalar(ScalarValue::from("café")),
526                ColumnarValue::Scalar(ScalarValue::from("ae")),
527                ColumnarValue::Scalar(ScalarValue::from("AE"))
528            ],
529            Ok(Some("cAfé")),
530            &str,
531            Utf8,
532            StringArray
533        );
534        // Utf8View input should produce Utf8View output
535        test_function!(
536            TranslateFunc::new(),
537            vec![
538                ColumnarValue::Scalar(ScalarValue::Utf8View(Some("12345".into()))),
539                ColumnarValue::Scalar(ScalarValue::from("143")),
540                ColumnarValue::Scalar(ScalarValue::from("ax"))
541            ],
542            Ok(Some("a2x5")),
543            &str,
544            Utf8View,
545            StringViewArray
546        );
547        // Null Utf8View input
548        test_function!(
549            TranslateFunc::new(),
550            vec![
551                ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
552                ColumnarValue::Scalar(ScalarValue::from("143")),
553                ColumnarValue::Scalar(ScalarValue::from("ax"))
554            ],
555            Ok(None),
556            &str,
557            Utf8View,
558            StringViewArray
559        );
560        // Non-ASCII Utf8View input
561        test_function!(
562            TranslateFunc::new(),
563            vec![
564                ColumnarValue::Scalar(ScalarValue::Utf8View(Some("é2íñ5".into()))),
565                ColumnarValue::Scalar(ScalarValue::from("éñí")),
566                ColumnarValue::Scalar(ScalarValue::from("óü"))
567            ],
568            Ok(Some("ó2ü5")),
569            &str,
570            Utf8View,
571            StringViewArray
572        );
573
574        #[cfg(not(feature = "unicode_expressions"))]
575        test_function!(
576            TranslateFunc::new(),
577            vec![
578                ColumnarValue::Scalar(ScalarValue::from("12345")),
579                ColumnarValue::Scalar(ScalarValue::from("143")),
580                ColumnarValue::Scalar(ScalarValue::from("ax")),
581            ],
582            internal_err!(
583                "function translate requires compilation with feature flag: unicode_expressions."
584            ),
585            &str,
586            Utf8,
587            StringArray
588        );
589
590        Ok(())
591    }
592
593    #[test]
594    fn test_array_args_with_nulls() -> Result<()> {
595        let string_array = Arc::new(StringArray::from(vec![
596            Some("café!"),
597            Some("abc"),
598            Some("abc"),
599        ])) as ArrayRef;
600        let from_array =
601            Arc::new(StringArray::from(vec![Some("!"), Some("a"), None])) as ArrayRef;
602        let to_array =
603            Arc::new(StringArray::from(vec![Some(""), Some("x"), Some("y")])) as ArrayRef;
604
605        let result = super::invoke_translate(&[string_array, from_array, to_array])?;
606        let result = result.as_any().downcast_ref::<StringArray>().unwrap();
607
608        assert_eq!(result.len(), 3);
609        assert_eq!(result.value(0), "café");
610        assert_eq!(result.value(1), "xbc");
611        assert!(result.is_null(2));
612
613        Ok(())
614    }
615}
datafusion_functions/unicode/translate.rs

datafusion_functions/unicode/
translate.rs
