danepitkin commented on code in PR #43729:
URL: https://github.com/apache/arrow/pull/43729#discussion_r1733433462
##########
python/pyarrow/tests/test_table.py:
##########
@@ -3357,3 +3357,167 @@ 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)
+
+
[email protected]
+def cuda_context():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ return cuda.Context(0)
+
+
[email protected]
+def schema():
+ return pa.schema([pa.field('c0', pa.int16()), pa.field('c1', pa.int32())])
+
+
[email protected]
+def cpu_arrays():
+ return [pa.array([1, 2, 3, 4, 5], pa.int16()),
+ pa.array([-10, -5, 0, 1, 10], pa.int32())]
+
+
[email protected]
+def cuda_arrays(cuda_context, cpu_arrays):
+ return [arr.copy_to(cuda_context.memory_manager) for arr in cpu_arrays]
+
+
[email protected]
+def cpu_recordbatch(cpu_arrays, schema):
+ return pa.record_batch(cpu_arrays, schema=schema)
+
+
[email protected]
+def cuda_recordbatch(cuda_context, cpu_recordbatch):
+ return cpu_recordbatch.copy_to(cuda_context.memory_manager)
+
+
+def verify_recordbatch_on_cuda_device(batch, expected_schema):
+ batch.validate()
+ assert batch.device_type == pa.DeviceAllocationType.CUDA
+ assert batch.is_cpu is False
+ assert batch.num_columns == len(batch.column_names)
+ assert str(batch) in repr(batch)
+ for c in batch.columns:
+ assert c.device_type == pa.DeviceAllocationType.CUDA
+ assert batch.schema == expected_schema
+
+
+def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch,
+ cuda_arrays, schema):
+ verify_recordbatch_on_cuda_device(cuda_recordbatch, expected_schema=schema)
+ assert cuda_recordbatch.num_rows == 5
+
+ # add_column() test
+ col = pa.Array.from_buffers(
+ pa.int8(), 5,
+ [None, cuda_context.buffer_from_data(np.array([6, 7, 8, 9, 10],
+ dtype=np.int8))])
+ new_batch = cuda_recordbatch.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(RuntimeError):
+ # Default array conversion not allowed for non-cpu record batch
+ cuda_recordbatch.add_column(2, 'c2', [1, 1, 1, 1, 1])
Review Comment:
It segfaults as well. Great catch! I'll add a check for this.
--
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]
