jorisvandenbossche commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1725623390
##########
python/pyarrow/tests/test_table.py:
##########
@@ -3357,3 +3357,138 @@ def test_invalid_non_join_column():
with pytest.raises(pa.lib.ArrowInvalid) as excinfo:
t2.join(t1, 'id', join_type='inner')
assert exp_error_msg in str(excinfo.value)
+
+
+def test_recordbatch_non_cpu():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ def create_data_on_cpu_device():
+ return [
+ np.array(range(5), dtype=np.int16),
+ np.array([-10, -5, 0, 1, 10], dtype=np.int32)
+ ]
+
+ def create_arrays_on_cuda_device():
+ cpu_data = create_data_on_cpu_device()
+ cuda_data = [ctx.buffer_from_data(x) for x in cpu_data]
+ cuda_arrays = [
+ pa.Array.from_buffers(pa.int16(), 5, [None, cuda_data[0]]),
+ pa.Array.from_buffers(pa.int32(), 5, [None, cuda_data[1]]),
+ ]
+ return cuda_arrays
+
+ def create_recordbatch_on_cuda_device():
+ cuda_arrays = create_arrays_on_cuda_device()
+ return pa.record_batch(cuda_arrays, ['c0', 'c1'])
+
+ def verify_recordbatch_on_cuda_device(rb, column_names=['c0', 'c1']):
+ rb.validate()
+ assert rb.device_type == pa.DeviceAllocationType.CUDA
+ assert rb.is_cpu is False
+ assert rb.num_columns == 2
+ assert rb.num_rows == 5
+ assert str(rb) in repr(rb)
+ for c in rb.columns:
+ assert c.device_type == pa.DeviceAllocationType.CUDA
+ assert rb.column_names == column_names
+
+ batch = create_recordbatch_on_cuda_device()
+ verify_recordbatch_on_cuda_device(batch)
+
+ # add_column() test
+ col = pa.Array.from_buffers(
+ pa.int8(), 5, [None, ctx.buffer_from_data(np.array([6, 7, 8, 9, 10],
dtype=np.int8))])
+ new_batch = batch.add_column(2, 'c2', col)
+ assert len(new_batch.columns) == 3
+ for c in new_batch.columns:
+ assert c.device_type == pa.DeviceAllocationType.CUDA
+ with pytest.raises(NotImplementedError):
+ # Default array conversion builds on CPU only
+ batch.add_column(2, 'c2', [1, 1, 1, 1, 1])
Review Comment:
I don't think that is too big of an assumption, but I was wondering where
the error is coming from right now (and why it is a NotImplementedError)
--
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]
