perf: Optimise low-level null scans and arg_max for bools (when chunked) (#22897)

Author: alexander-beedie

crates/polars-core/src/utils/mod.rs

@@ -10,7 +10,6 @@ mod schema;
 
 pub use any_value::*;
 use arrow::bitmap::Bitmap;
-use arrow::bitmap::bitmask::BitMask;
 pub use arrow::legacy::utils::*;
 pub use arrow::trusted_len::TrustMyLength;
 use flatten::*;
@@ -1200,26 +1199,27 @@ pub(crate) fn index_to_chunked_index_rev<
     )
 }
 
-pub(crate) fn first_non_null<'a, I>(iter: I) -> Option<usize>
+pub fn first_non_null<'a, I>(iter: I) -> Option<usize>
 where
     I: Iterator<Item = Option<&'a Bitmap>>,
 {
     let mut offset = 0;
     for validity in iter {
-        if let Some(validity) = validity {
-            let mask = BitMask::from_bitmap(validity);
-            if let Some(n) = mask.nth_set_bit_idx(0, 0) {
+        if let Some(mask) = validity {
+            let len_mask = mask.len();
+            let n = mask.leading_zeros();
+            if n < len_mask {
                 return Some(offset + n);
             }
-            offset += validity.len()
+            offset += len_mask
         } else {
             return Some(offset);
         }
     }
     None
 }
 
-pub(crate) fn last_non_null<'a, I>(iter: I, len: usize) -> Option<usize>
+pub fn last_non_null<'a, I>(iter: I, len: usize) -> Option<usize>
 where
     I: DoubleEndedIterator<Item = Option<&'a Bitmap>>,
 {
@@ -1228,15 +1228,15 @@ where
     }
     let mut offset = 0;
     for validity in iter.rev() {
-        if let Some(validity) = validity {
-            let mask = BitMask::from_bitmap(validity);
-            if let Some(n) = mask.nth_set_bit_idx_rev(0, mask.len()) {
-                let mask_start = len - offset - mask.len();
-                return Some(mask_start + n);
+        if let Some(mask) = validity {
+            let len_mask = mask.len();
+            let n = mask.trailing_zeros();
+            if n < len_mask {
+                return Some(len - offset - n - 1);
             }
-            offset += validity.len()
+            offset += len_mask;
         } else {
-            return Some(len - 1 - offset);
+            return Some(len - offset - 1);
         }
     }
     None

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

@@ -1,11 +1,11 @@
 use argminmax::ArgMinMax;
 use arrow::array::Array;
-use arrow::legacy::bit_util::*;
+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::with_match_physical_numeric_polars_type;
+use polars_core::utils::first_non_null;
 
 use super::*;
 
@@ -167,14 +167,14 @@ where
 }
 
 pub(crate) fn arg_max_bool(ca: &BooleanChunked) -> Option<usize> {
-    if ca.null_count() == ca.len() {
+    let null_count = ca.null_count();
+    if null_count == ca.len() {
         None
-    }
-    // don't check for any, that on itself is already an argmax search
-    else if ca.null_count() == 0 && ca.chunks().len() == 1 {
-        let arr = ca.downcast_iter().next().unwrap();
-        let mask = arr.values();
-        Some(first_set_bit(mask))
+    } 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()
@@ -202,7 +202,6 @@ where
         IsSorted::Descending => float_arg_max_sorted_descending(ca),
         _ => unreachable!(),
     };
-
     Some(out)
 }
 

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

@@ -80,6 +80,12 @@ pub fn index_of(series: &Series, needle: Scalar) -> PolarsResult<Option<usize>>
 
     // Series is not null, and the value is null:
     if needle.is_null() {
+        let null_count = series.null_count();
+        if null_count == 0 {
+            return Ok(None);
+        } else if null_count == series.len() {
+            return Ok(Some(0));
+        }
         let mut index = 0;
         for chunk in series.chunks() {
             let length = chunk.len();

py-polars/tests/unit/operations/test_index_of.py

@@ -11,10 +11,11 @@
 
 import polars as pl
 from polars.exceptions import InvalidOperationError
+from polars.testing import assert_frame_equal
+from polars.testing.parametric import series
 
 if TYPE_CHECKING:
     from polars._typing import IntoExpr
-from polars.testing import assert_frame_equal
 
 
 def isnan(value: object) -> bool:
@@ -351,3 +352,14 @@ def test_categorical(convert_to_literal: bool) -> None:
     ]:
         for value in expected_values:
             assert_index_of(s, value, convert_to_literal=convert_to_literal)
+
+
+@given(s=series(name="s", allow_chunks=True, max_size=10))
+def test_index_of_null_parametric(s: pl.Series) -> None:
+    idx_null = s.index_of(None)
+    if s.len() == 0:
+        assert idx_null is None
+    elif s.null_count() == 0:
+        assert idx_null is None
+    elif s.null_count() == len(s):
+        assert idx_null == 0

py-polars/tests/unit/series/test_series.py

@@ -1437,18 +1437,16 @@ def test_arg_sort() -> None:
         (pl.Series(["c", "b", "a"], dtype=pl.Categorical), 0, 2),
         (pl.Series([None, "c", "b", None, "a"], dtype=pl.Categorical), 1, 4),
         (pl.Series(["c", "b", "a"], dtype=pl.Categorical(ordering="lexical")), 2, 0),
-        (
-            pl.Series(
-                [None, "c", "b", None, "a"], dtype=pl.Categorical(ordering="lexical")
-            ),
-            4,
-            1,
-        ),
+        (pl.Series("s", [None, "c", "b", None, "a"], pl.Categorical("lexical")), 4, 1),
     ],
 )
 def test_arg_min_arg_max(series: pl.Series, argmin: int, argmax: int) -> None:
-    assert series.arg_min() == argmin
-    assert series.arg_max() == argmax
+    assert series.arg_min() == argmin, (
+        f"values: {series.to_list()}, expected {argmin} got {series.arg_min()}"
+    )
+    assert series.arg_max() == argmax, (
+        f"values: {series.to_list()}, expected {argmax} got {series.arg_max()}"
+    )
 
 
 @pytest.mark.parametrize(
@@ -2216,7 +2214,6 @@ def test_construction_large_nested_u64_17231() -> None:
 
     values = [{"f0": [9223372036854775808]}]
     dtype = pl.Struct({"f0": pl.List(pl.UInt64)})
-
     assert pl.Series(values, dtype=dtype).to_list() == values