@@ -364,3 +364,34 @@ def run(ctx):
364364 exa = MockExaEnvironment (meta )
365365 with pytest .raises (TypeError ):
366366 result = executor .run ([Group ([(1 ,), (2 ,), (3 ,), (4 ,), (5 ,), (6 ,)])], exa )
367+
368+
369+ def test_context_parameters ():
370+ def udf_wrapper ():
371+ def run (ctx ):
372+ ctx .emit (ctx [0 ], ctx .t1 )
373+ ctx .emit (ctx [1 ], ctx .t2 )
374+ ctx .emit (ctx [2 ], ctx .t3 )
375+
376+ input_columns = [Column ("t1" , int , "INTEGER" ),
377+ Column ("t2" , int , "INTEGER" ),
378+ Column ("t3" , int , "INTEGER" )]
379+ output_columns = [Column ("o1" , int , "INTEGER" ),
380+ Column ("o2" , int , "INTEGER" )]
381+ meta = MockMetaData (
382+ script_code_wrapper_function = udf_wrapper ,
383+ input_type = "SET" ,
384+ input_columns = input_columns ,
385+ output_type = "EMITS" ,
386+ output_columns = output_columns
387+ )
388+ input_data = [(1 , 2 , 3 ), (4 , 5 , 6 )]
389+ exa = MockExaEnvironment (meta )
390+ executor = UDFMockExecutor ()
391+ result = executor .run ([Group (input_data )], exa )
392+ for i , group in enumerate (result ):
393+ result_row = group .rows
394+ assert len (result_row ) == len (input_columns )
395+ for j in range (len (result_row )):
396+ assert len (result_row [j ]) == len (output_columns )
397+ assert input_data [i ][j ] == result_row [j ][0 ] == result_row [j ][1 ]
0 commit comments