Skip to content

Commit

Permalink
Speed up filter_run_end_array
Browse files Browse the repository at this point in the history
  • Loading branch information
Dandandan committed Nov 10, 2024
1 parent 28878d3 commit 527ac2f
Showing 1 changed file with 9 additions and 15 deletions.
24 changes: 9 additions & 15 deletions arrow-select/src/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -431,17 +431,16 @@ where
R::Native: AddAssign,
{
let run_ends: &RunEndBuffer<R::Native> = array.run_ends();
let mut values_filter = BooleanBufferBuilder::new(run_ends.len());
let mut new_run_ends = vec![R::default_value(); run_ends.len()];

let mut start = 0u64;
let mut i = 0;
let mut j = 0;
let mut count = R::default_value();
let filter_values = predicate.filter.values();

for mut end in run_ends.inner().into_iter().map(|i| (*i).into() as u64) {
let pred: BooleanArray = BooleanBuffer::collect_bool(run_ends.len(), |i| {
let mut keep = false;

let mut end = run_ends.inner()[i].into() as u64;
let difference = end.saturating_sub(filter_values.len() as u64);
end -= difference;

Expand All @@ -450,23 +449,18 @@ where
count += R::Native::from(pred);
keep |= pred
}

// this is to avoid branching
new_run_ends[i] = count;
i += keep as usize;
new_run_ends[j] = count;
j += keep as usize;

values_filter.append(keep);
start = end;
}

new_run_ends.truncate(i);
keep
}
).into();

if values_filter.is_empty() {
new_run_ends.clear();
}
new_run_ends.truncate(j);

let values = array.values();
let pred = BooleanArray::new(values_filter.finish(), None);
let values = filter(&values, &pred)?;

let run_ends = PrimitiveArray::<R>::new(new_run_ends.into(), None);
Expand Down

0 comments on commit 527ac2f

Please sign in to comment.