schedule work

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-cxx/src/read.rs

@@ -4,18 +4,11 @@
 use std::sync::Arc;
 
 use anyhow::Result;
-use arrow_array::RecordBatch;
 use arrow_array::RecordBatchReader;
-use arrow_array::cast::AsArray;
 use arrow_array::ffi::FFI_ArrowSchema;
 use arrow_array::ffi_stream::FFI_ArrowArrayStream;
-use arrow_schema::ArrowError;
-use arrow_schema::DataType;
 use arrow_schema::Schema;
 use arrow_schema::SchemaRef;
-use futures::StreamExt;
-use futures::stream::TryStreamExt;
-use vortex::array::arrow::IntoArrowArray;
 use vortex::buffer::Buffer;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::io::runtime::BlockingRuntime;
@@ -158,19 +151,9 @@ pub(crate) fn scan_builder_into_threadsafe_cloneable_reader(
             Arc::new(arrow_schema)
         }
     };
-    let data_type = DataType::Struct(schema.fields().clone());
-
     let stream = builder
         .inner
-        .into_stream()?
-        .map(move |result| {
-            result.and_then(|chunk| {
-                chunk
-                    .into_arrow(&data_type)
-                    .map(|struct_array| RecordBatch::from(struct_array.as_struct()))
-            })
-        })
-        .map_err(|e| ArrowError::ExternalError(Box::new(e)));
+        .into_record_batch_stream(Arc::clone(&schema))?;
 
     let iter = RUNTIME.block_on_stream_thread_safe(|_h| stream);
     let rbr = RecordBatchIteratorAdapter::new(iter, schema);

vortex-datafusion/src/persistent/opener.rs

@@ -37,6 +37,7 @@ use vortex::array::arrow::ArrowArrayExecutor;
 use vortex::error::VortexError;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::io::InstrumentedReadAt;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::LayoutReader;
 use vortex::layout::scan::scan_builder::ScanBuilder;
 use vortex::metrics::Label;
@@ -362,6 +363,7 @@ impl FileOpener for VortexOpener {
             }
 
             let stream_schema = Arc::new(stream_schema);
+            let handle = session.handle();
 
             let stream = scan_builder
                 .with_metrics_registry(metrics_registry)
@@ -370,10 +372,15 @@ impl FileOpener for VortexOpener {
                 .with_ordered(has_output_ordering)
                 .into_stream()
                 .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))?
-                .map(move |chunk| {
-                    let mut ctx = session.create_execution_ctx();
-                    chunk.and_then(|chunk| {
-                        chunk.execute_record_batch(stream_schema.as_ref(), &mut ctx)
+                .then(move |chunk| {
+                    let session = session.clone();
+                    let stream_schema = Arc::clone(&stream_schema);
+                    let handle = handle.clone();
+                    handle.spawn_blocking(move || {
+                        let mut ctx = session.create_execution_ctx();
+                        chunk.and_then(|chunk| {
+                            chunk.execute_record_batch(stream_schema.as_ref(), &mut ctx)
+                        })
                     })
                 })
                 .map_ok(move |rb| {

vortex-datafusion/src/v2/source.rs

@@ -429,15 +429,15 @@ impl DataSource for VortexDataSource {
             let handle = session.handle();
             let stream = scan_streams
                 .try_flatten_unordered(Some(num_partitions.get() * 2))
-                .map(move |result| {
+                .then(move |result| {
                     let session = session.clone();
                     let schema = Arc::clone(&projected_schema);
-                    handle.spawn_cpu(move || {
+                    let handle = handle.clone();
+                    handle.spawn_blocking(move || {
                         let mut ctx = session.create_execution_ctx();
                         result.and_then(|chunk| chunk.execute_record_batch(&schema, &mut ctx))
                     })
                 })
-                .buffered(num_partitions.get())
                 .map(|result| result.map_err(|e| DataFusionError::External(Box::new(e))));
 
             // Apply leftover projection (expressions that couldn't be pushed into Vortex).

vortex-layout/src/scan/arrow.rs

@@ -18,6 +18,7 @@ use vortex_array::VortexSessionExecute;
 use vortex_array::arrow::ArrowArrayExecutor;
 use vortex_error::VortexResult;
 use vortex_io::runtime::BlockingRuntime;
+use vortex_io::session::RuntimeSessionExt;
 
 use crate::scan::scan_builder::ScanBuilder;
 
@@ -52,13 +53,22 @@ impl ScanBuilder {
     ) -> VortexResult<impl Stream<Item = Result<RecordBatch, ArrowError>> + Send + 'static> {
         let data_type = DataType::Struct(schema.fields().clone());
         let session = self.session().clone();
+        let handle = session.handle();
+        let concurrency = std::thread::available_parallelism()
+            .map(|n| n.get())
+            .unwrap_or(1);
 
         let stream = self
             .into_stream()?
             .map(move |chunk| {
-                let mut ctx = session.create_execution_ctx();
-                chunk.and_then(|chunk| to_record_batch(chunk, &data_type, &mut ctx))
+                let session = session.clone();
+                let data_type = data_type.clone();
+                handle.spawn_blocking(move || {
+                    let mut ctx = session.create_execution_ctx();
+                    chunk.and_then(|chunk| to_record_batch(chunk, &data_type, &mut ctx))
+                })
             })
+            .buffered(concurrency)
             .map_err(|e| ArrowError::ExternalError(Box::new(e)));
 
         Ok(stream)