[GitHub] [arrow-datafusion] alamb commented on a change in pull request #789: WIP: Implement streaming versions of Dataframe.collect methods

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

alamb commented on a change in pull request #789:
URL: https://github.com/apache/arrow-datafusion/pull/789#discussion_r678247263



##########
File path: datafusion/src/dataframe.rs
##########
@@ -222,6 +223,21 @@ pub trait DataFrame: Send + Sync {
     /// ```
     async fn collect(&self) -> Result<Vec<RecordBatch>>;
 
+    /// Executes this DataFrame and returns a stream over a single partition
+    ///
+    /// ```
+    /// # use datafusion::prelude::*;
+    /// # use datafusion::error::Result;
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<()> {
+    /// let mut ctx = ExecutionContext::new();
+    /// let df = ctx.read_csv("tests/example.csv", CsvReadOptions::new())?;
+    /// let stream = df.collect_stream().await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    async fn collect_stream(&self) -> Result<SendableRecordBatchStream>;

Review comment:
       What if we called this something like `execute` rather than `collect_stream`? 
   
   ```
       async fn execute_stream(&self) -> Result<SendableRecordBatchStream>;
   ```
   
   This would mirror the naming of `ExecutionPlan::execute` and might make it clearer that `collect` means collect into a Vec and `execute` means get a stream

##########
File path: datafusion/src/execution/dataframe_impl.rs
##########
@@ -149,8 +152,19 @@ impl DataFrame for DataFrameImpl {
         Ok(collect(plan).await?)
     }
 
-    // Convert the logical plan represented by this DataFrame into a physical plan and
-    // execute it
+    /// Convert the logical plan represented by this DataFrame into a physical plan and
+    /// execute it, returning a stream over a single partition
+    async fn collect_stream(&self) -> Result<SendableRecordBatchStream> {
+        let state = self.ctx_state.lock().unwrap().clone();
+        let ctx = ExecutionContext::from(Arc::new(Mutex::new(state)));
+        let plan = ctx.optimize(&self.plan)?;
+        let plan = ctx.create_physical_plan(&plan)?;
+        collect_stream(plan).await
+    }
+
+    /// Convert the logical plan represented by this DataFrame into a physical plan and
+    /// execute it, collecting all resulting batches into memory while maintaining
+    /// partitioning
     async fn collect_partitioned(&self) -> Result<Vec<Vec<RecordBatch>>> {
         let state = self.ctx_state.lock().unwrap().clone();

Review comment:
       You could probably rewrite `collect_partitioned` to be in terms of `collect_stream_partitioned`:
   
   ```rust
   collect(self.collect_stream_partitioned().await?)
   ``` 
   or something like 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