jorisvandenbossche commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1741893868
##########
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 am not entirely following here. It seems you are hardcoding in the message
that the received device type is CPU, but that doesn't follow from `device_type
!= CDeviceAllocationType_kCPU` (that only asserts the device type from the
calling record batch, not from the passed array?
--
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]
