datafusion_functions/string/
rtrim.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::array::{ArrayRef, AsArray};
19use arrow::datatypes::DataType;
20use std::sync::Arc;
21
22use crate::string::common::*;
23use crate::utils::make_scalar_function;
24use datafusion_common::types::logical_string;
25use datafusion_common::{Result, exec_err};
26use datafusion_expr::function::Hint;
27use datafusion_expr::{
28    Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs,
29    ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility,
30};
31use datafusion_macros::user_doc;
32
33/// Returns the longest string with trailing characters removed. If the characters are not specified, spaces are removed.
34/// rtrim('testxxzx', 'xyz') = 'test'
35fn rtrim(args: &[ArrayRef]) -> Result<ArrayRef> {
36    let args = if args.len() > 1 {
37        let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?;
38        vec![Arc::clone(&args[0]), arg1]
39    } else {
40        args.to_owned()
41    };
42    match args[0].data_type() {
43        DataType::Utf8 => general_trim::<i32, TrimRight>(&args, false),
44        DataType::LargeUtf8 => general_trim::<i64, TrimRight>(&args, false),
45        DataType::Utf8View => general_trim::<i32, TrimRight>(&args, true),
46        DataType::Dictionary(_, _) => {
47            let dictionary = args[0].as_any_dictionary();
48            let trimmed = rtrim(&[Arc::clone(dictionary.values())])?;
49            Ok(dictionary.with_values(trimmed))
50        }
51        other => exec_err!(
52            "Unsupported data type {other:?} for function rtrim, expected Utf8, LargeUtf8 or Utf8View."
53        ),
54    }
55}
56
57#[user_doc(
58    doc_section(label = "String Functions"),
59    description = "Trims the specified trim string from the end of a string. If no trim string is provided, all spaces are removed from the end of the input string.",
60    syntax_example = "rtrim(str[, trim_str])",
61    alternative_syntax = "trim(TRAILING trim_str FROM str)",
62    sql_example = r#"```sql
63> select rtrim('  datafusion  ');
64+-------------------------------+
65| rtrim(Utf8("  datafusion  ")) |
66+-------------------------------+
67|   datafusion                  |
68+-------------------------------+
69> select rtrim('___datafusion___', '_');
70+-------------------------------------------+
71| rtrim(Utf8("___datafusion___"),Utf8("_")) |
72+-------------------------------------------+
73| ___datafusion                             |
74+-------------------------------------------+
75```"#,
76    standard_argument(name = "str", prefix = "String"),
77    argument(
78        name = "trim_str",
79        description = "String expression to trim from the end of the input string. Can be a constant, column, or function, and any combination of arithmetic operators. _Default is a space._"
80    ),
81    related_udf(name = "btrim"),
82    related_udf(name = "ltrim")
83)]
84#[derive(Debug, PartialEq, Eq, Hash)]
85pub struct RtrimFunc {
86    signature: Signature,
87}
88
89impl Default for RtrimFunc {
90    fn default() -> Self {
91        Self::new()
92    }
93}
94
95impl RtrimFunc {
96    pub fn new() -> Self {
97        Self {
98            signature: Signature::one_of(
99                vec![
100                    TypeSignature::Coercible(vec![
101                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
102                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
103                    ]),
104                    TypeSignature::Coercible(vec![
105                        Coercion::new_exact(TypeSignatureClass::Native(logical_string()))
106                            .with_encoding_preservation(
107                                EncodingPreservation::dictionary(),
108                            ),
109                    ]),
110                ],
111                Volatility::Immutable,
112            ),
113        }
114    }
115}
116
117impl ScalarUDFImpl for RtrimFunc {
118    fn name(&self) -> &str {
119        "rtrim"
120    }
121
122    fn signature(&self) -> &Signature {
123        &self.signature
124    }
125
126    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
127        Ok(arg_types[0].clone())
128    }
129
130    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
131        make_scalar_function(rtrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args)
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::rtrim::RtrimFunc;
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            RtrimFunc::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            RtrimFunc::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            RtrimFunc::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            RtrimFunc::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("alphabet")),
197            &str,
198            Utf8View,
199            StringViewArray
200        );
201        test_function!(
202            RtrimFunc::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            RtrimFunc::new(),
217            vec![
218                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
219                    "alphabetalphabetxxx"
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            RtrimFunc::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            RtrimFunc::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            RtrimFunc::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            RtrimFunc::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("alphabet")),
267            &str,
268            Utf8,
269            StringArray
270        );
271        test_function!(
272            RtrimFunc::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/rtrim.rs

datafusion_functions/string/
rtrim.rs
