forked from pola-rs/polars
-
Notifications
You must be signed in to change notification settings - Fork 0
/
npy.rs
50 lines (45 loc) · 1.55 KB
/
npy.rs
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
use ndarray::IntoDimension;
use numpy::{
npyffi::{self, flags, types::npy_intp},
ToNpyDims, PY_ARRAY_API,
};
use numpy::{Element, PyArray1};
use polars::chunked_array::builder::alloc;
use pyo3::prelude::*;
use std::{mem, ptr};
/// Create an empty numpy array arrows 64 byte alignment
///
/// # Safety
/// All elements in the array are non initialized
///
/// The array is also writable from Python.
pub unsafe fn aligned_array<T: Element>(py: Python<'_>, size: usize) -> (&PyArray1<T>, *mut T) {
let t_size = std::mem::size_of::<T>();
let capacity = size * t_size;
let ptr = alloc::allocate_aligned::<u8>(capacity).as_ptr() as *mut T;
let mut buf = Vec::from_raw_parts(ptr, 0, capacity);
buf.set_len(size);
// modified from
// numpy-0.10.0/src/array.rs:375
let len = buf.len();
let buffer_ptr = buf.as_mut_ptr();
let dims = [len].into_dimension();
let strides = [mem::size_of::<T>() as npy_intp];
let ptr = PY_ARRAY_API.PyArray_New(
PY_ARRAY_API.get_type_object(npyffi::NpyTypes::PyArray_Type),
dims.ndim_cint(),
dims.as_dims_ptr(),
T::npy_type() as i32,
strides.as_ptr() as *mut _, // strides
buffer_ptr as _, // data
mem::size_of::<T>() as i32, // itemsize
flags::NPY_ARRAY_OUT_ARRAY, // flag
ptr::null_mut(), //obj
);
mem::forget(buf);
(PyArray1::from_owned_ptr(py, ptr), buffer_ptr)
}
pub unsafe fn vec_from_ptr<T>(ptr: usize, len: usize) -> Vec<T> {
let ptr = ptr as *mut T;
Vec::from_raw_parts(ptr, len, len)
}