perf: improve join-asof materialization (pola-rs#14884)

crates/polars-core/Cargo.toml

@@ -80,7 +80,6 @@ round_series = []
 checked_arithmetic = []
 is_first_distinct = []
 is_last_distinct = []
-asof_join = []
 dot_product = []
 row_hash = []
 reinterpret = []
@@ -133,7 +132,6 @@ docs-selection = [
   "checked_arithmetic",
   "is_first_distinct",
   "is_last_distinct",
-  "asof_join",
   "dot_product",
   "row_hash",
   "rolling_window",

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

@@ -1,3 +1,4 @@
+use arrow::bitmap::MutableBitmap;
 use polars_utils::sync::SyncPtr;
 
 use super::*;
@@ -89,3 +90,31 @@ fn flatten_par_impl<T: Send + Sync + Copy>(
     }
     out
 }
+
+pub fn flatten_nullable<S: AsRef<[NullableIdxSize]> + Send + Sync>(
+    bufs: &[S],
+) -> PrimitiveArray<IdxSize> {
+    let a = || flatten_par(bufs);
+    let b = || {
+        let cap = bufs.iter().map(|s| s.as_ref().len()).sum::<usize>();
+        let mut validity = MutableBitmap::with_capacity(cap);
+        validity.extend_constant(cap, true);
+
+        let mut count = 0usize;
+        for s in bufs {
+            let s = s.as_ref();
+
+            for id in s {
+                if id.is_null_idx() {
+                    unsafe { validity.set_bit_unchecked(count, false) };
+                }
+
+                count += 1;
+            }
+        }
+        validity.freeze()
+    };
+
+    let (a, b) = POOL.join(a, b);
+    PrimitiveArray::from_vec(bytemuck::cast_vec::<_, IdxSize>(a)).with_validity(Some(b))
+}

crates/polars-ops/Cargo.toml

@@ -113,7 +113,7 @@ top_k = []
 pivot = ["polars-core/reinterpret"]
 cross_join = []
 chunked_ids = []
-asof_join = ["polars-core/asof_join"]
+asof_join = []
 semi_anti_join = []
 array_any_all = ["dtype-array"]
 array_count = ["dtype-array"]