apacheGH-43727: [Python] RecordBatch failse gracefully on non-cpu dev…

…ices
danepitkin · Aug 16, 2024 · 561239e · 561239e
1 parent d801dae
commit 561239e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/python/pyarrow/lib.pxd b/python/pyarrow/lib.pxd
@@ -519,6 +519,7 @@ cdef class Table(_Tabular):
         CTable* table
 
     cdef void init(self, const shared_ptr[CTable]& table)
+    cdef void _assert_cpu(self) noexcept
 
 
 cdef class RecordBatch(_Tabular):
@@ -528,6 +529,7 @@ cdef class RecordBatch(_Tabular):
         Schema _schema
 
     cdef void init(self, const shared_ptr[CRecordBatch]& table)
+    cdef void _assert_cpu(self) except *
 
 
 cdef class Device(_Weakrefable):

diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
@@ -3905,6 +3905,10 @@ cdef class RecordBatch(_Tabular):
         """
         return self.device_type == DeviceAllocationType.CPU
 
+    cdef void _assert_cpu(self) except *:
+        if self.sp_batch.get().device_type() != CDeviceAllocationType_kCPU:
+            raise NotImplementedError("Implemented only for data on CPU device")
+
 
 def _reconstruct_record_batch(columns, schema):
     """
@@ -5679,6 +5683,9 @@ cdef class Table(_Tabular):
         """
         return self.to_reader().__arrow_c_stream__(requested_schema)
 
+    cdef void _assert_cpu(self) noexcept:
+        return  # TODO GH-43728
+
 
 def _reconstruct_table(arrays, schema):
     """