Add function to wrap any stream

Author: alamb

datafusion-examples/examples/thread_pools.rs

@@ -24,7 +24,6 @@
 //! issues described in  the [Architecture section] such as throttled bandwidth
 //! due to congestion control and increased latencies for processing network
 //! messages.
-
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::error::Result;
 use datafusion::execution::SendableRecordBatchStream;
@@ -144,8 +143,6 @@ async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()>
 
 /// Demonstrates how to run queries on a different runtime than the current run
 /// and how to handle IO operations.
-///
-
 async fn different_runtime_advanced() -> Result<()> {
     // In this example, we will configure access to a remote object store
     // over the network during the plan

datafusion/physical-plan/src/cross_rt_stream.rs

@@ -103,7 +103,7 @@ where
 {
     /// Create new stream based on an existing stream that transports [`Result`]s.
     ///
-    /// Also receives an executor that actually executes the underlying stream as well as a converter that convets
+    /// Also receives an executor that actually executes the underlying stream as well as a converter that converts
     /// [`executor::JobError`] to the error type of the stream (so we can send potential crashes/panics).
     pub fn new_with_error_stream<S, C>(
         stream: S,

datafusion/physical-plan/src/dedicated_executor.rs

@@ -25,10 +25,7 @@ use crate::io_object_store::IoObjectStore;
 use crate::stream::RecordBatchStreamAdapter;
 use crate::SendableRecordBatchStream;
 use datafusion_common::DataFusionError;
-use futures::{
-    future::{BoxFuture, Shared},
-    Future, FutureExt, TryFutureExt,
-};
+use futures::{future::{BoxFuture, Shared}, Future, FutureExt, Stream, TryFutureExt};
 use log::{info, warn};
 use object_store::ObjectStore;
 use parking_lot::RwLock;
@@ -59,6 +56,8 @@ impl From<Builder> for DedicatedExecutorBuilder {
 /// tasks on the same threadpool by running futures (and any `tasks` that are
 /// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
 ///
+/// DedicatedExecutor can be `clone`ed and all clones share the same threadpool.
+///
 /// TODO add note about `io_thread`
 ///
 /// TODO: things we use in InfluxData
@@ -119,7 +118,7 @@ impl From<Builder> for DedicatedExecutorBuilder {
 /// happens when a runtime is dropped from within an asynchronous
 /// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
 ///
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct DedicatedExecutor {
     state: Arc<RwLock<State>>,
 }
@@ -222,8 +221,8 @@ impl DedicatedExecutor {
     /// 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))
+    pub fn wrap_object_store(&self, object_store: Arc<dyn ObjectStore>) -> Arc<IoObjectStore> {
+        Arc::new(IoObjectStore::new(self.clone(), object_store))
     }
 
     /// Returns a SendableRecordBatchStream that will run on this executor's thread pool
@@ -237,6 +236,28 @@ impl DedicatedExecutor {
         Box::pin(RecordBatchStreamAdapter::new(schema, cross_rt_stream))
     }
 
+    /// Runs an stream that produces Results on the executor's thread pool
+    ///
+    /// Ths stream must produce Results so that any errors on the dedicated
+    /// executor (like a panic or shutdown) can be communicated back.
+    ///
+    /// # Arguments:
+    /// - stream: the stream to run on this dedicated executor
+    /// - converter: a function that converts a [`JobError`] to the error type of the stream
+    pub fn run_stream<X, E, S, C>(
+        &self,
+        stream: S,
+        converter: C,
+    ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+    where
+        X: Send + 'static,
+        E: Send + 'static,
+        S: Stream<Item = Result<X, E>> + Send + 'static,
+        C: Fn(JobError) -> E + Send + 'static,
+    {
+        CrossRtStream::new_with_error_stream(stream, self.clone(), converter)
+    }
+
     /// Registers `handle` as the IO runtime for this thread
     ///
     /// This sets a thread-local variable
@@ -306,6 +327,7 @@ impl<T> Future for DropGuard<T> {
 /// The state is only used by the "outer" API, not by the newly created runtime. The new runtime waits for
 /// [`start_shutdown`](Self::start_shutdown) and signals the completion via
 /// [`completed_shutdown`](Self::completed_shutdown) (for which is owns the sender side).
+#[derive(Debug)]
 struct State {
     /// Runtime handle.
     ///

datafusion/physical-plan/src/io_object_store.rs

@@ -20,10 +20,12 @@ use std::sync::Arc;
 use crate::DedicatedExecutor;
 use async_trait::async_trait;
 use futures::stream::BoxStream;
+use futures::StreamExt;
 use object_store::{
     path::Path, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
     ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result,
 };
+use crate::dedicated_executor::JobError;
 
 /// 'ObjectStore' that wraps an inner `ObjectStore` and wraps all the underlying
 /// methods with [`DedicatedExecutor::spawn_io`] so that they are run on the Tokio Runtime
@@ -33,12 +35,14 @@ use object_store::{
 ///
 #[derive(Debug)]
 pub struct IoObjectStore {
+    executor: DedicatedExecutor,
     inner: Arc<dyn ObjectStore>,
 }
 
 impl IoObjectStore {
-    pub fn new(object_store: Arc<dyn ObjectStore>) -> Self {
+    pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
         Self {
+            executor,
             inner: object_store,
         }
     }
@@ -50,9 +54,15 @@ impl std::fmt::Display for IoObjectStore {
     }
 }
 
+fn convert_error(e: JobError) -> object_store::Error {
+    object_store::Error::Generic {
+        store: "IoObjectStore",
+        source: Box::new(e)
+    }
+}
+
 #[async_trait]
 impl ObjectStore for IoObjectStore {
-    /// TODO wrap the resulting stream in CrossRTStream
     async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
         let location = location.clone();
         let store = Arc::clone(&self.inner);
@@ -83,7 +93,12 @@ impl ObjectStore for IoObjectStore {
     }
 
     fn list(&self, prefix: Option<&Path>) -> BoxStream<'_, Result<ObjectMeta>> {
-        self.inner.list(prefix)
+        // run the inner list on the dedicated executor
+        let inner_stream = self.inner.list(prefix);
+
+        inner_stream
+        //self.executor.run_stream(inner_stream, convert_error)
+         //   .boxed()
     }
 
     async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
@@ -121,4 +136,6 @@ impl ObjectStore for IoObjectStore {
         })
         .await
     }
+
+
 }