[Python] NumPy 2 ABI changes (elsize)

[seberg]

Parent issue:

• [Python] Compatibility with NumPy 2.0 #39532

Describe the enhancement requested

As soon as we merge gh- in NumPy, it would be nice to have a small fix in here. I suspect the following changes should do things (the first two are obvious, the last is the annoying one).

The tricky remaining thing is that itemsize_ would be nice to be an intp/(Py_)ssize_t or int64 now and I am not sure if that has knock-on effects

diff --git a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc
index cb9cbe5b9..e29f5b31f 100644
--- a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc
+++ b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc
@@ -278,3 +278,3 @@ Status PyArray_NewFromPool(int nd, npy_intp* dims, PyArray_Descr* descr, MemoryP
   // * Get better performance through custom allocators
-  int64_t total_size = descr->elsize;
+  int64_t total_size = PyDataType_ELSIZE(descr);
   for (int i = 0; i < nd; ++i) {
diff --git a/python/pyarrow/src/arrow/python/numpy_convert.cc b/python/pyarrow/src/arrow/python/numpy_convert.cc
index dfee88c09..f46487a57 100644
--- a/python/pyarrow/src/arrow/python/numpy_convert.cc
+++ b/python/pyarrow/src/arrow/python/numpy_convert.cc
@@ -48,3 +48,3 @@ NumPyBuffer::NumPyBuffer(PyObject* ao) : Buffer(nullptr, 0) {
     data_ = const_cast<const uint8_t*>(ptr);
-    size_ = PyArray_SIZE(ndarray) * PyArray_DESCR(ndarray)->elsize;
+    size_ = PyArray_NBYTES(ndarray);
     capacity_ = size_;
diff --git a/python/pyarrow/src/arrow/python/numpy_interop.h b/python/pyarrow/src/arrow/python/numpy_interop.h
index ce7baed25..1de28c20d 100644
--- a/python/pyarrow/src/arrow/python/numpy_interop.h
+++ b/python/pyarrow/src/arrow/python/numpy_interop.h
@@ -69,2 +69,7 @@
 
+// Backported NumPy 2 API.  and can be removed if 2 is required.
+#if NPY_ABI_VERSION < 0x02000000
+  #define PyDataType_ELSIZE(descr) ((descr)->elsize)
+#endif
+
 namespace arrow {
diff --git a/python/pyarrow/src/arrow/python/numpy_to_arrow.cc b/python/pyarrow/src/arrow/python/numpy_to_arrow.cc
index 8903df31b..04572536c 100644
--- a/python/pyarrow/src/arrow/python/numpy_to_arrow.cc
+++ b/python/pyarrow/src/arrow/python/numpy_to_arrow.cc
@@ -198,3 +198,3 @@ class NumPyConverter {
     length_ = static_cast<int64_t>(PyArray_SIZE(arr_));
-    itemsize_ = static_cast<int>(PyArray_DESCR(arr_)->elsize);
+    itemsize_ = static_cast<int64_t>(PyArray_ITEMSIZE(arr_));
     stride_ = static_cast<int64_t>(PyArray_STRIDES(arr_)[0]);
@@ -298,3 +298,3 @@ class NumPyConverter {
   int64_t stride_;
-  int itemsize_;
+  int64_t itemsize_;

---

@jorisvandenbossche do you know this would affect pandas? I am now worried that pandas wheels will fail due to importing arrow. OTOH, it seems like you must have solved that issue (or it is a non-issue).

Component(s)

C++, Python

[jorisvandenbossche]

I am now worried that pandas wheels will fail due to importing arrow. OTOH, it seems like you must have solved that issue (or it is a non-issue).

PyArrow is not a build dependency, only a run dependency for pandas, so in case numpy breaks something in pyarrow, that wouldn't directly break pandas (only that functionality in pandas that requires pyarrow). Or am I misunderstanding your question?

[seberg]

Aha, so I guess we should get pandas wheels (as I suppose we have them and even if tests fail against 2.0, the new wheels tests will work when run against NumPy 1.x). But some things using arrow will be broken until we also have pyarrow wheels (with these fixes).

It is unclear how much of an effect that may be, but if you are not super worried I guess we'll have to see (if you are super worried, we can just stick to keeping elsize an int as few notice the other changes).

[jorisvandenbossche]

We do have nightly wheels building with numpy 2.0, so once your PR lands on the numpy side and we can do those fixes here, we can rather quickly have pyarrow nightly wheels that should work.

[jorisvandenbossche]

Issue resolved by pull request 40418
#40418