@@ -96,44 +96,33 @@ impl ScalarUDFImpl for ArrayHas {
9696 }
9797
9898 fn invoke ( & self , args : & [ ColumnarValue ] ) -> Result < ColumnarValue > {
99- // Always return null if the second argumet is null
100- // i.e. array_has(array, null) -> null
101- if let ColumnarValue :: Scalar ( s) = & args[ 1 ] {
102- if s. is_null ( ) {
103- return Ok ( ColumnarValue :: Scalar ( ScalarValue :: Boolean ( None ) ) ) ;
99+ match & args[ 1 ] {
100+ ColumnarValue :: Array ( array_needle) => {
101+ // the needle is already an array, convert the haystack to an array of the same length
102+ let haystack = args[ 0 ] . to_owned ( ) . into_array ( array_needle. len ( ) ) ?;
103+ let array = array_has_inner_for_array ( & haystack, array_needle) ?;
104+ Ok ( ColumnarValue :: Array ( array) )
104105 }
105- }
106-
107- // first, identify if any of the arguments is an Array. If yes, store its `len`,
108- // as any scalar will need to be converted to an array of len `len`.
109- let len = args
110- . iter ( )
111- . fold ( Option :: < usize > :: None , |acc, arg| match arg {
112- ColumnarValue :: Scalar ( _) => acc,
113- ColumnarValue :: Array ( a) => Some ( a. len ( ) ) ,
114- } ) ;
115-
116- let is_scalar = len. is_none ( ) ;
117-
118- let result = match args[ 1 ] {
119- ColumnarValue :: Array ( _) => {
120- let args = ColumnarValue :: values_to_arrays ( args) ?;
121- array_has_inner_for_array ( & args[ 0 ] , & args[ 1 ] )
122- }
123- ColumnarValue :: Scalar ( _) => {
106+ ColumnarValue :: Scalar ( scalar_needle) => {
107+ // Always return null if the second argument is null
108+ // i.e. array_has(array, null) -> null
109+ if scalar_needle. is_null ( ) {
110+ return Ok ( ColumnarValue :: Scalar ( ScalarValue :: Boolean ( None ) ) ) ;
111+ }
112+
113+ // since the needle is a scalar, convert it to an array of size 1
124114 let haystack = args[ 0 ] . to_owned ( ) . into_array ( 1 ) ?;
125- let needle = args [ 1 ] . to_owned ( ) . into_array ( 1 ) ?;
115+ let needle = scalar_needle . to_array_of_size ( 1 ) ?;
126116 let needle = Scalar :: new ( needle) ;
127- array_has_inner_for_scalar ( & haystack, & needle)
117+ let array = array_has_inner_for_scalar ( & haystack, & needle) ?;
118+ if let ColumnarValue :: Scalar ( _) = & args[ 0 ] {
119+ // If both inputs are scalar, keeps output as scalar
120+ let scalar_value = ScalarValue :: try_from_array ( & array, 0 ) ?;
121+ Ok ( ColumnarValue :: Scalar ( scalar_value) )
122+ } else {
123+ Ok ( ColumnarValue :: Array ( array) )
124+ }
128125 }
129- } ;
130-
131- if is_scalar {
132- // If all inputs are scalar, keeps output as scalar
133- let result = result. and_then ( |arr| ScalarValue :: try_from_array ( & arr, 0 ) ) ;
134- result. map ( ColumnarValue :: Scalar )
135- } else {
136- result. map ( ColumnarValue :: Array )
137126 }
138127 }
139128
@@ -218,10 +207,9 @@ fn array_has_dispatch_for_scalar<O: OffsetSizeTrait>(
218207 }
219208 let sliced_array = eq_array. slice ( start, length) ;
220209 // For nested list, check number of nulls
221- if sliced_array. null_count ( ) = = length {
222- continue ;
210+ if sliced_array. null_count ( ) ! = length {
211+ final_contained [ i ] = Some ( sliced_array . true_count ( ) > 0 ) ;
223212 }
224- final_contained[ i] = Some ( sliced_array. true_count ( ) > 0 ) ;
225213 }
226214
227215 Ok ( Arc :: new ( BooleanArray :: from ( final_contained) ) )
0 commit comments