ARROW-1424: [Python] Add CUDA support to pyarrow

cpp/cmake_modules/FindArrowCuda.cmake

@@ -0,0 +1,124 @@
+# 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.
+
+# - Find ARROW CUDA (arrow/gpu/cuda_api.h, libarrow_gpu.a, libarrow_gpu.so)
+#
+# This module requires Arrow from which it uses
+#   ARROW_FOUND
+#   ARROW_SEARCH_HEADER_PATHS
+#   ARROW_SEARCH_LIB_PATH
+#   ARROW_HOME
+#
+# This module defines
+#  ARROW_CUDA_INCLUDE_DIR, directory containing headers
+#  ARROW_CUDA_LIBS, directory containing arrow libraries
+#  ARROW_CUDA_STATIC_LIB, path to libarrow.a
+#  ARROW_CUDA_SHARED_LIB, path to libarrow's shared library
+#  ARROW_CUDA_SHARED_IMP_LIB, path to libarrow's import library (MSVC only)
+#  ARROW_CUDA_FOUND, whether arrow has been found
+
+#
+# TODO(ARROW-3209): rename arrow/gpu to arrow/cuda, arrow_gpu to arrow_cuda
+#
+
+include(FindPkgConfig)
+include(GNUInstallDirs)
+
+if (NOT DEFINED ARROW_FOUND)
+  if (ArrowCuda_FIND_REQUIRED)
+    find_package(Arrow REQUIRED)
+  else()
+    find_package(Arrow)
+  endif()
+endif()
+
+if (NOT ARROW_FOUND)
+  set(ARROW_CUDA_FOUND FALSE)
+  return()
+endif()
+
+find_path(ARROW_CUDA_INCLUDE_DIR arrow/gpu/cuda_api.h PATHS
+  ${ARROW_SEARCH_HEADER_PATHS}
+  NO_DEFAULT_PATH
+  )
+
+if (NOT (ARROW_CUDA_INCLUDE_DIR STREQUAL ARROW_INCLUDE_DIR))
+  set(ARROW_CUDA_WARN_MSG "Mismatch of Arrow and Arrow CUDA include directories:")
+  set(ARROW_CUDA_WARN_MSG "${ARROW_CUDA_WARN_MSG}  ARROW_INCLUDE_DIR=${ARROW_INCLUDE_DIR}")
+  set(ARROW_CUDA_WARN_MSG "${ARROW_CUDA_WARN_MSG}  ARROW_CUDA_INCLUDE_DIR=${ARROW_CUDA_INCLUDE_DIR}")
+  message(WARNING ${ARROW_CUDA_WARN_MSG})
+endif()
+
+find_library(ARROW_CUDA_LIB_PATH NAMES arrow_gpu
+  PATHS
+  ${ARROW_SEARCH_LIB_PATH}
+  NO_DEFAULT_PATH)
+get_filename_component(ARROW_CUDA_LIBS ${ARROW_CUDA_LIB_PATH} DIRECTORY)
+
+if (MSVC)
+  find_library(ARROW_CUDA_SHARED_LIBRARIES NAMES arrow_gpu
+    PATHS ${ARROW_HOME} NO_DEFAULT_PATH
+    PATH_SUFFIXES "bin" )
+  get_filename_component(ARROW_CUDA_SHARED_LIBS ${ARROW_CUDA_SHARED_LIBRARIES} PATH )
+endif()
+
+
+if (ARROW_CUDA_INCLUDE_DIR AND ARROW_CUDA_LIBS)
+  set(ARROW_CUDA_FOUND TRUE)
+  set(ARROW_CUDA_LIB_NAME arrow_gpu)
+  if (MSVC)
+    set(ARROW_CUDA_STATIC_LIB ${ARROW_CUDA_LIBS}/${ARROW_CUDA_LIB_NAME}${ARROW_MSVC_STATIC_LIB_SUFFIX}${CMAKE_STATIC_LIBRARY_SUFFIX})
+    set(ARROW_CUDA_SHARED_LIB ${ARROW_CUDA_SHARED_LIBS}/${ARROW_CUDA_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
+    set(ARROW_CUDA_SHARED_IMP_LIB ${ARROW_CUDA_LIBS}/${ARROW_CUDA_LIB_NAME}.lib)
+  else()
+    set(ARROW_CUDA_STATIC_LIB ${ARROW_LIBS}/lib${ARROW_CUDA_LIB_NAME}.a)
+    set(ARROW_CUDA_SHARED_LIB ${ARROW_LIBS}/lib${ARROW_CUDA_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
+  endif()
+endif()
+
+if (ARROW_CUDA_FOUND)
+  if (NOT ArrowCuda_FIND_QUIETLY)
+    message(STATUS "Found the Arrow CUDA library: ${ARROW_CUDA_LIB_PATH}")
+  endif()
+else()
+  if (NOT ArrowCuda_FIND_QUIETLY)
+    set(ARROW_CUDA_ERR_MSG "Could not find the Arrow CUDA library. Looked for headers")
+    set(ARROW_CUDA_ERR_MSG "${ARROW_CUDA_ERR_MSG} in ${ARROW_SEARCH_HEADER_PATHS}, and for libs")
+    set(ARROW_CUDA_ERR_MSG "${ARROW_CUDA_ERR_MSG} in ${ARROW_SEARCH_LIB_PATH}")
+    if (ArrowCuda_FIND_REQUIRED)
+      message(FATAL_ERROR "${ARROW_CUDA_ERR_MSG}")
+    else(ArrowCuda_FIND_REQUIRED)
+      message(STATUS "${ARROW_CUDA_ERR_MSG}")
+    endif (ArrowCuda_FIND_REQUIRED)
+  endif ()
+  set(ARROW_CUDA_FOUND FALSE)
+endif ()
+
+if (MSVC)
+  mark_as_advanced(
+    ARROW_CUDA_INCLUDE_DIR
+    ARROW_CUDA_STATIC_LIB
+    ARROW_CUDA_SHARED_LIB
+    ARROW_CUDA_SHARED_IMP_LIB
+  )
+else()
+  mark_as_advanced(
+    ARROW_CUDA_INCLUDE_DIR
+    ARROW_CUDA_STATIC_LIB
+    ARROW_CUDA_SHARED_LIB
+  )
+endif()

python/CMakeLists.txt

@@ -17,6 +17,9 @@
 #
 # Includes code assembled from BSD/MIT/Apache-licensed code from some 3rd-party
 # projects, including Kudu, Impala, and libdynd. See python/LICENSE.txt
+#
+# TODO(ARROW-3209): rename arrow_gpu to arrow_cuda
+#
 
 cmake_minimum_required(VERSION 2.7)
 project(pyarrow)
@@ -58,6 +61,9 @@ endif()
 
 # Top level cmake dir
 if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
+  option(PYARROW_BUILD_CUDA
+    "Build the PyArrow CUDA support"
+    OFF)
   option(PYARROW_BUILD_PARQUET
     "Build the PyArrow Parquet integration"
     OFF)
@@ -337,6 +343,32 @@ set(LINK_LIBS
   arrow_python_shared
 )
 
+if (PYARROW_BUILD_CUDA)
+  ## Arrow CUDA
+  find_package(ArrowCuda)
+  if(NOT ARROW_CUDA_FOUND)
+    message(FATAL_ERROR "Unable to locate Arrow CUDA libraries")
+  else()
+    if (PYARROW_BUNDLE_ARROW_CPP)
+      bundle_arrow_lib(ARROW_CUDA_SHARED_LIB
+        ABI_VERSION ${ARROW_ABI_VERSION}
+        SO_VERSION ${ARROW_SO_VERSION})
+      if (MSVC)
+        bundle_arrow_implib(ARROW_CUDA_SHARED_IMP_LIB)
+      endif()
+    endif()
+    if (MSVC)
+      ADD_THIRDPARTY_LIB(arrow_gpu
+        SHARED_LIB ${ARROW_CUDA_SHARED_IMP_LIB})
+    else()
+      ADD_THIRDPARTY_LIB(arrow_gpu
+        SHARED_LIB ${ARROW_CUDA_SHARED_LIB})
+    endif()
+    set(LINK_LIBS ${LINK_LIBS} arrow_gpu_shared)
+    set(CYTHON_EXTENSIONS ${CYTHON_EXTENSIONS} _cuda)
+  endif()
+endif()
+
 if (PYARROW_BUILD_PARQUET)
   ## Parquet
   find_package(Parquet)

python/pyarrow/_cuda.pxd

@@ -0,0 +1,62 @@
+# 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.
+
+from pyarrow.lib cimport *
+from pyarrow.includes.common cimport *
+from pyarrow.includes.libarrow cimport *
+from pyarrow.includes.libarrow_cuda cimport *
+
+
+cdef class Context:
+    cdef:
+        shared_ptr[CCudaContext] context
+        int device_number
+
+    cdef void init(self, const shared_ptr[CCudaContext]& ctx)
+
+
+cdef class IpcMemHandle:
+    cdef:
+        shared_ptr[CCudaIpcMemHandle] handle
+
+    cdef void init(self, shared_ptr[CCudaIpcMemHandle]& h)
+
+
+cdef class CudaBuffer(Buffer):
+    cdef:
+        shared_ptr[CCudaBuffer] cuda_buffer
+
+    cdef void init_cuda(self, const shared_ptr[CCudaBuffer]& buffer)
+
+
+cdef class HostBuffer(Buffer):
+    cdef:
+        shared_ptr[CCudaHostBuffer] host_buffer
+
+    cdef void init_host(self, const shared_ptr[CCudaHostBuffer]& buffer)
+
+
+cdef class BufferReader(NativeFile):
+    cdef:
+        CCudaBufferReader* reader
+        CudaBuffer buffer
+
+
+cdef class BufferWriter(NativeFile):
+    cdef:
+        CCudaBufferWriter* writer
+        CudaBuffer buffer