datafusion_functions/unicode/

rpad.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::utils::{make_scalar_function, utf8_to_str_type};
use DataType::{LargeUtf8, Utf8, Utf8View};
use arrow::array::{
    ArrayRef, AsArray, GenericStringArray, GenericStringBuilder, Int64Array,
    OffsetSizeTrait, StringArrayType, StringViewArray,
};
use arrow::datatypes::DataType;
use datafusion_common::DataFusionError;
use datafusion_common::cast::as_int64_array;
use datafusion_common::{Result, exec_err};
use datafusion_expr::TypeSignature::Exact;
use datafusion_expr::{
    ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
};
use datafusion_macros::user_doc;
use std::any::Any;
use std::fmt::Write;
use std::sync::Arc;
use unicode_segmentation::UnicodeSegmentation;

#[user_doc(
    doc_section(label = "String Functions"),
    description = "Pads the right side of a string with another string to a specified string length.",
    syntax_example = "rpad(str, n[, padding_str])",
    sql_example = r#"```sql
>  select rpad('datafusion', 20, '_-');
+-----------------------------------------------+
| rpad(Utf8("datafusion"),Int64(20),Utf8("_-")) |
+-----------------------------------------------+
| datafusion_-_-_-_-_-                          |
+-----------------------------------------------+
```"#,
    standard_argument(name = "str", prefix = "String"),
    argument(name = "n", description = "String length to pad to."),
    argument(
        name = "padding_str",
        description = "String expression to pad with. Can be a constant, column, or function, and any combination of string operators. _Default is a space._"
    ),
    related_udf(name = "lpad")
)]
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct RPadFunc {
    signature: Signature,
}

impl Default for RPadFunc {
    fn default() -> Self {
        Self::new()
    }
}

impl RPadFunc {
    pub fn new() -> Self {
        use DataType::*;
        Self {
            signature: Signature::one_of(
                vec![
                    Exact(vec![Utf8View, Int64]),
                    Exact(vec![Utf8View, Int64, Utf8View]),
                    Exact(vec![Utf8View, Int64, Utf8]),
                    Exact(vec![Utf8View, Int64, LargeUtf8]),
                    Exact(vec![Utf8, Int64]),
                    Exact(vec![Utf8, Int64, Utf8View]),
                    Exact(vec![Utf8, Int64, Utf8]),
                    Exact(vec![Utf8, Int64, LargeUtf8]),
                    Exact(vec![LargeUtf8, Int64]),
                    Exact(vec![LargeUtf8, Int64, Utf8View]),
                    Exact(vec![LargeUtf8, Int64, Utf8]),
                    Exact(vec![LargeUtf8, Int64, LargeUtf8]),
                ],
                Volatility::Immutable,
            ),
        }
    }
}

impl ScalarUDFImpl for RPadFunc {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "rpad"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
        utf8_to_str_type(&arg_types[0], "rpad")
    }

    fn invoke_with_args(
        &self,
        args: datafusion_expr::ScalarFunctionArgs,
    ) -> Result<ColumnarValue> {
        let args = &args.args;
        match (
            args.len(),
            args[0].data_type(),
            args.get(2).map(|arg| arg.data_type()),
        ) {
            (2, Utf8 | Utf8View, _) => {
                make_scalar_function(rpad::<i32, i32>, vec![])(args)
            }
            (2, LargeUtf8, _) => make_scalar_function(rpad::<i64, i64>, vec![])(args),
            (3, Utf8 | Utf8View, Some(Utf8 | Utf8View)) => {
                make_scalar_function(rpad::<i32, i32>, vec![])(args)
            }
            (3, LargeUtf8, Some(LargeUtf8)) => {
                make_scalar_function(rpad::<i64, i64>, vec![])(args)
            }
            (3, Utf8 | Utf8View, Some(LargeUtf8)) => {
                make_scalar_function(rpad::<i32, i64>, vec![])(args)
            }
            (3, LargeUtf8, Some(Utf8 | Utf8View)) => {
                make_scalar_function(rpad::<i64, i32>, vec![])(args)
            }
            (_, _, _) => {
                exec_err!("Unsupported combination of data types for function rpad")
            }
        }
    }

    fn documentation(&self) -> Option<&Documentation> {
        self.doc()
    }
}

fn rpad<StringArrayLen: OffsetSizeTrait, FillArrayLen: OffsetSizeTrait>(
    args: &[ArrayRef],
) -> Result<ArrayRef> {
    if args.len() < 2 || args.len() > 3 {
        return exec_err!(
            "rpad was called with {} arguments. It requires 2 or 3 arguments.",
            args.len()
        );
    }

    let length_array = as_int64_array(&args[1])?;
    match (
        args.len(),
        args[0].data_type(),
        args.get(2).map(|arg| arg.data_type()),
    ) {
        (2, Utf8View, _) => {
            rpad_impl::<&StringViewArray, &StringViewArray, StringArrayLen>(
                &args[0].as_string_view(),
                length_array,
                None,
            )
        }
        (3, Utf8View, Some(Utf8View)) => {
            rpad_impl::<&StringViewArray, &StringViewArray, StringArrayLen>(
                &args[0].as_string_view(),
                length_array,
                Some(args[2].as_string_view()),
            )
        }
        (3, Utf8View, Some(Utf8 | LargeUtf8)) => {
            rpad_impl::<&StringViewArray, &GenericStringArray<FillArrayLen>, StringArrayLen>(
                &args[0].as_string_view(),
                length_array,
                Some(args[2].as_string::<FillArrayLen>()),
            )
        }
        (3, Utf8 | LargeUtf8, Some(Utf8View)) => rpad_impl::<
            &GenericStringArray<StringArrayLen>,
            &StringViewArray,
            StringArrayLen,
        >(
            &args[0].as_string::<StringArrayLen>(),
            length_array,
            Some(args[2].as_string_view()),
        ),
        (_, _, _) => rpad_impl::<
            &GenericStringArray<StringArrayLen>,
            &GenericStringArray<FillArrayLen>,
            StringArrayLen,
        >(
            &args[0].as_string::<StringArrayLen>(),
            length_array,
            args.get(2).map(|arg| arg.as_string::<FillArrayLen>()),
        ),
    }
}

/// Extends the string to length 'length' by appending the characters fill (a space by default). If the string is already longer than length then it is truncated.
/// rpad('hi', 5, 'xy') = 'hixyx'
fn rpad_impl<'a, StringArrType, FillArrType, StringArrayLen>(
    string_array: &StringArrType,
    length_array: &Int64Array,
    fill_array: Option<FillArrType>,
) -> Result<ArrayRef>
where
    StringArrType: StringArrayType<'a>,
    FillArrType: StringArrayType<'a>,
    StringArrayLen: OffsetSizeTrait,
{
    let mut builder: GenericStringBuilder<StringArrayLen> = GenericStringBuilder::new();
    let mut graphemes_buf = Vec::new();
    let mut fill_chars_buf = Vec::new();

    match fill_array {
        None => {
            string_array.iter().zip(length_array.iter()).try_for_each(
                |(string, length)| -> Result<(), DataFusionError> {
                    match (string, length) {
                        (Some(string), Some(length)) => {
                            if length > i32::MAX as i64 {
                                return exec_err!(
                                    "rpad requested length {} too large",
                                    length
                                );
                            }
                            let length = if length < 0 { 0 } else { length as usize };
                            if length == 0 {
                                builder.append_value("");
                            } else {
                                // Reuse buffer by clearing and refilling
                                graphemes_buf.clear();
                                graphemes_buf.extend(string.graphemes(true));

                                if length < graphemes_buf.len() {
                                    builder
                                        .append_value(graphemes_buf[..length].concat());
                                } else {
                                    builder.write_str(string)?;
                                    builder.write_str(
                                        &" ".repeat(length - graphemes_buf.len()),
                                    )?;
                                    builder.append_value("");
                                }
                            }
                        }
                        _ => builder.append_null(),
                    }
                    Ok(())
                },
            )?;
        }
        Some(fill_array) => {
            string_array
                .iter()
                .zip(length_array.iter())
                .zip(fill_array.iter())
                .try_for_each(
                    |((string, length), fill)| -> Result<(), DataFusionError> {
                        match (string, length, fill) {
                            (Some(string), Some(length), Some(fill)) => {
                                if length > i32::MAX as i64 {
                                    return exec_err!(
                                        "rpad requested length {} too large",
                                        length
                                    );
                                }
                                let length = if length < 0 { 0 } else { length as usize };
                                // Reuse buffer by clearing and refilling
                                graphemes_buf.clear();
                                graphemes_buf.extend(string.graphemes(true));

                                if length < graphemes_buf.len() {
                                    builder
                                        .append_value(graphemes_buf[..length].concat());
                                } else if fill.is_empty() {
                                    builder.append_value(string);
                                } else {
                                    builder.write_str(string)?;
                                    // Reuse fill_chars_buf by clearing and refilling
                                    fill_chars_buf.clear();
                                    fill_chars_buf.extend(fill.chars());
                                    for l in 0..length - graphemes_buf.len() {
                                        let c = *fill_chars_buf
                                            .get(l % fill_chars_buf.len())
                                            .unwrap();
                                        builder.write_char(c)?;
                                    }
                                    builder.append_value("");
                                }
                            }
                            _ => builder.append_null(),
                        }
                        Ok(())
                    },
                )?;
        }
    }

    Ok(Arc::new(builder.finish()) as ArrayRef)
}

#[cfg(test)]
mod tests {
    use arrow::array::{Array, StringArray};
    use arrow::datatypes::DataType::Utf8;

    use datafusion_common::{Result, ScalarValue};
    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};

    use crate::unicode::rpad::RPadFunc;
    use crate::utils::test::test_function;

    #[test]
    fn test_functions() -> Result<()> {
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("josé")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
            ],
            Ok(Some("josé ")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
            ],
            Ok(Some("hi   ")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(0i64)),
            ],
            Ok(Some("")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::Int64(None)),
            ],
            Ok(None),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
            ],
            Ok(None),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
                ColumnarValue::Scalar(ScalarValue::from("xy")),
            ],
            Ok(Some("hixyx")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(21i64)),
                ColumnarValue::Scalar(ScalarValue::from("abcdef")),
            ],
            Ok(Some("hiabcdefabcdefabcdefa")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
                ColumnarValue::Scalar(ScalarValue::from(" ")),
            ],
            Ok(Some("hi   ")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
                ColumnarValue::Scalar(ScalarValue::from("")),
            ],
            Ok(Some("hi")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
                ColumnarValue::Scalar(ScalarValue::from("xy")),
            ],
            Ok(None),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::Int64(None)),
                ColumnarValue::Scalar(ScalarValue::from("xy")),
            ],
            Ok(None),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("hi")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
            ],
            Ok(None),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("josé")),
                ColumnarValue::Scalar(ScalarValue::from(10i64)),
                ColumnarValue::Scalar(ScalarValue::from("xy")),
            ],
            Ok(Some("joséxyxyxy")),
            &str,
            Utf8,
            StringArray
        );
        test_function!(
            RPadFunc::new(),
            vec![
                ColumnarValue::Scalar(ScalarValue::from("josé")),
                ColumnarValue::Scalar(ScalarValue::from(10i64)),
                ColumnarValue::Scalar(ScalarValue::from("éñ")),
            ],
            Ok(Some("josééñéñéñ")),
            &str,
            Utf8,
            StringArray
        );
        #[cfg(not(feature = "unicode_expressions"))]
        test_function!(
            RPadFunc::new(),
            &[
                ColumnarValue::Scalar(ScalarValue::from("josé")),
                ColumnarValue::Scalar(ScalarValue::from(5i64)),
            ],
            internal_err!(
                "function rpad requires compilation with feature flag: unicode_expressions."
            ),
            &str,
            Utf8,
            StringArray
        );

        Ok(())
    }
}