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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
