polars-python 0.52.0

Enable running Polars workloads in Python
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

#[cfg(all(
    not(feature = "default_alloc"),
    target_family = "unix",
    not(target_os = "emscripten"),
))]
#[global_allocator]
static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[cfg(all(
    not(feature = "default_alloc"),
    any(not(target_family = "unix"), target_os = "emscripten"),
))]
#[global_allocator]
static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

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 {
    unsafe { std::alloc::alloc(Layout::from_size_align_unchecked(size, align)) }
}

unsafe extern "C" fn dealloc(ptr: *mut u8, size: usize, align: usize) {
    unsafe { std::alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
}

unsafe extern "C" fn alloc_zeroed(size: usize, align: usize) -> *mut u8 {
    unsafe { 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 {
    unsafe {
        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,
            ),
        )
    }
}