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::transform_leaf_type_preserving_encoding;
19use arrow::array::{ArrayRef, AsArray, Int32Array, StringArrayType};
20use arrow::datatypes::DataType;
21use arrow::error::ArrowError;
22use datafusion_common::types::logical_string;
23use datafusion_common::utils::take_function_args;
24use datafusion_common::{Result, ScalarValue, internal_err};
25use datafusion_expr::{
26    ColumnarValue, Documentation, EncodingPreservation, TypeSignatureClass,
27};
28use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
29use datafusion_expr_common::signature::Coercion;
30use datafusion_macros::user_doc;
31use std::sync::Arc;
32
33#[user_doc(
34    doc_section(label = "String Functions"),
35    description = "Returns the first Unicode scalar value of a string.",
36    syntax_example = "ascii(str)",
37    sql_example = r#"```sql
38> select ascii('abc');
39+--------------------+
40| ascii(Utf8("abc")) |
41+--------------------+
42| 97                 |
43+--------------------+
44> select ascii('🚀');
45+-------------------+
46| ascii(Utf8("🚀")) |
47+-------------------+
48| 128640            |
49+-------------------+
50```"#,
51    standard_argument(name = "str", prefix = "String"),
52    related_udf(name = "chr")
53)]
54#[derive(Debug, PartialEq, Eq, Hash)]
55pub struct AsciiFunc {
56    signature: Signature,
57}
58
59impl Default for AsciiFunc {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl AsciiFunc {
66    pub fn new() -> Self {
67        Self {
68            signature: Signature::coercible(
69                vec![
70                    Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
71                        .with_encoding_preservation(EncodingPreservation::dictionary()),
72                ],
73                Volatility::Immutable,
74            ),
75        }
76    }
77}
78
79impl ScalarUDFImpl for AsciiFunc {
80    fn name(&self) -> &str {
81        "ascii"
82    }
83
84    fn signature(&self) -> &Signature {
85        &self.signature
86    }
87
88    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
89        transform_leaf_type_preserving_encoding(&arg_types[0], &|_| Ok(DataType::Int32))
90    }
91
92    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
93        let [arg] = take_function_args(self.name(), args.args)?;
94
95        match arg {
96            ColumnarValue::Scalar(scalar) => {
97                Ok(ColumnarValue::Scalar(ascii_scalar(&scalar)?))
98            }
99            ColumnarValue::Array(array) => Ok(ColumnarValue::Array(ascii(&[array])?)),
100        }
101    }
102
103    fn documentation(&self) -> Option<&Documentation> {
104        self.doc()
105    }
106}
107
108fn ascii_scalar(scalar: &ScalarValue) -> Result<ScalarValue> {
109    match scalar {
110        ScalarValue::Utf8(value)
111        | ScalarValue::LargeUtf8(value)
112        | ScalarValue::Utf8View(value) => {
113            Ok(ScalarValue::Int32(value.as_deref().map(first_char_code)))
114        }
115        ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary(
116            key_type.clone(),
117            Box::new(ascii_scalar(value)?),
118        )),
119        _ => internal_err!(
120            "Unexpected data type {:?} for function ascii",
121            scalar.data_type()
122        ),
123    }
124}
125
126/// Returns the Unicode scalar value of the first character of `s`, or 0 when
127/// `s` is empty. Reads the leading byte first so the common all-ASCII case
128/// avoids constructing a `char` iterator and decoding a multi-byte sequence.
129#[inline]
130fn first_char_code(s: &str) -> i32 {
131    match s.as_bytes().first() {
132        None => 0,
133        // ASCII byte: the codepoint equals the byte value.
134        Some(&b) if b < 0x80 => b as i32,
135        // Leading byte of a multi-byte sequence: decode the first char.
136        Some(_) => s.chars().next().map_or(0, |c| c as i32),
137    }
138}
139
140fn calculate_ascii<'a, V>(array: &V) -> Result<ArrayRef, ArrowError>
141where
142    V: StringArrayType<'a, Item = &'a str>,
143{
144    let len = array.len();
145    let nulls = array.nulls().cloned();
146
147    // Split the null-handling out of the hot loop: when there is no null
148    // buffer every index is valid, so we can skip the per-element null check
149    // and use unchecked accessors.
150    let values: Vec<i32> = match nulls {
151        Some(ref n) => (0..len)
152            .map(|i| {
153                if n.is_null(i) {
154                    0
155                } else {
156                    // SAFETY: `n.is_null(i)` was false, so `i` is a valid,
157                    // non-null index.
158                    let s = unsafe { array.value_unchecked(i) };
159                    first_char_code(s)
160                }
161            })
162            .collect(),
163        None => (0..len)
164            .map(|i| {
165                // SAFETY: no null buffer means every index in `0..len` is valid.
166                let s = unsafe { array.value_unchecked(i) };
167                first_char_code(s)
168            })
169            .collect(),
170    };
171
172    let array = Int32Array::new(values.into(), nulls);
173
174    Ok(Arc::new(array))
175}
176
177/// Returns the numeric code of the first character of the argument.
178pub fn ascii(args: &[ArrayRef]) -> Result<ArrayRef> {
179    match args[0].data_type() {
180        DataType::Utf8 => {
181            let string_array = args[0].as_string::<i32>();
182            Ok(calculate_ascii(&string_array)?)
183        }
184        DataType::LargeUtf8 => {
185            let string_array = args[0].as_string::<i64>();
186            Ok(calculate_ascii(&string_array)?)
187        }
188        DataType::Utf8View => {
189            let string_array = args[0].as_string_view();
190            Ok(calculate_ascii(&string_array)?)
191        }
192        DataType::Dictionary(_, _) => {
193            let dictionary = args[0].as_any_dictionary();
194            let converted = ascii(&[Arc::clone(dictionary.values())])?;
195            Ok(dictionary.with_values(converted))
196        }
197        _ => internal_err!("Unsupported data type"),
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use crate::string::ascii::AsciiFunc;
204    use crate::utils::test::test_function;
205    use arrow::array::{Array, Int32Array};
206    use arrow::datatypes::DataType::Int32;
207    use datafusion_common::{Result, ScalarValue};
208    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
209
210    macro_rules! test_ascii {
211        ($INPUT:expr, $EXPECTED:expr) => {
212            test_function!(
213                AsciiFunc::new(),
214                vec![ColumnarValue::Scalar(ScalarValue::Utf8($INPUT))],
215                $EXPECTED,
216                i32,
217                Int32,
218                Int32Array
219            );
220
221            test_function!(
222                AsciiFunc::new(),
223                vec![ColumnarValue::Scalar(ScalarValue::LargeUtf8($INPUT))],
224                $EXPECTED,
225                i32,
226                Int32,
227                Int32Array
228            );
229
230            test_function!(
231                AsciiFunc::new(),
232                vec![ColumnarValue::Scalar(ScalarValue::Utf8View($INPUT))],
233                $EXPECTED,
234                i32,
235                Int32,
236                Int32Array
237            );
238        };
239    }
240
241    #[test]
242    fn test_functions() -> Result<()> {
243        test_ascii!(Some(String::from("x")), Ok(Some(120)));
244        test_ascii!(Some(String::from("a")), Ok(Some(97)));
245        test_ascii!(Some(String::from("")), Ok(Some(0)));
246        test_ascii!(Some(String::from("🚀")), Ok(Some(128640)));
247        test_ascii!(Some(String::from("\n")), Ok(Some(10)));
248        test_ascii!(Some(String::from("\t")), Ok(Some(9)));
249        test_ascii!(None, Ok(None));
250        Ok(())
251    }
252}
datafusion_functions/string/ascii.rs

datafusion_functions/string/
ascii.rs
