wrap result stream

Author: alamb

datafusion-examples/examples/thread_pools.rs

@@ -28,7 +28,7 @@
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::error::Result;
 use datafusion::execution::SendableRecordBatchStream;
-use datafusion::physical_plan::DedicatedExecutor;
+use datafusion::physical_plan::{dedicated_executor, DedicatedExecutor};
 use datafusion::prelude::*;
 use futures::stream::StreamExt;
 use object_store::http::HttpBuilder;
@@ -172,6 +172,7 @@ async fn different_runtime_advanced() -> Result<()> {
 
     let http_store = DedicatedExecutor::wrap_object_store(http_store);
 
+    // Tell datafusion about processing http:// urls with this wrapped object store
     ctx.register_object_store(&base_url, http_store);
 
     // Plan (and execute) the query on the dedicated runtime
@@ -184,19 +185,28 @@ async fn different_runtime_advanced() -> Result<()> {
                 .await?;
             let stream: SendableRecordBatchStream = df.execute_stream().await?;
 
-            Ok(stream) as Result<SendableRecordBatchStream>
+            Ok(stream) as Result<_>
         }).await??;
 
-    // We have now planned the query on the dedicated runtime, but we still need to
-    // drive the stream (aka call `next()` to get the results.
+    // We have now planned the query on the dedicated runtime, Yay! but we still need to
+    // drive the stream (aka call `next()` to get the results).
 
-    // as mentioned above,  calling `next()` (including indirectly by using
-    // FlightDataEncoder to convert the results to flight to send it over the
-    // network), will *still* result in the CPU work (and a bunch of spawned
-    // tasks) being done on the runtime calling next() (aka the current runtime)
-    // and not on the dedicated runtime.
-
-    // to drive the stream on the dedicated runtime, we need to wrap it using a XXX stream function
+    // However, as mentioned above, calling  `next()` resolves the Stream (and
+    // any work it may do) on a thread in the current (default) runtime.
+    //
+    // To drive the stream on the dedicated runtime, we need to wrap it using a
+    // `DedicatedExecutor::wrap_stream` stream function
+    //
+    // Note if you don't do this you will likely see a panic about `No IO runtime registered.`
+    // because the threads in the current (main) tokio runtime have not had the IO runtime
+    // installed
+    let mut stream = dedicated_executor.run_sendable_record_batch_stream(stream);
+
+    // Note you can run other streams on the DedicatedExecutor as well using the
+    // DedicatedExecutor:YYYXXX function. This is helpful for example, if you
+    // need to do non trivial CPU work on the results of the stream (e.g.
+    // calling a FlightDataEncoder to convert the results to flight to send it
+    // over the network),
 
     while let Some(batch) = stream.next().await {
         println!("{}", pretty_format_batches(&[batch?]).unwrap());

datafusion/physical-plan/src/dedicated_executor.rs

@@ -218,12 +218,16 @@ impl DedicatedExecutor {
 
     /// Returns an ObjectStore instance that will always perform I/O work on the
     /// IO_RUNTIME.
+    ///
+    /// Note that this object store will only work correctly if run on this
+    /// dedicated executor. If you try and use it on another executor, it will
+    /// panic with "no IO runtime registered" type error.
     pub fn wrap_object_store(object_store: Arc<dyn ObjectStore>) -> Arc<IoObjectStore> {
         Arc::new(IoObjectStore::new(object_store))
     }
 
     /// Returns a SendableRecordBatchStream that will run on this executor's thread pool
-    pub fn wrap_stream(
+    pub fn run_sendable_record_batch_stream(
         &self,
         stream: SendableRecordBatchStream,
     ) -> SendableRecordBatchStream {