functions::array::array_remove

Author: nevi-me

src/functions/array.rs

@@ -87,6 +87,9 @@ impl ArrayFunctions {
         }
         Ok(b.finish())
     }
+
+    /// Locates the position of the first occurrence of the given value in the given array.
+    /// Returns 0 if element is not found, otherwise a 1-based index with the position in the array.
     fn array_position<T>(array: &ListArray, val: T::Native) -> Result<Int32Array, ArrowError>
     where
         T: ArrowPrimitiveType + ArrowNumericType,
@@ -114,7 +117,38 @@ impl ArrayFunctions {
         }
         Ok(b.finish())
     }
-    fn array_remove() {}
+
+    /// Remove all elements that equal the given element in the array
+    fn array_remove<T>(array: &ListArray, val: T::Native) -> Result<ListArray, ArrowError>
+    where
+        T: ArrowPrimitiveType + ArrowNumericType,
+        T::Native: std::cmp::PartialEq<T::Native>,
+    {
+        let values_builder: PrimitiveBuilder<T> = PrimitiveBuilder::new(array.values().len());
+        let mut b = ListBuilder::new(values_builder);
+        // get array datatype so we can downcast appropriately
+        let data_type = array.value_type();
+        for i in 0..array.len() {
+            if array.is_null(i) {
+                b.append(true)?
+            } else {
+                let values = array.values();
+                let values = values.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
+                let values = values.value_slice(
+                    array.value_offset(i) as usize,
+                    array.value_length(i) as usize,
+                );
+                values.iter().for_each(|x| {
+                    // append value if it should not be removed
+                    if x != &val {
+                        b.values().append_value(*x).unwrap();
+                    }
+                });
+                b.append(true)?;
+            }
+        }
+        Ok(b.finish())
+    }
     fn array_repeat() {}
     fn array_sort() {}
     fn array_union() {}
@@ -250,4 +284,34 @@ mod tests {
         assert_eq!(2, bools.value(4));
         assert_eq!(0, bools.value(5));
     }
+
+    #[test]
+    fn test_array_remove() {
+        // Construct a value array
+        let value_data =
+            Int64Array::from(vec![0, 0, 0, 1, 2, 1, 3, 4, 5, 1, 3, 2, 3, 2, 8, 3]).data();
+
+        let value_offsets = Buffer::from(&[0, 3, 6, 8, 12, 14, 16].to_byte_slice());
+
+        // Construct a list array from the above two
+        let list_data_type = DataType::List(Box::new(DataType::Int64));
+        let list_data = ArrayData::builder(list_data_type.clone())
+            .len(6)
+            .add_buffer(value_offsets.clone())
+            .add_child_data(value_data.clone())
+            .build();
+        let list_array = ListArray::from(list_data);
+
+        let b = ArrayFunctions::array_remove::<Int64Type>(&list_array, 2).unwrap();
+        let values = b.values();
+        let values = values.as_any().downcast_ref::<PrimitiveArray<Int64Type>>().unwrap();
+
+        assert_eq!(13, values.len());
+        assert_eq!(0, b.value_offset(0));
+        assert_eq!(3, b.value_offset(1));
+        assert_eq!(5, b.value_offset(2));
+        assert_eq!(7, b.value_offset(3));
+        assert_eq!(10, b.value_offset(4));
+        assert_eq!(11, b.value_offset(5));
+    }
 }