datafusion_functions/unicode/
character_length.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::{
19    make_scalar_function, transform_leaf_type_preserving_encoding, utf8_to_int_type,
20};
21use arrow::array::{
22    Array, ArrayRef, ArrowPrimitiveType, AsArray, OffsetSizeTrait, PrimitiveArray,
23    StringArrayType,
24};
25use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type};
26use datafusion_common::Result;
27use datafusion_common::types::{NativeType, logical_string};
28use datafusion_expr::{
29    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
30    ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
31};
32use datafusion_macros::user_doc;
33use std::sync::Arc;
34
35#[user_doc(
36    doc_section(label = "String Functions"),
37    description = "Returns the number of characters in a string.",
38    syntax_example = "character_length(str)",
39    sql_example = r#"```sql
40> select character_length('Ångström');
41+------------------------------------+
42| character_length(Utf8("Ångström")) |
43+------------------------------------+
44| 8                                  |
45+------------------------------------+
46```"#,
47    standard_argument(name = "str", prefix = "String"),
48    related_udf(name = "bit_length"),
49    related_udf(name = "octet_length")
50)]
51#[derive(Debug, PartialEq, Eq, Hash)]
52pub struct CharacterLengthFunc {
53    signature: Signature,
54    aliases: Vec<String>,
55}
56
57impl Default for CharacterLengthFunc {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl CharacterLengthFunc {
64    pub fn new() -> Self {
65        Self {
66            signature: Signature::coercible(
67                vec![
68                    Coercion::new_implicit(
69                        TypeSignatureClass::Native(logical_string()),
70                        vec![TypeSignatureClass::Any],
71                        NativeType::String,
72                    )
73                    .with_encoding_preservation(EncodingPreservation::dictionary()),
74                ],
75                Volatility::Immutable,
76            ),
77            aliases: vec![String::from("length"), String::from("char_length")],
78        }
79    }
80}
81
82impl ScalarUDFImpl for CharacterLengthFunc {
83    fn name(&self) -> &str {
84        "character_length"
85    }
86
87    fn signature(&self) -> &Signature {
88        &self.signature
89    }
90
91    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
92        transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| {
93            utf8_to_int_type(data_type, "character_length")
94        })
95    }
96
97    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
98        make_scalar_function(character_length, vec![])(&args.args)
99    }
100
101    fn aliases(&self) -> &[String] {
102        &self.aliases
103    }
104
105    fn documentation(&self) -> Option<&Documentation> {
106        self.doc()
107    }
108}
109
110/// Returns number of characters in the string.
111/// character_length('josé') = 4
112/// The implementation counts UTF-8 code points to count the number of characters
113fn character_length(args: &[ArrayRef]) -> Result<ArrayRef> {
114    match args[0].data_type() {
115        DataType::Utf8 => {
116            let string_array = args[0].as_string::<i32>();
117            character_length_general::<Int32Type, _>(&string_array)
118        }
119        DataType::LargeUtf8 => {
120            let string_array = args[0].as_string::<i64>();
121            character_length_general::<Int64Type, _>(&string_array)
122        }
123        DataType::Utf8View => {
124            let string_array = args[0].as_string_view();
125            character_length_general::<Int32Type, _>(&string_array)
126        }
127        DataType::Dictionary(_, _) => {
128            let dictionary = args[0].as_any_dictionary();
129            let converted = character_length(&[Arc::clone(dictionary.values())])?;
130            Ok(dictionary.with_values(converted))
131        }
132        _ => unreachable!("CharacterLengthFunc"),
133    }
134}
135
136fn character_length_general<'a, T, V>(array: &V) -> Result<ArrayRef>
137where
138    T: ArrowPrimitiveType,
139    T::Native: OffsetSizeTrait,
140    V: StringArrayType<'a>,
141{
142    // String characters are variable length encoded in UTF-8, counting the
143    // number of chars requires expensive decoding, however checking if the
144    // string is ASCII only is relatively cheap.
145    // If strings are ASCII only, count bytes instead.
146    let is_array_ascii_only = array.is_ascii();
147    let nulls = array.nulls().cloned();
148    let array = {
149        if is_array_ascii_only {
150            let values: Vec<_> = (0..array.len())
151                .map(|i| {
152                    // Safety: we are iterating with array.len() so the index is always valid
153                    let value = unsafe { array.value_unchecked(i) };
154                    T::Native::usize_as(value.len())
155                })
156                .collect();
157            PrimitiveArray::<T>::new(values.into(), nulls)
158        } else {
159            let values: Vec<_> = (0..array.len())
160                .map(|i| {
161                    // Safety: we are iterating with array.len() so the index is always valid
162                    if array.is_null(i) {
163                        T::default_value()
164                    } else {
165                        let value = unsafe { array.value_unchecked(i) };
166                        if value.is_empty() {
167                            T::default_value()
168                        } else if value.is_ascii() {
169                            T::Native::usize_as(value.len())
170                        } else {
171                            T::Native::usize_as(value.chars().count())
172                        }
173                    }
174                })
175                .collect();
176            PrimitiveArray::<T>::new(values.into(), nulls)
177        }
178    };
179
180    Ok(Arc::new(array))
181}
182
183#[cfg(test)]
184mod tests {
185    use crate::unicode::character_length::CharacterLengthFunc;
186    use crate::utils::test::test_function;
187    use arrow::array::{Array, Int32Array, Int64Array};
188    use arrow::datatypes::DataType::{Int32, Int64};
189    use datafusion_common::{Result, ScalarValue};
190    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
191
192    macro_rules! test_character_length {
193        ($INPUT:expr, $EXPECTED:expr) => {
194            test_function!(
195                CharacterLengthFunc::new(),
196                vec![ColumnarValue::Scalar(ScalarValue::Utf8($INPUT))],
197                $EXPECTED,
198                i32,
199                Int32,
200                Int32Array
201            );
202
203            test_function!(
204                CharacterLengthFunc::new(),
205                vec![ColumnarValue::Scalar(ScalarValue::LargeUtf8($INPUT))],
206                $EXPECTED,
207                i64,
208                Int64,
209                Int64Array
210            );
211
212            test_function!(
213                CharacterLengthFunc::new(),
214                vec![ColumnarValue::Scalar(ScalarValue::Utf8View($INPUT))],
215                $EXPECTED,
216                i32,
217                Int32,
218                Int32Array
219            );
220        };
221    }
222
223    #[test]
224    fn test_functions() -> Result<()> {
225        #[cfg(feature = "unicode_expressions")]
226        {
227            test_character_length!(Some(String::from("chars")), Ok(Some(5)));
228            test_character_length!(Some(String::from("josé")), Ok(Some(4)));
229            // test long strings (more than 12 bytes for StringView)
230            test_character_length!(Some(String::from("joséjoséjoséjosé")), Ok(Some(16)));
231            test_character_length!(Some(String::from("")), Ok(Some(0)));
232            test_character_length!(None, Ok(None));
233        }
234
235        #[cfg(not(feature = "unicode_expressions"))]
236        test_function!(
237            CharacterLengthFunc::new(),
238            &[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
239                String::from("josé")
240            )))],
241            internal_err!(
242                "function character_length requires compilation with feature flag: unicode_expressions."
243            ),
244            i32,
245            Int32,
246            Int32Array
247        );
248
249        Ok(())
250    }
251}
datafusion_functions/unicode/character_length.rs

datafusion_functions/unicode/
character_length.rs
