diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index e666ec70f94..53a69536816 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -391,7 +391,14 @@ class ArrayPrinter : public PrettyPrinter { } Status Print(const Array& array) { - RETURN_NOT_OK(VisitArrayInline(array, this)); + 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())); + RETURN_NOT_OK(VisitArrayInline(*array_cpu, this)); + } else { + RETURN_NOT_OK(VisitArrayInline(array, this)); + } Flush(); return Status::OK(); } diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index a9af1419ae4..daf4adc33e5 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -1241,6 +1241,9 @@ cdef class Array(_PandasConvertible): """ Render a "pretty-printed" string representation of the Array. + Note: for data on a non-CPU device, the full array is copied to CPU + memory. + Parameters ---------- indent : int, default 2 @@ -1261,8 +1264,6 @@ cdef class Array(_PandasConvertible): If the array should be rendered as a single line of text or if each element should be on its own line. """ - self._assert_cpu() - cdef: c_string result PrettyPrintOptions options diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 62a3eebb8c0..27e78e3396e 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -4035,6 +4035,8 @@ def test_non_cpu_array(): assert arr.is_cpu is False assert len(arr) == 4 assert arr.slice(2, 2).offset == 2 + assert repr(arr) + assert str(arr) # TODO support DLPack for CUDA with pytest.raises(NotImplementedError): @@ -4065,10 +4067,6 @@ def test_non_cpu_array(): arr.get_total_buffer_size() with pytest.raises(NotImplementedError): [i for i in iter(arr)] - with pytest.raises(NotImplementedError): - repr(arr) - with pytest.raises(NotImplementedError): - str(arr) with pytest.raises(NotImplementedError): arr == arr2 with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 400db2643bd..5e833d5371d 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -944,3 +944,57 @@ def test_device_interface_batch_array(): with pytest.raises(ValueError, match="ArrowSchema describes non-struct type"): pa.RecordBatch._import_from_c_device(ptr_array, ptr_schema) + + +def test_print_array(): + batch = make_recordbatch(10) + cbuf = cuda.serialize_record_batch(batch, global_context) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + arr = batch["f0"] + carr = cbatch["f0"] + assert str(carr) == str(arr) + + batch = make_recordbatch(100) + cbuf = cuda.serialize_record_batch(batch, global_context) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + arr = batch["f0"] + carr = cbatch["f0"] + assert str(carr) == str(arr) + + +def make_chunked_array(n_elements_per_chunk, n_chunks): + arrs = [] + carrs = [] + for _ in range(n_chunks): + batch = make_recordbatch(n_elements_per_chunk) + cbuf = cuda.serialize_record_batch(batch, global_context) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + arrs.append(batch["f0"]) + carrs.append(cbatch["f0"]) + + return pa.chunked_array(arrs), pa.chunked_array(carrs) + + +def test_print_chunked_array(): + arr, carr = make_chunked_array(10, 3) + assert str(carr) == str(arr) + + arr, carr = make_chunked_array(100, 20) + assert str(carr) == str(arr) + + +def test_print_record_batch(): + batch = make_recordbatch(10) + cbuf = cuda.serialize_record_batch(batch, global_context) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + assert str(cbatch) == str(batch) + + batch = make_recordbatch(100) + cbuf = cuda.serialize_record_batch(batch, global_context) + cbatch = cuda.read_record_batch(cbuf, batch.schema) + assert str(cbatch) == str(batch) + + +def test_print_table(): + _, table, _, ctable = make_table_cuda() + assert str(ctable) == str(table)