feat(python): allow internal api to get pointer to values buffer

py-polars/polars/internals/series/series.py

@@ -312,6 +312,16 @@ def _from_pandas(
             pandas_to_pyseries(name, values, nan_to_none=nan_to_none)
         )
 
+    def _get_ptr(self) -> int:
+        """
+        Get a pointer to the start of the values buffer of a numeric Series.
+
+        This will raise an error if the
+        ``Series`` contains multiple chunks
+
+        """
+        return self._s.get_ptr()
+
     @property
     def dtype(self) -> type[DataType]:
         """

py-polars/src/error.rs

@@ -68,3 +68,10 @@ create_exception!(exceptions, ShapeError, PyException);
 create_exception!(exceptions, SchemaError, PyException);
 create_exception!(exceptions, DuplicateError, PyException);
 create_exception!(exceptions, InvalidOperationError, PyException);
+
+#[macro_export]
+macro_rules! raise_err(
+    ($msg:expr, $err:ident) => {{
+        Err(PolarsError::$err($msg.into())).map_err(PyPolarsErr::from)?;
+    }}
+);

py-polars/src/npy.rs

@@ -4,8 +4,15 @@ use ndarray::IntoDimension;
 use numpy::npyffi::types::npy_intp;
 use numpy::npyffi::{self, flags};
 use numpy::{Element, PyArray1, ToNpyDims, PY_ARRAY_API};
+use polars_core::prelude::*;
 use polars_core::utils::arrow::types::NativeType;
+use polars_core::with_match_physical_numeric_polars_type;
 use pyo3::prelude::*;
+use pyo3::types::PyTuple;
+
+use crate::error::PyPolarsErr;
+use crate::raise_err;
+use crate::series::PySeries;
 
 /// Create an empty numpy array arrows 64 byte alignment
 ///
@@ -53,3 +60,82 @@ pub fn get_refcnt<T>(pyarray: &PyArray1<T>) -> isize {
         refcnt
     }
 }
+
+macro_rules! impl_ufuncs {
+    ($name:ident, $type:ident, $unsafe_from_ptr_method:ident) => {
+        #[pymethods]
+        impl PySeries {
+            // applies a ufunc by accepting a lambda out: ufunc(*args, out=out)
+            // the out array is allocated in this method, send to Python and once the ufunc is applied
+            // ownership is taken by Rust again to prevent memory leak.
+            // if the ufunc fails, we first must take ownership back.
+            pub fn $name(&self, lambda: &PyAny) -> PyResult<PySeries> {
+                // numpy array object, and a *mut ptr
+                Python::with_gil(|py| {
+                    let size = self.len();
+                    let (out_array, av) =
+                        unsafe { aligned_array::<<$type as PolarsNumericType>::Native>(py, size) };
+
+                    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!(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!(get_refcnt(out_array), 3);
+
+                            let validity = self.series.chunks()[0].validity().cloned();
+                            let ca = ChunkedArray::<$type>::new_from_owned_with_null_bitmap(
+                                self.name(),
+                                av,
+                                validity,
+                            );
+                            PySeries::new(ca.into_series())
+                        }
+                        Err(e) => {
+                            // return error information
+                            return Err(e);
+                        }
+                    };
+
+                    Ok(s)
+                })
+            }
+        }
+    };
+}
+impl_ufuncs!(apply_ufunc_f32, Float32Type, unsafe_from_ptr_f32);
+impl_ufuncs!(apply_ufunc_f64, Float64Type, unsafe_from_ptr_f64);
+impl_ufuncs!(apply_ufunc_u8, UInt8Type, unsafe_from_ptr_u8);
+impl_ufuncs!(apply_ufunc_u16, UInt16Type, unsafe_from_ptr_u16);
+impl_ufuncs!(apply_ufunc_u32, UInt32Type, unsafe_from_ptr_u32);
+impl_ufuncs!(apply_ufunc_u64, UInt64Type, unsafe_from_ptr_u64);
+impl_ufuncs!(apply_ufunc_i8, Int8Type, unsafe_from_ptr_i8);
+impl_ufuncs!(apply_ufunc_i16, Int16Type, unsafe_from_ptr_i16);
+impl_ufuncs!(apply_ufunc_i32, Int32Type, unsafe_from_ptr_i32);
+impl_ufuncs!(apply_ufunc_i64, Int64Type, unsafe_from_ptr_i64);
+
+fn get_ptr<T: PolarsNumericType>(ca: &ChunkedArray<T>) -> usize {
+    let arr = ca.downcast_iter().next().unwrap();
+    arr.values().as_ptr() as usize
+}
+
+#[pymethods]
+impl PySeries {
+    pub fn get_ptr(&self) -> PyResult<usize> {
+        let s = self.series.to_physical_repr();
+        let arrays = s.chunks();
+        if arrays.len() != 1 {
+            let msg = "Only can take pointer, if the 'series' contains a single chunk";
+            raise_err!(msg, ComputeError);
+        }
+        Ok(with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
+            let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
+            get_ptr(ca)
+        }))
+    }
+}

py-polars/src/series.rs

@@ -4,18 +4,18 @@ use polars_core::series::IsSorted;
 use polars_core::utils::{flatten_series, CustomIterTools};
 use pyo3::exceptions::{PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::{PyBytes, PyList, PyTuple};
+use pyo3::types::{PyBytes, PyList};
 use pyo3::Python;
 
 use crate::apply::series::{call_lambda_and_extract, ApplyLambda};
 use crate::arrow_interop::to_rust::array_to_rust;
 use crate::dataframe::PyDataFrame;
 use crate::error::PyPolarsErr;
 use crate::list_construction::py_seq_to_list;
-use crate::npy::{aligned_array, get_refcnt};
 use crate::prelude::*;
 use crate::set::set_at_idx;
 use crate::{apply_method_all_arrow_series2, arrow_interop};
+
 #[pyclass]
 #[repr(transparent)]
 #[derive(Clone)]
@@ -1142,64 +1142,6 @@ impl PySeries {
     }
 }
 
-macro_rules! impl_ufuncs {
-    ($name:ident, $type:ident, $unsafe_from_ptr_method:ident) => {
-        #[pymethods]
-        impl PySeries {
-            // applies a ufunc by accepting a lambda out: ufunc(*args, out=out)
-            // the out array is allocated in this method, send to Python and once the ufunc is applied
-            // ownership is taken by Rust again to prevent memory leak.
-            // if the ufunc fails, we first must take ownership back.
-            pub fn $name(&self, lambda: &PyAny) -> PyResult<PySeries> {
-                // numpy array object, and a *mut ptr
-                Python::with_gil(|py| {
-                    let size = self.len();
-                    let (out_array, av) =
-                        unsafe { aligned_array::<<$type as PolarsNumericType>::Native>(py, size) };
-
-                    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!(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!(get_refcnt(out_array), 3);
-
-                            let validity = self.series.chunks()[0].validity().cloned();
-                            let ca = ChunkedArray::<$type>::new_from_owned_with_null_bitmap(
-                                self.name(),
-                                av,
-                                validity,
-                            );
-                            PySeries::new(ca.into_series())
-                        }
-                        Err(e) => {
-                            // return error information
-                            return Err(e);
-                        }
-                    };
-
-                    Ok(s)
-                })
-            }
-        }
-    };
-}
-impl_ufuncs!(apply_ufunc_f32, Float32Type, unsafe_from_ptr_f32);
-impl_ufuncs!(apply_ufunc_f64, Float64Type, unsafe_from_ptr_f64);
-impl_ufuncs!(apply_ufunc_u8, UInt8Type, unsafe_from_ptr_u8);
-impl_ufuncs!(apply_ufunc_u16, UInt16Type, unsafe_from_ptr_u16);
-impl_ufuncs!(apply_ufunc_u32, UInt32Type, unsafe_from_ptr_u32);
-impl_ufuncs!(apply_ufunc_u64, UInt64Type, unsafe_from_ptr_u64);
-impl_ufuncs!(apply_ufunc_i8, Int8Type, unsafe_from_ptr_i8);
-impl_ufuncs!(apply_ufunc_i16, Int16Type, unsafe_from_ptr_i16);
-impl_ufuncs!(apply_ufunc_i32, Int32Type, unsafe_from_ptr_i32);
-impl_ufuncs!(apply_ufunc_i64, Int64Type, unsafe_from_ptr_i64);
-
 macro_rules! impl_set_with_mask {
     ($name:ident, $native:ty, $cast:ident, $variant:ident) => {
         fn $name(

py-polars/tests/unit/test_series.py

@@ -2544,3 +2544,15 @@ def test_item() -> None:
     s = pl.Series("a", [])
     with pytest.raises(ValueError):
         s.item()
+
+
+def test_ptr() -> None:
+    # not much to test on the ptr value itself.
+    s = pl.Series([1, None, 3])
+
+    ptr = s._get_ptr()
+    assert isinstance(ptr, int)
+    s2 = s.append(pl.Series([1, 2]))
+
+    ptr2 = s2.rechunk()._get_ptr()
+    assert ptr != ptr2