Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Fix
Browse files Browse the repository at this point in the history
  • Loading branch information
sundy-li committed Feb 28, 2022
1 parent 2299cb4 commit 5342941
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 11 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ target-tarpaulin
venv
lcov.info
Cargo.lock
example.arrow
fixtures
settings.json
dev/
Expand Down
43 changes: 32 additions & 11 deletions src/compute/filter.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
//! Contains operators to filter arrays such as [`filter`].
use std::ops::{Shl, Shr};
use std::ops::Shr;

use crate::array::growable::{make_growable, Growable};
use crate::bitmap::utils::{BitChunk, BitChunkIterExact, BitChunksExact};
Expand Down Expand Up @@ -32,9 +32,19 @@ where
.zip(mask_chunks.by_ref())
.for_each(|(chunk, validity_chunk)| {
let ones_iter = BitChunkOnes::new(validity_chunk);
for pos in ones_iter {
dst.write(chunk[pos]);
dst = dst.add(1);

let (size, _) = ones_iter.size_hint();
if size == T::Simd::LANES {
// Fast path: all lanes are set
unsafe {
std::ptr::copy(chunk.as_ptr(), dst, size);
dst = dst.add(size);
}
} else {
for pos in ones_iter {
dst.write(chunk[pos]);
dst = dst.add(1);
}
}
});

Expand Down Expand Up @@ -82,14 +92,25 @@ where
.zip(mask_chunks.by_ref())
.for_each(|((chunk, validity_chunk), mask_chunk)| {
let ones_iter = BitChunkOnes::new(mask_chunk);
let (size, _) = ones_iter.size_hint();

for pos in ones_iter {
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(),
);
if size == T::Simd::LANES {
// Fast path: all lanes are set
unsafe {
std::ptr::copy(chunk.as_ptr(), dst, size);
dst = dst.add(size);
new_validity.extend_from_slice(validity_chunk.to_ne_bytes().as_ref(), 0, size);
}
} else {
for pos in ones_iter {
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(),
);
}
}
});

Expand Down

0 comments on commit 5342941

Please sign in to comment.