GH-49437: [Python][GPU] Support legacy and native Numba CUDA context handles

.github/workflows/cuda_extra.yml

@@ -90,10 +90,6 @@ jobs:
             ubuntu: 24.04
             image: ubuntu-cuda-python
             title: AMD64 Ubuntu 24 CUDA 12.9.0 Python
-          - cuda: 11.7.1
-            ubuntu: 22.04
-            image: ubuntu-cuda-python
-            title: AMD64 Ubuntu 22 CUDA 11.7.1 Python
     env:
       ARCHERY_DEBUG: 1
       ARROW_ENABLE_TIMING_TESTS: OFF

ci/scripts/install_numba.sh

@@ -34,6 +34,12 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then
   . "${ARROW_PYTHON_VENV}/bin/activate"
 fi
 
+# Numba-CUDA 0.30.4 uses np.row_stack, which was removed in NumPy 2.5.
+# TODO: Remove after https://github.com/NVIDIA/numba-cuda/pull/916 is released.
+if [ "$#" -ge 2 ]; then
+  pip install "numpy<2.5"
+fi
+
 if [ "${numba}" = "master" ]; then
   pip install https://github.com/numba/numba/archive/main.tar.gz#egg=numba
 elif [ "${numba}" = "latest" ]; then

python/pyarrow/_cuda.pyx

@@ -16,13 +16,43 @@
 # under the License.
 
 
+import ctypes
+
+from pyarrow.vendored.version import Version
+
 from pyarrow.lib cimport *
 from pyarrow.includes.libarrow_cuda cimport *
 from pyarrow.lib import allocate_buffer, as_buffer, ArrowTypeError
 from pyarrow.util import get_contiguous_span
 cimport cpython as cp
 
 
+_NUMBA_CUDA_NATIVE_CONTEXT_VERSION = Version("0.28")
+
+
+def _numba_context_handle_value(handle):
+    """Return the integer value of a legacy or native Numba context handle."""
+    if hasattr(handle, "value"):
+        return handle.value
+    return int(handle)
+
+
+def _make_numba_context_handle(uintptr_t handle):
+    """Create the context handle representation expected by Numba."""
+    import numba.cuda
+
+    # numba-cuda 0.28 replaced ctypes context handles with CUContext type.
+    if getattr(numba.cuda, "implementation", None) == "NVIDIA":
+        import numba_cuda
+
+        version = Version(numba_cuda.__version__)
+        if version >= _NUMBA_CUDA_NATIVE_CONTEXT_VERSION:
+            from cuda.bindings.driver import CUcontext
+            return CUcontext(handle)
+
+    return ctypes.c_void_p(handle)
+
+
 cdef class Context(_Weakrefable):
     """
     CUDA driver context.
@@ -81,7 +111,7 @@ cdef class Context(_Weakrefable):
             import numba.cuda
             context = numba.cuda.current_context()
         return Context(device_number=context.device.id,
-                       handle=context.handle.value)
+                       handle=_numba_context_handle_value(context.handle))
 
     def to_numba(self):
         """
@@ -92,10 +122,9 @@ cdef class Context(_Weakrefable):
         context : numba.cuda.cudadrv.driver.Context
           Numba CUDA context instance.
         """
-        import ctypes
         import numba.cuda
         device = numba.cuda.gpus[self.device_number]
-        handle = ctypes.c_void_p(self.handle)
+        handle = _make_numba_context_handle(self.handle)
         context = numba.cuda.cudadrv.driver.Context(device, handle)
 
         class DummyPendingDeallocs(object):

python/pyarrow/tests/test_cuda_numba_interop.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import ctypes
+
 import pytest
 import pyarrow as pa
 try:
@@ -50,8 +52,14 @@ def teardown_module(module):
                          ids=context_choice_ids)
 def test_context(c):
     ctx, nb_ctx = context_choices[c]
-    assert ctx.handle == nb_ctx.handle.value
-    assert ctx.handle == ctx.to_numba().handle.value
+    converted = ctx.to_numba()
+    assert type(converted.handle) is type(nb_ctx.handle)
+    if isinstance(converted.handle, ctypes.c_void_p):
+        assert converted.handle.value == ctx.handle
+    else:
+        from cuda.bindings import driver
+        assert type(converted.handle) is driver.CUcontext
+        assert int(converted.handle) == ctx.handle
     ctx2 = cuda.Context.from_numba(nb_ctx)
     assert ctx.handle == ctx2.handle
     size = 10
@@ -203,7 +211,7 @@ def test_numba_context(c, dtype):
     with nb_cuda.gpus[0]:
         arr, cbuf = make_random_buffer(size, target='device',
                                        dtype=dtype, ctx=ctx)
-        assert cbuf.context.handle == nb_ctx.handle.value
+        assert cbuf.context.handle == cuda.Context.from_numba(nb_ctx).handle
         mem = cbuf.to_numba()
         darr = DeviceNDArray(arr.shape, arr.strides, arr.dtype, gpu_data=mem)
         np.testing.assert_equal(darr.copy_to_host(), arr)