diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 4575c03076e..4dc00ea2aba 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -25,7 +25,20 @@ use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::array::RunEndArraySlotsExt; use crate::compute::take::take_indices_unchecked; -const FILTER_TAKE_THRESHOLD: f64 = 0.1; + +/// Takes directly below this average number of selected rows per source run. +/// +/// Take scales with the selection size; run filtering scans every run. [#1969] introduced this +/// heuristic without a focused threshold benchmark. A larger value favors take. +/// +/// [#1969]: https://github.com/vortex-data/vortex/pull/1969 +const TAKE_SELECTED_ROWS_PER_RUN_THRESHOLD: f64 = 0.1; + +/// Takes directly below this row count to avoid fixed run and value filtering costs. +/// [#1969] introduced the cutoff. A larger value favors take. +/// +/// [#1969]: https://github.com/vortex-data/vortex/pull/1969 +const MIN_RUN_FILTER_SELECTED_ROWS: usize = 25; impl FilterKernel for RunEnd { fn filter( @@ -36,87 +49,98 @@ impl FilterKernel for RunEnd { let mask_values = mask .values() .vortex_expect("FilterKernel precondition: mask is Mask::Values"); - - let runs_ratio = mask_values.true_count() as f64 / array.ends().len() as f64; - - if runs_ratio < FILTER_TAKE_THRESHOLD || mask_values.true_count() < 25 { - Ok(Some(take_indices_unchecked( + let selected_rows = mask_values.true_count(); + let source_run_count = array.ends().len(); + let selected_rows_per_run = selected_rows as f64 / source_run_count as f64; + let use_direct_take = selected_rows_per_run < TAKE_SELECTED_ROWS_PER_RUN_THRESHOLD + || selected_rows < MIN_RUN_FILTER_SELECTED_ROWS; + + if use_direct_take { + return Ok(Some(take_indices_unchecked( array, mask_values.indices(), &Validity::NonNullable, ctx, - )?)) - } else { - let primitive_run_ends = array.ends().clone().execute::(ctx)?; - let (run_ends, values_mask) = - match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { - filter_run_end_primitive( - primitive_run_ends.as_slice::

(), - array.offset() as u64, - array.len() as u64, - mask_values.bit_buffer(), - )? - }); - let values = array.values().filter(values_mask)?; - - // SAFETY: guaranteed by implementation of filter_run_end_primitive - unsafe { - Ok(Some( - RunEnd::new_unchecked( - run_ends.into_array(), - values, - 0, - mask_values.true_count(), - ) - .into_array(), - )) - } + )?)); } + + let primitive_run_ends = array.ends().clone().execute::(ctx)?; + let (filtered_run_ends, values_mask) = + match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| { + filter_run_end_primitive( + primitive_run_ends.as_slice::

(), + array.offset() as u64, + array.len() as u64, + mask_values.bit_buffer(), + )? + }); + let filtered_values = array.values().filter(values_mask)?; + + // SAFETY: `filter_run_end_primitive` returns one strictly increasing end for each retained + // run value, with the final end equal to `selected_rows`. + let filtered = unsafe { + RunEnd::new_unchecked( + filtered_run_ends.into_array(), + filtered_values, + 0, + selected_rows, + ) + }; + + Ok(Some(filtered.into_array())) } } -// Code adapted from apache arrow-rs https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425 +/// Recomputes cumulative run ends and selects the run values retained by `mask`. +/// +/// The caller supplies validated RunEnd metadata for `offset..offset + length` and a mask containing +/// exactly `length` bits. The returned ends are strictly increasing and contain one entry for each +/// selected run value. +/// +/// Adapted from the [Apache Arrow Rust implementation](https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425). pub fn filter_run_end_primitive + AsPrimitive>( run_ends: &[R], offset: u64, length: u64, mask: &BitBuffer, ) -> VortexResult<(PrimitiveArray, Mask)> { - let mut new_run_ends = buffer_mut![R::zero(); run_ends.len()]; + let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()]; - let mut start = 0u64; - let mut j = 0; - let mut count = R::zero(); + let mut run_start = 0u64; + let mut retained_run_count = 0; + let mut filtered_end = R::zero(); - let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| { - let end = min(run_ends[i].as_() - offset, length); + let values_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |run_idx| { + let run_end = min(run_ends[run_idx].as_() - offset, length); - // SIMD popcount of the predicate bits in this run. The range matches the - // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`. - let start_usize = start + // Bulk popcount is SIMD-capable and avoids per-bit reads. The input contract and clamp prove + // `run_start_idx <= run_end_idx <= mask.len()`. + let run_start_idx = run_start .try_into() .vortex_expect("run start index must fit in usize"); - let end_usize = end + let run_end_idx = run_end .try_into() .vortex_expect("run end index must fit in usize"); - let run_trues = mask.count_range(start_usize, end_usize); - count += ::from(run_trues) + let selected_in_run = mask.count_range(run_start_idx, run_end_idx); + filtered_end += ::from(selected_in_run) .vortex_expect("run popcount must fit in run-end native type"); - let keep = run_trues > 0; + let retain_run = selected_in_run > 0; - // this is to avoid branching - new_run_ends[j] = count; - j += keep as usize; + // Always write the current end, then advance only for a retained run. This keeps the loop + // branchless. + filtered_run_ends[retained_run_count] = filtered_end; + retained_run_count += retain_run as usize; - start = end; - keep + run_start = run_end; + retain_run }) .into(); - new_run_ends.truncate(j); + filtered_run_ends.truncate(retained_run_count); + Ok(( - PrimitiveArray::new(new_run_ends, Validity::NonNullable), - new_mask, + PrimitiveArray::new(filtered_run_ends, Validity::NonNullable), + values_mask, )) }