ARROW-17435: [CI][Python][CUDA] Install Numba for CUDA interop tests (a…

…pache#13899)

Checked locally on a Ubuntu 20.04 host with:
```
archery docker run ubuntu-cuda-python
```

Authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>

.env

@@ -67,6 +67,7 @@ KARTOTHEK=latest
 LLVM=13
 MAVEN=3.5.4
 NODE=16
+NUMBA=latest
 NUMPY=latest
 PANDAS=latest
 PYTHON=3.8

ci/docker/linux-apt-python-3.dockerfile

@@ -39,6 +39,12 @@ RUN pip install \
     -r arrow/python/requirements-build.txt \
     -r arrow/python/requirements-test.txt
 
+ARG numba
+COPY ci/scripts/install_numba.sh /arrow/ci/scripts/
+RUN if [ "${numba}" != "" ]; then \
+        /arrow/ci/scripts/install_numba.sh ${numba} \
+    ; fi
+
 ENV ARROW_PYTHON=ON \
     ARROW_BUILD_STATIC=OFF \
     ARROW_BUILD_TESTS=OFF \

ci/scripts/install_numba.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+set -e
+
+if [ "$#" -ne 1 ]; then
+  echo "Usage: $0 <numba version>"
+  exit 1
+fi
+
+numba=$1
+
+if [ "${numba}" = "master" ]; then
+  pip install https://github.com/numba/numba/archive/main.tar.gz#egg=numba
+elif [ "${numba}" = "latest" ]; then
+  pip install numba
+else
+  pip install numba==${numba}
+fi

ci/scripts/python_test.sh

@@ -54,4 +54,4 @@ export PYARROW_TEST_ORC
 export PYARROW_TEST_PARQUET
 export PYARROW_TEST_S3
 
-pytest -r s -v ${PYTEST_ARGS} --pyargs pyarrow
+pytest -r s ${PYTEST_ARGS} --pyargs pyarrow

docker-compose.yml

@@ -487,7 +487,7 @@ services:
         llvm: ${LLVM}
     shm_size: *shm-size
     ulimits: *ulimits
-    environment:
+    environment: &cuda-environment
       <<: *ccache
       ARROW_BUILD_STATIC: "OFF"
       ARROW_CUDA: "ON"
@@ -775,10 +775,10 @@ services:
         - ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-python-3
       args:
         base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
+        numba: ${NUMBA}
     shm_size: *shm-size
     environment:
-      <<: *ccache
-      ARROW_CUDA: "ON"
+      <<: [ *ccache, *cuda-environment ]
     volumes: *ubuntu-volumes
     command: &python-command >
       /bin/bash -c "

python/setup.cfg

@@ -28,6 +28,8 @@ build-dir  = doc/_build
 addopts = --ignore=scripts
 filterwarnings =
     error:The SparseDataFrame:FutureWarning
+# Get a debug traceback when a test takes a really long time
+faulthandler_timeout = 300
 
 [pep8]
 ignore = E211,E225,E226,E227,E402,W504