perf: Add first_(true|false)_idx to BooleanChunked and use in bool arg_(min|max)

crates/polars-arrow/src/legacy/bit_util.rs

@@ -1,81 +1,8 @@
-use crate::bitmap::Bitmap;
 /// Forked from Arrow until their API stabilizes.
 ///
 /// Note that the bound checks are optimized away.
 ///
-use crate::bitmap::utils::{BitChunkIterExact, BitChunks, BitChunksExact};
-
-fn first_set_bit_impl<I>(mut mask_chunks: I) -> usize
-where
-    I: BitChunkIterExact<u64>,
-{
-    let mut total = 0usize;
-    const SIZE: u32 = 64;
-    for chunk in &mut mask_chunks {
-        let pos = chunk.trailing_zeros();
-        if pos != SIZE {
-            return total + pos as usize;
-        } else {
-            total += SIZE as usize
-        }
-    }
-    if let Some(pos) = mask_chunks.remainder_iter().position(|v| v) {
-        total += pos;
-        return total;
-    }
-    // all null, return the first
-    0
-}
-
-pub fn first_set_bit(mask: &Bitmap) -> usize {
-    if mask.unset_bits() == 0 || mask.unset_bits() == mask.len() {
-        return 0;
-    }
-    let (slice, offset, length) = mask.as_slice();
-    if offset == 0 {
-        let mask_chunks = BitChunksExact::<u64>::new(slice, length);
-        first_set_bit_impl(mask_chunks)
-    } else {
-        let mask_chunks = mask.chunks::<u64>();
-        first_set_bit_impl(mask_chunks)
-    }
-}
-
-fn first_unset_bit_impl<I>(mut mask_chunks: I) -> usize
-where
-    I: BitChunkIterExact<u64>,
-{
-    let mut total = 0usize;
-    const SIZE: u32 = 64;
-    for chunk in &mut mask_chunks {
-        let pos = chunk.trailing_ones();
-        if pos != SIZE {
-            return total + pos as usize;
-        } else {
-            total += SIZE as usize
-        }
-    }
-    if let Some(pos) = mask_chunks.remainder_iter().position(|v| !v) {
-        total += pos;
-        return total;
-    }
-    // all null, return the first
-    0
-}
-
-pub fn first_unset_bit(mask: &Bitmap) -> usize {
-    if mask.unset_bits() == 0 || mask.unset_bits() == mask.len() {
-        return 0;
-    }
-    let (slice, offset, length) = mask.as_slice();
-    if offset == 0 {
-        let mask_chunks = BitChunksExact::<u64>::new(slice, length);
-        first_unset_bit_impl(mask_chunks)
-    } else {
-        let mask_chunks = mask.chunks::<u64>();
-        first_unset_bit_impl(mask_chunks)
-    }
-}
+use crate::bitmap::utils::{BitChunkIterExact, BitChunks};
 
 pub fn find_first_true_false_null(
     mut bit_chunks: BitChunks<u64>,

crates/polars-core/src/chunked_array/ops/bits.rs

@@ -1,5 +1,53 @@
 use super::BooleanChunked;
 
+fn first_true_idx_impl(ca: &BooleanChunked, invert: bool) -> Option<usize> {
+    let null_count = ca.null_count();
+    if null_count == ca.len() {
+        return None;
+    }
+
+    if (ca.is_sorted_ascending_flag() && invert) || (ca.is_sorted_descending_flag() && !invert) {
+        return ca.first_non_null();
+    }
+
+    let invert_mask = if invert { u64::MAX } else { 0 };
+    let mut offset = 0;
+    for arr in ca.downcast_iter() {
+        let values = arr.values();
+        if let Some(validity) = arr.validity() {
+            let mut x_it = values.fast_iter_u56();
+            let mut v_it = validity.fast_iter_u56();
+            for (x, v) in x_it.by_ref().zip(v_it.by_ref()) {
+                let n = ((x ^ invert_mask) & v).trailing_zeros() as usize;
+                if n < 56 {
+                    return Some(offset + n);
+                }
+                offset += 56;
+            }
+
+            let (x, rest_len) = x_it.remainder();
+            let (v, _rest_len) = v_it.remainder();
+            let n = ((x ^ invert_mask) & v).trailing_zeros() as usize;
+            if n < rest_len {
+                return Some(offset + n);
+            }
+            offset += rest_len;
+        } else {
+            let n = if invert {
+                values.leading_ones()
+            } else {
+                values.leading_zeros()
+            };
+            if n < values.len() {
+                return Some(offset + n);
+            }
+            offset += values.len();
+        }
+    }
+
+    None
+}
+
 impl BooleanChunked {
     pub fn num_trues(&self) -> usize {
         self.downcast_iter()
@@ -18,4 +66,12 @@ impl BooleanChunked {
             })
             .sum()
     }
+
+    pub fn first_true_idx(&self) -> Option<usize> {
+        first_true_idx_impl(self, false)
+    }
+
+    pub fn first_false_idx(&self) -> Option<usize> {
+        first_true_idx_impl(self, true)
+    }
 }

crates/polars-ops/src/series/ops/arg_min_max.rs

@@ -1,11 +1,9 @@
 use argminmax::ArgMinMax;
 use arrow::array::Array;
-use arrow::legacy::bit_util::first_unset_bit;
 use polars_core::chunked_array::ops::float_sorted_arg_max::{
     float_arg_max_sorted_ascending, float_arg_max_sorted_descending,
 };
 use polars_core::series::IsSorted;
-use polars_core::utils::first_non_null;
 
 use super::*;
 
@@ -166,29 +164,8 @@ where
     }
 }
 
-pub(crate) fn arg_max_bool(ca: &BooleanChunked) -> Option<usize> {
-    let null_count = ca.null_count();
-    if null_count == ca.len() {
-        None
-    } else if null_count == 0 {
-        // if no null values we only have True/False, which can be downcast to
-        // operate on as bitmap 1/0, allowing for a fast-path; if `first_non_null`
-        // returns None this implies all values are zero (eg: False), so we return 0
-        first_non_null(ca.downcast_iter().map(|arr| Some(arr.values()))).or(Some(0))
-    } else {
-        let mut first_false_idx: Option<usize> = None;
-        ca.iter()
-            .enumerate()
-            .find_map(|(idx, val)| match val {
-                Some(true) => Some(idx),
-                Some(false) if first_false_idx.is_none() => {
-                    first_false_idx = Some(idx);
-                    None
-                },
-                _ => None,
-            })
-            .or(first_false_idx)
-    }
+fn arg_max_bool(ca: &BooleanChunked) -> Option<usize> {
+    ca.first_true_idx().or_else(|| ca.first_false_idx())
 }
 
 /// # Safety
@@ -206,26 +183,7 @@ where
 }
 
 fn arg_min_bool(ca: &BooleanChunked) -> Option<usize> {
-    if ca.null_count() == ca.len() {
-        None
-    } else if ca.null_count() == 0 && ca.chunks().len() == 1 {
-        let arr = ca.downcast_iter().next().unwrap();
-        let mask = arr.values();
-        Some(first_unset_bit(mask))
-    } else {
-        let mut first_true_idx: Option<usize> = None;
-        ca.iter()
-            .enumerate()
-            .find_map(|(idx, val)| match val {
-                Some(false) => Some(idx),
-                Some(true) if first_true_idx.is_none() => {
-                    first_true_idx = Some(idx);
-                    None
-                },
-                _ => None,
-            })
-            .or(first_true_idx)
-    }
+    ca.first_false_idx().or_else(|| ca.first_true_idx())
 }
 
 fn arg_min_str(ca: &StringChunked) -> Option<usize> {

crates/polars-ops/src/series/ops/index_of.rs

@@ -44,65 +44,6 @@ where
     index_of_value::<_, PrimitiveArray<T::Native>>(ca, value)
 }
 
-fn index_of_bool(ca: &BooleanChunked, value: bool) -> Option<usize> {
-    let mut index = 0;
-    for chunk in ca.downcast_iter() {
-        let num_possible_trues = chunk.values().set_bits();
-
-        // All values are !value.
-        if (value && num_possible_trues == 0) || (!value && num_possible_trues == chunk.len()) {
-            index += chunk.len();
-            continue;
-        }
-
-        match chunk.validity() {
-            None => {
-                // A lack of a validity bitmap means there are no nulls, so we
-                // can simplify our logic and use a faster code path:
-                let n = if value {
-                    chunk.values().leading_zeros()
-                } else {
-                    chunk.values().leading_ones()
-                };
-
-                if n != chunk.len() {
-                    return Some(index + n);
-                }
-                index += n;
-            },
-            Some(validity) => {
-                // All values are none.
-                if validity.set_bits() == 0 {
-                    index += chunk.len();
-                    continue;
-                }
-
-                let mask = if value { 0 } else { u64::MAX };
-                let mut validity = validity.fast_iter_u56();
-                let mut values = chunk.values().fast_iter_u56();
-                for (is_null, value) in validity.by_ref().zip(values.by_ref()) {
-                    let eq_mask = is_null & (value ^ mask);
-                    if eq_mask != 0 {
-                        return Some(index + eq_mask.trailing_zeros() as usize);
-                    }
-                    index += 56;
-                }
-
-                let (is_null, l1) = validity.remainder();
-                let (value, l2) = values.remainder();
-                assert_eq!(l1, l2);
-
-                let eq_mask = is_null & (value ^ mask);
-                if eq_mask != 0 {
-                    return Some(index + eq_mask.trailing_zeros() as usize);
-                }
-                index += l1;
-            },
-        }
-    }
-    None
-}
-
 /// Try casting the value to the correct type, then call
 /// index_of_numeric_value().
 macro_rules! try_index_of_numeric_ca {
@@ -140,28 +81,28 @@ pub fn index_of(series: &Series, needle: Scalar) -> PolarsResult<Option<usize>>
             return Ok(Some(0));
         }
 
-        let mut index = 0;
+        let mut offset = 0;
         for chunk in series.chunks() {
             let length = chunk.len();
             if let Some(bitmap) = chunk.validity() {
                 let leading_ones = bitmap.leading_ones();
                 if leading_ones < length {
-                    return Ok(Some(index + leading_ones));
+                    return Ok(Some(offset + leading_ones));
                 }
-            } else {
-                index += length;
             }
+            offset += length;
         }
         return Ok(None);
     }
 
     use DataType as DT;
     match series.dtype().to_physical() {
         DT::Null => unreachable!("handled above"),
-        DT::Boolean => Ok(index_of_bool(
-            series.bool()?,
-            needle.value().extract_bool().unwrap(),
-        )),
+        DT::Boolean => Ok(if needle.value().extract_bool().unwrap() {
+            series.bool().unwrap().first_true_idx()
+        } else {
+            series.bool().unwrap().first_false_idx()
+        }),
         dt if dt.is_primitive_numeric() => {
             let series = series.to_physical_repr();
             Ok(downcast_as_macro_arg_physical!(