diff --git a/py-polars/polars/ffi.py b/py-polars/polars/ffi.py index 1aaf0513fc90..59f8b0a76f12 100644 --- a/py-polars/polars/ffi.py +++ b/py-polars/polars/ffi.py @@ -2,6 +2,12 @@ from numpy import ctypeslib import ctypes from typing import Any +from .polars import ( + aligned_array_f32, + aligned_array_f64, + aligned_array_i32, + aligned_array_i64, +) def ptr_to_numpy(ptr: int, len: int, ptr_type: Any) -> np.ndarray: @@ -12,7 +18,7 @@ def ptr_to_numpy(ptr: int, len: int, ptr_type: Any) -> np.ndarray: ptr C/Rust ptr casted to usize len - Lenght of the array values + Length of the array values ptr_type Example: f32: ctypes.c_float) diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs index 45629bd2a5ef..8e7cdc72fe0c 100644 --- a/py-polars/src/lib.rs +++ b/py-polars/src/lib.rs @@ -1,14 +1,35 @@ +use numpy::PyArray1; use pyo3::prelude::*; +use pyo3::wrap_pyfunction; pub mod dataframe; pub mod error; +pub mod npy; pub mod series; use crate::{dataframe::PyDataFrame, series::PySeries}; +macro_rules! create_aligned_buffer { + ($name:ident, $type:ty) => { + #[pyfunction] + pub fn $name(size: usize) -> Py> { + npy::aligned_array::<$type>(size) + } + }; +} + +create_aligned_buffer!(aligned_array_f32, f32); +create_aligned_buffer!(aligned_array_f64, f64); +create_aligned_buffer!(aligned_array_i32, i32); +create_aligned_buffer!(aligned_array_i64, i64); + #[pymodule] fn polars(_py: Python, m: &PyModule) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; + m.add_class::().unwrap(); + m.add_class::().unwrap(); + m.add_wrapped(wrap_pyfunction!(aligned_array_f32)).unwrap(); + m.add_wrapped(wrap_pyfunction!(aligned_array_f64)).unwrap(); + m.add_wrapped(wrap_pyfunction!(aligned_array_i32)).unwrap(); + m.add_wrapped(wrap_pyfunction!(aligned_array_i64)).unwrap(); Ok(()) } diff --git a/py-polars/src/npy.rs b/py-polars/src/npy.rs new file mode 100644 index 000000000000..7b7fb1dac611 --- /dev/null +++ b/py-polars/src/npy.rs @@ -0,0 +1,12 @@ +use numpy::{Element, PyArray1}; +use polars::prelude::*; +use pyo3::prelude::*; + +/// Create an empty numpy array arrows 64 byte alignment +pub fn aligned_array(size: usize) -> Py> { + let mut buf: Vec = Vec::with_capacity_aligned(size); + unsafe { buf.set_len(size) } + let gil = Python::acquire_gil(); + let python = gil.python(); + PyArray1::from_vec(python, buf).to_owned() +} diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs index 73d482f3db32..151bdd6a37c3 100644 --- a/py-polars/src/series.rs +++ b/py-polars/src/series.rs @@ -1,5 +1,5 @@ use crate::error::PyPolarsEr; -use numpy::{Element, PyArray1}; +use numpy::PyArray1; use polars::prelude::*; use pyo3::types::PyList; use pyo3::{exceptions::RuntimeError, prelude::*}; diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py index 3f1271550415..7c927f1fd482 100644 --- a/py-polars/tests/test_series.py +++ b/py-polars/tests/test_series.py @@ -1,4 +1,5 @@ from polars import Series +from polars.ffi import aligned_array_f32 import numpy as np @@ -133,3 +134,19 @@ def test_view(): a = Series("a", [1.0, 2.0, 3.0]) assert isinstance(a.view(), np.ndarray) assert np.all(a.view() == np.array([1, 2, 3])) + + +def test_numpy_interface(): + a = aligned_array_f32(10) + assert a.dtype == np.float32 + assert a.shape == (10,) + pointer, read_only_flag = a.__array_interface__["data"] + # set read only flag to False + a.__array_interface__["data"] = (pointer, False) + # the __array_interface is used to create a new array (pointing to the same memory) + b = np.array(a) + # now the memory is writeable + b[0] = 1 + + # TODO: sent pointer to Rust and take ownership of array. + # https://stackoverflow.com/questions/37988849/safer-way-to-expose-a-c-allocated-memory-buffer-using-numpy-ctypes