Skip to content

Commit

Permalink
simd selection
Browse files Browse the repository at this point in the history
  • Loading branch information
platoneko committed Feb 28, 2022
1 parent eb581cd commit e599767
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 10 deletions.
35 changes: 25 additions & 10 deletions common/datavalues/src/columns/primitive/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -189,17 +189,32 @@ impl<T: PrimitiveType> Column for PrimitiveColumn<T> {
if length == self.len() {
return Arc::new(self.clone());
}
let iter = self
.values()
.iter()
.zip(filter.values().iter())
.filter(|(_, f)| *f)
.map(|(v, _)| *v);

let values: Vec<T> = iter.collect();
let col = PrimitiveColumn {
values: values.into(),
};
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);
}
}
offset += MASK_BITS;
}

for i in offset..self.len() {
if filter.get_data(i) {
res.push(values[i]);
}
}

let col = PrimitiveColumn { values: res.into() };

Arc::new(col)
}
Expand Down
1 change: 1 addition & 0 deletions common/datavalues/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#![feature(generic_associated_types)]
#![feature(trusted_len)]
#![feature(core_intrinsics)]

#[macro_use]
mod macros;
Expand Down
45 changes: 45 additions & 0 deletions common/datavalues/tests/it/columns/primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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| i < 100 || i > 800)),
expect: (0..N)
.map(|i| i as i32)
.filter(|&i| i < 100 || i > 800)
.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
);
}
}

0 comments on commit e599767

Please sign in to comment.