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/10/22 01:09:48 UTC
[PR] GH-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
zeroshade opened a new pull request, #44495:
URL: https://github.com/apache/arrow/pull/44495
<!--
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
Adding functions to `bridge.h`/`bridge.cc` to implement helpers for managing the Async C Data interfaces
<!--
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?
Two functions added to bridge.h:
1. `CreateAsyncDeviceStreamHandler` populates a `ArrowAsyncDeviceStreamHandler` and an `Executor` to provide a future that resolves to an `AsyncRecordBatchGenerator` to produce record batches as they are pushed asynchronously. The `ArrowAsyncDeviceStreamHandler` can then be passed to any asynchronous producer.
2. `ExportAsyncRecordBatchReader` takes a record batch generator and a schema, along with an `ArrowAsyncDeviceStreamHandler` to use for calling the callbacks to push data as it is available from the generator.
<!--
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?
Unit tests are added (currently only one test, more tests to be added)
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
4. 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
<!--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1811601567
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
+ // code is an errno-compatible error code
+ int code() const { return code_; }
+ // returns any metadata that was returned with the error, likely in a
+ // key-value format similar to ArrowSchema metadata
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchGenerator to provide an interface for the consumer to retrieve data as
+/// it is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
Review Comment:
This directly maps to the proposal though; you can only request a number of batches and not a number of bytes. Maybe we should discuss it 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1813013801
##########
cpp/src/arrow/c/abi.h:
##########
@@ -228,6 +228,214 @@ struct ArrowDeviceArrayStream {
#endif // ARROW_C_DEVICE_STREAM_INTERFACE
+#ifndef ARROW_C_ASYNC_STREAM_INTERFACE
+# define ARROW_C_ASYNC_STREAM_INTERFACE
+
+// ArrowAsyncTask represents available data from a producer that was passed to
+// an invocation of `on_next_task` on the ArrowAsyncDeviceStreamHandler.
+//
+// The reason for this Task approach instead of the Async interface returning
+// the Array directly is to allow for more complex thread handling and reducing
+// context switching and data transfers between CPU cores (e.g. from one L1/L2
+// cache to another) if desired.
+//
+// For example, the `on_next_task` callback can be called when data is ready, while
+// the producer puts potential "decoding" logic in the `ArrowAsyncTask` object. This
+// allows for the producer to manage the I/O on one thread which calls `on_next_task`
+// and the consumer can determine when the decoding (producer logic in the `extract_data`
+// callback of the task) occurs and on which thread, to avoid a CPU core transfer
+// (data staying in the L2 cache).
+struct ArrowAsyncTask {
+ // This callback should populate the ArrowDeviceArray associated with this task.
+ // The order of ArrowAsyncTasks provided by the producer enables a consumer to
+ // ensure the order of data to process.
+ //
+ // This function is expected to be synchronous, but should not perform any blocking
+ // I/O. Ideally it should be as cheap as possible so as to not tie up the consumer
+ // thread unnecessarily.
+ //
+ // Returns: 0 if successful, errno-compatible error otherwise.
+ //
+ // If a non-0 value is returned then it should be followed by a call to `on_error`
+ // on the appropriate ArrowAsyncDeviceStreamHandler. This is because it's highly
+ // likely that whatever is calling this function may be entirely disconnected from
+ // the current control flow. Indicating an error here with a non-zero return allows
+ // the current flow to be aware of the error occurring, while still allowing any
+ // logging or error handling to still be centralized in the `on_error` callback of
+ // the original Async handler.
+ //
+ // Rather than a release callback, any required cleanup should be performed as part
+ // of the invocation of `extract_data`. Ownership of the Array is passed to the consumer
+ // calling this, and so it must be released separately.
+ //
+ // It is only valid to call this method exactly once.
+ int (*extract_data)(struct ArrowAsyncTask* self, struct ArrowDeviceArray* out);
Review Comment:
Got it! It makes sense to propagate it as early as possible like you've done already 🙂
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "bkietz (via GitHub)" <gi...@apache.org>.
bkietz commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1833031853
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
+ // code is an errno-compatible error code
+ int code() const { return code_; }
+ // returns any metadata that was returned with the error, likely in a
+ // key-value format similar to ArrowSchema metadata
+ const std::string& ErrorMetadata() const { return metadata_; }
Review Comment:
```suggestion
const std::string& ErrorMetadataString() const { return metadata_; }
std::shared_ptr<KeyValueMetadata> ErrorMetadata() const;
```
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "bkietz (via GitHub)" <gi...@apache.org>.
bkietz commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1837058002
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,347 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct TaskWithMetadata {
+ ArrowAsyncTask task_;
+ std::shared_ptr<KeyValueMetadata> metadata_;
+ };
+
+ struct State {
+ State(uint64_t queue_size, DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ Result<RecordBatchWithMetadata> next() {
+ TaskWithMetadata task;
+ {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cv_.wait(lock,
+ [&] { return !error_.ok() || !batches_.empty() || end_of_stream_; });
+ if (!error_.ok()) {
+ return error_;
+ }
+
+ if (batches_.empty() && end_of_stream_) {
+ return IterationEnd<RecordBatchWithMetadata>();
+ }
+
+ task = std::move(batches_.front());
+ batches_.pop();
+ }
+
+ producer_->request(producer_, 1);
+ ArrowDeviceArray out;
+ if (task.task_.extract_data(&task.task_, &out) != 0) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cv_.wait(lock, [&] { return !error_.ok(); });
+ return error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(auto batch, ImportDeviceRecordBatch(&out, schema_, mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.metadata_)};
+ }
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<TaskWithMetadata> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
Review Comment:
This one must've missed the batch
```suggestion
AsyncRecordBatchIterator(uint64_t queue_size, DeviceMemoryMapper mapper)
```
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1837116617
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,347 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct TaskWithMetadata {
+ ArrowAsyncTask task_;
+ std::shared_ptr<KeyValueMetadata> metadata_;
+ };
+
+ struct State {
+ State(uint64_t queue_size, DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ Result<RecordBatchWithMetadata> next() {
+ TaskWithMetadata task;
+ {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cv_.wait(lock,
+ [&] { return !error_.ok() || !batches_.empty() || end_of_stream_; });
+ if (!error_.ok()) {
+ return error_;
+ }
+
+ if (batches_.empty() && end_of_stream_) {
+ return IterationEnd<RecordBatchWithMetadata>();
+ }
+
+ task = std::move(batches_.front());
+ batches_.pop();
+ }
+
+ producer_->request(producer_, 1);
+ ArrowDeviceArray out;
+ if (task.task_.extract_data(&task.task_, &out) != 0) {
+ std::unique_lock<std::mutex> lock(mutex_);
+ cv_.wait(lock, [&] { return !error_.ok(); });
+ return error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(auto batch, ImportDeviceRecordBatch(&out, schema_, mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.metadata_)};
+ }
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<TaskWithMetadata> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
Review Comment:
I missed one? Darn 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809754998
##########
cpp/src/arrow/record_batch.h:
##########
@@ -17,6 +17,7 @@
#pragma once
+#include <chrono>
Review Comment:
remnant of a previous version, i'll remove 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809769176
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
Review Comment:
for the iterator yea, that was my intent. this is only exposed via `MakeBackgroundGenerator` so that kinda enforces 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade merged PR #44495:
URL: https://github.com/apache/arrow/pull/44495
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809756941
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
+ return !state_->error_.ok() || state_->cancelled_ ||
+ state_->pending_requests_ > 0;
+ });
+ }
+
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->cancelled_) {
+ return Status::Cancelled("Consumer cancelled");
+ }
+
+ if (state_->pending_requests_ > 0) {
+ state_->pending_requests_--;
+ lock.unlock();
+
+ ArrowAsyncTask task;
+ task.private_data = new PrivateTaskData{state_, record};
+ task.extract_data = AsyncProducer::extract_data;
+
+ if (int status = handler_->on_next_task(handler_, &task, nullptr) != 0) {
+ delete reinterpret_cast<PrivateTaskData*>(task.private_data);
+ return Status::UnknownError("Received error from handler::on_next_task ", status);
+ }
+ }
+
+ return Status::OK();
+ }
+
+ static void request(struct ArrowAsyncProducer* producer, int64_t n) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (self->cancelled_) {
Review Comment:
Do we also want to check `error_`? (Does it make sense to fold `error_` and `cancelled_` together?)
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
Review Comment:
Are we supposed to have a blocking call in consuming an AsyncGenerator?
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1811679662
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
+ // code is an errno-compatible error code
+ int code() const { return code_; }
+ // returns any metadata that was returned with the error, likely in a
+ // key-value format similar to ArrowSchema metadata
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchGenerator to provide an interface for the consumer to retrieve data as
+/// it is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
Review Comment:
Got it! I don't think that complexity is needed then...often that type of option can be set outside this framework anyway (e.g., in the postgres driver one could set the batch size and manipulate this number to approximate what I suggested).
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "conbench-apache-arrow[bot] (via GitHub)" <gi...@apache.org>.
conbench-apache-arrow[bot] commented on PR #44495:
URL: https://github.com/apache/arrow/pull/44495#issuecomment-2469699165
After merging your PR, Conbench analyzed the 3 benchmarking runs that have been run so far on merge-commit aab7d81aeec1e8f106bdda953cdeb00f7f78b355.
There were no benchmark performance regressions. 🎉
The [full Conbench report](https://github.com/apache/arrow/runs/32844522626) has more details. It also includes information about 8 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809754786
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ std::string ToString() const override { return message_; }
+ int code() const { return code_; }
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchReader to provide an interface for the consumer to retrieve data as it
+/// is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
+/// \param[in] mapper mapping from device type and ID to memory manager
+/// \return Future that resolves to either an error or AsyncRecordBatchGenerator once a
+/// schema is available or an error is received.
+ARROW_EXPORT
+Future<AsyncRecordBatchGenerator> CreateAsyncDeviceStreamHandler(
+ struct ArrowAsyncDeviceStreamHandler* handler, internal::Executor* executor,
+ uint64_t queue_size = 5, const DeviceMemoryMapper mapper = DefaultDeviceMemoryMapper);
Review Comment:
So, I originally did this, but ended up with a segfault as the `mapper` didn't stick around for the async side of things. So I needed it to not be a reference and I just `std::move` it around.
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809761188
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ std::string ToString() const override { return message_; }
+ int code() const { return code_; }
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchReader to provide an interface for the consumer to retrieve data as it
+/// is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
+/// \param[in] mapper mapping from device type and ID to memory manager
+/// \return Future that resolves to either an error or AsyncRecordBatchGenerator once a
+/// schema is available or an error is received.
+ARROW_EXPORT
+Future<AsyncRecordBatchGenerator> CreateAsyncDeviceStreamHandler(
+ struct ArrowAsyncDeviceStreamHandler* handler, internal::Executor* executor,
+ uint64_t queue_size = 5, const DeviceMemoryMapper mapper = DefaultDeviceMemoryMapper);
Review Comment:
ah interesting...In that case maybe just make it `T` instead of `const T`
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809769877
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
Review Comment:
I suppose the way it "should" work is that we never call into the AsyncGenerator unless we have a pending request? But we'd then need to handle more of the control flow here ourselves
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809765677
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
Review Comment:
I honestly couldn't figure out a better way around this to manage the backpressure from the calls to `request`. I'm open to suggestions for a better route here other than simply saying this should not be run on the main thread *shrug*
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1811426098
##########
cpp/src/arrow/c/abi.h:
##########
@@ -228,6 +228,214 @@ struct ArrowDeviceArrayStream {
#endif // ARROW_C_DEVICE_STREAM_INTERFACE
+#ifndef ARROW_C_ASYNC_STREAM_INTERFACE
+# define ARROW_C_ASYNC_STREAM_INTERFACE
+
+// ArrowAsyncTask represents available data from a producer that was passed to
+// an invocation of `on_next_task` on the ArrowAsyncDeviceStreamHandler.
+//
+// The reason for this Task approach instead of the Async interface returning
+// the Array directly is to allow for more complex thread handling and reducing
+// context switching and data transfers between CPU cores (e.g. from one L1/L2
+// cache to another) if desired.
+//
+// For example, the `on_next_task` callback can be called when data is ready, while
+// the producer puts potential "decoding" logic in the `ArrowAsyncTask` object. This
+// allows for the producer to manage the I/O on one thread which calls `on_next_task`
+// and the consumer can determine when the decoding (producer logic in the `extract_data`
+// callback of the task) occurs and on which thread, to avoid a CPU core transfer
+// (data staying in the L2 cache).
+struct ArrowAsyncTask {
+ // This callback should populate the ArrowDeviceArray associated with this task.
+ // The order of ArrowAsyncTasks provided by the producer enables a consumer to
+ // ensure the order of data to process.
+ //
+ // This function is expected to be synchronous, but should not perform any blocking
+ // I/O. Ideally it should be as cheap as possible so as to not tie up the consumer
+ // thread unnecessarily.
+ //
+ // Returns: 0 if successful, errno-compatible error otherwise.
+ //
+ // If a non-0 value is returned then it should be followed by a call to `on_error`
+ // on the appropriate ArrowAsyncDeviceStreamHandler. This is because it's highly
+ // likely that whatever is calling this function may be entirely disconnected from
+ // the current control flow. Indicating an error here with a non-zero return allows
+ // the current flow to be aware of the error occurring, while still allowing any
+ // logging or error handling to still be centralized in the `on_error` callback of
+ // the original Async handler.
+ //
+ // Rather than a release callback, any required cleanup should be performed as part
+ // of the invocation of `extract_data`. Ownership of the Array is passed to the consumer
+ // calling this, and so it must be released separately.
+ //
+ // It is only valid to call this method exactly once.
+ int (*extract_data)(struct ArrowAsyncTask* self, struct ArrowDeviceArray* out);
Review Comment:
An issue for the format PR, but was it discussed to allow `const char* metadata` here, too for batch-specific metadata? Specifically, an IPC scanner could propagate
https://github.com/apache/arrow/blob/bcb4653c6387a2b22df52a3bbc91317607abdccc/format/Message.fbs#L154
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
+ // code is an errno-compatible error code
+ int code() const { return code_; }
+ // returns any metadata that was returned with the error, likely in a
+ // key-value format similar to ArrowSchema metadata
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchGenerator to provide an interface for the consumer to retrieve data as
+/// it is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
Review Comment:
Not important for proof-of-concept, but it might be a more meaningful parameter as a number of bytes.
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809766336
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
+ return !state_->error_.ok() || state_->cancelled_ ||
+ state_->pending_requests_ > 0;
+ });
+ }
+
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->cancelled_) {
+ return Status::Cancelled("Consumer cancelled");
+ }
+
+ if (state_->pending_requests_ > 0) {
+ state_->pending_requests_--;
+ lock.unlock();
+
+ ArrowAsyncTask task;
+ task.private_data = new PrivateTaskData{state_, record};
+ task.extract_data = AsyncProducer::extract_data;
+
+ if (int status = handler_->on_next_task(handler_, &task, nullptr) != 0) {
+ delete reinterpret_cast<PrivateTaskData*>(task.private_data);
+ return Status::UnknownError("Received error from handler::on_next_task ", status);
+ }
+ }
+
+ return Status::OK();
+ }
+
+ static void request(struct ArrowAsyncProducer* producer, int64_t n) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (self->cancelled_) {
Review Comment:
That's a good point, since cancelled just results in us pushing out a `Cancelled` status, i can fold them together.
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809767566
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
Review Comment:
I guess blocking calls are OK due to MakeBackgroundGenerator here?
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
Review Comment:
nit: I don't think you need this check? You can always wait on the CV with the condition
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1813430432
##########
cpp/src/arrow/c/bridge_test.cc:
##########
Review Comment:
added tests to propagate errors
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809747368
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ std::string ToString() const override { return message_; }
+ int code() const { return code_; }
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchReader to provide an interface for the consumer to retrieve data as it
+/// is pushed to the handler.
+///
+/// \param[in,out] handler C struct to be populated
+/// \param[in] executor the executor to use for waiting and populating record batches
+/// \param[in] queue_size initial number of record batches to request for queueing
+/// \param[in] mapper mapping from device type and ID to memory manager
+/// \return Future that resolves to either an error or AsyncRecordBatchGenerator once a
+/// schema is available or an error is received.
+ARROW_EXPORT
+Future<AsyncRecordBatchGenerator> CreateAsyncDeviceStreamHandler(
+ struct ArrowAsyncDeviceStreamHandler* handler, internal::Executor* executor,
+ uint64_t queue_size = 5, const DeviceMemoryMapper mapper = DefaultDeviceMemoryMapper);
Review Comment:
```suggestion
uint64_t queue_size = 5, const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper);
```
Most functions seem to declare it this way
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
Review Comment:
nit: docstring? (at least on the fields; presumably code is supposed to be an errno?)
##########
cpp/src/arrow/c/bridge_test.cc:
##########
Review Comment:
Can we add a test that also propagates errors? (Both: before the schema, and after the schema)
##########
cpp/src/arrow/record_batch.h:
##########
@@ -17,6 +17,7 @@
#pragma once
+#include <chrono>
Review Comment:
Did you mean to add this here?
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ std::string ToString() const override { return message_; }
+ int code() const { return code_; }
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
Review Comment:
This really needs to not be internal at some point
##########
cpp/src/arrow/c/bridge_test.cc:
##########
@@ -5311,4 +5313,71 @@ TEST_F(TestArrayDeviceStreamRoundtrip, ChunkedArrayRoundtripEmpty) {
});
}
+class TestAsyncDeviceArrayStreamRoundTrip : 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);
+ }
+};
+
+TEST_F(TestAsyncDeviceArrayStreamRoundTrip, 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});
+
+ struct ArrowAsyncDeviceStreamHandler handler;
+ auto fut_gen = CreateAsyncDeviceStreamHandler(&handler, internal::GetCpuThreadPool(), 1,
+ TestDeviceArrayRoundtrip::DeviceMapper);
+ ASSERT_FALSE(fut_gen.is_finished());
+
+ ASSERT_OK_AND_ASSIGN(auto fut, internal::GetCpuThreadPool()->Submit([&]() {
+ return ExportAsyncRecordBatchReader(orig_schema, MakeVectorGenerator(batches),
+ device->device_type(), &handler);
+ }));
+
+ ASSERT_OK_AND_ASSIGN(auto generator, fut_gen.result());
+ AssertSchemaEqual(*orig_schema, *generator.schema);
+
+ auto collect_fut = CollectAsyncGenerator(generator.generator);
+ ASSERT_OK_AND_ASSIGN(auto results, collect_fut.result());
+ ASSERT_OK(fut.status());
Review Comment:
nit but I believe we have specific macros for futures that also make sure to time out (ASSERT_FINISHES_OK etc)
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,75 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ std::string ToString() const override { return message_; }
+ int code() const { return code_; }
+ const std::string& ErrorMetadata() const { return metadata_; }
+
+ private:
+ int code_{0};
+ std::string message_;
+ std::string metadata_;
+};
+
+struct AsyncRecordBatchGenerator {
+ std::shared_ptr<Schema> schema;
+ DeviceAllocationType device_type;
+ AsyncGenerator<RecordBatchWithMetadata> generator;
+};
+
+namespace internal {
+class Executor;
+}
+
+/// \brief Create an AsyncRecordBatchReader and populate a corresponding handler to pass
+/// to a producer
+///
+/// The ArrowAsyncDeviceStreamHandler struct is intended to have its callbacks populated
+/// and then be passed to a producer to call the appropriate callbacks when data is ready.
+/// This inverts the traditional flow of control, and so we construct a corresponding
+/// AsyncRecordBatchReader to provide an interface for the consumer to retrieve data as it
Review Comment:
```suggestion
/// AsyncRecordBatchGenerator to provide an interface for the consumer to retrieve data as it
```
?
##########
cpp/src/arrow/c/bridge_test.cc:
##########
@@ -5311,4 +5313,71 @@ TEST_F(TestArrayDeviceStreamRoundtrip, ChunkedArrayRoundtripEmpty) {
});
}
+class TestAsyncDeviceArrayStreamRoundTrip : 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);
+ }
+};
+
+TEST_F(TestAsyncDeviceArrayStreamRoundTrip, 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});
+
+ struct ArrowAsyncDeviceStreamHandler handler;
+ auto fut_gen = CreateAsyncDeviceStreamHandler(&handler, internal::GetCpuThreadPool(), 1,
+ TestDeviceArrayRoundtrip::DeviceMapper);
+ ASSERT_FALSE(fut_gen.is_finished());
+
+ ASSERT_OK_AND_ASSIGN(auto fut, internal::GetCpuThreadPool()->Submit([&]() {
+ return ExportAsyncRecordBatchReader(orig_schema, MakeVectorGenerator(batches),
+ device->device_type(), &handler);
+ }));
+
+ ASSERT_OK_AND_ASSIGN(auto generator, fut_gen.result());
+ AssertSchemaEqual(*orig_schema, *generator.schema);
Review Comment:
```suggestion
ASSERT_NO_FATAL_FAILURE(AssertSchemaEqual(*orig_schema, *generator.schema));
```
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1809769125
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,333 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->error_.ok()) {
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ }
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema,
+ const char* additional_metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message, metadata),
+ std::string(message));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ bool cancelled_{false};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema, nullptr) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
Review Comment:
I guess this works 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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #44495:
URL: https://github.com/apache/arrow/pull/44495#issuecomment-2468686623
@bkietz I addressed your suggestions!
@paleolimbot I've updated the doc briefs to include `EXPERIMENTAL: `
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1812975249
##########
cpp/src/arrow/c/abi.h:
##########
@@ -228,6 +228,214 @@ struct ArrowDeviceArrayStream {
#endif // ARROW_C_DEVICE_STREAM_INTERFACE
+#ifndef ARROW_C_ASYNC_STREAM_INTERFACE
+# define ARROW_C_ASYNC_STREAM_INTERFACE
+
+// ArrowAsyncTask represents available data from a producer that was passed to
+// an invocation of `on_next_task` on the ArrowAsyncDeviceStreamHandler.
+//
+// The reason for this Task approach instead of the Async interface returning
+// the Array directly is to allow for more complex thread handling and reducing
+// context switching and data transfers between CPU cores (e.g. from one L1/L2
+// cache to another) if desired.
+//
+// For example, the `on_next_task` callback can be called when data is ready, while
+// the producer puts potential "decoding" logic in the `ArrowAsyncTask` object. This
+// allows for the producer to manage the I/O on one thread which calls `on_next_task`
+// and the consumer can determine when the decoding (producer logic in the `extract_data`
+// callback of the task) occurs and on which thread, to avoid a CPU core transfer
+// (data staying in the L2 cache).
+struct ArrowAsyncTask {
+ // This callback should populate the ArrowDeviceArray associated with this task.
+ // The order of ArrowAsyncTasks provided by the producer enables a consumer to
+ // ensure the order of data to process.
+ //
+ // This function is expected to be synchronous, but should not perform any blocking
+ // I/O. Ideally it should be as cheap as possible so as to not tie up the consumer
+ // thread unnecessarily.
+ //
+ // Returns: 0 if successful, errno-compatible error otherwise.
+ //
+ // If a non-0 value is returned then it should be followed by a call to `on_error`
+ // on the appropriate ArrowAsyncDeviceStreamHandler. This is because it's highly
+ // likely that whatever is calling this function may be entirely disconnected from
+ // the current control flow. Indicating an error here with a non-zero return allows
+ // the current flow to be aware of the error occurring, while still allowing any
+ // logging or error handling to still be centralized in the `on_error` callback of
+ // the original Async handler.
+ //
+ // Rather than a release callback, any required cleanup should be performed as part
+ // of the invocation of `extract_data`. Ownership of the Array is passed to the consumer
+ // calling this, and so it must be released separately.
+ //
+ // It is only valid to call this method exactly once.
+ int (*extract_data)(struct ArrowAsyncTask* self, struct ArrowDeviceArray* out);
Review Comment:
Any custom metadata was already passed alongside the `ArrowAsyncTask` when `on_next_task` was called. So it wouldn't make sense to also add it here too. You can see in the code how I propagate the metadata along with the task itself.
I wouldn't be opposed to moving the metadata to the `extract_data` call itself, I just think it shouldn't be on both sides.
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1833564614
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
Review Comment:
nit: `const std::string&` to avoid the copy?
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "lidavidm (via GitHub)" <gi...@apache.org>.
lidavidm commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1833773008
##########
cpp/src/arrow/c/bridge.h:
##########
@@ -406,4 +407,81 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
/// @}
+/// \defgroup c-async-stream-interface Functions for working with the async C data
+/// interface.
+///
+/// @{
+
+/// \brief AsyncErrorDetail is a StatusDetail that contains an error code and message
+/// from an asynchronous operation.
+class AsyncErrorDetail : public StatusDetail {
+ public:
+ AsyncErrorDetail(int code, std::string message, std::string metadata)
+ : code_(code), message_(std::move(message)), metadata_(std::move(metadata)) {}
+ const char* type_id() const override { return "AsyncErrorDetail"; }
+ // ToString just returns the error message that was returned with the error
+ std::string ToString() const override { return message_; }
Review Comment:
oh wait this is ToString and it just happens to be trivial so maybe not
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "zeroshade (via GitHub)" <gi...@apache.org>.
zeroshade commented on PR #44495:
URL: https://github.com/apache/arrow/pull/44495#issuecomment-2461747050
I've updated this with the updated changes to the Async structs. Can I get an updated review?
CC @lidavidm @bkietz @pitrou @paleolimbot
--
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-43631: [C++] Add C++ implementation of Async C Data Interface [arrow]
Posted by "bkietz (via GitHub)" <gi...@apache.org>.
bkietz commented on code in PR #44495:
URL: https://github.com/apache/arrow/pull/44495#discussion_r1834704974
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
Review Comment:
Since this block mostly uses State's members, it should probably be a member function of State
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
Review Comment:
Nit: instead of pair, could we use
```c++
struct TaskWithMetadata {
ArrowAsyncTask task;
std::shared_ptr<KeyValueMetadata>> metadata;
};
```
and rename this
```suggestion
std::queue<TaskWithMetadata> tasks_;
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
Review Comment:
```suggestion
AsyncRecordBatchIterator(uint64_t queue_size, DeviceMemoryMapper mapper)
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
Review Comment:
Nit
```suggestion
return IterationEnd<RecordBatchWithMetadata>();
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
Review Comment:
```suggestion
State(uint64_t queue_size, DeviceMemoryMapper mapper)
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::string message_str, metadata_str;
+ if (message != nullptr) {
+ message_str = message;
+ }
+ if (metadata != nullptr) {
+ metadata_str = metadata;
+ }
+
+ Status error = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message_str, std::move(metadata_str)),
+ std::move(message_str));
+
+ if (!private_data->fut_iterator_.is_finished()) {
+ private_data->fut_iterator_.MarkFinished(error);
+ return;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = error;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
+ return !state_->error_.ok() || state_->pending_requests_ > 0;
+ });
+ }
+
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->pending_requests_ > 0) {
+ state_->pending_requests_--;
+ lock.unlock();
+
+ ArrowAsyncTask task;
+ task.private_data = new PrivateTaskData{state_, record};
+ task.extract_data = AsyncProducer::extract_data;
+
+ if (int status = handler_->on_next_task(handler_, &task, nullptr) != 0) {
+ delete reinterpret_cast<PrivateTaskData*>(task.private_data);
+ return Status::UnknownError("Received error from handler::on_next_task ", status);
+ }
+ }
+
+ return Status::OK();
+ }
+
+ static void request(struct ArrowAsyncProducer* producer, int64_t n) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (!self->error_.ok()) {
+ return;
+ }
+ self->pending_requests_ += n;
+ }
+ self->cv_.notify_all();
+ }
+
+ static void cancel(struct ArrowAsyncProducer* producer) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (!self->error_.ok()) {
+ return;
+ }
+ self->error_ = Status::Cancelled("Consumer requested cancellation");
+ }
+ self->cv_.notify_all();
+ }
+
+ static int extract_data(struct ArrowAsyncTask* task, struct ArrowDeviceArray* out) {
+ auto private_data = reinterpret_cast<PrivateTaskData*>(task->private_data);
+ int ret = 0;
+ if (out != nullptr) {
+ auto status = ExportDeviceRecordBatch(*private_data->record_,
+ private_data->record_->GetSyncEvent(), out);
+ if (!status.ok()) {
+ std::lock_guard<std::mutex> lock(private_data->producer_->mutex_);
+ private_data->producer_->error_ = status;
+ }
+ }
+
+ delete private_data;
+ return ret;
+ }
+
+ struct ArrowAsyncDeviceStreamHandler* handler_;
+ std::shared_ptr<State> state_;
+};
+
+} // namespace
+
+Future<AsyncRecordBatchGenerator> CreateAsyncDeviceStreamHandler(
+ struct ArrowAsyncDeviceStreamHandler* handler, internal::Executor* executor,
+ uint64_t queue_size, DeviceMemoryMapper mapper) {
+ auto iterator =
+ std::make_shared<AsyncRecordBatchIterator>(queue_size, std::move(mapper));
+ return AsyncRecordBatchIterator::Make(*iterator, handler)
+ .Then([executor](std::shared_ptr<AsyncRecordBatchIterator::State> state)
+ -> Result<AsyncRecordBatchGenerator> {
+ AsyncRecordBatchGenerator gen{state->schema_, state->device_type_, nullptr};
+ auto it =
+ Iterator<RecordBatchWithMetadata>(AsyncRecordBatchIterator{std::move(state)});
+ ARROW_ASSIGN_OR_RAISE(gen.generator,
+ MakeBackgroundGenerator(std::move(it), executor));
+ return gen;
+ });
+}
+
+Future<> ExportAsyncRecordBatchReader(
+ std::shared_ptr<Schema> schema,
+ AsyncGenerator<std::shared_ptr<RecordBatch>> generator,
+ DeviceAllocationType device_type, struct ArrowAsyncDeviceStreamHandler* handler) {
+ if (!schema) {
+ handler->on_error(handler, EINVAL, "Schema is null", nullptr);
+ handler->release(handler);
+ return Future<>::MakeFinished(Status::Invalid("Schema is null"));
+ }
+
+ struct ArrowSchema c_schema;
+ SchemaExportGuard guard(&c_schema);
+
+ auto status = ExportSchema(*schema, &c_schema);
+ if (!status.ok()) {
+ handler->on_error(handler, EINVAL, status.message().c_str(), nullptr);
+ handler->release(handler);
+ return Future<>::MakeFinished(status);
+ }
+
+ return VisitAsyncGenerator(generator, AsyncProducer{device_type, &c_schema, handler})
+ .Then(
+ [handler]() -> Future<> {
Review Comment:
I think this should be able to just return Status since we don't have any true continuations in here
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
Review Comment:
```suggestion
private_data->state_->error_ = std::move(maybe_decoded).status();
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::string message_str, metadata_str;
+ if (message != nullptr) {
+ message_str = message;
+ }
+ if (metadata != nullptr) {
+ metadata_str = metadata;
+ }
+
+ Status error = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message_str, std::move(metadata_str)),
+ std::move(message_str));
+
+ if (!private_data->fut_iterator_.is_finished()) {
+ private_data->fut_iterator_.MarkFinished(error);
+ return;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = error;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
+ }
+
+ std::shared_ptr<State> state_;
+};
+
+struct AsyncProducer {
+ struct State {
+ struct ArrowAsyncProducer producer_;
+
+ std::mutex mutex_;
+ std::condition_variable cv_;
+ uint64_t pending_requests_{0};
+ Status error_{Status::OK()};
+ };
+
+ AsyncProducer(DeviceAllocationType device_type, struct ArrowSchema* schema,
+ struct ArrowAsyncDeviceStreamHandler* handler)
+ : handler_{handler}, state_{std::make_shared<State>()} {
+ state_->producer_.device_type = static_cast<ArrowDeviceType>(device_type);
+ state_->producer_.private_data = reinterpret_cast<void*>(state_.get());
+ state_->producer_.request = AsyncProducer::request;
+ state_->producer_.cancel = AsyncProducer::cancel;
+ handler_->producer = &state_->producer_;
+
+ if (int status = handler_->on_schema(handler_, schema) != 0) {
+ state_->error_ =
+ Status::UnknownError("Received error from handler::on_schema ", status);
+ }
+ }
+
+ struct PrivateTaskData {
+ PrivateTaskData(std::shared_ptr<State> producer, std::shared_ptr<RecordBatch> record)
+ : producer_{std::move(producer)}, record_(std::move(record)) {}
+
+ std::shared_ptr<State> producer_;
+ std::shared_ptr<RecordBatch> record_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateTaskData);
+ };
+
+ Status operator()(const std::shared_ptr<RecordBatch>& record) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ if (state_->pending_requests_ == 0) {
+ state_->cv_.wait(lock, [this]() -> bool {
+ return !state_->error_.ok() || state_->pending_requests_ > 0;
+ });
+ }
+
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->pending_requests_ > 0) {
+ state_->pending_requests_--;
+ lock.unlock();
+
+ ArrowAsyncTask task;
+ task.private_data = new PrivateTaskData{state_, record};
+ task.extract_data = AsyncProducer::extract_data;
+
+ if (int status = handler_->on_next_task(handler_, &task, nullptr) != 0) {
+ delete reinterpret_cast<PrivateTaskData*>(task.private_data);
+ return Status::UnknownError("Received error from handler::on_next_task ", status);
+ }
+ }
+
+ return Status::OK();
+ }
+
+ static void request(struct ArrowAsyncProducer* producer, int64_t n) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (!self->error_.ok()) {
+ return;
+ }
+ self->pending_requests_ += n;
+ }
+ self->cv_.notify_all();
+ }
+
+ static void cancel(struct ArrowAsyncProducer* producer) {
+ auto* self = reinterpret_cast<State*>(producer->private_data);
+ {
+ std::lock_guard<std::mutex> lock(self->mutex_);
+ if (!self->error_.ok()) {
+ return;
+ }
+ self->error_ = Status::Cancelled("Consumer requested cancellation");
+ }
+ self->cv_.notify_all();
+ }
+
+ static int extract_data(struct ArrowAsyncTask* task, struct ArrowDeviceArray* out) {
+ auto private_data = reinterpret_cast<PrivateTaskData*>(task->private_data);
Review Comment:
In lieu of `defer` :smile:
```suggestion
std::unique_ptr private_data{reinterpret_cast<PrivateTaskData*>(task->private_data)};
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::string message_str, metadata_str;
+ if (message != nullptr) {
+ message_str = message;
+ }
+ if (metadata != nullptr) {
+ metadata_str = metadata;
+ }
+
+ Status error = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message_str, std::move(metadata_str)),
+ std::move(message_str));
+
+ if (!private_data->fut_iterator_.is_finished()) {
+ private_data->fut_iterator_.MarkFinished(error);
+ return;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = error;
Review Comment:
```suggestion
private_data->state_->error_ = std::move(error);
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
Review Comment:
```suggestion
private_data->state_->schema_ = maybe_schema.MoveValueUnsafe();
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
Review Comment:
```suggestion
kvmetadata = std::move(maybe_decoded->metadata);
```
##########
cpp/src/arrow/c/bridge.cc:
##########
@@ -2511,4 +2516,345 @@ Result<std::shared_ptr<ChunkedArray>> ImportDeviceChunkedArray(
return ImportChunked</*IsDevice=*/true>(stream, mapper);
}
+namespace {
+
+class AsyncRecordBatchIterator {
+ public:
+ struct State {
+ State(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : queue_size_{queue_size}, mapper_{std::move(mapper)} {}
+
+ const uint64_t queue_size_;
+ const DeviceMemoryMapper mapper_;
+ ArrowAsyncProducer* producer_;
+ DeviceAllocationType device_type_;
+
+ std::mutex mutex_;
+ std::shared_ptr<Schema> schema_;
+ std::condition_variable cv_;
+ std::queue<std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>>> batches_;
+ bool end_of_stream_ = false;
+ Status error_{Status::OK()};
+ };
+
+ AsyncRecordBatchIterator(uint64_t queue_size, const DeviceMemoryMapper mapper)
+ : state_{std::make_shared<State>(queue_size, std::move(mapper))} {}
+
+ explicit AsyncRecordBatchIterator(std::shared_ptr<State> state)
+ : state_{std::move(state)} {}
+
+ const std::shared_ptr<Schema>& schema() const { return state_->schema_; }
+
+ DeviceAllocationType device_type() const { return state_->device_type_; }
+
+ Result<RecordBatchWithMetadata> Next() {
+ std::pair<ArrowAsyncTask, std::shared_ptr<KeyValueMetadata>> task;
+ {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] {
+ return !state_->error_.ok() || !state_->batches_.empty() ||
+ state_->end_of_stream_;
+ });
+ if (!state_->error_.ok()) {
+ return state_->error_;
+ }
+
+ if (state_->batches_.empty() && state_->end_of_stream_) {
+ return RecordBatchWithMetadata{nullptr, nullptr};
+ }
+
+ task = state_->batches_.front();
+ state_->batches_.pop();
+ }
+
+ state_->producer_->request(state_->producer_, 1);
+ ArrowDeviceArray out;
+ if (task.first.extract_data(&task.first, &out) != 0) {
+ std::unique_lock<std::mutex> lock(state_->mutex_);
+ state_->cv_.wait(lock, [&] { return !state_->error_.ok(); });
+ return state_->error_;
+ }
+
+ ARROW_ASSIGN_OR_RAISE(
+ auto batch, ImportDeviceRecordBatch(&out, state_->schema_, state_->mapper_));
+ return RecordBatchWithMetadata{std::move(batch), std::move(task.second)};
+ }
+
+ static Future<std::shared_ptr<AsyncRecordBatchIterator::State>> Make(
+ AsyncRecordBatchIterator& iterator, struct ArrowAsyncDeviceStreamHandler* handler) {
+ auto iterator_fut = Future<std::shared_ptr<AsyncRecordBatchIterator::State>>::Make();
+
+ auto private_data = new PrivateData{iterator.state_};
+ private_data->fut_iterator_ = iterator_fut;
+
+ handler->private_data = private_data;
+ handler->on_schema = on_schema;
+ handler->on_next_task = on_next_task;
+ handler->on_error = on_error;
+ handler->release = release;
+ return iterator_fut;
+ }
+
+ private:
+ struct PrivateData {
+ explicit PrivateData(std::shared_ptr<State> state) : state_(std::move(state)) {}
+
+ std::shared_ptr<State> state_;
+ Future<std::shared_ptr<AsyncRecordBatchIterator::State>> fut_iterator_;
+ ARROW_DISALLOW_COPY_AND_ASSIGN(PrivateData);
+ };
+
+ static int on_schema(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ if (self->producer != nullptr) {
+ private_data->state_->producer_ = self->producer;
+ private_data->state_->device_type_ =
+ static_cast<DeviceAllocationType>(self->producer->device_type);
+ }
+
+ auto maybe_schema = ImportSchema(stream_schema);
+ if (!maybe_schema.ok()) {
+ private_data->fut_iterator_.MarkFinished(maybe_schema.status());
+ return EINVAL;
+ }
+
+ auto schema = maybe_schema.MoveValueUnsafe();
+ private_data->state_->schema_ = schema;
+ private_data->fut_iterator_.MarkFinished(private_data->state_);
+ self->producer->request(self->producer,
+ static_cast<int64_t>(private_data->state_->queue_size_));
+ return 0;
+ }
+
+ static int on_next_task(ArrowAsyncDeviceStreamHandler* self, ArrowAsyncTask* task,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+
+ if (task == nullptr) {
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->end_of_stream_ = true;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ std::shared_ptr<KeyValueMetadata> kvmetadata;
+ if (metadata != nullptr) {
+ auto maybe_decoded = DecodeMetadata(metadata);
+ if (!maybe_decoded.ok()) {
+ private_data->state_->error_ = maybe_decoded.status();
+ private_data->state_->cv_.notify_one();
+ return EINVAL;
+ }
+
+ kvmetadata = maybe_decoded->metadata;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->batches_.emplace(*task, std::move(kvmetadata));
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ return 0;
+ }
+
+ static void on_error(ArrowAsyncDeviceStreamHandler* self, int code, const char* message,
+ const char* metadata) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ std::string message_str, metadata_str;
+ if (message != nullptr) {
+ message_str = message;
+ }
+ if (metadata != nullptr) {
+ metadata_str = metadata;
+ }
+
+ Status error = Status::FromDetailAndArgs(
+ StatusCode::UnknownError,
+ std::make_shared<AsyncErrorDetail>(code, message_str, std::move(metadata_str)),
+ std::move(message_str));
+
+ if (!private_data->fut_iterator_.is_finished()) {
+ private_data->fut_iterator_.MarkFinished(error);
+ return;
+ }
+
+ std::unique_lock<std::mutex> lock(private_data->state_->mutex_);
+ private_data->state_->error_ = error;
+ lock.unlock();
+ private_data->state_->cv_.notify_one();
+ }
+
+ static void release(ArrowAsyncDeviceStreamHandler* self) {
+ auto* private_data = reinterpret_cast<PrivateData*>(self->private_data);
+ delete private_data;
Review Comment:
nit
```suggestion
delete reinterpret_cast<PrivateData*>(self->private_data);
```
--
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