diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 4c6a5c0668f..85097b2ccae 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -431,17 +431,16 @@ where R::Native: AddAssign, { let run_ends: &RunEndBuffer = 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; @@ -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::::new(new_run_ends.into(), None);