diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 0f5ac43d54a..10f8855f2f3 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -18,6 +18,7 @@ use futures::Stream; use futures::StreamExt; use futures::future::BoxFuture; use itertools::Itertools; +use kanal::AsyncSender; use num_traits::AsPrimitive; use parking_lot::Mutex; use static_assertions::assert_impl_all; @@ -52,6 +53,7 @@ use vortex::scalar_fn::fns::binary::Binary; use vortex::scalar_fn::fns::operators::Operator; use vortex::scalar_fn::fns::pack::Pack; use vortex::scan::DataSource; +use vortex::scan::Partition; use vortex::scan::ScanRequest; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::parallelism::get_available_parallelism; @@ -211,6 +213,59 @@ pub fn bind(input: &BindInputRef, result: &mut BindResultRef) -> VortexResult>; +type PartitionSender = AsyncSender; + +async fn on_partition( + sender: PartitionSender, + partition: PartitionResult, + pending: Arc, +) { + let partition = match partition { + Ok(partition) => partition, + Err(e) => { + let _ = sender.send(Err(e)).await; + return; + } + }; + + let cache = Arc::new(ConversionCache { + file_index: partition.index(), + ..Default::default() + }); + + let splits = match partition.execute_splits() { + Ok(splits) => splits, + Err(e) => { + let _ = sender.send(Err(e)).await; + return; + } + }; + let handles = splits + .into_iter() + .map(|task| RUNTIME.handle().spawn(task)) + .collect::>(); + for handle in handles { + match handle.await { + Ok(Some(array)) => { + pending.fetch_add(1, Ordering::Relaxed); + if sender.send(Ok((array, Arc::clone(&cache)))).await.is_err() { + // Exit early if the receiver has been dropped, which happens + // when the scan is complete or if an error has occurred in + // another partition. + return; + } + } + // split is filtered + Ok(None) => {} + Err(e) => { + let _ = sender.send(Err(e)).await; + return; + } + } + } +} + pub fn init_global(init_input: &TableInitInput) -> VortexResult { debug!(input=?init_input, "table function global input"); @@ -274,63 +329,28 @@ pub fn init_global(init_input: &TableInitInput) -> VortexResult partition, - Err(e) => { - let _ = tx.send(Err(e)).await; - return; - } - }; - - let cache = Arc::new(ConversionCache { - file_index: partition.index(), - ..Default::default() - }); - - let mut stream = match partition.execute() { - Ok(s) => s, - Err(e) => { - let _ = tx.send(Err(e)).await; - return; - } - }; - while let Some(item) = stream.next().await { - pending.fetch_add(1, Ordering::Relaxed); - if tx - .send(item.map(|a| (a, Arc::clone(&cache)))) - .await - .is_err() - { - // Exit early if the receiver has been dropped, which happens when the - // scan is complete or if an error has occurred in another partition. - return; - } - } - }) + RUNTIME + .handle() + .spawn(async move { on_partition(sender, partition, pending).await }) }) .buffer_unordered(num_workers); - let iterator = RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, rx)); + let iterator = + RUNTIME.block_on_stream_thread_safe(|_handle| scan_driver_stream(stream, receiver)); let aggregates = bind_data.aggregates.clone(); let partials = build_partials( diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 94ce7acf277..19c4cfec911 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -36,6 +36,7 @@ use vortex_scan::Partition; use vortex_scan::PartitionRef; use vortex_scan::PartitionStream; use vortex_scan::ScanRequest; +use vortex_scan::SplitTask; use vortex_scan::selection::Selection; use vortex_session::VortexSession; @@ -341,6 +342,18 @@ impl Partition for LayoutReaderSplit { dtype, stream, ))) } + + fn execute_splits(self: Box) -> VortexResult> { + ScanBuilder::new(self.session, self.reader) + .with_row_range(self.row_range) + .with_selection(self.selection) + .with_projection(self.projection) + .with_some_filter(self.filter) + .with_some_limit(self.limit) + .with_some_metrics_registry(self.metrics_registry) + .with_ordered(self.ordered) + .build() + } } /// A scan that produces no data, only empty arrays with the correct row count. @@ -401,4 +414,11 @@ impl Partition for Empty { stream::iter(iter), ))) } + + fn execute_splits(self: Box) -> VortexResult> { + let scalar = Scalar::default_value(&self.dtype); + let row_count = usize::try_from(self.row_count)?; + let array = ConstantArray::new(scalar.clone(), row_count).into_array(); + Ok(vec![futures::future::ready(Ok(Some(array))).boxed()]) + } } diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index d9de768f271..eb509e40c34 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -56,6 +56,7 @@ use vortex_scan::Partition; use vortex_scan::PartitionRef; use vortex_scan::PartitionStream; use vortex_scan::ScanRequest; +use vortex_scan::SplitTask; use vortex_scan::selection::Selection; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; @@ -596,6 +597,23 @@ impl Partition for MultiLayoutPartition { dtype, stream, ))) } + + fn execute_splits(self: Box) -> VortexResult> { + let request = self.request; + let filter = request.filter?; + let mut builder = ScanBuilder::new(self.session, self.reader) + .with_selection(request.selection) + .with_projection(request.projection) + .with_some_filter(filter) + .with_some_limit(request.limit) + .with_ordered(request.ordered); + + if let Some(row_range) = request.row_range { + builder = builder.with_row_range(row_range); + } + + builder.build() + } } #[cfg(test)] diff --git a/vortex-scan/src/lib.rs b/vortex-scan/src/lib.rs index 3aba97cda8a..499fca63c81 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -31,8 +31,10 @@ use std::ops::Range; use std::sync::Arc; use async_trait::async_trait; +use futures::future::BoxFuture; use futures::stream::BoxStream; use selection::Selection; +use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::expr::Expression; @@ -208,4 +210,10 @@ pub trait Partition: 'static + Send { /// operations should be spawned onto the runtime to enable parallel execution across /// threads. fn execute(self: Box) -> VortexResult; + + /// Prepare independent split tasks for a partition and execute it. + fn execute_splits(self: Box) -> VortexResult>; } + +/// Split task of a partition. Future yields array of a split. +pub type SplitTask = BoxFuture<'static, VortexResult>>;