perf: Reduce memcopy in parquet

crates/polars-io/src/mmap.rs

@@ -1,6 +1,5 @@
 use std::fs::File;
 use std::io::{BufReader, Cursor, Read, Seek};
-use std::sync::Arc;
 
 use polars_core::config::verbose;
 use polars_utils::mmap::{MMapSemaphore, MemSlice};
@@ -82,14 +81,27 @@ impl std::ops::Deref for ReaderBytes<'_> {
     }
 }
 
-/// Require 'static to force the caller to do any transmute as it's usually much
-/// clearer to see there whether it's sound.
+/// There are some places that perform manual lifetime management after transmuting `ReaderBytes`
+/// to have a `'static` inner lifetime. The advantage to doing this is that it lets you construct a
+/// `MemSlice` from the `ReaderBytes` in a zero-copy manner regardless of the underlying enum
+/// variant.
 impl ReaderBytes<'static> {
-    pub fn into_mem_slice(self) -> MemSlice {
+    /// Construct a `MemSlice` in a zero-copy manner from the underlying bytes, with the assumption
+    /// that the underlying bytes have a `'static` lifetime. This is marked as unsafe despite having
+    /// a `'static` inner lifetime, as the `Owned(Vec<u8>)` variant is not covered by the lifetime
+    /// guarantee.
+    ///
+    ///  # Safety
+    /// `Self` outlives the returned `MemSlice` if this enum variant is an `Owned(Vec<u8>)`.
+    pub unsafe fn to_static_slice(&self) -> MemSlice {
         match self {
             ReaderBytes::Borrowed(v) => MemSlice::from_static(v),
-            ReaderBytes::Owned(v) => MemSlice::from_vec(v),
-            ReaderBytes::Mapped(v, _) => MemSlice::from_mmap(Arc::new(v)),
+            ReaderBytes::Owned(v) => MemSlice::from_static(unsafe {
+                std::mem::transmute::<&[u8], &'static [u8]>(v.as_slice())
+            }),
+            ReaderBytes::Mapped(v, _) => unsafe {
+                MemSlice::from_static(std::mem::transmute::<&[u8], &'static [u8]>(v.as_ref()))
+            },
         }
     }
 }

crates/polars-io/src/parquet/read/read_impl.rs

@@ -15,7 +15,6 @@ use polars_parquet::parquet::statistics::Statistics;
 use polars_parquet::read::{
     self, ColumnChunkMetadata, FileMetadata, Filter, PhysicalType, RowGroupMetadata,
 };
-use polars_utils::mmap::MemSlice;
 use rayon::prelude::*;
 
 #[cfg(feature = "cloud")]
@@ -908,10 +907,9 @@ pub fn read_parquet<R: MmapBytesReader>(
     }
 
     let reader = ReaderBytes::from(&mut reader);
-    let store = mmap::ColumnStore::Local(
-        unsafe { std::mem::transmute::<ReaderBytes<'_>, ReaderBytes<'static>>(reader) }
-            .into_mem_slice(),
-    );
+    let store = mmap::ColumnStore::Local(unsafe {
+        std::mem::transmute::<ReaderBytes<'_>, ReaderBytes<'static>>(reader).to_static_slice()
+    });
 
     let dfs = rg_to_dfs(
         &store,
@@ -959,9 +957,9 @@ impl FetchRowGroupsFromMmapReader {
 
     fn fetch_row_groups(&mut self, _row_groups: Range<usize>) -> PolarsResult<ColumnStore> {
         // @TODO: we can something smarter here with mmap
-        Ok(mmap::ColumnStore::Local(MemSlice::from_vec(
-            self.0.deref().to_vec(),
-        )))
+        Ok(mmap::ColumnStore::Local(unsafe {
+            self.0.to_static_slice()
+        }))
     }
 }