Skip to content

Commit

Permalink
Update REVERSE scalar function to support Utf8View (#11973)
Browse files Browse the repository at this point in the history
  • Loading branch information
Omega359 authored Aug 15, 2024
1 parent 4baa901 commit 06bcf33
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 51 deletions.
106 changes: 57 additions & 49 deletions datafusion/functions/src/unicode/reverse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,14 @@
use std::any::Any;
use std::sync::Arc;

use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
use arrow::array::{
Array, ArrayAccessor, ArrayIter, ArrayRef, AsArray, GenericStringArray,
OffsetSizeTrait,
};
use arrow::datatypes::DataType;

use datafusion_common::cast::as_generic_string_array;
use datafusion_common::{exec_err, Result};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
use DataType::{LargeUtf8, Utf8, Utf8View};

use crate::utils::{make_scalar_function, utf8_to_str_type};

Expand All @@ -44,7 +46,7 @@ impl ReverseFunc {
Self {
signature: Signature::uniform(
1,
vec![Utf8, LargeUtf8],
vec![Utf8View, Utf8, LargeUtf8],
Volatility::Immutable,
),
}
Expand All @@ -70,8 +72,8 @@ impl ScalarUDFImpl for ReverseFunc {

fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
match args[0].data_type() {
DataType::Utf8 => make_scalar_function(reverse::<i32>, vec![])(args),
DataType::LargeUtf8 => make_scalar_function(reverse::<i64>, vec![])(args),
Utf8 | Utf8View => make_scalar_function(reverse::<i32>, vec![])(args),
LargeUtf8 => make_scalar_function(reverse::<i64>, vec![])(args),
other => {
exec_err!("Unsupported data type {other:?} for function reverse")
}
Expand All @@ -83,10 +85,17 @@ impl ScalarUDFImpl for ReverseFunc {
/// reverse('abcde') = 'edcba'
/// The implementation uses UTF-8 code points as characters
pub fn reverse<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
let string_array = as_generic_string_array::<T>(&args[0])?;
if args[0].data_type() == &Utf8View {
reverse_impl::<T, _>(args[0].as_string_view())
} else {
reverse_impl::<T, _>(args[0].as_string::<T>())
}
}

let result = string_array
.iter()
fn reverse_impl<'a, T: OffsetSizeTrait, V: ArrayAccessor<Item = &'a str>>(
string_array: V,
) -> Result<ArrayRef> {
let result = ArrayIter::new(string_array)
.map(|string| string.map(|string: &str| string.chars().rev().collect::<String>()))
.collect::<GenericStringArray<T>>();

Expand All @@ -95,59 +104,58 @@ pub fn reverse<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {

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

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

use crate::unicode::reverse::ReverseFunc;
use crate::utils::test::test_function;

macro_rules! test_reverse {
($INPUT:expr, $EXPECTED:expr) => {
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8($INPUT))],
$EXPECTED,
&str,
Utf8,
StringArray
);

test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::LargeUtf8($INPUT))],
$EXPECTED,
&str,
LargeUtf8,
LargeStringArray
);

test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8View($INPUT))],
$EXPECTED,
&str,
Utf8,
StringArray
);
};
}

#[test]
fn test_functions() -> Result<()> {
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::from("abcde"))],
Ok(Some("edcba")),
&str,
Utf8,
StringArray
);
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::from("loẅks"))],
Ok(Some("sk̈wol")),
&str,
Utf8,
StringArray
);
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::from("loẅks"))],
Ok(Some("sk̈wol")),
&str,
Utf8,
StringArray
);
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::Utf8(None))],
Ok(None),
&str,
Utf8,
StringArray
);
test_reverse!(Some("abcde".into()), Ok(Some("edcba")));
test_reverse!(Some("loẅks".into()), Ok(Some("sk̈wol")));
test_reverse!(Some("loẅks".into()), Ok(Some("sk̈wol")));
test_reverse!(None, Ok(None));
#[cfg(not(feature = "unicode_expressions"))]
test_function!(
ReverseFunc::new(),
&[ColumnarValue::Scalar(ScalarValue::from("abcde"))],
test_reverse!(
Some("abcde".into()),
internal_err!(
"function reverse requires compilation with feature flag: unicode_expressions."
),
&str,
Utf8,
StringArray
);

Ok(())
Expand Down
30 changes: 30 additions & 0 deletions datafusion/sqllogictest/test_files/functions.slt
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,16 @@ SELECT reverse('abcde')
----
edcba

query T
SELECT reverse(arrow_cast('abcde', 'LargeUtf8'))
----
edcba

query T
SELECT reverse(arrow_cast('abcde', 'Utf8View'))
----
edcba

query T
SELECT reverse(arrow_cast('abcde', 'Dictionary(Int32, Utf8)'))
----
Expand All @@ -244,11 +254,31 @@ SELECT reverse('loẅks')
----
sk̈wol

query T
SELECT reverse(arrow_cast('loẅks', 'LargeUtf8'))
----
sk̈wol

query T
SELECT reverse(arrow_cast('loẅks', 'Utf8View'))
----
sk̈wol

query T
SELECT reverse(NULL)
----
NULL

query T
SELECT reverse(arrow_cast(NULL, 'LargeUtf8'))
----
NULL

query T
SELECT reverse(arrow_cast(NULL, 'Utf8View'))
----
NULL

query T
SELECT right('abcde', -2)
----
Expand Down
3 changes: 1 addition & 2 deletions datafusion/sqllogictest/test_files/string_view.slt
Original file line number Diff line number Diff line change
Expand Up @@ -890,14 +890,13 @@ logical_plan
03)----TableScan: test projection=[column1_utf8view, column2_utf8view]

## Ensure no casts for REVERSE
## TODO file ticket
query TT
EXPLAIN SELECT
REVERSE(column1_utf8view) as c1
FROM test;
----
logical_plan
01)Projection: reverse(CAST(test.column1_utf8view AS Utf8)) AS c1
01)Projection: reverse(test.column1_utf8view) AS c1
02)--TableScan: test projection=[column1_utf8view]


Expand Down

0 comments on commit 06bcf33

Please sign in to comment.