perf(rust, python): optimize 'arg_where' (pola-rs#7039)

josemasar · Feb 21, 2023 · f1bcfe3 · f1bcfe3
1 parent 34e1681
commit f1bcfe3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 12 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/arg_where.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/arg_where.rs
@@ -1,4 +1,4 @@
-use polars_arrow::trusted_len::PushUnchecked;
+use polars_core::utils::arrow::bitmap::utils::SlicesIterator;
 
 use super::*;
 
@@ -10,23 +10,27 @@ pub(super) fn arg_where(s: &mut [Series]) -> PolarsResult<Option<Series>> {
     } else {
         let capacity = predicate.sum().unwrap();
         let mut out = Vec::with_capacity(capacity as usize);
-        let mut cnt = 0 as IdxSize;
+        let mut total_offset = 0;
 
         predicate.downcast_iter().for_each(|arr| {
             let values = match arr.validity() {
-                Some(validity) => validity & arr.values(),
-                None => arr.values().clone(),
+                Some(validity) if validity.unset_bits() > 0 => validity & arr.values(),
+                _ => arr.values().clone(),
             };
 
-            // todo! could use chunkiter from arrow here
-            for bit in values.iter() {
-                if bit {
-                    // safety:
-                    // we allocated enough slots
-                    unsafe { out.push_unchecked(cnt) }
+            for (offset, len) in SlicesIterator::new(&values) {
+                // law of small numbers optimization
+                if len == 1 {
+                    out.push(offset as IdxSize)
+                } else {
+                    let offset = (offset + total_offset) as IdxSize;
+                    let len = len as IdxSize;
+                    let iter = offset..offset + len;
+                    out.extend(iter)
                 }
-                cnt += 1;
             }
+
+            total_offset += arr.len();
         });
         let arr = Box::new(IdxArr::from_vec(out)) as ArrayRef;
         unsafe {

diff --git a/polars/polars-lazy/polars-plan/src/dsl/functions.rs b/polars/polars-lazy/polars-plan/src/dsl/functions.rs
@@ -1078,7 +1078,6 @@ pub fn arg_where<E: Into<Expr>>(condition: E) -> Expr {
         function: FunctionExpr::ArgWhere,
         options: FunctionOptions {
             collect_groups: ApplyOptions::ApplyGroups,
-            fmt_str: "arg_where",
             ..Default::default()
         },
     }