(arrow) branch main updated: GH-47371, GH-48281: [Python][CI] Fix Numba-CUDA interop (#48284)

[apitrou]

This is an automated email from the ASF dual-hosted git repository.

apitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new 57cb17259c GH-47371, GH-48281: [Python][CI] Fix Numba-CUDA interop 
(#48284)
57cb17259c is described below

commit 57cb17259cdbebec0741dfc20aff210f32a80b1e
Author: Graham Markall <535640+gmark...@users.noreply.github.com>
AuthorDate: Tue Dec 2 16:15:04 2025 +0000

    GH-47371, GH-48281: [Python][CI] Fix Numba-CUDA interop (#48284)
    
    ### Rationale for this change
    
    The extra that specifies the CUDA version for Numba-CUDA is required for 
the CUDA Python bindings to be installed as a Numba-CUDA dependency. Presently 
the Numba-CUDA interop tests are skipped, and it would be preferable to enable 
them again and ensure the test suite doesn't skip them due to lack of 
dependencies.
    
    ### What changes are included in this PR?
    
    - Addition of the appropriate extra to the Numba-CUDA installation, so that 
the bindings will be installed in the CI environment.
    - Removal of the skip for the Numba-CUDA interop tests.
    
    ### Are these changes tested?
    
    Yes (assuming the CI setup works correctly again with these changes)
    
    ### Are there any user-facing changes?
    
    No.
    * GitHub Issue: #47371
    * GitHub Issue: #48281
    
    Authored-by: Graham Markall <gmark...@nvidia.com>
    Signed-off-by: Antoine Pitrou <anto...@python.org>
---
 ci/docker/linux-apt-python-3.dockerfile         |  3 ++-
 ci/scripts/install_numba.sh                     | 21 ++++++++++++++-------
 compose.yaml                                    |  1 +
 python/pyarrow/tests/test_cuda_numba_interop.py |  3 ---
 4 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/ci/docker/linux-apt-python-3.dockerfile 
b/ci/docker/linux-apt-python-3.dockerfile
index d68bed2628..d5e7bf0088 100644
--- a/ci/docker/linux-apt-python-3.dockerfile
+++ b/ci/docker/linux-apt-python-3.dockerfile
@@ -33,9 +33,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \
 
 ARG numba
 ARG numba_cuda
+ARG cuda
 COPY ci/scripts/install_numba.sh /arrow/ci/scripts/
 RUN if [ "${numba}" != "" ]; then \
-        /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \
+        /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} ${cuda} \
     ; fi
 
 ENV ARROW_ACERO=ON \
diff --git a/ci/scripts/install_numba.sh b/ci/scripts/install_numba.sh
index 053579bd6d..91c36ca571 100755
--- a/ci/scripts/install_numba.sh
+++ b/ci/scripts/install_numba.sh
@@ -19,8 +19,8 @@
 
 set -e
 
-if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then
-  echo "Usage: $0 <numba version> [numba-cuda version]"
+if [ "$#" -ne 1 ] && [ "$#" -ne 2 ] && [ "$#" -ne 3 ]; then
+  echo "Usage: $0 <numba version> [numba-cuda version] [CUDA version]"
   exit 1
 fi
 
@@ -34,8 +34,6 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then
   . "${ARROW_PYTHON_VENV}/bin/activate"
 fi
 
-# TODO: GH-47371 (install Python CUDA bindings explicitly)
-
 if [ "${numba}" = "master" ]; then
   pip install https://github.com/numba/numba/archive/main.tar.gz#egg=numba
 elif [ "${numba}" = "latest" ]; then
@@ -50,10 +48,19 @@ fi
 
 numba_cuda=$2
 
+DEFAULT_CUDA_VERSION="12"
+
+if [ "$#" -eq 3 ]; then
+  # Extract the CUDA major version only
+  cuda_version="${3%%.*}"
+else
+  cuda_version="${DEFAULT_CUDA_VERSION}"
+fi
+
 if [ "${numba_cuda}" = "master" ]; then
-  pip install 
https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda
+  pip install "numba-cuda[cu${cuda_version}] @ 
https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz";
 elif [ "${numba_cuda}" = "latest" ]; then
-  pip install numba-cuda
+  pip install "numba-cuda[cu${cuda_version}]"
 else
-  pip install "numba-cuda==${numba_cuda}"
+  pip install "numba-cuda[cu${cuda_version}]==${numba_cuda}"
 fi
diff --git a/compose.yaml b/compose.yaml
index e83f1446cd..8abc9ad660 100644
--- a/compose.yaml
+++ b/compose.yaml
@@ -999,6 +999,7 @@ services:
         base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
         numba: ${NUMBA}
         numba_cuda: ${NUMBA_CUDA}
+        cuda: ${CUDA}
     shm_size: *shm-size
     environment:
       <<: [*common, *ccache, *sccache]
diff --git a/python/pyarrow/tests/test_cuda_numba_interop.py 
b/python/pyarrow/tests/test_cuda_numba_interop.py
index 0ce25bb608..876f3c7f76 100644
--- a/python/pyarrow/tests/test_cuda_numba_interop.py
+++ b/python/pyarrow/tests/test_cuda_numba_interop.py
@@ -22,9 +22,6 @@ try:
 except ImportError:
     pytestmark = pytest.mark.numpy
 
-pytest.skip("Numba integration tests broken by Numba API changes, see 
GH-48265",
-            allow_module_level=True)
-
 dtypes = ['uint8', 'int16', 'float32']
 cuda = pytest.importorskip("pyarrow.cuda")
 nb_cuda = pytest.importorskip("numba.cuda")