datafusion_functions/string/
octet_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 arrow::compute::kernels::length::length;
19use arrow::datatypes::DataType;
20
21use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type};
22use datafusion_common::types::logical_string;
23use datafusion_common::utils::take_function_args;
24use datafusion_common::{Result, ScalarValue};
25use datafusion_expr::{
26    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
27    ScalarUDFImpl, Signature, TypeSignatureClass, Volatility,
28};
29use datafusion_macros::user_doc;
30
31#[user_doc(
32    doc_section(label = "String Functions"),
33    description = "Returns the length of a string in bytes.",
34    syntax_example = "octet_length(str)",
35    sql_example = r#"```sql
36> select octet_length('Ångström');
37+--------------------------------+
38| octet_length(Utf8("Ångström")) |
39+--------------------------------+
40| 10                             |
41+--------------------------------+
42```"#,
43    standard_argument(name = "str", prefix = "String"),
44    related_udf(name = "bit_length"),
45    related_udf(name = "length")
46)]
47#[derive(Debug, PartialEq, Eq, Hash)]
48pub struct OctetLengthFunc {
49    signature: Signature,
50}
51
52impl Default for OctetLengthFunc {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl OctetLengthFunc {
59    pub fn new() -> Self {
60        Self {
61            signature: Signature::coercible(
62                vec![
63                    Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
64                        .with_encoding_preservation(EncodingPreservation::dictionary()),
65                ],
66                Volatility::Immutable,
67            ),
68        }
69    }
70}
71
72impl ScalarUDFImpl for OctetLengthFunc {
73    fn name(&self) -> &str {
74        "octet_length"
75    }
76
77    fn signature(&self) -> &Signature {
78        &self.signature
79    }
80
81    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
82        transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| {
83            utf8_to_int_type(data_type, "octet_length")
84        })
85    }
86
87    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
88        let [array] = take_function_args(self.name(), &args.args)?;
89
90        match array {
91            ColumnarValue::Array(v) => Ok(ColumnarValue::Array(length(v.as_ref())?)),
92            ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(octet_length_scalar(v))),
93        }
94    }
95
96    fn documentation(&self) -> Option<&Documentation> {
97        self.doc()
98    }
99}
100
101fn octet_length_scalar(value: &ScalarValue) -> ScalarValue {
102    match value {
103        ScalarValue::Utf8(v) => ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)),
104        ScalarValue::LargeUtf8(v) => {
105            ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64))
106        }
107        ScalarValue::Utf8View(v) => {
108            ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32))
109        }
110        ScalarValue::Dictionary(key_type, value) => ScalarValue::Dictionary(
111            key_type.clone(),
112            Box::new(octet_length_scalar(value)),
113        ),
114        _ => unreachable!("OctetLengthFunc"),
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use std::sync::Arc;
121
122    use arrow::array::{Array, Int32Array, StringArray};
123    use arrow::datatypes::DataType::Int32;
124
125    use datafusion_common::ScalarValue;
126    use datafusion_common::{Result, exec_err};
127    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
128
129    use crate::string::octet_length::OctetLengthFunc;
130    use crate::utils::test::test_function;
131
132    #[test]
133    fn test_functions() -> Result<()> {
134        test_function!(
135            OctetLengthFunc::new(),
136            vec![ColumnarValue::Scalar(ScalarValue::Int32(Some(12)))],
137            exec_err!(
138                "The OCTET_LENGTH function can only accept strings, but got Int32."
139            ),
140            i32,
141            Int32,
142            Int32Array
143        );
144        test_function!(
145            OctetLengthFunc::new(),
146            vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
147                String::from("chars"),
148                String::from("chars2"),
149            ])))],
150            Ok(Some(5)),
151            i32,
152            Int32,
153            Int32Array
154        );
155        test_function!(
156            OctetLengthFunc::new(),
157            vec![
158                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("chars")))),
159                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("chars"))))
160            ],
161            exec_err!("octet_length function requires 1 argument, got 2"),
162            i32,
163            Int32,
164            Int32Array
165        );
166        test_function!(
167            OctetLengthFunc::new(),
168            vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
169                String::from("chars")
170            )))],
171            Ok(Some(5)),
172            i32,
173            Int32,
174            Int32Array
175        );
176        test_function!(
177            OctetLengthFunc::new(),
178            vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
179                String::from("josé")
180            )))],
181            Ok(Some(5)),
182            i32,
183            Int32,
184            Int32Array
185        );
186        test_function!(
187            OctetLengthFunc::new(),
188            vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
189                String::from("")
190            )))],
191            Ok(Some(0)),
192            i32,
193            Int32,
194            Int32Array
195        );
196        test_function!(
197            OctetLengthFunc::new(),
198            vec![ColumnarValue::Scalar(ScalarValue::Utf8(None))],
199            Ok(None),
200            i32,
201            Int32,
202            Int32Array
203        );
204        test_function!(
205            OctetLengthFunc::new(),
206            vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
207                String::from("joséjoséjoséjosé")
208            )))],
209            Ok(Some(20)),
210            i32,
211            Int32,
212            Int32Array
213        );
214        test_function!(
215            OctetLengthFunc::new(),
216            vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
217                String::from("josé")
218            )))],
219            Ok(Some(5)),
220            i32,
221            Int32,
222            Int32Array
223        );
224        test_function!(
225            OctetLengthFunc::new(),
226            vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
227                String::from("")
228            )))],
229            Ok(Some(0)),
230            i32,
231            Int32,
232            Int32Array
233        );
234
235        Ok(())
236    }
237}
datafusion_functions/string/octet_length.rs

datafusion_functions/string/
octet_length.rs
