keep null bitmap with ufunc

py-polars/polars/series.py

@@ -691,6 +691,8 @@ def view(self) -> np.ndarray:
         return array
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        if self._s.n_chunks() > 0:
+            self._s.rechunk(in_place=True)
 
         if method == "__call__":
             args = []
@@ -704,7 +706,9 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             (out, ptr) = aligned_array_f64(self.len())
             kwargs["out"] = out
             ufunc(*args, **kwargs)
-            return wrap_s(series_from_ptr_f64(self.name, ptr, self.len()))
+            # get method for current dtype
+            f = getattr(self._s, f"unsafe_from_ptr_{self.dtype}")
+            return wrap_s(f(ptr, self.len()))
 
         else:
             return NotImplemented

py-polars/src/series.rs

@@ -1,5 +1,7 @@
 use crate::error::PyPolarsEr;
+use crate::npy;
 use numpy::PyArray1;
+use polars::chunked_array::builder::get_bitmap;
 use polars::prelude::*;
 use pyo3::types::PyList;
 use pyo3::{exceptions::RuntimeError, prelude::*, Python};
@@ -318,6 +320,15 @@ impl PySeries {
             _ => todo!(),
         }
     }
+
+    pub fn unsafe_from_ptr_f64(&self, ptr: usize, len: usize) -> Self {
+        let v: Vec<f64> = unsafe { npy::vec_from_ptr(ptr, len) };
+        let av = AlignedVec::new(v).unwrap();
+        let (null_count, null_bitmap) = get_bitmap(self.series.chunks()[0].as_ref());
+        let ca =
+            ChunkedArray::new_from_owned_with_null_bitmap(self.name(), av, null_bitmap, null_count);
+        Self::new(Series::Float64(ca))
+    }
 }
 
 macro_rules! impl_cast {