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 <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>

Author: gmarkall

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 \

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

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]

python/pyarrow/tests/test_cuda_numba_interop.py

@@ -22,9 +22,6 @@
 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")