Fix

src/compute/filter.rs

@@ -1,13 +1,17 @@
 //! Contains operators to filter arrays such as [`filter`].
+use std::ops::{Shl, Shr};
+
 use crate::array::growable::{make_growable, Growable};
-use crate::bitmap::utils::{BitChunkIterExact, BitChunksExact};
+use crate::bitmap::utils::{BitChunk, BitChunkIterExact, BitChunksExact};
 use crate::bitmap::{utils::SlicesIterator, Bitmap, MutableBitmap};
 use crate::chunk::Chunk;
 use crate::datatypes::DataType;
 use crate::error::Result;
 use crate::types::simd::{NativeSimd, Simd};
 use crate::types::BitChunkOnes;
 use crate::{array::*, types::NativeType};
+use num_traits::One;
+use num_traits::Zero;
 
 /// Function that can filter arbitrary arrays
 pub type Filter<'a> = Box<dyn Fn(&dyn Array) -> Box<dyn Array> + 'a + Send + Sync>;
@@ -22,36 +26,18 @@ where
 {
     let mut chunks = values.chunks_exact(T::Simd::LANES);
     let mut new = Vec::<T>::with_capacity(filter_count);
+    let mut dst = new.as_mut_ptr();
     chunks
         .by_ref()
         .zip(mask_chunks.by_ref())
         .for_each(|(chunk, validity_chunk)| {
             let ones_iter = BitChunkOnes::new(validity_chunk);
-            let mut start_pos = 0;
-            let mut continuous_one = 0;
-
             for pos in ones_iter {
-                if continuous_one == 0 {
-                    start_pos = pos;
-                    continuous_one = 1;
-                    continue;
-                }
-
-                if pos == start_pos + continuous_one {
-                    continuous_one += 1;
-                } else {
-                    new.extend_from_slice(&chunk[start_pos..start_pos + continuous_one]);
-                    start_pos = pos;
-                    continuous_one = 1;
-                }
-            }
-
-            if continuous_one > 0 {
-                new.extend_from_slice(&chunk[start_pos..start_pos + continuous_one]);
+                dst.write(chunk[pos]);
+                dst = dst.add(1);
             }
         });
 
-    let mut dst = new.as_mut_ptr().add(new.len());
     chunks
         .remainder()
         .iter()
@@ -87,47 +73,26 @@ where
     let mut validity_chunks = validity.chunks::<<T::Simd as NativeSimd>::Chunk>();
 
     let mut new = Vec::<T>::with_capacity(filter_count);
+    let mut dst = new.as_mut_ptr().add(new.len());
     let mut new_validity = MutableBitmap::with_capacity(filter_count);
+
     chunks
         .by_ref()
         .zip(validity_chunks.by_ref())
         .zip(mask_chunks.by_ref())
         .for_each(|((chunk, validity_chunk), mask_chunk)| {
             let ones_iter = BitChunkOnes::new(mask_chunk);
-            let mut start_pos = 0;
-            let mut continuous_one = 0;
 
             for pos in ones_iter {
-                if continuous_one == 0 {
-                    start_pos = pos;
-                    continuous_one = 1;
-                    continue;
-                }
-
-                if pos == start_pos + continuous_one {
-                    continuous_one += 1;
-                } else {
-                    new.extend_from_slice(&chunk[start_pos..start_pos + continuous_one]);
-                    new_validity.extend_from_slice(
-                        validity_chunk.to_ne_bytes().as_ref(),
-                        start_pos,
-                        continuous_one,
-                    );
-                    continuous_one = 1;
-                }
-            }
-
-            if continuous_one > 0 {
-                new.extend_from_slice(&chunk[start_pos..start_pos + continuous_one]);
-                new_validity.extend_from_slice(
-                    validity_chunk.to_ne_bytes().as_ref(),
-                    start_pos,
-                    continuous_one,
+                dst.write(chunk[pos]);
+                dst = dst.add(1);
+                new_validity.push(
+                    (validity_chunk.shr(pos) & <<<T as Simd>::Simd as NativeSimd>::Chunk>::one())
+                        > <<<T as Simd>::Simd as NativeSimd>::Chunk>::zero(),
                 );
             }
         });
 
-    let mut dst = new.as_mut_ptr().add(new.len());
     chunks
         .remainder()
         .iter()