| 681 | |
| 682 | @pytest.mark.numpy |
| 683 | def test_non_cpu_buffer(pickle_module): |
| 684 | cuda = pytest.importorskip("pyarrow.cuda") |
| 685 | ctx = cuda.Context(0) |
| 686 | |
| 687 | data = np.array([b'testing']) |
| 688 | cuda_buf = ctx.buffer_from_data(data) |
| 689 | arr = pa.FixedSizeBinaryArray.from_buffers(pa.binary(7), 1, [None, cuda_buf]) |
| 690 | buf_on_gpu = arr.buffers()[1] |
| 691 | |
| 692 | assert buf_on_gpu.size == cuda_buf.size |
| 693 | assert buf_on_gpu.address == cuda_buf.address |
| 694 | assert buf_on_gpu.is_cpu == cuda_buf.is_cpu |
| 695 | assert buf_on_gpu.is_mutable |
| 696 | |
| 697 | repr1 = "<pyarrow.Buffer address=" |
| 698 | repr2 = "size=7 is_cpu=False is_mutable=True>" |
| 699 | assert repr1 in repr(buf_on_gpu) |
| 700 | assert repr2 in repr(buf_on_gpu) |
| 701 | |
| 702 | buf_on_gpu_sliced = buf_on_gpu.slice(2) |
| 703 | cuda_sliced = cuda.CudaBuffer.from_buffer(buf_on_gpu_sliced) |
| 704 | assert cuda_sliced.to_pybytes() == b'sting' |
| 705 | |
| 706 | buf_on_gpu_sliced = buf_on_gpu[2:4] |
| 707 | cuda_sliced = cuda.CudaBuffer.from_buffer(buf_on_gpu_sliced) |
| 708 | assert cuda_sliced.to_pybytes() == b'st' |
| 709 | |
| 710 | # Sliced buffers with same address |
| 711 | assert buf_on_gpu_sliced.equals(cuda_buf[2:4]) |
| 712 | |
| 713 | # Buffers on different devices |
| 714 | msg_device = "Device on which the data resides differs between buffers" |
| 715 | with pytest.raises(ValueError, match=msg_device): |
| 716 | buf_on_gpu.equals(pa.py_buffer(data)) |
| 717 | |
| 718 | msg = "Implemented only for data on CPU device" |
| 719 | # Buffers with different addresses |
| 720 | arr_short = np.array([b'sting']) |
| 721 | cuda_buf_short = ctx.buffer_from_data(arr_short) |
| 722 | with pytest.raises(NotImplementedError, match=msg): |
| 723 | buf_on_gpu_sliced.equals(cuda_buf_short) |
| 724 | arr_short = pa.FixedSizeBinaryArray.from_buffers( |
| 725 | pa.binary(5), 1, [None, cuda_buf_short] |
| 726 | ) |
| 727 | buf_on_gpu_short = arr_short.buffers()[1] |
| 728 | with pytest.raises(NotImplementedError, match=msg): |
| 729 | buf_on_gpu_sliced.equals(buf_on_gpu_short) |
| 730 | |
| 731 | with pytest.raises(NotImplementedError, match=msg): |
| 732 | buf_on_gpu.hex() |
| 733 | |
| 734 | with pytest.raises(NotImplementedError, match=msg): |
| 735 | cuda_buf.hex() |
| 736 | |
| 737 | with pytest.raises(NotImplementedError, match=msg): |
| 738 | buf_on_gpu[1] |
| 739 | |
| 740 | with pytest.raises(NotImplementedError, match=msg): |