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::{ArrayRef, AsArray, Int32Array, StringArrayType};
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: StringArrayType<'a, Item = &'a str>,
107{
108    let values: Vec<_> = (0..array.len())
109        .map(|i| {
110            if array.is_null(i) {
111                0
112            } else {
113                let s = array.value(i);
114                s.chars().next().map_or(0, |c| c as i32)
115            }
116        })
117        .collect();
118
119    let array = Int32Array::new(values.into(), array.nulls().cloned());
120
121    Ok(Arc::new(array))
122}
123
124/// Returns the numeric code of the first character of the argument.
125pub fn ascii(args: &[ArrayRef]) -> Result<ArrayRef> {
126    match args[0].data_type() {
127        DataType::Utf8 => {
128            let string_array = args[0].as_string::<i32>();
129            Ok(calculate_ascii(string_array)?)
130        }
131        DataType::LargeUtf8 => {
132            let string_array = args[0].as_string::<i64>();
133            Ok(calculate_ascii(string_array)?)
134        }
135        DataType::Utf8View => {
136            let string_array = args[0].as_string_view();
137            Ok(calculate_ascii(string_array)?)
138        }
139        _ => internal_err!("Unsupported data type"),
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use crate::string::ascii::AsciiFunc;
146    use crate::utils::test::test_function;
147    use arrow::array::{Array, Int32Array};
148    use arrow::datatypes::DataType::Int32;
149    use datafusion_common::{Result, ScalarValue};
150    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
151
152    macro_rules! test_ascii {
153        ($INPUT:expr, $EXPECTED:expr) => {
154            test_function!(
155                AsciiFunc::new(),
156                vec![ColumnarValue::Scalar(ScalarValue::Utf8($INPUT))],
157                $EXPECTED,
158                i32,
159                Int32,
160                Int32Array
161            );
162
163            test_function!(
164                AsciiFunc::new(),
165                vec![ColumnarValue::Scalar(ScalarValue::LargeUtf8($INPUT))],
166                $EXPECTED,
167                i32,
168                Int32,
169                Int32Array
170            );
171
172            test_function!(
173                AsciiFunc::new(),
174                vec![ColumnarValue::Scalar(ScalarValue::Utf8View($INPUT))],
175                $EXPECTED,
176                i32,
177                Int32,
178                Int32Array
179            );
180        };
181    }
182
183    #[test]
184    fn test_functions() -> Result<()> {
185        test_ascii!(Some(String::from("x")), Ok(Some(120)));
186        test_ascii!(Some(String::from("a")), Ok(Some(97)));
187        test_ascii!(Some(String::from("")), Ok(Some(0)));
188        test_ascii!(Some(String::from("🚀")), Ok(Some(128640)));
189        test_ascii!(None, Ok(None));
190        Ok(())
191    }
192}
datafusion_functions/string/ascii.rs

datafusion_functions/string/
ascii.rs
