datafusion_functions/string/
btrim.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::string::common::*;
19use crate::utils::make_scalar_function;
20use arrow::array::{ArrayRef, AsArray};
21use arrow::datatypes::DataType;
22use datafusion_common::types::logical_string;
23use datafusion_common::{Result, exec_err};
24use datafusion_expr::function::Hint;
25use datafusion_expr::{
26    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
27    ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility,
28};
29use datafusion_macros::user_doc;
30use std::sync::Arc;
31
32/// Returns the longest string with leading and trailing characters removed. If the characters are not specified, spaces are removed.
33/// btrim('xyxtrimyyx', 'xyz') = 'trim'
34fn btrim(args: &[ArrayRef]) -> Result<ArrayRef> {
35    let args = if args.len() > 1 {
36        let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?;
37        vec![Arc::clone(&args[0]), arg1]
38    } else {
39        args.to_owned()
40    };
41    match args[0].data_type() {
42        DataType::Utf8 => general_trim::<i32, TrimBoth>(&args, false),
43        DataType::LargeUtf8 => general_trim::<i64, TrimBoth>(&args, false),
44        DataType::Utf8View => general_trim::<i32, TrimBoth>(&args, true),
45        DataType::Dictionary(_, _) => {
46            let dictionary = args[0].as_any_dictionary();
47            let trimmed = btrim(&[Arc::clone(dictionary.values())])?;
48            Ok(dictionary.with_values(trimmed))
49        }
50        other => exec_err!(
51            "Unsupported data type {other:?} for function btrim, expected Utf8, LargeUtf8 or Utf8View."
52        ),
53    }
54}
55
56#[user_doc(
57    doc_section(label = "String Functions"),
58    description = "Trims the specified trim string from the start and end of a string. If no trim string is provided, all spaces are removed from the start and end of the input string.",
59    syntax_example = "btrim(str[, trim_str])",
60    sql_example = r#"```sql
61> select btrim('__datafusion____', '_');
62+-------------------------------------------+
63| btrim(Utf8("__datafusion____"),Utf8("_")) |
64+-------------------------------------------+
65| datafusion                                |
66+-------------------------------------------+
67```"#,
68    standard_argument(name = "str", prefix = "String"),
69    argument(
70        name = "trim_str",
71        description = r"String expression to operate on. Can be a constant, column, or function, and any combination of operators. _Default is a space._"
72    ),
73    alternative_syntax = "trim(BOTH trim_str FROM str)",
74    alternative_syntax = "trim(trim_str FROM str)",
75    related_udf(name = "ltrim"),
76    related_udf(name = "rtrim")
77)]
78#[derive(Debug, PartialEq, Eq, Hash)]
79pub struct BTrimFunc {
80    signature: Signature,
81    aliases: Vec<String>,
82}
83
84impl Default for BTrimFunc {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl BTrimFunc {
91    pub fn new() -> Self {
92        Self {
93            signature: Signature::one_of(
94                vec![
95                    TypeSignature::Coercible(vec![
96                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
97                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
98                    ]),
99                    TypeSignature::Coercible(vec![
100                        Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
101                            .with_encoding_preservation(
102                                EncodingPreservation::dictionary(),
103                            ),
104                    ]),
105                ],
106                Volatility::Immutable,
107            ),
108            aliases: vec![String::from("trim")],
109        }
110    }
111}
112
113impl ScalarUDFImpl for BTrimFunc {
114    fn name(&self) -> &str {
115        "btrim"
116    }
117
118    fn signature(&self) -> &Signature {
119        &self.signature
120    }
121
122    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
123        Ok(arg_types[0].clone())
124    }
125
126    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
127        make_scalar_function(btrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args)
128    }
129
130    fn aliases(&self) -> &[String] {
131        &self.aliases
132    }
133
134    fn documentation(&self) -> Option<&Documentation> {
135        self.doc()
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use arrow::array::{Array, StringArray, StringViewArray};
142    use arrow::datatypes::DataType::{Utf8, Utf8View};
143
144    use datafusion_common::{Result, ScalarValue};
145    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
146
147    use crate::string::btrim::BTrimFunc;
148    use crate::utils::test::test_function;
149
150    #[test]
151    fn test_functions() {
152        // String view cases for checking normal logic
153        test_function!(
154            BTrimFunc::new(),
155            vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
156                String::from("alphabet  ")
157            )))],
158            Ok(Some("alphabet")),
159            &str,
160            Utf8View,
161            StringViewArray
162        );
163        test_function!(
164            BTrimFunc::new(),
165            vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
166                String::from("  alphabet  ")
167            ))),],
168            Ok(Some("alphabet")),
169            &str,
170            Utf8View,
171            StringViewArray
172        );
173        test_function!(
174            BTrimFunc::new(),
175            vec![
176                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
177                    "alphabet"
178                )))),
179                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("t")))),
180            ],
181            Ok(Some("alphabe")),
182            &str,
183            Utf8View,
184            StringViewArray
185        );
186        test_function!(
187            BTrimFunc::new(),
188            vec![
189                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
190                    "alphabet"
191                )))),
192                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
193                    "alphabe"
194                )))),
195            ],
196            Ok(Some("t")),
197            &str,
198            Utf8View,
199            StringViewArray
200        );
201        test_function!(
202            BTrimFunc::new(),
203            vec![
204                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
205                    "alphabet"
206                )))),
207                ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
208            ],
209            Ok(None),
210            &str,
211            Utf8View,
212            StringViewArray
213        );
214        // Special string view case for checking unlined output(len > 12)
215        test_function!(
216            BTrimFunc::new(),
217            vec![
218                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
219                    "xxxalphabetalphabetxxx"
220                )))),
221                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("x")))),
222            ],
223            Ok(Some("alphabetalphabet")),
224            &str,
225            Utf8View,
226            StringViewArray
227        );
228        // String cases
229        test_function!(
230            BTrimFunc::new(),
231            vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
232                String::from("alphabet  ")
233            ))),],
234            Ok(Some("alphabet")),
235            &str,
236            Utf8,
237            StringArray
238        );
239        test_function!(
240            BTrimFunc::new(),
241            vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
242                String::from("alphabet  ")
243            ))),],
244            Ok(Some("alphabet")),
245            &str,
246            Utf8,
247            StringArray
248        );
249        test_function!(
250            BTrimFunc::new(),
251            vec![
252                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
253                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("t")))),
254            ],
255            Ok(Some("alphabe")),
256            &str,
257            Utf8,
258            StringArray
259        );
260        test_function!(
261            BTrimFunc::new(),
262            vec![
263                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
264                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabe")))),
265            ],
266            Ok(Some("t")),
267            &str,
268            Utf8,
269            StringArray
270        );
271        test_function!(
272            BTrimFunc::new(),
273            vec![
274                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
275                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
276            ],
277            Ok(None),
278            &str,
279            Utf8,
280            StringArray
281        );
282    }
283}
datafusion_functions/string/btrim.rs

datafusion_functions/string/
btrim.rs
