diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index eea2ce3dca0..7d1ff2e1c06 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -442,13 +442,12 @@ impl Stream for LazyScanStream { match &mut self.state { LazyScanState::Builder(builder) => { let builder = builder.take().vortex_expect("polled after completion"); - // Only `prepare` runs on the blocking pool. Split tasks are built lazily by - // the stream below, so that the configured concurrency — not the split count - // — bounds how many splits have registered reads. + // Preparation is CPU work. Split tasks are built lazily by the stream + // below, so configured concurrency bounds registered split reads. let task = builder .session .handle() - .spawn_blocking(move || builder.prepare()); + .spawn_cpu(move || builder.prepare()); self.state = LazyScanState::Preparing(task); } LazyScanState::Preparing(task) => match ready!(Pin::new(task).poll(cx)) { @@ -497,6 +496,7 @@ mod test { use std::time::Duration; use futures::Stream; + use futures::future::BoxFuture; use futures::task::noop_waker_ref; use parking_lot::Mutex; use rstest::rstest; @@ -519,7 +519,10 @@ mod test { use vortex_array::expr::root; use vortex_error::VortexResult; use vortex_error::vortex_err; + use vortex_io::runtime::AbortHandleRef; use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::Executor; + use vortex_io::runtime::Handle; use vortex_io::runtime::single::SingleThreadRuntime; use vortex_mask::Mask; @@ -678,6 +681,56 @@ mod test { assert_eq!(calls.load(Ordering::Relaxed), 0); } + struct CountingExecutor { + runtime: tokio::runtime::Handle, + cpu: AtomicUsize, + blocking: AtomicUsize, + } + + impl Executor for CountingExecutor { + fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { + Executor::spawn(&self.runtime, fut) + } + + fn spawn_io(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { + Executor::spawn_io(&self.runtime, fut) + } + + fn spawn_cpu(&self, task: Box) -> AbortHandleRef { + self.cpu.fetch_add(1, Ordering::Relaxed); + Executor::spawn_cpu(&self.runtime, task) + } + + fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { + self.blocking.fetch_add(1, Ordering::Relaxed); + Executor::spawn_blocking_io(&self.runtime, task) + } + } + + #[tokio::test] + async fn into_stream_schedules_prepare_on_cpu_executor() -> VortexResult<()> { + let executor = Arc::new(CountingExecutor { + runtime: tokio::runtime::Handle::current(), + cpu: AtomicUsize::new(0), + blocking: AtomicUsize::new(0), + }); + let runtime: Arc = Arc::::clone(&executor); + let session = session_with_handle(Handle::new(Arc::downgrade(&runtime))); + let calls = Arc::new(AtomicUsize::new(0)); + let reader = Arc::new(CountingLayoutReader::new(Arc::clone(&calls))); + let mut stream = ScanBuilder::new(session, reader).into_stream()?; + + assert_eq!(executor.cpu.load(Ordering::Relaxed), 0); + assert_eq!(executor.blocking.load(Ordering::Relaxed), 0); + let mut cx = Context::from_waker(noop_waker_ref()); + assert!(Pin::new(&mut stream).poll_next(&mut cx).is_pending()); + assert_eq!(executor.cpu.load(Ordering::Relaxed), 1); + assert_eq!(executor.blocking.load(Ordering::Relaxed), 0); + // The current-thread runtime has not polled the scheduled CPU task yet. + assert_eq!(calls.load(Ordering::Relaxed), 0); + Ok(()) + } + #[derive(Debug)] struct SplittingLayoutReader { name: Arc,