ISSUE-4264: simd selection

common/datavalues/src/columns/primitive/mod.rs

@@ -20,6 +20,7 @@ use std::sync::Arc;
 
 use common_arrow::arrow::array::Array;
 use common_arrow::arrow::array::PrimitiveArray;
+use common_arrow::arrow::bitmap::utils::BitChunkIterExact;
 use common_arrow::arrow::bitmap::Bitmap;
 use common_arrow::arrow::buffer::Buffer;
 use common_arrow::arrow::compute::arity::unary;
@@ -189,17 +190,50 @@ impl<T: PrimitiveType> Column for PrimitiveColumn<T> {
         if length == self.len() {
             return Arc::new(self.clone());
         }
-        let iter = self
-            .values()
+
+        let mut new = Vec::<T>::with_capacity(length);
+        let mut dst = new.as_mut_ptr();
+
+        const CHUNK_SIZE: usize = 64;
+        let mut chunks = self.values().chunks_exact(CHUNK_SIZE);
+        let mut mask_chunks = filter.values().chunks::<u64>();
+
+        chunks
+            .by_ref()
+            .zip(mask_chunks.by_ref())
+            .for_each(|(chunk, mut mask)| {
+                if mask == u64::MAX {
+                    unsafe {
+                        std::ptr::copy(chunk.as_ptr(), dst, CHUNK_SIZE);
+                        dst = dst.add(CHUNK_SIZE);
+                    }
+                } else {
+                    while mask != 0 {
+                        let n = mask.trailing_zeros() as usize;
+                        unsafe {
+                            dst.write(chunk[n]);
+                            dst = dst.add(1);
+                        }
+                        mask = mask & (mask - 1);
+                    }
+                }
+            });
+
+        chunks
+            .remainder()
             .iter()
-            .zip(filter.values().iter())
-            .filter(|(_, f)| *f)
-            .map(|(v, _)| *v);
+            .zip(mask_chunks.remainder_iter())
+            .for_each(|(value, is_selected)| {
+                if is_selected {
+                    unsafe {
+                        dst.write(*value);
+                        dst = dst.add(1);
+                    }
+                }
+            });
 
-        let values: Vec<T> = iter.collect();
-        let col = PrimitiveColumn {
-            values: values.into(),
-        };
+        unsafe { new.set_len(length) };
+        let col = PrimitiveColumn { values: new.into() };
 
         Arc::new(col)
     }

common/datavalues/tests/it/columns/primitive.rs

@@ -54,3 +54,48 @@ fn test_const_column() {
     let c = ConstColumn::new(Series::from_data(vec![PI]), 24).arc();
     println!("{:?}", c);
 }
+
+#[test]
+fn test_filter_column() {
+    const N: usize = 1000;
+    let it = (0..N).map(|i| i as i32);
+    let data_column: PrimitiveColumn<i32> = Int32Column::from_iterator(it);
+
+    struct Test {
+        filter: BooleanColumn,
+        expect: Vec<i32>,
+    }
+
+    let tests: Vec<Test> = vec![
+        Test {
+            filter: BooleanColumn::from_iterator((0..N).map(|_| true)),
+            expect: (0..N).map(|i| i as i32).collect(),
+        },
+        Test {
+            filter: BooleanColumn::from_iterator((0..N).map(|_| false)),
+            expect: vec![],
+        },
+        Test {
+            filter: BooleanColumn::from_iterator((0..N).map(|i| i % 10 == 0)),
+            expect: (0..N).map(|i| i as i32).filter(|i| i % 10 == 0).collect(),
+        },
+        Test {
+            filter: BooleanColumn::from_iterator((0..N).map(|i| !(100..=800).contains(&i))),
+            expect: (0..N)
+                .map(|i| i as i32)
+                .filter(|&i| !(100..=800).contains(&i))
+                .collect(),
+        },
+    ];
+
+    for test in tests {
+        let res = data_column.filter(&test.filter);
+        assert_eq!(
+            res.as_any()
+                .downcast_ref::<PrimitiveColumn<i32>>()
+                .unwrap()
+                .values(),
+            test.expect
+        );
+    }
+}