perf: Call large munmap's in background thread

crates/polars-io/src/csv/read/parser.rs

@@ -5,6 +5,7 @@ use num_traits::Pow;
 use polars_core::prelude::*;
 use polars_core::{POOL, config};
 use polars_error::feature_gated;
+use polars_utils::mmap::MMapSemaphore;
 use polars_utils::select::select_unpredictable;
 use rayon::prelude::*;
 
@@ -38,7 +39,7 @@ pub fn count_rows(
         polars_utils::open_file(path)?
     };
 
-    let mmap = unsafe { memmap::Mmap::map(&file).unwrap() };
+    let mmap = MMapSemaphore::new_from_file(&file).unwrap();
     let owned = &mut vec![];
     let reader_bytes = maybe_decompress_bytes(mmap.as_ref(), owned)?;
 

crates/polars-io/src/ipc/ipc_reader_async.rs

@@ -7,6 +7,7 @@ use polars_core::datatypes::IDX_DTYPE;
 use polars_core::frame::DataFrame;
 use polars_core::schema::{Schema, SchemaExt};
 use polars_error::{PolarsResult, polars_bail, polars_err, to_compute_err};
+use polars_utils::mmap::MMapSemaphore;
 use polars_utils::pl_str::PlSmallStr;
 
 use crate::RowIndex;
@@ -136,7 +137,7 @@ impl IpcReaderAsync {
         // TODO: Only download what is needed rather than the entire file by
         // making use of the projection, row limit, predicate and such.
         let file = tokio::task::block_in_place(|| self.cache_entry.try_open_check_latest())?;
-        let bytes = unsafe { memmap::Mmap::map(&file) }.unwrap();
+        let bytes = MMapSemaphore::new_from_file(&file).unwrap();
 
         let projection = match options.projection.as_deref() {
             Some(projection) => {
@@ -185,7 +186,7 @@ impl IpcReaderAsync {
         // TODO: Only download what is needed rather than the entire file by
         // making use of the projection, row limit, predicate and such.
         let file = tokio::task::block_in_place(|| self.cache_entry.try_open_check_latest())?;
-        let bytes = unsafe { memmap::Mmap::map(&file) }.unwrap();
+        let bytes = MMapSemaphore::new_from_file(&file).unwrap();
         get_row_count(&mut std::io::Cursor::new(bytes.as_ref()))
     }
 }

crates/polars-utils/src/mmap.rs

@@ -1,5 +1,7 @@
 use std::fs::File;
 use std::io;
+use std::mem::ManuallyDrop;
+use std::sync::LazyLock;
 
 pub use memmap::Mmap;
 
@@ -154,6 +156,7 @@ use polars_error::PolarsResult;
 #[cfg(target_family = "unix")]
 use polars_error::polars_bail;
 pub use private::MemSlice;
+use rayon::{ThreadPool, ThreadPoolBuilder};
 
 /// A cursor over a [`MemSlice`].
 #[derive(Debug, Clone)]
@@ -258,6 +261,20 @@ impl io::Seek for MemReader {
     }
 }
 
+pub static UNMAP_POOL: LazyLock<ThreadPool> = LazyLock::new(|| {
+    let thread_name = std::env::var("POLARS_THREAD_NAME").unwrap_or_else(|_| "polars".to_string());
+    ThreadPoolBuilder::new()
+        .num_threads(
+            std::thread::available_parallelism()
+                .unwrap_or(std::num::NonZeroUsize::new(1).unwrap())
+                .get()
+                .min(4),
+        )
+        .thread_name(move |i| format!("{}-unmap-{}", thread_name, i))
+        .build()
+        .expect("could not spawn threads")
+});
+
 // Keep track of memory mapped files so we don't write to them while reading
 // Use a btree as it uses less memory than a hashmap and this thing never shrinks.
 // Write handle in Windows is exclusive, so this is only necessary in Unix.
@@ -270,7 +287,34 @@ static MEMORY_MAPPED_FILES: std::sync::LazyLock<
 pub struct MMapSemaphore {
     #[cfg(target_family = "unix")]
     key: (u64, u64),
-    mmap: Mmap,
+    mmap: ManuallyDrop<Mmap>,
+}
+
+impl Drop for MMapSemaphore {
+    fn drop(&mut self) {
+        #[cfg(target_family = "unix")]
+        {
+            let mut guard = MEMORY_MAPPED_FILES.lock().unwrap();
+            if let std::collections::btree_map::Entry::Occupied(mut e) = guard.entry(self.key) {
+                let v = e.get_mut();
+                *v -= 1;
+
+                if *v == 0 {
+                    e.remove_entry();
+                }
+            }
+        }
+
+        unsafe {
+            let mmap = ManuallyDrop::take(&mut self.mmap);
+            // If the unmap is 1 MiB or bigger, we do it in a background thread.
+            if self.mmap.len() >= 1024 * 1024 {
+                UNMAP_POOL.spawn(move || drop(mmap));
+            } else {
+                drop(mmap);
+            }
+        }
+    }
 }
 
 impl MMapSemaphore {
@@ -293,11 +337,16 @@ impl MMapSemaphore {
                 std::collections::btree_map::Entry::Occupied(mut e) => *e.get_mut() += 1,
                 std::collections::btree_map::Entry::Vacant(e) => _ = e.insert(1),
             }
-            Ok(Self { key, mmap })
+            Ok(Self {
+                key,
+                mmap: ManuallyDrop::new(mmap),
+            })
         }
 
         #[cfg(not(target_family = "unix"))]
-        Ok(Self { mmap })
+        Ok(Self {
+            mmap: ManuallyDrop::new(mmap),
+        })
     }
 
     pub fn new_from_file(file: &File) -> PolarsResult<MMapSemaphore> {
@@ -316,21 +365,6 @@ impl AsRef<[u8]> for MMapSemaphore {
     }
 }
 
-#[cfg(target_family = "unix")]
-impl Drop for MMapSemaphore {
-    fn drop(&mut self) {
-        let mut guard = MEMORY_MAPPED_FILES.lock().unwrap();
-        if let std::collections::btree_map::Entry::Occupied(mut e) = guard.entry(self.key) {
-            let v = e.get_mut();
-            *v -= 1;
-
-            if *v == 0 {
-                e.remove_entry();
-            }
-        }
-    }
-}
-
 pub fn ensure_not_mapped(
     #[cfg_attr(not(target_family = "unix"), allow(unused))] file_md: &std::fs::Metadata,
 ) -> PolarsResult<()> {