neilconway commented on code in PR #22508:
URL: https://github.com/apache/datafusion/pull/22508#discussion_r3299375652


##########
datafusion/functions-nested/src/remove.rs:
##########
@@ -210,7 +210,9 @@ impl ScalarUDFImpl for ArrayRemoveN {
         &self,
         args: datafusion_expr::ReturnFieldArgs,
     ) -> Result<FieldRef> {
-        Ok(Arc::clone(&args.arg_fields[0]))
+        let array_field = args.arg_fields[0].as_ref().clone();
+        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
+        Ok(Arc::new(array_field.with_nullable(nullable)))

Review Comment:
   Seems like we want the same change to `return_field_from_args` for 
`array_remove` and `array_remove_all`? Looking around, it seems like 
`array_any_match` has a similar bug. Can you fix these, either as part of this 
PR or in a separate PR?



##########
datafusion/functions-nested/src/remove.rs:
##########
@@ -512,30 +521,47 @@ mod tests {
     fn test_array_remove_n_nullability() {
         for nullability in [true, false] {
             for item_nullability in [true, false] {
-                let input_field = Arc::new(Field::new(
-                    "num",
-                    DataType::new_list(DataType::Int32, item_nullability),
-                    nullability,
-                ));
-                let args_fields = vec![
-                    Arc::clone(&input_field),
-                    Arc::new(Field::new("a", DataType::Int32, false)),
-                    Arc::new(Field::new("b", DataType::Int64, false)),
-                ];
-                let scalar_args = vec![
-                    None,
-                    Some(&ScalarValue::Int32(Some(1))),
-                    Some(&ScalarValue::Int64(Some(1))),
-                ];
-
-                let result = ArrayRemoveN::new()
-                    .return_field_from_args(ReturnFieldArgs {
-                        arg_fields: &args_fields,
-                        scalar_arguments: &scalar_args,
-                    })
-                    .unwrap();
-
-                assert_eq!(result, input_field);
+                for element_nullability in [true, false] {
+                    for count_nullability in [true, false] {

Review Comment:
   You could potentially clean up the deep nesting with `iproduct!`:
   
   ```
       for (array_nullable, item_nullable, element_nullable, count_nullable) in
           iproduct!(bools, bools, bools, bools)
   ```
   
   But feel free to ignore if you think that's too clever.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to