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::{exec_err, Result};
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)]
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    let result = string_array_iter
152        .zip(from_array_iter)
153        .zip(to_array_iter)
154        .map(|((string, from), to)| match (string, from, to) {
155            (Some(string), Some(from), Some(to)) => {
156                // create a hashmap of [char, index] to change from O(n) to O(1) for from list
157                let from_map: HashMap<&str, usize> = from
158                    .graphemes(true)
159                    .collect::<Vec<&str>>()
160                    .iter()
161                    .enumerate()
162                    .map(|(index, c)| (c.to_owned(), index))
163                    .collect();
164
165                let to = to.graphemes(true).collect::<Vec<&str>>();
166
167                Some(
168                    string
169                        .graphemes(true)
170                        .collect::<Vec<&str>>()
171                        .iter()
172                        .flat_map(|c| match from_map.get(*c) {
173                            Some(n) => to.get(*n).copied(),
174                            None => Some(*c),
175                        })
176                        .collect::<Vec<&str>>()
177                        .concat(),
178                )
179            }
180            _ => None,
181        })
182        .collect::<GenericStringArray<T>>();
183
184    Ok(Arc::new(result) as ArrayRef)
185}
186
187#[cfg(test)]
188mod tests {
189    use arrow::array::{Array, StringArray};
190    use arrow::datatypes::DataType::Utf8;
191
192    use datafusion_common::{Result, ScalarValue};
193    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
194
195    use crate::unicode::translate::TranslateFunc;
196    use crate::utils::test::test_function;
197
198    #[test]
199    fn test_functions() -> Result<()> {
200        test_function!(
201            TranslateFunc::new(),
202            vec![
203                ColumnarValue::Scalar(ScalarValue::from("12345")),
204                ColumnarValue::Scalar(ScalarValue::from("143")),
205                ColumnarValue::Scalar(ScalarValue::from("ax"))
206            ],
207            Ok(Some("a2x5")),
208            &str,
209            Utf8,
210            StringArray
211        );
212        test_function!(
213            TranslateFunc::new(),
214            vec![
215                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
216                ColumnarValue::Scalar(ScalarValue::from("143")),
217                ColumnarValue::Scalar(ScalarValue::from("ax"))
218            ],
219            Ok(None),
220            &str,
221            Utf8,
222            StringArray
223        );
224        test_function!(
225            TranslateFunc::new(),
226            vec![
227                ColumnarValue::Scalar(ScalarValue::from("12345")),
228                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
229                ColumnarValue::Scalar(ScalarValue::from("ax"))
230            ],
231            Ok(None),
232            &str,
233            Utf8,
234            StringArray
235        );
236        test_function!(
237            TranslateFunc::new(),
238            vec![
239                ColumnarValue::Scalar(ScalarValue::from("12345")),
240                ColumnarValue::Scalar(ScalarValue::from("143")),
241                ColumnarValue::Scalar(ScalarValue::Utf8(None))
242            ],
243            Ok(None),
244            &str,
245            Utf8,
246            StringArray
247        );
248        test_function!(
249            TranslateFunc::new(),
250            vec![
251                ColumnarValue::Scalar(ScalarValue::from("é2íñ5")),
252                ColumnarValue::Scalar(ScalarValue::from("éñí")),
253                ColumnarValue::Scalar(ScalarValue::from("óü")),
254            ],
255            Ok(Some("ó2ü5")),
256            &str,
257            Utf8,
258            StringArray
259        );
260        #[cfg(not(feature = "unicode_expressions"))]
261        test_function!(
262            TranslateFunc::new(),
263            vec![
264                ColumnarValue::Scalar(ScalarValue::from("12345")),
265                ColumnarValue::Scalar(ScalarValue::from("143")),
266                ColumnarValue::Scalar(ScalarValue::from("ax")),
267            ],
268            internal_err!(
269                "function translate requires compilation with feature flag: unicode_expressions."
270            ),
271            &str,
272            Utf8,
273            StringArray
274        );
275
276        Ok(())
277    }
278}
datafusion_functions/unicode/translate.rs

datafusion_functions/unicode/
translate.rs
