Add example for using a separate threadpool for CPU bound work

Author: alamb

datafusion-examples/examples/thread_pools.rs

@@ -0,0 +1,167 @@
+// 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.
+
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use futures::stream::StreamExt;
+use datafusion::execution::SendableRecordBatchStream;
+//! This example shows how to use a separate thread pool (tokio [`Runtime`])) to
+//! run the CPU intensive parts of DataFusion plans.
+//!
+//! Running DataFusion plans that perform I/O, such as reading parquet files
+//! directly from remote object storage (e.g. AWS S3) without care will result
+//! in running CPU intensive jobs on the same thread pool, which can lead to the
+//! issues described in  the [Architecture section] such as throttled bandwidth
+//! due to congestion control and increased latencies for processing network
+//! messages.
+
+
+/// 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 actually creates a [`Runtime`] and installs it as
+/// as the "current" runtime (on which any `async` futures, streams and tasks
+/// are run).
+#[tokio::main]
+async fn main() -> Result<()> {
+    let ctx = SessionContext::new()
+        // enabling URL table means we can select directly from
+        // paths in SQL queries.
+        .enable_url_table();
+    let sql = format!(
+        "SELECT * FROM '{}/alltypes_plain.parquet'",
+        datafusion::test_util::parquet_test_data()
+    );
+
+    // Run the same query on the same runtime. Note that calling `await` here
+    // will effectively run the future (in this case the `async` function) on
+    // the current runtime.
+    same_runtime(&ctx, &sql).await?;
+
+    // Run the same query on a different runtime. Note that we are still calling
+    // `await` here, so the the `async` function still runs on the current runtime.
+    // We use the `DedicatedExecutor` to run the query on a different runtime.
+    different_runtime(ctx, sql).await?;
+    Ok(())
+}
+
+/// Run queries directly on the current tokio `Runtime`
+///
+/// This is now 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` on a DataFrame returns 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 that begin doing CPU intensive work 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 obvious, calling the `next()` function often will
+    // 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 (default) Runtime.
+    while let Some(batch) = stream.next().await {
+        println!("{}", pretty_format_batches(&[batch?]).unwrap());
+    }
+    Ok(())
+}
+
+/// Demonstrates how to run queries on a **different** runtime than the current one
+///
+/// This is typically how you should run DataFusion queries from a network
+/// server or when processing data from a remote object store.
+async fn different_runtime(ctx: SessionContext, sql: String) -> Result<()> {
+    // First, we need a new runtime, which we can create with the tokio builder
+    // however, since we are already in the context of another runtime
+    // (installed by #[tokio::main]) we create a new thread for the runtime
+    tokio::task::spawn_blocking(move || {
+        std::thread::spawn(move || thread_entry(ctx, sql.to_string()))
+            .join()
+            .expect("thread did not panic")
+    })
+    .await
+    .expect("task did not panic")
+}
+
+/// This is the entry point of thread that we started our second runtime on
+fn thread_entry(ctx: SessionContext, sql: String) -> Result<()> {
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        // only enable the time driver (not the I/O driver), meaning this
+        // runtime will not be able to perform network I/O
+        .enable_time()
+        .build()?;
+
+    // Now we can run the actual code we want on a different runtime
+    runtime.block_on(async move { different_runtime_inner(ctx, sql).await })
+}
+
+/// this is run on a new, separate runtime
+async fn different_runtime_inner(ctx: SessionContext, sql: String) -> Result<()> {
+    // Setup execution as before
+    let df = ctx.sql(&sql).await?;
+
+    let mut stream = df.execute_stream().await?;
+
+    //XXX  Note at this point, calling next() will be run on our new threadpool. However, this will also spawn the catalog and object store requests on the same threadpool as well!
+
+    // While this will mean we don't interfere with handling of other network requests, it will mean tht the network requests that happen as part of query processing will still be running on the same threadpool
+
+    // TODO show this working
+
+    // To avoid this, all IO access, both catalog and data (e.g. object_store) must be spawned on to their own runtime, like this:
+    // TODO....
+
+    //
+    // care is required to avoid calling `next()` (aka polling) from the default IO thread (even if planning / execution is run on that other thread)
+    // Best practice is to do all of DataFusion planning / execution on a separate pool. Note that some care is required for remote catalogs such as iceberg that
+    // themselves do network IO
+    // TODO figure out how to cause an erorr due to io thread
+
+    while let Some(batch) = stream.next().await {
+        println!("{}", pretty_format_batches(&[batch?]).unwrap());
+    }
+
+    //
+    // You can use the DataFusion "DedicatedExecutor" in this case to handle most of this automatically for you.
+
+    // Note that special care must be taken when running Datafusion plans that do async io to transfer the work to their own thread pools.
+
+    // Using separate Runtime will avoid other requests being messed up but it won't really help requests made from DataSources such as
+    // reading parquet files from object_store.
+    //
+    // Thus this runtime also disables IO so that we are sure there is no IO work being done on it.
+
+    Ok(())
+}

datafusion/common-runtime/Cargo.toml

@@ -38,6 +38,8 @@ path = "src/lib.rs"
 [dependencies]
 log = { workspace = true }
 tokio = { workspace = true }
+parking_lot = { workspace = true }
+futures = { workspace = true }
 
 [dev-dependencies]
-tokio = { version = "1.36", features = ["rt", "rt-multi-thread", "time"] }
+tokio = { version = "1.36", features = ["rt", "rt-multi-thread", "time", "net"] }