@@ -62,7 +62,7 @@ fn rs_width_doc() -> Documentation {
6262 "RS_Width(raster: Raster)" . to_string ( ) ,
6363 )
6464 . with_argument ( "raster" , "Raster: Input raster" )
65- . with_sql_example ( "SELECT RS_Width(raster )" . to_string ( ) )
65+ . with_sql_example ( "SELECT RS_Width(RS_Example() )" . to_string ( ) )
6666 . build ( )
6767}
6868
@@ -73,7 +73,7 @@ fn rs_height_doc() -> Documentation {
7373 "RS_Height(raster: Raster)" . to_string ( ) ,
7474 )
7575 . with_argument ( "raster" , "Raster: Input raster" )
76- . with_sql_example ( "SELECT RS_Height(raster )" . to_string ( ) )
76+ . with_sql_example ( "SELECT RS_Height(RS_Example() )" . to_string ( ) )
7777 . build ( )
7878}
7979
@@ -130,11 +130,13 @@ impl SedonaScalarKernel for RsSize {
130130#[ cfg( test) ]
131131mod tests {
132132 use super :: * ;
133- use arrow_array:: { Array , UInt64Array } ;
133+ use arrow_array:: UInt64Array ;
134134 use datafusion_expr:: ScalarUDF ;
135135 use rstest:: rstest;
136136 use sedona_schema:: datatypes:: RASTER ;
137+ use sedona_testing:: compare:: assert_array_equal;
137138 use sedona_testing:: rasters:: generate_test_rasters;
139+ use sedona_testing:: testers:: ScalarUdfTester ;
138140
139141 #[ test]
140142 fn udf_size ( ) {
@@ -149,35 +151,21 @@ mod tests {
149151
150152 #[ rstest]
151153 fn udf_invoke ( #[ values( SizeType :: Width , SizeType :: Height ) ] st : SizeType ) {
152- let kernel = RsSize {
153- size_type : st. clone ( ) ,
154+ let udf = match st {
155+ SizeType :: Height => rs_height_udf ( ) ,
156+ SizeType :: Width => rs_width_udf ( ) ,
154157 } ;
155- // 3 rasters, second one is null
156- let rasters = generate_test_rasters ( 3 , Some ( 1 ) ) . unwrap ( ) ;
157-
158- // Create the UDF and invoke it
159- let args = [ ColumnarValue :: Array ( Arc :: new ( rasters) ) ] ;
160- let arg_types = vec ! [ RASTER ] ;
161-
162- let result = kernel. invoke_batch ( & arg_types, & args) . unwrap ( ) ;
163-
164- // Check the result
165- if let ColumnarValue :: Array ( result_array) = result {
166- let size_array = result_array. as_any ( ) . downcast_ref :: < UInt64Array > ( ) . unwrap ( ) ;
158+ let tester = ScalarUdfTester :: new ( udf. into ( ) , vec ! [ RASTER ] ) ;
167159
168- assert_eq ! ( size_array. len( ) , 3 ) ;
160+ let rasters = generate_test_rasters ( 3 , Some ( 1 ) ) . unwrap ( ) ;
161+ let expected_values = match st {
162+ SizeType :: Height => vec ! [ Some ( 2 ) , None , Some ( 4 ) ] ,
163+ SizeType :: Width => vec ! [ Some ( 1 ) , None , Some ( 3 ) ] ,
164+ } ;
165+ let expected: Arc < dyn arrow_array:: Array > = Arc :: new ( UInt64Array :: from ( expected_values) ) ;
169166
170- match st. clone ( ) {
171- SizeType :: Width => assert_eq ! ( size_array. value( 0 ) , 1 ) , // First raster width
172- SizeType :: Height => assert_eq ! ( size_array. value( 0 ) , 2 ) , // First raster height
173- }
174- assert ! ( size_array. is_null( 1 ) ) ; // Second raster is null
175- match st. clone ( ) {
176- SizeType :: Width => assert_eq ! ( size_array. value( 2 ) , 3 ) , // Third raster width
177- SizeType :: Height => assert_eq ! ( size_array. value( 2 ) , 4 ) , // Third raster height
178- }
179- } else {
180- panic ! ( "Expected array result" ) ;
181- }
167+ // Check scalars
168+ let result = tester. invoke_array ( Arc :: new ( rasters) ) . unwrap ( ) ;
169+ assert_array_equal ( & result, & expected) ;
182170 }
183171}
0 commit comments