Fix filter UB and add fast path

arrow/src/compute/kernels/filter.rs

@@ -197,6 +197,10 @@ pub fn build_filter(filter: &BooleanArray) -> Result<Filter> {
     let chunks = iter.collect::<Vec<_>>();
 
     Ok(Box::new(move |array: &ArrayData| {
+        if filter_count == array.len() {
+            return array.clone();
+        }
+
         let mut mutable = MutableArrayData::new(vec![array], false, filter_count);
         chunks
             .iter()
@@ -205,6 +209,22 @@ pub fn build_filter(filter: &BooleanArray) -> Result<Filter> {
     }))
 }
 
+/// Remove null values by do a bitmask AND operation with null bits and the boolean bits.
+fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray {
+    let array_data = filter.data_ref();
+    let null_bitmap = array_data.null_buffer().unwrap();
+    let mask = filter.values();
+    let offset = filter.offset();
+
+    let new_mask = buffer_bin_and(mask, offset, null_bitmap, offset, filter.len());
+
+    let array_data = ArrayData::builder(DataType::Boolean)
+        .len(filter.len())
+        .add_buffer(new_mask)
+        .build();
+    BooleanArray::from(array_data)
+}
+
 /// Filters an [Array], returning elements matching the filter (i.e. where the values are true).
 ///
 /// # Example
@@ -225,38 +245,40 @@ pub fn filter(array: &Array, filter: &BooleanArray) -> Result<ArrayRef> {
     if filter.null_count() > 0 {
         // this greatly simplifies subsequent filtering code
         // now we only have a boolean mask to deal with
-        let array_data = filter.data_ref();
-        let null_bitmap = array_data.null_buffer().unwrap();
-        let mask = filter.values();
-        let offset = filter.offset();
-
-        let new_mask = buffer_bin_and(mask, offset, null_bitmap, offset, filter.len());
-
-        let array_data = ArrayData::builder(DataType::Boolean)
-            .len(filter.len())
-            .add_buffer(new_mask)
-            .build();
-        let filter = BooleanArray::from(array_data);
+        let filter = prep_null_mask_filter(filter);
         // fully qualified syntax, because we have an argument with the same name
         return crate::compute::kernels::filter::filter(array, &filter);
     }
 
     let iter = SlicesIterator::new(filter);
-
-    let mut mutable =
-        MutableArrayData::new(vec![array.data_ref()], false, iter.filter_count);
-    iter.for_each(|(start, end)| mutable.extend(0, start, end));
-    let data = mutable.freeze();
-    Ok(make_array(data))
+    if iter.filter_count == array.len() {
+        let data = array.data().clone();
+        Ok(make_array(data))
+    } else {
+        let mut mutable =
+            MutableArrayData::new(vec![array.data_ref()], false, iter.filter_count);
+        iter.for_each(|(start, end)| mutable.extend(0, start, end));
+        let data = mutable.freeze();
+        Ok(make_array(data))
+    }
 }
 
 /// Returns a new [RecordBatch] with arrays containing only values matching the filter.
-/// WARNING: the nulls of `filter` are ignored and the value on its slot is considered.
-/// Therefore, it is considered undefined behavior to pass `filter` with null values.
 pub fn filter_record_batch(
     record_batch: &RecordBatch,
     filter: &BooleanArray,
 ) -> Result<RecordBatch> {
+    if filter.null_count() > 0 {
+        // this greatly simplifies subsequent filtering code
+        // now we only have a boolean mask to deal with
+        let filter = prep_null_mask_filter(filter);
+        // fully qualified syntax, because we have an argument with the same name
+        return crate::compute::kernels::filter::filter_record_batch(
+            record_batch,
+            &filter,
+        );
+    }
+
     let filter = build_filter(filter)?;
     let filtered_arrays = record_batch
         .columns()
@@ -625,4 +647,18 @@ mod tests {
         assert_eq!(out_arr0, out_arr1);
         Ok(())
     }
+
+    #[test]
+    fn test_fast_path() -> Result<()> {
+        let a: PrimitiveArray<Int64Type> =
+            PrimitiveArray::from(vec![Some(1), Some(2), None]);
+        let mask = BooleanArray::from(vec![true, true, true]);
+        let out = filter(&a, &mask)?;
+        let b = out
+            .as_any()
+            .downcast_ref::<PrimitiveArray<Int64Type>>()
+            .unwrap();
+        assert_eq!(&a, b);
+        Ok(())
+    }
 }