Support ufuncs in combination with PyPy. (#1917)

py-polars/src/npy.rs

@@ -51,3 +51,15 @@ 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)
 }
+
+/// Get reference counter for numpy arrays.
+///   - For CPython: Get reference counter.
+///   - For PyPy: Reference counters for a live PyPy object = refcnt + 2 << 60.
+pub fn get_refcnt<T>(pyarray: &PyArray1<T>) -> isize {
+    let refcnt = pyarray.get_refcnt();
+    if refcnt >= (2 << 60) {
+        refcnt - (2 << 60)
+    } else {
+        refcnt
+    }
+}

py-polars/src/series.rs

@@ -5,7 +5,11 @@ use crate::datatypes::PyDataType;
 use crate::error::PyPolarsEr;
 use crate::list_construction::py_seq_to_list;
 use crate::utils::{downsample_str_to_rule, reinterpret, str_to_polarstype};
-use crate::{arrow_interop, npy::aligned_array, prelude::*};
+use crate::{
+    arrow_interop,
+    npy::{aligned_array, get_refcnt},
+    prelude::*,
+};
 use numpy::PyArray1;
 use polars_core::utils::CustomIterTools;
 use pyo3::types::{PyBytes, PyList, PyTuple};
@@ -1379,17 +1383,17 @@ macro_rules! impl_ufuncs {
                 let (out_array, av) =
                     unsafe { aligned_array::<<$type as PolarsNumericType>::Native>(py, size) };
 
-                debug_assert_eq!(out_array.get_refcnt(), 1);
+                debug_assert_eq!(get_refcnt(out_array), 1);
                 // inserting it in a tuple increase the reference count by 1.
                 let args = PyTuple::new(py, &[out_array]);
-                debug_assert_eq!(out_array.get_refcnt(), 2);
+                debug_assert_eq!(get_refcnt(out_array), 2);
 
                 // whatever the result, we must take the leaked memory ownership back
                 let s = match lambda.call1(args) {
                     Ok(_) => {
                         // if this assert fails, the lambda has taken a reference to the object, so we must panic
                         // args and the lambda return have a reference, making a total of 3
-                        assert_eq!(out_array.get_refcnt(), 3);
+                        assert_eq!(get_refcnt(out_array), 3);
 
                         let validity = self.series.chunks()[0].validity().cloned();
                         let ca = ChunkedArray::<$type>::new_from_owned_with_null_bitmap(