@@ -792,3 +792,155 @@ def test_IPC(size):
792792 p .start ()
793793 p .join ()
794794 assert p .exitcode == 0
795+
796+
797+ def _arr_copy_to_host (carr ):
798+ # TODO replace below with copy to device when exposed in python
799+ buffers = []
800+ for cbuf in carr .buffers ():
801+ if cbuf is None :
802+ buffers .append (None )
803+ else :
804+ buf = global_context .foreign_buffer (
805+ cbuf .address , cbuf .size , cbuf
806+ ).copy_to_host ()
807+ buffers .append (buf )
808+
809+ child = pa .Array .from_buffers (carr .type .value_type , 3 , buffers [2 :])
810+ new = pa .Array .from_buffers (carr .type , 2 , buffers [:2 ], children = [child ])
811+ return new
812+
813+
814+ def test_device_interface_array ():
815+ cffi = pytest .importorskip ("pyarrow.cffi" )
816+ ffi = cffi .ffi
817+
818+ c_schema = ffi .new ("struct ArrowSchema*" )
819+ ptr_schema = int (ffi .cast ("uintptr_t" , c_schema ))
820+ c_array = ffi .new ("struct ArrowDeviceArray*" )
821+ ptr_array = int (ffi .cast ("uintptr_t" , c_array ))
822+
823+ typ = pa .list_ (pa .int32 ())
824+ arr = pa .array ([[1 ], [2 , 42 ]], type = typ )
825+
826+ # TODO replace below with copy to device when exposed in python
827+ cbuffers = []
828+ for buf in arr .buffers ():
829+ if buf is None :
830+ cbuffers .append (None )
831+ else :
832+ cbuf = global_context .new_buffer (buf .size )
833+ cbuf .copy_from_host (buf , position = 0 , nbytes = buf .size )
834+ cbuffers .append (cbuf )
835+
836+ carr = pa .Array .from_buffers (typ , 2 , cbuffers [:2 ], children = [
837+ pa .Array .from_buffers (typ .value_type , 3 , cbuffers [2 :])
838+ ])
839+
840+ # Type is known up front
841+ carr ._export_to_c_device (ptr_array )
842+
843+ # verify exported struct
844+ assert c_array .device_type == 2 # ARROW_DEVICE_CUDA 2
845+ assert c_array .device_id == global_context .device_number
846+ assert c_array .array .length == 2
847+
848+ # Delete recreate C++ object from exported pointer
849+ del carr
850+ carr_new = pa .Array ._import_from_c_device (ptr_array , typ )
851+ assert carr_new .type == pa .list_ (pa .int32 ())
852+ arr_new = _arr_copy_to_host (carr_new )
853+ assert arr_new .equals (arr )
854+
855+ del carr_new
856+ # Now released
857+ with pytest .raises (ValueError , match = "Cannot import released ArrowArray" ):
858+ pa .Array ._import_from_c_device (ptr_array , typ )
859+
860+ # Schema is exported and imported at the same time
861+ carr = pa .Array .from_buffers (typ , 2 , cbuffers [:2 ], children = [
862+ pa .Array .from_buffers (typ .value_type , 3 , cbuffers [2 :])
863+ ])
864+ carr ._export_to_c_device (ptr_array , ptr_schema )
865+ # Delete and recreate C++ objects from exported pointers
866+ del carr
867+ carr_new = pa .Array ._import_from_c_device (ptr_array , ptr_schema )
868+ assert carr_new .type == pa .list_ (pa .int32 ())
869+ arr_new = _arr_copy_to_host (carr_new )
870+ assert arr_new .equals (arr )
871+
872+ del carr_new
873+ # Now released
874+ with pytest .raises (ValueError , match = "Cannot import released ArrowSchema" ):
875+ pa .Array ._import_from_c_device (ptr_array , ptr_schema )
876+
877+
878+ def _batch_copy_to_host (cbatch ):
879+ # TODO replace below with copy to device when exposed in python
880+ arrs = []
881+ for col in cbatch .columns :
882+ buffers = [
883+ global_context .foreign_buffer (buf .address , buf .size , buf ).copy_to_host ()
884+ if buf is not None else None
885+ for buf in col .buffers ()
886+ ]
887+ new = pa .Array .from_buffers (col .type , len (col ), buffers )
888+ arrs .append (new )
889+
890+ return pa .RecordBatch .from_arrays (arrs , schema = cbatch .schema )
891+
892+
893+ def test_device_interface_batch_array ():
894+ cffi = pytest .importorskip ("pyarrow.cffi" )
895+ ffi = cffi .ffi
896+
897+ c_schema = ffi .new ("struct ArrowSchema*" )
898+ ptr_schema = int (ffi .cast ("uintptr_t" , c_schema ))
899+ c_array = ffi .new ("struct ArrowDeviceArray*" )
900+ ptr_array = int (ffi .cast ("uintptr_t" , c_array ))
901+
902+ batch = make_recordbatch (10 )
903+ schema = batch .schema
904+ cbuf = cuda .serialize_record_batch (batch , global_context )
905+ cbatch = cuda .read_record_batch (cbuf , schema )
906+
907+ # Schema is known up front
908+ cbatch ._export_to_c_device (ptr_array )
909+
910+ # verify exported struct
911+ assert c_array .device_type == 2 # ARROW_DEVICE_CUDA 2
912+ assert c_array .device_id == global_context .device_number
913+ assert c_array .array .length == 10
914+
915+ # Delete recreate C++ object from exported pointer
916+ del cbatch
917+ cbatch_new = pa .RecordBatch ._import_from_c_device (ptr_array , schema )
918+ assert cbatch_new .schema == schema
919+ batch_new = _batch_copy_to_host (cbatch_new )
920+ assert batch_new .equals (batch )
921+
922+ del cbatch_new
923+ # Now released
924+ with pytest .raises (ValueError , match = "Cannot import released ArrowArray" ):
925+ pa .RecordBatch ._import_from_c_device (ptr_array , schema )
926+
927+ # Schema is exported and imported at the same time
928+ cbatch = cuda .read_record_batch (cbuf , schema )
929+ cbatch ._export_to_c_device (ptr_array , ptr_schema )
930+ # Delete and recreate C++ objects from exported pointers
931+ del cbatch
932+ cbatch_new = pa .RecordBatch ._import_from_c_device (ptr_array , ptr_schema )
933+ assert cbatch_new .schema == schema
934+ batch_new = _batch_copy_to_host (cbatch_new )
935+ assert batch_new .equals (batch )
936+
937+ del cbatch_new
938+ # Now released
939+ with pytest .raises (ValueError , match = "Cannot import released ArrowSchema" ):
940+ pa .RecordBatch ._import_from_c_device (ptr_array , ptr_schema )
941+
942+ # Not a struct type
943+ pa .int32 ()._export_to_c (ptr_schema )
944+ with pytest .raises (ValueError ,
945+ match = "ArrowSchema describes non-struct type" ):
946+ pa .RecordBatch ._import_from_c_device (ptr_array , ptr_schema )
0 commit comments