From 7890c001cfc30eb8c08d2ee76f07a81130b93bd5 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Thu, 6 Jun 2024 14:04:41 +0200 Subject: [PATCH 1/8] GH-41664: PrettyPrint non-cpu data by copying slice to default CPU device --- cpp/src/arrow/pretty_print.cc | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index e666ec70f94..0c66f6e8f54 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -31,6 +31,7 @@ #include #include "arrow/array.h" +#include "arrow/array/concatenate.h" #include "arrow/chunked_array.h" #include "arrow/record_batch.h" #include "arrow/status.h" @@ -413,18 +414,37 @@ Status ArrayPrinter::WriteValidityBitmap(const Array& array) { } } +Result> CopyStartEndToCPU(const Array& arr, int window) { + std::shared_ptr arr_sliced; + if (arr.length() > (2 * window + 1)) { + ARROW_ASSIGN_OR_RAISE( + auto arr_start, + arr.Slice(0, window + 1)->CopyTo(default_cpu_memory_manager())); + ARROW_ASSIGN_OR_RAISE( + auto arr_end, + arr.Slice(arr.length() - window - 1)->CopyTo(default_cpu_memory_manager())); + ARROW_ASSIGN_OR_RAISE(arr_sliced, Concatenate({arr_start, arr_end})); + } else { + ARROW_ASSIGN_OR_RAISE(arr_sliced, arr.CopyTo(default_cpu_memory_manager())); + } + return arr_sliced; +} + } // namespace Status PrettyPrint(const Array& arr, int indent, std::ostream* sink) { PrettyPrintOptions options; options.indent = indent; - ArrayPrinter printer(options, sink); - return printer.Print(arr); + return PrettyPrint(arr, options, sink); } Status PrettyPrint(const Array& arr, const PrettyPrintOptions& options, std::ostream* sink) { ArrayPrinter printer(options, sink); + if (arr.device_type() != DeviceAllocationType::kCPU) { + ARROW_ASSIGN_OR_RAISE(auto arr_sliced, CopyStartEndToCPU(arr, options.window)); + return printer.Print(*arr_sliced); + } return printer.Print(arr); } From 050fd9d7d89e4a93c1cb5f2ae36aa22f288cb947 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 12 Jun 2024 14:27:12 +0000 Subject: [PATCH 2/8] fix for chunked array + add python tests --- cpp/src/arrow/pretty_print.cc | 3 +- python/pyarrow/tests/test_cuda.py | 54 +++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index 0c66f6e8f54..bdcf89c37f5 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -496,8 +496,7 @@ Status PrettyPrint(const ChunkedArray& chunked_arr, const PrettyPrintOptions& op } else { PrettyPrintOptions chunk_options = options; chunk_options.indent += options.indent_size; - ArrayPrinter printer(chunk_options, sink); - RETURN_NOT_OK(printer.Print(*chunked_arr.chunk(i))); + RETURN_NOT_OK(PrettyPrint(*chunked_arr.chunk(i), chunk_options, sink)); } } if (!options.skip_new_lines) { diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 43cd16a3cf6..cb068e02c70 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -792,3 +792,57 @@ def test_IPC(size): p.start() p.join() assert p.exitcode == 0 + + +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) From 1c09d52ab57764102f204df1db4c110b29180639 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 12 Jun 2024 16:33:05 +0200 Subject: [PATCH 3/8] formatting --- cpp/src/arrow/pretty_print.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index bdcf89c37f5..bdc0082b3d5 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -417,9 +417,8 @@ Status ArrayPrinter::WriteValidityBitmap(const Array& array) { Result> CopyStartEndToCPU(const Array& arr, int window) { std::shared_ptr arr_sliced; if (arr.length() > (2 * window + 1)) { - ARROW_ASSIGN_OR_RAISE( - auto arr_start, - arr.Slice(0, window + 1)->CopyTo(default_cpu_memory_manager())); + ARROW_ASSIGN_OR_RAISE(auto arr_start, + arr.Slice(0, window + 1)->CopyTo(default_cpu_memory_manager())); ARROW_ASSIGN_OR_RAISE( auto arr_end, arr.Slice(arr.length() - window - 1)->CopyTo(default_cpu_memory_manager())); From 595ad8982ae19d09a7ee0d2eabbd1dbce28ca7c8 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Thu, 20 Jun 2024 18:55:47 +0200 Subject: [PATCH 4/8] simplify approach (just copy full array instead of slice+concat) --- cpp/src/arrow/pretty_print.cc | 20 ++------------------ 1 file changed, 2 insertions(+), 18 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index bdc0082b3d5..7a8fa081707 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -31,7 +31,6 @@ #include #include "arrow/array.h" -#include "arrow/array/concatenate.h" #include "arrow/chunked_array.h" #include "arrow/record_batch.h" #include "arrow/status.h" @@ -414,21 +413,6 @@ Status ArrayPrinter::WriteValidityBitmap(const Array& array) { } } -Result> CopyStartEndToCPU(const Array& arr, int window) { - std::shared_ptr arr_sliced; - if (arr.length() > (2 * window + 1)) { - ARROW_ASSIGN_OR_RAISE(auto arr_start, - arr.Slice(0, window + 1)->CopyTo(default_cpu_memory_manager())); - ARROW_ASSIGN_OR_RAISE( - auto arr_end, - arr.Slice(arr.length() - window - 1)->CopyTo(default_cpu_memory_manager())); - ARROW_ASSIGN_OR_RAISE(arr_sliced, Concatenate({arr_start, arr_end})); - } else { - ARROW_ASSIGN_OR_RAISE(arr_sliced, arr.CopyTo(default_cpu_memory_manager())); - } - return arr_sliced; -} - } // namespace Status PrettyPrint(const Array& arr, int indent, std::ostream* sink) { @@ -441,8 +425,8 @@ Status PrettyPrint(const Array& arr, const PrettyPrintOptions& options, std::ostream* sink) { ArrayPrinter printer(options, sink); if (arr.device_type() != DeviceAllocationType::kCPU) { - ARROW_ASSIGN_OR_RAISE(auto arr_sliced, CopyStartEndToCPU(arr, options.window)); - return printer.Print(*arr_sliced); + ARROW_ASSIGN_OR_RAISE(auto arr_copied, arr.CopyTo(default_cpu_memory_manager())); + return printer.Print(*arr_copied); } return printer.Print(arr); } From 555f8b6b5fa76ce431155bbb67bd89c4e33c8663 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Tue, 25 Jun 2024 18:34:30 +0200 Subject: [PATCH 5/8] move device check into ArrayPrinter::Print --- cpp/src/arrow/pretty_print.cc | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index 7a8fa081707..a90da466f5b 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -391,7 +391,12 @@ class ArrayPrinter : public PrettyPrinter { } Status Print(const Array& array) { - RETURN_NOT_OK(VisitArrayInline(array, this)); + if (array.device_type() != DeviceAllocationType::kCPU) { + 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(); } @@ -418,16 +423,13 @@ Status ArrayPrinter::WriteValidityBitmap(const Array& array) { Status PrettyPrint(const Array& arr, int indent, std::ostream* sink) { PrettyPrintOptions options; options.indent = indent; - return PrettyPrint(arr, options, sink); + ArrayPrinter printer(options, sink); + return printer.Print(arr); } Status PrettyPrint(const Array& arr, const PrettyPrintOptions& options, std::ostream* sink) { ArrayPrinter printer(options, sink); - if (arr.device_type() != DeviceAllocationType::kCPU) { - ARROW_ASSIGN_OR_RAISE(auto arr_copied, arr.CopyTo(default_cpu_memory_manager())); - return printer.Print(*arr_copied); - } return printer.Print(arr); } @@ -479,7 +481,8 @@ Status PrettyPrint(const ChunkedArray& chunked_arr, const PrettyPrintOptions& op } else { PrettyPrintOptions chunk_options = options; chunk_options.indent += options.indent_size; - RETURN_NOT_OK(PrettyPrint(*chunked_arr.chunk(i), chunk_options, sink)); + ArrayPrinter printer(chunk_options, sink); + RETURN_NOT_OK(printer.Print(*chunked_arr.chunk(i))); } } if (!options.skip_new_lines) { From 1b64ccf1d491389032b8d62d0c743b7949ed8ab8 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 26 Jun 2024 11:07:53 +0200 Subject: [PATCH 6/8] remove assert_cpu now repr work --- python/pyarrow/array.pxi | 5 +++-- python/pyarrow/tests/test_array.py | 6 ++---- 2 files changed, 5 insertions(+), 6 deletions(-) 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): From fe82b3e4ec98d6d1ff2fa263fa32197c075a1fc4 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 26 Jun 2024 11:25:31 +0200 Subject: [PATCH 7/8] add TODO comment --- cpp/src/arrow/pretty_print.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index a90da466f5b..33dc23d8d0d 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -392,6 +392,8 @@ class ArrayPrinter : public PrettyPrinter { Status Print(const Array& array) { if (array.device_type() != DeviceAllocationType::kCPU) { + // TODO: ideally we only copy start/end slice based on the window size that is being + // printed (requires CopyTo for slices https://github.com/apache/arrow/issues/43055) ARROW_ASSIGN_OR_RAISE(auto array_cpu, array.CopyTo(default_cpu_memory_manager())); RETURN_NOT_OK(VisitArrayInline(*array_cpu, this)); } else { From c65c7a56808d781de6f66921cb38a57b51361262 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 26 Jun 2024 16:41:42 +0200 Subject: [PATCH 8/8] Update cpp/src/arrow/pretty_print.cc Co-authored-by: Felipe Oliveira Carvalho --- cpp/src/arrow/pretty_print.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index 33dc23d8d0d..53a69536816 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -392,8 +392,8 @@ class ArrayPrinter : public PrettyPrinter { Status Print(const Array& array) { if (array.device_type() != DeviceAllocationType::kCPU) { - // TODO: ideally we only copy start/end slice based on the window size that is being - // printed (requires CopyTo for slices https://github.com/apache/arrow/issues/43055) + // 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 {