danepitkin commented on code in PR #43974:
URL: https://github.com/apache/arrow/pull/43974#discussion_r1752524817
##########
python/pyarrow/tests/test_table.py:
##########
@@ -3586,31 +3601,32 @@ def verify_cuda_recordbatch(batch, expected_schema):
def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch,
cuda_arrays, schema):
verify_cuda_recordbatch(cuda_recordbatch, expected_schema=schema)
+
+ # shape test
assert cuda_recordbatch.shape == (5, 2)
# columns() test
assert len(cuda_recordbatch.columns) == 2
# add_column(), set_column() test
for fn in [cuda_recordbatch.add_column, cuda_recordbatch.set_column]:
- col = pa.array([6, 7, 8, 9, 10],
pa.int8()).copy_to(cuda_context.memory_manager)
+ col = pa.array([1] * cuda_recordbatch.num_rows, pa.int8()
+ ).copy_to(cuda_context.memory_manager)
new_batch = fn(2, 'c2', col)
- assert len(new_batch.columns) == 3
- for c in new_batch.columns:
- assert c.device_type == pa.DeviceAllocationType.CUDA
+ verify_cuda_recordbatch(
+ new_batch, expected_schema=schema.append(pa.field('c2',
pa.int8())))
err_msg = ("Got column on device <DeviceAllocationType.CPU: 1>, "
"but expected <DeviceAllocationType.CUDA: 2>.")
with pytest.raises(TypeError, match=err_msg):
- fn(2, 'c2', [1, 1, 1, 1, 1])
+ fn(2, 'c2', [1] * cuda_recordbatch.num_rows)
# remove_column() test
new_batch = cuda_recordbatch.remove_column(1)
verify_cuda_recordbatch(new_batch, expected_schema=schema.remove(1))
# drop_columns() test
- new_batch = cuda_recordbatch.drop_columns(['c0', 'c1'])
- assert len(new_batch.columns) == 0
- assert new_batch.device_type == pa.DeviceAllocationType.CUDA
+ new_batch = cuda_recordbatch.drop_columns(['c1'])
+ verify_cuda_recordbatch(new_batch, expected_schema=schema.remove(1))
Review Comment:
Turns out, the empty batch will still say CUDA device.
--
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]
