jorisvandenbossche commented on code in PR #41889:
URL: https://github.com/apache/arrow/pull/41889#discussion_r1632751914
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,57 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer(pickle_module):
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ arr = np.arange(4, dtype=np.int32)
+ cuda_buf = ctx.buffer_from_data(arr)
+
+ arr = pa.Array.from_buffers(pa.int32(), 4, [None, cuda_buf])
+ buf_on_gpu = arr.buffers()[1]
+
+ assert buf_on_gpu.size == cuda_buf.size
+ assert buf_on_gpu.address == cuda_buf.address
+ assert buf_on_gpu.is_cpu == cuda_buf.is_cpu
+
+ repr1 = "<pyarrow.Buffer address="
+ repr2 = "size=16 is_cpu=False is_mutable=True>"
+ assert repr1 in repr(buf_on_gpu)
+ assert repr2 in repr(buf_on_gpu)
+
+ msg = "Implemented only for data on CPU device"
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu.hex()
+
+ with pytest.raises(NotImplementedError, match=msg):
+ cuda_buf.hex()
+
+ assert buf_on_gpu.is_mutable
+
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu[1]
+
+ with pytest.raises(NotImplementedError, match=msg):
+ cuda_buf[1]
+
+ with pytest.raises(NotImplementedError, match=msg):
+ pickle_module.dumps(buf_on_gpu, protocol=4)
+
+ buf = pa.py_buffer(b'testing')
+ arr = pa.FixedSizeBinaryArray.from_buffers(pa.binary(7), 1, [None, buf])
+ buf_on_gpu = arr.buffers()[1]
+ buf_on_gpu_sliced = buf_on_gpu.slice(2)
+
+ buf = pa.py_buffer(b'sting')
+ arr = pa.FixedSizeBinaryArray.from_buffers(pa.binary(5), 1, [None, buf])
+ buf_on_gpu_expected = arr.buffers()[1]
+
+ assert buf_on_gpu_sliced.equals(buf_on_gpu_expected)
+ assert buf.equals(buf_on_gpu_expected)
Review Comment:
> I am not sure if I can check that the `slice` method works correctly as I
can't use the `to_pybytes()`
But to_pybytes should work on a CudaBuffer, I think?
> `equals` doesn't seem to work.
I also don't see this explicitly implemented for CudaBuffer on the C++ side
(so if that uses the base Buffer::Equals implementation, this is probably
expected to crash because it is comparing memory as if it is on the CPU). So
let's protect from calling that for now on non-CPU, but should maybe open an
issue about that whether it is worth implementing CudaBuffer::Equals.
--
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]
