[GitHub] [arrow-datafusion] edrevo commented on a change in pull request #535: Make BallistaContext::collect streaming

[GitBox <gi...@apache.org>]

edrevo commented on a change in pull request #535:
URL: https://github.com/apache/arrow-datafusion/pull/535#discussion_r648949167



##########
File path: ballista/rust/client/src/context.rs
##########
@@ -68,6 +74,32 @@ impl BallistaContextState {
     }
 }
 
+struct WrappedStream {
+    stream: Pin<Box<dyn Stream<Item = ArrowResult<RecordBatch>> + Send + Sync>>,
+    schema: SchemaRef,
+}
+
+impl RecordBatchStream for WrappedStream {
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}
+
+impl Stream for WrappedStream {
+    type Item = ArrowResult<RecordBatch>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.stream.poll_next_unpin(cx)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.stream.size_hint()
+    }
+}
+

Review comment:
       I was surprised I couldn't find anything like this. If there is a similar struct please do let me know and I'll use that one.
   
   Also, since this is a pretty general wrapper, if you want me to move this to another place and make it public, I can do that too.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org