Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ISSUE-4264: simd selection #4271

Merged
merged 6 commits into from
Mar 7, 2022
Merged

ISSUE-4264: simd selection #4271

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e599767
simd selection
platoneko Feb 28, 2022
9071bc3
make clippy happy
platoneko Feb 28, 2022
9b4ccf1
directly copy via ptr
platoneko Feb 28, 2022
43d5d5e
use extra loop to consume offset
platoneko Mar 6, 2022
ad8060c
Merge branch 'main' into simd-select
mergify[bot] Mar 7, 2022
bd08a7c
Merge branch 'main' into simd-select
mergify[bot] Mar 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 40 additions & 24 deletions common/datavalues/src/columns/primitive/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -190,34 +191,49 @@ impl<T: PrimitiveType> Column for PrimitiveColumn<T> {
return Arc::new(self.clone());
}

let mut res = Vec::<T>::with_capacity(length);
let mut offset = 0;
let values = self.values();

const MASK_BITS: usize = 64;
for mut mask in filter.values().chunks::<u64>() {
if mask == u64::MAX {
res.extend(&values[offset..offset + MASK_BITS]);
} else {
while mask != 0 {
let n = std::intrinsics::cttz(mask) as usize;
res.push(values[offset + n]);
mask = mask & (mask - 1);
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>();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a faster path for mask_chunks, if the filter has no offsets, it could use chunks_exact to have better performance. See: https://github.com/jorgecarleitao/arrow2/blob/main/src/compute/filter.rs#L142-L147

There are two approaches:

  1. Use generic dispatch like BitChunkIterExact

  2. if offset > 0, use an extra loop to consume the offset, this is called Header Loops, then we can continue the Main Loops && Tail Loops.

The second one may have better performance because chunks.iter() must merge two u64 values to generate one merged u64, see: merge_reversed in arrow2


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);
}
}
}
offset += MASK_BITS;
}
});

for (&v, _) in values
chunks
.remainder()
.iter()
.zip(filter.values().iter())
.skip(offset)
.filter(|(_, f)| *f)
{
res.push(v);
}
.zip(mask_chunks.remainder_iter())
.for_each(|(value, is_selected)| {
if is_selected {
unsafe {
dst.write(*value);
dst = dst.add(1);
}
}
});

let col = PrimitiveColumn { values: res.into() };
unsafe { new.set_len(length) };
let col = PrimitiveColumn { values: new.into() };

Arc::new(col)
}
Expand Down
1 change: 0 additions & 1 deletion common/datavalues/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@

#![feature(generic_associated_types)]
#![feature(trusted_len)]
#![feature(core_intrinsics)]

#[macro_use]
mod macros;
Expand Down
4 changes: 2 additions & 2 deletions common/datavalues/tests/it/columns/primitive.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,10 +80,10 @@ fn test_filter_column() {
expect: (0..N).map(|i| i as i32).filter(|i| i % 10 == 0).collect(),
},
Test {
filter: BooleanColumn::from_iterator((0..N).map(|i| i < 100 || i > 800)),
filter: BooleanColumn::from_iterator((0..N).map(|i| !(100..=800).contains(&i))),
expect: (0..N)
.map(|i| i as i32)
.filter(|&i| i < 100 || i > 800)
.filter(|&i| !(100..=800).contains(&i))
.collect(),
},
];
Expand Down