feat(python): Expose allocator to capsule (#17817)

py-polars/src/allocator.rs

@@ -48,3 +48,61 @@ static ALLOC: MiMalloc = MiMalloc;
     not(allocator = "mimalloc"),
 ))]
 static ALLOC: TracemallocAllocator<Jemalloc> = TracemallocAllocator::new(Jemalloc);
+
+use std::alloc::Layout;
+use std::ffi::{c_char, c_void};
+
+use pyo3::ffi::PyCapsule_New;
+use pyo3::{Bound, PyAny, PyResult, Python};
+
+unsafe extern "C" fn alloc(size: usize, align: usize) -> *mut u8 {
+    std::alloc::alloc(Layout::from_size_align_unchecked(size, align))
+}
+
+unsafe extern "C" fn dealloc(ptr: *mut u8, size: usize, align: usize) {
+    std::alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align))
+}
+
+unsafe extern "C" fn alloc_zeroed(size: usize, align: usize) -> *mut u8 {
+    std::alloc::alloc_zeroed(Layout::from_size_align_unchecked(size, align))
+}
+
+unsafe extern "C" fn realloc(ptr: *mut u8, size: usize, align: usize, new_size: usize) -> *mut u8 {
+    std::alloc::realloc(
+        ptr,
+        Layout::from_size_align_unchecked(size, align),
+        new_size,
+    )
+}
+
+#[repr(C)]
+struct AllocatorCapsule {
+    alloc: unsafe extern "C" fn(usize, usize) -> *mut u8,
+    dealloc: unsafe extern "C" fn(*mut u8, usize, usize),
+    alloc_zeroed: unsafe extern "C" fn(usize, usize) -> *mut u8,
+    realloc: unsafe extern "C" fn(*mut u8, usize, usize, usize) -> *mut u8,
+}
+
+static ALLOCATOR_CAPSULE: AllocatorCapsule = AllocatorCapsule {
+    alloc,
+    alloc_zeroed,
+    dealloc,
+    realloc,
+};
+
+static ALLOCATOR_CAPSULE_NAME: &[u8] = b"polars.polars._allocator\0";
+
+pub fn create_allocator_capsule(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
+    unsafe {
+        Bound::from_owned_ptr_or_err(
+            py,
+            PyCapsule_New(
+                &ALLOCATOR_CAPSULE as *const AllocatorCapsule
+            // Users of this capsule is not allowed to modify it.
+            as *mut c_void,
+                ALLOCATOR_CAPSULE_NAME.as_ptr() as *const c_char,
+                None,
+            ),
+        )
+    }
+}

py-polars/src/lib.rs

@@ -44,6 +44,7 @@ mod utils;
 use pyo3::prelude::*;
 use pyo3::{wrap_pyfunction, wrap_pymodule};
 
+use crate::allocator::create_allocator_capsule;
 #[cfg(feature = "csv")]
 use crate::batched_csv::PyBatchedCsv;
 use crate::conversion::Wrap;
@@ -419,6 +420,9 @@ fn polars(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(functions::register_plugin_function))
         .unwrap();
 
+    // Capsules
+    m.add("_allocator", create_allocator_capsule(py)?)?;
+
     Ok(())
 }