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

datafusion_functions/string/
rtrim.rs
