danepitkin commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1742503750
##########
python/pyarrow/table.pxi:
##########
@@ -2800,17 +2801,26 @@ cdef class RecordBatch(_Tabular):
shared_ptr[CRecordBatch] c_batch
Field c_field
Array c_arr
+ CDeviceAllocationType device_type =
self.sp_batch.get().device_type()
if isinstance(column, Array):
c_arr = column
else:
- if not self.is_cpu:
- raise RuntimeError("A pa.Array() object is required when "
- "adding columns to a RecordBatch on a "
- f"non-cpu device. Got {type(column)!r} "
- "instead.")
+ if device_type != CDeviceAllocationType_kCPU:
+ cpu_device_type = _wrap_device_allocation_type(
+ CDeviceAllocationType_kCPU)
+ raise TypeError("The column must be allocated on the same "
+ "device as the RecordBatch. Got column on "
+ f"device {cpu_device_type!r}, but expected "
+ f"{self.device_type!r}.")
Review Comment:
I like this better! Will update
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
