[python] implement get interface for single value retrieval

pola-rs · Sep 16, 2020 · df6268d · df6268d
1 parent 1d0eded
commit df6268d
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 10 deletions.
diff --git a/py-polars/polars/series.py b/py-polars/polars/series.py
@@ -236,8 +236,10 @@ def __getitem__(self, item):
             if stride != 1:
                 return NotImplemented
             return self.slice(start, stop - start)
-        # TODO: implement
-        return NotImplemented
+        f = getattr(self._s, f"get_{self.dtype}", None)
+        if f is None:
+            return NotImplemented
+        return f(item)
 
     @property
     def dtype(self):

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -340,17 +340,34 @@ impl PySeries {
             _ => todo!(),
         }
     }
+}
 
-    pub fn unsafe_from_ptr_f64(&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::Float64(ca))
-    }
+macro_rules! impl_get {
+    ($name:ident, $series_variant:ident, $type:ty) => {
+        #[pymethods]
+        impl PySeries {
+            pub fn $name(&self, index: usize) -> Option<$type> {
+                if let Series::$series_variant(ca) = &self.series {
+                    ca.get(index)
+                } else {
+                    None
+                }
+            }
+        }
+    };
 }
 
+impl_get!(get_f32, Float32, f32);
+impl_get!(get_f64, Float64, f64);
+impl_get!(get_u8, UInt8, u8);
+impl_get!(get_u16, UInt16, u16);
+impl_get!(get_u32, UInt32, u32);
+impl_get!(get_u64, UInt64, u64);
+impl_get!(get_i8, Int8, i8);
+impl_get!(get_i16, Int16, i16);
+impl_get!(get_i32, Int32, i32);
+impl_get!(get_i64, Int64, i64);
+
 macro_rules! impl_unsafe_from_ptr {
     ($name:ident, $series_variant:ident) => {
         #[pymethods]
@@ -372,6 +389,7 @@ macro_rules! impl_unsafe_from_ptr {
 }
 
 impl_unsafe_from_ptr!(unsafe_from_ptr_f32, Float32);
+impl_unsafe_from_ptr!(unsafe_from_ptr_f64, Float64);
 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);

diff --git a/py-polars/tests/test_series.py b/py-polars/tests/test_series.py
@@ -137,6 +137,7 @@ def test_view():
 
 
 def test_numpy_interface():
+    # this isn't used anymore.
     a, ptr = aligned_array_f32(10)
     assert a.dtype == np.float32
     assert a.shape == (10,)
@@ -157,3 +158,14 @@ def test_ufunc():
     b = np.multiply(a, 4)
     assert isinstance(b, Series)
     assert b == [4, 8, 12, 16]
+
+    # test if null bitmask is preserved
+    a = Series("a", [1.0, None, 3.0], nullable=True)
+    b = np.exp(a)
+    assert b.null_count() == 1
+
+
+def test_get():
+    a = Series("a", [1, 2, 3])
+    assert a[0] == 1
+    assert a[:2] == [1, 2]