apacheGH-43507: [C++] Use ViewOrCopyTo instead of CopyTo when pretty …

…printing non-CPU data
jorisvandenbossche · Aug 1, 2024 · 1760a86 · 1760a86
1 parent d4d92e4
commit 1760a86
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc
@@ -394,7 +394,8 @@ class ArrayPrinter : public PrettyPrinter {
     if (array.device_type() != DeviceAllocationType::kCPU) {
       // GH-43055: ideally we only copy start/end slices from non-CPU memory
       // based on the window size that is being printed
-      ARROW_ASSIGN_OR_RAISE(auto array_cpu, array.CopyTo(default_cpu_memory_manager()));
+      ARROW_ASSIGN_OR_RAISE(auto array_cpu,
+                            array.ViewOrCopyTo(default_cpu_memory_manager()));
       RETURN_NOT_OK(VisitArrayInline(*array_cpu, this));
     } else {
       RETURN_NOT_OK(VisitArrayInline(array, this));

diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py
@@ -973,6 +973,17 @@ def test_print_array():
     assert str(carr) == str(arr)
 
 
+@pytest.mark.parametrize("size", [10, 100])
+def test_print_array_host(size):
+    buf = cuda.new_host_buffer(size*8)
+    np_arr = np.frombuffer(buf, dtype=np.int64)
+    np_arr[:] = range(size)
+
+    arr = pa.array(range(size), pa.int64())
+    carr = pa.Array.from_buffers(pa.int64(), size, [None, buf])
+    assert str(carr) == str(arr)
+
+
 def make_chunked_array(n_elements_per_chunk, n_chunks):
     arrs = []
     carrs = []