(
input_list: ArrayRef,
expected_list: GenericListArray<i32>,
element_to_remove: ScalarValue,
)
| 843 | } |
| 844 | |
| 845 | fn assert_array_remove_all( |
| 846 | input_list: ArrayRef, |
| 847 | expected_list: GenericListArray<i32>, |
| 848 | element_to_remove: ScalarValue, |
| 849 | ) { |
| 850 | assert_eq!(input_list.data_type(), expected_list.data_type()); |
| 851 | assert_eq!(expected_list.value_type(), element_to_remove.data_type()); |
| 852 | let input_list_len = input_list.len(); |
| 853 | let input_list_data_type = input_list.data_type().clone(); |
| 854 | |
| 855 | let udf = ArrayRemoveAll::new(); |
| 856 | let args_fields = vec![ |
| 857 | Arc::new(Field::new("num", input_list.data_type().clone(), false)), |
| 858 | Arc::new(Field::new( |
| 859 | "el", |
| 860 | element_to_remove.data_type(), |
| 861 | element_to_remove.is_null(), |
| 862 | )), |
| 863 | ]; |
| 864 | let scalar_args = vec![None, Some(&element_to_remove)]; |
| 865 | |
| 866 | let return_field = udf |
| 867 | .return_field_from_args(ReturnFieldArgs { |
| 868 | arg_fields: &args_fields, |
| 869 | scalar_arguments: &scalar_args, |
| 870 | }) |
| 871 | .unwrap(); |
| 872 | |
| 873 | let result = udf |
| 874 | .invoke_with_args(ScalarFunctionArgs { |
| 875 | args: vec![ |
| 876 | ColumnarValue::Array(input_list), |
| 877 | ColumnarValue::Scalar(element_to_remove), |
| 878 | ], |
| 879 | arg_fields: args_fields, |
| 880 | number_rows: input_list_len, |
| 881 | return_field, |
| 882 | config_options: Arc::new(Default::default()), |
| 883 | }) |
| 884 | .unwrap(); |
| 885 | |
| 886 | assert_eq!(result.data_type(), input_list_data_type); |
| 887 | match result { |
| 888 | ColumnarValue::Array(array) => { |
| 889 | let result_list = array.as_list::<i32>(); |
| 890 | assert_eq!(result_list, &expected_list); |
| 891 | } |
| 892 | _ => panic!("Expected ColumnarValue::Array"), |
| 893 | } |
| 894 | } |
| 895 | } |
searching dependent graphs…