datafusion_functions/string/
ascii.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 crate::utils::make_scalar_function;
19use arrow::array::{ArrayAccessor, ArrayIter, ArrayRef, AsArray, Int32Array};
20use arrow::datatypes::DataType;
21use arrow::error::ArrowError;
22use datafusion_common::types::logical_string;
23use datafusion_common::{internal_err, Result};
24use datafusion_expr::{ColumnarValue, Documentation, TypeSignatureClass};
25use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
26use datafusion_expr_common::signature::Coercion;
27use datafusion_macros::user_doc;
28use std::any::Any;
29use std::sync::Arc;
30
31#[user_doc(
32    doc_section(label = "String Functions"),
33    description = "Returns the Unicode character code of the first character in a string.",
34    syntax_example = "ascii(str)",
35    sql_example = r#"```sql
36> select ascii('abc');
37+--------------------+
38| ascii(Utf8("abc")) |
39+--------------------+
40| 97                 |
41+--------------------+
42> select ascii('🚀');
43+-------------------+
44| ascii(Utf8("🚀")) |
45+-------------------+
46| 128640            |
47+-------------------+
48```"#,
49    standard_argument(name = "str", prefix = "String"),
50    related_udf(name = "chr")
51)]
52#[derive(Debug)]
53pub struct AsciiFunc {
54    signature: Signature,
55}
56
57impl Default for AsciiFunc {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl AsciiFunc {
64    pub fn new() -> Self {
65        Self {
66            signature: Signature::coercible(
67                vec![Coercion::new_exact(TypeSignatureClass::Native(
68                    logical_string(),
69                ))],
70                Volatility::Immutable,
71            ),
72        }
73    }
74}
75
76impl ScalarUDFImpl for AsciiFunc {
77    fn as_any(&self) -> &dyn Any {
78        self
79    }
80
81    fn name(&self) -> &str {
82        "ascii"
83    }
84
85    fn signature(&self) -> &Signature {
86        &self.signature
87    }
88
89    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
90        use DataType::*;
91
92        Ok(Int32)
93    }
94
95    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
96        make_scalar_function(ascii, vec![])(&args.args)
97    }
98
99    fn documentation(&self) -> Option<&Documentation> {
100        self.doc()
101    }
102}
103
104fn calculate_ascii<'a, V>(array: V) -> Result<ArrayRef, ArrowError>
105where
106    V: ArrayAccessor<Item = &'a str>,
107{
108    let iter = ArrayIter::new(array);
109    let result = iter
110        .map(|string| {
111            string.map(|s| {
112                let mut chars = s.chars();
113                chars.next().map_or(0, |v| v as i32)
114            })
115        })
116        .collect::<Int32Array>();
117
118    Ok(Arc::new(result) as ArrayRef)
119}
120
121/// Returns the numeric code of the first character of the argument.
122pub fn ascii(args: &[ArrayRef]) -> Result<ArrayRef> {
123    match args[0].data_type() {
124        DataType::Utf8 => {
125            let string_array = args[0].as_string::<i32>();
126            Ok(calculate_ascii(string_array)?)
127        }
128        DataType::LargeUtf8 => {
129            let string_array = args[0].as_string::<i64>();
130            Ok(calculate_ascii(string_array)?)
131        }
132        DataType::Utf8View => {
133            let string_array = args[0].as_string_view();
134            Ok(calculate_ascii(string_array)?)
135        }
136        _ => internal_err!("Unsupported data type"),
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use crate::string::ascii::AsciiFunc;
143    use crate::utils::test::test_function;
144    use arrow::array::{Array, Int32Array};
145    use arrow::datatypes::DataType::Int32;
146    use datafusion_common::{Result, ScalarValue};
147    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
148
149    macro_rules! test_ascii {
150        ($INPUT:expr, $EXPECTED:expr) => {
151            test_function!(
152                AsciiFunc::new(),
153                vec![ColumnarValue::Scalar(ScalarValue::Utf8($INPUT))],
154                $EXPECTED,
155                i32,
156                Int32,
157                Int32Array
158            );
159
160            test_function!(
161                AsciiFunc::new(),
162                vec![ColumnarValue::Scalar(ScalarValue::LargeUtf8($INPUT))],
163                $EXPECTED,
164                i32,
165                Int32,
166                Int32Array
167            );
168
169            test_function!(
170                AsciiFunc::new(),
171                vec![ColumnarValue::Scalar(ScalarValue::Utf8View($INPUT))],
172                $EXPECTED,
173                i32,
174                Int32,
175                Int32Array
176            );
177        };
178    }
179
180    #[test]
181    fn test_functions() -> Result<()> {
182        test_ascii!(Some(String::from("x")), Ok(Some(120)));
183        test_ascii!(Some(String::from("a")), Ok(Some(97)));
184        test_ascii!(Some(String::from("")), Ok(Some(0)));
185        test_ascii!(None, Ok(None));
186        Ok(())
187    }
188}
datafusion_functions/string/ascii.rs

datafusion_functions/string/
ascii.rs
