You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2025/01/24 20:35:34 UTC
[PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
alamb opened a new pull request, #14286:
URL: https://github.com/apache/datafusion/pull/14286
Note: This PR contains a (substantial) example and supporting code. It has no changes to the core.
## Which issue does this PR close?
- Closes https://github.com/apache/datafusion/issues/12393
- Note this is new version version of https://github.com/apache/datafusion/pull/13424
## Rationale for this change
I have heard from multiple people multiple times that the specifics of using multiple threadpools for separate CPU and IO work in DataFusion is confusing.
However, this is a key detail for building high performance engines which process data directly from remote storage, which I think is a very important capability for DataFusion
My [past attempts](https://github.com/apache/datafusion/pull/13424) to make this example have been bogged down trying to get consensus on details of how to transfer results across streams, the wisdom of wrapping streams, and other details.
However, there have been no other actual alternatives proposed (aka no actual code that I could write an example with). So while the approach of wrapping streams may be a bit ugly, it is working for us in production at InfluxDB and I believe @adriangb said they are using the same strategy at Pydantic.
In my opinion, the community needs **something** that works, even if it is not optimial, which is this example.
I would personally like to merge it in so it is easier to find (and not locked in a PR) and iterate afterwards
## What changes are included in this PR?
1. `thread_pools.rs` example
2. a `dedicated_executor.rs` module that
Note that the DedicatedExecutor code is orginally from
1. InfluxDB 3.0 [(source link)](https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor), largely written by @tustvold and @crepererum
3. The `IoObjectStore` wrapper is based on work from @matthewmturne in https://github.com/datafusion-contrib/datafusion-dft/pull/247
Note that I have purposely avoided any changes to the DataFusion crates (such as adding `DedicatedExecutor` to `physical-plan`), but I have written the code with tests with an eye towards doing exactly such a thing one we have some experience / feedback on the approach. This is tracked by
- https://github.com/apache/datafusion/issues/14285
Right now, I think the most important thing is to get this example in for people to try and confirm if it helps them
## Are these changes tested?
Yes the example is run as part of CI and there are tests
TODO: Verify that the tests in the examples run in CI
## Are there any user-facing changes?
Not really
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613969732
> I think this is great Andrew. For what it's worth if this were packaged up in some installable way (even if it had to be from git, etc.) I'm sure we'd be super happy to can our custom stuff and use this :)
Thank you 🙏 If this is the case I would be happy to make a PR to put `DedicatedExecutor` into the core. If had some evidence that people would actually use this I could make a more compelling story to put it into the core (despite the hesitations from @tustvold )
@adriangb Is there any way you can test / verify that the structures in this crate would work for you (like can you temporarily copy/paste dedicated_executor.rs into your tree?)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613970625
Yes I'll do just that and report back. I don't think it should block merging this as an example though.
I understand the hesitation to put it in core, but unless there's something better that's plug and play I think it's better to have this than nothing.
Could an alternative be to put it in a contrib crate or a new workspace package (so that CI also runs on it and it can have tests), add whatever APIs core needs to make it easier to plug in and document the setup? In my mind unless it's set up by default there isn't much more value in having it in core vs any other installable package.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "16pierre (via GitHub)" <gi...@apache.org>.
16pierre commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2086641885
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
How much room is there for clients to mess this up silently, aka, how hard is it to introduce silent regressions in production where we no longer use the expected (correct) threadpool ?
In particular, it seems quite easy to mess up the client-side "stream driving" part.
This is the kind of regressions on which I'd want loud failures early in the development process, at CI-time for example. So here I'd find it desirable to have strengthened guarantees, for example, going from `you will likely see a panic` to `you're **guaranteed** to see a panic`
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2100915637
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
> would be cool if such patterns were shipped out-of-the-box in DF.
Agreed -- it would make a great example
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2634552246
> Assuming this was a mistake we made at InfluxData, we had issues with plan time concurrency that were caused by this also, they were just much harder to track down as they didn't manifest as obviously (it took ~3 months).
For anyone else that is interested, our particular usecase was many concurrent queries (as many as the system can handle) with very latency sensitive workloads (10s of ms) and requires network accesses during planng of queries
FWIW this example shows how to do avoid doing remote catalog access on the CPU thread
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953482293
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
> replace all possible prefixes
The list of registered object stores is empty by default and you only explicitly register the stores you are using.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2654986803
> Is there extra work I'd need to do to get this to work with the `ParquetSink`? This isn't a full reproducer but just a quick copy of the relevant section of the `do_put_statement_ingest()` method in my flightsql server.
>
> I get
>
> > flight_sql: writing data to disk
> > thread 'tokio-runtime-worker' panicked at /Users/dja/.rustup/toolchains/stable-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/ops/function.rs:250:5:
> > A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
>
> ```rust
> let data_sink = ParquetSink::new(file_sink_config, table_options);
>
> let record_batch_stream =
> FlightRecordBatchStream::new_from_flight_data(flight_data_stream.map_err(|e| e.into()));
>
> let stream = Box::pin(RecordBatchStreamAdapter::new(
> schema,
> record_batch_stream.map_err(|e| DataFusionError::External(Box::new(e))),
> ));
>
> // This displays "IoObjectStore"
> info!(
> "{}",
> self.datafusion_state
> .runtime_env()
> .object_store(object_store_url.as_ref())
> .unwrap()
> );
>
> info!("writing data to disk");
> let rows_written = self
> .dedicated_rt
> .spawn(async move { data_sink.write_all(stream, &datafusion.task_ctx()).await })
> .await
> .map_err(|e| Status::data_loss(format!("error processing stream: {e}")))?
> .map_err(|e| Status::data_loss(format!("data write failed: {e}")))?;
> info!("wrote {rows_written} rows");
> ```
Maybe you need to run `data_sink.write_all(stream, &datafusion.task_ctx()).await })` somehow on the iO thread? I am nto sure what it is doing
What is the panic error message?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2661181898
Alright so I was able to strip down a reproducer for the JoinError panic but it's a little unsatisfying. I absolutely tried to start very simple, but
- I was unable to trigger the issue with the `MockStore`, the simplest I could trigger it with is running MinIO in a local docker container
- I was unable to trigger the issue with one or a small number of small record batches. The issue only crops up with longer streams of data.
I imagine it could be slightly simpler but at least for me this is the easiest thing I could come up with that consistently triggers the issue.
https://github.com/djanderson/parquet-sink-dedicated-exec-repro
Would or anyone else potentially be able to see if this reproduces the error for you? I haven't given up on running the issue down but because of the nature of what does and doesn't trigger the issue, help from more knowledgeable folks would be hugely appreciated.
This is what a small number of record batches looks like, which is a **_failure to reproduce the issue_**:
Important things to note in this case are:
- the file is successfully written to minio
- the client gets a successful response
<img width="1727" alt="image" src="https://github.com/user-attachments/assets/23e84298-fe74-4849-8c38-401015a93022" />
This is what a larger number of record batches looks like, which demonstrates the tokio panic on the server. Important things to note in this case are:
- the file fails to be written to minio
- the client gets an error response
<img width="1727" alt="image" src="https://github.com/user-attachments/assets/808e87c0-1e08-4cfb-bb17-40d3f60741f2" />
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707438841
> Sounds like a prime candidate for that object-store-utils crate :)
Yeah for sure. Now we just need to figure out where to put that (in apache or not 🤔 )
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2086779525
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
There is a lot of room for error, unfortunately
Thankfully, thanks to @tustvold and @ion-elgreco there is now a *much* better solution in ObjectStore 0.12.1 (just released today in fact)
https://github.com/apache/arrow-rs-object-store/pull/332
If you have time to update this example to use that API that would be great. Otherwise I hope to do so over the next week or two
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2706604532
> Hey all, sorry for the noise on this thread lately, I've been working to try and understand how to apply this example to an actual DataFusion-based server, and not having much luck.
@djanderson -- I played around with your example / repro and I think you may be hitting a gRPC / flight client timeout (not related to what the server is doing). I wrote up my findings here
- https://github.com/djanderson/parquet-sink-dedicated-exec-repro/issues/1
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2616139591
This looks great. I'll try to reproduce my issue and try it out.
And I have to dig a bit deeper into the topic to understand the concerns about this approach.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1936924661
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
Review Comment:
As this is a wrapper it needs to override every method and call through to the underlying impl, otherwise it will de-optimise stores with specialized impls
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "Tom-Newton (via GitHub)" <gi...@apache.org>.
Tom-Newton commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613996363
We are still having quite significant problems, but we can't reproduce it reliably and I haven't been personally working on it recently.
Anecdotally we think it's more frequent when:
1. Reading between between regions (e.g. compute in Azure US South Central reading Azure US East blob storage).
2. When spawning a large number of parallel jobs that all read the same thing using delta-rs.
3. Reading one particular table where we need to read a larger data volume. delta-rs via object-store is still only reading metadata though (order 100MB) and the larger data volumes loaded using the `pyarrow` Azure filesystem don't seem to suffer the same problem.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1970921927
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
+ self.inner
+ .as_mut()
+ .expect("paths that take put inner back")
+ .put_part(data)
Review Comment:
Wrap `IoMultipartUpload::put_part()` in `spawn_io_static`, otherwise this will panic.
```suggestion
DedicatedExecutor::spawn_io_static(
self.inner
.as_mut()
.expect("paths that take put inner back")
.put_part(data),
)
.boxed()
```
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "16pierre (via GitHub)" <gi...@apache.org>.
16pierre commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2086641885
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
How much room is there for clients to mess this up silently, aka, how hard is it to introduce silent regressions in production where we no longer use the expected (correct) threadpool ?
In particular, it seems quite easy to forget the client-side "stream driving" part.
This is the kind of regressions on which I'd want loud failures early in the development process, at CI-time for example. So here I'd find it desirable to have strengthened guarantees, for example, going from `you will likely see a panic` to `you're **guaranteed** to see a panic`
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "slyons (via GitHub)" <gi...@apache.org>.
slyons commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953478390
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
Doesn't this mean that if we want to run IO for all possible object_stores in another runtime we'd have to replace all possible prefixes?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2661294213
Incidentally, by disabling the dedicated executor in this reproducer repo, it seems to also (finally) demonstrates the problem we're looking to solve in the first place: I can consistently get it to cause a client timeout and stream RST with full data loss, even with both server and client running in release mode.
The benefit to this over other attempts at reproducing the problem (including my own previous attempts) is that this is a very simple flight server and client, straight from the examples... the type of thing a new user of the ecosystem might reasonably write to test out the "FDAP" stack. There are no artificial timeouts and the only slight cheat may be that I am pushing fully random over the pipe just to ensure that most of the data doesn't get compressed away so that there's some significant disk IO.
I need to see if this also works with the localstack S3 that @JanKaul used in his work, I had no idea that existed!
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2619873876
For your information, I was able to reproduce the IO stall error with a very [simple example](https://github.com/JanKaul/cpu-io-executor/blob/main/src/single_runtime.rs). In case anybody is interested and wants to try it out. I will test the dedicated executor tomorrow.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707110927
> My other theory was that DataFusion might start requests but not consume them in time, thus leading to timeout errors even though the response was ready. However, I was not able to find any evidence that this was actually happening.
I'd expect precisely something like this if CPU is blocking IO tasks because the IO struggles to make progress and eventually times out. Couple that with retries and you can end up in scenarios where it's hard to tell what's causing what.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707387482
After reading all the relevant feedback on this PR here is my summary:
1. Some people have found it valuable to have an ObjectStore which does I/O on a different tokio Runtime
2. The DedicatedExecutor API seems overly complicated for what it is doing.
Thus my suggested next steps are:
1. For the near term / workaround, break out the object store wrapper to pass work on different runtimes . Perhaps based on the [one in delta-rs](https://github.com/delta-io/delta-rs/blob/e30ab7e366eb209718c87acb6974a815503181bc/crates/core/src/storage/mod.rs#L116-L120)
2. Update this example to that wrapper and use a tokio::Runtime directly rather than a DedicatedExecutor
3. Work on a more general purpose IO API (see https://github.com/apache/datafusion/issues/14854 from @Xuanwo) which might make it easier to get this right without the shim / spawn io
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1931411853
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
In the S3 implementation of ObjectStore, put_part is 'tokio spawned' on the runtime of its caller (it uses a JoinSet and currently uses spawn and not spawn_on):
https://github.com/apache/arrow-rs/blob/6aaff7e38a573f797b31f89f869c3706cbe26e37/object_store/src/upload.rs#L211
If I understand the implementation and integration pattern correctly, this means that the put_part future for S3 will be spawned on the dedicated (CPU only) executor and not the IO runtime. We may have to patch the object store crate to get true IO/CPU isolation in this case.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2614047765
@alamb ah good point, I missed that! Definitely good to add, will have a better look at where these payload streams are collected
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613988466
Based on the previous discussions, and draft PRs, I ended up with this Object store wrapper to spawn the io tasks in a different handle: https://github.com/delta-io/delta-rs/blob/main/crates%2Fcore%2Fsrc%2Fstorage%2Fmod.rs#L116-L124
I do like the approach proposed in this PR, and worthwhile to have it be part of core.
Regarding testing it, since this wrapper above was added, I haven't heard anyone mentioning these issues in delta-rs anymore outside of @Tom-Newton, did you end up resolving it on your end? I forgot to follow up on it
Slightly offtopic, but I wonder if there was anything done with surfacing the true error better? Because `error decoding response body` on its own is not clear enough if you don't have the context of why it can happen, I assume the body is incomplete because the request is timed out, if that's the case shouldn't that be the error message instead?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1929545060
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
Review Comment:
Here is the IoObjectStore that implements `ObjectStore` and wraps another ObjectStore that does IO on a different thread.
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
Review Comment:
this is a monster amount of code to put in an example -- it is a large amount of tests and documentation and I wrote it with an eye towards eventually putting it into datafusion itself
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
+ self.inner
+ .as_mut()
+ .expect("paths that take put inner back")
+ .put_part(data)
+ }
+
+ async fn complete(&mut self) -> object_store::Result<PutResult> {
+ // because we are running the task on a different runtime, we
+ // can't send references to &self. Thus take out of self.inner
+ // to run on io thread
+ let mut inner = self.inner.take().expect("paths that take put inner back");
+
+ let (inner, result) = self
+ .dedicated_executor
+ .spawn_io(async move {
+ let result = inner.as_mut().complete().await;
+ (inner, result)
+ })
+ .await;
+ self.inner = Some(inner);
+ result
+ }
+
+ async fn abort(&mut self) -> object_store::Result<()> {
+ let mut inner = self.inner.take().expect("paths that take put inner back");
+
+ let (inner, result) = self
+ .dedicated_executor
+ .spawn_io(async move {
+ let result = inner.as_mut().abort().await;
+ (inner, result)
+ })
+ .await;
+ self.inner = Some(inner);
+ result
+ }
+}
+
+#[cfg(test)]
+#[allow(unused_qualifications)]
+mod tests {
+ use super::*;
+ use bytes::Bytes;
+ use futures::StreamExt;
+ use object_store::memory::InMemory;
+ use std::fmt::Formatter;
+ use std::{
+ panic::panic_any,
+ sync::{Arc, Barrier},
+ time::Duration,
+ };
+ use tokio::{net::TcpListener, sync::Barrier as AsyncBarrier};
+
+ /// Wait for the barrier and then return `result`
+ async fn do_work(result: usize, barrier: Arc<Barrier>) -> usize {
+ barrier.wait();
+ result
+ }
+
+ /// Wait for the barrier and then return `result`
+ async fn do_work_async(result: usize, barrier: Arc<AsyncBarrier>) -> usize {
+ barrier.wait().await;
+ result
+ }
+
+ /// Returns a DedicatedExecutor with a single thread and IO enabled
+ fn exec() -> DedicatedExecutor {
+ exec_with_threads(1)
+ }
+
+ fn exec2() -> DedicatedExecutor {
+ exec_with_threads(2)
+ }
+
+ fn exec_with_threads(threads: usize) -> DedicatedExecutor {
+ let mut runtime_builder = Builder::new_multi_thread();
+ runtime_builder.worker_threads(threads);
+ runtime_builder.enable_all();
+
+ DedicatedExecutorBuilder::from(runtime_builder)
+ .with_name("Test DedicatedExecutor")
+ .build()
+ }
+
+ async fn test_io_runtime_multi_thread_impl(dedicated: DedicatedExecutor) {
+ let io_runtime_id = std::thread::current().id();
+ dedicated
+ .spawn(async move {
+ let dedicated_id = std::thread::current().id();
+ let spawned = DedicatedExecutor::spawn_io_static(async move {
+ std::thread::current().id()
+ })
+ .await;
+
+ assert_ne!(dedicated_id, spawned);
+ assert_eq!(io_runtime_id, spawned);
+ })
+ .await
+ .unwrap();
+ }
+
+ #[tokio::test]
+ async fn basic() {
+ let barrier = Arc::new(Barrier::new(2));
+
+ let exec = exec();
+ let dedicated_task = exec.spawn(do_work(42, Arc::clone(&barrier)));
+
+ // Note the dedicated task will never complete if it runs on
+ // the main tokio thread (as this test is not using the
+ // 'multithreaded' version of the executor and the call to
+ // barrier.wait actually blocks the tokio thread)
+ barrier.wait();
+
+ // should be able to get the result
+ assert_eq!(dedicated_task.await.unwrap(), 42);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn basic_clone() {
+ let barrier = Arc::new(Barrier::new(2));
+ let exec = exec();
+ // Run task on clone should work fine
+ let dedicated_task = exec.clone().spawn(do_work(42, Arc::clone(&barrier)));
+ barrier.wait();
+ assert_eq!(dedicated_task.await.unwrap(), 42);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn drop_empty_exec() {
+ exec();
+ }
+
+ #[tokio::test]
+ async fn drop_clone() {
+ let barrier = Arc::new(Barrier::new(2));
+ let exec = exec();
+
+ drop(exec.clone());
+
+ let task = exec.spawn(do_work(42, Arc::clone(&barrier)));
+ barrier.wait();
+ assert_eq!(task.await.unwrap(), 42);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ #[should_panic(expected = "foo")]
+ async fn just_panic() {
+ struct S(DedicatedExecutor);
+
+ impl Drop for S {
+ fn drop(&mut self) {
+ self.0.join().now_or_never();
+ }
+ }
+
+ let exec = exec();
+ let _s = S(exec);
+
+ // this must not lead to a double-panic and SIGILL
+ panic!("foo")
+ }
+
+ #[tokio::test]
+ async fn multi_task() {
+ let barrier = Arc::new(Barrier::new(3));
+
+ // make an executor with two threads
+ let exec = exec2();
+ let dedicated_task1 = exec.spawn(do_work(11, Arc::clone(&barrier)));
+ let dedicated_task2 = exec.spawn(do_work(42, Arc::clone(&barrier)));
+
+ // block main thread until completion of other two tasks
+ barrier.wait();
+
+ // should be able to get the result
+ assert_eq!(dedicated_task1.await.unwrap(), 11);
+ assert_eq!(dedicated_task2.await.unwrap(), 42);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn tokio_spawn() {
+ let exec = exec2();
+
+ // spawn a task that spawns to other tasks and ensure they run on the dedicated
+ // executor
+ let dedicated_task = exec.spawn(async move {
+ // spawn separate tasks
+ let t1 = tokio::task::spawn(async { 25usize });
+ t1.await.unwrap()
+ });
+
+ // Validate the inner task ran to completion (aka it did not panic)
+ assert_eq!(dedicated_task.await.unwrap(), 25);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn panic_on_executor_str() {
+ let exec = exec();
+ let dedicated_task = exec.spawn(async move {
+ if true {
+ panic!("At the disco, on the dedicated task scheduler");
+ } else {
+ 42
+ }
+ });
+
+ // should not be able to get the result
+ let err = dedicated_task.await.unwrap_err();
+ assert_eq!(
+ err.to_string(),
+ "Panic: At the disco, on the dedicated task scheduler",
+ );
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn panic_on_executor_string() {
+ let exec = exec();
+ let dedicated_task = exec.spawn(async move {
+ if true {
+ panic!("{} {}", 1, 2);
+ } else {
+ 42
+ }
+ });
+
+ // should not be able to get the result
+ let err = dedicated_task.await.unwrap_err();
+ assert_eq!(err.to_string(), "Panic: 1 2",);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn panic_on_executor_other() {
+ let exec = exec();
+ let dedicated_task = exec.spawn(async move {
+ if true {
+ panic_any(1)
+ } else {
+ 42
+ }
+ });
+
+ // should not be able to get the result
+ let err = dedicated_task.await.unwrap_err();
+ assert_eq!(err.to_string(), "Panic: unknown internal error",);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn executor_shutdown_while_task_running() {
+ let barrier_1 = Arc::new(Barrier::new(2));
+ let captured_1 = Arc::clone(&barrier_1);
+ let barrier_2 = Arc::new(Barrier::new(2));
+ let captured_2 = Arc::clone(&barrier_2);
+
+ let exec = exec();
+ let dedicated_task = exec.spawn(async move {
+ captured_1.wait();
+ do_work(42, captured_2).await
+ });
+ barrier_1.wait();
+
+ exec.shutdown();
+ // block main thread until completion of the outstanding task
+ barrier_2.wait();
+
+ // task should complete successfully
+ assert_eq!(dedicated_task.await.unwrap(), 42);
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn executor_submit_task_after_shutdown() {
+ let exec = exec();
+
+ // Simulate trying to submit tasks once executor has shutdown
+ exec.shutdown();
+ let dedicated_task = exec.spawn(async { 11 });
+
+ // task should complete, but return an error
+ let err = dedicated_task.await.unwrap_err();
+ assert_eq!(
+ err.to_string(),
+ "Worker thread gone, executor was likely shut down"
+ );
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn executor_submit_task_after_clone_shutdown() {
+ let exec = exec();
+
+ // shutdown the clone (but not the exec)
+ exec.clone().join().await;
+
+ // Simulate trying to submit tasks once executor has shutdown
+ let dedicated_task = exec.spawn(async { 11 });
+
+ // task should complete, but return an error
+ let err = dedicated_task.await.unwrap_err();
+ assert_eq!(
+ err.to_string(),
+ "Worker thread gone, executor was likely shut down"
+ );
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn executor_join() {
+ let exec = exec();
+ // test it doesn't hang
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn executor_join2() {
+ let exec = exec();
+ // test it doesn't hang
+ exec.join().await;
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ #[allow(clippy::redundant_clone)]
+ async fn executor_clone_join() {
+ let exec = exec();
+ // test it doesn't hang
+ exec.clone().join().await;
+ exec.clone().join().await;
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn drop_receiver() {
+ // create empty executor
+ let exec = exec();
+
+ // create first blocked task
+ let barrier1_pre = Arc::new(AsyncBarrier::new(2));
+ let barrier1_pre_captured = Arc::clone(&barrier1_pre);
+ let barrier1_post = Arc::new(AsyncBarrier::new(2));
+ let barrier1_post_captured = Arc::clone(&barrier1_post);
+ let dedicated_task1 = exec.spawn(async move {
+ barrier1_pre_captured.wait().await;
+ do_work_async(11, barrier1_post_captured).await
+ });
+ barrier1_pre.wait().await;
+
+ // create second blocked task
+ let barrier2_pre = Arc::new(AsyncBarrier::new(2));
+ let barrier2_pre_captured = Arc::clone(&barrier2_pre);
+ let barrier2_post = Arc::new(AsyncBarrier::new(2));
+ let barrier2_post_captured = Arc::clone(&barrier2_post);
+ let dedicated_task2 = exec.spawn(async move {
+ barrier2_pre_captured.wait().await;
+ do_work_async(22, barrier2_post_captured).await
+ });
+ barrier2_pre.wait().await;
+
+ // cancel task
+ drop(dedicated_task1);
+
+ // cancelation might take a short while
+ tokio::time::timeout(Duration::from_secs(1), async {
+ loop {
+ if Arc::strong_count(&barrier1_post) == 1 {
+ return;
+ }
+ tokio::time::sleep(Duration::from_millis(10)).await
+ }
+ })
+ .await
+ .unwrap();
+
+ // unblock other task
+ barrier2_post.wait().await;
+ assert_eq!(dedicated_task2.await.unwrap(), 22);
+ tokio::time::timeout(Duration::from_secs(1), async {
+ loop {
+ if Arc::strong_count(&barrier2_post) == 1 {
+ return;
+ }
+ tokio::time::sleep(Duration::from_millis(10)).await
+ }
+ })
+ .await
+ .unwrap();
+
+ exec.join().await;
+ }
+
+ #[tokio::test]
+ async fn test_io_runtime_multi_thread() {
+ let mut runtime_builder = tokio::runtime::Builder::new_multi_thread();
+ runtime_builder.worker_threads(1);
+
+ let dedicated = DedicatedExecutorBuilder::from(runtime_builder)
+ .with_name("Test DedicatedExecutor")
+ .build();
+ test_io_runtime_multi_thread_impl(dedicated).await;
+ }
+
+ #[tokio::test]
+ async fn test_io_runtime_current_thread() {
+ let runtime_builder = tokio::runtime::Builder::new_current_thread();
+
+ let dedicated = DedicatedExecutorBuilder::new_from_builder(runtime_builder)
+ .with_name("Test DedicatedExecutor")
+ .build();
+ test_io_runtime_multi_thread_impl(dedicated).await;
+ }
+
+ #[tokio::test]
+ async fn test_that_default_executor_prevents_io() {
+ let exec = DedicatedExecutorBuilder::new().build();
+
+ let io_disabled = exec
+ .spawn(async move {
+ // the only way (I've found) to test if IO is enabled is to use it and observe if tokio panics
+ TcpListener::bind("127.0.0.1:0")
+ .catch_unwind()
+ .await
+ .is_err()
+ })
+ .await
+ .unwrap();
+
+ assert!(io_disabled)
+ }
+
+ #[tokio::test]
+ async fn test_happy_path() {
+ let rt_io = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .enable_all()
+ .build()
+ .unwrap();
+
+ let io_thread_id = rt_io
+ .spawn(async move { std::thread::current().id() })
+ .await
+ .unwrap();
+ let parent_thread_id = std::thread::current().id();
+ assert_ne!(io_thread_id, parent_thread_id);
+
+ DedicatedExecutor::register_io_runtime(Some(rt_io.handle().clone()));
+
+ let measured_thread_id =
+ DedicatedExecutor::spawn_io_static(
+ async move { std::thread::current().id() },
+ )
+ .await;
+ assert_eq!(measured_thread_id, io_thread_id);
+
+ rt_io.shutdown_background();
+ }
+
+ #[tokio::test]
+ #[should_panic(expected = "IO runtime registered")]
+ async fn test_panic_if_no_runtime_registered() {
+ DedicatedExecutor::spawn_io_static(futures::future::ready(())).await;
+ }
+
+ #[tokio::test]
+ #[should_panic(expected = "IO runtime was shut down")]
+ async fn test_io_runtime_down() {
+ let rt_io = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(1)
+ .enable_all()
+ .build()
+ .unwrap();
+
+ DedicatedExecutor::register_io_runtime(Some(rt_io.handle().clone()));
+
+ tokio::task::spawn_blocking(move || {
+ rt_io.shutdown_timeout(Duration::from_secs(1));
+ })
+ .await
+ .unwrap();
+
+ DedicatedExecutor::spawn_io_static(futures::future::ready(())).await;
+ }
+
+ // -----------------------------------
+ // ----- Tests for IoObjectStore ------
Review Comment:
Here are tests that prove the IO is actually happening on the correct thread (both the original requests as well as the streams that are returned)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1929544551
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
Review Comment:
This is the example that shows the best practice for separating CPU and IO
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2706983051
I have also made a reproducer showing how timeouts happen with slower internet connections / certain files (because DataFusion is making single large requests which can't be completed within a timeout):
- https://github.com/apache/datafusion/issues/15067
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2708187238
> > > Sounds like a prime candidate for that object-store-utils crate :)
> >
> >
> > Yeah for sure. Now we just need to figure out where to put that (in apache or not 🤔 )
>
> I think it can be beneficial in Apache but I don't know these processes tbh.
>
> It would also be great to add caching object store implementations in there as well, such as this one https://github.com/slatedb/slatedb/blob/main/src%2Fcached_object_store%2Fobject_store.rs (or one of the other many out there :p)
- I filed a ticket for discussion about a crate with more object store utilities: https://github.com/apache/arrow-rs/issues/7251
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "MarcPerrier (via GitHub)" <gi...@apache.org>.
MarcPerrier commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2097962631
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
Review Comment:
We are currently trying to implement a similar pattern and I wondered why the IO_RUNTIME is given through a thread local and not through an inner variable. Is it to have access to stronger encapsulation and/or useful behaviour (like spawn_io_static) or is there a correctness reason 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.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb closed pull request #14286: Example for using a separate threadpool for CPU bound work (try 2)
URL: https://github.com/apache/datafusion/pull/14286
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1943824052
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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_cpu_sendable_record_batch_stream(stream);
+
+ // Note you can run other streams on the DedicatedExecutor as well using the
+ // same 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),
Review Comment:
I'm guessing you mean you can send the `FlightDataEncoder` via `run_cpu_stream` instead of "using the same function", `run_cpu_sendable_record_batch_stream`. At least, I can't quite make that work. `FlightDataEncoder` has a method `.known_schema()`, but not `.schema()`, so it appears to not implement the `RecordBatchStream` trait.
I got it compiling like this (`batch_stream` is my `SendableRecordBatchStream`):
```rust
let flight_data_encoder = FlightDataEncoderBuilder::new()
.with_schema(schema)
.with_options(write_options)
.build(batch_stream.map_err(|e| FlightError::ExternalError(Box::new(e))));
let flight_data_stream = self
.dedicated_executor
.run_cpu_stream(flight_data_encoder, |job_error| {
FlightError::from_external_error(Box::new(job_error))
})
.map_err(Status::from)
.boxed();
return Ok(Response::new(flight_data_stream));
```
but lmk if there's a cleaner way to pass the encoder to the dedicated executor, as this comment suggests.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2708240200
FWIW the new HttpClient abstraction, introduced in ObjectStore 0.12, provides a potentially nicer way to spawn IO on to a separate runtime - https://github.com/apache/arrow-rs/pull/7253
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707200014
> > My other theory was that DataFusion might start requests but not consume them in time, thus leading to timeout errors even though the response was ready. However, I was not able to find any evidence that this was actually happening.
>
> I'd expect precisely something like this if CPU is blocking IO tasks because the IO struggles to make progress and eventually times out. Couple that with retries and you can end up in scenarios where it's hard to tell what's causing what.
One thing that I noticed is that DataFusion uses `ObjectStore::get_ranges` which returns buffered `Bytes` -- so from the network perspective I think the data is being consumed 🤔
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1931656149
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
I think this just needs to be wrapped in spawn_io, the overhead of the additional tokio task will be irrelevant when compared to the cross RT and IO overheads.
_The fact this is so hard to spot is a good example of why this pattern is really fragile_
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2626852759
> As this problem seems to only occur when executing the physical plan.
Assuming this was a mistake we made at InfluxData, we had issues with plan time concurrency that were caused by this also, they were just much harder to track down (it took ~3 months).
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707451759
> > Sounds like a prime candidate for that object-store-utils crate :)
>
> Yeah for sure. Now we just need to figure out where to put that (in apache or not 🤔 )
I think it can be beneficial in Apache but I don't know these processes tbh.
It would also be great to add caching object store implementations in there as well, such as this one https://github.com/slatedb/slatedb/blob/main/src%2Fcached_object_store%2Fobject_store.rs (or one of the other many out there :p)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953490590
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
I guess technically there's a `file://` handler registered by default, so I lied :)
https://github.com/apache/datafusion/blob/82461b74e0be01be8ac17d878dec53140e4209f0/datafusion/execution/src/object_store.rs#L201
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953482293
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
> replace all possible prefixes
The list of registered object stores is empty by default and you only explicitly register the stores you are using.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2917020530
BTW the SpawnService is what should be used now: https://github.com/apache/arrow-rs-object-store/pull/332
Sadly, the docs are broken for the current version of object_store so I can't link the docs to it.
I'll try and find time to update this PR (or make a new one) with the API
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2624213017
I took a look at your example @JanKaul, thank you for this as I think it very nicely demonstrates the challenge of shimming IO at the object store interface. Unfortunately I think this may be a different issue from the runtime stall issue.
At a high level the code is doing this
```
object_store.get().into_stream().map(|x| {
sleep(Duration::from_seconds(2));
}).try_collect()
```
The problem is this introduces 2 second pauses between polling the streaming request from object storage. Regardless of runtime setup, holding a request across a CPU bound task will result in timeouts, as backpressure will eventually cause the sender to timeout.
There are two broad strategies to avoid this:
1. Add an unbounded queue on the stream output
2. Make separate requests each time for the work that can be performed, e.g. first 1MB, perform work, then fetch next MB
Option 1. will potentially buffer the entire input stream if the consumer is running behind, but should ensure the request runs to completion. Option 2. is I think the better approach, and is what code should probably be doing (FWIW this is what the parquet reader does).
_As an aside I also took the time to port over the custom executor approach, to show what that looks like - https://github.com/JanKaul/cpu-io-executor/pull/1. I still think this is a cleaner approach, although much like the DedicatedExecutor in this PR, it doesn't help with the issue this example is running into_
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1935223288
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
It shouldn't matter that they're written out of order as they're assembled based on the index captured when the future is created. Could you perhaps share your code?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2624431385
Thanks a lot for having a look! I'll try to adopt the strategy 2 to simulate the behavior of the parquet reader.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "davidhewitt (via GitHub)" <gi...@apache.org>.
davidhewitt commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2615483188
> This is neat -- I actually like that it is in terms of just a tokio runtime rather than somehthing like `DedicatedExecutor`. I might try and work on IoObjectStore to follow that pattern instead
I think that makes sense to me too. When I see the current `IoObjectStore` containing a `DedicatedExecutor` and then calling `spawn_io`, the abstraction feels muddled. Having just a handle to a runtime to spawn works on seems simpler and generalises better.
> In the interests of avoiding confusion, as my objections appear to have gotten a little misinterpreted, I'd like to clarify the fact this approach comes with non-trivial overheads is **not** what concerns me with this approach. Rather that we know from experience at InfluxData that this pattern is fragile, easy to mess up, and leads to emergent behaviour that is highly non-trivial to reproduce and debug.
Just to put the brakes on rushing this into core too quickly, I have to support @tustvold in raising concern; we have had plenty of issues at Pydantic with this pattern where we have missed cases where we should have spawned IO (or CPU work) onto the other runtime.
It seems to me like most folks agree research into schedulers for the CPU work that can sit within the single tokio runtime would be _much_ easier to integrate for downstream use cases, probably at the cost of complexity (overheads?) within datafusion itself.
We have already vendored this pattern and have it working for us, you don't need to rush this through for us.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1943824052
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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_cpu_sendable_record_batch_stream(stream);
+
+ // Note you can run other streams on the DedicatedExecutor as well using the
+ // same 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),
Review Comment:
I'm guessing you mean you can send the `FlightDataEncoder` via `run_cpu_stream` instead of "using the same function", `run_cpu_sendable_record_batch_stream`. At least, I can't quite make that work. `FlightDataEncoder` has a method `.known_schema()`, but not `.schema()`, so it appears to not implement the `RecordBatchStream` trait.
I got it compiling like this (`batch_stream` is my `SendableRecordBatchStream`):
```rust
let flight_data_encoder = FlightDataEncoderBuilder::new()
.with_schema(schema)
.with_options(write_options)
.build(batch_stream.map_err(|e| FlightError::ExternalError(Box::new(e))));
let flight_data_stream = self
.dedicated_executor
.run_cpu_stream(flight_data_encoder, |job_error| {
FlightError::from_external_error(Box::new(job_error))
})
.map_err(Status::from)
.boxed();
return Ok(Response::new(flight_data_stream));
```
but lmk if there's a cleaner way to pass the encoder to the dedicated executor, as this comment suggests.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613975126
In the interests of avoiding confusion, as my objections appear to have gotten a little misinterpreted, I'd like to clarify the fact this approach comes with non-trivial overheads is **not** what concerns me with this approach. Rather that we know from experience at InfluxData that this pattern is fragile, easy to mess up, and leads to emergent behaviour that is highly non-trivial to reproduce and debug.
That being said as Andrew says, nobody has emerged who is able/willing to resolve this with a more holistic approach, e.g. something closer to what polars/DuckDb/Hyper are doing to separate IO/compute, and so proceeding with something is better than nothing. I just hoped someone might step up.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1929824035
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
Review Comment:
I'll note that this clobbers any `on_thread_start` set on `runtime_builder`. Unfortunately it is not stored publicly on `RuntimeBuilder` so it is not possible to pull off the existing setting and wrap it. I suggest we add an `on_thread_start` method to `DedicatedExecutorBuilder`?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2706993616
I have one more theory I want to chase down this afternoon and then I will write up my thoughts on next steps here
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2086999273
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
@alamb, while working on: https://github.com/delta-io/delta-rs/pull/3426 to get it in delta-rs, I realized that on some object_stores (HDFS for example) you still might want to do a wrapper style because they don't have a HttpConnector to be configured
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2621226726
I tried the dedicated executor in my [test example](https://github.com/JanKaul/cpu-io-executor/blob/main/src/two_runtimes.rs). I'm not entirely sure if I'm using it wrong or maybe my test is too contrived. But I'm still getting a single Timeout Error. This is already much better compared to when using a single runtime, which produces many TimeOut Errors.
I invite everybody to have a look, maybe you can spot a bug.
I will try the dedicated executor with a real example. The problem is that it is difficult to reproduce because the issue only occurred at a certain size.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1931411853
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
In the S3 implementation of ObjectStore, put_part is 'tokio spawned' on the runtime of its caller (it uses a JoinSet and currently uses spawn and not spawn_on):
https://github.com/apache/arrow-rs/blob/6aaff7e38a573f797b31f89f869c3706cbe26e37/object_store/src/upload.rs#L211
If I understand the implementation and integration pattern correctly, this means that the put_part future for S3 will be spawned on the dedicated (CPU only) executor and not the IO runtime. We may have to patch the object store crate to get true IO/CPU utilization in this case.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1934929386
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
I'm trying to wrap put_part with spawn_io, but I'm encountering some issues. For multipart uploads, I believe the part_idx needs to be incremented in the same order that PutPayload instances are generated from the stream being written. The indirection introduced by sending the put_part future to the I/O runtime via spawn_io seems to cause the payloads to be written out of order.
As a result, I'm seeing failures in complete() that indicate parts with a size smaller than 5MB are being sent before the last part.
I don't see actually anything in the standard S3MultipartUpload implementation that prevents this behavior, but I haven't encountered any failures related to this in the past, so I may be mistaken.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1936870081
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
Thank you, the test was very helpful and clearly illustrates that even though the async part uploads run concurrently, the synchronous part idx assignment happens serially on the same thread. I accidentally changed this behavior by wrapping inner in an arc mutex (to fix some borrow checker issues) and having all of inner put_parts() contend for the mutex, which broke the ordering guarantees. Fixed it.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1936870081
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
Thank you, the test was very helpful and clearly illustrates that even though the async part uploads run concurrently, the synchronous part idx assignment happens serially on the same thread. I accidentally changed this behavior by wrapping inner in an arc mutex and having all of inner put_parts() contend for the mutex, which broke the ordering guarantees. Fixed it.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2655439836
I checked out this PR and tried to add a test to reproduce the issue. Including here although it's definitely too dependency heavy for a unit test:
<details>
#[tokio::test]
async fn test_parquet_sink_io_doesnt_panic() {
use arrow::array as arrow_array; // needed for record_batch! macro
use datafusion::datasource::file_format::parquet::ParquetSink;
use datafusion::datasource::listing::ListingTableUrl;
use datafusion::datasource::physical_plan::FileSink;
use datafusion::datasource::physical_plan::FileSinkConfig;
use datafusion::execution::object_store::ObjectStoreUrl;
use datafusion::prelude::SessionContext;
use datafusion_expr::dml::InsertOp;
let exec = DedicatedExecutorBuilder::new().build();
let store = MockStore::create().await;
let store = exec.wrap_object_store_for_io(store);
let batch = record_batch!(("col", Int32, vec![1, 2, 3])).unwrap();
let schema = batch.schema();
let stream_adapter = RecordBatchStreamAdapter::new(
batch.schema(),
futures::stream::iter(vec![Ok(batch)]),
);
let stream: SendableRecordBatchStream = Box::pin(stream_adapter);
let ctx = SessionContext::new();
let object_store_url = ObjectStoreUrl::local_filesystem();
ctx.register_object_store(object_store_url.as_ref(), store);
// Configure sink
let file_sink_config = FileSinkConfig {
object_store_url,
file_groups: vec![],
table_paths: vec![ListingTableUrl::parse("/test.parquet").unwrap()],
output_schema: schema,
table_partition_cols: vec![],
insert_op: InsertOp::Overwrite,
keep_partition_by_columns: false,
file_extension: String::from("parquet"),
};
let table_options = Default::default();
let data_sink = ParquetSink::new(file_sink_config, table_options);
// Execute write on dedicated runtime
exec.spawn(
async move { data_sink.write_all(stream, &ctx.task_ctx()).await.unwrap() },
)
.await
.unwrap();
exec.join().await
}
</details>
In the end it did not reproduce the issue I'm seeing locally. I noticed that @mertak-synnada recently refactored `ParquetSink.write_all()` so I updated my flightsql server to 45.0.0 to pull in his changes, but... no joy.
Will keep investigating but no clear leads right now, other than that I'm running my object store against a local MinIO container.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1965798308
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
@rohitrastogi , could you suggest the edit to this PR that resolved your issue with multipart uploads (or at least drop in a patch as a comment)? I'm just now realizing that the issue I'm seeing (https://github.com/apache/datafusion/pull/14286#issuecomment-2661181898) is the exact same that you were discussing... took me a while to make the connection, but now making sense that there is some threshold above which objectstore switches from a single "put" to a multipart, and that is exactly when I'm seeing the IO error, as you already reported.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707071207
> I have one more theory I want to chase down this afternoon and then I will write up my thoughts on next steps here
My other theory was that DataFusion might start requests but not consume them in time, thus leading to timeout errors even though the response was ready. However, I was not able to find any evidence that this was actually happening
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2707431726
> After reading all the relevant feedback on this PR here is my summary:
> 1. Some people have found it valuable to have an ObjectStore which does I/O on a different tokio Runtime
> 2. The DedicatedExecutor API seems overly complicated for what it is doing.
>
> Thus my suggested next steps are:
> 1. For the near term / workaround, break out the object store wrapper to pass work on different runtimes . Perhaps based on the [one in delta-rs](https://github.com/delta-io/delta-rs/blob/e30ab7e366eb209718c87acb6974a815503181bc/crates/core/src/storage/mod.rs#L116-L120)
> 2. Update this example to that wrapper and use a tokio::Runtime directly rather than a DedicatedExecutor
> 3. Work on a more general purpose IO API (see https://github.com/apache/datafusion/issues/14854 from @Xuanwo) which might make it easier to get this right without the shim / spawn io
Sounds like a prime candidate for that object-store-utils crate :)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1968617896
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
Review Comment:
```suggestion
pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>> + use<T>
```
Use new [RPIT lifetime precise capturing](https://doc.rust-lang.org/nightly/edition-guide/rust-2024/rpit-lifetime-capture.html) for compatibility with newly released rust 2024 edition. Otherwise this example fails to compile.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2614037782
> Based on the previous discussions, and draft PRs, I ended up with this Object store wrapper to spawn the io tasks in a different handle: https://github.com/delta-io/delta-rs/blob/main/crates%2Fcore%2Fsrc%2Fstorage%2Fmod.rs#L116-L124
This is neat -- I actually like that it is in terms of just a tokio runtime rather than somehthing like `DedicatedExecutor`. I might try and work on IoObjectStore to follow that pattern instead
One thing I noticed is that it doesn't actually shim the get result stream (so while the intial GET request is on the other runtime, when reading data from the stream that will still happen on the same runtime)
I handled this by wrapping the stream like this:
```rust
/// Wrap the result stream if necessary
fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
match payload {
GetResultPayload::File(_, _) => payload,
GetResultPayload::Stream(stream) => {
let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
GetResultPayload::Stream(new_stream)
}
}
}
```
It might be an interesting thing to try in delta.rs :thin'
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1929544864
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
Review Comment:
This pattern is @tustvold 's core concern as I understand -- that the overhead of transferring data back/forth between runtimes by wrapping streams could be avoided with lower level hooks for IO
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1929826332
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
Review Comment:
It'd be nice if there was also a way to warn users that we are clobbering their configuration...
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2626715496
@tustvold I've adopted the [example](https://github.com/JanKaul/cpu-io-executor/blob/main/src/buffered_stream.rs) to somewhat resemble the parquet reader. It tries to mimic its behavior by reading the stream from the object store and creating a stream of batched data. It would be great if you could have a look. The goal is to mimic the parquet readers behavior from an async perspective. To be able to test different approaches.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "rohitrastogi (via GitHub)" <gi...@apache.org>.
rohitrastogi commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1934929386
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
I'm trying to wrap put_part with spawn_io, but I'm encountering some issues. For multipart uploads, I believe the part_idx needs to be incremented in the same order that PutPayload instances are generated from the stream being written. The indirection introduced by sending the put_part future to the I/O runtime via spawn_io seems to cause the payloads to be written out of order.
As a result, I'm seeing failures in complete() that indicate parts with a size smaller than 5MB are being sent before the last part.
I don't see anything in the standard S3MultipartUpload implementation that prevents this behavior, but I haven't encountered any failures related to this in the past, so I may be mistaken.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2684139203
Hey all, sorry for the noise on this thread lately, I've been working to try and understand how to apply this example to an actual DataFusion-based server, and not having much luck.
I now have the dedicated executor from this PR fully integrated and working in my [executor starvation reproducer](https://github.com/djanderson/parquet-sink-dedicated-exec-repro), however, I'm not observing a significant improvement. I see others mentioning that a similar approach has worked okay for them in the past, which makes me question what I'm seeing. I could really use another set of eyes on it.
I'm intentionally exercising the ingest path. It seems to more consistently exhibit the starvation issues, and the result of hitting the issue on ingest is potential data loss.... definitely something a database should avoid 😬
Does it show what I think it shows, or _might_ it be some other phenomenon? If it shows 1) consistent executor starvation and 2) that this PR doesn't seem to improve it and 3) that we therefore have no working example as a community of how to ingest data into a server-deployed datafusion-based database, would we consider that a significant issue? Maybe even worth calling out as a [priority on the roadmap](https://github.com/apache/datafusion/issues/14580)?
As always, willing to keep testing with the threadpool approach or other suggestions. LMK.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1970921927
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
+ self.inner
+ .as_mut()
+ .expect("paths that take put inner back")
+ .put_part(data)
Review Comment:
```suggestion
DedicatedExecutor::spawn_io_static(
self.inner
.as_mut()
.expect("paths that take put inner back")
.put_part(data),
)
.boxed()
```
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2685741602
> Hey all, sorry for the noise on this thread lately, I've been working to try and understand how to apply this example to an actual DataFusion-based server, and not having much luck.
>
> I now have the dedicated executor from this PR fully integrated and working in my [executor starvation reproducer](https://github.com/djanderson/parquet-sink-dedicated-exec-repro), however, I'm not observing a significant improvement. I see others mentioning that a similar approach has worked okay for them in the past, which makes me question what I'm seeing. I could really use another set of eyes on it.
>
> I'm intentionally exercising the ingest path. It seems to more consistently exhibit the starvation issues, and the result of hitting the issue on ingest is potential data loss.... definitely something a database should avoid 😬
>
> Does it show what I think it shows, or _might_ it be some other phenomenon? If it shows 1) consistent executor starvation and 2) that this PR doesn't seem to improve it and 3) that we therefore have no working example as a community of how to ingest data into a server-deployed datafusion-based database, would we consider that a significant issue? Maybe even worth calling out as a [priority on the roadmap](https://github.com/apache/datafusion/issues/14580)?
>
> As always, willing to keep testing with the threadpool approach or other suggestions. LMK.
Thank you @djanderson -- I hadn't seen https://github.com/djanderson/parquet-sink-dedicated-exec-repro
I hope to get back to this issue in the next week or two and try and come up with the next steps based on all the feedback so far.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2720973005
After thinking more about this more, what i am currently thinking is:
1. Wait until https://github.com/apache/arrow-rs/pull/7253 is available in DataFusion (likely April)
2. Update this example to use that and another runtime directly (rather than dedicated executor, etc)
In parallel I think discussions on how to get some more optimized implemenations for object stores can proceed:
- https://github.com/apache/arrow-rs/issues/7251
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "ion-elgreco (via GitHub)" <gi...@apache.org>.
ion-elgreco commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2086999273
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
@alamb, while working on: https://github.com/delta-io/delta-rs/pull/3426 to get it in delta-rs, I realized that one some stores you still might want to do a wrapper style because they don't have a HttpConnector to be configured
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2654795222
Is there extra work I'd need to do to get this to work with the `ParquetSink`? This isn't a full reproducer but just a quick copy of the relevant section of the `do_put_statement_ingest()` method in my flightsql server.
I get
> flight_sql: writing data to disk
thread 'tokio-runtime-worker' panicked at /Users/dja/.rustup/toolchains/stable-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/ops/function.rs:250:5:
A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
```rust
let data_sink = ParquetSink::new(file_sink_config, table_options);
let record_batch_stream =
FlightRecordBatchStream::new_from_flight_data(flight_data_stream.map_err(|e| e.into()));
let stream = Box::pin(RecordBatchStreamAdapter::new(
schema,
record_batch_stream.map_err(|e| DataFusionError::External(Box::new(e))),
));
// This displays "IoObjectStore"
info!(
"{}",
self.datafusion_state
.runtime_env()
.object_store(object_store_url.as_ref())
.unwrap()
);
info!("writing data to disk");
let rows_written = self
.dedicated_rt
.spawn(async move { data_sink.write_all(stream, &datafusion.task_ctx()).await })
.await
.map_err(|e| Status::data_loss(format!("error processing stream: {e}")))?
.map_err(|e| Status::data_loss(format!("data write failed: {e}")))?;
info!("wrote {rows_written} rows");
```
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "slyons (via GitHub)" <gi...@apache.org>.
slyons commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953478390
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
Doesn't this mean that if we want to run IO for all possible object_stores in another runtime we'd have to replace all possible prefixes?
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2954152637
Here is an updated version of the example, using the new ObjectStore API:
- https://github.com/apache/datafusion/pull/16331
Let's continue the conversation there
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2098291930
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
Review Comment:
The reason it is a thread local / static is that the original location of this code (influxdb_iox`) has several IO calls where we didn't have access to the `DedicatedExecutor` and we needed to avoid plumbing a reference down all the way
If you can do it with an inner reference I think that is much better
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "16pierre (via GitHub)" <gi...@apache.org>.
16pierre commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r2099735037
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
+
+ // Plan and begin to execute the query on the dedicated runtime
+ let stream = dedicated_executor
+ .spawn(async move {
+ // Plan / execute the query
+ let url = "https://github.com/apache/arrow-testing/raw/master/data/csv/aggregate_test_100.csv";
+ let df = ctx
+ .sql(&format!("SELECT c1,c2,c3 FROM '{url}' LIMIT 5"))
+ .await?;
+
+ // since we wrapped the object store, and I/O will actually happen
+ // on the current runtime.
+ let stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ Ok(stream) as Result<_>
+ }).await??;
+
+ // We have now planned the query on the dedicated executor Yay! However,
+ // most applications will still drive the stream (aka call `next()` to get
+ // the results) from the current runtime, for example to send results back
+ // over arrow-flight.
+
+ // However, as mentioned above, calling `next()` drives 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
+ // the `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
Review Comment:
> the challenge is a way to actually achieve this in practice in a way that doesn't involve boiling the ocean
Agreed solving the full problem is non trivial. Looping back on my previous comment `I wonder if there's room for forcing functions like wrappers...`, after inspecting the [influxdb3_core](https://github.com/influxdata/influxdb3_core) codebase (thanks for opensourcing btw), `IOxSessionContext` seems to roughly match what I had in mind here: by wrapping the DF session with a custom session with correctly configured threadpools, we can avoid exposing the raw DF session execution methods, and achieve forcing functions on our client callsites.
We'll probably mimick this pattern with a homemade wrapper; would be cool if such patterns were shipped out-of-the-box in DF.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2655024463
> What is the panic error message?
Panic error message is just
> error=status: DataLoss, message: "data write failed: IO error: Error joining spawned task: task 93 panicked with message \"A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.\"
> Maybe you need to run data_sink.write_all(stream, &datafusion.task_ctx()).await }) somehow on the iO thread? I am nto sure what it is doing
Yeah, without dedicated executor, this is where I hit executor starvation, so I can't wrap the whole call in spawn_io, I don't think... otherwise it would be the same as no having the dedicated executor at all.
I thought this would work first review because ParquetSink.write_all [retrieves the object_store from the runtime_env](https://github.com/apache/datafusion/blob/branch-44/datafusion/core/src/datasource/file_format/parquet.rs#L793) and and uses it to initialize the async arrow writer.
Maybe there's some interaction with the `file_write_tasks` JoinSet that I don't understand?
https://github.com/apache/datafusion/blob/branch-44/datafusion/core/src/datasource/file_format/parquet.rs#L865
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "djanderson (via GitHub)" <gi...@apache.org>.
djanderson commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1953490590
##########
datafusion-examples/examples/thread_pools.rs:
##########
@@ -0,0 +1,238 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! This example shows how to use separate thread pools (tokio [`Runtime`]))s to
+//! run the IO and CPU intensive parts of DataFusion plans.
+//!
+//! # Background
+//!
+//! DataFusion, by default, plans and executes all operations (both CPU and IO)
+//! on the same thread pool. This makes it fast and easy to get started, but
+//! can cause issues when running at scale, especially when fetching and operating
+//! on data directly from remote sources.
+//!
+//! Specifically, DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) on the same threadpool
+//! as CPU intensive work, can lead to the issues described in the
+//! [Architecture section] such as throttled network bandwidth (due to congestion
+//! control) and increased latencies or timeouts while processing network
+//! messages.
+//!
+//! It is possible, but more complex, as shows in this example, to separate
+//! the IO and CPU bound work on separate runtimes to avoid these issues.
+use crate::thread_pools_lib::dedicated_executor::DedicatedExecutor;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use object_store::http::HttpBuilder;
+use object_store::ObjectStore;
+use std::sync::Arc;
+use url::Url;
+
+mod thread_pools_lib;
+
+/// Normally, you don't need to worry about the details of the tokio runtime,
+/// but for this example it is important to understand how the [`Runtime`]s work.
+///
+/// There is a "current" runtime that is installed in a thread local variable
+/// that is used by the `tokio::spawn` function.
+///
+/// The `#[tokio::main]` macro creates a [`Runtime`] and installs it as
+/// as the "current" runtime in a thread local variable, on which any `async`
+/// [`Future`], [`Stream]`s and [`Task]`s are run.
+#[tokio::main]
+async fn main() -> Result<()> {
+ // The first two examples read local files, so enable the URL table feature
+ // so we can treat filenames as tables in SQL.
+ let ctx = SessionContext::new().enable_url_table();
+ let sql = format!(
+ "SELECT * FROM '{}/alltypes_plain.parquet'",
+ datafusion::test_util::parquet_test_data()
+ );
+
+ // Run a query on the current runtime. Calling `await` means the future
+ // (in this case the `async` function and all spawned work in DataFusion
+ // plans) on the current runtime.
+ same_runtime(&ctx, &sql).await?;
+
+ // Run the same query but this time on a different runtime. Since we call
+ // `await` here, the `async` function itself runs on the current runtime,
+ // but internally `different_runtime_basic` uses a `DedicatedExecutor` to
+ // run the execute the DataFusion plan on a different Runtime.
+ different_runtime_basic(ctx, sql).await?;
+
+ // Run the same query on a different runtime, including remote IO.
+ //
+ // This is best practice for production systems
+ different_runtime_advanced().await?;
+
+ Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is how most examples in DataFusion are written and works well for
+/// development and local query processing.
+async fn same_runtime(ctx: &SessionContext, sql: &str) -> Result<()> {
+ // Calling .sql is an async function as it may also do network
+ // I/O, for example to contact a remote catalog or do an object store LIST
+ let df = ctx.sql(sql).await?;
+
+ // While many examples call `collect` or `show()`, those methods buffers the
+ // results. Internally DataFusion generates output a RecordBatch at a time
+
+ // Calling `execute_stream` return a `SendableRecordBatchStream`. Depending
+ // on the plan, this may also do network I/O, for example to begin reading a
+ // parquet file from a remote object store as well. It is also possible that
+ // this function call spawns tasks as well.
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` drives the plan, producing new `RecordBatch`es using the
+ // current runtime (and typically also the current thread).
+ //
+ // Perhaps somewhat non obviously, calling the `next()` function can also
+ // result in other tasks being spawned on the current runtime (e.g. for
+ // `RepartitionExec` to read data from each of its input partitions in
+ // parallel).
+ //
+ // Executing the plan like this results in all CPU intensive work
+ // running on same Runtime, in this case whichever one ran the work
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(())
+}
+
+/// Run queries on a **different** runtime than the current one
+///
+/// This is an intermediate example for explanatory purposes. Production systems
+/// should follow the recommendations on [`different_runtime_advanced`] when
+/// running DataFusion queries from a network server or when processing data
+/// from a remote object store.
+async fn different_runtime_basic(ctx: SessionContext, sql: String) -> Result<()> {
+ // Since we are already in the context of runtime (installed by
+ // #[tokio::main]), First, we need a new runtime, which is managed by
+ // a DedicatedExecutor (a library in this example)
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // Now, we run the query on the new runtime
+ dedicated_executor
+ .spawn(async move {
+ // this closure runs on the different thread pool
+ let df = ctx.sql(&sql).await?;
+ let mut stream: SendableRecordBatchStream = df.execute_stream().await?;
+
+ // Calling `next()` to drive the plan in this closure drives the
+ // execution on the other thread pool
+ //
+ // NOTE any IO run by this plan (for example, reading from an
+ // `ObjectStore`) will be done on this new thread pool as well.
+ while let Some(batch) = stream.next().await {
+ println!("{}", pretty_format_batches(&[batch?]).unwrap());
+ }
+ Ok(()) as Result<()>
+ })
+ // even though we are `await`ing here on the "current" pool, internally
+ // the DedicatedExecutor runs the work on the separate thread pool and
+ // the `await` simply notifies it when the work is done
+ .await??;
+
+ // When done with a DedicatedExecutor, it should be shut down cleanly to give
+ // any outstanding tasks a chance to complete.
+ dedicated_executor.join().await;
+
+ Ok(())
+}
+
+/// Demonstrates running queries so that
+/// 1. IO operations happen on the current thread pool
+/// 2. CPU bound tasks happen on a different thread pool
+async fn different_runtime_advanced() -> Result<()> {
+ // In this example, we will query a file via https, reading
+ // the data directly from the plan
+
+ let ctx = SessionContext::new().enable_url_table();
+
+ // setup http object store
+ let base_url = Url::parse("https://github.com").unwrap();
+ let http_store: Arc<dyn ObjectStore> =
+ Arc::new(HttpBuilder::new().with_url(base_url.clone()).build()?);
+
+ let dedicated_executor = DedicatedExecutor::builder().build();
+
+ // By default, the object store will use the runtime that calls `await` for
+ // IO operations. As shown above, using a DedicatedExecutor will run the
+ // plan (and all its IO on the same runtime).
+ //
+ // To avoid this, we can wrap the object store to run on the "IO" runtime
+ //
+ // You can use the DedicatedExecutor::spawn_io method to run other IO
+ // operations.
+ //
+ // Note if we don't do this the example fails with an error like
+ //
+ // ctx.register_object_store(&base_url, http_store);
+ // A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers.
+ let http_store = dedicated_executor.wrap_object_store_for_io(http_store);
+
+ // Tell DataDusion to process `http://` urls with this wrapped object store
+ ctx.register_object_store(&base_url, http_store);
Review Comment:
I guess technically there's a `file://` handler registered by default, so I lied :)
https://github.com/apache/datafusion/blob/82461b74e0be01be8ac17d878dec53140e4209f0/datafusion/execution/src/object_store.rs#L201
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2733529188
Converting to a draft until we hav spawn service
- https://github.com/apache/arrow-rs/pull/7253
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2619967575
> It seems to me like most folks agree research into schedulers for the CPU work that can sit within the single tokio runtime would be _much_ easier to integrate for downstream use cases, probably at the cost of complexity (overheads?) within datafusion itself.
@davidhewitt I agree in principle. I think any time one has to cross from sync code to async code the API is going to get challenging. The upside of the current DataFusion approach with a single tokio runtime is it just "works" (for some defintion of work)
The downside is that squeezing out the maximum performance is tough
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613975012
Is there anyone willing to test out this approach in their application so we can get some confirmation that it avoids some of the timeouts reported while doing heavy processing of data from object store?
Some people I think who have reported seeing this issue:
- @djanderson I think you have seen issues - https://github.com/datafusion-contrib/datafusion-dft/pull/248#issuecomment-2491819394
- @rohitrastogi https://github.com/datafusion-contrib/datafusion-dft/pull/248#issuecomment-2519678264
- @ion-elgreco saw it on https://github.com/apache/arrow-rs/issues/5882
- @JanKaul mentioned he saw something similar at the Amsterdam Meetup
I am willing to help make a draft PR to whatever repo to hook `DedicatedExecutor` up, if someone is willing to then test on a case that was seeing object store timeouts before
Also FYI @ozankabak as we discussed this topic recently.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2613936275
I think this is great Andrew. For what it's worth if this were packaged up in some installable way (even if it had to be from git, etc.) I'm sure we'd be super happy to can our custom stuff and use this :)
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "adriangb (via GitHub)" <gi...@apache.org>.
adriangb commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2614524613
I've successfully made a PR to integrate this internally. It was pretty straightforward. We'll have to scrutinize a bit to see if we can tell if anything is missing (this is very error prone code that can go bad in subtle ways) and if all goes well will report back after this is in production for a couple of days.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2626724671
I think you misunderstand, the parquet reader does not do streaming fetches at all, it uses get_range and get_ranges.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#discussion_r1935223288
##########
datafusion-examples/examples/thread_pools_lib/dedicated_executor.rs:
##########
@@ -0,0 +1,1778 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [DedicatedExecutor] for running CPU-bound tasks on a separate tokio runtime.
+
+use crate::SendableRecordBatchStream;
+use async_trait::async_trait;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion_common::error::GenericError;
+use datafusion_common::DataFusionError;
+use futures::stream::BoxStream;
+use futures::{
+ future::{BoxFuture, Shared},
+ Future, FutureExt, Stream, StreamExt, TryFutureExt,
+};
+use log::{info, warn};
+use object_store::path::Path;
+use object_store::{
+ GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
+ ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, UploadPart,
+};
+use std::cell::RefCell;
+use std::pin::Pin;
+use std::sync::RwLock;
+use std::task::{Context, Poll};
+use std::{fmt::Display, sync::Arc, time::Duration};
+use tokio::runtime::Builder;
+use tokio::task::JoinHandle;
+use tokio::{
+ runtime::Handle,
+ sync::{oneshot::error::RecvError, Notify},
+ task::JoinSet,
+};
+use tokio_stream::wrappers::ReceiverStream;
+
+/// Create a [`DedicatedExecutorBuilder`] from a tokio [`Builder`]
+impl From<Builder> for DedicatedExecutorBuilder {
+ fn from(value: Builder) -> Self {
+ Self::new_from_builder(value)
+ }
+}
+
+/// Manages a separate tokio [`Runtime`] (thread pool) for executing CPU bound
+/// tasks such as DataFusion `ExecutionPlans`.
+///
+/// See [`DedicatedExecutorBuilder`] for creating a new instance.
+///
+/// A `DedicatedExecutor` can helps avoid issues when runnnig IO and CPU bound tasks on the
+/// same thread pool by running futures (and any `tasks` that are
+/// `tokio::task::spawned` by them) on a separate tokio [`Executor`].
+///
+/// `DedicatedExecutor`s can be `clone`ed and all clones share the same thread pool.
+///
+/// Since the primary use for a `DedicatedExecutor` is offloading CPU bound
+/// work, IO work can not be performed on tasks launched in the Executor.
+///
+/// To perform IO, see:
+/// - [`Self::spawn_io`]
+/// - [`Self::wrap_object_store`]
+///
+/// When [`DedicatedExecutorBuilder::build`] is called, a reference to the
+/// "current" tokio runtime will be stored and used, via [`register_io_runtime`] by all
+/// threads spawned by the executor. Any I/O done by threads in this
+/// [`DedicatedExecutor`] should use [`spawn_io`], which will run them on the
+/// I/O runtime.
+///
+/// # TODO examples
+///
+/// # Background
+///
+/// Tokio has the notion of the "current" runtime, which runs the current future
+/// and any tasks spawned by it. Typically, this is the runtime created by
+/// `tokio::main` and is used for the main application logic and I/O handling
+///
+/// For CPU bound work, such as DataFusion plan execution, it is important to
+/// run on a separate thread pool to avoid blocking the I/O handling for extended
+/// periods of time in order to avoid long poll latencies (which decreases the
+/// throughput of small requests under concurrent load).
+///
+/// # IO Scheduling
+///
+/// I/O, such as network calls, should not be performed on the runtime managed
+/// by [`DedicatedExecutor`]. As tokio is a cooperative scheduler, long-running
+/// CPU tasks will not be preempted and can therefore starve servicing of other
+/// tasks. This manifests in long poll-latencies, where a task is ready to run
+/// but isn't being scheduled to run. For CPU-bound work this isn't a problem as
+/// there is no external party waiting on a response, however, for I/O tasks,
+/// long poll latencies can prevent timely servicing of IO, which can have a
+/// significant detrimental effect.
+///
+/// # Details
+///
+/// The worker thread priority is set to low so that such tasks do
+/// not starve other more important tasks (such as answering health checks)
+///
+/// Follows the example from stack overflow and spawns a new
+/// thread to install a Tokio runtime "context"
+/// <https://stackoverflow.com/questions/62536566>
+///
+/// # Trouble Shooting:
+///
+/// ## "No IO runtime registered. Call `register_io_runtime`/`register_current_runtime_for_io` in current thread!
+///
+/// This means that IO was attempted on a tokio runtime that was not registered
+/// for IO. One solution is to run the task using [DedicatedExecutor::spawn].
+///
+/// ## "Cannot drop a runtime in a context where blocking is not allowed"`
+///
+/// If you try to use this structure from an async context you see something like
+/// thread 'test_builder_plan' panicked at 'Cannot
+/// drop a runtime in a context where blocking is not allowed it means This
+/// happens when a runtime is dropped from within an asynchronous
+/// context.', .../tokio-1.4.0/src/runtime/blocking/shutdown.rs:51:21
+///
+/// # Notes
+/// This code is derived from code originally written for [InfluxDB 3.0]
+///
+/// [InfluxDB 3.0]: https://github.com/influxdata/influxdb3_core/tree/6fcbb004232738d55655f32f4ad2385523d10696/executor
+#[derive(Clone, Debug)]
+pub struct DedicatedExecutor {
+ /// State for managing Tokio Runtime Handle for CPU tasks
+ state: Arc<RwLock<State>>,
+}
+
+impl DedicatedExecutor {
+ /// Create a new builder to crate a [`DedicatedExecutor`]
+ pub fn builder() -> DedicatedExecutorBuilder {
+ DedicatedExecutorBuilder::new()
+ }
+
+ /// Runs the specified [`Future`] (and any tasks it spawns) on the thread
+ /// pool managed by this `DedicatedExecutor`.
+ ///
+ /// See The struct documentation for more details
+ ///
+ /// The specified task is added to the tokio executor immediately and
+ /// compete for the thread pool's resources.
+ ///
+ /// # Behavior on `Drop`
+ ///
+ /// UNLIKE [`tokio::task::spawn`], the returned future is **cancelled** when
+ /// it is dropped. Thus, you need ensure the returned future lives until it
+ /// completes (call `await`) or you wish to cancel it.
+ pub fn spawn<T>(&self, task: T) -> impl Future<Output = Result<T::Output, JobError>>
+ where
+ T: Future + Send + 'static,
+ T::Output: Send + 'static,
+ {
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.handle.clone()
+ };
+
+ let Some(handle) = handle else {
+ return futures::future::err(JobError::WorkerGone).boxed();
+ };
+
+ // use JoinSet implement "cancel on drop"
+ let mut join_set = JoinSet::new();
+ join_set.spawn_on(task, &handle);
+ async move {
+ join_set
+ .join_next()
+ .await
+ .expect("just spawned task")
+ .map_err(|e| match e.try_into_panic() {
+ Ok(e) => {
+ let s = if let Some(s) = e.downcast_ref::<String>() {
+ s.clone()
+ } else if let Some(s) = e.downcast_ref::<&str>() {
+ s.to_string()
+ } else {
+ "unknown internal error".to_string()
+ };
+
+ JobError::Panic { msg: s }
+ }
+ Err(_) => JobError::WorkerGone,
+ })
+ }
+ .boxed()
+ }
+
+ /// signals shutdown of this executor and any Clones
+ pub fn shutdown(&self) {
+ // hang up the channel which will cause the dedicated thread
+ // to quit
+ let mut state = self.state.write().expect("lock not poisoned");
+ state.handle = None;
+ state.start_shutdown.notify_one();
+ }
+
+ /// Stops all subsequent task executions, and waits for the worker
+ /// thread to complete. Note this will shutdown all clones of this
+ /// `DedicatedExecutor` as well.
+ ///
+ /// Only the first all to `join` will actually wait for the
+ /// executing thread to complete. All other calls to join will
+ /// complete immediately.
+ ///
+ /// # Panic / Drop
+ /// [`DedicatedExecutor`] implements shutdown on [`Drop`]. You should just use this behavior and NOT call
+ /// [`join`](Self::join) manually during [`Drop`] or panics because this might lead to another panic, see
+ /// <https://github.com/rust-lang/futures-rs/issues/2575>.
+ pub async fn join(&self) {
+ self.shutdown();
+
+ // get handle mutex is held
+ let handle = {
+ let state = self.state.read().expect("lock not poisoned");
+ state.completed_shutdown.clone()
+ };
+
+ // wait for completion while not holding the mutex to avoid
+ // deadlocks
+ handle.await.expect("Thread died?")
+ }
+
+ /// 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_for_io(
+ &self,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Arc<IoObjectStore> {
+ Arc::new(IoObjectStore::new(self.clone(), object_store))
+ }
+
+ /// Runs the [`SendableRecordBatchStream`] on the CPU thread pool
+ ///
+ /// This is a convenience method around [`Self::run_cpu_stream`]
+ pub fn run_cpu_sendable_record_batch_stream(
+ &self,
+ stream: SendableRecordBatchStream,
+ ) -> SendableRecordBatchStream {
+ let schema = stream.schema();
+ let stream = self.run_cpu_stream(stream, |job_error| {
+ let job_error: GenericError = Box::new(job_error);
+ DataFusionError::from(job_error)
+ .context("Running RecordBatchStream on DedicatedExecutor")
+ });
+
+ Box::pin(RecordBatchStreamAdapter::new(schema, stream))
+ }
+
+ /// Runs a stream on the CPU thread pool
+ ///
+ /// # Note
+ /// 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_cpu_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,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ // make a copy to send any job error results back
+ let error_tx = tx.clone();
+
+ // This task will run on the CPU runtime
+ let task = self.spawn(async move {
+ // drive the stream forward on the CPU runtime, sending results
+ // back to the original (presumably IO) runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ });
+
+ // fire up a task on the current runtime which transfers results back
+ // from the CPU runtime to the calling runtime
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ if let Err(e) = task.await {
+ // error running task, try and report it back. An error sending
+ // means the receiver was dropped so there is nowhere to
+ // report errors. Thus ignored via ok()
+ error_tx.send(Err(converter(e))).await.ok();
+ }
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Runs a stream on the IO thread pool
+ ///
+ /// Ths stream must produce Results so that any errors on the dedicated
+ /// executor (like a panic or shutdown) can be communicated back.
+ ///
+ /// Note this has a slightly different API compared to
+ /// [`Self::run_cpu_stream`] because the DedicatedExecutor doesn't monitor
+ /// the CPU thread pool (that came elsewhere)
+ ///
+ /// # Arguments:
+ /// - stream: the stream to run on this dedicated executor
+ pub fn run_io_stream<X, E, S>(
+ &self,
+ stream: S,
+ ) -> impl Stream<Item = Result<X, E>> + Send + 'static
+ where
+ X: Send + 'static,
+ E: Send + 'static,
+ S: Stream<Item = Result<X, E>> + Send + 'static,
+ {
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+
+ let mut set = JoinSet::new();
+ set.spawn(Self::spawn_io_static(async move {
+ // drive the stream forward on the IO runtime, sending results
+ // back to the original runtime
+ let mut stream = Box::pin(stream);
+ while let Some(result) = stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ }));
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ }
+
+ /// Registers `handle` as the IO runtime for this thread
+ ///
+ /// Users should not need to call this function as it is handled by
+ /// [`DedicatedExecutorBuilder`]
+ ///
+ /// # Notes
+ ///
+ /// This sets a thread-local variable
+ ///
+ /// See [`spawn_io`](Self::spawn_io) for more details
+ pub fn register_io_runtime(handle: Option<Handle>) {
+ IO_RUNTIME.set(handle)
+ }
+
+ /// Runs `fut` on IO runtime of this DedicatedExecutor
+ pub async fn spawn_io<Fut>(&self, fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ Self::spawn_io_static(fut).await
+ }
+
+ /// Runs `fut` on the runtime most recently registered by *any* DedicatedExecutor.
+ ///
+ /// When possible, it is preferred to use [`Self::spawn_io`]
+ ///
+ /// This functon is provided, similarly to tokio's [`Handle::current()`] to
+ /// avoid having to thread a `DedicatedExecutor` throughout your program.
+ ///
+ /// # Panic
+ /// Needs a IO runtime [registered](register_io_runtime).
+ pub async fn spawn_io_static<Fut>(fut: Fut) -> Fut::Output
+ where
+ Fut: Future + Send + 'static,
+ Fut::Output: Send,
+ {
+ let h = IO_RUNTIME.with_borrow(|h| h.clone()).expect(
+ "No IO runtime registered. If you hit this panic, it likely \
+ means a DataFusion plan or other CPU bound work is running on the \
+ a tokio threadpool used for IO. Try spawning the work using \
+ `DedicatedExecutor::spawn` or for tests `DedicatedExecutor::register_current_runtime_for_io`",
+ );
+ DropGuard(h.spawn(fut)).await
+ }
+}
+
+/// A wrapper around a receiver stream and task that ensures the inner
+/// task is cancelled on drop
+struct StreamAndTask<T> {
+ inner: ReceiverStream<T>,
+ /// Task which produces no output. On drop the outstanding task is cancelled
+ #[expect(dead_code)]
+ set: JoinSet<()>,
+}
+
+impl<T> Stream for StreamAndTask<T> {
+ type Item = T;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ self.inner.poll_next_unpin(cx)
+ }
+}
+
+thread_local! {
+ /// Tokio runtime `Handle` for doing network (I/O) operations, see [`spawn_io`]
+ pub static IO_RUNTIME: RefCell<Option<Handle>> = const { RefCell::new(None) };
+}
+
+struct DropGuard<T>(JoinHandle<T>);
+impl<T> Drop for DropGuard<T> {
+ fn drop(&mut self) {
+ self.0.abort()
+ }
+}
+
+impl<T> Future for DropGuard<T> {
+ type Output = T;
+
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ Poll::Ready(match std::task::ready!(self.0.poll_unpin(cx)) {
+ Ok(v) => v,
+ Err(e) if e.is_cancelled() => panic!("IO runtime was shut down"),
+ Err(e) => std::panic::resume_unwind(e.into_panic()),
+ })
+ }
+}
+
+/// Runs futures (and any `tasks` that are `tokio::task::spawned` by
+/// them) on a separate tokio Executor.
+///
+/// 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 for CPU tasks
+ ///
+ /// This is `None` when the executor is shutting down.
+ handle: Option<Handle>,
+
+ /// If notified, the executor tokio runtime will begin to shutdown.
+ ///
+ /// We could implement this by checking `handle.is_none()` in regular intervals but requires regular wake-ups and
+ /// locking of the state. Just using a proper async signal is nicer.
+ start_shutdown: Arc<Notify>,
+
+ /// Receiver side indicating that shutdown is complete.
+ completed_shutdown: Shared<BoxFuture<'static, Result<(), Arc<RecvError>>>>,
+
+ /// The inner thread that can be used to join during drop.
+ thread: Option<std::thread::JoinHandle<()>>,
+}
+
+/// IMPORTANT: Implement `Drop` for [`State`], NOT for [`DedicatedExecutor`],
+/// because the executor can be cloned and clones share their inner state.
+impl Drop for State {
+ fn drop(&mut self) {
+ if self.handle.is_some() {
+ warn!("DedicatedExecutor dropped without calling shutdown()");
+ self.handle = None;
+ self.start_shutdown.notify_one();
+ }
+
+ // do NOT poll the shared future if we are panicking due to https://github.com/rust-lang/futures-rs/issues/2575
+ if !std::thread::panicking()
+ && self.completed_shutdown.clone().now_or_never().is_none()
+ {
+ warn!("DedicatedExecutor dropped without waiting for worker termination",);
+ }
+
+ // join thread but don't care about the results
+ self.thread.take().expect("not dropped yet").join().ok();
+ }
+}
+
+const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60 * 5);
+
+/// Potential error returned when polling [`DedicatedExecutor::spawn`].
+#[derive(Debug)]
+pub enum JobError {
+ WorkerGone,
+ Panic { msg: String },
+}
+
+impl Display for JobError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ JobError::WorkerGone => {
+ write!(f, "Worker thread gone, executor was likely shut down")
+ }
+ JobError::Panic { msg } => write!(f, "Panic: {}", msg),
+ }
+ }
+}
+
+impl std::error::Error for JobError {}
+
+/// Builder for [`DedicatedExecutor`]
+pub struct DedicatedExecutorBuilder {
+ /// Name given to all execution threads. Defaults to "DedicatedExecutor"
+ name: String,
+ /// Builder for tokio runtime. Defaults to multi-threaded builder
+ runtime_builder: Builder,
+}
+
+impl From<JobError> for DataFusionError {
+ fn from(value: JobError) -> Self {
+ DataFusionError::External(Box::new(value))
+ .context("JobError from DedicatedExecutor")
+ }
+}
+
+impl DedicatedExecutorBuilder {
+ /// Create a new `DedicatedExecutorBuilder` with default values
+ ///
+ /// Note that by default this `DedicatedExecutor` will not be able to
+ /// perform network I/O.
+ pub fn new() -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder: Builder::new_multi_thread(),
+ }
+ }
+
+ /// Create a new `DedicatedExecutorBuilder` from a pre-existing tokio
+ /// runtime [`Builder`].
+ ///
+ /// This method permits customizing the tokio [`Executor`] used for the
+ /// [`DedicatedExecutor`]
+ pub fn new_from_builder(runtime_builder: Builder) -> Self {
+ Self {
+ name: String::from("DedicatedExecutor"),
+ runtime_builder,
+ }
+ }
+
+ /// Set the name of the dedicated executor (appear in the names of each thread).
+ ///
+ /// Defaults to "DedicatedExecutor"
+ #[cfg(test)]
+ pub fn with_name(mut self, name: impl Into<String>) -> Self {
+ self.name = name.into();
+ self
+ }
+
+ /// Set the number of worker threads. Defaults to the tokio default (the
+ /// number of virtual CPUs)
+ #[allow(dead_code)]
+ pub fn with_worker_threads(mut self, num_threads: usize) -> Self {
+ self.runtime_builder.worker_threads(num_threads);
+ self
+ }
+
+ /// Creates a new `DedicatedExecutor` with a dedicated tokio
+ /// executor that is separate from the thread pool created via
+ /// `[tokio::main]` or similar.
+ ///
+ /// Note: If [`DedicatedExecutorBuilder::build`] is called from an existing
+ /// tokio runtime, it will assume that the existing runtime should be used
+ /// for I/O.
+ ///
+ /// See the documentation on [`DedicatedExecutor`] for more details.
+ pub fn build(self) -> DedicatedExecutor {
+ let Self {
+ name,
+ runtime_builder,
+ } = self;
+
+ let notify_shutdown = Arc::new(Notify::new());
+ let notify_shutdown_captured = Arc::clone(¬ify_shutdown);
+
+ let (tx_shutdown, rx_shutdown) = tokio::sync::oneshot::channel();
+ let (tx_handle, rx_handle) = std::sync::mpsc::channel();
+
+ let io_handle = Handle::try_current().ok();
+ let thread = std::thread::Builder::new()
+ .name(format!("{name} driver"))
+ .spawn(move || {
+ // also register the IO runtime for the current thread, since it might be used as well (esp. for the
+ // current thread RT)
+ DedicatedExecutor::register_io_runtime(io_handle.clone());
+
+ info!("Creating DedicatedExecutor",);
+
+ let mut runtime_builder = runtime_builder;
+ let runtime = runtime_builder
+ .on_thread_start(move || {
+ DedicatedExecutor::register_io_runtime(io_handle.clone())
+ })
+ .build()
+ .expect("Creating tokio runtime");
+
+ runtime.block_on(async move {
+ // Enable the "notified" receiver BEFORE sending the runtime handle back to the constructor thread
+ // (i.e .the one that runs `new`) to avoid the potential (but unlikely) race that the shutdown is
+ // started right after the constructor finishes and the new runtime calls
+ // `notify_shutdown_captured.notified().await`.
+ //
+ // Tokio provides an API for that by calling `enable` on the `notified` future (this requires
+ // pinning though).
+ let shutdown = notify_shutdown_captured.notified();
+ let mut shutdown = std::pin::pin!(shutdown);
+ shutdown.as_mut().enable();
+
+ if tx_handle.send(Handle::current()).is_err() {
+ return;
+ }
+ shutdown.await;
+ });
+
+ runtime.shutdown_timeout(SHUTDOWN_TIMEOUT);
+
+ // send shutdown "done" signal
+ tx_shutdown.send(()).ok();
+ })
+ .expect("executor setup");
+
+ let handle = rx_handle.recv().expect("driver started");
+
+ let state = State {
+ handle: Some(handle),
+ start_shutdown: notify_shutdown,
+ completed_shutdown: rx_shutdown.map_err(Arc::new).boxed().shared(),
+ thread: Some(thread),
+ };
+
+ DedicatedExecutor {
+ state: Arc::new(RwLock::new(state)),
+ }
+ }
+}
+
+/// Wraps an inner [`ObjectStore`] so that all underlying methods are run on
+/// the Tokio Runtime dedicated to doing IO
+///
+/// # See Also
+///
+/// [`DedicatedExecutor::spawn_io`] for more details
+#[derive(Debug)]
+pub struct IoObjectStore {
+ dedicated_executor: DedicatedExecutor,
+ inner: Arc<dyn ObjectStore>,
+}
+
+impl IoObjectStore {
+ pub fn new(executor: DedicatedExecutor, object_store: Arc<dyn ObjectStore>) -> Self {
+ Self {
+ dedicated_executor: executor,
+ inner: object_store,
+ }
+ }
+
+ /// Wrap the result stream if necessary
+ fn wrap_if_necessary(&self, payload: GetResultPayload) -> GetResultPayload {
+ match payload {
+ GetResultPayload::File(_, _) => payload,
+ GetResultPayload::Stream(stream) => {
+ let new_stream = self.dedicated_executor.run_io_stream(stream).boxed();
+ GetResultPayload::Stream(new_stream)
+ }
+ }
+ }
+}
+
+impl Display for IoObjectStore {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ writeln!(f, "IoObjectStore")
+ }
+}
+
+#[async_trait]
+impl ObjectStore for IoObjectStore {
+ /// Return GetResult for the location and options.
+ async fn get_opts(
+ &self,
+ location: &Path,
+ options: GetOptions,
+ ) -> object_store::Result<GetResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let mut results = self
+ .dedicated_executor
+ .spawn_io(async move { store.get_opts(&location, options).await })
+ .await?;
+
+ // As the results can be a stream too, we must wrap that too
+ results.payload = self.wrap_if_necessary(results.payload);
+ Ok(results)
+ }
+
+ async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn copy_if_not_exists(
+ &self,
+ from: &Path,
+ to: &Path,
+ ) -> object_store::Result<()> {
+ let from = from.clone();
+ let to = to.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.copy(&from, &to).await })
+ .await
+ }
+
+ async fn delete(&self, location: &Path) -> object_store::Result<()> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.delete(&location).await })
+ .await
+ }
+
+ fn list(
+ &self,
+ prefix: Option<&Path>,
+ ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
+ // run the inner list on the dedicated executor
+ //
+ // This requires some fiddling as we can't pass the result of list
+ // across the runtime, so start another task that drives list on the IO
+ // thread and sends results back to the executor thread.
+ let (tx, rx) = tokio::sync::mpsc::channel(1);
+ let prefix_captured = prefix.cloned();
+ let inner_captured = Arc::clone(&self.inner);
+ let executor_captured = self.dedicated_executor.clone();
+ let mut set = JoinSet::new();
+ set.spawn(async move {
+ executor_captured
+ .spawn_io(async move {
+ // run and buffer on the IO stream
+ let mut list_stream = inner_captured.list(prefix_captured.as_ref());
+ while let Some(result) = list_stream.next().await {
+ // try to send to the sender, if error means the
+ // receiver has been closed and we terminate early
+ if tx.send(result).await.is_err() {
+ return;
+ }
+ }
+ })
+ .await;
+ });
+
+ StreamAndTask {
+ inner: ReceiverStream::new(rx),
+ set,
+ }
+ .boxed()
+ }
+
+ async fn list_with_delimiter(
+ &self,
+ prefix: Option<&Path>,
+ ) -> object_store::Result<ListResult> {
+ let prefix = prefix.cloned();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.list_with_delimiter(prefix.as_ref()).await })
+ .await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOpts,
+ ) -> object_store::Result<Box<dyn MultipartUpload>> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ let result = self
+ .dedicated_executor
+ .spawn_io(async move { store.put_multipart_opts(&location, opts).await })
+ .await?;
+
+ // the resulting object has async functions too which we must wrap
+ Ok(Box::new(IoMultipartUpload::new(
+ self.dedicated_executor.clone(),
+ result,
+ )))
+ }
+
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> object_store::Result<PutResult> {
+ let location = location.clone();
+ let store = Arc::clone(&self.inner);
+ self.dedicated_executor
+ .spawn_io(async move { store.put_opts(&location, payload, opts).await })
+ .await
+ }
+}
+
+/// A wrapper around a `MultipartUpload` that runs the async functions on the
+/// IO thread of the DedicatedExecutor
+#[derive(Debug)]
+struct IoMultipartUpload {
+ dedicated_executor: DedicatedExecutor,
+ /// Inner upload (needs to be an option so we can send in closure)
+ inner: Option<Box<dyn MultipartUpload>>,
+}
+impl IoMultipartUpload {
+ fn new(
+ dedicated_executor: DedicatedExecutor,
+ inner: Box<dyn MultipartUpload>,
+ ) -> Self {
+ Self {
+ dedicated_executor,
+ inner: Some(inner),
+ }
+ }
+}
+
+#[async_trait]
+impl MultipartUpload for IoMultipartUpload {
+ fn put_part(&mut self, data: PutPayload) -> UploadPart {
Review Comment:
It shouldn't matter that they're written out of order as they're assembled based on the index captured when the future is created. Could you perhaps share your code?
Edit: I created an explicit test in object_store to show this - https://github.com/apache/arrow-rs/pull/7047
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Example for using a separate threadpool for CPU bound work (try 2) [datafusion]
Posted by "JanKaul (via GitHub)" <gi...@apache.org>.
JanKaul commented on PR #14286:
URL: https://github.com/apache/datafusion/pull/14286#issuecomment-2626849059
Thanks for having a look. I have to dig a bit deeper to really understand where the issue is.
I think this PR is a very pragmatic approach that solves the most important use case. As this problem seems to only occur when executing the physical plan.
--
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@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org