diff --git a/.env b/.env index 81638c67efd..6b9bea19f94 100644 --- a/.env +++ b/.env @@ -67,6 +67,7 @@ LLVM=18 MAVEN=3.8.7 NODE=20 NUMBA=latest +NUMBA_CUDA=latest 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 e215976d448..d68bed26288 100644 --- a/ci/docker/linux-apt-python-3.dockerfile +++ b/ci/docker/linux-apt-python-3.dockerfile @@ -32,9 +32,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \ -r arrow/python/requirements-test.txt ARG numba +ARG numba_cuda COPY ci/scripts/install_numba.sh /arrow/ci/scripts/ RUN if [ "${numba}" != "" ]; then \ - /arrow/ci/scripts/install_numba.sh ${numba} \ + /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \ ; fi ENV ARROW_ACERO=ON \ diff --git a/ci/scripts/install_numba.sh b/ci/scripts/install_numba.sh index 87c256e4b8a..22e0df2b3c6 100755 --- a/ci/scripts/install_numba.sh +++ b/ci/scripts/install_numba.sh @@ -19,8 +19,8 @@ set -e -if [ "$#" -ne 1 ]; then - echo "Usage: $0 " +if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then + echo "Usage: $0 [numba-cuda version]" exit 1 fi @@ -41,3 +41,17 @@ elif [ "${numba}" = "latest" ]; then else pip install "numba==${numba}" fi + +if [ "$#" -eq 1 ]; then + exit 0 +fi + +numba_cuda=$2 + +if [ "${numba_cuda}" = "master" ]; then + pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda +elif [ "${numba_cuda}" = "latest" ]; then + pip install numba-cuda +else + pip install "numba-cuda==${numba_cuda}" +fi diff --git a/docker-compose.yml b/docker-compose.yml index 528e4cc8fce..7bdd82c5757 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -560,10 +560,11 @@ services: ubuntu-cuda-cpp: # Usage: - # docker compose build cuda-cpp - # docker compose run --rm cuda-cpp + # docker compose build ubuntu-cuda-cpp + # docker compose run --rm ubuntu-cuda-cpp # Parameters: # ARCH: amd64 + # UBUNTU: 22.04, 24.04 # CUDA: image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp build: @@ -946,13 +947,15 @@ services: ubuntu-cuda-python: # Usage: - # docker compose build cuda-cpp - # docker compose build cuda-python - # docker compose run --rm cuda-python + # docker compose build ubuntu-cuda-cpp + # docker compose build ubuntu-cuda-python + # docker compose run --rm ubuntu-cuda-python # Parameters: # ARCH: amd64 # CUDA: # UBUNTU: 22.04, 24.04 + # NUMBA: master, latest, + # NUMBA_CUDA: master, latest, image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-python-3 build: context: . @@ -962,6 +965,7 @@ services: args: base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp numba: ${NUMBA} + numba_cuda: ${NUMBA_CUDA} shm_size: *shm-size environment: <<: [*common, *ccache, *sccache] diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx index 340f7e35c73..96e5c139492 100644 --- a/python/pyarrow/_cuda.pyx +++ b/python/pyarrow/_cuda.pyx @@ -451,9 +451,9 @@ cdef class CudaBuffer(Buffer): Device buffer as a view of numba MemoryPointer. """ ctx = Context.from_numba(mem.context) - if mem.device_pointer.value is None and mem.size==0: + if mem.device_pointer_value is None and mem.size==0: return ctx.new_buffer(0) - return ctx.foreign_buffer(mem.device_pointer.value, mem.size, base=mem) + return ctx.foreign_buffer(mem.device_pointer_value, mem.size, base=mem) def to_numba(self): """Return numba memory pointer of CudaBuffer instance.