@@ -19,17 +19,21 @@ use std::any::Any;
19
19
use std:: cmp:: { max, Ordering } ;
20
20
use std:: sync:: Arc ;
21
21
22
- use arrow:: array:: { ArrayRef , GenericStringArray , OffsetSizeTrait } ;
22
+ use arrow:: array:: {
23
+ Array , ArrayAccessor , ArrayIter , ArrayRef , GenericStringArray , Int64Array ,
24
+ OffsetSizeTrait ,
25
+ } ;
23
26
use arrow:: datatypes:: DataType ;
24
27
25
- use datafusion_common:: cast:: { as_generic_string_array, as_int64_array} ;
28
+ use crate :: utils:: { make_scalar_function, utf8_to_str_type} ;
29
+ use datafusion_common:: cast:: {
30
+ as_generic_string_array, as_int64_array, as_string_view_array,
31
+ } ;
26
32
use datafusion_common:: exec_err;
27
33
use datafusion_common:: Result ;
28
34
use datafusion_expr:: TypeSignature :: Exact ;
29
35
use datafusion_expr:: { ColumnarValue , ScalarUDFImpl , Signature , Volatility } ;
30
36
31
- use crate :: utils:: { make_scalar_function, utf8_to_str_type} ;
32
-
33
37
#[ derive( Debug ) ]
34
38
pub struct RightFunc {
35
39
signature : Signature ,
@@ -46,7 +50,11 @@ impl RightFunc {
46
50
use DataType :: * ;
47
51
Self {
48
52
signature : Signature :: one_of (
49
- vec ! [ Exact ( vec![ Utf8 , Int64 ] ) , Exact ( vec![ LargeUtf8 , Int64 ] ) ] ,
53
+ vec ! [
54
+ Exact ( vec![ Utf8View , Int64 ] ) ,
55
+ Exact ( vec![ Utf8 , Int64 ] ) ,
56
+ Exact ( vec![ LargeUtf8 , Int64 ] ) ,
57
+ ] ,
50
58
Volatility :: Immutable ,
51
59
) ,
52
60
}
@@ -72,9 +80,14 @@ impl ScalarUDFImpl for RightFunc {
72
80
73
81
fn invoke ( & self , args : & [ ColumnarValue ] ) -> Result < ColumnarValue > {
74
82
match args[ 0 ] . data_type ( ) {
75
- DataType :: Utf8 => make_scalar_function ( right :: < i32 > , vec ! [ ] ) ( args) ,
83
+ DataType :: Utf8 | DataType :: Utf8View => {
84
+ make_scalar_function ( right :: < i32 > , vec ! [ ] ) ( args)
85
+ }
76
86
DataType :: LargeUtf8 => make_scalar_function ( right :: < i64 > , vec ! [ ] ) ( args) ,
77
- other => exec_err ! ( "Unsupported data type {other:?} for function right" ) ,
87
+ other => exec_err ! (
88
+ "Unsupported data type {other:?} for function right,\
89
+ expected Utf8View, Utf8 or LargeUtf8."
90
+ ) ,
78
91
}
79
92
}
80
93
}
@@ -83,11 +96,26 @@ impl ScalarUDFImpl for RightFunc {
83
96
/// right('abcde', 2) = 'de'
84
97
/// The implementation uses UTF-8 code points as characters
85
98
pub fn right < T : OffsetSizeTrait > ( args : & [ ArrayRef ] ) -> Result < ArrayRef > {
86
- let string_array = as_generic_string_array :: < T > ( & args[ 0 ] ) ?;
87
99
let n_array = as_int64_array ( & args[ 1 ] ) ?;
100
+ if args[ 0 ] . data_type ( ) == & DataType :: Utf8View {
101
+ // string_view_right(args)
102
+ let string_array = as_string_view_array ( & args[ 0 ] ) ?;
103
+ right_impl :: < T , _ > ( & mut string_array. iter ( ) , n_array)
104
+ } else {
105
+ // string_right::<T>(args)
106
+ let string_array = & as_generic_string_array :: < T > ( & args[ 0 ] ) ?;
107
+ right_impl :: < T , _ > ( & mut string_array. iter ( ) , n_array)
108
+ }
109
+ }
88
110
89
- let result = string_array
90
- . iter ( )
111
+ // Currently the return type can only be Utf8 or LargeUtf8, to reach fully support, we need
112
+ // to edit the `get_optimal_return_type` in utils.rs to make the udfs be able to return Utf8View
113
+ // See https://github.com/apache/datafusion/issues/11790#issuecomment-2283777166
114
+ fn right_impl < ' a , T : OffsetSizeTrait , V : ArrayAccessor < Item = & ' a str > > (
115
+ string_array_iter : & mut ArrayIter < V > ,
116
+ n_array : & Int64Array ,
117
+ ) -> Result < ArrayRef > {
118
+ let result = string_array_iter
91
119
. zip ( n_array. iter ( ) )
92
120
. map ( |( string, n) | match ( string, n) {
93
121
( Some ( string) , Some ( n) ) => match n. cmp ( & 0 ) {
0 commit comments