alamb commented on code in PR #3464:
URL: https://github.com/apache/arrow-rs/pull/3464#discussion_r1091936930
##########
arrow-flight/src/client.rs:
##########
@@ -531,3 +549,142 @@ impl FlightClient {
request
}
}
+
+// splits the input stream into an invallable flight data stream and errors
errors
+//
+// TODO generify
+fn split_stream(
+ input_stream: BoxStream<'static, Result<FlightData>>,
+) -> (SplitStreamOk, SplitStreamErr) {
+ let inner = SplitStream {
+ input_stream,
+ next_ok: None,
+ next_err: None,
+ done: false,
+ };
+ let inner = Arc::new(Mutex::new(inner));
+
+ let ok_stream = SplitStreamOk {
+ inner: Arc::clone(&inner),
+ };
+
+ let err_stream = SplitStreamErr {
+ inner: Arc::clone(&inner),
+ };
+
+ (ok_stream, err_stream)
+}
+
+struct SplitStream {
+ input_stream: BoxStream<'static, Result<FlightData>>,
+ next_ok: Option<FlightData>,
+ next_err: Option<FlightError>,
+ done: bool,
+}
+
+impl SplitStream {
+ // returns the next ok item ready if any
+ fn poll_next_ok(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Option<FlightData>> {
+ loop {
+ if let Some(flight_data) = self.next_ok.take() {
+ return Poll::Ready(Some(flight_data));
+ }
+
+ if self.done {
+ return Poll::Ready(None);
+ }
+
+ // try to get another item from the inner stream
+ if !self.maybe_read(cx) {
+ return Poll::Pending;
+ }
+ }
+ }
+
+ // returns the next err item
+ fn poll_next_err(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> std::task::Poll<Option<FlightError>> {
+ loop {
+ if let Some(e) = self.next_err.take() {
+ return Poll::Ready(Some(e));
+ }
+
+ if self.done {
+ return Poll::Ready(None);
+ }
+
+ // try to get another item from the inner stream
+ if !self.maybe_read(cx) {
+ return Poll::Pending;
+ }
+ }
+ }
+
+ // if we have space for both ok and error, take next from inner stream
+ // returns true if read an item false otherwise
+ fn maybe_read(&mut self, cx: &mut std::task::Context<'_>) -> bool {
+ // if there is space for ok and err, take next
+ if self.next_ok.is_some() || self.next_err.is_some() {
+ // can't take next until there is space
+ return false;
+ }
+
+ let next = match self.input_stream.poll_next_unpin(cx) {
+ Poll::Pending => return false,
+ Poll::Ready(next) => next,
+ };
+
+ match next {
+ Some(Ok(flight_data)) => {
+ self.next_ok = Some(flight_data);
+ }
+ Some(Err(e)) => {
+ self.next_err = Some(e);
+ // stop reading once we see an error
+ self.done = true;
+ }
+ None => {
+ self.done = true;
+ }
+ };
+
+ true
+ }
+}
+
+/// returns only the OK responses from a stream of results
+struct SplitStreamErr {
Review Comment:
This makes sense @Dandandan 👍 thank you for the suggestion
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
