From 260fef9d3a679533c5a00a9d3dc3dcde4b4695df Mon Sep 17 00:00:00 2001
From: Graham Markall <gmarkall@nvidia.com>
Date: Tue, 19 Aug 2025 14:59:24 +0100
Subject: [PATCH] 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
---
 .env                                    |  2 +-
 ci/docker/linux-apt-python-3.dockerfile |  3 ++-
 ci/scripts/install_numba.sh             | 21 ++++++++++++++++-----
 docker-compose.yml                      |  1 +
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/.env b/.env
index 84f3890b254..2f3741442fe 100644
--- a/.env
+++ b/.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
diff --git a/ci/docker/linux-apt-python-3.dockerfile b/ci/docker/linux-apt-python-3.dockerfile
index d68bed26288..d5e7bf00886 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 22e0df2b3c6..e8951b05166 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
 
@@ -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
diff --git a/docker-compose.yml b/docker-compose.yml
index 4a40e382d8d..695f85a9f41 100644
--- a/docker-compose.yml
+++ b/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]