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 std::any::Any;
19use std::sync::Arc;
20
21use arrow::array::{
22    ArrayAccessor, ArrayIter, ArrayRef, AsArray, GenericStringArray, OffsetSizeTrait,
23};
24use arrow::datatypes::DataType;
25use datafusion_common::HashMap;
26use unicode_segmentation::UnicodeSegmentation;
27
28use crate::utils::{make_scalar_function, utf8_to_str_type};
29use datafusion_common::{Result, exec_err};
30use datafusion_expr::TypeSignature::Exact;
31use datafusion_expr::{
32    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
33};
34use datafusion_macros::user_doc;
35
36#[user_doc(
37    doc_section(label = "String Functions"),
38    description = "Translates characters in a string to specified translation characters.",
39    syntax_example = "translate(str, chars, translation)",
40    sql_example = r#"```sql
41> select translate('twice', 'wic', 'her');
42+--------------------------------------------------+
43| translate(Utf8("twice"),Utf8("wic"),Utf8("her")) |
44+--------------------------------------------------+
45| there                                            |
46+--------------------------------------------------+
47```"#,
48    standard_argument(name = "str", prefix = "String"),
49    argument(name = "chars", description = "Characters to translate."),
50    argument(
51        name = "translation",
52        description = "Translation characters. Translation characters replace only characters at the same position in the **chars** string."
53    )
54)]
55#[derive(Debug, PartialEq, Eq, Hash)]
56pub struct TranslateFunc {
57    signature: Signature,
58}
59
60impl Default for TranslateFunc {
61    fn default() -> Self {
62        Self::new()
63    }
64}
65
66impl TranslateFunc {
67    pub fn new() -> Self {
68        use DataType::*;
69        Self {
70            signature: Signature::one_of(
71                vec![
72                    Exact(vec![Utf8View, Utf8, Utf8]),
73                    Exact(vec![Utf8, Utf8, Utf8]),
74                ],
75                Volatility::Immutable,
76            ),
77        }
78    }
79}
80
81impl ScalarUDFImpl for TranslateFunc {
82    fn as_any(&self) -> &dyn Any {
83        self
84    }
85
86    fn name(&self) -> &str {
87        "translate"
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
95        utf8_to_str_type(&arg_types[0], "translate")
96    }
97
98    fn invoke_with_args(
99        &self,
100        args: datafusion_expr::ScalarFunctionArgs,
101    ) -> Result<ColumnarValue> {
102        make_scalar_function(invoke_translate, vec![])(&args.args)
103    }
104
105    fn documentation(&self) -> Option<&Documentation> {
106        self.doc()
107    }
108}
109
110fn invoke_translate(args: &[ArrayRef]) -> Result<ArrayRef> {
111    match args[0].data_type() {
112        DataType::Utf8View => {
113            let string_array = args[0].as_string_view();
114            let from_array = args[1].as_string::<i32>();
115            let to_array = args[2].as_string::<i32>();
116            translate::<i32, _, _>(string_array, from_array, to_array)
117        }
118        DataType::Utf8 => {
119            let string_array = args[0].as_string::<i32>();
120            let from_array = args[1].as_string::<i32>();
121            let to_array = args[2].as_string::<i32>();
122            translate::<i32, _, _>(string_array, from_array, to_array)
123        }
124        DataType::LargeUtf8 => {
125            let string_array = args[0].as_string::<i64>();
126            let from_array = args[1].as_string::<i64>();
127            let to_array = args[2].as_string::<i64>();
128            translate::<i64, _, _>(string_array, from_array, to_array)
129        }
130        other => {
131            exec_err!("Unsupported data type {other:?} for function translate")
132        }
133    }
134}
135
136/// Replaces each character in string that matches a character in the from set with the corresponding character in the to set. If from is longer than to, occurrences of the extra characters in from are deleted.
137/// translate('12345', '143', 'ax') = 'a2x5'
138fn translate<'a, T: OffsetSizeTrait, V, B>(
139    string_array: V,
140    from_array: B,
141    to_array: B,
142) -> Result<ArrayRef>
143where
144    V: ArrayAccessor<Item = &'a str>,
145    B: ArrayAccessor<Item = &'a str>,
146{
147    let string_array_iter = ArrayIter::new(string_array);
148    let from_array_iter = ArrayIter::new(from_array);
149    let to_array_iter = ArrayIter::new(to_array);
150
151    // Reusable buffers to avoid allocating for each row
152    let mut from_map: HashMap<&str, usize> = HashMap::new();
153    let mut from_graphemes: Vec<&str> = Vec::new();
154    let mut to_graphemes: Vec<&str> = Vec::new();
155    let mut string_graphemes: Vec<&str> = Vec::new();
156    let mut result_graphemes: Vec<&str> = Vec::new();
157
158    let result = string_array_iter
159        .zip(from_array_iter)
160        .zip(to_array_iter)
161        .map(|((string, from), to)| match (string, from, to) {
162            (Some(string), Some(from), Some(to)) => {
163                // Clear and reuse buffers
164                from_map.clear();
165                from_graphemes.clear();
166                to_graphemes.clear();
167                string_graphemes.clear();
168                result_graphemes.clear();
169
170                // Build from_map using reusable buffer
171                from_graphemes.extend(from.graphemes(true));
172                for (index, c) in from_graphemes.iter().enumerate() {
173                    // Ignore characters that already exist in from_map, else insert
174                    from_map.entry(*c).or_insert(index);
175                }
176
177                // Build to_graphemes
178                to_graphemes.extend(to.graphemes(true));
179
180                // Process string and build result
181                string_graphemes.extend(string.graphemes(true));
182                for c in &string_graphemes {
183                    match from_map.get(*c) {
184                        Some(n) => {
185                            if let Some(replacement) = to_graphemes.get(*n) {
186                                result_graphemes.push(*replacement);
187                            }
188                        }
189                        None => result_graphemes.push(*c),
190                    }
191                }
192
193                Some(result_graphemes.concat())
194            }
195            _ => None,
196        })
197        .collect::<GenericStringArray<T>>();
198
199    Ok(Arc::new(result) as ArrayRef)
200}
201
202#[cfg(test)]
203mod tests {
204    use arrow::array::{Array, StringArray};
205    use arrow::datatypes::DataType::Utf8;
206
207    use datafusion_common::{Result, ScalarValue};
208    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
209
210    use crate::unicode::translate::TranslateFunc;
211    use crate::utils::test::test_function;
212
213    #[test]
214    fn test_functions() -> Result<()> {
215        test_function!(
216            TranslateFunc::new(),
217            vec![
218                ColumnarValue::Scalar(ScalarValue::from("12345")),
219                ColumnarValue::Scalar(ScalarValue::from("143")),
220                ColumnarValue::Scalar(ScalarValue::from("ax"))
221            ],
222            Ok(Some("a2x5")),
223            &str,
224            Utf8,
225            StringArray
226        );
227        test_function!(
228            TranslateFunc::new(),
229            vec![
230                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
231                ColumnarValue::Scalar(ScalarValue::from("143")),
232                ColumnarValue::Scalar(ScalarValue::from("ax"))
233            ],
234            Ok(None),
235            &str,
236            Utf8,
237            StringArray
238        );
239        test_function!(
240            TranslateFunc::new(),
241            vec![
242                ColumnarValue::Scalar(ScalarValue::from("12345")),
243                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
244                ColumnarValue::Scalar(ScalarValue::from("ax"))
245            ],
246            Ok(None),
247            &str,
248            Utf8,
249            StringArray
250        );
251        test_function!(
252            TranslateFunc::new(),
253            vec![
254                ColumnarValue::Scalar(ScalarValue::from("12345")),
255                ColumnarValue::Scalar(ScalarValue::from("143")),
256                ColumnarValue::Scalar(ScalarValue::Utf8(None))
257            ],
258            Ok(None),
259            &str,
260            Utf8,
261            StringArray
262        );
263        test_function!(
264            TranslateFunc::new(),
265            vec![
266                ColumnarValue::Scalar(ScalarValue::from("abcabc")),
267                ColumnarValue::Scalar(ScalarValue::from("aa")),
268                ColumnarValue::Scalar(ScalarValue::from("de"))
269            ],
270            Ok(Some("dbcdbc")),
271            &str,
272            Utf8,
273            StringArray
274        );
275        test_function!(
276            TranslateFunc::new(),
277            vec![
278                ColumnarValue::Scalar(ScalarValue::from("é2íñ5")),
279                ColumnarValue::Scalar(ScalarValue::from("éñí")),
280                ColumnarValue::Scalar(ScalarValue::from("óü")),
281            ],
282            Ok(Some("ó2ü5")),
283            &str,
284            Utf8,
285            StringArray
286        );
287        #[cfg(not(feature = "unicode_expressions"))]
288        test_function!(
289            TranslateFunc::new(),
290            vec![
291                ColumnarValue::Scalar(ScalarValue::from("12345")),
292                ColumnarValue::Scalar(ScalarValue::from("143")),
293                ColumnarValue::Scalar(ScalarValue::from("ax")),
294            ],
295            internal_err!(
296                "function translate requires compilation with feature flag: unicode_expressions."
297            ),
298            &str,
299            Utf8,
300            StringArray
301        );
302
303        Ok(())
304    }
305}
datafusion_functions/unicode/translate.rs

datafusion_functions/unicode/
translate.rs
