westonpace commented on a change in pull request #9643:
URL: https://github.com/apache/arrow/pull/9643#discussion_r589800665
##########
File path: cpp/src/arrow/util/async_generator.h
##########
@@ -87,6 +103,143 @@ Future<std::vector<T>>
CollectAsyncGenerator(AsyncGenerator<T> generator) {
return Loop(LoopBody{std::move(generator), std::move(vec)});
}
+template <typename T, typename V>
+class MappingGenerator {
+ public:
+ MappingGenerator(AsyncGenerator<T> source, std::function<Future<V>(const
T&)> map)
+ : state_(std::make_shared<State>(std::move(source), std::move(map))) {}
+
+ Future<V> operator()() {
+ auto future = Future<V>::Make();
+ bool should_trigger;
+ {
+ auto guard = state_->mutex.Lock();
+ if (state_->finished) {
+ return Future<V>::MakeFinished(IterationTraits<V>::End());
+ }
+ should_trigger = state_->waiting_jobs.empty();
+ state_->waiting_jobs.push_back(future);
+ }
+ if (should_trigger) {
+ state_->source().AddCallback(Callback{state_});
+ }
+ return future;
+ }
+
+ private:
+ struct State {
+ State(AsyncGenerator<T> source, std::function<Future<V>(const T&)> map)
+ : source(std::move(source)),
+ map(std::move(map)),
+ waiting_jobs(),
+ mutex(),
+ finished(false) {}
+
+ AsyncGenerator<T> source;
+ std::function<Future<V>(const T&)> map;
+ std::deque<Future<V>> waiting_jobs;
+ util::Mutex mutex;
+ bool finished;
+ };
+
+ struct Callback;
+
+ struct MappedCallback {
+ void operator()(const Result<V>& maybe_next) {
+ bool end = !maybe_next.ok() || IterationTraits<V>::IsEnd(*maybe_next);
+ bool should_purge = false;
+ if (end) {
+ {
+ auto guard = state->mutex.Lock();
+ should_purge = !state->finished;
+ state->finished = true;
+ }
+ }
+ sink.MarkFinished(maybe_next);
+ if (should_purge) {
+ while (!state->waiting_jobs.empty()) {
+ state->waiting_jobs.front().MarkFinished(IterationTraits<V>::End());
+ state->waiting_jobs.pop_front();
+ }
+ }
+ }
+ std::shared_ptr<State> state;
+ Future<V> sink;
+ };
+
+ struct Callback {
+ void operator()(const Result<T>& maybe_next) {
+ Future<V> sink;
+ bool end = !maybe_next.ok() || IterationTraits<T>::IsEnd(*maybe_next);
+ bool should_purge = false;
+ bool should_trigger;
+ {
+ auto guard = state->mutex.Lock();
+ if (end) {
+ should_purge = !state->finished;
+ state->finished = true;
+ }
+ sink = state->waiting_jobs.front();
+ state->waiting_jobs.pop_front();
+ should_trigger = !end && !state->waiting_jobs.empty();
+ }
+ if (should_purge) {
+ while (!state->waiting_jobs.empty()) {
+ state->waiting_jobs.front().MarkFinished(IterationTraits<V>::End());
+ state->waiting_jobs.pop_front();
+ }
+ }
+ if (should_trigger) {
Review comment:
Ok, I added an explanation of async-reentrant to the top of the file
(couldn't find a better place) and added "\see MakeXYZ" to "class XYZGenerator".
##########
File path: cpp/src/arrow/util/vector.h
##########
@@ -81,5 +84,53 @@ std::vector<T> FilterVector(std::vector<T> values,
Predicate&& predicate) {
return values;
}
+/// \brief Like MapIterator, but where the function can fail.
Review comment:
Done.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
