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

cpp/cmake_modules/FindArrow.cmake

@@ -79,6 +79,12 @@ find_library(ARROW_PYTHON_LIB_PATH NAMES arrow_python
   NO_DEFAULT_PATH)
 get_filename_component(ARROW_PYTHON_LIBS ${ARROW_PYTHON_LIB_PATH} DIRECTORY)
 
+find_library(ARROW_GPU_LIB_PATH NAMES arrow_gpu
+  PATHS
+  ${ARROW_SEARCH_LIB_PATH}
+  NO_DEFAULT_PATH)
+get_filename_component(ARROW_GPU_LIBS ${ARROW_GPU_LIB_PATH} DIRECTORY)
+
 if (MSVC)
   SET(CMAKE_FIND_LIBRARY_SUFFIXES ".lib" ".dll")
 
@@ -95,6 +101,13 @@ if (MSVC)
     PATH_SUFFIXES "bin" )
   get_filename_component(ARROW_SHARED_LIBS ${ARROW_SHARED_LIBRARIES} PATH )
   get_filename_component(ARROW_PYTHON_SHARED_LIBS ${ARROW_PYTHON_SHARED_LIBRARIES} PATH )
+
+  if (PYARROW_BUILD_CUDA)
+    find_library(ARROW_GPU_SHARED_LIBRARIES NAMES arrow_gpu
+      PATHS ${ARROW_HOME} NO_DEFAULT_PATH
+      PATH_SUFFIXES "bin" )
+    get_filename_component(ARROW_GPU_SHARED_LIBS ${ARROW_GPU_SHARED_LIBRARIES} PATH )
+  endif()
 endif ()
 
 if (ARROW_INCLUDE_DIR AND ARROW_LIBS)
@@ -117,10 +130,41 @@ if (ARROW_INCLUDE_DIR AND ARROW_LIBS)
   endif()
 endif()
 
+if (PYARROW_BUILD_CUDA AND ARROW_GPU_LIBS)
+  set(ARROW_GPU_FOUND TRUE)
+  set(ARROW_GPU_LIB_NAME arrow_gpu)
+  if (MSVC)
+    set(ARROW_GPU_STATIC_LIB ${ARROW_GPU_LIBS}/${ARROW_GPU_LIB_NAME}${ARROW_MSVC_STATIC_LIB_SUFFIX}${CMAKE_STATIC_LIBRARY_SUFFIX})
+    set(ARROW_GPU_SHARED_LIB ${ARROW_GPU_SHARED_LIBS}/${ARROW_GPU_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
+    set(ARROW_GPU_SHARED_IMP_LIB ${ARROW_GPU_LIBS}/${ARROW_GPU_LIB_NAME}.lib)
+  else()
+    set(ARROW_GPU_STATIC_LIB ${ARROW_LIBS}/lib${ARROW_GPU_LIB_NAME}.a)
+    set(ARROW_GPU_SHARED_LIB ${ARROW_LIBS}/lib${ARROW_GPU_LIB_NAME}${CMAKE_SHARED_LIBRARY_SUFFIX})
+  endif()
+endif()
+
+
 if (ARROW_FOUND)
   if (NOT Arrow_FIND_QUIETLY)
     message(STATUS "Found the Arrow core library: ${ARROW_LIB_PATH}")
     message(STATUS "Found the Arrow Python library: ${ARROW_PYTHON_LIB_PATH}")
+    if (PYARROW_BUILD_CUDA)
+      if (ARROW_GPU_FOUND)
+	message(STATUS "Found the Arrow GPU library: ${ARROW_GPU_LIB_PATH}")
+      else()
+	set(ARROW_ERR_MSG "Could not find the Arrow GPU library. Looked for libs")
+	set(ARROW_ERR_MSG "${ARROW_ERR_MSG} in ${ARROW_SEARCH_LIB_PATH}")
+	if (Arrow_FIND_REQUIRED)
+	  message(FATAL_ERROR "${ARROW_ERR_MSG}")
+	else (Arrow_FIND_REQUIRED)
+	  message(STATUS "${ARROW_ERR_MSG}")
+	endif()
+	set(ARROW_GPU_FOUND FALSE)
+      endif()
+    else()
+      message(STATUS "Found but not using the Arrow GPU library: ${ARROW_GPU_LIB_PATH}")
+      set(ARROW_GPU_FOUND FALSE)
+    endif()
   endif ()
 else ()
   if (NOT Arrow_FIND_QUIETLY)
@@ -134,6 +178,7 @@ else ()
     endif (Arrow_FIND_REQUIRED)
   endif ()
   set(ARROW_FOUND FALSE)
+  set(ARROW_GPU_FOUND FALSE)
 endif ()
 
 if (MSVC)
@@ -145,6 +190,9 @@ if (MSVC)
     ARROW_PYTHON_STATIC_LIB
     ARROW_PYTHON_SHARED_LIB
     ARROW_PYTHON_SHARED_IMP_LIB
+    ARROW_GPU_STATIC_LIB
+    ARROW_GPU_SHARED_LIB
+    ARROW_GPU_SHARED_IMP_LIB
   )
 else()
   mark_as_advanced(
@@ -153,5 +201,7 @@ else()
     ARROW_SHARED_LIB
     ARROW_PYTHON_STATIC_LIB
     ARROW_PYTHON_SHARED_LIB
+    ARROW_GPU_STATIC_LIB
+    ARROW_GPU_SHARED_LIB
   )
 endif()

python/CMakeLists.txt

@@ -58,6 +58,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)
@@ -275,6 +278,12 @@ if (PYARROW_BUNDLE_ARROW_CPP)
     ABI_VERSION ${ARROW_ABI_VERSION}
     SO_VERSION ${ARROW_SO_VERSION})
 
+  if (ARROW_GPU_FOUND)
+    bundle_arrow_lib(ARROW_GPU_SHARED_LIB
+      ABI_VERSION ${ARROW_ABI_VERSION}
+      SO_VERSION ${ARROW_SO_VERSION})
+  endif()
+
   # boost
   if (PYARROW_BOOST_USE_SHARED AND PYARROW_BUNDLE_BOOST)
     set(Boost_USE_STATIC_LIBS OFF)
@@ -305,6 +314,9 @@ if (PYARROW_BUNDLE_ARROW_CPP)
   if (MSVC)
     bundle_arrow_implib(ARROW_SHARED_IMP_LIB)
     bundle_arrow_implib(ARROW_PYTHON_SHARED_IMP_LIB)
+    if (ARROW_GPU_FOUND)
+      bundle_arrow_implib(ARROW_GPU_SHARED_IMP_LIB)
+    endif()
   endif()
 endif()
 
@@ -313,11 +325,19 @@ if (MSVC)
     SHARED_LIB ${ARROW_SHARED_IMP_LIB})
   ADD_THIRDPARTY_LIB(arrow_python
     SHARED_LIB ${ARROW_PYTHON_SHARED_IMP_LIB})
+  if (ARROW_GPU_FOUND)
+    ADD_THIRDPARTY_LIB(arrow_gpu
+      SHARED_LIB ${ARROW_GPU_SHARED_IMP_LIB})
+  endif()
 else()
   ADD_THIRDPARTY_LIB(arrow
     SHARED_LIB ${ARROW_SHARED_LIB})
   ADD_THIRDPARTY_LIB(arrow_python
     SHARED_LIB ${ARROW_PYTHON_SHARED_LIB})
+  if (ARROW_GPU_FOUND)
+    ADD_THIRDPARTY_LIB(arrow_gpu
+      SHARED_LIB ${ARROW_GPU_SHARED_LIB})
+  endif()
 endif()
 
 ############################################################
@@ -330,12 +350,18 @@ endif()
 
 set(CYTHON_EXTENSIONS
   lib
-)
+  )
 
 set(LINK_LIBS
   arrow_shared
   arrow_python_shared
-)
+  )
+
+
+if (ARROW_GPU_FOUND)
+  set(LINK_LIBS ${LINK_LIBS} arrow_gpu_shared)
+  set(CYTHON_EXTENSIONS ${CYTHON_EXTENSIONS} _cuda)
+endif()
 
 if (PYARROW_BUILD_PARQUET)
   ## 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