jorisvandenbossche commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1741920065
##########
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:
Alternatively, we could also not do that check here, and let the values be
coerced to a pyarrow.Array with the line just below. In case you are eg passing
a list of numbers, that will be converted to a CPU pyarrow.Array, and then we
will still get the error below from checking the device.
It means we would do this conversion "needlessly" (but it also only happens
when you get an error), but it would simplify the code quite a bit because this
whole `if` block can be removed.
That still gives the same error message as what I was complaining about
above, but that is maybe worth the simpler code.
(it would in practice also enable that someone could pass a non-pyarrow
object that has a `__arrow_c_device_array__` method)
--
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]
