Reuse the IoDispatcher across DataFusion instances (#1299)

Author: gatesn

Cargo.lock

@@ -4,7 +4,11 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use tokio::runtime::Builder;
 
 fn benchmark(c: &mut Criterion) {
-    let runtime = Builder::new_current_thread().enable_all().build().unwrap();
+    let runtime = Builder::new_current_thread()
+        .thread_name("benchmark-tpch")
+        .enable_all()
+        .build()
+        .unwrap();
 
     // Run TPC-H data gen.
     let data_dir = DBGen::new(DBGenOptions::default()).generate().unwrap();

bench-vortex/benches/tpch.rs

@@ -2,7 +2,7 @@ use std::fs::File;
 use std::ops::Range;
 use std::path::{Path, PathBuf};
 use std::process::Command;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use arrow_array::types::Int64Type;
 use arrow_array::{
@@ -27,11 +27,16 @@ use vortex::arrow::FromArrowType;
 use vortex::compress::CompressionStrategy;
 use vortex::dtype::DType;
 use vortex::error::VortexResult;
-use vortex::file::{LayoutContext, LayoutDeserializer, VortexFileWriter, VortexReadBuilder};
+use vortex::file::{
+    IoDispatcher, LayoutContext, LayoutDeserializer, VortexFileWriter, VortexReadBuilder,
+};
 use vortex::io::{ObjectStoreReadAt, TokioFile, VortexReadAt, VortexWrite};
 use vortex::sampling_compressor::{SamplingCompressor, ALL_ENCODINGS_CONTEXT};
 use vortex::{Array, IntoArray, IntoCanonical};
 
+static DISPATCHER: LazyLock<Arc<IoDispatcher>> =
+    LazyLock::new(|| Arc::new(IoDispatcher::new_tokio(1)));
+
 pub const BATCH_SIZE: usize = 65_536;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -51,6 +56,7 @@ pub async fn open_vortex(path: &Path) -> VortexResult<Array> {
             LayoutContext::default().into(),
         ),
     )
+    .with_io_dispatcher(DISPATCHER.clone())
     .build()
     .await?
     .read_all()

bench-vortex/src/reader.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use arrow_array::RecordBatch;
 use arrow_schema::SchemaRef;
@@ -9,9 +9,15 @@ use futures::{FutureExt as _, StreamExt, TryStreamExt};
 use object_store::ObjectStore;
 use vortex_array::Context;
 use vortex_expr::datafusion::convert_expr_to_vortex;
-use vortex_file::{LayoutContext, LayoutDeserializer, Projection, RowFilter, VortexReadBuilder};
+use vortex_file::{
+    IoDispatcher, LayoutContext, LayoutDeserializer, Projection, RowFilter, VortexReadBuilder,
+};
 use vortex_io::ObjectStoreReadAt;
 
+/// Share an IO dispatcher across all DataFusion instances.
+static IO_DISPATCHER: LazyLock<Arc<IoDispatcher>> =
+    LazyLock::new(|| Arc::new(IoDispatcher::new_tokio(1)));
+
 pub struct VortexFileOpener {
     pub ctx: Arc<Context>,
     pub object_store: Arc<dyn ObjectStore>,
@@ -28,7 +34,8 @@ impl FileOpener for VortexFileOpener {
         let mut builder = VortexReadBuilder::new(
             read_at,
             LayoutDeserializer::new(self.ctx.clone(), Arc::new(LayoutContext::default())),
-        );
+        )
+        .with_io_dispatcher(IO_DISPATCHER.clone());
 
         let row_filter = self
             .predicate

vortex-datafusion/src/persistent/opener.rs

@@ -26,13 +26,13 @@ impl TokioDispatcher {
         let threads: Vec<_> = (0..num_threads)
             .map(|tid| {
                 let worker_thread = std::thread::Builder::new();
-                let worker_thread = worker_thread.name(format!("tokio-dispatch-{tid}"));
                 let rx: flume::Receiver<Box<dyn TokioSpawn + Send>> = rx.clone();
 
                 worker_thread
                     .spawn(move || {
                         // Create a runtime-per-thread
                         let rt = tokio::runtime::Builder::new_current_thread()
+                            .thread_name(format!("tokio-dispatch-{tid}"))
                             .enable_all()
                             .build()
                             .unwrap_or_else(|e| {

vortex-file/src/dispatcher/tokio.rs

@@ -71,7 +71,7 @@ pub struct VortexReadBuilder<R> {
     size: Option<u64>,
     row_mask: Option<Array>,
     row_filter: Option<RowFilter>,
-    io_dispatcher: Option<IoDispatcher>,
+    io_dispatcher: Option<Arc<IoDispatcher>>,
 }
 
 impl<R: VortexReadAt> VortexReadBuilder<R> {
@@ -112,7 +112,7 @@ impl<R: VortexReadAt> VortexReadBuilder<R> {
         self
     }
 
-    pub fn with_io_dispatcher(mut self, dispatcher: IoDispatcher) -> Self {
+    pub fn with_io_dispatcher(mut self, dispatcher: Arc<IoDispatcher>) -> Self {
         self.io_dispatcher = Some(dispatcher);
         self
     }
@@ -176,7 +176,7 @@ impl<R: VortexReadAt> VortexReadBuilder<R> {
         // Default: fallback to single-threaded tokio dispatcher.
         let io_dispatcher = self
             .io_dispatcher
-            .unwrap_or_else(|| IoDispatcher::new_tokio(1));
+            .unwrap_or_else(|| Arc::new(IoDispatcher::new_tokio(1)));
 
         Ok(VortexFileArrayStream::new(
             self.read_at,

vortex-file/src/read/builder/mod.rs

@@ -35,7 +35,7 @@ pub struct VortexFileArrayStream<R> {
     messages_cache: Arc<RwLock<LayoutMessageCache>>,
     state: Option<StreamingState>,
     input: R,
-    dispatcher: IoDispatcher,
+    dispatcher: Arc<IoDispatcher>,
 }
 
 impl<R: VortexReadAt> VortexFileArrayStream<R> {
@@ -48,7 +48,7 @@ impl<R: VortexReadAt> VortexFileArrayStream<R> {
         dtype: DType,
         row_count: u64,
         row_mask: Option<RowMask>,
-        dispatcher: IoDispatcher,
+        dispatcher: Arc<IoDispatcher>,
     ) -> Self {
         let mask_iterator = if let Some(fr) = filter_reader {
             Box::new(FilteringRowSplitIterator::new(fr, row_count, row_mask)) as MaskIteratorRef