[python] #38 Create 64 byte aligned writable numpy array

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)

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<PyArray1<$type>> {
+            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::<PySeries>()?;
-    m.add_class::<PyDataFrame>()?;
+    m.add_class::<PySeries>().unwrap();
+    m.add_class::<PyDataFrame>().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(())
 }

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<T: Element>(size: usize) -> Py<PyArray1<T>> {
+    let mut buf: Vec<T> = 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()
+}

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::*};

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