jorisvandenbossche commented on code in PR #41889:
URL: https://github.com/apache/arrow/pull/41889#discussion_r1625665673
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,48 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ import numpy as np
Review Comment:
```suggestion
```
(already imported at the top)
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,48 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ import numpy as np
+ 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
+
+ msg = "Implemented only for data on CPU device"
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu.hex()
+
+ assert buf_on_gpu.is_mutable
+
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu[1]
+
+ buf = pa.py_buffer(b'testing')
+ arr_offsets = pa.array([0, 7], type=pa.int32())
+ offsets_buf = arr_offsets.buffers()[1]
Review Comment:
```suggestion
offsets_buf = pa.py_buffer(np.array([0, 7], np.int32))
```
(maybe a bit simpler)
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,48 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ import numpy as np
+ 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
+
+ msg = "Implemented only for data on CPU device"
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu.hex()
Review Comment:
Can you test this with `cuda_buf.hex()` as well?
(Because the CudaBuffer class is subclassing, it therefore has this method
available as well from the parent class, and it doesn't seem to be tested in
test_cuda.py right now)
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,48 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ import numpy as np
+ 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
+
+ msg = "Implemented only for data on CPU device"
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu.hex()
+
+ assert buf_on_gpu.is_mutable
+
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu[1]
+
+ buf = pa.py_buffer(b'testing')
+ arr_offsets = pa.array([0, 7], type=pa.int32())
+ offsets_buf = arr_offsets.buffers()[1]
Review Comment:
Or even simpler would be to use a FixedSizeBinaryArray, then you don't need
offsets to start with
##########
python/pyarrow/tests/test_io.py:
##########
@@ -669,6 +669,48 @@ def test_allocate_buffer_resizable():
assert buf.size == 200
+def test_non_cpu_buffer():
+ cuda = pytest.importorskip("pyarrow.cuda")
+ ctx = cuda.Context(0)
+
+ import numpy as np
+ 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
+
+ msg = "Implemented only for data on CPU device"
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu.hex()
+
+ assert buf_on_gpu.is_mutable
+
+ with pytest.raises(NotImplementedError, match=msg):
+ buf_on_gpu[1]
Review Comment:
Same here for `cuda_buf` (although it _might_ work in that case, as I see it
has a `getitem` defined)
--
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]
