(
input_list: ArrayRef,
expected_list: GenericListArray<i32>,
element_to_remove: ScalarValue,
)
| 632 | } |
| 633 | |
| 634 | fn assert_array_remove( |
| 635 | input_list: ArrayRef, |
| 636 | expected_list: GenericListArray<i32>, |
| 637 | element_to_remove: ScalarValue, |
| 638 | ) { |
| 639 | assert_eq!(input_list.data_type(), expected_list.data_type()); |
| 640 | assert_eq!(expected_list.value_type(), element_to_remove.data_type()); |
| 641 | let input_list_len = input_list.len(); |
| 642 | let input_list_data_type = input_list.data_type().clone(); |
| 643 | |
| 644 | let udf = ArrayRemove::new(); |
| 645 | let args_fields = vec![ |
| 646 | Arc::new(Field::new("num", input_list.data_type().clone(), false)), |
| 647 | Arc::new(Field::new( |
| 648 | "el", |
| 649 | element_to_remove.data_type(), |
| 650 | element_to_remove.is_null(), |
| 651 | )), |
| 652 | ]; |
| 653 | let scalar_args = vec![None, Some(&element_to_remove)]; |
| 654 | |
| 655 | let return_field = udf |
| 656 | .return_field_from_args(ReturnFieldArgs { |
| 657 | arg_fields: &args_fields, |
| 658 | scalar_arguments: &scalar_args, |
| 659 | }) |
| 660 | .unwrap(); |
| 661 | |
| 662 | let result = udf |
| 663 | .invoke_with_args(ScalarFunctionArgs { |
| 664 | args: vec![ |
| 665 | ColumnarValue::Array(input_list), |
| 666 | ColumnarValue::Scalar(element_to_remove), |
| 667 | ], |
| 668 | arg_fields: args_fields, |
| 669 | number_rows: input_list_len, |
| 670 | return_field, |
| 671 | config_options: Arc::new(Default::default()), |
| 672 | }) |
| 673 | .unwrap(); |
| 674 | |
| 675 | assert_eq!(result.data_type(), input_list_data_type); |
| 676 | match result { |
| 677 | ColumnarValue::Array(array) => { |
| 678 | let result_list = array.as_list::<i32>(); |
| 679 | assert_eq!(result_list, &expected_list); |
| 680 | } |
| 681 | _ => panic!("Expected ColumnarValue::Array"), |
| 682 | } |
| 683 | } |
| 684 | |
| 685 | #[test] |
| 686 | fn test_array_remove_n_non_nullable() { |
searching dependent graphs…