fixup import from pyarrow._cuda

jorisvandenbossche · jorisvandenbossche · commit 83b0d587f62f · 2024-03-19T15:24:30.000Z
diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx
@@ -965,7 +965,8 @@ def read_record_batch(object buffer, object schema, *,
     return pyarrow_wrap_batch(batch)
 
 
-def _import_device_array(in_ptr, type):
+def _import_device_array_cuda(in_ptr, type):
+    # equivalent to the definition in array.pxi but using CudaDefaultMemoryMapper
     cdef:
         void* c_ptr = _as_c_pointer(in_ptr)
         void* c_type_ptr
@@ -978,19 +979,20 @@ def _import_device_array(in_ptr, type):
         with nogil:
             c_array = GetResultValue(
                 ImportDeviceArray(<ArrowDeviceArray*> c_ptr,
-                                    <ArrowSchema*> c_type_ptr,
-                                    CudaDefaultMemoryMapper)
+                                  <ArrowSchema*> c_type_ptr,
+                                  CudaDefaultMemoryMapper)
             )
     else:
         with nogil:
             c_array = GetResultValue(
                 ImportDeviceArray(<ArrowDeviceArray*> c_ptr, c_type,
-                                    CudaDefaultMemoryMapper)
+                                  CudaDefaultMemoryMapper)
             )
     return pyarrow_wrap_array(c_array)
 
 
-def _import_device_recordbatch(in_ptr, schema):
+def _import_device_recordbatch_cuda(in_ptr, schema):
+    # equivalent to the definition in table.pxi but using CudaDefaultMemoryMapper
     cdef:
         void* c_ptr = _as_c_pointer(in_ptr)
         void* c_schema_ptr
@@ -1003,11 +1005,13 @@ def _import_device_recordbatch(in_ptr, schema):
         with nogil:
             c_batch = GetResultValue(ImportDeviceRecordBatch(
                 <ArrowDeviceArray*> c_ptr, <ArrowSchema*> c_schema_ptr,
-                CudaDefaultMemoryMapper))
+                CudaDefaultMemoryMapper)
+            )
     else:
         with nogil:
             c_batch = GetResultValue(ImportDeviceRecordBatch(
-                <ArrowDeviceArray*> c_ptr, c_schema, CudaDefaultMemoryMapper))
+                <ArrowDeviceArray*> c_ptr, c_schema, CudaDefaultMemoryMapper)
+            )
     return pyarrow_wrap_batch(c_batch)
 
 
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
@@ -22,21 +22,7 @@ import warnings
 from cython import sizeof
 
 
-def _import_device_array(in_ptr, type):
-    """
-    Import Array from a C ArrowDeviceArray struct, given its pointer
-    and the imported array type.
-
-    Parameters
-    ----------
-    in_ptr: int
-        The raw pointer to a C ArrowDeviceArray struct.
-    type: DataType or int
-        Either a DataType object, or the raw pointer to a C ArrowSchema
-        struct.
-
-    This is a low-level function intended for expert users.
-    """
+def _import_device_array_cpu(in_ptr, type):
     cdef:
         void* c_ptr = _as_c_pointer(in_ptr)
         void* c_type_ptr
@@ -62,9 +48,11 @@ def _import_device_array(in_ptr, type):
 
 
 try:
-    from pyarrow._cuda import _import_device_array
+    from pyarrow._cuda import _import_device_array_cuda
+
+    _import_device_array = _import_device_array_cuda
 except ImportError:
-    pass
+    _import_device_array = _import_device_array_cpu
 
 
 cdef _sequence_to_array(object sequence, object mask, object size,
diff --git a/python/pyarrow/includes/libarrow_cuda.pxd b/python/pyarrow/includes/libarrow_cuda.pxd
@@ -93,8 +93,9 @@ cdef extern from "arrow/gpu/cuda_api.h" namespace "arrow::cuda" nogil:
     CResult[shared_ptr[CCudaHostBuffer]] AllocateCudaHostBuffer(
         int device_number, const int64_t size)
 
-    CResult[shared_ptr[CMemoryManager]] CudaDefaultMemoryMapper" DefaultMemoryMapper"(
-        ArrowDeviceType device_type, int64_t device_id)
+    CResult[shared_ptr[CMemoryManager]] \
+        CudaDefaultMemoryMapper" arrow::cuda::DefaultMemoryMapper"(
+            ArrowDeviceType device_type, int64_t device_id)
 
     # Cuda prefix is added to avoid picking up arrow::cuda functions
     # from arrow namespace.
diff --git a/python/pyarrow/lib.pyx b/python/pyarrow/lib.pyx
@@ -26,6 +26,8 @@ import os
 import sys
 
 from cython.operator cimport dereference as deref
+from cython cimport binding
+
 from pyarrow.includes.libarrow cimport *
 from pyarrow.includes.libarrow_python cimport *
 from pyarrow.includes.common cimport PyObject_to_object
@@ -162,6 +164,9 @@ include "pandas-shim.pxi"
 # Memory pools and allocation
 include "memory.pxi"
 
+# File IO
+include "io.pxi"
+
 # DataType, Field, Schema
 include "types.pxi"
 
@@ -183,9 +188,6 @@ include "tensor.pxi"
 # DLPack
 include "_dlpack.pxi"
 
-# File IO
-include "io.pxi"
-
 # IPC / Messaging
 include "ipc.pxi"
 
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
@@ -21,7 +21,7 @@ import warnings
 from cython import sizeof
 
 
-def _import_device_recordbatch(in_ptr, schema):
+def _import_device_recordbatch_cpu(in_ptr, schema):
     cdef:
         void* c_ptr = _as_c_pointer(in_ptr)
         void* c_schema_ptr
@@ -43,9 +43,11 @@ def _import_device_recordbatch(in_ptr, schema):
 
 
 try:
-    from pyarrow._cuda import _import_device_recordbatch
+    from pyarrow._cuda import _import_device_recordbatch_cuda
+
+    _import_device_recordbatch = _import_device_recordbatch_cuda
 except ImportError:
-    pass
+    _import_device_recordbatch = _import_device_recordbatch_cpu
 
 
 cdef class ChunkedArray(_PandasConvertible):
