GH-39771: [C++][Device] Generic CopyBatchTo/CopyArrayTo memory types#39772
GH-39771: [C++][Device] Generic CopyBatchTo/CopyArrayTo memory types#39772zeroshade merged 8 commits intoapache:mainfrom
Conversation
|
|
|
Wow. This is what I wanted when I wrote #39164. diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc
index bd2c2b716d..a462dd2beb 100644
--- a/cpp/src/arrow/ipc/read_write_test.cc
+++ b/cpp/src/arrow/ipc/read_write_test.cc
@@ -1336,31 +1336,11 @@ class CopyCollectListener : public CollectListener {
Status OnRecordBatchWithMetadataDecoded(
RecordBatchWithMetadata record_batch_with_metadata) override {
- auto& record_batch = record_batch_with_metadata.batch;
- for (auto column_data : record_batch->column_data()) {
- ARROW_RETURN_NOT_OK(CopyArrayData(column_data));
- }
+ ARROW_ASSIGN_OR_RAISE(record_batch_with_metadata.batch,
+ CopyBatchTo(record_batch_with_metadata.batch,
+ default_cpu_memory_manager()));
return CollectListener::OnRecordBatchWithMetadataDecoded(record_batch_with_metadata);
}
-
- private:
- Status CopyArrayData(std::shared_ptr<ArrayData> data) {
- auto& buffers = data->buffers;
- for (size_t i = 0; i < buffers.size(); ++i) {
- auto& buffer = buffers[i];
- if (!buffer) {
- continue;
- }
- ARROW_ASSIGN_OR_RAISE(buffers[i], Buffer::Copy(buffer, buffer->memory_manager()));
- }
- for (auto child_data : data->child_data) {
- ARROW_RETURN_NOT_OK(CopyArrayData(child_data));
- }
- if (data->dictionary) {
- ARROW_RETURN_NOT_OK(CopyArrayData(data->dictionary));
- }
- return Status::OK();
- }
};
struct StreamDecoderWriterHelper : public StreamWriterHelper { |
cpp/src/arrow/device.h
Outdated
There was a problem hiding this comment.
Nit, but I think this could simply be a member of Array.
Also, do we want a similar CopyOrView? If the array, or part thereof, is already located on the right device, we probably don't want to incur any pointless memory copy.
There was a problem hiding this comment.
I can make those changes, sure
There was a problem hiding this comment.
CopyOrView creates weird expectations about the returned values: if you copy, the caller can have unique ownership of the array, if it's a view, you have to share. Another possible solution to the problem @pitrou brought up is failing the operation if source and destination are the same device, forcing the caller to be aware of data location before it calls the copy function. The compromise solution: provide CopyOrView in terms of a CopyArrayTo that returns a bad status when called for arrays on the same device.
(I understand that the memory model in Arrow is that every buffer is immutable and always shared, but we might want these rules to be relaxed at some point.)
There was a problem hiding this comment.
This might be a bad idea without a guarantee that all buffers making up an array or RecordBatch are on the same device. :/
There was a problem hiding this comment.
If the caller requires unique ownership of the array, then they should just call CopyTo in the first place.
We don't want to fail if the source and destination are the same device because maybe they are just intentionally doing a full copy to get unique ownership, but also because the device is tracked on a per-buffer basis, not at the array level. Having a CopyOrView allows the user to attempt to avoid copies if they don't need unique ownership of the buffers, so i don't think it's weird expectations about the returned values. It's very clear what the expectations are:
- If you need unique ownership, call
Copy - If you don't need unique ownership and want to avoid copies if possible then use
ViewOrCopy(for example if you're using Cuda Host memory you can get a CPU accessible view without having to manually perform a copy)
This might be a bad idea without a guarantee that all buffers making up an array or RecordBatch are on the same device. :/
One way to ensure they are all on the same device is to call this function! 😄
There was a problem hiding this comment.
You're right. I'm thinking about ownership in a weird way (at least relative to how Arrow works).
c45c73d to
4e698d2
Compare
|
@pitrou I've added all your suggestions in the latest changes |
|
After merging your PR, Conbench analyzed the 6 benchmarking runs that have been run so far on merge-commit 87b515e. There were no benchmark performance regressions. 🎉 The full Conbench report has more details. It also includes information about 5 possible false positives for unstable benchmarks that are known to sometimes produce them. |
…types (apache#39772) ### Rationale for this change Right now our MemoryManager interfaces operate solely at the buffer level and we do not provide any higher level facilities to copy an entire array or record batch between memory types. We should implement CopyArrayTo and CopyBatchTo functions which recursively utilize the buffer level copying to create a new Array/RecordBatch whose buffers have been copied to the destination memory manager. ### What changes are included in this PR? Exposing a `CopyArrayTo` and `CopyBatchTo` function for copying entire Array or RecordBatches between memory types. ### Are these changes tested? Tests are still being written but will be added. * Closes: apache#39771 Authored-by: Matt Topol <zotthewizard@gmail.com> Signed-off-by: Matt Topol <zotthewizard@gmail.com>
Rationale for this change
Right now our MemoryManager interfaces operate solely at the buffer level and we do not provide any higher level facilities to copy an entire array or record batch between memory types. We should implement CopyArrayTo and CopyBatchTo functions which recursively utilize the buffer level copying to create a new Array/RecordBatch whose buffers have been copied to the destination memory manager.
What changes are included in this PR?
Exposing a
CopyArrayToandCopyBatchTofunction for copying entire Array or RecordBatches between memory types.Are these changes tested?
Tests are still being written but will be added.