fix: Allow search_sorted on boolean series (#18387)

pola-rs · Aug 27, 2024 · 950e4c3 · 950e4c3
1 parent 6269126
commit 950e4c3
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/crates/polars-ops/src/series/ops/search_sorted.rs b/crates/polars-ops/src/series/ops/search_sorted.rs
@@ -21,6 +21,28 @@ pub fn search_sorted(
             let idx = binary_search_ca(&ca, search_values.iter(), side, descending);
             Ok(IdxCa::new_vec(s.name(), idx))
         },
+        DataType::Boolean => {
+            let ca = s.bool().unwrap();
+            let search_values = search_values.bool()?;
+
+            let mut none_pos = None;
+            let mut false_pos = None;
+            let mut true_pos = None;
+            let idxs = search_values
+                .iter()
+                .map(|v| {
+                    let cache = match v {
+                        None => &mut none_pos,
+                        Some(false) => &mut false_pos,
+                        Some(true) => &mut true_pos,
+                    };
+                    *cache.get_or_insert_with(|| {
+                        binary_search_ca(ca, [v].into_iter(), side, descending)[0]
+                    })
+                })
+                .collect();
+            Ok(IdxCa::new_vec(s.name(), idxs))
+        },
         DataType::Binary => {
             let ca = s.binary().unwrap();