ufunc unsafe from ptr for all types; #38

pola-rs · Sep 16, 2020 · 5c954dc · 5c954dc
1 parent 6bf2b8c
commit 5c954dc
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/py-polars/polars/series.py b/py-polars/polars/series.py
@@ -52,6 +52,10 @@ def __init__(
                 self._s = PySeries.new_i64(name, values)
             elif dtype == np.int32:
                 self._s = PySeries.new_i32(name, values)
+            elif dtype == np.int16:
+                self._s = PySeries.new_i16(name, values)
+            elif dtype == np.int8:
+                self._s = PySeries.new_i8(name, values)
             elif dtype == np.float32:
                 self._s = PySeries.new_f32(name, values)
             elif dtype == np.float64:
@@ -60,10 +64,14 @@ def __init__(
                 self._s = PySeries.new_str(name, values)
             elif dtype == np.bool:
                 self._s = PySeries.new_bool(name, values)
+            elif dtype == np.uint8:
+                self._s = PySeries.new_u8(name, values)
+            elif dtype == np.uint16:
+                self._s = PySeries.new_u16(name, values)
             elif dtype == np.uint32:
                 self._s = PySeries.new_u32(name, values)
             elif dtype == np.uint64:
-                self._s = PySeries.new_u32(name, np.array(values, dtype=np.uint32))
+                self._s = PySeries.new_u64(name, values)
             else:
                 raise ValueError(f"dtype: {dtype} not known")
 

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -36,12 +36,16 @@ macro_rules! init_method {
     };
 }
 
+init_method!(new_i8, i8);
 init_method!(new_i32, i32);
 init_method!(new_i64, i64);
 init_method!(new_f32, f32);
 init_method!(new_f64, f64);
 init_method!(new_bool, bool);
+init_method!(new_u8, u8);
+init_method!(new_u16, u16);
 init_method!(new_u32, u32);
+init_method!(new_u64, u64);
 init_method!(new_date32, i32);
 init_method!(new_date64, i64);
 init_method!(new_duration_ns, i64);
@@ -322,7 +326,7 @@ impl PySeries {
     }
 
     pub fn unsafe_from_ptr_f64(&self, ptr: usize, len: usize) -> Self {
-        let v: Vec<f64> = unsafe { npy::vec_from_ptr(ptr, len) };
+        let v = 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 =
@@ -331,6 +335,36 @@ impl PySeries {
     }
 }
 
+macro_rules! impl_unsafe_from_ptr {
+    ($name:ident, $series_variant:ident) => {
+        #[pymethods]
+        impl PySeries {
+            pub fn $name(&self, ptr: usize, len: usize) -> Self {
+                let v = 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::$series_variant(ca))
+            }
+        }
+    };
+}
+
+impl_unsafe_from_ptr!(unsafe_from_ptr_f32, Float32);
+impl_unsafe_from_ptr!(unsafe_from_ptr_u8, UInt8);
+impl_unsafe_from_ptr!(unsafe_from_ptr_u16, UInt16);
+impl_unsafe_from_ptr!(unsafe_from_ptr_u32, UInt32);
+impl_unsafe_from_ptr!(unsafe_from_ptr_u64, UInt64);
+impl_unsafe_from_ptr!(unsafe_from_ptr_i8, Int8);
+impl_unsafe_from_ptr!(unsafe_from_ptr_i16, Int16);
+impl_unsafe_from_ptr!(unsafe_from_ptr_i32, Int32);
+impl_unsafe_from_ptr!(unsafe_from_ptr_i64, Int64);
+
 macro_rules! impl_cast {
     ($name:ident, $type:ty) => {
         #[pymethods]