jorisvandenbossche commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1741912427
##########
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:
Ah, but this is in the case that you are passing something that is _not_ an
pyarrow array (and therefore we assume it is always on the CPU?)
Maybe we can clarify that in the error message? I think it would be more
useful to get a message saying something like "when adding a column to a
non-CPU batch, can only do that by passing a pyarrow.Array (of the same
device), but got {type(column)}"
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
