datafusion_functions/string/
chr.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::sync::Arc;
19
20use arrow::array::{Array, ArrayRef, Int64Array};
21use arrow::datatypes::DataType;
22use arrow::datatypes::DataType::Int64;
23use arrow::datatypes::DataType::Utf8;
24
25use crate::strings::GenericStringArrayBuilder;
26use datafusion_common::cast::as_int64_array;
27use datafusion_common::utils::take_function_args;
28use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
29use datafusion_expr::{ColumnarValue, Documentation, Volatility};
30use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
31use datafusion_macros::user_doc;
32
33/// Returns the character with the given code.
34/// chr(65) = 'A'
35fn chr_array(integer_array: &Int64Array) -> Result<ArrayRef> {
36    let len = integer_array.len();
37    let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(
38        len, // 1 byte per character, assuming that is the common case
39        len,
40    );
41
42    let mut buf = [0u8; 4];
43    let nulls = integer_array.nulls();
44
45    if let Some(n) = nulls {
46        for i in 0..len {
47            if n.is_null(i) {
48                builder.append_placeholder();
49                continue;
50            }
51            // SAFETY: bounds + null check above.
52            let integer = unsafe { integer_array.value_unchecked(i) };
53            if let Ok(u) = u32::try_from(integer)
54                && let Some(c) = core::char::from_u32(u)
55            {
56                builder.append_value(c.encode_utf8(&mut buf));
57                continue;
58            }
59            return exec_err!("invalid Unicode scalar value: {integer}");
60        }
61    } else {
62        for i in 0..len {
63            // SAFETY: no null buffer means every index is valid.
64            let integer = unsafe { integer_array.value_unchecked(i) };
65            if let Ok(u) = u32::try_from(integer)
66                && let Some(c) = core::char::from_u32(u)
67            {
68                builder.append_value(c.encode_utf8(&mut buf));
69                continue;
70            }
71            return exec_err!("invalid Unicode scalar value: {integer}");
72        }
73    }
74
75    Ok(Arc::new(builder.finish(nulls.cloned())?) as ArrayRef)
76}
77
78#[user_doc(
79    doc_section(label = "String Functions"),
80    description = "Returns a string containing the character with the specified Unicode scalar value.",
81    syntax_example = "chr(expression)",
82    sql_example = r#"```sql
83> select chr(128640);
84+--------------------+
85| chr(Int64(128640)) |
86+--------------------+
87| 🚀                 |
88+--------------------+
89```"#,
90    standard_argument(name = "expression", prefix = "String"),
91    related_udf(name = "ascii")
92)]
93#[derive(Debug, PartialEq, Eq, Hash)]
94pub struct ChrFunc {
95    signature: Signature,
96}
97
98impl Default for ChrFunc {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104impl ChrFunc {
105    pub fn new() -> Self {
106        Self {
107            signature: Signature::uniform(1, vec![Int64], Volatility::Immutable),
108        }
109    }
110}
111
112impl ScalarUDFImpl for ChrFunc {
113    fn name(&self) -> &str {
114        "chr"
115    }
116
117    fn signature(&self) -> &Signature {
118        &self.signature
119    }
120
121    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
122        Ok(Utf8)
123    }
124
125    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
126        let [arg] = take_function_args(self.name(), args.args)?;
127
128        match arg {
129            ColumnarValue::Scalar(ScalarValue::Int64(Some(code_point))) => {
130                if let Ok(u) = u32::try_from(code_point)
131                    && let Some(c) = core::char::from_u32(u)
132                {
133                    Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
134                        c.to_string(),
135                    ))))
136                } else {
137                    exec_err!("invalid Unicode scalar value: {code_point}")
138                }
139            }
140            ColumnarValue::Scalar(ScalarValue::Int64(None)) => {
141                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
142            }
143            ColumnarValue::Array(array) => {
144                let integer_array = as_int64_array(&array)?;
145                Ok(ColumnarValue::Array(chr_array(integer_array)?))
146            }
147            other => internal_err!(
148                "Unexpected data type {:?} for function chr",
149                other.data_type()
150            ),
151        }
152    }
153
154    fn documentation(&self) -> Option<&Documentation> {
155        self.doc()
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    use arrow::array::{Array, StringArray};
164    use arrow::datatypes::Field;
165    use datafusion_common::assert_contains;
166    use datafusion_common::config::ConfigOptions;
167
168    fn invoke_chr(arg: ColumnarValue, number_rows: usize) -> Result<ColumnarValue> {
169        ChrFunc::new().invoke_with_args(ScalarFunctionArgs {
170            args: vec![arg],
171            arg_fields: vec![Field::new("a", Int64, true).into()],
172            number_rows,
173            return_field: Field::new("f", Utf8, true).into(),
174            config_options: Arc::new(ConfigOptions::default()),
175        })
176    }
177
178    #[test]
179    fn test_chr_normal() {
180        let input = Arc::new(Int64Array::from(vec![
181            Some(0),        // \u{0000}
182            Some(65),       // A
183            Some(66),       // B
184            Some(67),       // C
185            Some(128640),   // 🚀
186            Some(8364),     // €
187            Some(945),      // α
188            None,           // NULL
189            Some(32),       // space
190            Some(10),       // newline
191            Some(9),        // tab
192            Some(0x10FFFF), // 0x10FFFF, the largest Unicode code point
193        ]));
194
195        let result = invoke_chr(ColumnarValue::Array(input), 12).unwrap();
196        let ColumnarValue::Array(arr) = result else {
197            panic!("Expected array");
198        };
199        let string_array = arr.as_any().downcast_ref::<StringArray>().unwrap();
200
201        let expected = [
202            "\u{0000}",
203            "A",
204            "B",
205            "C",
206            "🚀",
207            "€",
208            "α",
209            "",
210            " ",
211            "\n",
212            "\t",
213            "\u{10ffff}",
214        ];
215
216        assert_eq!(string_array.len(), expected.len());
217        assert_eq!(string_array.null_count(), 1);
218        assert!(string_array.is_null(7));
219        for (i, e) in expected.iter().enumerate() {
220            if i == 7 {
221                continue;
222            }
223            assert!(!string_array.is_null(i));
224            assert_eq!(string_array.value(i), *e);
225        }
226    }
227
228    #[test]
229    fn test_chr_error() {
230        let input = Arc::new(Int64Array::from(vec![i64::MAX]));
231        let result = invoke_chr(ColumnarValue::Array(input), 1);
232        assert!(result.is_err());
233        assert_contains!(
234            result.err().unwrap().to_string(),
235            "invalid Unicode scalar value: 9223372036854775807"
236        );
237
238        let input = Arc::new(Int64Array::from(vec![0x10FFFF + 1]));
239        let result = invoke_chr(ColumnarValue::Array(input), 1);
240        assert!(result.is_err());
241        assert_contains!(
242            result.err().unwrap().to_string(),
243            "invalid Unicode scalar value: 1114112"
244        );
245
246        let input = Arc::new(Int64Array::from(vec![0xD800 + 1]));
247        let result = invoke_chr(ColumnarValue::Array(input), 1);
248        assert!(result.is_err());
249        assert_contains!(
250            result.err().unwrap().to_string(),
251            "invalid Unicode scalar value: 55297"
252        );
253
254        let input = Arc::new(Int64Array::from(vec![i64::MIN + 2i64]));
255        let result = invoke_chr(ColumnarValue::Array(input), 1);
256        assert!(result.is_err());
257        assert_contains!(
258            result.err().unwrap().to_string(),
259            "invalid Unicode scalar value: -9223372036854775806"
260        );
261
262        let input = Arc::new(Int64Array::from(vec![-1]));
263        let result = invoke_chr(ColumnarValue::Array(input), 1);
264        assert!(result.is_err());
265        assert_contains!(
266            result.err().unwrap().to_string(),
267            "invalid Unicode scalar value: -1"
268        );
269
270        let input = Arc::new(Int64Array::from(vec![65, -1, 66]));
271        let result = invoke_chr(ColumnarValue::Array(input), 3);
272        assert!(result.is_err());
273        assert_contains!(
274            result.err().unwrap().to_string(),
275            "invalid Unicode scalar value: -1"
276        );
277    }
278
279    #[test]
280    fn test_chr_empty() {
281        let input = Arc::new(Int64Array::from(Vec::<i64>::new()));
282        let result = invoke_chr(ColumnarValue::Array(input), 0).unwrap();
283        let ColumnarValue::Array(arr) = result else {
284            panic!("Expected array");
285        };
286        let string_array = arr.as_any().downcast_ref::<StringArray>().unwrap();
287        assert_eq!(string_array.len(), 0);
288    }
289
290    #[test]
291    fn test_chr_scalar() {
292        let result =
293            invoke_chr(ColumnarValue::Scalar(ScalarValue::Int64(Some(65))), 1).unwrap();
294
295        match result {
296            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
297                assert_eq!(s, "A");
298            }
299            other => panic!("Unexpected result: {other:?}"),
300        }
301    }
302
303    #[test]
304    fn test_chr_scalar_null() {
305        let result =
306            invoke_chr(ColumnarValue::Scalar(ScalarValue::Int64(None)), 1).unwrap();
307
308        match result {
309            ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {}
310            other => panic!("Unexpected result: {other:?}"),
311        }
312    }
313}
datafusion_functions/string/chr.rs

datafusion_functions/string/
chr.rs
