Skip to content

Commit

Permalink
add unit tests for xTrim.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Rachelint
Rachelint committed Sep 9, 2024
1 parent c559cc1 commit ea47890
Show file tree
Hide file tree
Showing 4 changed files with 379 additions and 13 deletions.
137 changes: 136 additions & 1 deletion datafusion/functions/src/string/btrim.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,11 @@ impl ScalarUDFImpl for BTrimFunc {
}

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
utf8_to_str_type(&arg_types[0], "btrim")
if arg_types[0] == DataType::Utf8View {
Ok(DataType::Utf8View)
} else {
utf8_to_str_type(&arg_types[0], "btrim")
}
}

fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
Expand All @@ -106,3 +110,134 @@ impl ScalarUDFImpl for BTrimFunc {
&self.aliases
}
}

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


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

use crate::string::btrim::BTrimFunc;
use crate::utils::test::test_function;

#[test]
fn test_functions() {
test_function!(
BTrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet")),
&str,
Utf8View,
StringViewArray
);
test_function!(
BTrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
String::from(" alphabet ")
))),],
Ok(Some("alphabet")),
&str,
Utf8View,
StringViewArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("t")))),
],
Ok(Some("alphabe")),
&str,
Utf8View,
StringViewArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabe"
)))),
],
Ok(Some("t")),
&str,
Utf8View,
StringViewArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
],
Ok(None),
&str,
Utf8View,
StringViewArray
);
test_function!(
BTrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
);
test_function!(
BTrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("t")))),
],
Ok(Some("alphabe")),
&str,
Utf8,
StringArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabe")))),
],
Ok(Some("t")),
&str,
Utf8,
StringArray
);
test_function!(
BTrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(None)),
],
Ok(None),
&str,
Utf8,
StringArray
);
}
}
2 changes: 1 addition & 1 deletion datafusion/functions/src/string/common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ fn string_view_trim<'a, T: OffsetSizeTrait>(
if characters_array.is_null(0) {
return Ok(new_null_array(
// The schema is expecting utf8 as null
&DataType::Utf8,
&DataType::Utf8View,
string_view_array.len(),
));
}
Expand Down
137 changes: 136 additions & 1 deletion datafusion/functions/src/string/ltrim.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,11 @@ impl ScalarUDFImpl for LtrimFunc {
}

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
utf8_to_str_type(&arg_types[0], "ltrim")
if arg_types[0] == DataType::Utf8View {
Ok(DataType::Utf8View)
} else {
utf8_to_str_type(&arg_types[0], "ltrim")
}
}

fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
Expand All @@ -101,3 +105,134 @@ impl ScalarUDFImpl for LtrimFunc {
}
}
}

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


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

use crate::string::ltrim::LtrimFunc;
use crate::utils::test::test_function;

#[test]
fn test_functions() {
test_function!(
LtrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet ")),
&str,
Utf8View,
StringViewArray
);
test_function!(
LtrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
String::from(" alphabet ")
))),],
Ok(Some("alphabet ")),
&str,
Utf8View,
StringViewArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("t")))),
],
Ok(Some("alphabet")),
&str,
Utf8View,
StringViewArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabe"
)))),
],
Ok(Some("t")),
&str,
Utf8View,
StringViewArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
"alphabet"
)))),
ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
],
Ok(None),
&str,
Utf8View,
StringViewArray
);
test_function!(
LtrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet ")),
&str,
Utf8,
StringArray
);
test_function!(
LtrimFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
String::from("alphabet ")
))),],
Ok(Some("alphabet ")),
&str,
Utf8,
StringArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("t")))),
],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabe")))),
],
Ok(Some("t")),
&str,
Utf8,
StringArray
);
test_function!(
LtrimFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
ColumnarValue::Scalar(ScalarValue::Utf8(None)),
],
Ok(None),
&str,
Utf8,
StringArray
);
}
}
Loading

0 comments on commit ea47890

Please sign in to comment.