GH-47371: [Python] Fix installation of Numba-CUDA in CI environment

As per the Numba-CUDA docs, `numba-cuda` should be installed with a
variant depending on the CUDA version. Failing to specify the variant
leads to `numba-cuda` being installed without the CUDA toolkit wheels
that it needs. This commit adds the variant specifier to the
installation.

I think Arrow is using CUDA 11.7 in CI images, so I've changed the
Numba-CUDA version to 0.18.1, which is the last version to support CUDA
11 (the Numba-CUDA policy is to support all minor versions within the
last two major versions of CUDA). IF Arrow moves to CUDA 12 for CI
images, this could be changed back to `latest` in future.

Numba-CUDA installation docs: https://nvidia.github.io/numba-cuda/user/installation.html

Author: gmarkall

.env

@@ -67,7 +67,7 @@ LLVM=18
 MAVEN=3.8.7
 NODE=20
 NUMBA=latest
-NUMBA_CUDA=latest
+NUMBA_CUDA=0.18.1
 NUMPY=latest
 PANDAS=latest
 PYTHON=3.9

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
 
@@ -48,10 +48,21 @@ fi
 
 numba_cuda=$2
 
+if [ "$#" -eq 3 ]; then
+  cuda=$3
+else
+  # Default to CUDA 11
+  cuda=11
+fi
+
+# Variants are cu11, cu12, etc. depending on CUDA version
+variant="cu${cuda%%.*}"
+
 if [ "${numba_cuda}" = "master" ]; then
-  pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda
+  pip install "numba-cuda[$variant] @ https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz"
 elif [ "${numba_cuda}" = "latest" ]; then
-  pip install numba-cuda
+  pip install numba-cuda[$variant]
 else
-  pip install "numba-cuda==${numba_cuda}"
+  echo installing for $variant
+  pip install "numba-cuda[$variant]==${numba_cuda}"
 fi

docker-compose.yml

@@ -966,6 +966,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]