polars-python 0.54.4

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

use std::alloc::Layout;
use std::ffi::{c_char, c_void};

use pyo3::ffi::PyCapsule_New;
use pyo3::{Bound, PyAny, PyResult, Python, pyfunction};

#[global_allocator]
static ALLOC: polars_ooc::Allocator = polars_ooc::Allocator;

#[pyfunction]
pub fn _estimate_memory_usage() -> u64 {
    polars_ooc::estimate_memory_usage()
}

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,
            ),
        )
    }
}