datafusion_functions/string/
ends_with.rsuse std::any::Any;
use std::sync::{Arc, OnceLock};
use arrow::array::ArrayRef;
use arrow::datatypes::DataType;
use crate::utils::make_scalar_function;
use datafusion_common::{internal_err, Result};
use datafusion_expr::scalar_doc_sections::DOC_SECTION_STRING;
use datafusion_expr::{ColumnarValue, Documentation, Volatility};
use datafusion_expr::{ScalarUDFImpl, Signature};
#[derive(Debug)]
pub struct EndsWithFunc {
signature: Signature,
}
impl Default for EndsWithFunc {
fn default() -> Self {
EndsWithFunc::new()
}
}
impl EndsWithFunc {
pub fn new() -> Self {
Self {
signature: Signature::string(2, Volatility::Immutable),
}
}
}
impl ScalarUDFImpl for EndsWithFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"ends_with"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Boolean)
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
match args[0].data_type() {
DataType::Utf8View | DataType::Utf8 | DataType::LargeUtf8 => {
make_scalar_function(ends_with, vec![])(args)
}
other => {
internal_err!("Unsupported data type {other:?} for function ends_with. Expected Utf8, LargeUtf8 or Utf8View")?
}
}
}
fn documentation(&self) -> Option<&Documentation> {
Some(get_ends_with_doc())
}
}
static DOCUMENTATION: OnceLock<Documentation> = OnceLock::new();
fn get_ends_with_doc() -> &'static Documentation {
DOCUMENTATION.get_or_init(|| {
Documentation::builder()
.with_doc_section(DOC_SECTION_STRING)
.with_description("Tests if a string ends with a substring.")
.with_syntax_example("ends_with(str, substr)")
.with_sql_example(
r#"```sql
> select ends_with('datafusion', 'soin');
+--------------------------------------------+
| ends_with(Utf8("datafusion"),Utf8("soin")) |
+--------------------------------------------+
| false |
+--------------------------------------------+
> select ends_with('datafusion', 'sion');
+--------------------------------------------+
| ends_with(Utf8("datafusion"),Utf8("sion")) |
+--------------------------------------------+
| true |
+--------------------------------------------+
```"#,
)
.with_standard_argument("str", Some("String"))
.with_argument("substr", "Substring to test for.")
.build()
.unwrap()
})
}
pub fn ends_with(args: &[ArrayRef]) -> Result<ArrayRef> {
let result = arrow::compute::kernels::comparison::ends_with(&args[0], &args[1])?;
Ok(Arc::new(result) as ArrayRef)
}
#[cfg(test)]
mod tests {
use arrow::array::{Array, BooleanArray};
use arrow::datatypes::DataType::Boolean;
use datafusion_common::Result;
use datafusion_common::ScalarValue;
use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
use crate::string::ends_with::EndsWithFunc;
use crate::utils::test::test_function;
#[test]
fn test_functions() -> Result<()> {
test_function!(
EndsWithFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::from("alphabet")),
ColumnarValue::Scalar(ScalarValue::from("alph")),
],
Ok(Some(false)),
bool,
Boolean,
BooleanArray
);
test_function!(
EndsWithFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::from("alphabet")),
ColumnarValue::Scalar(ScalarValue::from("bet")),
],
Ok(Some(true)),
bool,
Boolean,
BooleanArray
);
test_function!(
EndsWithFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(None)),
ColumnarValue::Scalar(ScalarValue::from("alph")),
],
Ok(None),
bool,
Boolean,
BooleanArray
);
test_function!(
EndsWithFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::from("alphabet")),
ColumnarValue::Scalar(ScalarValue::Utf8(None)),
],
Ok(None),
bool,
Boolean,
BooleanArray
);
Ok(())
}
}