rust take ownership of output array: #38

py-polars/polars/ffi.py

@@ -8,6 +8,7 @@
     aligned_array_f64,
     aligned_array_i32,
     aligned_array_i64,
+    series_from_ptr_f64,
 )
 
 # https://stackoverflow.com/questions/4355524/getting-data-from-ctypes-array-into-numpy

py-polars/polars/series.py

@@ -2,7 +2,7 @@
 from .polars import PySeries
 import numpy as np
 from typing import Optional, List, Sequence, Union, Any
-from .ffi import ptr_to_numpy, aligned_array_f64
+from .ffi import ptr_to_numpy, aligned_array_f64, series_from_ptr_f64
 import ctypes
 from numbers import Number
 
@@ -704,7 +704,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             (out, ptr) = aligned_array_f64(self.len())
             kwargs["out"] = out
             ufunc(*args, **kwargs)
-            return kwargs["out"]
+            return wrap_s(series_from_ptr_f64(self.name, ptr, self.len()))
 
         else:
             return NotImplemented

py-polars/src/lib.rs

@@ -8,6 +8,8 @@ pub mod npy;
 pub mod series;
 
 use crate::{dataframe::PyDataFrame, series::PySeries};
+use polars::chunked_array::builder::AlignedVec;
+use polars::prelude::*;
 
 macro_rules! create_aligned_buffer {
     ($name:ident, $type:ty) => {
@@ -24,6 +26,14 @@ create_aligned_buffer!(aligned_array_f64, f64);
 create_aligned_buffer!(aligned_array_i32, i32);
 create_aligned_buffer!(aligned_array_i64, i64);
 
+#[pyfunction]
+pub fn series_from_ptr_f64(name: &str, ptr: usize, len: usize) -> PySeries {
+    let v: Vec<f64> = unsafe { npy::vec_from_ptr(ptr, len) };
+    let av = AlignedVec::new(v).unwrap();
+    let ca = ChunkedArray::new_from_aligned_vec(name, av);
+    PySeries::new(Series::Float64(ca))
+}
+
 #[pymodule]
 fn polars(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PySeries>().unwrap();
@@ -32,5 +42,7 @@ fn polars(_py: Python, m: &PyModule) -> PyResult<()> {
     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();
+    m.add_wrapped(wrap_pyfunction!(series_from_ptr_f64))
+        .unwrap();
     Ok(())
 }

py-polars/src/npy.rs

@@ -107,3 +107,8 @@ pub fn aligned_array<T: Element>(size: usize) -> (Py<PyArray1<T>>, *mut T) {
         (PyArray1::from_owned_ptr(py, ptr).to_owned(), 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)
+}

py-polars/tests/test_series.py

@@ -137,7 +137,7 @@ def test_view():
 
 
 def test_numpy_interface():
-    a = aligned_array_f32(10)
+    a, ptr = aligned_array_f32(10)
     assert a.dtype == np.float32
     assert a.shape == (10,)
     pointer, read_only_flag = a.__array_interface__["data"]
@@ -154,5 +154,6 @@ def test_numpy_interface():
 
 def test_ufunc():
     a = Series("a", [1.0, 2.0, 3.0, 4.0])
-    print(np.multiply(a, 4))
-    assert False
+    b = np.multiply(a, 4)
+    assert isinstance(b, Series)
+    assert b == [4, 8, 12, 16]