You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "zeroshade (via GitHub)" <gi...@apache.org> on 2024/03/26 22:31:55 UTC
[PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
zeroshade opened a new pull request, #40807:
URL: https://github.com/apache/arrow/pull/40807
<!--
Thanks for opening a pull request!
If this is your first pull request you can find detailed information on how
to contribute here:
* [New Contributor's Guide](https://arrow.apache.org/docs/dev/developers/guide/step_by_step/pr_lifecycle.html#reviews-and-merge-of-the-pull-request)
* [Contributing Overview](https://arrow.apache.org/docs/dev/developers/overview.html)
If this is not a [minor PR](https://github.com/apache/arrow/blob/main/CONTRIBUTING.md#Minor-Fixes). Could you open an issue for this pull request on GitHub? https://github.com/apache/arrow/issues/new/choose
Opening GitHub issues ahead of time contributes to the [Openness](http://theapacheway.com/open/#:~:text=Openness%20allows%20new%20users%20the,must%20happen%20in%20the%20open.) of the Apache Arrow project.
Then could you also rename the pull request title in the following format?
GH-${GITHUB_ISSUE_ID}: [${COMPONENT}] ${SUMMARY}
or
MINOR: [${COMPONENT}] ${SUMMARY}
In the case of PARQUET issues on JIRA the title also supports:
PARQUET-${JIRA_ISSUE_ID}: [${COMPONENT}] ${SUMMARY}
-->
### Rationale for this change
The original PRs for adding support for importing and exporting the new C Device interface (#36488 / #36489) only added support for the Arrays themselves, not for the stream structure. We should support both.
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
### What changes are included in this PR?
Adding parallel functions for Import/Export of streams that accept `ArrowDeviceArrayStream`.
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
### Are these changes tested?
Test writing in progress, wanted to get this up for review while I write tests.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
### Are there any user-facing changes?
No, only new functions have been added.
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please uncomment the line below and explain which changes are breaking.
-->
<!-- **This PR includes breaking changes to public APIs.** -->
<!--
Please uncomment the line below (and provide explanation) if the changes fix either (a) a security vulnerability, (b) a bug that caused incorrect or invalid data to be produced, or (c) a bug that causes a crash (even when the API contract is upheld). We use this to highlight fixes to issues that may affect users without their knowledge. For this reason, fixing bugs that cause errors don't count, since those are usually obvious.
-->
<!-- **This PR contains a "Critical Fix".** -->
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2036265239
Needs some tests?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1560099063
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
Review Comment:
Okay, I've propagated the event and device_type throughout the record batch and the `Make` functions (defaulting it to nullptr) which should allow us to ensure this is all correct now. I'll update the documentation comments accordingly
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1611654558
##########
cpp/src/arrow/record_batch.h:
##########
@@ -23,6 +23,7 @@
#include <vector>
#include "arrow/compare.h"
+#include "arrow/device.h"
Review Comment:
I think at least moving `DeviceAllocationType` to `type_fwd.h` would make sense.
As for `Device::SyncEvent`, that's indeed a problem with nested classes (and a good reason to avoid them :-)). That could be fixed by moving `SyncEvent` out of `Device` and adding a compatibility alias.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1544855424
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
Review Comment:
Yea, at the moment I don't have an implementation there. I put a note on the `ImportDeviceRecordBatchReader` documentation if you look at bridge.h:
```
/// We are not yet bubbling the sync events from the buffers up to
/// the `GetSyncEvent` method of an imported RecordBatch. This will be added in a future
/// update.
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1550652287
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -84,6 +92,14 @@ inline void ArrowArrayMove(struct ArrowArray* src, struct ArrowArray* dest) {
ArrowArrayMarkReleased(src);
}
+inline void ArrowDeviceArrayMove(struct ArrowDeviceArray* src,
+ struct ArrowDeviceArray* dest) {
+ assert(dest != src);
+ assert(!ArrowDeviceArrayIsReleased(src));
+ memcpy(dest, src, sizeof(struct ArrowDeviceArray));
Review Comment:
Why is `ArrowDeviceArrayMove` in the opposite direction of `memcpy`, `memmove`...?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2052,7 +2065,19 @@ Status ExportStreamNext(const std::shared_ptr<ChunkedArray>& src, int64_t i,
}
}
-template <typename T>
+Status ExportDeviceStreamNext(const std::shared_ptr<ChunkedArray>& src, int64_t i,
+ std::shared_ptr<Device::SyncEvent> sync,
+ struct ArrowDeviceArray* out_array) {
+ if (i >= src->num_chunks()) {
+ // End of stream
+ ArrowArrayMarkReleased(&out_array->array);
+ return Status::OK();
+ } else {
+ return ExportDeviceArray(*src->chunk(static_cast<int>(i)), sync, out_array);
Review Comment:
`std::move(sync)` to save on an atomic ref-count update.
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -84,6 +92,14 @@ inline void ArrowArrayMove(struct ArrowArray* src, struct ArrowArray* dest) {
ArrowArrayMarkReleased(src);
}
+inline void ArrowDeviceArrayMove(struct ArrowDeviceArray* src,
+ struct ArrowDeviceArray* dest) {
+ assert(dest != src);
+ assert(!ArrowDeviceArrayIsReleased(src));
+ memcpy(dest, src, sizeof(struct ArrowDeviceArray));
Review Comment:
Oh... the others are inverted too.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1552004964
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
Is it ok to keep a mere reference here?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551965664
##########
cpp/src/arrow/record_batch.h:
##########
@@ -23,6 +23,7 @@
#include <vector>
#include "arrow/compare.h"
+#include "arrow/device.h"
Review Comment:
Can we instead add the necessary declarations to `type_fwd.h`?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553869528
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2067,14 +2092,20 @@ class ExportedArrayStream {
ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
};
- explicit ExportedArrayStream(struct ArrowArrayStream* stream) : stream_(stream) {}
+ explicit ExportedArrayStream(StreamType* stream) : stream_(stream) {}
Status GetSchema(struct ArrowSchema* out_schema) {
return ExportStreamSchema(reader(), out_schema);
}
- Status GetNext(struct ArrowArray* out_array) {
- return ExportStreamNext(reader(), next_batch_num(), out_array);
+ Status GetNext(ArrayType* out_array) {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowArray>) {
+ return ExportStreamNext(reader(), next_batch_num(), out_array);
+ } else if constexpr (std::is_same_v<T, ChunkedArray>) {
+ return ExportDeviceStreamNext(reader(), next_batch_num(), nullptr, out_array);
Review Comment:
Right now it's a placeholder so I don't need to change the API when we decide how we want to handle sync events with chunked arrays in the future. If you think it'd be better to just leave it out of the signature for now and leave a comment for the future instead (since this isn't part of the public interface anyways) then I can do that.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2088499517
From what I can see in GDB, this is happening with an empty schema (no fields), but then I don't understand why the size check in recordbatch.cc is not preventing this:
```
if (boxed_columns_.size() > 0) {
device_type_ = boxed_columns_[0]->device_type();
}
```
(but also no idea why we are ending up creating a record batch with an empty schema for this test)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2088957599
@jorisvandenbossche i solved the crashing python test. Essentially I inadvertently exposed a bug in `MakeArrayOfNull` when handling nested extension types :smile:
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551972371
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
Well, this is an abstract class, so there's no constructor here :-)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2078195035
@jorisvandenbossche do you know if the failures of `test_extension_type.py` are actually related to this PR's changes at all? I can't seem to find out enough info to pin it down
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1591178649
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(buf->device_type()));
Review Comment:
Added the debug checks in the ArrayData constructors
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2088365810
> @paleolimbot said he was able to replicate the failure with his own debug python build
It would be more accurate to say that I ran the tests with a debug build of Python since I have a decent set up for that now, and got:
```
> raise CalledProcessError(retcode, cmd)
E subprocess.CalledProcessError: Command '['/Users/dewey/Desktop/rscratch/arrow/.venv/bin/python', 'setup.py', 'build_ext', '--inplace']' returned non-zero exit status 1.
/opt/homebrew/Cellar/python@3.12/3.12.2_1/Frameworks/Python.framework/Versions/3.12/lib/python3.12/subprocess.py:413: CalledProcessError
```
...with at least one segfaulting Python subprocess. I don't have a lot of experience with this type of test but if it is not on `main`, is reproducible, and has never been seen before it is probably a result of this change.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1586297468
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(buf->device_type()));
Review Comment:
Should this only be a debug check? You can create an array with buffers from a different device, or do we disallow that upon creation?
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -623,6 +667,16 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return null_sync_event_;
+}
+
+DeviceAllocationType RecordBatch::device_type() const {
+ return DeviceAllocationType::kCPU;
Review Comment:
Why is this not returning `device_type_`?
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out);
+
+/// \brief Export C++ ChunkedArray using the C device data interface format.
+///
+/// The resulting ArrowDeviceArrayStream keeps the chunked array data and buffers
+/// alive until its release callback is called by the consumer.
+///
+/// \param[in] chunked_array ChunkedArray object to export
+/// \param[in] device_type the device type the data is located on
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
Review Comment:
Is there a reason the device type is an argument here, while it's not for the plain array export (`ExportDeviceArray`) or for the RecordBatchReader export just above
##########
cpp/src/arrow/array/array_base.h:
##########
@@ -224,6 +224,8 @@ class ARROW_EXPORT Array {
/// \return Status
Status ValidateFull() const;
+ DeviceAllocationType device_type() const { return data_->device_type(); }
Review Comment:
Can you add a doc comment?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1586401760
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(buf->device_type()));
Review Comment:
We don't currently disallow it on creation. We don't do any checking of the buffer device types on ArrayData or Array creation currently. It feels wrong to do so in a non-debug context but it might make sense to add debug checks to confirm the device types of all the buffers match in the constructor. Thoughts?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1754640983
##########
cpp/src/arrow/record_batch.h:
##########
@@ -23,6 +23,7 @@
#include <vector>
#include "arrow/compare.h"
+#include "arrow/device.h"
Review Comment:
> I think at least moving DeviceAllocationType to type_fwd.h would make sense.
I've done that here https://github.com/apache/arrow/pull/43853
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade merged PR #40807:
URL: https://github.com/apache/arrow/pull/40807
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1540286553
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2151,53 +2188,96 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ return ExportedArrayStream<RecordBatchReader, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array), out);
+ return ExportedArrayStream<ChunkedArray, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(reader), out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(chunked_array),
+ out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <typename StreamTraits, typename ArrayTraits>
class ArrayStreamReader {
+ protected:
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper& mapper = DefaultDeviceMapper)
+ : mapper_{mapper} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
}
return status;
}
+ Result<std::shared_ptr<RecordBatch>> ImportRecordBatchInternal(
+ ArrayType* array, std::shared_ptr<Schema> schema) {
Review Comment:
Would this be possible to do this as an overload instead of an `if constexpr`?
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,37 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \note it is assumed that the output pointer has already be zeroed out before
+/// calling this function.
Review Comment:
Is this true for other export methods? I don't think this is a general expectation for any exporting function I've seen, although it *is* safer (several months ago I responded to somebody who thought they had a memory leak, but it turns out they were accidentally reusing an `ArrowSchema`).
If it is something worth adding here, it would be sufficient that the structure pointed to by the output pointer is marked as released (not necessarily completely zeroed out).
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual std::shared_ptr<Device::SyncEvent> GetSyncEvent();
Review Comment:
Is it reasonable to expect that when a member `Array` is accessed or returned that a user of the API will be able to walk back up to the `RecordBatch` and check the `SyncEvent`? (I don't have a better idea...giving each `Array` a copy of the shared pointer also seems like it might have problems).
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -115,6 +147,14 @@ inline void ArrowArrayStreamMove(struct ArrowArrayStream* src,
ArrowArrayStreamMarkReleased(src);
}
+inline void ArrowDeviceArrayStreamMove(struct ArrowDeviceArrayStream* src,
+ struct ArrowDeviceArrayStream* dest) {
+ assert(dest != src);
Review Comment:
I also see `ARROW_C_ASSERT()` and `DCHECK()`...is there a preferred one of these to use for consistency?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551891352
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -84,6 +92,14 @@ inline void ArrowArrayMove(struct ArrowArray* src, struct ArrowArray* dest) {
ArrowArrayMarkReleased(src);
}
+inline void ArrowDeviceArrayMove(struct ArrowDeviceArray* src,
+ struct ArrowDeviceArray* dest) {
+ assert(dest != src);
Review Comment:
Every so slightly out-of-scope, but technically this file is missing `#include <assert.h>`.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551979138
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -84,6 +92,14 @@ inline void ArrowArrayMove(struct ArrowArray* src, struct ArrowArray* dest) {
ArrowArrayMarkReleased(src);
}
+inline void ArrowDeviceArrayMove(struct ArrowDeviceArray* src,
+ struct ArrowDeviceArray* dest) {
+ assert(dest != src);
Review Comment:
Good point, we should add it.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551995013
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2041,6 +2041,19 @@ Status ExportStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t i
}
}
+Status ExportDeviceStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t,
Review Comment:
If the `int64_t` is not used, just remove it from the signature?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551883203
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
It seems like for now this will have to be supplied by the caller on construction? That seems like a better intermediate state than returning a value that could lead to segfaults?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553873043
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
it *should* be okay since it's an `std::function`, I wanted to avoid performing a copy if it wasn't necessary.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553855438
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2041,6 +2041,19 @@ Status ExportStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t i
}
}
+Status ExportDeviceStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t,
Review Comment:
It's in the signature so that we can leverage overloading instead of having to add `if constexpr` checks to pass the `int64_t` or not depending on whether we're using a `RecordBatchReader` or ` ChunkedArray`. Removing it from the signature would prevent being able to use the overloading.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551974105
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
Review Comment:
Can we mark this EXPERIMENTAL until we're sure this is the right API at the right place?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541369648
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual std::shared_ptr<Device::SyncEvent> GetSyncEvent();
Review Comment:
Got it! Making sure the buffers had access to the event is what I was worried about. I can't speak to how this would be used at the `Array` or `RecordBatch` level but it might be helpful to explain that in a comment for those of us that don't run into sync events that often.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1544857636
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2151,53 +2188,96 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ return ExportedArrayStream<RecordBatchReader, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array), out);
+ return ExportedArrayStream<ChunkedArray, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(reader), out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(chunked_array),
+ out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <typename StreamTraits, typename ArrayTraits>
class ArrayStreamReader {
+ protected:
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper& mapper = DefaultDeviceMapper)
+ : mapper_{mapper} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
}
return status;
}
+ Result<std::shared_ptr<RecordBatch>> ImportRecordBatchInternal(
+ ArrayType* array, std::shared_ptr<Schema> schema) {
Review Comment:
well i already did it. I'll wait for more feedback and see if anyone has an opinion on this :)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551967477
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
Review Comment:
Does `ARROW_DEVICE_CPU` even exist as a symbol?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551878346
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
Review Comment:
Would it make sense to add `GetSyncEvent()` when it actually does something? It seems like this is perhaps guessing at a future API that we don't know will exist yet or that we are not sure will be used?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551997153
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2067,14 +2092,20 @@ class ExportedArrayStream {
ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
};
- explicit ExportedArrayStream(struct ArrowArrayStream* stream) : stream_(stream) {}
+ explicit ExportedArrayStream(StreamType* stream) : stream_(stream) {}
Status GetSchema(struct ArrowSchema* out_schema) {
return ExportStreamSchema(reader(), out_schema);
}
- Status GetNext(struct ArrowArray* out_array) {
- return ExportStreamNext(reader(), next_batch_num(), out_array);
+ Status GetNext(ArrayType* out_array) {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowArray>) {
+ return ExportStreamNext(reader(), next_batch_num(), out_array);
+ } else if constexpr (std::is_same_v<T, ChunkedArray>) {
+ return ExportDeviceStreamNext(reader(), next_batch_num(), nullptr, out_array);
Review Comment:
Why pass `nullptr` for the sync event? You could remove it from the function signature instead.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "conbench-apache-arrow[bot] (via GitHub)" <gi...@apache.org>.
conbench-apache-arrow[bot] commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2123785178
After merging your PR, Conbench analyzed the 7 benchmarking runs that have been run so far on merge-commit 8169d6e719453acd0e7ca1b6f784d800cca4f113.
There were no benchmark performance regressions. 🎉
The [full Conbench report](https://github.com/apache/arrow/runs/25257512991) has more details. It also includes information about 9 possible false positives for unstable benchmarks that are known to sometimes produce them.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1587554861
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -623,6 +667,16 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return null_sync_event_;
+}
+
+DeviceAllocationType RecordBatch::device_type() const {
+ return DeviceAllocationType::kCPU;
Review Comment:
Ah, I missed the distinction between RecordBatch and SimpleRecordBatch in the diff
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1560096246
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
I've propagated the device_type throughout the readers now and to the static `Make` methods for constructing the readers, so now this should be correct.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1550973742
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
Review Comment:
Is the idea that longer term there would be a generic implementation here in RecordBatch that checks the sync events of its underlying buffers? Because in practice, we don't subclass RecordBatch in Arrow C++ to add CUDA support, so that method cannot be overriden
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1580251600
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2067,14 +2092,20 @@ class ExportedArrayStream {
ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
};
- explicit ExportedArrayStream(struct ArrowArrayStream* stream) : stream_(stream) {}
+ explicit ExportedArrayStream(StreamType* stream) : stream_(stream) {}
Status GetSchema(struct ArrowSchema* out_schema) {
return ExportStreamSchema(reader(), out_schema);
}
- Status GetNext(struct ArrowArray* out_array) {
- return ExportStreamNext(reader(), next_batch_num(), out_array);
+ Status GetNext(ArrayType* out_array) {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowArray>) {
+ return ExportStreamNext(reader(), next_batch_num(), out_array);
+ } else if constexpr (std::is_same_v<T, ChunkedArray>) {
+ return ExportDeviceStreamNext(reader(), next_batch_num(), nullptr, out_array);
Review Comment:
Just a small vote from me that I think it might be less confusing to leave it out of the signature for now if it is not doing anything.
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,36 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ int type = 0;
Review Comment:
Would a comment (either here or in the header) be appropriate here to explain that this is walking the tree of buffers to ensure that all of them are allocated on the same device?
For `type == 0` specifically, it looks like you are using this as a sentinel for "not yet assigned"? Would it be appropriate to define this somewhere with a descriptive name?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2041,6 +2041,19 @@ Status ExportStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t i
}
}
+Status ExportDeviceStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t,
Review Comment:
Maybe worth giving it a name (e.g., `int64_t i`) and a comment explaining that to the unwary reader?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541314162
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -115,6 +147,14 @@ inline void ArrowArrayStreamMove(struct ArrowArrayStream* src,
ArrowArrayStreamMarkReleased(src);
}
+inline void ArrowDeviceArrayStreamMove(struct ArrowDeviceArrayStream* src,
+ struct ArrowDeviceArrayStream* dest) {
+ assert(dest != src);
Review Comment:
I was just maintaining the consistency that the existing non-device functions had
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553838834
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
Review Comment:
good point, it's a macro that exists only in `abi.h`, i'll fix the comment to reference the DeviceAllocationType
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1556407921
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
Let's not overdo this either.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551825314
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
As of yet, I haven't implemented figuring out what type to send here. Since we aren't storing the device type at the record batch level, this would need to spin through the buffers of its columns which could *potentially* be from multiple devices if a user did something weird. So I'm not sure what the best solution here is. I'd rather not add a `RecordBatch` level device_type member, but I don't yet know how to best handle the situation if a consumer does something bad
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2086853686
@paleolimbot Damn i was hoping that the rebase I just did would solve those segfaults :frowning: Guess i'm gonna have to dig in and figure out what the hell is segfaulting there
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1611641179
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,42 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
Review Comment:
This condition implies that, conversely, in non-debug mode we could immediately return when we encounter a buffer? Instead of continue looping on all buffers and children...
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551975341
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
Review Comment:
Make this EXPERIMENTAL as well?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2088492861
The Python failures seems definitely related. Fetching this PR in my local dev setup, I see the same segfault running the python tests. GDB backtrace:
```
$ gdb --args python -m pytest python/pyarrow/tests/test_extension_type.py::test_parquet_extension_with_nested_storage
...
Thread 1 "python" received signal SIGSEGV, Segmentation fault.
arrow::ArrayData::device_type (this=0x0) at /home/joris/scipy/repos/arrow/cpp/src/arrow/array/data.cc:234
234 for (const auto& buf : buffers) {
(gdb) bt
#0 arrow::ArrayData::device_type (this=0x0) at /home/joris/scipy/repos/arrow/cpp/src/arrow/array/data.cc:234
#1 0x00007ffff4a51b60 in arrow::ArrayData::device_type (this=0x555556d4e6c0)
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr_base.h:1665
#2 0x00007ffff4879610 in arrow::Array::device_type (this=<optimized out>)
at /home/joris/scipy/repos/arrow/cpp/src/arrow/array/array_base.h:227
#3 arrow::SimpleRecordBatch::SimpleRecordBatch (this=0x555556ea58b0, schema=..., num_rows=<optimized out>, columns=..., sync_event=...)
at /home/joris/scipy/repos/arrow/cpp/src/arrow/record_batch.cc:69
#4 0x00007ffff4890323 in std::_Construct<arrow::SimpleRecordBatch, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (__p=<optimized out>)
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/stl_construct.h:119
#5 std::allocator_traits<std::allocator<void> >::construct<arrow::SimpleRecordBatch, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (
__p=<optimized out>) at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/alloc_traits.h:635
#6 std::_Sp_counted_ptr_inplace<arrow::SimpleRecordBatch, std::allocator<void>, (__gnu_cxx::_Lock_policy)2>::_Sp_counted_ptr_inplace<std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (__a=..., this=<optimized out>)
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr_base.h:604
#7 std::__shared_count<(__gnu_cxx::_Lock_policy)2>::__shared_count<arrow::SimpleRecordBatch, std::allocator<void>, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (__a=..., __p=<optimized out>, this=<optimized out>)
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr_base.h:971
#8 std::__shared_ptr<arrow::SimpleRecordBatch, (__gnu_cxx::_Lock_policy)2>::__shared_ptr<std::allocator<void>, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (__tag=..., this=<optimized out>)
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr_base.h:1712
#9 std::shared_ptr<arrow::SimpleRecordBatch>::shared_ptr<std::allocator<void>, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > (__tag=...,
this=<optimized out>) at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr.h:464
#10 std::make_shared<arrow::SimpleRecordBatch, std::shared_ptr<arrow::Schema>, long&, std::vector<std::shared_ptr<arrow::Array>, std::allocator<std::shared_ptr<arrow::Array> > >, std::shared_ptr<arrow::Device::SyncEvent> > ()
at /home/joris/conda/envs/arrow-dev/x86_64-conda-linux-gnu/include/c++/12.3.0/bits/shared_ptr.h:1010
#11 arrow::RecordBatch::Make (schema=..., num_rows=num_rows@entry=0, columns=..., sync_event=...)
at /home/joris/scipy/repos/arrow/cpp/src/arrow/record_batch.cc:217
#12 0x00007fff48dcb9ae in arrow::dataset::(anonymous namespace)::FragmentToBatches (fragment=..., options=...)
at /home/joris/scipy/repos/arrow/cpp/src/arrow/dataset/scanner.cc:316
#13 0x00007fff48dcc364 in operator() (fragment=..., __closure=<optimized out>)
at /home/joris/scipy/repos/arrow/cpp/src/arrow/dataset/scanner.cc:333
```
But not directly understanding why it is crashing on this. The problem seems to be that the ArrayData itself is a nullptr.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1543728000
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
Review Comment:
Potentially related to my other comment, but the existing `ExportDeviceArray` has a `sync` keyword that the user of this API needs to provide. Is there a reason those methods don't have that?
The returned `ArrowDeviceArrayStream` itself doesn't have a sync event member, but the `ArrowDeviceArray`s that it will return still have that. The user shouldn't pass the sync event to set in those arrays?
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
Review Comment:
I don't fully understand this part. AFAIK a RecordBatch currently is agnostic the device of its buffers. So for example you can have a RecordBatch backed by buffers that live in CUDA memory, but then this method will always hardcoded return NULL, which is not correct in that case?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583750847
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
Review Comment:
Was there a backlash about adding `kUNDEFINED` to `DeviceAllocationType`?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2340,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper mapper_;
};
-class ArrayStreamBatchReader : public RecordBatchReader, public ArrayStreamReader {
+template <typename StreamTraits, typename ArrayTraits>
+class ArrayStreamBatchReader : public RecordBatchReader,
+ public ArrayStreamReader<StreamTraits, ArrayTraits> {
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamBatchReader(struct ArrowArrayStream* stream)
- : ArrayStreamReader(stream) {}
+ explicit ArrayStreamBatchReader(
+ StreamType* stream, const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper)
+ : ArrayStreamReader<StreamTraits, ArrayTraits>(stream, mapper) {}
Status Init() {
- ARROW_ASSIGN_OR_RAISE(schema_, ReadSchema());
+ ARROW_ASSIGN_OR_RAISE(schema_, this->ReadSchema());
return Status::OK();
}
std::shared_ptr<Schema> schema() const override { return schema_; }
Status ReadNext(std::shared_ptr<RecordBatch>* batch) override {
- ARROW_RETURN_NOT_OK(CheckNotReleased());
+ ARROW_RETURN_NOT_OK(this->CheckNotReleased());
- struct ArrowArray c_array;
- ARROW_RETURN_NOT_OK(ReadNextArrayInternal(&c_array));
+ ArrayType c_array;
+ ARROW_RETURN_NOT_OK(this->ReadNextArrayInternal(&c_array));
- if (ArrowArrayIsReleased(&c_array)) {
+ if (ArrayTraits::IsReleasedFunc(&c_array)) {
// End of stream
batch->reset();
return Status::OK();
} else {
- return ImportRecordBatch(&c_array, schema_).Value(batch);
+ return this->ImportRecordBatchInternal(&c_array, schema_).Value(batch);
}
}
Status Close() override {
- ReleaseStream();
+ this->ReleaseStream();
return Status::OK();
}
+ DeviceAllocationType device_type() const override { return this->get_device_type(); }
+
private:
std::shared_ptr<Schema> schema_;
};
-class ArrayStreamArrayReader : public ArrayStreamReader {
+template <typename StreamTraits, typename ArrayTraits>
+class ArrayStreamArrayReader : public ArrayStreamReader<StreamTraits, ArrayTraits> {
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamArrayReader(struct ArrowArrayStream* stream)
- : ArrayStreamReader(stream) {}
+ explicit ArrayStreamArrayReader(
+ StreamType* stream, const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper)
Review Comment:
I would recommend avoiding default param values on internal abstractions. Not having them forces us to think about the correct value to pass on every callsite.
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -167,30 +188,41 @@ class SimpleRecordBatch : public RecordBatch {
return RecordBatch::Validate();
}
+ std::shared_ptr<Device::SyncEvent> GetSyncEvent() const override { return sync_event_; }
Review Comment:
Are there any override that would have to create one `Device::SyncEvent` per call? If not, this should return a `const ...&` to avoid the non-trivial cost of incrementing a ref-count on every call.
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out);
+
+/// \brief Export C++ ChunkedArray using the c device data interface format.
Review Comment:
```suggestion
/// \brief Export C++ ChunkedArray using the C device data interface format.
```
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -167,30 +188,41 @@ class SimpleRecordBatch : public RecordBatch {
return RecordBatch::Validate();
}
+ std::shared_ptr<Device::SyncEvent> GetSyncEvent() const override { return sync_event_; }
+
+ DeviceAllocationType device_type() const override { return device_type_; }
+
private:
std::vector<std::shared_ptr<ArrayData>> columns_;
// Caching boxed array data
mutable std::vector<std::shared_ptr<Array>> boxed_columns_;
+
+ // the type of device that the buffers for columns are allocated on.
+ // all columns should be on the same type of device.
+ DeviceAllocationType device_type_;
+ std::shared_ptr<Device::SyncEvent> sync_event_;
};
RecordBatch::RecordBatch(const std::shared_ptr<Schema>& schema, int64_t num_rows)
: schema_(schema), num_rows_(num_rows) {}
std::shared_ptr<RecordBatch> RecordBatch::Make(
std::shared_ptr<Schema> schema, int64_t num_rows,
- std::vector<std::shared_ptr<Array>> columns) {
+ std::vector<std::shared_ptr<Array>> columns,
+ std::shared_ptr<Device::SyncEvent> sync_event) {
DCHECK_EQ(schema->num_fields(), static_cast<int>(columns.size()));
return std::make_shared<SimpleRecordBatch>(std::move(schema), num_rows,
- std::move(columns));
+ std::move(columns), sync_event);
}
std::shared_ptr<RecordBatch> RecordBatch::Make(
std::shared_ptr<Schema> schema, int64_t num_rows,
- std::vector<std::shared_ptr<ArrayData>> columns) {
+ std::vector<std::shared_ptr<ArrayData>> columns, DeviceAllocationType device_type,
+ std::shared_ptr<Device::SyncEvent> sync_event) {
DCHECK_EQ(schema->num_fields(), static_cast<int>(columns.size()));
return std::make_shared<SimpleRecordBatch>(std::move(schema), num_rows,
- std::move(columns));
+ std::move(columns), device_type, sync_event);
Review Comment:
`std::move(sync_event)`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1587559683
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
+ } else {
+ DCHECK_EQ(type, static_cast<int>(buf->device_type()));
Review Comment:
> it might make sense to add debug checks to confirm the device types of all the buffers match in the constructor
I think that would make sense yes, if we don't want to support that (I don't know if there would actually be use cases for allowing it, in which case I would rather expect a normal error here when trying to export it, saying it requires a single device type)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2096303239
@jorisvandenbossche @pitrou any further comments?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1557622935
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
Yes, exactly.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1556477280
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
> Let's not overdo this either.
So just removing the `&` is what you have in mind?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553909516
##########
cpp/src/arrow/record_batch.h:
##########
@@ -23,6 +23,7 @@
#include <vector>
#include "arrow/compare.h"
+#include "arrow/device.h"
Review Comment:
We would have to move the nested `Device::SyncEvent` outside of the `Device` class to do that since you can't forward declare an inner class like that. It would also require pushing the entire `DeviceAllocationType` declaration to `type_fwd.h` which I'm not sure if we want to do. So I don't think we can do this.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1556406405
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
Would be nice to be able to enforce that only move-only functions [1] be accepted here.
[1] https://github.com/ofats/any_invocable
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2079607152
> Just to check my understanding of this change: sync events are now propagated from device arrays to record batches (in addition to all the buffers in all the arrays in the record batch); however, one would have to specifically specify a sync event explicitly if constructing a record batch from (say) some arrays (or else it will return nullptr, meaning "the batch is synchronized" or "I'm handling synchronization myself").
Yep. That's the current situational understanding.
> It seems as though you are punting on sync events for the Array for now...is walking all the buffers and checking for the same underlying non-null sync event pointer prohibitively expensive?
I didn't want to require the same underlying non-null sync event pointer in all children as that would make it difficult to manipulate / update / construct record batches from arrays gathered from multiple places. But we don't currently have any sort of API for defining how to merge / join events into a single event and I didn't want to try to tackle that in this PR.
essentially, it's not prohibitively expensive, we just don't want that to be a requirement
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1581183170
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2041,6 +2041,19 @@ Status ExportStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t i
}
}
+Status ExportDeviceStreamNext(const std::shared_ptr<RecordBatchReader>& src, int64_t,
Review Comment:
Sure, added.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "bkietz (via GitHub)" <gi...@apache.org>.
bkietz commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1586297491
##########
python/pyarrow/tests/test_cffi.py:
##########
@@ -45,7 +45,7 @@
ValueError, match="Cannot import released ArrowArray")
assert_stream_released = pytest.raises(
- ValueError, match="Cannot import released ArrowArrayStream")
+ ValueError, match="Cannot import released Arrow Stream")
Review Comment:
nit: I suppose this is to be generic since it might be an ArrowDeviceArrayStream instead, but it's inconsistent with the above asserters
##########
cpp/src/arrow/record_batch.h:
##########
@@ -260,6 +270,18 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief EXPERIMENTAL: Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in CPU memory, then this
+ /// will return null. If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual const std::shared_ptr<Device::SyncEvent>& GetSyncEvent() const;
Review Comment:
Is it desirable to give this a default implementation at all? Seems it could instead be pure virtual
##########
cpp/src/arrow/c/bridge_test.cc:
##########
@@ -4746,4 +4750,531 @@ TEST_F(TestArrayStreamRoundtrip, ChunkedArrayRoundtripEmpty) {
});
}
+////////////////////////////////////////////////////////////////////////////
+// Array device stream export tests
+
+class TestArrayDeviceStreamExport : public BaseArrayStreamTest {
+ public:
+ void AssertStreamSchema(struct ArrowDeviceArrayStream* c_stream,
+ const Schema& expected) {
+ struct ArrowSchema c_schema;
+ ASSERT_EQ(0, c_stream->get_schema(c_stream, &c_schema));
+
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_FALSE(ArrowSchemaIsReleased(&c_schema));
+ ASSERT_OK_AND_ASSIGN(auto schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(expected, *schema, /*check_metadata=*/true);
+ }
+
+ void AssertStreamEnd(struct ArrowDeviceArrayStream* c_stream) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_TRUE(ArrowDeviceArrayIsReleased(&c_array));
+ }
+
+ void AssertStreamNext(struct ArrowDeviceArrayStream* c_stream,
+ const RecordBatch& expected) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_FALSE(ArrowDeviceArrayIsReleased(&c_array));
+
+ ASSERT_OK_AND_ASSIGN(auto batch,
+ ImportDeviceRecordBatch(&c_array, expected.schema(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(expected, *batch);
+ }
+
+ void AssertStreamNext(struct ArrowDeviceArrayStream* c_stream, const Array& expected) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_FALSE(ArrowDeviceArrayIsReleased(&c_array));
+
+ ASSERT_OK_AND_ASSIGN(auto array,
+ ImportDeviceArray(&c_array, expected.type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertArraysEqual(expected, *array);
+ }
+
+ static Result<std::shared_ptr<ArrayData>> ToDeviceData(
+ const std::shared_ptr<MemoryManager>& mm, const ArrayData& data) {
+ arrow::BufferVector buffers;
+ for (const auto& buf : data.buffers) {
+ if (buf) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, mm->CopyBuffer(buf, mm));
+ buffers.push_back(dest);
+ } else {
+ buffers.push_back(nullptr);
+ }
+ }
+
+ arrow::ArrayDataVector children;
+ for (const auto& child : data.child_data) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+ children.push_back(dest);
+ }
+
+ return ArrayData::Make(data.type, data.length, buffers, children, data.null_count,
+ data.offset);
+ }
+
+ static Result<std::shared_ptr<Array>> ToDevice(const std::shared_ptr<MemoryManager>& mm,
+ const ArrayData& data) {
+ ARROW_ASSIGN_OR_RAISE(auto result, ToDeviceData(mm, data));
+ return MakeArray(result);
+ }
+};
+
+TEST_F(TestArrayDeviceStreamExport, Empty) {
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader,
+ RecordBatchReader::Make(batches, schema,
+ static_cast<DeviceAllocationType>(kMyDeviceType)));
+
+ struct ArrowDeviceArrayStream c_stream;
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ AssertStreamSchema(&c_stream, *schema);
+ AssertStreamEnd(&c_stream);
+ AssertStreamEnd(&c_stream);
+}
+
+TEST_F(TestArrayDeviceStreamExport, Simple) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ RecordBatchReader::Make(batches, schema, device->device_type()));
+
+ struct ArrowDeviceArrayStream c_stream;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ AssertStreamSchema(&c_stream, *schema);
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ AssertStreamNext(&c_stream, *batches[0]);
+ AssertStreamNext(&c_stream, *batches[1]);
+ AssertStreamEnd(&c_stream);
+ AssertStreamEnd(&c_stream);
+}
+
+TEST_F(TestArrayDeviceStreamExport, ArrayLifetime) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ RecordBatchReader::Make(batches, schema, device->device_type()));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+ struct ArrowDeviceArray c_array0, c_array1;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array0));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array1));
+ AssertStreamEnd(&c_stream);
+ }
+
+ DeviceArrayExportGuard guard0(&c_array0), guard1(&c_array1);
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(*schema, *got_schema, /*check_metadata=*/true);
+ }
+
+ ASSERT_EQ(kMyDeviceType, c_array0.device_type);
+ ASSERT_EQ(kMyDeviceType, c_array1.device_type);
+
+ ASSERT_GT(pool_->bytes_allocated(), orig_allocated_);
+ ASSERT_OK_AND_ASSIGN(
+ auto batch,
+ ImportDeviceRecordBatch(&c_array1, schema, TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(*batches[1], *batch);
+ ASSERT_EQ(device->device_type(), batch->device_type());
+ ASSERT_OK_AND_ASSIGN(
+ batch,
+ ImportDeviceRecordBatch(&c_array0, schema, TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(*batches[0], *batch);
+ ASSERT_EQ(device->device_type(), batch->device_type());
+}
+
+TEST_F(TestArrayDeviceStreamExport, Errors) {
+ auto reader =
+ std::make_shared<FailingRecordBatchReader>(Status::Invalid("some example error"));
+
+ struct ArrowDeviceArrayStream c_stream;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ struct ArrowSchema c_schema;
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_FALSE(ArrowSchemaIsReleased(&c_schema));
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(schema, arrow::schema({}), /*check_metadata=*/true);
+ }
+
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(EINVAL, c_stream.get_next(&c_stream, &c_array));
+}
+
+TEST_F(TestArrayDeviceStreamExport, ChunkedArrayExportEmpty) {
+ ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make({}, int32()));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+
+ ASSERT_OK(ExportDeviceChunkedArray(
+ chunked_array, static_cast<DeviceAllocationType>(kMyDeviceType), &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ AssertStreamEnd(&c_stream);
+ }
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_type, ImportType(&c_schema));
+ AssertTypeEqual(*chunked_array->type(), *got_type);
+ }
+}
+
+TEST_F(TestArrayDeviceStreamExport, ChunkedArrayExport) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+
+ ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make({arr1, arr2}));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+ struct ArrowDeviceArray c_array0, c_array1;
+
+ ASSERT_OK(ExportDeviceChunkedArray(chunked_array, device->device_type(), &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array0));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array1));
+ AssertStreamEnd(&c_stream);
+ }
+
+ DeviceArrayExportGuard guard0(&c_array0), guard1(&c_array1);
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_type, ImportType(&c_schema));
+ AssertTypeEqual(*chunked_array->type(), *got_type);
+ }
+
+ ASSERT_EQ(kMyDeviceType, c_array0.device_type);
+ ASSERT_EQ(kMyDeviceType, c_array1.device_type);
+
+ ASSERT_GT(pool_->bytes_allocated(), orig_allocated_);
+ ASSERT_OK_AND_ASSIGN(auto array,
+ ImportDeviceArray(&c_array0, chunked_array->type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ ASSERT_EQ(device->device_type(), array->device_type());
+ AssertArraysEqual(*chunked_array->chunk(0), *array);
+ ASSERT_OK_AND_ASSIGN(array, ImportDeviceArray(&c_array1, chunked_array->type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ ASSERT_EQ(device->device_type(), array->device_type());
+ AssertArraysEqual(*chunked_array->chunk(1), *array);
+}
+
+////////////////////////////////////////////////////////////////////////////
+// Array device stream roundtrip tests
+
+class TestArrayDeviceStreamRoundtrip : public BaseArrayStreamTest {
+ public:
+ static Result<std::shared_ptr<ArrayData>> ToDeviceData(
+ const std::shared_ptr<MemoryManager>& mm, const ArrayData& data) {
+ arrow::BufferVector buffers;
+ for (const auto& buf : data.buffers) {
+ if (buf) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, mm->CopyBuffer(buf, mm));
+ buffers.push_back(dest);
+ } else {
+ buffers.push_back(nullptr);
+ }
+ }
+
+ arrow::ArrayDataVector children;
+ for (const auto& child : data.child_data) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+ children.push_back(dest);
+ }
+
+ return ArrayData::Make(data.type, data.length, buffers, children, data.null_count,
+ data.offset);
+ }
+
+ static Result<std::shared_ptr<Array>> ToDevice(const std::shared_ptr<MemoryManager>& mm,
+ const ArrayData& data) {
+ ARROW_ASSIGN_OR_RAISE(auto result, ToDeviceData(mm, data));
+ return MakeArray(result);
+ }
+
+ void Roundtrip(std::shared_ptr<RecordBatchReader>* reader,
+ struct ArrowDeviceArrayStream* c_stream) {
+ ASSERT_OK(ExportDeviceRecordBatchReader(*reader, c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(c_stream));
+
+ ASSERT_OK_AND_ASSIGN(
+ auto got_reader,
+ ImportDeviceRecordBatchReader(c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ *reader = std::move(got_reader);
+ }
+
+ void Roundtrip(
+ std::shared_ptr<RecordBatchReader> reader,
+ std::function<void(const std::shared_ptr<RecordBatchReader>&)> check_func) {
+ ArrowDeviceArrayStream c_stream;
+
+ // NOTE: ReleaseCallback<> is not immediately usable with ArrowDeviceArayStream
+ // because get_next and get_schema need the original private_data.
+ std::weak_ptr<RecordBatchReader> weak_reader(reader);
+ ASSERT_EQ(weak_reader.use_count(), 1); // Expiration check will fail otherwise
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(std::move(reader), &c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ {
+ ASSERT_OK_AND_ASSIGN(auto new_reader,
+ ImportDeviceRecordBatchReader(
+ &c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ // stream was moved
+ ASSERT_TRUE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_FALSE(weak_reader.expired());
+
+ check_func(new_reader);
+ }
+ // Stream was released when `new_reader` was destroyed
+ ASSERT_TRUE(weak_reader.expired());
+ }
+
+ void Roundtrip(std::shared_ptr<ChunkedArray> src,
+ std::function<void(const std::shared_ptr<ChunkedArray>&)> check_func) {
+ ArrowDeviceArrayStream c_stream;
+
+ // One original copy to compare the result, one copy held by the stream
+ std::weak_ptr<ChunkedArray> weak_src(src);
+ int64_t initial_use_count = weak_src.use_count();
+
+ ASSERT_OK(ExportDeviceChunkedArray(
+ std::move(src), static_cast<DeviceAllocationType>(kMyDeviceType), &c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ {
+ ASSERT_OK_AND_ASSIGN(
+ auto dst,
+ ImportDeviceChunkedArray(&c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ // Stream was moved, consumed, and released
+ ASSERT_TRUE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ // Stream was released by ImportDeviceChunkedArray but original copy remains
+ ASSERT_EQ(weak_src.use_count(), initial_use_count - 1);
+
+ check_func(dst);
+ }
+ }
+
+ void AssertReaderNext(const std::shared_ptr<RecordBatchReader>& reader,
+ const RecordBatch& expected) {
+ ASSERT_OK_AND_ASSIGN(auto batch, reader->Next());
+ ASSERT_NE(batch, nullptr);
+ ASSERT_EQ(static_cast<DeviceAllocationType>(kMyDeviceType), batch->device_type());
+ AssertBatchesEqual(expected, *batch);
+ }
+
+ void AssertReaderEnd(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_OK_AND_ASSIGN(auto batch, reader->Next());
+ ASSERT_EQ(batch, nullptr);
+ }
+
+ void AssertReaderClosed(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_THAT(reader->Next(),
+ Raises(StatusCode::Invalid, ::testing::HasSubstr("already been closed")));
+ }
+
+ void AssertReaderClose(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_OK(reader->Close());
+ AssertReaderClosed(reader);
+ }
+};
+
+TEST_F(TestArrayDeviceStreamRoundtrip, Simple) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto orig_schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(orig_schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader, RecordBatchReader::Make(batches, orig_schema, device->device_type()));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ AssertSchemaEqual(*orig_schema, *reader->schema(), /*check_metadata=*/true);
+ AssertReaderNext(reader, *batches[0]);
+ AssertReaderNext(reader, *batches[1]);
+ AssertReaderEnd(reader);
+ AssertReaderEnd(reader);
+ AssertReaderClose(reader);
+ });
+}
+
+TEST_F(TestArrayDeviceStreamRoundtrip, CloseEarly) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto orig_schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(orig_schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader, RecordBatchReader::Make(batches, orig_schema, device->device_type()));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ AssertReaderNext(reader, *batches[0]);
+ AssertReaderClose(reader);
+ });
+}
+
+TEST_F(TestArrayDeviceStreamRoundtrip, Errors) {
+ auto reader = std::make_shared<FailingRecordBatchReader>(
+ Status::Invalid("roundtrip error example"));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ auto status = reader->Next().status();
+ ASSERT_RAISES(Invalid, status);
+ ASSERT_THAT(status.message(), ::testing::HasSubstr("roundtrip error example"));
+ });
+}
+
+TEST_F(TestArrayDeviceStreamRoundtrip, SchemaError) {
+ struct StreamState {
+ bool released = false;
+
+ static const char* GetLastError(struct ArrowDeviceArrayStream* stream) {
+ return "Expected error";
+ }
+
+ static int GetSchema(struct ArrowDeviceArrayStream* stream,
+ struct ArrowSchema* schema) {
+ return EIO;
+ }
+
+ static int GetNext(struct ArrowDeviceArrayStream* stream,
+ struct ArrowDeviceArray* array) {
+ return EINVAL;
+ }
+
+ static void Release(struct ArrowDeviceArrayStream* stream) {
+ reinterpret_cast<StreamState*>(stream->private_data)->released = true;
+ std::memset(stream, 0, sizeof(*stream));
+ }
+ } state;
+ struct ArrowDeviceArrayStream stream = {};
+ stream.get_last_error = &StreamState::GetLastError;
+ stream.get_schema = &StreamState::GetSchema;
+ stream.get_next = &StreamState::GetNext;
+ stream.release = &StreamState::Release;
+ stream.private_data = &state;
Review Comment:
Empty lambdas are convertible to function pointers:
```suggestion
struct ArrowDeviceArrayStream stream = {};
stream.get_last_error = [](struct ArrowDeviceArrayStream* stream) {
return "Expected error";
};
stream.get_schema = [](struct ArrowDeviceArrayStream* stream,
struct ArrowSchema* schema) {
return EIO;
};
stream.get_next = [](struct ArrowDeviceArrayStream* stream,
struct ArrowDeviceArray* array) {
return EINVAL;
};
stream.release = [](struct ArrowDeviceArrayStream* stream) {
*static_cast<bool*>(stream->private_data) = true;
std::memset(stream, 0, sizeof(*stream));
};
bool released = false;
stream.private_data = &released;
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2150,53 +2186,102 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ memset(out, 0, sizeof(struct ArrowArrayStream));
+ return ExportedArrayStream<RecordBatchReader, false>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array), out);
+ memset(out, 0, sizeof(struct ArrowArrayStream));
+ return ExportedArrayStream<ChunkedArray, false>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ memset(out, 0, sizeof(struct ArrowDeviceArrayStream));
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, true>::Make(std::move(reader), out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ memset(out, 0, sizeof(struct ArrowDeviceArrayStream));
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, true>::Make(std::move(chunked_array), out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <bool IsDevice>
class ArrayStreamReader {
+ protected:
+ using StreamTraits =
+ std::conditional_t<IsDevice, internal::ArrayDeviceStreamExportTraits,
+ internal::ArrayStreamExportTraits>;
+ using StreamType = typename StreamTraits::CType;
+ using ArrayTraits = std::conditional_t<IsDevice, internal::ArrayDeviceExportTraits,
+ internal::ArrayExportTraits>;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper mapper = DefaultDeviceMemoryMapper)
+ : mapper_{std::move(mapper)} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
Review Comment:
The comment above suggests we don't need to check whether the array is released either
```suggestion
if (!status.ok()) {
ArrayTraits::ReleaseFunc(array);
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2335,138 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
Review Comment:
```suggestion
if constexpr (IsDevice) {
```
##########
cpp/src/arrow/record_batch.h:
##########
@@ -260,6 +270,18 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief EXPERIMENTAL: Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in CPU memory, then this
+ /// will return null. If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual const std::shared_ptr<Device::SyncEvent>& GetSyncEvent() const;
+
+ virtual DeviceAllocationType device_type() const;
Review Comment:
Same here, now that we always specify the device type in the constructor
##########
cpp/src/arrow/c/bridge_test.cc:
##########
@@ -4746,4 +4750,531 @@ TEST_F(TestArrayStreamRoundtrip, ChunkedArrayRoundtripEmpty) {
});
}
+////////////////////////////////////////////////////////////////////////////
+// Array device stream export tests
+
+class TestArrayDeviceStreamExport : public BaseArrayStreamTest {
+ public:
+ void AssertStreamSchema(struct ArrowDeviceArrayStream* c_stream,
+ const Schema& expected) {
+ struct ArrowSchema c_schema;
+ ASSERT_EQ(0, c_stream->get_schema(c_stream, &c_schema));
+
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_FALSE(ArrowSchemaIsReleased(&c_schema));
+ ASSERT_OK_AND_ASSIGN(auto schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(expected, *schema, /*check_metadata=*/true);
+ }
+
+ void AssertStreamEnd(struct ArrowDeviceArrayStream* c_stream) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_TRUE(ArrowDeviceArrayIsReleased(&c_array));
+ }
+
+ void AssertStreamNext(struct ArrowDeviceArrayStream* c_stream,
+ const RecordBatch& expected) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_FALSE(ArrowDeviceArrayIsReleased(&c_array));
+
+ ASSERT_OK_AND_ASSIGN(auto batch,
+ ImportDeviceRecordBatch(&c_array, expected.schema(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(expected, *batch);
+ }
+
+ void AssertStreamNext(struct ArrowDeviceArrayStream* c_stream, const Array& expected) {
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(0, c_stream->get_next(c_stream, &c_array));
+
+ DeviceArrayExportGuard guard(&c_array);
+ ASSERT_FALSE(ArrowDeviceArrayIsReleased(&c_array));
+
+ ASSERT_OK_AND_ASSIGN(auto array,
+ ImportDeviceArray(&c_array, expected.type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertArraysEqual(expected, *array);
+ }
+
+ static Result<std::shared_ptr<ArrayData>> ToDeviceData(
+ const std::shared_ptr<MemoryManager>& mm, const ArrayData& data) {
+ arrow::BufferVector buffers;
+ for (const auto& buf : data.buffers) {
+ if (buf) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, mm->CopyBuffer(buf, mm));
+ buffers.push_back(dest);
+ } else {
+ buffers.push_back(nullptr);
+ }
+ }
+
+ arrow::ArrayDataVector children;
+ for (const auto& child : data.child_data) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+ children.push_back(dest);
+ }
+
+ return ArrayData::Make(data.type, data.length, buffers, children, data.null_count,
+ data.offset);
+ }
+
+ static Result<std::shared_ptr<Array>> ToDevice(const std::shared_ptr<MemoryManager>& mm,
+ const ArrayData& data) {
+ ARROW_ASSIGN_OR_RAISE(auto result, ToDeviceData(mm, data));
+ return MakeArray(result);
+ }
+};
+
+TEST_F(TestArrayDeviceStreamExport, Empty) {
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader,
+ RecordBatchReader::Make(batches, schema,
+ static_cast<DeviceAllocationType>(kMyDeviceType)));
+
+ struct ArrowDeviceArrayStream c_stream;
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ AssertStreamSchema(&c_stream, *schema);
+ AssertStreamEnd(&c_stream);
+ AssertStreamEnd(&c_stream);
+}
+
+TEST_F(TestArrayDeviceStreamExport, Simple) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ RecordBatchReader::Make(batches, schema, device->device_type()));
+
+ struct ArrowDeviceArrayStream c_stream;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ AssertStreamSchema(&c_stream, *schema);
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ AssertStreamNext(&c_stream, *batches[0]);
+ AssertStreamNext(&c_stream, *batches[1]);
+ AssertStreamEnd(&c_stream);
+ AssertStreamEnd(&c_stream);
+}
+
+TEST_F(TestArrayDeviceStreamExport, ArrayLifetime) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(auto reader,
+ RecordBatchReader::Make(batches, schema, device->device_type()));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+ struct ArrowDeviceArray c_array0, c_array1;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array0));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array1));
+ AssertStreamEnd(&c_stream);
+ }
+
+ DeviceArrayExportGuard guard0(&c_array0), guard1(&c_array1);
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(*schema, *got_schema, /*check_metadata=*/true);
+ }
+
+ ASSERT_EQ(kMyDeviceType, c_array0.device_type);
+ ASSERT_EQ(kMyDeviceType, c_array1.device_type);
+
+ ASSERT_GT(pool_->bytes_allocated(), orig_allocated_);
+ ASSERT_OK_AND_ASSIGN(
+ auto batch,
+ ImportDeviceRecordBatch(&c_array1, schema, TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(*batches[1], *batch);
+ ASSERT_EQ(device->device_type(), batch->device_type());
+ ASSERT_OK_AND_ASSIGN(
+ batch,
+ ImportDeviceRecordBatch(&c_array0, schema, TestDeviceArrayRoundtrip::DeviceMapper));
+ AssertBatchesEqual(*batches[0], *batch);
+ ASSERT_EQ(device->device_type(), batch->device_type());
+}
+
+TEST_F(TestArrayDeviceStreamExport, Errors) {
+ auto reader =
+ std::make_shared<FailingRecordBatchReader>(Status::Invalid("some example error"));
+
+ struct ArrowDeviceArrayStream c_stream;
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(reader, &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ struct ArrowSchema c_schema;
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_FALSE(ArrowSchemaIsReleased(&c_schema));
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto schema, ImportSchema(&c_schema));
+ AssertSchemaEqual(schema, arrow::schema({}), /*check_metadata=*/true);
+ }
+
+ struct ArrowDeviceArray c_array;
+ ASSERT_EQ(EINVAL, c_stream.get_next(&c_stream, &c_array));
+}
+
+TEST_F(TestArrayDeviceStreamExport, ChunkedArrayExportEmpty) {
+ ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make({}, int32()));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+
+ ASSERT_OK(ExportDeviceChunkedArray(
+ chunked_array, static_cast<DeviceAllocationType>(kMyDeviceType), &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ AssertStreamEnd(&c_stream);
+ }
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_type, ImportType(&c_schema));
+ AssertTypeEqual(*chunked_array->type(), *got_type);
+ }
+}
+
+TEST_F(TestArrayDeviceStreamExport, ChunkedArrayExport) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+
+ ASSERT_OK_AND_ASSIGN(auto chunked_array, ChunkedArray::Make({arr1, arr2}));
+
+ struct ArrowDeviceArrayStream c_stream;
+ struct ArrowSchema c_schema;
+ struct ArrowDeviceArray c_array0, c_array1;
+
+ ASSERT_OK(ExportDeviceChunkedArray(chunked_array, device->device_type(), &c_stream));
+ DeviceArrayStreamExportGuard guard(&c_stream);
+
+ {
+ DeviceArrayStreamExportGuard guard(&c_stream);
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ ASSERT_EQ(0, c_stream.get_schema(&c_stream, &c_schema));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array0));
+ ASSERT_EQ(0, c_stream.get_next(&c_stream, &c_array1));
+ AssertStreamEnd(&c_stream);
+ }
+
+ DeviceArrayExportGuard guard0(&c_array0), guard1(&c_array1);
+
+ {
+ SchemaExportGuard schema_guard(&c_schema);
+ ASSERT_OK_AND_ASSIGN(auto got_type, ImportType(&c_schema));
+ AssertTypeEqual(*chunked_array->type(), *got_type);
+ }
+
+ ASSERT_EQ(kMyDeviceType, c_array0.device_type);
+ ASSERT_EQ(kMyDeviceType, c_array1.device_type);
+
+ ASSERT_GT(pool_->bytes_allocated(), orig_allocated_);
+ ASSERT_OK_AND_ASSIGN(auto array,
+ ImportDeviceArray(&c_array0, chunked_array->type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ ASSERT_EQ(device->device_type(), array->device_type());
+ AssertArraysEqual(*chunked_array->chunk(0), *array);
+ ASSERT_OK_AND_ASSIGN(array, ImportDeviceArray(&c_array1, chunked_array->type(),
+ TestDeviceArrayRoundtrip::DeviceMapper));
+ ASSERT_EQ(device->device_type(), array->device_type());
+ AssertArraysEqual(*chunked_array->chunk(1), *array);
+}
+
+////////////////////////////////////////////////////////////////////////////
+// Array device stream roundtrip tests
+
+class TestArrayDeviceStreamRoundtrip : public BaseArrayStreamTest {
+ public:
+ static Result<std::shared_ptr<ArrayData>> ToDeviceData(
+ const std::shared_ptr<MemoryManager>& mm, const ArrayData& data) {
+ arrow::BufferVector buffers;
+ for (const auto& buf : data.buffers) {
+ if (buf) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, mm->CopyBuffer(buf, mm));
+ buffers.push_back(dest);
+ } else {
+ buffers.push_back(nullptr);
+ }
+ }
+
+ arrow::ArrayDataVector children;
+ for (const auto& child : data.child_data) {
+ ARROW_ASSIGN_OR_RAISE(auto dest, ToDeviceData(mm, *child));
+ children.push_back(dest);
+ }
+
+ return ArrayData::Make(data.type, data.length, buffers, children, data.null_count,
+ data.offset);
+ }
+
+ static Result<std::shared_ptr<Array>> ToDevice(const std::shared_ptr<MemoryManager>& mm,
+ const ArrayData& data) {
+ ARROW_ASSIGN_OR_RAISE(auto result, ToDeviceData(mm, data));
+ return MakeArray(result);
+ }
+
+ void Roundtrip(std::shared_ptr<RecordBatchReader>* reader,
+ struct ArrowDeviceArrayStream* c_stream) {
+ ASSERT_OK(ExportDeviceRecordBatchReader(*reader, c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(c_stream));
+
+ ASSERT_OK_AND_ASSIGN(
+ auto got_reader,
+ ImportDeviceRecordBatchReader(c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ *reader = std::move(got_reader);
+ }
+
+ void Roundtrip(
+ std::shared_ptr<RecordBatchReader> reader,
+ std::function<void(const std::shared_ptr<RecordBatchReader>&)> check_func) {
+ ArrowDeviceArrayStream c_stream;
+
+ // NOTE: ReleaseCallback<> is not immediately usable with ArrowDeviceArayStream
+ // because get_next and get_schema need the original private_data.
+ std::weak_ptr<RecordBatchReader> weak_reader(reader);
+ ASSERT_EQ(weak_reader.use_count(), 1); // Expiration check will fail otherwise
+
+ ASSERT_OK(ExportDeviceRecordBatchReader(std::move(reader), &c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ {
+ ASSERT_OK_AND_ASSIGN(auto new_reader,
+ ImportDeviceRecordBatchReader(
+ &c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ // stream was moved
+ ASSERT_TRUE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_FALSE(weak_reader.expired());
+
+ check_func(new_reader);
+ }
+ // Stream was released when `new_reader` was destroyed
+ ASSERT_TRUE(weak_reader.expired());
+ }
+
+ void Roundtrip(std::shared_ptr<ChunkedArray> src,
+ std::function<void(const std::shared_ptr<ChunkedArray>&)> check_func) {
+ ArrowDeviceArrayStream c_stream;
+
+ // One original copy to compare the result, one copy held by the stream
+ std::weak_ptr<ChunkedArray> weak_src(src);
+ int64_t initial_use_count = weak_src.use_count();
+
+ ASSERT_OK(ExportDeviceChunkedArray(
+ std::move(src), static_cast<DeviceAllocationType>(kMyDeviceType), &c_stream));
+ ASSERT_FALSE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+ ASSERT_EQ(kMyDeviceType, c_stream.device_type);
+
+ {
+ ASSERT_OK_AND_ASSIGN(
+ auto dst,
+ ImportDeviceChunkedArray(&c_stream, TestDeviceArrayRoundtrip::DeviceMapper));
+ // Stream was moved, consumed, and released
+ ASSERT_TRUE(ArrowDeviceArrayStreamIsReleased(&c_stream));
+
+ // Stream was released by ImportDeviceChunkedArray but original copy remains
+ ASSERT_EQ(weak_src.use_count(), initial_use_count - 1);
+
+ check_func(dst);
+ }
+ }
+
+ void AssertReaderNext(const std::shared_ptr<RecordBatchReader>& reader,
+ const RecordBatch& expected) {
+ ASSERT_OK_AND_ASSIGN(auto batch, reader->Next());
+ ASSERT_NE(batch, nullptr);
+ ASSERT_EQ(static_cast<DeviceAllocationType>(kMyDeviceType), batch->device_type());
+ AssertBatchesEqual(expected, *batch);
+ }
+
+ void AssertReaderEnd(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_OK_AND_ASSIGN(auto batch, reader->Next());
+ ASSERT_EQ(batch, nullptr);
+ }
+
+ void AssertReaderClosed(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_THAT(reader->Next(),
+ Raises(StatusCode::Invalid, ::testing::HasSubstr("already been closed")));
+ }
+
+ void AssertReaderClose(const std::shared_ptr<RecordBatchReader>& reader) {
+ ASSERT_OK(reader->Close());
+ AssertReaderClosed(reader);
+ }
+};
+
+TEST_F(TestArrayDeviceStreamRoundtrip, Simple) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto orig_schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(orig_schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader, RecordBatchReader::Make(batches, orig_schema, device->device_type()));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ AssertSchemaEqual(*orig_schema, *reader->schema(), /*check_metadata=*/true);
+ AssertReaderNext(reader, *batches[0]);
+ AssertReaderNext(reader, *batches[1]);
+ AssertReaderEnd(reader);
+ AssertReaderEnd(reader);
+ AssertReaderClose(reader);
+ });
+}
+
+TEST_F(TestArrayDeviceStreamRoundtrip, CloseEarly) {
+ std::shared_ptr<Device> device = std::make_shared<MyDevice>(1);
+ auto mm = device->default_memory_manager();
+
+ ASSERT_OK_AND_ASSIGN(auto arr1,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[1, 2]")->data()));
+ ASSERT_EQ(device->device_type(), arr1->device_type());
+ ASSERT_OK_AND_ASSIGN(auto arr2,
+ ToDevice(mm, *ArrayFromJSON(int32(), "[4, 5, null]")->data()));
+ ASSERT_EQ(device->device_type(), arr2->device_type());
+ auto orig_schema = arrow::schema({field("ints", int32())});
+ auto batches = MakeBatches(orig_schema, {arr1, arr2});
+ ASSERT_OK_AND_ASSIGN(
+ auto reader, RecordBatchReader::Make(batches, orig_schema, device->device_type()));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ AssertReaderNext(reader, *batches[0]);
+ AssertReaderClose(reader);
+ });
+}
+
+TEST_F(TestArrayDeviceStreamRoundtrip, Errors) {
+ auto reader = std::make_shared<FailingRecordBatchReader>(
+ Status::Invalid("roundtrip error example"));
+
+ Roundtrip(std::move(reader), [&](const std::shared_ptr<RecordBatchReader>& reader) {
+ auto status = reader->Next().status();
+ ASSERT_RAISES(Invalid, status);
+ ASSERT_THAT(status.message(), ::testing::HasSubstr("roundtrip error example"));
Review Comment:
```suggestion
EXPECT_THAT(reader->Next(), Raises(StatusCode::Invalid, ::testing::HasSubstr("roundtrip error example")));
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1586410971
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out);
+
+/// \brief Export C++ ChunkedArray using the C device data interface format.
+///
+/// The resulting ArrowDeviceArrayStream keeps the chunked array data and buffers
+/// alive until its release callback is called by the consumer.
+///
+/// \param[in] chunked_array ChunkedArray object to export
+/// \param[in] device_type the device type the data is located on
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
Review Comment:
we don't currently track device type on a ChunkedArray since each individual chunk could potentially be on a different device. So rather than adding device type to ChunkedArray also, I'm taking it as a parameter
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541423331
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2151,53 +2188,96 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ return ExportedArrayStream<RecordBatchReader, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array), out);
+ return ExportedArrayStream<ChunkedArray, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(reader), out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(chunked_array),
+ out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <typename StreamTraits, typename ArrayTraits>
class ArrayStreamReader {
+ protected:
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper& mapper = DefaultDeviceMapper)
+ : mapper_{mapper} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
}
return status;
}
+ Result<std::shared_ptr<RecordBatch>> ImportRecordBatchInternal(
+ ArrayType* array, std::shared_ptr<Schema> schema) {
Review Comment:
sure, though it seems a bit weird having the class be able to handle both cases regardless of the template params
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541836007
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,37 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \note it is assumed that the output pointer has already be zeroed out before
+/// calling this function.
Review Comment:
I added it here
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541313200
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,37 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \note it is assumed that the output pointer has already be zeroed out before
+/// calling this function.
Review Comment:
I added this note because I noticed that the other export methods also did not, themselves, do a memset and expected that the output pointer has already been zeroed out.
> If it is something worth adding here, it would be sufficient that the structure pointed to by the output pointer is marked as released (not necessarily completely zeroed out).
That's fair, though for some interoperability situations (like the connections with Go/CGO) we ended up having to add a trampoline that calls memset to zero it out in order to be safe. But i'd be fine with just saying this should be released
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541356282
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,37 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \note it is assumed that the output pointer has already be zeroed out before
+/// calling this function.
Review Comment:
> I added this note because I noticed that the other export methods also did not, themselves, do a memset and expected that the output pointer has already been zeroed out.
It would be safer to fix that on Arrow C++'s side (and zero out the memory if that is indeed an assumption it needs to make). Requiring the consumer to zero out the struct (or even set the release callback to null) would be a breaking change and I am not sure that is what you want to do in this PR (even if it is a good idea).
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541325263
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual std::shared_ptr<Device::SyncEvent> GetSyncEvent();
Review Comment:
technically, the `SyncEvent` is held at the buffer levels currently (and should be the same pointer for each relevant buffer). A record batch which can return the sync event should be somehow tied to all of its columns (ideally the same event for all the columns)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1541388661
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
+ ///
+ /// \return null or a Device::SyncEvent
+ virtual std::shared_ptr<Device::SyncEvent> GetSyncEvent();
Review Comment:
sure thing, i'll add some comments for that
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2021581865
:warning: GitHub issue #40078 **has been automatically assigned in GitHub** to PR creator.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2059906938
Hey everyone, I've updated this with tests now. Please give it another pass when you have a chance, thanks!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1581184868
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2067,14 +2092,20 @@ class ExportedArrayStream {
ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
};
- explicit ExportedArrayStream(struct ArrowArrayStream* stream) : stream_(stream) {}
+ explicit ExportedArrayStream(StreamType* stream) : stream_(stream) {}
Status GetSchema(struct ArrowSchema* out_schema) {
return ExportStreamSchema(reader(), out_schema);
}
- Status GetNext(struct ArrowArray* out_array) {
- return ExportStreamNext(reader(), next_batch_num(), out_array);
+ Status GetNext(ArrayType* out_array) {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowArray>) {
+ return ExportStreamNext(reader(), next_batch_num(), out_array);
+ } else if constexpr (std::is_same_v<T, ChunkedArray>) {
+ return ExportDeviceStreamNext(reader(), next_batch_num(), nullptr, out_array);
Review Comment:
okay. I'll leave it out for now
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2079653168
> When exported over the C Data interface it is a requirement, no? (Or at least, a better option than exporting a device array that is potentially yet synchronized with a nullptr sync event?)
Yea, the sync event is at the top level of the ArrowDeviceArray struct to represent that the data for a single ArrowArray should always be tied to a single sync event, or have no synchronization required. Thus any joining / handling of multiple sync events needs to be managed by the producer so it can generate a *single* sync event if synchronization is needed.
As such, this code will use that sync event when constructing imported record batches so that the event can be propagated. But once you are manually constructing record batches from existing Arrow Arrays, all bets are off, which is why we just take a sync event as an argument when constructing the record batches.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583821785
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2340,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper mapper_;
};
-class ArrayStreamBatchReader : public RecordBatchReader, public ArrayStreamReader {
+template <typename StreamTraits, typename ArrayTraits>
+class ArrayStreamBatchReader : public RecordBatchReader,
+ public ArrayStreamReader<StreamTraits, ArrayTraits> {
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamBatchReader(struct ArrowArrayStream* stream)
- : ArrayStreamReader(stream) {}
+ explicit ArrayStreamBatchReader(
+ StreamType* stream, const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper)
+ : ArrayStreamReader<StreamTraits, ArrayTraits>(stream, mapper) {}
Status Init() {
- ARROW_ASSIGN_OR_RAISE(schema_, ReadSchema());
+ ARROW_ASSIGN_OR_RAISE(schema_, this->ReadSchema());
return Status::OK();
}
std::shared_ptr<Schema> schema() const override { return schema_; }
Status ReadNext(std::shared_ptr<RecordBatch>* batch) override {
- ARROW_RETURN_NOT_OK(CheckNotReleased());
+ ARROW_RETURN_NOT_OK(this->CheckNotReleased());
- struct ArrowArray c_array;
- ARROW_RETURN_NOT_OK(ReadNextArrayInternal(&c_array));
+ ArrayType c_array;
+ ARROW_RETURN_NOT_OK(this->ReadNextArrayInternal(&c_array));
- if (ArrowArrayIsReleased(&c_array)) {
+ if (ArrayTraits::IsReleasedFunc(&c_array)) {
// End of stream
batch->reset();
return Status::OK();
} else {
- return ImportRecordBatch(&c_array, schema_).Value(batch);
+ return this->ImportRecordBatchInternal(&c_array, schema_).Value(batch);
}
}
Status Close() override {
- ReleaseStream();
+ this->ReleaseStream();
return Status::OK();
}
+ DeviceAllocationType device_type() const override { return this->get_device_type(); }
+
private:
std::shared_ptr<Schema> schema_;
};
-class ArrayStreamArrayReader : public ArrayStreamReader {
+template <typename StreamTraits, typename ArrayTraits>
+class ArrayStreamArrayReader : public ArrayStreamReader<StreamTraits, ArrayTraits> {
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamArrayReader(struct ArrowArrayStream* stream)
- : ArrayStreamReader(stream) {}
+ explicit ArrayStreamArrayReader(
+ StreamType* stream, const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper)
Review Comment:
In the vast majority of cases, using the `DefaultDeviceMemoryMapper` is the correct value to pass on any callsite. I also don't want to break any existing callers / users of this function.
It's only when you start introducing non-CPU data that you'd want to pass something else.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583825043
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -167,30 +188,41 @@ class SimpleRecordBatch : public RecordBatch {
return RecordBatch::Validate();
}
+ std::shared_ptr<Device::SyncEvent> GetSyncEvent() const override { return sync_event_; }
Review Comment:
In theory I can't think of one, but that doesn't mean there *can't* be a situation like that.
But that's still a good point. I'd rather default to having it return `const ...&` for now, and update this in the future if necessary.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1586416413
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -623,6 +667,16 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return null_sync_event_;
+}
+
+DeviceAllocationType RecordBatch::device_type() const {
+ return DeviceAllocationType::kCPU;
Review Comment:
because `device_type_` was added to the `SimpleRecordBatch` derived class from RecordBatch, it's not in the base class. I didn't want to force implementations to do it that way if they potentially could pull the device_type from elsewhere.
So the default implementation just returns CPU, but @bkietz brought up a good point https://github.com/apache/arrow/pull/40807#discussion_r1586298099 where it might make sense to just make this pure virtual instead of having a default impl
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1611642632
##########
cpp/src/arrow/array/data.h:
##########
@@ -358,6 +369,16 @@ struct ARROW_EXPORT ArrayData {
/// \see GetNullCount
int64_t ComputeLogicalNullCount() const;
+ /// \brief Returns the device_type of the underlying buffers and children
Review Comment:
Nit, but we tend to use infinitives in docstring ("return", not "returns")
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553814507
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
+ /// on a device, then if synchronization is needed before accessing the
+ /// data the returned sync event will allow for it.
Review Comment:
It was either that or introducing a subclass, yes.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1553831336
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
So the suggestion is we *should* add a `device_type` member to the `RecordBatch` class and have it be provided?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "jorisvandenbossche (via GitHub)" <gi...@apache.org>.
jorisvandenbossche commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1550982861
##########
cpp/src/arrow/record_batch.h:
##########
@@ -300,6 +311,11 @@ class ARROW_EXPORT RecordBatchReader {
/// \brief finalize reader
virtual Status Close() { return Status::OK(); }
+ /// \brief Get the device type for record batches this reader produces
+ ///
+ /// default implementation is to return ARROW_DEVICE_CPU
+ virtual DeviceAllocationType device_type() const { return DeviceAllocationType::kCPU; }
Review Comment:
This is going to be wrong if we import a RecordBatchReader from a non-CPU ArrowDeviceArrayStream? Or is that right now not yet possible?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551987917
##########
cpp/src/arrow/c/helpers.h:
##########
@@ -115,6 +147,14 @@ inline void ArrowArrayStreamMove(struct ArrowArrayStream* src,
ArrowArrayStreamMarkReleased(src);
}
+inline void ArrowDeviceArrayStreamMove(struct ArrowDeviceArrayStream* src,
+ struct ArrowDeviceArrayStream* dest) {
+ assert(dest != src);
Review Comment:
This is a C file, so any C++ inclusion is out of luck. As for `ARROW_C_ASSERT`, good question: it seems meant to also trigger in non-debug builds.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551964200
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
+ /// should return null (the default impl). If the data for this batch is
Review Comment:
```suggestion
/// will return null. If the data for this batch is
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1551963277
##########
cpp/src/arrow/record_batch.h:
##########
@@ -254,6 +255,16 @@ class ARROW_EXPORT RecordBatch {
/// \return Status
virtual Status ValidateFull() const;
+ /// \brief Return a top-level sync event object for this record batch
+ ///
+ /// If all of the data for this record batch is in host memory, then this
Review Comment:
```suggestion
/// If all of the data for this record batch is in CPU memory, then this
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1543526576
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2151,53 +2188,96 @@ class ExportedArrayStream {
int64_t next_batch_num() { return private_data()->batch_num_++; }
- struct ArrowArrayStream* stream_;
+ StreamType* stream_;
};
} // namespace
Status ExportRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<RecordBatchReader>::Make(std::move(reader), out);
+ return ExportedArrayStream<RecordBatchReader, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(reader), out);
}
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out) {
- return ExportedArrayStream<ChunkedArray>::Make(std::move(chunked_array), out);
+ return ExportedArrayStream<ChunkedArray, struct ArrowArrayStream,
+ struct ArrowArray>::Make(std::move(chunked_array), out);
+}
+
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(reader->device_type());
+ return ExportedArrayStream<RecordBatchReader, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(reader), out);
+}
+
+Status ExportDeviceChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
+ DeviceAllocationType device_type,
+ struct ArrowDeviceArrayStream* out) {
+ out->device_type = static_cast<ArrowDeviceType>(device_type);
+ return ExportedArrayStream<ChunkedArray, struct ArrowDeviceArrayStream,
+ struct ArrowDeviceArray>::Make(std::move(chunked_array),
+ out);
}
//////////////////////////////////////////////////////////////////////////
// C stream import
namespace {
+template <typename StreamTraits, typename ArrayTraits>
class ArrayStreamReader {
+ protected:
+ using StreamType = typename StreamTraits::CType;
+ using ArrayType = typename ArrayTraits::CType;
+
public:
- explicit ArrayStreamReader(struct ArrowArrayStream* stream) {
- ArrowArrayStreamMove(stream, &stream_);
- DCHECK(!ArrowArrayStreamIsReleased(&stream_));
+ explicit ArrayStreamReader(StreamType* stream,
+ const DeviceMemoryMapper& mapper = DefaultDeviceMapper)
+ : mapper_{mapper} {
+ StreamTraits::MoveFunc(stream, &stream_);
+ DCHECK(!StreamTraits::IsReleasedFunc(&stream_));
}
~ArrayStreamReader() { ReleaseStream(); }
void ReleaseStream() {
- if (!ArrowArrayStreamIsReleased(&stream_)) {
- ArrowArrayStreamRelease(&stream_);
- }
- DCHECK(ArrowArrayStreamIsReleased(&stream_));
+ // all our trait release funcs check IsReleased so we don't
+ // need to repeat it here
+ StreamTraits::ReleaseFunc(&stream_);
+ DCHECK(StreamTraits::IsReleasedFunc(&stream_));
}
protected:
- Status ReadNextArrayInternal(struct ArrowArray* array) {
- ArrowArrayMarkReleased(array);
+ Status ReadNextArrayInternal(ArrayType* array) {
+ ArrayTraits::MarkReleased(array);
Status status = StatusFromCError(stream_.get_next(&stream_, array));
- if (!status.ok() && !ArrowArrayIsReleased(array)) {
- ArrowArrayRelease(array);
+ if (!status.ok() && !ArrayTraits::IsReleasedFunc(array)) {
+ ArrayTraits::ReleaseFunc(array);
}
return status;
}
+ Result<std::shared_ptr<RecordBatch>> ImportRecordBatchInternal(
+ ArrayType* array, std::shared_ptr<Schema> schema) {
Review Comment:
Your call! That was feedback I got the last time I tried to use `if constexpr` in this section of code 🙂
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1544856868
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
Review Comment:
This is why I put the `GetSyncEvent` on the `RecordBatch` object since there could potentially be a different sync event for each record batch returned by the reader. Right now it currently is hardcoded to return null, but the `GetSyncEvent` function is virtual, so someone could potentially have a `RecordBatchReader` that returns a class which inherits from `RecordBatch` that implements `GetSyncEvent` to return the correct event object etc.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1556023666
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
I don't get why it being a `std::function` makes it ok. A `std::function` can own arbitrary state (for example a closure's captured variables).
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1560011135
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2250,101 +2336,126 @@ class ArrayStreamReader {
return {code, last_error ? std::string(last_error) : ""};
}
+ DeviceAllocationType get_device_type() const {
+ if constexpr (std::is_same_v<ArrayType, struct ArrowDeviceArray>) {
+ return static_cast<DeviceAllocationType>(stream_.device_type);
+ } else {
+ return DeviceAllocationType::kCPU;
+ }
+ }
+
private:
- mutable struct ArrowArrayStream stream_;
+ mutable StreamType stream_;
+ const DeviceMemoryMapper& mapper_;
Review Comment:
After doing a bit of research, not only is @pitrou correct that the current version isn't necessarily safe, but in most cases it's preferable and more performant to pass an `std::function` by value and use `std::move` instead. So I'll update this.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1561601329
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -321,6 +321,31 @@ ARROW_EXPORT
Status ExportChunkedArray(std::shared_ptr<ChunkedArray> chunked_array,
struct ArrowArrayStream* out);
+/// \brief Export C++ RecordBatchReader using the C device stream interface
+///
+/// The resulting ArrowDeviceArrayStream struct keeps the record batch reader
+/// alive until its release callback is called by the consumer. The device
+/// type is determined by calling device_type() on the RecordBatchReader.
+///
+/// \param[in] reader RecordBatchReader object to export
+/// \param[out] out C struct to export the stream to
+ARROW_EXPORT
+Status ExportDeviceRecordBatchReader(std::shared_ptr<RecordBatchReader> reader,
Review Comment:
I've made a bunch of changes propagating the sync event and the device type now, so `GetSyncEvent` now does something :)
Thoughts?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1581172015
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,36 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ int type = 0;
Review Comment:
For the `type == 0`, we explicitly don't have an "UNASSIGNED" value in the enum or the macros because we don't want that to be used or seen as a potentially valid value for the C struct or the enum. Notice that there's no situation in which we actually return the 0 here.
I'd be fine with defining it *internally* to the library as a sentinel constant, but I don't want it to ever end up externally viewable and was trying to keep it that way.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2080046370
Anyone have any further comments here?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2079641888
> essentially, it's not prohibitively expensive, we just don't want that to be a requirement
When exported over the C Data interface it *is* a requirement, no? (Or at least, a better option than exporting a device array that is potentially yet synchronized with a `nullptr` sync event?)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1581194144
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,36 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ int type = 0;
Review Comment:
A comment is OK too if this is the only place it's used 🙂
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583822347
##########
cpp/src/arrow/array/data.cc:
##########
@@ -224,6 +224,41 @@ int64_t ArrayData::ComputeLogicalNullCount() const {
return ArraySpan(*this).ComputeLogicalNullCount();
}
+DeviceAllocationType ArrayData::device_type() const {
+ // we're using 0 as a sentinel value for NOT YET ASSIGNED
+ // there is explicitly no constant DeviceAllocationType to represent
+ // the "UNASSIGNED" case as it is invalid for data to not have an
+ // assigned device type. If it's still 0 at the end, then we return
+ // CPU as the allocation device type
+ int type = 0;
+ for (const auto& buf : buffers) {
+ if (!buf) continue;
+ if (type == 0) {
+ type = static_cast<int>(buf->device_type());
Review Comment:
I personally am opposed to adding a `kUNDEFINED` to `DeviceAllocationType` because there is no valid case where `kUNDEFINED` would be an acceptable value to be returned. So I don't want to make every consumer have to check for it
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583832510
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -664,7 +667,9 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
-const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const { return nullptr; }
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return nullptr;
Review Comment:
For this case, you should have a `null_sync_event_` member variable because it's illegal to return a reference to a `std::shared_ptr<Device::SyncEvent>` created on the stack implicitly from `nullptr`.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583834010
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -664,7 +667,9 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
-const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const { return nullptr; }
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return nullptr;
Review Comment:
Bugs would manifest in usage like this:
```cpp
auto& sync_event = batch.GetSyncEvent();
// pass the & to something that receives a `shared_ptr<...>&`...
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #40807:
URL: https://github.com/apache/arrow/pull/40807#discussion_r1583844222
##########
cpp/src/arrow/record_batch.cc:
##########
@@ -664,7 +667,9 @@ Status RecordBatch::ValidateFull() const {
return ValidateBatch(*this, /*full_validation=*/true);
}
-const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const { return nullptr; }
+const std::shared_ptr<Device::SyncEvent>& RecordBatch::GetSyncEvent() const {
+ return nullptr;
Review Comment:
that's probably why i didn't have it returning a const reference in the first place haha. I'll make something static that I can return const refs to
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-40078: [C++] Import/Export ArrowDeviceArrayStream [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #40807:
URL: https://github.com/apache/arrow/pull/40807#issuecomment-2087247081
I was able to narrow the failure down to attempting to read a parquet file containing extension types with registered extensions causing the segfault. i.e.: either this line https://github.com/zeroshade/arrow/blob/main/python/pyarrow/tests/test_extension_type.py#L1051 or the line after it calling validate.
This is almost certainly *not* being caused by this change, which shouldn't touch anything around that code, and @paleolimbot said he was able to replicate the failure with his own debug python build that is separate from this change. So I don't think it's related or relevant to this PR.
@jorisvandenbossche @AlenkaF would either of you be able to take a look and see if you can figure out why there's a segfault happening at that line in the test and if I should file a separate issue to track that?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org