From 94bf00a63cb1aae088b325e307f2756d957fffdf Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 20 Apr 2026 12:30:27 +0100 Subject: [PATCH 01/17] Support limit and filter for ScanBuilder and RepeatedScan Signed-off-by: Adam Gutglick --- vortex-bench/src/datasets/tpch_l_comment.rs | 2 +- .../src/persistent/access_plan.rs | 5 +- vortex-datafusion/src/persistent/opener.rs | 20 +- vortex-file/src/file.rs | 3 +- vortex-file/src/tests.rs | 2 +- vortex-layout/src/scan/arrow.rs | 11 +- vortex-layout/src/scan/limit.rs | 74 +++++ vortex-layout/src/scan/mod.rs | 1 + vortex-layout/src/scan/repeated_scan.rs | 120 +++++--- vortex-layout/src/scan/scan_builder.rs | 288 +++++++++++------- vortex-layout/src/scan/tasks.rs | 19 +- 11 files changed, 370 insertions(+), 175 deletions(-) create mode 100644 vortex-layout/src/scan/limit.rs diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index ebe6ff2e96a..619cff09b28 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use anyhow::Result; use async_trait::async_trait; +use futures::StreamExt; use futures::TryStreamExt; use glob::glob; use vortex::array::ArrayRef; @@ -86,7 +87,6 @@ impl Dataset for TPCHLCommentChunked { Ok(canonical.into_array()) } }) - .into_array_stream()? .try_collect() .await?; chunks.extend(file_chunks); diff --git a/vortex-datafusion/src/persistent/access_plan.rs b/vortex-datafusion/src/persistent/access_plan.rs index ad7bf941a2c..0e14b6730b0 100644 --- a/vortex-datafusion/src/persistent/access_plan.rs +++ b/vortex-datafusion/src/persistent/access_plan.rs @@ -58,10 +58,7 @@ impl VortexAccessPlan { /// /// This is used internally by the file opener after it has translated a /// `PartitionedFile` into a Vortex scan. - pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder - where - A: 'static + Send, - { + pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder { let Self { selection } = self; if let Some(selection) = selection { diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..ef03be17242 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -436,18 +436,20 @@ impl FileOpener for VortexOpener { .with_projection(scan_projection) .with_some_filter(filter) .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(); - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) + chunk.and_then(|chunk| { + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) }) - .into_stream() - .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index c9a71f85c55..94a3b70494a 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -11,7 +11,6 @@ use std::sync::Arc; use std::sync::OnceLock; use itertools::Itertools; -use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::Expression; @@ -175,7 +174,7 @@ impl VortexFile { /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and /// execution options. - pub fn scan(&self) -> VortexResult> { + pub fn scan(&self) -> VortexResult { Ok(ScanBuilder::new( self.session.clone(), self.layout_reader()?, diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index e665ab18d50..b7550949c8c 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1278,7 +1278,7 @@ async fn write_nullable_top_level_struct() { async fn round_trip( array: &ArrayRef, - f: impl Fn(ScanBuilder) -> VortexResult>, + f: impl Fn(ScanBuilder) -> VortexResult, ) -> VortexResult { let mut writer = vec![]; SESSION diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 663b29d9501..85b4653e40f 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -10,6 +10,7 @@ use arrow_schema::ArrowError; use arrow_schema::Field; use arrow_schema::SchemaRef; use futures::Stream; +use futures::StreamExt; use futures::TryStreamExt; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -20,7 +21,7 @@ use vortex_io::runtime::BlockingRuntime; use crate::scan::scan_builder::ScanBuilder; -impl ScanBuilder { +impl ScanBuilder { /// Creates a new `RecordBatchReader` from the scan builder. /// /// The `schema` parameter is used to define the schema of the resulting record batches. In @@ -35,11 +36,11 @@ impl ScanBuilder { let session = self.session().clone(); let iter = self + .into_iter(runtime)? .map(move |chunk| { let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) + chunk.and_then(|chunk| to_record_batch(chunk, &struct_field, &mut ctx)) }) - .into_iter(runtime)? .map(|result| result.map_err(|e| ArrowError::ExternalError(Box::new(e)))); Ok(RecordBatchIteratorAdapter { iter, schema }) @@ -53,11 +54,11 @@ impl ScanBuilder { let session = self.session().clone(); let stream = self + .into_stream()? .map(move |chunk| { let mut ctx = session.create_execution_ctx(); - to_record_batch(chunk, &struct_field, &mut ctx) + chunk.and_then(|chunk| to_record_batch(chunk, &struct_field, &mut ctx)) }) - .into_stream()? .map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(stream) diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs new file mode 100644 index 00000000000..cc99d0f6c6b --- /dev/null +++ b/vortex-layout/src/scan/limit.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +use futures::Stream; +use futures::StreamExt; +use futures::stream; +use futures::stream::BoxStream; +use vortex_array::ArrayRef; +use vortex_error::VortexResult; + +pub(crate) fn limit_array_stream( + stream: S, + limit: Option, +) -> BoxStream<'static, VortexResult> +where + S: Stream> + Send + 'static, +{ + match limit { + Some(limit) => RowLimitedStream::new(stream.boxed(), limit).boxed(), + None => stream.boxed(), + } +} + +struct RowLimitedStream { + inner: BoxStream<'static, VortexResult>, + remaining: u64, +} + +impl RowLimitedStream { + fn new(inner: BoxStream<'static, VortexResult>, remaining: u64) -> Self { + Self { inner, remaining } + } + + fn abort_pending(&mut self) { + let inner = std::mem::replace(&mut self.inner, stream::empty().boxed()); + drop(inner); + } +} + +impl Stream for RowLimitedStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.remaining == 0 { + return Poll::Ready(None); + } + + match self.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => { + let chunk_len = chunk.len() as u64; + if chunk_len <= self.remaining { + self.remaining -= chunk_len; + if self.remaining == 0 { + self.abort_pending(); + } + Poll::Ready(Some(Ok(chunk))) + } else { + let limit = match usize::try_from(self.remaining) { + Ok(limit) => limit, + Err(_) => unreachable!("remaining rows cannot exceed the current chunk"), + }; + self.remaining = 0; + self.abort_pending(); + Poll::Ready(Some(chunk.slice(0..limit))) + } + } + other => other, + } + } +} diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index 98fd1918a42..e08e4cb3f31 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -4,6 +4,7 @@ pub mod arrow; mod filter; pub mod layout; +mod limit; pub mod multi; pub mod repeated_scan; pub mod scan_builder; diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 681f33639bc..302fa5c3052 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -6,8 +6,9 @@ use std::iter; use std::ops::Range; use std::sync::Arc; +use futures::FutureExt; use futures::Stream; -use futures::future::BoxFuture; +use futures::StreamExt; use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; @@ -27,15 +28,15 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; +use crate::scan::limit::limit_array_stream; use crate::scan::splits::Splits; -use crate::scan::tasks::TaskContext; -use crate::scan::tasks::split_exec; +use crate::scan::tasks::{split_exec, TaskContext, TaskFuture}; /// A projected subset (by indices, range, and filter) of rows from a Vortex data source. /// /// The method of this struct enable, possibly concurrent, scanning of multiple row ranges of this /// data source. -pub struct RepeatedScan { +pub struct RepeatedScan { session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -49,15 +50,13 @@ pub struct RepeatedScan { splits: Splits, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, - /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - map_fn: Arc VortexResult + Send + Sync>, /// Maximal number of rows to read (after filtering) limit: Option, /// The dtype of the projected arrays. dtype: DType, } -impl RepeatedScan { +impl RepeatedScan { pub fn dtype(&self) -> &DType { &self.dtype } @@ -81,9 +80,6 @@ impl RepeatedScan { let stream = self.execute_stream(row_range)?; Ok(ArrayStreamAdapter::new(dtype, stream)) } -} - -impl RepeatedScan { /// Constructor just to allow `scan_builder` to create a `RepeatedScan`. #[expect( clippy::too_many_arguments, @@ -99,7 +95,6 @@ impl RepeatedScan { selection: Selection, splits: Splits, concurrency: usize, - map_fn: Arc VortexResult + Send + Sync>, limit: Option, dtype: DType, ) -> Self { @@ -113,16 +108,12 @@ impl RepeatedScan { selection, splits, concurrency, - map_fn, limit, dtype, } } - pub fn execute( - &self, - row_range: Option>, - ) -> VortexResult>>>> { + fn split_ranges(&self, row_range: Option>) -> Vec> { let selection_range: Option> = match &self.selection { Selection::IncludeByIndex(buf) if !buf.is_empty() => { Some(buf[0]..buf[buf.len() - 1] + 1) @@ -135,14 +126,14 @@ impl RepeatedScan { let row_range = intersect_ranges(self.row_range.as_ref(), row_range); let row_range = intersect_ranges(row_range.as_ref(), selection_range); - let ranges = match &self.splits { + match &self.splits { Splits::Natural(vec) => { debug_assert!(vec.is_sorted()); let splits_iter = match row_range { None => Either::Left(vec.iter().copied()), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } let lo = vec.partition_point(|&x| x < range.start); let hi = vec.partition_point(|&x| x < range.end); @@ -154,33 +145,43 @@ impl RepeatedScan { } }; - Either::Left(splits_iter.tuple_windows().map(|(start, end)| start..end)) + splits_iter + .tuple_windows() + .map(|(start, end)| start..end) + .collect() } - Splits::Ranges(ranges) => Either::Right(match row_range { - None => Either::Left(ranges.iter().cloned()), + Splits::Ranges(ranges) => match row_range { + None => ranges.to_vec(), Some(range) => { if range.is_empty() { - return Ok(Vec::new()); + return Vec::new(); } - Either::Right(ranges.iter().filter_map(move |r| { - let start = cmp::max(r.start, range.start); - let end = cmp::min(r.end, range.end); - (start < end).then_some(start..end) - })) + ranges + .iter() + .filter_map(move |r| { + let start = cmp::max(r.start, range.start); + let end = cmp::min(r.end, range.end); + (start < end).then_some(start..end) + }) + .collect() } - }), - }; + }, + } + } - let mut limit = self.limit; + pub(crate) fn execute( + &self, + row_range: Option>, + ) -> VortexResult>>> { + let mut limit = self.limit.filter(|_| self.filter.is_none()); let mut tasks = Vec::new(); let ctx = Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), reader: Arc::clone(&self.layout_reader), projection: self.projection.clone(), - mapper: Arc::clone(&self.map_fn), }); - for range in ranges { + for range in self.split_ranges(row_range) { let row_mask = self.selection.row_mask(&range); if row_mask.mask().all_false() { continue; @@ -195,25 +196,66 @@ impl RepeatedScan { Ok(tasks) } - pub fn execute_stream( + pub(crate) fn execute_stream( &self, row_range: Option>, - ) -> VortexResult> + Send + 'static + use> { - use futures::StreamExt; - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = self.concurrency * num_workers; + ) -> VortexResult> + Send + 'static> { let handle = self.session.handle(); + if self.filter.is_some() && self.limit.is_some() { + let ctx = Arc::new(TaskContext { + filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), + reader: Arc::clone(&self.layout_reader), + projection: self.projection.clone(), + }); + let selection = self.selection.clone(); + let ordered = self.ordered; + let limit = self.limit; + let stream = futures::stream::iter(self.split_ranges(row_range)).filter_map(move |range| { + let row_mask = selection.row_mask(&range); + if row_mask.mask().all_false() { + return async { None }.boxed(); + } + + let ctx = Arc::clone(&ctx); + let handle = handle.clone(); + async move { + Some( + async move { + let task = split_exec(ctx, row_mask, None)?; + handle.spawn(task).await + } + .boxed(), + ) + } + .boxed() + }); + let stream = if ordered { + stream.buffered(1).boxed() + } else { + stream.buffer_unordered(1).boxed() + }; + + return Ok(limit_array_stream( + stream.filter_map(|chunk| async move { chunk.transpose() }), + limit, + )); + } + + let num_workers = get_available_parallelism().unwrap_or(1); + let concurrency = self.concurrency * num_workers; let stream = futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); - let stream = if self.ordered { stream.buffered(concurrency).boxed() } else { stream.buffer_unordered(concurrency).boxed() }; - Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) + Ok(limit_array_stream( + stream.filter_map(|chunk| async move { chunk.transpose() }), + self.limit, + )) } } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..092b7cbdadd 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -10,7 +10,6 @@ use std::task::ready; use futures::Stream; use futures::StreamExt; -use futures::future::BoxFuture; use futures::stream::BoxStream; use itertools::Itertools; use vortex_array::ArrayRef; @@ -21,21 +20,17 @@ use vortex_array::expr::analysis::referenced_field_paths; use vortex_array::expr::root; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; -use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_io::runtime::BlockingRuntime; -use vortex_io::runtime::Handle; use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_metrics::MetricsRegistry; use vortex_scan::selection::Selection; use vortex_session::VortexSession; -use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReader; use crate::LayoutReaderRef; @@ -56,7 +51,7 @@ use crate::scan::splits::attempt_split_ranges; /// Projection and filter expressions are optimized against the reader dtype during /// [`prepare`](Self::prepare). Work is divided by the configured [`SplitBy`] strategy or by /// explicit selection ranges. -pub struct ScanBuilder { +pub struct ScanBuilder { session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -72,11 +67,7 @@ pub struct ScanBuilder { split_by: SplitBy, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, - /// Function to apply to each [`ArrayRef`] within the spawned split tasks. - map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, - /// Should we try to prune the file (using stats) on open. - file_stats: Option>, /// Maximal number of rows to read (after filtering) limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, @@ -84,7 +75,7 @@ pub struct ScanBuilder { row_offset: u64, } -impl ScanBuilder { +impl ScanBuilder { /// Create a scan builder over `layout_reader` using `session` for runtime and execution state. pub fn new(session: VortexSession, layout_reader: Arc) -> Self { Self { @@ -99,9 +90,7 @@ impl ScanBuilder { // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. concurrency: 4, - map_fn: Arc::new(Ok), metrics_registry: None, - file_stats: None, limit: None, row_offset: 0, } @@ -130,7 +119,7 @@ impl ScanBuilder { } } -impl ScanBuilder { +impl ScanBuilder { /// Add a filter expression evaluated against the projected row ranges. pub fn with_filter(mut self, filter: Expression) -> Self { self.filter = Some(filter); @@ -237,38 +226,10 @@ impl ScanBuilder { &self.session } - /// Map each split of the scan. The function will be run on the spawned task. - pub fn map( - self, - map_fn: impl Fn(A) -> VortexResult + 'static + Send + Sync, - ) -> ScanBuilder { - let old_map_fn = self.map_fn; - ScanBuilder { - session: self.session, - layout_reader: self.layout_reader, - projection: self.projection, - filter: self.filter, - ordered: self.ordered, - row_range: self.row_range, - selection: self.selection, - split_by: self.split_by, - concurrency: self.concurrency, - metrics_registry: self.metrics_registry, - file_stats: self.file_stats, - limit: self.limit, - row_offset: self.row_offset, - map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)), - } - } - /// Optimize expressions, compute split ranges, and return an executable repeated scan. - pub fn prepare(self) -> VortexResult> { + pub fn prepare(self) -> VortexResult { let dtype = self.dtype()?; - if self.filter.is_some() && self.limit.is_some() { - vortex_bail!("Vortex doesn't support scans with both a filter and a limit") - } - // Spin up the root layout reader, and wrap it in a FilterLayoutReader to perform // conjunction splitting if a filter is provided. let mut layout_reader = self.layout_reader; @@ -319,26 +280,15 @@ impl ScanBuilder { self.selection, splits, self.concurrency, - self.map_fn, self.limit, dtype, )) } - /// Constructs a task per row split of the scan, returned as a vector of futures. - pub fn build(self) -> VortexResult>>>> { - // The ultimate short circuit - if self.limit.is_some_and(|l| l == 0) { - return Ok(vec![]); - } - - self.prepare()?.execute(None) - } - /// Returns a [`Stream`] with tasks spawned onto the session's runtime handle. pub fn into_stream( self, - ) -> VortexResult> + Send + 'static + use> { + ) -> VortexResult> + Send + 'static> { Ok(LazyScanStream::new(self)) } @@ -346,82 +296,55 @@ impl ScanBuilder { pub fn into_iter( self, runtime: &B, - ) -> VortexResult> + 'static> { + ) -> VortexResult> + 'static> { let stream = self.into_stream()?; Ok(runtime.block_on_stream(stream)) } } -enum LazyScanState { - Builder(Option>>), - Preparing(PreparingScan), - Stream(BoxStream<'static, VortexResult>), +enum LazyScanState { + Builder(Option>), + Preparing(PreparingScan), + Stream(BoxStream<'static, VortexResult>), Error(Option), } -type PreparedScanTasks = Vec>>>; - -struct PreparingScan { - ordered: bool, - concurrency: usize, - handle: Handle, - task: Task>>, +struct PreparingScan { + task: Task>, } -struct LazyScanStream { - state: LazyScanState, +struct LazyScanStream { + state: LazyScanState, } -impl LazyScanStream { - fn new(builder: ScanBuilder) -> Self { +impl LazyScanStream { + fn new(builder: ScanBuilder) -> Self { Self { state: LazyScanState::Builder(Some(Box::new(builder))), } } } -impl Unpin for LazyScanStream {} +impl Unpin for LazyScanStream {} -impl Stream for LazyScanStream { - type Item = VortexResult; +impl Stream for LazyScanStream { + type Item = VortexResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { match &mut self.state { LazyScanState::Builder(builder) => { let builder = builder.take().vortex_expect("polled after completion"); - let ordered = builder.ordered; - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = builder.concurrency * num_workers; let handle = builder.session.handle(); - let task = handle.spawn_blocking(move || { - builder.prepare().and_then(|scan| scan.execute(None)) - }); - self.state = LazyScanState::Preparing(PreparingScan { - ordered, - concurrency, - handle, - task, - }); + let task = handle.spawn_blocking(move || builder.prepare()); + self.state = LazyScanState::Preparing(PreparingScan { task }); } LazyScanState::Preparing(preparing) => { match ready!(Pin::new(&mut preparing.task).poll(cx)) { - Ok(tasks) => { - let ordered = preparing.ordered; - let concurrency = preparing.concurrency; - let handle = preparing.handle.clone(); - let stream = - futures::stream::iter(tasks).map(move |task| handle.spawn(task)); - let stream = if ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; - let stream = stream - .filter_map(|chunk| async move { chunk.transpose() }) - .boxed(); - self.state = LazyScanState::Stream(stream); - } + Ok(scan) => match scan.execute_stream(None) { + Ok(stream) => self.state = LazyScanState::Stream(stream.boxed()), + Err(err) => self.state = LazyScanState::Error(Some(err)), + }, Err(err) => self.state = LazyScanState::Error(Some(err)), } } @@ -466,6 +389,7 @@ mod test { use futures::Stream; use futures::task::noop_waker_ref; use parking_lot::Mutex; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; @@ -752,6 +676,166 @@ mod test { Ok(()) } + #[derive(Debug)] + struct FilteringLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + keep_row: fn(u64) -> bool, + } + + impl FilteringLayoutReader { + fn new(row_count: u64, keep_row: fn(u64) -> bool) -> Self { + Self { + name: Arc::from("filtering"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + keep_row, + } + } + } + + impl LayoutReader for FilteringLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + row_range: &Range, + splits: &mut BTreeSet, + ) -> VortexResult<()> { + for split in ((row_range.start + 2)..row_range.end).step_by(2) { + splits.insert(split); + } + splits.insert(row_range.end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let keep_row = self.keep_row; + let row_count = usize::try_from(row_range.end - row_range.start) + .map_err(|_| vortex_err!("row range must fit in usize"))?; + + Ok(MaskFuture::new(row_count, async move { + let input_mask = mask.await?; + let filtered = (row_range.start..row_range.end) + .enumerate() + .map(|(idx, row)| input_mask.value(idx) && keep_row(row)); + Ok(Mask::from_iter(filtered)) + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + + let array = PrimitiveArray::from_iter(start..end).into_array(); + array.filter(mask.await?) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + fn collect_scan_values(iter: I) -> VortexResult> + where + I: IntoIterator>, + { + let mut values = Vec::new(); + for chunk in iter { + #[expect(deprecated)] + let primitive = chunk?.to_primitive(); + values.extend(primitive.into_buffer::()); + } + Ok(values) + } + + fn drain_runtime(runtime: &SingleThreadRuntime) { + for _ in 0..4 { + let mut yielded = false; + runtime.block_on(futures::future::poll_fn(move |cx| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + })); + } + } + + #[test] + fn into_stream_limits_filtered_results() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = crate::scan::test::session_with_handle(runtime.handle()); + let reader = Arc::new(FilteringLayoutReader::new(8, |_| true)); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(3) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + assert_eq!(values, [0, 1, 2]); + Ok(()) + } + + #[test] + fn prepared_scan_limits_filtered_results() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = crate::scan::test::session_with_handle(runtime.handle()); + let reader = Arc::new(FilteringLayoutReader::new(8, |row| row % 2 == 1)); + + let scan = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(3) + .prepare()?; + let values = collect_scan_values(scan.execute_array_iter(None, &runtime)?)?; + drain_runtime(&runtime); + + assert_eq!(values, [1, 3, 5]); + Ok(()) + } + #[derive(Debug)] struct BlockingSplitsLayoutReader { name: Arc, diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index a86546e15ef..40d770a9903 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -19,7 +19,7 @@ use vortex_scan::row_mask::RowMask; use crate::LayoutReader; use crate::scan::filter::FilterExpr; -pub type TaskFuture = BoxFuture<'static, VortexResult>; +pub type TaskFuture = BoxFuture<'static, VortexResult>; /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this @@ -33,13 +33,12 @@ pub type TaskFuture = BoxFuture<'static, VortexResult>; /// The intersected row range is then further reduced via expression-based pruning. After pruning /// has eliminated more blocks, the full filter is executed over the remainder of the split. /// -/// This mask is then provided to the reader to perform a filtered projection over the split data, -/// finally mapping the Vortex columnar record batches into some result type `A`. -pub fn split_exec( - ctx: Arc>, +/// This mask is then provided to the reader to perform a filtered projection over the split data. +pub fn split_exec( + ctx: Arc, read_mask: RowMask, limit: Option<&mut u64>, -) -> VortexResult>> { +) -> VortexResult>> { let row_range = read_mask.row_range(); let row_mask = read_mask.mask().clone(); @@ -136,15 +135,13 @@ pub fn split_exec( ctx.reader .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; - let mapper = Arc::clone(&ctx.mapper); let array_fut = async move { let mask = filter_mask.await?; if mask.all_false() { return Ok(None); } - let array = projection_future.await?; - mapper(array).map(Some) + projection_future.await.map(Some) }; Ok(array_fut.boxed()) @@ -153,13 +150,11 @@ pub fn split_exec( /// Information needed to execute a single split task. /// /// Row selection is evaluated before creating a split task so it's not included -pub struct TaskContext { +pub struct TaskContext { /// The shared filter expression. pub filter: Option>, /// The layout reader. pub reader: Arc, /// The projection expression to apply to gather the scanned rows. pub projection: Expression, - /// Function that maps into an A. - pub mapper: Arc VortexResult + Send + Sync>, } From 08ae9cb5de0c9f4a6821479b9db92bf3dd3e4bcf Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 20 Apr 2026 12:55:39 +0100 Subject: [PATCH 02/17] schedule work Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/opener.rs | 30 ++++++++++++++-------- vortex-layout/src/scan/arrow.rs | 23 ++++++++++++++--- 2 files changed, 39 insertions(+), 14 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index ef03be17242..89205c7f9be 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -46,6 +46,7 @@ use vortex::error::VortexError; use vortex::error::VortexExpect; 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::layout::scan::split_by::SplitBy; @@ -430,7 +431,9 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + let handle = session.handle(); + let stream_target_field = + Arc::new(Field::new_struct("", stream_schema.fields().clone(), false)); let stream = scan_builder .with_metrics_registry(metrics_registry) .with_projection(scan_projection) @@ -438,16 +441,21 @@ 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| { - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(&stream_target_field), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) + .then(move |chunk| { + let session = session.clone(); + let stream_target_field = Arc::clone(&stream_target_field); + let handle = handle.clone(); + handle.spawn_blocking(move || { + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(stream_target_field.as_ref()), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) + }) }) }) .map_err(move |e: VortexError| { diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 85b4653e40f..f62bf4e920f 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -18,6 +18,7 @@ use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::session::RuntimeSessionExt; use crate::scan::scan_builder::ScanBuilder; @@ -50,15 +51,31 @@ impl ScanBuilder { self, schema: SchemaRef, ) -> VortexResult> + Send + 'static> { - let struct_field = Field::new_struct("", schema.fields().clone(), false); + let struct_field = Arc::new(Field::new_struct("", schema.fields().clone(), false)); 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, &struct_field, &mut ctx)) + let session = session.clone(); + let handle = handle.clone(); + let struct_field = Arc::clone(&struct_field); + async move { + handle + .spawn_blocking(move || { + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + to_record_batch(chunk, struct_field.as_ref(), &mut ctx) + }) + }) + .await + } }) + .buffered(concurrency) .map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(stream) From 19a9e7c540134a3ec13897e9860974757036ff2b Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 20 Apr 2026 12:57:53 +0100 Subject: [PATCH 03/17] layout Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/arrow.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index f62bf4e920f..9772580d14a 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -18,7 +18,6 @@ use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; -use vortex_io::session::RuntimeSessionExt; use crate::scan::scan_builder::ScanBuilder; @@ -53,10 +52,6 @@ impl ScanBuilder { ) -> VortexResult> + Send + 'static> { let struct_field = Arc::new(Field::new_struct("", schema.fields().clone(), false)); 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()? @@ -75,7 +70,6 @@ impl ScanBuilder { .await } }) - .buffered(concurrency) .map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(stream) From 0a1ae31954b301b046ddb197d179f82c9d66e17e Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 20 Apr 2026 16:32:55 +0100 Subject: [PATCH 04/17] bufferd batch Signed-off-by: Adam Gutglick --- vortex-bench/src/datasets/tpch_l_comment.rs | 9 ++++++--- vortex-datafusion/src/persistent/opener.rs | 20 ++++++++++++-------- vortex-datafusion/src/persistent/stream.rs | 10 +++++++--- vortex-layout/src/scan/arrow.rs | 6 ++++++ vortex-layout/src/scan/repeated_scan.rs | 18 ++++++++++++------ vortex-layout/src/scan/scan_builder.rs | 21 +++++++++++---------- vortex-python/src/file.rs | 2 +- vortex-python/src/scan.rs | 3 +-- 8 files changed, 56 insertions(+), 33 deletions(-) diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index 619cff09b28..049837b56df 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -79,12 +79,15 @@ impl Dataset for TPCHLCommentChunked { let file_chunks: Vec<_> = file .scan()? .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable)) + .into_stream()? .map({ let ctx = ctx.clone(); - move |a| { + move |result| { let mut ctx = ctx.clone(); - let canonical = a.execute::(&mut ctx)?; - Ok(canonical.into_array()) + result.and_then(|a| { + let canonical = a.execute::(&mut ctx)?; + Ok(canonical.into_array()) + }) } }) .try_collect() diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 89205c7f9be..e0b890c3329 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -431,9 +431,10 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } - let handle = session.handle(); let stream_target_field = Arc::new(Field::new_struct("", stream_schema.fields().clone(), false)); + let handle = session.handle(); + let file_location = file.object_meta.location.clone(); let stream = scan_builder .with_metrics_registry(metrics_registry) .with_projection(scan_projection) @@ -441,7 +442,7 @@ impl FileOpener for VortexOpener { .with_ordered(has_output_ordering) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .then(move |chunk| { + .map(move |chunk| { let session = session.clone(); let stream_target_field = Arc::clone(&stream_target_field); let handle = handle.clone(); @@ -458,12 +459,8 @@ impl FileOpener for VortexOpener { }) }) }) - .map_err(move |e: VortexError| { - DataFusionError::External(Box::new(e.with_context(format!( - "Failed to read Vortex file: {}", - file.object_meta.location - )))) - }) + .buffered(2) + .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) .map(move |batch| { let batch = if projector.projection().as_ref().is_empty() { batch @@ -582,6 +579,12 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: u64::try_from(midpoint_byte).vortex_expect("midpoint byte projection should fit into u64") } +fn vortex_file_read_error(path: &Path, error: VortexError) -> DataFusionError { + DataFusionError::External(Box::new( + error.with_context(format!("Failed to read Vortex file: {path}")), + )) +} + #[cfg(test)] mod tests { use std::fmt; @@ -614,6 +617,7 @@ mod tests { use datafusion_physical_expr::expressions as df_expr; use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::projection::ProjectionExpr; + use futures::TryStreamExt; use insta::assert_snapshot; use itertools::Itertools; use object_store::ObjectStore; diff --git a/vortex-datafusion/src/persistent/stream.rs b/vortex-datafusion/src/persistent/stream.rs index af2fbc8693e..9038ba81db8 100644 --- a/vortex-datafusion/src/persistent/stream.rs +++ b/vortex-datafusion/src/persistent/stream.rs @@ -17,14 +17,14 @@ use futures::stream::BoxStream; /// [`PartitionedFile`]: datafusion_datasource::PartitionedFile pub(crate) struct PrunableStream { file_pruner: FilePruner, - stream: BoxStream<'static, DFResult>, + stream: Option>>, } impl PrunableStream { pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult>) -> Self { Self { file_pruner, - stream, + stream: Some(stream), } } } @@ -34,9 +34,13 @@ impl Stream for PrunableStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.as_mut().file_pruner.should_prune()? { + self.stream.take(); Poll::Ready(None) } else { - self.stream.poll_next_unpin(cx) + match self.stream.as_mut() { + Some(stream) => stream.poll_next_unpin(cx), + None => Poll::Ready(None), + } } } } diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index 9772580d14a..f62bf4e920f 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -18,6 +18,7 @@ use vortex_array::VortexSessionExecute; use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::session::RuntimeSessionExt; use crate::scan::scan_builder::ScanBuilder; @@ -52,6 +53,10 @@ impl ScanBuilder { ) -> VortexResult> + Send + 'static> { let struct_field = Arc::new(Field::new_struct("", schema.fields().clone(), false)); 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()? @@ -70,6 +75,7 @@ impl ScanBuilder { .await } }) + .buffered(concurrency) .map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(stream) diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 302fa5c3052..eff69ffbaa9 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -50,7 +50,7 @@ pub struct RepeatedScan { splits: Splits, /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, - /// Maximal number of rows to read (after filtering) + /// Maximal number of rows to read (after filtering). limit: Option, /// The dtype of the projected arrays. dtype: DType, @@ -80,6 +80,7 @@ impl RepeatedScan { let stream = self.execute_stream(row_range)?; Ok(ArrayStreamAdapter::new(dtype, stream)) } + /// Constructor just to allow `scan_builder` to create a `RepeatedScan`. #[expect( clippy::too_many_arguments, @@ -173,24 +174,29 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult>>> { - let mut limit = self.limit.filter(|_| self.filter.is_none()); - let mut tasks = Vec::new(); let ctx = Arc::new(TaskContext { filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), reader: Arc::clone(&self.layout_reader), projection: self.projection.clone(), }); + let mut limit = self.limit.filter(|_| self.filter.is_none()); + let mut tasks = Vec::new(); + for range in self.split_ranges(row_range) { + if range.start >= range.end { + continue; + } + if limit.is_some_and(|l| l == 0) { + break; + } + let row_mask = self.selection.row_mask(&range); if row_mask.mask().all_false() { continue; } tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); - if limit.is_some_and(|l| l == 0) { - break; - } } Ok(tasks) diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 092b7cbdadd..cc2604e9a88 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -40,7 +40,7 @@ use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; -/// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. +/// Builder for scanning a [`LayoutReader`] into arrays, streams, or iterators. /// /// A scan has three independent row restriction mechanisms: /// @@ -68,7 +68,7 @@ pub struct ScanBuilder { /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, metrics_registry: Option>, - /// Maximal number of rows to read (after filtering) + /// Maximal number of rows to read after filtering. limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, /// but not by the scan [`Selection`] which remains relative. @@ -711,13 +711,14 @@ mod test { fn register_splits( &self, _field_mask: &[FieldMask], - row_range: &Range, - splits: &mut BTreeSet, + split_range: &SplitRange, + splits: &mut RowSplits, ) -> VortexResult<()> { + let row_range = split_range.row_range(); for split in ((row_range.start + 2)..row_range.end).step_by(2) { - splits.insert(split); + splits.push(split_range.row_offset() + split); } - splits.insert(row_range.end); + splits.push(split_range.root_row_range().end); Ok(()) } @@ -778,10 +779,10 @@ mod test { where I: IntoIterator>, { + let mut ctx = array_session().create_execution_ctx(); let mut values = Vec::new(); for chunk in iter { - #[expect(deprecated)] - let primitive = chunk?.to_primitive(); + let primitive = chunk?.execute::(&mut ctx)?; values.extend(primitive.into_buffer::()); } Ok(values) @@ -805,7 +806,7 @@ mod test { #[test] fn into_stream_limits_filtered_results() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); - let session = crate::scan::test::session_with_handle(runtime.handle()); + let session = session_with_handle(runtime.handle()); let reader = Arc::new(FilteringLayoutReader::new(8, |_| true)); let stream = ScanBuilder::new(session, reader) @@ -822,7 +823,7 @@ mod test { #[test] fn prepared_scan_limits_filtered_results() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); - let session = crate::scan::test::session_with_handle(runtime.handle()); + let session = session_with_handle(runtime.handle()); let reader = Arc::new(FilteringLayoutReader::new(8, |row| row % 2 == 1)); let scan = ScanBuilder::new(session, reader) diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 9550067b956..214b1a6244b 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -208,7 +208,7 @@ fn scan_builder( indices: Option, batch_size: Option, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult { let mut builder = vxf .scan()? .with_some_filter(expr) diff --git a/vortex-python/src/scan.rs b/vortex-python/src/scan.rs index cfdf77a6b8b..8ee95f7f8d5 100644 --- a/vortex-python/src/scan.rs +++ b/vortex-python/src/scan.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; -use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; use vortex::error::VortexResult; use vortex::layout::scan::repeated_scan::RepeatedScan; @@ -30,7 +29,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { #[pyclass(name = "RepeatedScan", module = "vortex", frozen)] pub struct PyRepeatedScan { - pub scan: Arc>, + pub scan: Arc, pub row_count: u64, } From d3a39deabb653ea5c97b407ce340ee64b692a427 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 13:16:04 +0100 Subject: [PATCH 05/17] Shared limit pushdown Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/opener.rs | 29 +++- vortex-layout/src/scan/layout.rs | 151 ++++++++++++++++++++- vortex-layout/src/scan/limit.rs | 93 +++++++++++++ vortex-layout/src/scan/multi.rs | 21 ++- 4 files changed, 288 insertions(+), 6 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index e0b890c3329..50009b10cd5 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -421,9 +421,7 @@ impl FileOpener for VortexOpener { }) .transpose()?; - if let Some(limit) = limit - && filter.is_none() - { + if let Some(limit) = limit { scan_builder = scan_builder.with_limit(limit); } @@ -1002,6 +1000,31 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_applies_limit_after_filtering() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "filtered-limit/file.vortex"; + let batch = record_batch!(( + "a", + Int32, + vec![Some(1), Some(2), Some(3), Some(4), Some(5), Some(6)] + )) + .unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let file = PartitionedFile::new(file_path.to_string(), data_size); + let table_schema = TableSchema::from_file_schema(batch.schema()); + let filter = logical2physical(&col("a").gt(lit(0_i32)), table_schema.table_schema()); + + let mut opener = make_opener(object_store, table_schema, Some(filter)); + opener.limit = Some(3); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 0b1f150c67e..4eb69cdc86c 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -40,6 +40,8 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; use crate::LayoutReaderRef; +use crate::scan::limit::SharedRowLimit; +use crate::scan::limit::limit_array_stream_shared; use crate::scan::scan_builder::ScanBuilder; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. @@ -157,6 +159,8 @@ impl DataSource for LayoutReaderDataSource { } } + let shared_limit = scan_request.limit.map(SharedRowLimit::new); + Ok(Box::new(LayoutReaderScan { reader: Arc::clone(&self.reader), session: self.session.clone(), @@ -164,6 +168,7 @@ impl DataSource for LayoutReaderDataSource { projection: scan_request.projection, filter: scan_request.filter, limit: scan_request.limit, + shared_limit, selection: scan_request.selection, ordered: scan_request.ordered, metrics_registry: self.metrics_registry.clone(), @@ -185,6 +190,7 @@ struct LayoutReaderScan { projection: Expression, filter: Option, limit: Option, + shared_limit: Option, ordered: bool, selection: Selection, metrics_registry: Option>, @@ -251,6 +257,7 @@ impl Stream for LayoutReaderScan { projection: this.projection.clone(), filter: this.filter.clone(), limit: split_limit, + shared_limit: this.shared_limit.clone(), ordered: this.ordered, row_range, selection: this.selection.clone(), @@ -278,6 +285,7 @@ struct LayoutReaderSplit { projection: Expression, filter: Option, limit: Option, + shared_limit: Option, ordered: bool, row_range: Range, selection: Selection, @@ -312,6 +320,7 @@ impl Partition for LayoutReaderSplit { } fn execute(self: Box) -> VortexResult { + let shared_limit = self.shared_limit.clone(); let builder = ScanBuilder::new(self.session, self.reader) .with_row_range(self.row_range) .with_selection(self.selection) @@ -324,7 +333,7 @@ impl Partition for LayoutReaderSplit { let dtype = builder.dtype()?; // Use into_stream() which creates a LazyScanStream that spawns individual I/O // tasks onto the runtime, enabling parallel execution across executor threads. - let stream = builder.into_stream()?; + let stream = limit_array_stream_shared(builder.into_stream()?, shared_limit); Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, @@ -391,3 +400,143 @@ impl Partition for Empty { ))) } } + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::Arc; + + use futures::TryStreamExt; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_mask::Mask; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; + + use super::LayoutReaderDataSource; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; + use crate::scan::test::session_with_handle; + + #[derive(Debug)] + struct TestLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + } + + impl TestLayoutReader { + fn new(row_count: u64) -> Self { + Self { + name: Arc::from("test"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + } + } + } + + impl LayoutReader for TestLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start)?; + let end = i32::try_from(row_range.end)?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask.await?) + })) + } + } + + #[test] + fn filtered_limit_is_global_across_scan_partitions() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let source = LayoutReaderDataSource::new(Arc::new(TestLayoutReader::new(6)), session) + .with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + assert_eq!(values, [0, 1, 2]); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs index cc99d0f6c6b..098749f628b 100644 --- a/vortex-layout/src/scan/limit.rs +++ b/vortex-layout/src/scan/limit.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; @@ -25,6 +28,45 @@ where } } +/// A row limit shared by streams that execute independent scan partitions. +#[derive(Clone)] +pub(crate) struct SharedRowLimit(Arc); + +impl SharedRowLimit { + pub(crate) fn new(limit: u64) -> Self { + Self(Arc::new(AtomicU64::new(limit))) + } + + fn reserve(&self, requested: u64) -> (u64, bool) { + let mut remaining = self.0.load(Ordering::Relaxed); + loop { + let reserved = remaining.min(requested); + match self.0.compare_exchange_weak( + remaining, + remaining - reserved, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return (reserved, reserved == remaining), + Err(actual) => remaining = actual, + } + } + } +} + +pub(crate) fn limit_array_stream_shared( + stream: S, + limit: Option, +) -> BoxStream<'static, VortexResult> +where + S: Stream> + Send + 'static, +{ + match limit { + Some(limit) => SharedRowLimitedStream::new(stream.boxed(), limit).boxed(), + None => stream.boxed(), + } +} + struct RowLimitedStream { inner: BoxStream<'static, VortexResult>, remaining: u64, @@ -72,3 +114,54 @@ impl Stream for RowLimitedStream { } } } + +struct SharedRowLimitedStream { + inner: BoxStream<'static, VortexResult>, + limit: SharedRowLimit, +} + +impl SharedRowLimitedStream { + fn new(inner: BoxStream<'static, VortexResult>, limit: SharedRowLimit) -> Self { + Self { inner, limit } + } + + fn abort_pending(&mut self) { + let inner = std::mem::replace(&mut self.inner, stream::empty().boxed()); + drop(inner); + } +} + +impl Stream for SharedRowLimitedStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(chunk))) => { + let chunk_len = chunk.len() as u64; + let (reserved, exhausted) = self.limit.reserve(chunk_len); + + if exhausted { + self.abort_pending(); + } + + if reserved == 0 { + if exhausted { + Poll::Ready(None) + } else { + Poll::Ready(Some(Ok(chunk))) + } + } else if reserved == chunk_len { + Poll::Ready(Some(Ok(chunk))) + } else { + let limit = match usize::try_from(reserved) { + Ok(limit) => limit, + Err(_) => unreachable!("reserved rows cannot exceed the current chunk"), + }; + self.abort_pending(); + Poll::Ready(Some(chunk.slice(0..limit))) + } + } + other => other, + } + } +} diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index cb516d2500d..cb6f04b98c9 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -57,6 +57,8 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; +use crate::scan::limit::SharedRowLimit; +use crate::scan::limit::limit_array_stream_shared; use crate::scan::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. @@ -300,10 +302,13 @@ impl DataSource for MultiLayoutDataSource { let dtype = scan_request.projection.return_dtype(&self.dtype)?; + let shared_limit = scan_request.limit.map(SharedRowLimit::new); + Ok(Box::new(MultiLayoutScan { session: self.session.clone(), dtype, request: scan_request, + shared_limit, ready, deferred, handle: self.session.handle(), @@ -320,6 +325,7 @@ struct MultiLayoutScan { session: VortexSession, dtype: DType, request: ScanRequest, + shared_limit: Option, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -345,6 +351,7 @@ impl DataSourceScan for MultiLayoutScan { session, dtype: _, request, + shared_limit, ready, deferred, handle, @@ -399,7 +406,13 @@ impl DataSourceScan for MultiLayoutScan { .chain(deferred_stream) .enumerate() .flat_map(move |(i, reader_result)| match reader_result { - Ok(reader) => reader_partition(i, reader, session.clone(), request.clone()), + Ok(reader) => reader_partition( + i, + reader, + session.clone(), + request.clone(), + shared_limit.clone(), + ), Err(e) => stream::once(async move { Err(e) }).boxed(), }) .boxed() @@ -416,6 +429,7 @@ fn reader_partition( reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, + shared_limit: Option, ) -> PartitionStream { let row_count = reader.row_count(); let row_range = request.row_range.clone().unwrap_or(0..row_count); @@ -461,6 +475,7 @@ fn reader_partition( row_range: Some(row_range), ..request }, + shared_limit, index: partition_idx, }) as PartitionRef) }) @@ -475,6 +490,7 @@ struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, + shared_limit: Option, index: usize, } @@ -510,6 +526,7 @@ impl Partition for MultiLayoutPartition { } fn execute(self: Box) -> VortexResult { + let shared_limit = self.shared_limit.clone(); let request = self.request; let mut builder = ScanBuilder::new(self.session, self.reader) .with_selection(request.selection) @@ -523,7 +540,7 @@ impl Partition for MultiLayoutPartition { } let dtype = builder.dtype()?; - let stream = builder.into_stream()?; + let stream = limit_array_stream_shared(builder.into_stream()?, shared_limit); Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, From f48334985609405695807ad68ca3ba86c7429f39 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 13:47:57 +0100 Subject: [PATCH 06/17] claude CR stuff Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/layout.rs | 9 +- vortex-layout/src/scan/limit.rs | 134 ++++++++++-------------- vortex-layout/src/scan/multi.rs | 9 +- vortex-layout/src/scan/repeated_scan.rs | 116 ++++++++++---------- 4 files changed, 133 insertions(+), 135 deletions(-) diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 4eb69cdc86c..201602e02d3 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -40,8 +40,9 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; use crate::LayoutReaderRef; +use crate::scan::limit::LimitedStream; +use crate::scan::limit::RowBudget; use crate::scan::limit::SharedRowLimit; -use crate::scan::limit::limit_array_stream_shared; use crate::scan::scan_builder::ScanBuilder; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. @@ -333,7 +334,11 @@ impl Partition for LayoutReaderSplit { let dtype = builder.dtype()?; // Use into_stream() which creates a LazyScanStream that spawns individual I/O // tasks onto the runtime, enabling parallel execution across executor threads. - let stream = limit_array_stream_shared(builder.into_stream()?, shared_limit); + let stream = builder.into_stream()?.boxed(); + let stream = match shared_limit { + Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), + None => stream, + }; Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs index 098749f628b..d7b478d438d 100644 --- a/vortex-layout/src/scan/limit.rs +++ b/vortex-layout/src/scan/limit.rs @@ -13,22 +13,15 @@ use futures::StreamExt; use futures::stream; use futures::stream::BoxStream; use vortex_array::ArrayRef; +use vortex_error::VortexExpect; use vortex_error::VortexResult; -pub(crate) fn limit_array_stream( - stream: S, - limit: Option, -) -> BoxStream<'static, VortexResult> -where - S: Stream> + Send + 'static, -{ - match limit { - Some(limit) => RowLimitedStream::new(stream.boxed(), limit).boxed(), - None => stream.boxed(), - } -} - /// A row limit shared by streams that execute independent scan partitions. +/// +/// The shared budget gives "at most `limit` rows in total" semantics without any ordering +/// guarantee across the streams that share it: whichever stream produces a chunk first claims +/// the budget first. It is therefore only correct for consumers that treat the combined output +/// as unordered (e.g. a bare `LIMIT n`), not for order-preserving cross-partition consumption. #[derive(Clone)] pub(crate) struct SharedRowLimit(Arc); @@ -52,99 +45,87 @@ impl SharedRowLimit { } } } -} -pub(crate) fn limit_array_stream_shared( - stream: S, - limit: Option, -) -> BoxStream<'static, VortexResult> -where - S: Stream> + Send + 'static, -{ - match limit { - Some(limit) => SharedRowLimitedStream::new(stream.boxed(), limit).boxed(), - None => stream.boxed(), + fn is_exhausted(&self) -> bool { + self.0.load(Ordering::Relaxed) == 0 } } -struct RowLimitedStream { - inner: BoxStream<'static, VortexResult>, - remaining: u64, -} - -impl RowLimitedStream { - fn new(inner: BoxStream<'static, VortexResult>, remaining: u64) -> Self { - Self { inner, remaining } - } - - fn abort_pending(&mut self) { - let inner = std::mem::replace(&mut self.inner, stream::empty().boxed()); - drop(inner); - } +/// The remaining rows a [`LimitedStream`] may emit, either privately or shared across streams. +pub(crate) enum RowBudget { + /// A budget private to a single stream. + Local(u64), + /// A budget shared across streams executing independent scan partitions. + Shared(SharedRowLimit), } -impl Stream for RowLimitedStream { - type Item = VortexResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.remaining == 0 { - return Poll::Ready(None); +impl RowBudget { + /// Reserve up to `requested` rows from the budget. + /// + /// Returns `(reserved, exhausted)` where `reserved <= requested` and `exhausted` is true + /// once the budget has reached zero. + fn reserve(&mut self, requested: u64) -> (u64, bool) { + match self { + RowBudget::Local(remaining) => { + let reserved = (*remaining).min(requested); + *remaining -= reserved; + (reserved, *remaining == 0) + } + RowBudget::Shared(shared) => shared.reserve(requested), } + } - match self.inner.as_mut().poll_next(cx) { - Poll::Ready(Some(Ok(chunk))) => { - let chunk_len = chunk.len() as u64; - if chunk_len <= self.remaining { - self.remaining -= chunk_len; - if self.remaining == 0 { - self.abort_pending(); - } - Poll::Ready(Some(Ok(chunk))) - } else { - let limit = match usize::try_from(self.remaining) { - Ok(limit) => limit, - Err(_) => unreachable!("remaining rows cannot exceed the current chunk"), - }; - self.remaining = 0; - self.abort_pending(); - Poll::Ready(Some(chunk.slice(0..limit))) - } - } - other => other, + fn is_exhausted(&self) -> bool { + match self { + RowBudget::Local(remaining) => *remaining == 0, + RowBudget::Shared(shared) => shared.is_exhausted(), } } } -struct SharedRowLimitedStream { +/// Wraps a stream, emitting chunks until its [`RowBudget`] is exhausted, then terminating. +pub(crate) struct LimitedStream { inner: BoxStream<'static, VortexResult>, - limit: SharedRowLimit, + budget: RowBudget, } -impl SharedRowLimitedStream { - fn new(inner: BoxStream<'static, VortexResult>, limit: SharedRowLimit) -> Self { - Self { inner, limit } +impl LimitedStream { + pub(crate) fn new( + inner: BoxStream<'static, VortexResult>, + budget: RowBudget, + ) -> Self { + Self { inner, budget } } + /// Drop the inner stream so no further work (including spawned split tasks) is polled. fn abort_pending(&mut self) { - let inner = std::mem::replace(&mut self.inner, stream::empty().boxed()); - drop(inner); + self.inner = stream::empty().boxed(); } } -impl Stream for SharedRowLimitedStream { +impl Stream for LimitedStream { type Item = VortexResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + // Avoid reading a chunk we have no budget for. For a shared budget this also stops a + // partition whose siblings already consumed the limit. + if self.budget.is_exhausted() { + self.abort_pending(); + return Poll::Ready(None); + } + match self.inner.as_mut().poll_next(cx) { Poll::Ready(Some(Ok(chunk))) => { let chunk_len = chunk.len() as u64; - let (reserved, exhausted) = self.limit.reserve(chunk_len); + let (reserved, exhausted) = self.budget.reserve(chunk_len); if exhausted { self.abort_pending(); } if reserved == 0 { + // Either the budget was already exhausted (stop), or this is an empty chunk + // while the budget still has room (pass it through). if exhausted { Poll::Ready(None) } else { @@ -153,11 +134,8 @@ impl Stream for SharedRowLimitedStream { } else if reserved == chunk_len { Poll::Ready(Some(Ok(chunk))) } else { - let limit = match usize::try_from(reserved) { - Ok(limit) => limit, - Err(_) => unreachable!("reserved rows cannot exceed the current chunk"), - }; - self.abort_pending(); + let limit = usize::try_from(reserved) + .vortex_expect("reserved rows are bounded by the chunk length"); Poll::Ready(Some(chunk.slice(0..limit))) } } diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index cb6f04b98c9..2c191dda626 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -57,8 +57,9 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; +use crate::scan::limit::LimitedStream; +use crate::scan::limit::RowBudget; use crate::scan::limit::SharedRowLimit; -use crate::scan::limit::limit_array_stream_shared; use crate::scan::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. @@ -540,7 +541,11 @@ impl Partition for MultiLayoutPartition { } let dtype = builder.dtype()?; - let stream = limit_array_stream_shared(builder.into_stream()?, shared_limit); + let stream = builder.into_stream()?.boxed(); + let stream = match shared_limit { + Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), + None => stream, + }; Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index eff69ffbaa9..693bb3e7aa4 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use futures::FutureExt; use futures::Stream; use futures::StreamExt; +use futures::stream::BoxStream; use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; @@ -21,6 +22,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_scan::selection::Selection; use vortex_session::VortexSession; @@ -28,7 +30,8 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; -use crate::scan::limit::limit_array_stream; +use crate::scan::limit::LimitedStream; +use crate::scan::limit::RowBudget; use crate::scan::splits::Splits; use crate::scan::tasks::{split_exec, TaskContext, TaskFuture}; @@ -170,15 +173,19 @@ impl RepeatedScan { } } + fn task_context(&self) -> Arc { + Arc::new(TaskContext { + filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), + reader: Arc::clone(&self.layout_reader), + projection: self.projection.clone(), + }) + } + pub(crate) fn execute( &self, row_range: Option>, ) -> VortexResult>>> { - let ctx = Arc::new(TaskContext { - filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), - reader: Arc::clone(&self.layout_reader), - projection: self.projection.clone(), - }); + let ctx = self.task_context(); let mut limit = self.limit.filter(|_| self.filter.is_none()); let mut tasks = Vec::new(); @@ -206,62 +213,65 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult> + Send + 'static> { + let num_workers = get_available_parallelism().unwrap_or(1); + let concurrency = self.concurrency * num_workers; let handle = self.session.handle(); + // With both a filter and a limit we cannot know each split's output row count ahead of + // time, so split tasks are built lazily as the stream is polled. `buffered`'s read-ahead + // (bounded by `concurrency`) registers IO for splits eagerly, but only as far as the + // limit requires: `limit_array_stream` drops the inner stream once the limit is reached, + // capping over-read at `concurrency` splits. if self.filter.is_some() && self.limit.is_some() { - let ctx = Arc::new(TaskContext { - filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))), - reader: Arc::clone(&self.layout_reader), - projection: self.projection.clone(), - }); + let ctx = self.task_context(); let selection = self.selection.clone(); - let ordered = self.ordered; - let limit = self.limit; - let stream = futures::stream::iter(self.split_ranges(row_range)).filter_map(move |range| { - let row_mask = selection.row_mask(&range); - if row_mask.mask().all_false() { - return async { None }.boxed(); - } - - let ctx = Arc::clone(&ctx); - let handle = handle.clone(); - async move { - Some( - async move { - let task = split_exec(ctx, row_mask, None)?; - handle.spawn(task).await - } - .boxed(), - ) - } - .boxed() - }); - let stream = if ordered { - stream.buffered(1).boxed() - } else { - stream.buffer_unordered(1).boxed() - }; + let tasks = + futures::stream::iter(self.split_ranges(row_range)).filter_map(move |range| { + // Build the row mask and split task synchronously so the IO system sees the + // split's ranges as soon as `buffered` pulls it, without cloning `selection`. + let row_mask = selection.row_mask(&range); + let spawned = (!row_mask.mask().all_false()).then(|| { + let task = split_exec(Arc::clone(&ctx), row_mask, None) + .unwrap_or_else(|err| async move { Err(err) }.boxed()); + handle.spawn(task) + }); + async move { spawned } + }); - return Ok(limit_array_stream( - stream.filter_map(|chunk| async move { chunk.transpose() }), - limit, - )); + return Ok(schedule(tasks, self.ordered, concurrency, self.limit)); } - let num_workers = get_available_parallelism().unwrap_or(1); - let concurrency = self.concurrency * num_workers; - let stream = + // No filter (or no limit): build every task eagerly so the IO system sees all split + // ranges up front. A no-filter limit is applied exactly per split inside `execute`. + let tasks = futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); - let stream = if self.ordered { - stream.buffered(concurrency).boxed() - } else { - stream.buffer_unordered(concurrency).boxed() - }; - Ok(limit_array_stream( - stream.filter_map(|chunk| async move { chunk.transpose() }), - self.limit, - )) + Ok(schedule(tasks, self.ordered, concurrency, self.limit)) + } +} + +/// Spawn-buffer a stream of split tasks, transposing empty splits away and applying `limit`. +fn schedule( + tasks: S, + ordered: bool, + concurrency: usize, + limit: Option, +) -> BoxStream<'static, VortexResult> +where + S: Stream>>> + Send + 'static, +{ + let stream = if ordered { + tasks.buffered(concurrency).boxed() + } else { + tasks.buffer_unordered(concurrency).boxed() + }; + let stream = stream + .filter_map(|chunk| async move { chunk.transpose() }) + .boxed(); + + match limit { + Some(limit) => LimitedStream::new(stream, RowBudget::Local(limit)).boxed(), + None => stream, } } From 2df10c2fbcacb87f75052e8739bb41aa98f27142 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 13:54:39 +0100 Subject: [PATCH 07/17] remove generic task Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/repeated_scan.rs | 2 +- vortex-layout/src/scan/tasks.rs | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 693bb3e7aa4..bbed9be40ff 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -184,7 +184,7 @@ impl RepeatedScan { pub(crate) fn execute( &self, row_range: Option>, - ) -> VortexResult>>> { + ) -> VortexResult> { let ctx = self.task_context(); let mut limit = self.limit.filter(|_| self.filter.is_none()); diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 40d770a9903..ad3c4d3e2da 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -19,7 +19,7 @@ use vortex_scan::row_mask::RowMask; use crate::LayoutReader; use crate::scan::filter::FilterExpr; -pub type TaskFuture = BoxFuture<'static, VortexResult>; +pub type TaskFuture = BoxFuture<'static, VortexResult>>; /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this @@ -33,12 +33,13 @@ pub type TaskFuture = BoxFuture<'static, VortexResult>; /// The intersected row range is then further reduced via expression-based pruning. After pruning /// has eliminated more blocks, the full filter is executed over the remainder of the split. /// -/// This mask is then provided to the reader to perform a filtered projection over the split data. +/// This mask is then provided to the reader to perform a filtered projection over the split data, +/// yielding the projected array (or `None` when the split selects no rows). pub fn split_exec( ctx: Arc, read_mask: RowMask, limit: Option<&mut u64>, -) -> VortexResult>> { +) -> VortexResult { let row_range = read_mask.row_range(); let row_mask = read_mask.mask().clone(); From d49284cec94a713cc00a0e8d640874856e946e5b Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 14:02:22 +0100 Subject: [PATCH 08/17] docs Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/layout.rs | 2 ++ vortex-layout/src/scan/limit.rs | 10 +++++----- vortex-layout/src/scan/multi.rs | 2 ++ vortex-layout/src/scan/repeated_scan.rs | 5 +---- 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 201602e02d3..d2145fab29f 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -335,6 +335,8 @@ impl Partition for LayoutReaderSplit { // Use into_stream() which creates a LazyScanStream that spawns individual I/O // tasks onto the runtime, enabling parallel execution across executor threads. let stream = builder.into_stream()?.boxed(); + // Caps total rows across all partitions; only correct for unordered consumption + // (see `SharedRowLimit`). let stream = match shared_limit { Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), None => stream, diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs index d7b478d438d..855c00e6ed2 100644 --- a/vortex-layout/src/scan/limit.rs +++ b/vortex-layout/src/scan/limit.rs @@ -16,12 +16,12 @@ use vortex_array::ArrayRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; -/// A row limit shared by streams that execute independent scan partitions. +/// A row limit shared by the streams executing one scan's independent partitions. /// -/// The shared budget gives "at most `limit` rows in total" semantics without any ordering -/// guarantee across the streams that share it: whichever stream produces a chunk first claims -/// the budget first. It is therefore only correct for consumers that treat the combined output -/// as unordered (e.g. a bare `LIMIT n`), not for order-preserving cross-partition consumption. +/// The single budget is claimed in completion order, not row order, so the combined output is +/// "any `limit` rows", not "the first `limit` in scan order": under concurrency a later partition +/// can drain the budget and starve an earlier one. Only sound for unordered consumers (a bare +/// `LIMIT n`); order-preserving consumers must use a per-partition `ScanBuilder::with_limit`. #[derive(Clone)] pub(crate) struct SharedRowLimit(Arc); diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index 2c191dda626..d78695884d7 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -542,6 +542,8 @@ impl Partition for MultiLayoutPartition { let dtype = builder.dtype()?; let stream = builder.into_stream()?.boxed(); + // Caps total rows across all partitions; only correct for unordered consumption + // (see `SharedRowLimit`). let stream = match shared_limit { Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), None => stream, diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index bbed9be40ff..073724b52a7 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -181,10 +181,7 @@ impl RepeatedScan { }) } - pub(crate) fn execute( - &self, - row_range: Option>, - ) -> VortexResult> { + pub(crate) fn execute(&self, row_range: Option>) -> VortexResult> { let ctx = self.task_context(); let mut limit = self.limit.filter(|_| self.filter.is_none()); From 28d6bc12be6ee6b3c6b021d530135cae320c6c3d Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 15:56:30 +0100 Subject: [PATCH 09/17] async start Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/scan_builder.rs | 110 +++++++++++++++++++++++-- 1 file changed, 104 insertions(+), 6 deletions(-) diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index cc2604e9a88..a6d98ffce85 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -286,6 +286,9 @@ impl ScanBuilder { } /// Returns a [`Stream`] with tasks spawned onto the session's runtime handle. + /// + /// Preparation and initial stream construction begin on the first poll. Errors from either + /// step are returned as the stream's next item. pub fn into_stream( self, ) -> VortexResult> + Send + 'static> { @@ -310,7 +313,7 @@ enum LazyScanState { } struct PreparingScan { - task: Task>, + task: Task>>>, } struct LazyScanStream { @@ -336,15 +339,20 @@ impl Stream for LazyScanStream { LazyScanState::Builder(builder) => { let builder = builder.take().vortex_expect("polled after completion"); let handle = builder.session.handle(); - let task = handle.spawn_blocking(move || builder.prepare()); + // IMPORTANT: Building the stream can synchronously walk the layout and + // register I/O for every split. Keep it with preparation in this blocking + // task: poll_next must only wait for and poll an already-constructed stream. + // This also keeps construction errors on the Preparing -> Error path rather + // than running construction on the caller's executor. + let task = handle.spawn_blocking(move || { + let scan = builder.prepare()?; + Ok(scan.execute_stream(None)?.boxed()) + }); self.state = LazyScanState::Preparing(PreparingScan { task }); } LazyScanState::Preparing(preparing) => { match ready!(Pin::new(&mut preparing.task).poll(cx)) { - Ok(scan) => match scan.execute_stream(None) { - Ok(stream) => self.state = LazyScanState::Stream(stream.boxed()), - Err(err) => self.state = LazyScanState::Error(Some(err)), - }, + Ok(stream) => self.state = LazyScanState::Stream(stream), Err(err) => self.state = LazyScanState::Error(Some(err)), } } @@ -411,6 +419,7 @@ mod test { use vortex_error::vortex_err; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_io::runtime::tokio::TokioRuntime; use vortex_mask::Mask; use super::ScanBuilder; @@ -571,6 +580,13 @@ mod test { dtype: DType, row_count: u64, register_splits_calls: Arc, + blocking_projection: Option, + } + + #[derive(Debug)] + struct BlockingProjection { + started: mpsc::Sender<()>, + gate: Arc>, } impl SplittingLayoutReader { @@ -580,8 +596,19 @@ mod test { dtype: DType::Primitive(PType::I32, Nullability::NonNullable), row_count: 4, register_splits_calls, + blocking_projection: None, } } + + fn with_blocking_projection( + register_splits_calls: Arc, + gate: Arc>, + started: mpsc::Sender<()>, + ) -> Self { + let mut reader = Self::new(register_splits_calls); + reader.blocking_projection = Some(BlockingProjection { started, gate }); + reader + } } impl LayoutReader for SplittingLayoutReader { @@ -634,6 +661,14 @@ mod test { _expr: &Expression, _mask: MaskFuture, ) -> VortexResult { + if let Some(blocking_projection) = &self.blocking_projection { + blocking_projection + .started + .send(()) + .map_err(|_| vortex_err!("test projection-start receiver dropped"))?; + let _guard = blocking_projection.gate.lock(); + } + let start = usize::try_from(row_range.start) .map_err(|_| vortex_err!("row_range.start must fit in usize"))?; let end = usize::try_from(row_range.end) @@ -676,6 +711,69 @@ mod test { Ok(()) } + #[tokio::test] + async fn into_stream_constructs_tasks_off_the_poller() -> VortexResult<()> { + let gate = Arc::new(Mutex::new(())); + let guard = gate.lock(); + let calls = Arc::new(AtomicUsize::new(0)); + let (started_send, started_recv) = mpsc::channel(); + let reader = Arc::new(SplittingLayoutReader::with_blocking_projection( + Arc::clone(&calls), + Arc::clone(&gate), + started_send, + )); + + let runtime = TokioRuntime::new(tokio::runtime::Handle::current()); + let session = session_with_handle(runtime.handle()); + let mut stream = ScanBuilder::new(session, reader).into_stream()?; + + let (poll_send, poll_recv) = mpsc::channel(); + let (release_send, release_recv) = mpsc::channel(); + let join = std::thread::spawn(move || { + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + let poll = Pin::new(&mut stream).poll_next(&mut cx); + let _ = poll_send.send(matches!(poll, Poll::Pending)); + let _ = release_recv.recv(); + }); + + let poll_result = poll_recv.recv_timeout(Duration::from_secs(1)); + let projection_started = started_recv.recv_timeout(Duration::from_secs(1)); + + // Release the task and join its caller before reporting a failed assertion. + drop(guard); + let _ = release_send.send(()); + drop(join.join()); + + assert!( + poll_result.is_ok_and(|poll_pending| poll_pending), + "first poll must return while scan task construction is blocked" + ); + projection_started + .map_err(|_| vortex_err!("stream construction did not begin in the background"))?; + assert_eq!(calls.load(Ordering::Relaxed), 1); + + Ok(()) + } + + #[tokio::test] + async fn into_stream_reports_stream_construction_errors() -> VortexResult<()> { + let range_start = i32::MAX as u64 + 1; + let reader = Arc::new(SplittingLayoutReader::new(Arc::new(AtomicUsize::new(0)))); + let session = session_with_handle(TokioRuntime::current()); + let mut stream = ScanBuilder::new(session, reader) + .with_row_range(range_start..range_start + 1) + .into_stream()?; + + assert!(matches!( + futures::StreamExt::next(&mut stream).await, + Some(Err(_)) + )); + assert!(futures::StreamExt::next(&mut stream).await.is_none()); + + Ok(()) + } + #[derive(Debug)] struct FilteringLayoutReader { name: Arc, From d768fe2b17d600d16b411c15a580298edead4257 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 10:45:37 +0100 Subject: [PATCH 10/17] more code? Signed-off-by: Adam Gutglick --- vortex-datafusion/src/v2/source.rs | 189 ++++++++++++- vortex-layout/src/scan/layout.rs | 312 +++++++++++++++++--- vortex-layout/src/scan/limit.rs | 143 ++-------- vortex-layout/src/scan/multi.rs | 360 +++++++++++++++++++++--- vortex-layout/src/scan/repeated_scan.rs | 142 ++++++++-- vortex-layout/src/scan/scan_builder.rs | 273 ++++++++++++++++++ vortex-layout/src/scan/tasks.rs | 258 +++++++++++------ vortex-scan/src/lib.rs | 9 +- 8 files changed, 1364 insertions(+), 322 deletions(-) diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 325f8eae92f..efcec06dc8e 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -97,10 +97,15 @@ use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; use futures::TryStreamExt; use futures::future::try_join_all; +use futures::stream::BoxStream; +use tokio_stream::wrappers::ReceiverStream; +use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; +use vortex::array::stream::SendableArrayStream; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::Nullability; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::expr::Expression; @@ -110,6 +115,7 @@ use vortex::expr::pack; use vortex::expr::root; use vortex::expr::stats::Precision; use vortex::expr::transform::replace; +use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; use vortex::scan::DataSourceRef; use vortex::scan::ScanRequest; @@ -402,6 +408,7 @@ impl DataSource for VortexDataSource { )); let session = self.session.clone(); let num_partitions = self.num_partitions; + let ordered = self.ordered; // Pre-build the leftover projector (if any) so we can apply it after batch conversion. let leftover_projector = self @@ -418,17 +425,18 @@ impl DataSource for VortexDataSource { .map_err(|e| DataFusionError::External(Box::new(e)))?; // Each split.execute() returns a lazy stream whose early polls do preparation - // work (expression resolution, layout traversal, first I/O spawns). We use - // try_flatten_unordered to poll multiple split streams concurrently so that - // the next split is already warm when the current one finishes. + // work (expression resolution, layout traversal, first I/O spawns). Both ordering + // modes flatten with cross-partition I/O look-ahead; the ordered path additionally + // preserves partition order so an ordered global limit cannot observe later rows first. let scan_streams = scan.partitions().map(|split_result| { let split = split_result?; split.execute() }); let handle = session.handle(); - let stream = scan_streams - .try_flatten_unordered(Some(num_partitions * 2)) + let chunks = + flatten_scan_streams(scan_streams, ordered, num_partitions * 2, handle.clone()); + let stream = chunks .map(move |result| { let session = session.clone(); let target_field = Arc::clone(&projected_target_field); @@ -654,6 +662,54 @@ impl DataSource for VortexDataSource { } } +fn flatten_scan_streams( + scan_streams: S, + ordered: bool, + concurrency: usize, + handle: Handle, +) -> BoxStream<'static, VortexResult> +where + S: futures::Stream> + Send + 'static, +{ + if !ordered { + return scan_streams + .try_flatten_unordered(Some(concurrency)) + .boxed(); + } + + // Ordered output must preserve global row order, but later partitions should still warm up + // (start their I/O) while an earlier one is draining. Each partition is drained by a spawned + // task into a bounded channel: spawning starts its I/O immediately, and the bounded channel + // back-pressures so read-ahead stays capped. `buffered` warms up to `concurrency` partitions + // ahead while `try_flatten` emits their chunks strictly in partition order. At most + // `concurrency * CHUNKS_AHEAD` chunks are buffered across the warming partitions. + const CHUNKS_AHEAD: usize = 2; + let lookahead = concurrency.max(1); + scan_streams + .map(move |stream_result| { + let handle = handle.clone(); + async move { + let mut stream = stream_result?; + let (tx, rx) = tokio::sync::mpsc::channel(CHUNKS_AHEAD); + handle + .spawn(async move { + while let Some(item) = stream.next().await { + // A send error means the consumer was dropped (for example once a + // global limit is reached), so stop draining and release the I/O. + if tx.send(item).await.is_err() { + break; + } + } + }) + .detach(); + Ok::<_, VortexError>(ReceiverStream::new(rx)) + } + }) + .buffered(lookahead) + .try_flatten() + .boxed() +} + /// Convert a Vortex [`Option`] to a DataFusion /// [`DataFusionPrecision`]. /// @@ -665,3 +721,126 @@ fn estimate_to_df_precision(est: &Precision) -> DFPrecision { Precision::Absent => DFPrecision::Absent, } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicU8; + use std::sync::atomic::Ordering; + use std::task::Poll; + + use futures::TryStreamExt; + use vortex::array::ArrayRef; + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::array_session; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::stream::ArrayStreamAdapter; + use vortex::array::stream::ArrayStreamExt; + use vortex::dtype::DType; + use vortex::dtype::Nullability; + use vortex::dtype::PType; + use vortex::error::VortexError; + use vortex::error::VortexResult; + use vortex::io::runtime::tokio::TokioRuntime; + + use super::flatten_scan_streams; + + fn i32_stream(dtype: &DType, chunks: Vec>) -> super::SendableArrayStream { + let items = chunks + .into_iter() + .map(|values| Ok(PrimitiveArray::from_iter(values).into_array())) + .collect::>>(); + ArrayStreamAdapter::new(dtype.clone(), futures::stream::iter(items)).boxed() + } + + fn collect_i32(chunks: Vec) -> VortexResult> { + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for chunk in chunks { + let primitive = chunk.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + Ok(values) + } + + #[tokio::test] + async fn ordered_flatten_preserves_partition_order() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let first = i32_stream(&dtype, vec![vec![0, 1], vec![2]]); + let second = i32_stream(&dtype, vec![vec![10, 11]]); + + let streams = + futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); + let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let chunks = flattened.try_collect::>().await?; + + assert_eq!(collect_i32(chunks)?, [0, 1, 2, 10, 11]); + Ok(()) + } + + #[tokio::test] + async fn ordered_flatten_warms_later_partition_while_first_pending() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + + // The first partition self-wakes `Pending` a few times, then yields one chunk and ends. + let first_done = Arc::new(AtomicBool::new(false)); + let first_done_for_stream = Arc::clone(&first_done); + let mut pending_polls = 3usize; + let mut emitted = false; + let first = ArrayStreamAdapter::new( + dtype.clone(), + futures::stream::poll_fn(move |cx| -> Poll>> { + if pending_polls > 0 { + pending_polls -= 1; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + if !emitted { + emitted = true; + return Poll::Ready(Some(Ok(PrimitiveArray::from_iter([0i32]).into_array()))); + } + first_done_for_stream.store(true, Ordering::SeqCst); + Poll::Ready(None) + }), + ) + .boxed(); + + // The second partition records, at its first poll, whether the first had already finished. + // 0 = not yet polled, 1 = first still running, 2 = first already done. + let observation = Arc::new(AtomicU8::new(0)); + let observation_for_stream = Arc::clone(&observation); + let first_done_for_second = Arc::clone(&first_done); + let mut first_poll = true; + let second = ArrayStreamAdapter::new( + dtype, + futures::stream::poll_fn(move |_| -> Poll>> { + if first_poll { + first_poll = false; + let seen = if first_done_for_second.load(Ordering::SeqCst) { + 2 + } else { + 1 + }; + observation_for_stream.store(seen, Ordering::SeqCst); + return Poll::Ready(Some(Ok(PrimitiveArray::from_iter([100i32]).into_array()))); + } + Poll::Ready(None) + }), + ) + .boxed(); + + let streams = + futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); + let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let chunks = flattened.try_collect::>().await?; + + // Global order is preserved: the first partition's row precedes the later partition's. + assert_eq!(collect_i32(chunks)?, [0, 100]); + // The later partition was polled (warmed) while the first was still pending, proving the + // ordered flatten keeps cross-partition look-ahead instead of serializing. + assert_eq!(observation.load(Ordering::SeqCst), 1); + Ok(()) + } +} diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index d2145fab29f..b018b6bd98c 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -40,9 +40,7 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; use crate::LayoutReaderRef; -use crate::scan::limit::LimitedStream; -use crate::scan::limit::RowBudget; -use crate::scan::limit::SharedRowLimit; +use crate::scan::limit::RowLimit; use crate::scan::scan_builder::ScanBuilder; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. @@ -160,7 +158,18 @@ impl DataSource for LayoutReaderDataSource { } } - let shared_limit = scan_request.limit.map(SharedRowLimit::new); + // An ordered limit must see earlier filtered rows before later external partitions are + // allowed to reserve any budget. Emit one partition for that path; its inner scan keeps + // the limit at mask level. Unordered partitions share a completion-order budget instead. + let ordered_limit = scan_request.ordered && scan_request.limit.is_some(); + let row_limit = (!scan_request.ordered) + .then(|| scan_request.limit.map(RowLimit::new)) + .flatten(); + let split_size = if ordered_limit { + row_range.end - row_range.start + } else { + self.split_max_row_count + }; Ok(Box::new(LayoutReaderScan { reader: Arc::clone(&self.reader), @@ -169,13 +178,13 @@ impl DataSource for LayoutReaderDataSource { projection: scan_request.projection, filter: scan_request.filter, limit: scan_request.limit, - shared_limit, + row_limit, selection: scan_request.selection, ordered: scan_request.ordered, metrics_registry: self.metrics_registry.clone(), next_row: row_range.start, end_row: row_range.end, - split_size: self.split_max_row_count, + split_size, })) } @@ -191,7 +200,7 @@ struct LayoutReaderScan { projection: Expression, filter: Option, limit: Option, - shared_limit: Option, + row_limit: Option, ordered: bool, selection: Selection, metrics_registry: Option>, @@ -232,33 +241,22 @@ impl Stream for LayoutReaderScan { if this.limit.is_some_and(|limit| limit == 0) { return Poll::Ready(None); } + if this.row_limit.as_ref().is_some_and(RowLimit::is_exhausted) { + return Poll::Ready(None); + } let split_end = this .next_row .saturating_add(this.split_size) .min(this.end_row); let row_range = this.next_row..split_end; - let split_rows = split_end - this.next_row; - - let split_limit = this.limit; - // Only decrement the remaining limit when there is no filter. With a filter, - // the actual output row count is unknown (could be anywhere from 0 to split_rows), - // so decrementing by split_rows would be too aggressive and could stop producing - // splits before the limit is reached. Instead, pass the full remaining limit to - // each split and let the engine enforce the exact limit at the stream level. - if this.filter.is_none() - && let Some(ref mut limit) = this.limit - { - *limit = limit.saturating_sub(split_rows); - } - let split = Box::new(LayoutReaderSplit { reader: Arc::clone(&this.reader), session: this.session.clone(), projection: this.projection.clone(), filter: this.filter.clone(), - limit: split_limit, - shared_limit: this.shared_limit.clone(), + limit: this.limit, + row_limit: this.row_limit.clone(), ordered: this.ordered, row_range, selection: this.selection.clone(), @@ -271,7 +269,10 @@ impl Stream for LayoutReaderScan { } fn size_hint(&self) -> (usize, Option) { - if self.next_row >= self.end_row { + if self.next_row >= self.end_row + || self.limit.is_some_and(|limit| limit == 0) + || self.row_limit.as_ref().is_some_and(RowLimit::is_exhausted) + { return (0, Some(0)); } let remaining_rows = self.end_row - self.next_row; @@ -286,7 +287,7 @@ struct LayoutReaderSplit { projection: Expression, filter: Option, limit: Option, - shared_limit: Option, + row_limit: Option, ordered: bool, row_range: Range, selection: Selection, @@ -309,7 +310,7 @@ impl Partition for LayoutReaderSplit { let row_count = self.selection.row_count(row_count); let row_count = self.limit.map_or(row_count, |limit| row_count.min(limit)); - if self.filter.is_some() { + if self.filter.is_some() || self.row_limit.is_some() { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -321,13 +322,13 @@ impl Partition for LayoutReaderSplit { } fn execute(self: Box) -> VortexResult { - let shared_limit = self.shared_limit.clone(); let builder = 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_row_limit(self.row_limit) .with_some_metrics_registry(self.metrics_registry) .with_ordered(self.ordered); @@ -335,12 +336,6 @@ impl Partition for LayoutReaderSplit { // Use into_stream() which creates a LazyScanStream that spawns individual I/O // tasks onto the runtime, enabling parallel execution across executor threads. let stream = builder.into_stream()?.boxed(); - // Caps total rows across all partitions; only correct for unordered consumption - // (see `SharedRowLimit`). - let stream = match shared_limit { - Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), - None => stream, - }; Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, @@ -412,8 +407,11 @@ impl Partition for Empty { mod tests { use std::ops::Range; use std::sync::Arc; + use std::task::Poll; + use futures::StreamExt; use futures::TryStreamExt; + use parking_lot::Mutex; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; @@ -444,6 +442,7 @@ mod tests { name: Arc, dtype: DType, row_count: u64, + projection_masks: Option>>>, } impl TestLayoutReader { @@ -452,8 +451,14 @@ mod tests { name: Arc::from("test"), dtype: DType::Primitive(PType::I32, Nullability::NonNullable), row_count, + projection_masks: None, } } + + fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = Some(projection_masks); + self + } } impl LayoutReader for TestLayoutReader { @@ -508,8 +513,124 @@ mod tests { mask: MaskFuture, ) -> VortexResult { let row_range = row_range.clone(); + let projection_masks = self.projection_masks.clone(); + + Ok(Box::pin(async move { + let mask = mask.await?; + if let Some(projection_masks) = projection_masks { + projection_masks.lock().push(mask.true_count()); + } + let start = i32::try_from(row_range.start)?; + let end = i32::try_from(row_range.end)?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask) + })) + } + } + + #[derive(Debug)] + struct DelayedFirstSplitReader { + name: Arc, + dtype: DType, + projection_ranges: Arc>>>, + filter_ranges: Arc>>>, + } + + impl DelayedFirstSplitReader { + fn new(projection_ranges: Arc>>>) -> Self { + Self { + name: Arc::from("delayed-first-split"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + projection_ranges, + filter_ranges: Arc::new(Mutex::new(Vec::new())), + } + } + + /// Row ranges whose filter was actually evaluated (recorded when the split's filter + /// future runs, not when it is merely scheduled). + fn filter_ranges(&self) -> Arc>>> { + Arc::clone(&self.filter_ranges) + } + } + + impl LayoutReader for DelayedFirstSplitReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 4 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + splits.push(split_range.row_offset() + row_range.start + 2); + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + self.filter_ranges.lock().push(row_range.clone()); + let delay = row_range.start == 0; + let len = mask.len(); + + Ok(MaskFuture::new(len, async move { + if delay { + let mut yielded = false; + futures::future::poll_fn(move |cx| { + if yielded { + Poll::Ready(()) + } else { + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + }) + .await; + } + mask.await + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let projection_ranges = Arc::clone(&self.projection_ranges); Ok(Box::pin(async move { + projection_ranges.lock().push(row_range.clone()); let start = i32::try_from(row_range.start)?; let end = i32::try_from(row_range.end)?; PrimitiveArray::from_iter(start..end) @@ -533,6 +654,7 @@ mod tests { ..Default::default() }))?; let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); let mut ctx = array_session().create_execution_ctx(); let mut values = Vec::new(); @@ -546,4 +668,128 @@ mod tests { assert_eq!(values, [0, 1, 2]); Ok(()) } + + #[test] + fn ordered_filtered_limit_waits_for_the_earlier_split() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_ranges = Arc::new(Mutex::new(Vec::new())); + let source = LayoutReaderDataSource::new( + Arc::new(DelayedFirstSplitReader::new(Arc::clone(&projection_ranges))), + session, + ) + .with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(1), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + assert_eq!(values, [0]); + let projection_ranges = projection_ranges.lock(); + assert_eq!(projection_ranges.len(), 1); + assert_eq!(projection_ranges[0], 0..2); + Ok(()) + } + + #[test] + fn ordered_filtered_limit_evaluates_later_split_filter_concurrently() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_ranges = Arc::new(Mutex::new(Vec::new())); + let reader = DelayedFirstSplitReader::new(Arc::clone(&projection_ranges)); + let filter_ranges = reader.filter_ranges(); + let source = + LayoutReaderDataSource::new(Arc::new(reader), session).with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(1), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + // Ordered LIMIT is still exact: only the first split's earliest row is projected. + assert_eq!(values, [0]); + let projection_ranges = projection_ranges.lock(); + assert_eq!(projection_ranges.len(), 1); + assert_eq!(projection_ranges[0], 0..2); + drop(projection_ranges); + + // But the later split's filter still runs while the delayed first split reserves, so + // prefetch is not disabled (serializing to concurrency=1 would only ever filter 0..2). + let filter_ranges = filter_ranges.lock(); + assert!( + filter_ranges.contains(&(0..2)) && filter_ranges.contains(&(2..4)), + "expected both splits' filters to be evaluated, got {filter_ranges:?}" + ); + Ok(()) + } + + #[test] + fn unordered_limit_never_projects_more_than_the_global_budget() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let source = LayoutReaderDataSource::new( + Arc::new( + TestLayoutReader::new(12).with_projection_masks(Arc::clone(&projection_masks)), + ), + session, + ) + .with_split_max_row_count(2); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: false, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + let chunks = runtime.block_on( + futures::stream::iter(partitions) + .map(|partition| partition.execute()) + .try_flatten_unordered(Some(6)) + .try_collect::>(), + )?; + + let mut ctx = array_session().create_execution_ctx(); + let values = chunks + .into_iter() + .map(|chunk| { + chunk + .execute::(&mut ctx) + .map(|primitive| primitive.into_buffer::()) + }) + .collect::>>()?; + let row_count = values.iter().map(|values| values.len()).sum::(); + + assert_eq!(row_count, 3); + assert_eq!(projection_masks.lock().iter().sum::(), 3); + Ok(()) + } } diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs index 855c00e6ed2..60dcf921432 100644 --- a/vortex-layout/src/scan/limit.rs +++ b/vortex-layout/src/scan/limit.rs @@ -1,145 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use futures::Stream; -use futures::StreamExt; -use futures::stream; -use futures::stream::BoxStream; -use vortex_array::ArrayRef; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; +use vortex_mask::Mask; -/// A row limit shared by the streams executing one scan's independent partitions. +/// A cloneable row limit shared by all work that can contribute rows to one scan. /// -/// The single budget is claimed in completion order, not row order, so the combined output is -/// "any `limit` rows", not "the first `limit` in scan order": under concurrency a later partition -/// can drain the budget and starve an earlier one. Only sound for unordered consumers (a bare -/// `LIMIT n`); order-preserving consumers must use a per-partition `ScanBuilder::with_limit`. +/// Rows are reserved from a selection mask before projection work is constructed. This keeps +/// rows that cannot be returned out of projection evaluation entirely. When a limit is shared by +/// concurrent unordered partitions, reservation order is completion order, so callers may return +/// any matching rows. Ordered limited scans instead serialize their external partitions before +/// sharing a `RowLimit`, preserving the first matching rows in scan order. #[derive(Clone)] -pub(crate) struct SharedRowLimit(Arc); +pub(crate) struct RowLimit(Arc); -impl SharedRowLimit { +impl RowLimit { pub(crate) fn new(limit: u64) -> Self { Self(Arc::new(AtomicU64::new(limit))) } - fn reserve(&self, requested: u64) -> (u64, bool) { + /// Reserve rows selected by `mask` and retain only the earliest granted rows in that mask. + pub(crate) fn limit(&self, mask: Mask) -> Mask { + let requested = u64::try_from(mask.true_count()).unwrap_or(u64::MAX); + let granted = self.reserve(requested); + let granted = usize::try_from(granted).unwrap_or(usize::MAX); + mask.limit(granted) + } + + pub(crate) fn is_exhausted(&self) -> bool { + self.0.load(Ordering::Relaxed) == 0 + } + + fn reserve(&self, requested: u64) -> u64 { let mut remaining = self.0.load(Ordering::Relaxed); loop { - let reserved = remaining.min(requested); + let granted = remaining.min(requested); match self.0.compare_exchange_weak( remaining, - remaining - reserved, + remaining - granted, Ordering::Relaxed, Ordering::Relaxed, ) { - Ok(_) => return (reserved, reserved == remaining), + Ok(_) => return granted, Err(actual) => remaining = actual, } } } - - fn is_exhausted(&self) -> bool { - self.0.load(Ordering::Relaxed) == 0 - } -} - -/// The remaining rows a [`LimitedStream`] may emit, either privately or shared across streams. -pub(crate) enum RowBudget { - /// A budget private to a single stream. - Local(u64), - /// A budget shared across streams executing independent scan partitions. - Shared(SharedRowLimit), -} - -impl RowBudget { - /// Reserve up to `requested` rows from the budget. - /// - /// Returns `(reserved, exhausted)` where `reserved <= requested` and `exhausted` is true - /// once the budget has reached zero. - fn reserve(&mut self, requested: u64) -> (u64, bool) { - match self { - RowBudget::Local(remaining) => { - let reserved = (*remaining).min(requested); - *remaining -= reserved; - (reserved, *remaining == 0) - } - RowBudget::Shared(shared) => shared.reserve(requested), - } - } - - fn is_exhausted(&self) -> bool { - match self { - RowBudget::Local(remaining) => *remaining == 0, - RowBudget::Shared(shared) => shared.is_exhausted(), - } - } -} - -/// Wraps a stream, emitting chunks until its [`RowBudget`] is exhausted, then terminating. -pub(crate) struct LimitedStream { - inner: BoxStream<'static, VortexResult>, - budget: RowBudget, -} - -impl LimitedStream { - pub(crate) fn new( - inner: BoxStream<'static, VortexResult>, - budget: RowBudget, - ) -> Self { - Self { inner, budget } - } - - /// Drop the inner stream so no further work (including spawned split tasks) is polled. - fn abort_pending(&mut self) { - self.inner = stream::empty().boxed(); - } -} - -impl Stream for LimitedStream { - type Item = VortexResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - // Avoid reading a chunk we have no budget for. For a shared budget this also stops a - // partition whose siblings already consumed the limit. - if self.budget.is_exhausted() { - self.abort_pending(); - return Poll::Ready(None); - } - - match self.inner.as_mut().poll_next(cx) { - Poll::Ready(Some(Ok(chunk))) => { - let chunk_len = chunk.len() as u64; - let (reserved, exhausted) = self.budget.reserve(chunk_len); - - if exhausted { - self.abort_pending(); - } - - if reserved == 0 { - // Either the budget was already exhausted (stop), or this is an empty chunk - // while the budget still has room (pass it through). - if exhausted { - Poll::Ready(None) - } else { - Poll::Ready(Some(Ok(chunk))) - } - } else if reserved == chunk_len { - Poll::Ready(Some(Ok(chunk))) - } else { - let limit = usize::try_from(reserved) - .vortex_expect("reserved rows are bounded by the chunk length"); - Poll::Ready(Some(chunk.slice(0..limit))) - } - } - other => other, - } - } } diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index d78695884d7..31d0c13e2ee 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -57,9 +57,7 @@ use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; -use crate::scan::limit::LimitedStream; -use crate::scan::limit::RowBudget; -use crate::scan::limit::SharedRowLimit; +use crate::scan::limit::RowLimit; use crate::scan::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. @@ -90,6 +88,7 @@ pub struct MultiLayoutDataSource { concurrency: usize, } +#[derive(Clone)] pub enum MultiLayoutChild { Opened { reader: LayoutReaderRef, @@ -289,6 +288,19 @@ impl DataSource for MultiLayoutDataSource { } async fn scan(&self, scan_request: ScanRequest) -> VortexResult { + let dtype = scan_request.projection.return_dtype(&self.dtype)?; + + if scan_request.ordered && scan_request.limit.is_some() { + // Ordered global limits must consume complete files in order. A single composite + // partition owns a local mask-level limit and scans each selected file sequentially. + return Ok(Box::new(OrderedMultiLayoutScan { + session: self.session.clone(), + dtype, + request: scan_request, + children: self.children.iter().cloned().enumerate().collect(), + })); + } + let mut ready = VecDeque::new(); let mut deferred = VecDeque::new(); @@ -301,15 +313,14 @@ impl DataSource for MultiLayoutDataSource { } } - let dtype = scan_request.projection.return_dtype(&self.dtype)?; - - let shared_limit = scan_request.limit.map(SharedRowLimit::new); + // Only unordered scans share a completion-order limit across external partitions. + let row_limit = scan_request.limit.map(RowLimit::new); Ok(Box::new(MultiLayoutScan { session: self.session.clone(), dtype, request: scan_request, - shared_limit, + row_limit, ready, deferred, handle: self.session.handle(), @@ -326,7 +337,7 @@ struct MultiLayoutScan { session: VortexSession, dtype: DType, request: ScanRequest, - shared_limit: Option, + row_limit: Option, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -352,7 +363,7 @@ impl DataSourceScan for MultiLayoutScan { session, dtype: _, request, - shared_limit, + row_limit, ready, deferred, handle, @@ -412,7 +423,7 @@ impl DataSourceScan for MultiLayoutScan { reader, session.clone(), request.clone(), - shared_limit.clone(), + row_limit.clone(), ), Err(e) => stream::once(async move { Err(e) }).boxed(), }) @@ -420,6 +431,126 @@ impl DataSourceScan for MultiLayoutScan { } } +/// An ordered scan with a global limit. It produces one partition so a later file cannot reserve +/// rows before every earlier selected file has been accounted for. +struct OrderedMultiLayoutScan { + session: VortexSession, + dtype: DType, + request: ScanRequest, + children: VecDeque<(usize, MultiLayoutChild)>, +} + +impl DataSourceScan for OrderedMultiLayoutScan { + fn dtype(&self) -> &DType { + &self.dtype + } + + fn partition_count(&self) -> Precision { + Precision::exact(1usize) + } + + fn partitions(self: Box) -> PartitionStream { + let Self { + session, + dtype, + request, + children, + } = *self; + + stream::once(async move { + Ok(Box::new(OrderedMultiLayoutPartition { + session, + dtype, + request, + children, + }) as PartitionRef) + }) + .boxed() + } +} + +/// A composite partition that opens and scans selected readers one at a time. +struct OrderedMultiLayoutPartition { + session: VortexSession, + dtype: DType, + request: ScanRequest, + children: VecDeque<(usize, MultiLayoutChild)>, +} + +impl Partition for OrderedMultiLayoutPartition { + fn as_any(&self) -> &dyn Any { + self + } + + fn index(&self) -> usize { + 0 + } + + fn row_count(&self) -> Precision { + Precision::inexact(self.request.limit.unwrap_or_default()) + } + + fn byte_size(&self) -> Precision { + Precision::Absent + } + + fn execute(self: Box) -> VortexResult { + let Self { + session, + dtype, + request, + children, + } = *self; + let Some(limit) = request.limit else { + vortex_bail!("ordered multi-layout partitions require a row limit"); + }; + let row_limit = RowLimit::new(limit); + + let stream = async_stream::try_stream! { + for (index, child) in children { + if row_limit.is_exhausted() { + break; + } + if !partition_is_selected(index, &request) { + continue; + } + + let reader = match child { + MultiLayoutChild::Opened { reader, .. } => reader, + MultiLayoutChild::Deferred { factory, .. } => { + let Some(reader) = factory + .open() + .instrument(tracing::info_span!("LayoutReaderFactory::open")) + .await? + else { + continue; + }; + reader + } + }; + + let mut partitions = reader_partition( + index, + reader, + session.clone(), + request.clone(), + Some(row_limit.clone()), + ); + while let Some(partition) = partitions.next().await { + let mut chunks = partition?.execute()?; + while let Some(chunk) = chunks.next().await { + yield chunk?; + } + } + } + }; + + Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( + dtype, stream, + ))) + } +} + /// Generates a partition stream for a single layout reader. /// /// Checks file-level pruning first (via `pruning_evaluation`). If the filter proves no rows @@ -430,30 +561,14 @@ fn reader_partition( reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, - shared_limit: Option, + row_limit: Option, ) -> PartitionStream { let row_count = reader.row_count(); let row_range = request.row_range.clone().unwrap_or(0..row_count); - let partition_idx_u64: u64 = partition_idx as u64; - if let Some(range) = &request.partition_range - && !range.contains(&partition_idx_u64) - { + if !partition_is_selected(partition_idx, &request) { return stream::empty().boxed(); - }; - match &request.partition_selection { - Selection::IncludeByIndex(buffer) => { - if buffer.as_slice().binary_search(&partition_idx_u64).is_err() { - return stream::empty().boxed(); - } - } - Selection::ExcludeByIndex(buffer) => { - if buffer.as_slice().binary_search(&partition_idx_u64).is_ok() { - return stream::empty().boxed(); - } - } - _ => {} - }; + } // Check file-level pruning: if the filter can be proven false for the entire row range // using file-level statistics, skip this reader entirely. @@ -476,13 +591,32 @@ fn reader_partition( row_range: Some(row_range), ..request }, - shared_limit, + row_limit, index: partition_idx, }) as PartitionRef) }) .boxed() } +fn partition_is_selected(partition_idx: usize, request: &ScanRequest) -> bool { + let partition_idx = partition_idx as u64; + if let Some(range) = &request.partition_range + && !range.contains(&partition_idx) + { + return false; + } + + match &request.partition_selection { + Selection::IncludeByIndex(buffer) => { + buffer.as_slice().binary_search(&partition_idx).is_ok() + } + Selection::ExcludeByIndex(buffer) => { + buffer.as_slice().binary_search(&partition_idx).is_err() + } + _ => true, + } +} + /// A partition backed by a single [`LayoutReaderRef`] and a row range. /// /// On `execute()`, creates a [`ScanBuilder`] over the row range, enabling @@ -491,7 +625,7 @@ struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, - shared_limit: Option, + row_limit: Option, index: usize, } @@ -515,7 +649,7 @@ impl Partition for MultiLayoutPartition { .limit .map_or(row_count, |limit| row_count.min(limit)); - if self.request.filter.is_some() { + if self.request.filter.is_some() || self.row_limit.is_some() { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -527,13 +661,13 @@ impl Partition for MultiLayoutPartition { } fn execute(self: Box) -> VortexResult { - let shared_limit = self.shared_limit.clone(); let request = self.request; let mut builder = ScanBuilder::new(self.session, self.reader) .with_selection(request.selection) .with_projection(request.projection) .with_some_filter(request.filter) .with_some_limit(request.limit) + .with_some_row_limit(self.row_limit) .with_ordered(request.ordered); if let Some(row_range) = request.row_range { @@ -542,12 +676,6 @@ impl Partition for MultiLayoutPartition { let dtype = builder.dtype()?; let stream = builder.into_stream()?.boxed(); - // Caps total rows across all partitions; only correct for unordered consumption - // (see `SharedRowLimit`). - let stream = match shared_limit { - Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(), - None => stream, - }; Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new( dtype, stream, @@ -557,11 +685,36 @@ impl Partition for MultiLayoutPartition { #[cfg(test)] mod tests { + use std::ops::Range; + use std::sync::Arc; + + use async_trait::async_trait; + use futures::TryStreamExt; use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_io::runtime::BlockingRuntime; + use vortex_io::runtime::single::SingleThreadRuntime; + use vortex_mask::Mask; + use vortex_scan::DataSource; + use vortex_scan::ScanRequest; use super::*; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; use crate::scan::test::new_session; + use crate::scan::test::session_with_handle; struct NeverOpened; @@ -593,4 +746,133 @@ mod tests { fn byte_size_precision(#[case] sizes: Vec>, #[case] expected: Precision) { assert_eq!(deferred_source(sizes).byte_size(), expected); } + + #[derive(Debug)] + struct TestLayoutReader { + name: Arc, + dtype: DType, + base: i32, + row_count: u64, + } + + impl TestLayoutReader { + fn new(name: &'static str, base: i32, row_count: u64) -> Self { + Self { + name: Arc::from(name), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + base, + row_count, + } + } + } + + impl LayoutReader for TestLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let base = self.base; + + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start)?; + let end = i32::try_from(row_range.end)?; + PrimitiveArray::from_iter((start..end).map(|value| base + value)) + .into_array() + .filter(mask.await?) + })) + } + } + + struct StaticReaderFactory { + reader: LayoutReaderRef, + } + + #[async_trait] + impl LayoutReaderFactory for StaticReaderFactory { + async fn open(&self) -> VortexResult> { + Ok(Some(Arc::clone(&self.reader))) + } + } + + #[test] + fn ordered_limit_scans_multiple_readers_sequentially() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let first: LayoutReaderRef = Arc::new(TestLayoutReader::new("first", 0, 2)); + let second: LayoutReaderRef = Arc::new(TestLayoutReader::new("second", 10, 2)); + let source = MultiLayoutDataSource::new_with_first( + first, + vec![Arc::new(StaticReaderFactory { reader: second })], + vec![], + &session, + ); + + let scan = runtime.block_on(source.scan(ScanRequest { + filter: Some(root()), + limit: Some(3), + ordered: true, + ..Default::default() + }))?; + let partitions = runtime.block_on(scan.partitions().try_collect::>())?; + assert_eq!(partitions.len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for partition in partitions { + for chunk in runtime.block_on_stream(partition.execute()?) { + let primitive = chunk?.execute::(&mut ctx)?; + values.extend(primitive.into_buffer::()); + } + } + + assert_eq!(values, [0, 1, 10]); + Ok(()) + } } diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 073724b52a7..0f9cc6655bd 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -30,10 +30,13 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; -use crate::scan::limit::LimitedStream; -use crate::scan::limit::RowBudget; +use crate::scan::limit::RowLimit; use crate::scan::splits::Splits; -use crate::scan::tasks::{split_exec, TaskContext, TaskFuture}; +use crate::scan::tasks::TaskContext; +use crate::scan::tasks::TaskFuture; +use crate::scan::tasks::filter_split; +use crate::scan::tasks::project_split; +use crate::scan::tasks::split_exec; /// A projected subset (by indices, range, and filter) of rows from a Vortex data source. /// @@ -55,6 +58,8 @@ pub struct RepeatedScan { concurrency: usize, /// Maximal number of rows to read (after filtering). limit: Option, + /// An optional row limit shared with sibling external partitions. + row_limit: Option, /// The dtype of the projected arrays. dtype: DType, } @@ -89,7 +94,7 @@ impl RepeatedScan { clippy::too_many_arguments, reason = "all arguments are needed for scan construction" )] - pub fn new( + pub(crate) fn new( session: VortexSession, layout_reader: LayoutReaderRef, projection: Expression, @@ -100,6 +105,7 @@ impl RepeatedScan { splits: Splits, concurrency: usize, limit: Option, + row_limit: Option, dtype: DType, ) -> Self { Self { @@ -113,6 +119,7 @@ impl RepeatedScan { splits, concurrency, limit, + row_limit, dtype, } } @@ -181,17 +188,20 @@ impl RepeatedScan { }) } - pub(crate) fn execute(&self, row_range: Option>) -> VortexResult> { + pub(crate) fn execute( + &self, + row_range: Option>, + row_limit: Option, + ) -> VortexResult> { let ctx = self.task_context(); - let mut limit = self.limit.filter(|_| self.filter.is_none()); let mut tasks = Vec::new(); for range in self.split_ranges(row_range) { if range.start >= range.end { continue; } - if limit.is_some_and(|l| l == 0) { + if row_limit.as_ref().is_some_and(RowLimit::is_exhausted) { break; } @@ -200,7 +210,7 @@ impl RepeatedScan { continue; } - tasks.push(split_exec(Arc::clone(&ctx), row_mask, limit.as_mut())?); + tasks.push(split_exec(Arc::clone(&ctx), row_mask, row_limit.clone())?); } Ok(tasks) @@ -211,48 +221,123 @@ impl RepeatedScan { row_range: Option>, ) -> VortexResult> + Send + 'static> { let num_workers = get_available_parallelism().unwrap_or(1); + let row_limit = self + .row_limit + .clone() + .or_else(|| self.limit.map(RowLimit::new)); let concurrency = self.concurrency * num_workers; let handle = self.session.handle(); // With both a filter and a limit we cannot know each split's output row count ahead of - // time, so split tasks are built lazily as the stream is polled. `buffered`'s read-ahead - // (bounded by `concurrency`) registers IO for splits eagerly, but only as far as the - // limit requires: `limit_array_stream` drops the inner stream once the limit is reached, - // capping over-read at `concurrency` splits. - if self.filter.is_some() && self.limit.is_some() { + // time, so split tasks are built lazily as the stream is polled. Once another task drains + // the shared budget, `take_while` prevents further task creation; already-prefetched + // tasks observe the exhausted budget after filtering and return `None` without projection. + if self.filter.is_some() + && let Some(row_limit) = row_limit.clone() + { + // An ordered LIMIT would be violated if a later split reserved rows before an earlier + // one, so ordered scans use a two-stage pipeline that keeps I/O concurrent while + // reserving strictly in split order. + if self.ordered { + return Ok(self.ordered_filtered_limit_stream(row_range, row_limit, concurrency)); + } + let ctx = self.task_context(); let selection = self.selection.clone(); - let tasks = - futures::stream::iter(self.split_ranges(row_range)).filter_map(move |range| { + let task_limit = row_limit.clone(); + let tasks = futures::stream::iter(self.split_ranges(row_range)) + .take_while(move |_| futures::future::ready(!task_limit.is_exhausted())) + .filter_map(move |range| { // Build the row mask and split task synchronously so the IO system sees the - // split's ranges as soon as `buffered` pulls it, without cloning `selection`. + // split's ranges as soon as `buffer_unordered` pulls it, without cloning + // `selection`. let row_mask = selection.row_mask(&range); let spawned = (!row_mask.mask().all_false()).then(|| { - let task = split_exec(Arc::clone(&ctx), row_mask, None) + let task = split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) .unwrap_or_else(|err| async move { Err(err) }.boxed()); handle.spawn(task) }); async move { spawned } }); - return Ok(schedule(tasks, self.ordered, concurrency, self.limit)); + return Ok(schedule(tasks, false, concurrency)); } // No filter (or no limit): build every task eagerly so the IO system sees all split - // ranges up front. A no-filter limit is applied exactly per split inside `execute`. - let tasks = - futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task)); + // ranges up front. A no-filter limit is applied to each selection mask inside `execute`. + let tasks = futures::stream::iter(self.execute(row_range, row_limit)?) + .map(move |task| handle.spawn(task)); - Ok(schedule(tasks, self.ordered, concurrency, self.limit)) + Ok(schedule(tasks, self.ordered, concurrency)) + } + + /// Ordered filtered scan with a shared row limit. + /// + /// I/O concurrency is decoupled from limit reservation so that an ordered `LIMIT` does not + /// force serial execution. Filters for a window of splits are evaluated concurrently (stage + /// one), rows are reserved against the shared limit strictly in split order on the seam + /// between the stages, and only the reserved masks are projected (stage two, also concurrent). + /// Because reservation happens in split order, the earliest matching rows always win the + /// budget; because projection only runs for reserved masks, no rows past the limit are + /// decoded. `take_while` stops feeding stage one once the limit is exhausted, bounding + /// speculative filter I/O to the pipeline depth. + fn ordered_filtered_limit_stream( + &self, + row_range: Option>, + row_limit: RowLimit, + concurrency: usize, + ) -> BoxStream<'static, VortexResult> { + let ctx = self.task_context(); + let handle = self.session.handle(); + + // Stage one: evaluate the filter for each split concurrently while preserving split order. + let filter_tasks = { + let ctx = Arc::clone(&ctx); + let handle = handle.clone(); + let selection = self.selection.clone(); + let take_limit = row_limit.clone(); + futures::stream::iter(self.split_ranges(row_range)) + .take_while(move |_| futures::future::ready(!take_limit.is_exhausted())) + .filter_map(move |range| { + // Build the row mask synchronously so the IO system sees the split's ranges as + // soon as `buffered` pulls it, without cloning `selection`. + let row_mask = selection.row_mask(&range); + let spawned = (!row_mask.mask().all_false()) + .then(|| handle.spawn(filter_split(Arc::clone(&ctx), row_mask))); + async move { spawned } + }) + .buffered(concurrency) + }; + + // Seam + stage two: reserve in split order (this map runs sequentially as the ordered + // stage-one stream is consumed), then spawn projection work for the reserved mask. + filter_tasks + .map(move |result| { + let task: TaskFuture = match result { + Ok((row_range, mask)) => { + let mask = row_limit.limit(mask); + if mask.all_false() { + async { Ok(None) }.boxed() + } else { + project_split(Arc::clone(&ctx), row_range, mask) + .unwrap_or_else(|err| async move { Err(err) }.boxed()) + } + } + Err(err) => async move { Err(err) }.boxed(), + }; + handle.spawn(task) + }) + .buffered(concurrency) + .filter_map(|chunk| async move { chunk.transpose() }) + .boxed() } } -/// Spawn-buffer a stream of split tasks, transposing empty splits away and applying `limit`. +/// Spawn-buffer a stream of split tasks and transpose empty splits away. fn schedule( tasks: S, ordered: bool, concurrency: usize, - limit: Option, ) -> BoxStream<'static, VortexResult> where S: Stream>>> + Send + 'static, @@ -262,14 +347,9 @@ where } else { tasks.buffer_unordered(concurrency).boxed() }; - let stream = stream + stream .filter_map(|chunk| async move { chunk.transpose() }) - .boxed(); - - match limit { - Some(limit) => LimitedStream::new(stream, RowBudget::Local(limit)).boxed(), - None => stream, - } + .boxed() } fn intersect_ranges(left: Option<&Range>, right: Option>) -> Option> { diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index a6d98ffce85..d92d42d54a4 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -35,6 +35,7 @@ use vortex_session::VortexSession; use crate::LayoutReader; use crate::LayoutReaderRef; use crate::layouts::row_idx::RowIdxLayoutReader; +use crate::scan::limit::RowLimit; use crate::scan::repeated_scan::RepeatedScan; use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; @@ -70,6 +71,8 @@ pub struct ScanBuilder { metrics_registry: Option>, /// Maximal number of rows to read after filtering. limit: Option, + /// A row limit shared with sibling external partitions, when the caller owns one. + row_limit: Option, /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression, /// but not by the scan [`Selection`] which remains relative. row_offset: u64, @@ -92,6 +95,7 @@ impl ScanBuilder { concurrency: 4, metrics_registry: None, limit: None, + row_limit: None, row_offset: 0, } } @@ -216,6 +220,12 @@ impl ScanBuilder { self } + /// Use a row limit supplied by the enclosing data source instead of creating a local one. + pub(crate) fn with_some_row_limit(mut self, row_limit: Option) -> Self { + self.row_limit = row_limit; + self + } + /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { self.projection.return_dtype(self.layout_reader.dtype()) @@ -281,6 +291,7 @@ impl ScanBuilder { splits, self.concurrency, self.limit, + self.row_limit, dtype, )) } @@ -873,6 +884,197 @@ mod test { } } + #[derive(Debug)] + struct ProjectionMaskLayoutReader { + name: Arc, + dtype: DType, + row_count: u64, + projection_masks: Arc>>, + projection_error: bool, + } + + impl ProjectionMaskLayoutReader { + fn new(row_count: u64, projection_masks: Arc>>) -> Self { + Self { + name: Arc::from("projection-mask"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + row_count, + projection_masks, + projection_error: false, + } + } + + fn with_projection_error(mut self) -> Self { + self.projection_error = true; + self + } + } + + impl LayoutReader for ProjectionMaskLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let projection_masks = Arc::clone(&self.projection_masks); + let projection_error = self.projection_error; + + Ok(Box::pin(async move { + let mask = mask.await?; + projection_masks.lock().push(mask.true_count()); + if projection_error { + return Err(vortex_err!("projection failed")); + } + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + #[derive(Debug)] + struct ErrorThenMatchLayoutReader { + name: Arc, + dtype: DType, + projection_masks: Arc>>, + } + + impl ErrorThenMatchLayoutReader { + fn new(projection_masks: Arc>>) -> Self { + Self { + name: Arc::from("error-then-match"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + projection_masks, + } + } + } + + impl LayoutReader for ErrorThenMatchLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 2 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + splits.push(split_range.row_offset() + row_range.start + 1); + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + if row_range.start == 0 { + let len = mask.len(); + return Ok(MaskFuture::new(len, async move { + Err(vortex_err!("first split filter failed")) + })); + } + Ok(mask) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + let projection_masks = Arc::clone(&self.projection_masks); + + Ok(Box::pin(async move { + let mask = mask.await?; + projection_masks.lock().push(mask.true_count()); + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + fn collect_scan_values(iter: I) -> VortexResult> where I: IntoIterator>, @@ -918,6 +1120,77 @@ mod test { Ok(()) } + #[test] + fn filtered_limit_limits_projection_mask_before_projection() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new(ProjectionMaskLayoutReader::new( + 100_000, + Arc::clone(&projection_masks), + )); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + + assert_eq!(values, [0]); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + + #[test] + fn filter_errors_are_stream_items_and_do_not_consume_the_limit() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new(ErrorThenMatchLayoutReader::new(Arc::clone( + &projection_masks, + ))); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + let Some(chunk) = iter.next() else { + return Err(vortex_err!( + "matching split was not polled after the filter error" + )); + }; + let mut ctx = array_session().create_execution_ctx(); + let primitive = chunk?.execute::(&mut ctx)?; + + assert_eq!(primitive.into_buffer::().as_slice(), [1]); + assert!(iter.next().is_none()); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + + #[test] + fn projection_errors_are_stream_items() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new( + ProjectionMaskLayoutReader::new(1, Arc::clone(&projection_masks)) + .with_projection_error(), + ); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + assert!(iter.next().is_none()); + assert_eq!(projection_masks.lock().as_slice(), [1]); + Ok(()) + } + #[test] fn prepared_scan_limits_filtered_results() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index ad3c4d3e2da..7a17869dcdf 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -4,6 +4,7 @@ //! Split scanning task implementation. use std::ops::BitAnd; +use std::ops::Range; use std::sync::Arc; use bit_vec::BitVec; @@ -12,12 +13,14 @@ use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::expr::Expression; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; use crate::LayoutReader; use crate::scan::filter::FilterExpr; +use crate::scan::limit::RowLimit; pub type TaskFuture = BoxFuture<'static, VortexResult>>; @@ -33,121 +36,190 @@ pub type TaskFuture = BoxFuture<'static, VortexResult>>; /// The intersected row range is then further reduced via expression-based pruning. After pruning /// has eliminated more blocks, the full filter is executed over the remainder of the split. /// -/// This mask is then provided to the reader to perform a filtered projection over the split data, -/// yielding the projected array (or `None` when the split selects no rows). +/// The final mask is limited before it is given to the reader to perform a filtered projection +/// over the split data, yielding the projected array (or `None` when the split selects no rows). +/// Limiting before projection prevents decode work for rows that the scan cannot return. pub fn split_exec( ctx: Arc, read_mask: RowMask, - limit: Option<&mut u64>, + row_limit: Option, ) -> VortexResult { let row_range = read_mask.row_range(); let row_mask = read_mask.mask().clone(); - let filter_mask = match ctx.filter.as_ref() { - // No filter == immediate mask - None => { - let row_mask = match limit { - Some(l) if *l == 0 => Mask::new_false(row_mask.len()), - Some(l) => { - let true_count = row_mask.true_count(); - let mask_limit = usize::try_from(*l) - .map(|l| l.min(true_count)) - .unwrap_or(true_count); - let row_mask = row_mask.limit(mask_limit); - *l -= mask_limit as u64; - row_mask - } - None => row_mask, - }; - - MaskFuture::ready(row_mask) - } - Some(filter) => { - // NOTE: it's very important that the pruning and filter evaluations are built OUTSIDE - // the future. Registering these row ranges eagerly is a hint to the IO system that - // we want to start prefetching the IO for this split. - let reader = Arc::clone(&ctx.reader); - let filter = Arc::clone(filter); - let row_range = row_range.clone(); - - MaskFuture::new(row_mask.len(), async move { - let mut mask = row_mask; - let mut dynamic_versions = vec![None; filter.conjuncts().len()]; - - // TODO(ngates): we could use FuturedUnordered to intersect the masks in parallel. - for (idx, conjunct) in filter.conjuncts().iter().enumerate() { - if mask.all_false() { - return Ok(mask); - } - - // Store the latest version of the dynamic expression prior to pruning. - // We will re-run the pruning later if the version has changed in the meantime. - dynamic_versions[idx] = filter.dynamic_updates(idx).map(|du| du.version()); - - let conjunct_mask = reader - .pruning_evaluation(&row_range, conjunct, mask.clone())? - .await?; - mask = mask.bitand(&conjunct_mask); - } - - // Now we loop through the conjuncts in the preferred order and evaluate them. - let mut remaining = BitVec::from_elem(filter.conjuncts().len(), true); - while let Some(idx) = filter.next_conjunct(&remaining) { - remaining.set(idx, false); - if mask.all_false() { - return Ok(mask); - } - - let conjunct = &filter.conjuncts()[idx]; - - // If the dynamic expression has changed since pruning, re-run the pruning. - // Store the dynamic update once to avoid TOCTOU race condition - let current_version = filter.dynamic_updates(idx).map(|du| du.version()); - if let Some(dv) = current_version - && dynamic_versions[idx].is_none_or(|v| v < dv) - { - // The dynamic expression has been updated, re-run the pruning. - dynamic_versions[idx] = Some(dv); - let conjunct_mask = reader - .pruning_evaluation(&row_range, conjunct, mask.clone())? - .await?; - mask = mask.bitand(&conjunct_mask); - } - if mask.all_false() { - return Ok(mask); - } - - let conjunct_mask = reader - .filter_evaluation(&row_range, conjunct, MaskFuture::ready(mask))? - .await?; - filter.report_selectivity(idx, conjunct_mask.density()); - - // Filter evaluations return a mask already intersected with the input mask. - mask = conjunct_mask; - } - - Ok(mask) - }) + let Some(filter) = ctx.filter.as_ref() else { + let row_mask = if let Some(limit) = row_limit { + limit.limit(row_mask) + } else { + row_mask + }; + if row_mask.all_false() { + return Ok(async { Ok(None) }.boxed()); } + + // With no filter, limit the selection before constructing projection work. + let projection = ctx.reader.projection_evaluation( + &row_range, + &ctx.projection, + MaskFuture::ready(row_mask), + )?; + return Ok(async move { projection.await.map(Some) }.boxed()); }; - // Step 4: execute the projection, only at the mask for rows which match the filter - let projection_future = - ctx.reader - .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; + let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); + + let Some(row_limit) = row_limit else { + // Without a limit, retain the existing eager projection setup so readers can prefetch + // projection work while the filter is being evaluated. + let projection = + ctx.reader + .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; + let array_fut = async move { + let mask = filter_mask.await?; + if mask.all_false() { + return Ok(None); + } + + projection.await.map(Some) + }; + return Ok(array_fut.boxed()); + }; let array_fut = async move { let mask = filter_mask.await?; + // A filter error above returns before reserving any rows. Once filtering has succeeded, + // reserve only the matching rows and construct projection work for that limited mask. + let mask = row_limit.limit(mask); if mask.all_false() { return Ok(None); } - projection_future.await.map(Some) + let projection = ctx.reader.projection_evaluation( + &row_range, + &ctx.projection, + MaskFuture::ready(mask), + )?; + projection.await.map(Some) }; Ok(array_fut.boxed()) } +/// Evaluate the filter for a split and return the matching mask together with its row range. +/// +/// This is the first stage of the ordered filtered-limit pipeline. It performs the same pruning +/// and filter I/O as [`split_exec`] (registered eagerly, outside the returned future, so the IO +/// system can prefetch while earlier splits are still reserving), but it neither reserves against +/// the limit nor projects. The caller reserves the returned mask in split order and then projects +/// it via [`project_split`], which keeps ordered `LIMIT` semantics without serializing I/O. +pub fn filter_split( + ctx: Arc, + read_mask: RowMask, +) -> BoxFuture<'static, VortexResult<(Range, Mask)>> { + let row_range = read_mask.row_range(); + let row_mask = read_mask.mask().clone(); + let filter = ctx + .filter + .as_ref() + .vortex_expect("filter_split requires a filtered scan"); + let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); + + async move { + let mask = filter_mask.await?; + Ok((row_range, mask)) + } + .boxed() +} + +/// Project an already-reserved, non-empty mask for a split. +/// +/// This is the second stage of the ordered filtered-limit pipeline, run after the caller has +/// reserved rows against the limit in split order (see [`filter_split`]). +pub fn project_split( + ctx: Arc, + row_range: Range, + mask: Mask, +) -> VortexResult { + let projection = + ctx.reader + .projection_evaluation(&row_range, &ctx.projection, MaskFuture::ready(mask))?; + Ok(async move { projection.await.map(Some) }.boxed()) +} + +/// Build the filtered mask for a split. +/// +/// The pruning and filter evaluations are constructed OUTSIDE the returned future on purpose: +/// registering these row ranges eagerly is a hint to the IO system that we want to start +/// prefetching the IO for this split. +fn build_filter_mask( + reader: &Arc, + filter: &Arc, + row_range: &Range, + row_mask: Mask, +) -> MaskFuture { + let reader = Arc::clone(reader); + let filter = Arc::clone(filter); + let filter_row_range = row_range.clone(); + MaskFuture::new(row_mask.len(), async move { + let mut mask = row_mask; + let mut dynamic_versions = vec![None; filter.conjuncts().len()]; + + // TODO(ngates): we could use FuturedUnordered to intersect the masks in parallel. + for (idx, conjunct) in filter.conjuncts().iter().enumerate() { + if mask.all_false() { + return Ok(mask); + } + + // Store the latest version of the dynamic expression prior to pruning. + // We will re-run the pruning later if the version has changed in the meantime. + dynamic_versions[idx] = filter.dynamic_updates(idx).map(|du| du.version()); + + let conjunct_mask = reader + .pruning_evaluation(&filter_row_range, conjunct, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); + } + + // Now we loop through the conjuncts in the preferred order and evaluate them. + let mut remaining = BitVec::from_elem(filter.conjuncts().len(), true); + while let Some(idx) = filter.next_conjunct(&remaining) { + remaining.set(idx, false); + if mask.all_false() { + return Ok(mask); + } + + let conjunct = &filter.conjuncts()[idx]; + + // If the dynamic expression has changed since pruning, re-run the pruning. + // Store the dynamic update once to avoid TOCTOU race condition. + let current_version = filter.dynamic_updates(idx).map(|du| du.version()); + if let Some(dv) = current_version + && dynamic_versions[idx].is_none_or(|v| v < dv) + { + // The dynamic expression has changed, re-run the pruning. + dynamic_versions[idx] = Some(dv); + let conjunct_mask = reader + .pruning_evaluation(&filter_row_range, conjunct, mask.clone())? + .await?; + mask = mask.bitand(&conjunct_mask); + } + if mask.all_false() { + return Ok(mask); + } + + let conjunct_mask = reader + .filter_evaluation(&filter_row_range, conjunct, MaskFuture::ready(mask))? + .await?; + filter.report_selectivity(idx, conjunct_mask.density()); + + // Filter evaluations return a mask already intersected with the input mask. + mask = conjunct_mask; + } + + Ok(mask) + }) +} + /// Information needed to execute a single split task. /// /// Row selection is evaluated before creating a split task so it's not included diff --git a/vortex-scan/src/lib.rs b/vortex-scan/src/lib.rs index 2f2e6e1c429..58877399422 100644 --- a/vortex-scan/src/lib.rs +++ b/vortex-scan/src/lib.rs @@ -139,10 +139,13 @@ pub struct ScanRequest { /// Partition range to scan, which allows readers to skip unwanted partitions. pub partition_range: Option>, /// Whether the scan should preserve row order. If false, the scan may produce rows in any - /// order, for example to enable parallel execution across partitions. + /// order, for example to enable parallel execution across partitions. An ordered global + /// limit intentionally serializes external partitions so the returned rows are the first + /// matching rows in scan order. pub ordered: bool, - /// Optional limit on the number of rows returned by scan. Limits are applied after all - /// filtering and row selection. + /// Optional global limit on the number of rows returned by the scan. Limits are applied after + /// filtering and row selection, but implementations must trim the resulting selection mask + /// before projection so they do not decode rows that cannot be returned. pub limit: Option, } From 54631a431d0ed20c176e0ebfb35236aad5fd2fbe Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 11:29:06 +0100 Subject: [PATCH 11/17] fix bug Signed-off-by: Adam Gutglick --- vortex-datafusion/src/v2/source.rs | 160 +++++++++++++++--------- vortex-layout/src/scan/repeated_scan.rs | 46 ++++--- vortex-layout/src/scan/tasks.rs | 91 +++++++++++--- 3 files changed, 203 insertions(+), 94 deletions(-) diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index efcec06dc8e..155ea463af9 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -105,7 +105,6 @@ use vortex::array::stream::SendableArrayStream; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::Nullability; -use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::expr::Expression; @@ -678,32 +677,38 @@ where } // Ordered output must preserve global row order, but later partitions should still warm up - // (start their I/O) while an earlier one is draining. Each partition is drained by a spawned - // task into a bounded channel: spawning starts its I/O immediately, and the bounded channel - // back-pressures so read-ahead stays capped. `buffered` warms up to `concurrency` partitions - // ahead while `try_flatten` emits their chunks strictly in partition order. At most - // `concurrency * CHUNKS_AHEAD` chunks are buffered across the warming partitions. - const CHUNKS_AHEAD: usize = 2; + // (start their I/O) while an earlier one is draining. `buffered` invokes this map while + // filling its window, so creating the channel and spawning the drain here starts each window + // member before `try_flatten` begins draining the first receiver. Each drain has a bounded + // channel for back-pressure. It can hold `CHUNK_BUFFER_CAPACITY` chunks plus one pending + // send, and at most `lookahead` drains are live at once. + const CHUNK_BUFFER_CAPACITY: usize = 2; let lookahead = concurrency.max(1); scan_streams .map(move |stream_result| { let handle = handle.clone(); - async move { - let mut stream = stream_result?; - let (tx, rx) = tokio::sync::mpsc::channel(CHUNKS_AHEAD); + let receiver = stream_result.map(|mut stream| { + let (tx, rx) = tokio::sync::mpsc::channel(CHUNK_BUFFER_CAPACITY); handle .spawn(async move { - while let Some(item) = stream.next().await { - // A send error means the consumer was dropped (for example once a - // global limit is reached), so stop draining and release the I/O. - if tx.send(item).await.is_err() { - break; + loop { + tokio::select! { + _ = tx.closed() => break, + item = stream.next() => { + let Some(item) = item else { + break; + }; + if tx.send(item).await.is_err() { + break; + } + } } } }) .detach(); - Ok::<_, VortexError>(ReceiverStream::new(rx)) - } + ReceiverStream::new(rx) + }); + futures::future::ready(receiver) }) .buffered(lookahead) .try_flatten() @@ -725,12 +730,14 @@ fn estimate_to_df_precision(est: &Precision) -> DFPrecision { #[cfg(test)] mod tests { use std::sync::Arc; - use std::sync::atomic::AtomicBool; - use std::sync::atomic::AtomicU8; - use std::sync::atomic::Ordering; + use std::time::Duration; + use std::pin::Pin; + use std::task::Context; use std::task::Poll; + use futures::Stream; use futures::TryStreamExt; + use tokio::sync::Notify; use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; @@ -743,6 +750,7 @@ mod tests { use vortex::dtype::PType; use vortex::error::VortexError; use vortex::error::VortexResult; + use vortex::error::vortex_err; use vortex::io::runtime::tokio::TokioRuntime; use super::flatten_scan_streams; @@ -765,6 +773,30 @@ mod tests { Ok(values) } + struct DropNotifier(Option>); + + impl Drop for DropNotifier { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + + struct DropNotifyingPendingStream { + started: Arc, + _drop_notifier: DropNotifier, + } + + impl Stream for DropNotifyingPendingStream { + type Item = VortexResult; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.started.notify_one(); + Poll::Pending + } + } + #[tokio::test] async fn ordered_flatten_preserves_partition_order() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); @@ -784,49 +816,28 @@ mod tests { async fn ordered_flatten_warms_later_partition_while_first_pending() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - // The first partition self-wakes `Pending` a few times, then yields one chunk and ends. - let first_done = Arc::new(AtomicBool::new(false)); - let first_done_for_stream = Arc::clone(&first_done); - let mut pending_polls = 3usize; - let mut emitted = false; + let first_started = Arc::new(Notify::new()); + let first_started_for_stream = Arc::clone(&first_started); + let release_first = Arc::new(Notify::new()); + let release_first_for_stream = Arc::clone(&release_first); let first = ArrayStreamAdapter::new( dtype.clone(), - futures::stream::poll_fn(move |cx| -> Poll>> { - if pending_polls > 0 { - pending_polls -= 1; - cx.waker().wake_by_ref(); - return Poll::Pending; - } - if !emitted { - emitted = true; - return Poll::Ready(Some(Ok(PrimitiveArray::from_iter([0i32]).into_array()))); - } - first_done_for_stream.store(true, Ordering::SeqCst); - Poll::Ready(None) + futures::stream::once(async move { + first_started_for_stream.notify_one(); + release_first_for_stream.notified().await; + Ok(PrimitiveArray::from_iter([0i32]).into_array()) }), ) .boxed(); - // The second partition records, at its first poll, whether the first had already finished. - // 0 = not yet polled, 1 = first still running, 2 = first already done. - let observation = Arc::new(AtomicU8::new(0)); - let observation_for_stream = Arc::clone(&observation); - let first_done_for_second = Arc::clone(&first_done); - let mut first_poll = true; + let first_started_for_second = Arc::clone(&first_started); + let release_first_for_second = Arc::clone(&release_first); let second = ArrayStreamAdapter::new( dtype, - futures::stream::poll_fn(move |_| -> Poll>> { - if first_poll { - first_poll = false; - let seen = if first_done_for_second.load(Ordering::SeqCst) { - 2 - } else { - 1 - }; - observation_for_stream.store(seen, Ordering::SeqCst); - return Poll::Ready(Some(Ok(PrimitiveArray::from_iter([100i32]).into_array()))); - } - Poll::Ready(None) + futures::stream::once(async move { + first_started_for_second.notified().await; + release_first_for_second.notify_one(); + Ok(PrimitiveArray::from_iter([100i32]).into_array()) }), ) .boxed(); @@ -834,13 +845,42 @@ mod tests { let streams = futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); - let chunks = flattened.try_collect::>().await?; + let chunks = tokio::time::timeout(Duration::from_secs(1), flattened.try_collect::>()) + .await + .map_err(|_| vortex_err!("ordered flatten did not warm the later partition"))??; - // Global order is preserved: the first partition's row precedes the later partition's. assert_eq!(collect_i32(chunks)?, [0, 100]); - // The later partition was polled (warmed) while the first was still pending, proving the - // ordered flatten keeps cross-partition look-ahead instead of serializing. - assert_eq!(observation.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[tokio::test] + async fn ordered_flatten_cancels_a_pending_drain_when_dropped() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let started = Arc::new(Notify::new()); + let (dropped_send, dropped_recv) = tokio::sync::oneshot::channel(); + let pending = ArrayStreamAdapter::new( + dtype, + DropNotifyingPendingStream { + started: Arc::clone(&started), + _drop_notifier: DropNotifier(Some(dropped_send)), + }, + ) + .boxed(); + + let streams = futures::stream::iter([Ok::<_, VortexError>(pending)]); + let mut flattened = flatten_scan_streams(streams, true, 1, TokioRuntime::current()); + let mut next = Box::pin(futures::StreamExt::next(&mut flattened)); + tokio::select! { + _ = started.notified() => {} + _ = &mut next => return Err(vortex_err!("pending drain unexpectedly produced a chunk")), + } + drop(next); + drop(flattened); + + let dropped = tokio::time::timeout(Duration::from_secs(1), dropped_recv) + .await + .map_err(|_| vortex_err!("dropping the flattened stream did not cancel its drain"))?; + dropped.map_err(|_| vortex_err!("pending source dropped without notifying the test"))?; Ok(()) } } diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 0f9cc6655bd..9f4e5b7875d 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -34,6 +34,7 @@ use crate::scan::limit::RowLimit; use crate::scan::splits::Splits; use crate::scan::tasks::TaskContext; use crate::scan::tasks::TaskFuture; +use crate::scan::tasks::TaskResult; use crate::scan::tasks::filter_split; use crate::scan::tasks::project_split; use crate::scan::tasks::split_exec; @@ -254,7 +255,7 @@ impl RepeatedScan { let row_mask = selection.row_mask(&range); let spawned = (!row_mask.mask().all_false()).then(|| { let task = split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) - .unwrap_or_else(|err| async move { Err(err) }.boxed()); + .unwrap_or_else(|err| async move { TaskResult::Terminal(err) }.boxed()); handle.spawn(task) }); async move { spawned } @@ -310,46 +311,59 @@ impl RepeatedScan { }; // Seam + stage two: reserve in split order (this map runs sequentially as the ordered - // stage-one stream is consumed), then spawn projection work for the reserved mask. - filter_tasks + // stage-one stream is consumed), then spawn projection work for the reserved mask. The + // second gate drops results that were already filtering when an earlier split exhausted + // the budget, including any later filter errors. + let stage_two_limit = row_limit.clone(); + let tasks = filter_tasks + .take_while(move |_| futures::future::ready(!stage_two_limit.is_exhausted())) .map(move |result| { let task: TaskFuture = match result { Ok((row_range, mask)) => { let mask = row_limit.limit(mask); if mask.all_false() { - async { Ok(None) }.boxed() + async { TaskResult::Array(None) }.boxed() } else { project_split(Arc::clone(&ctx), row_range, mask) - .unwrap_or_else(|err| async move { Err(err) }.boxed()) } } - Err(err) => async move { Err(err) }.boxed(), + Err(err) => async move { TaskResult::Recoverable(err) }.boxed(), }; handle.spawn(task) - }) - .buffered(concurrency) - .filter_map(|chunk| async move { chunk.transpose() }) - .boxed() + }); + schedule(tasks, true, concurrency) } } -/// Spawn-buffer a stream of split tasks and transpose empty splits away. +/// Spawn-buffer a stream of split tasks, preserving recoverable filter errors and stopping after +/// a projection error that occurred after a row limit reservation. fn schedule( tasks: S, ordered: bool, concurrency: usize, ) -> BoxStream<'static, VortexResult> where - S: Stream>>> + Send + 'static, + S: Stream> + Send + 'static, { - let stream = if ordered { + let mut stream = if ordered { tasks.buffered(concurrency).boxed() } else { tasks.buffer_unordered(concurrency).boxed() }; - stream - .filter_map(|chunk| async move { chunk.transpose() }) - .boxed() + async_stream::stream! { + while let Some(result) = stream.next().await { + match result { + TaskResult::Array(Some(array)) => yield Ok(array), + TaskResult::Array(None) => {} + TaskResult::Recoverable(err) => yield Err(err), + TaskResult::Terminal(err) => { + yield Err(err); + return; + } + } + } + } + .boxed() } fn intersect_ranges(left: Option<&Range>, right: Option>) -> Option> { diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 7a17869dcdf..a830930fb35 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -14,6 +14,7 @@ use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::expr::Expression; use vortex_error::VortexExpect; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; @@ -22,7 +23,21 @@ use crate::LayoutReader; use crate::scan::filter::FilterExpr; use crate::scan::limit::RowLimit; -pub type TaskFuture = BoxFuture<'static, VortexResult>>; +/// The result of a split task. +/// +/// Filter errors happen before a row limit reserves any rows, so callers may report them and +/// continue with later splits. Projection errors after reservation cannot safely release rows back +/// to a concurrent limit, so callers must report them and terminate the limited scan. +pub(crate) enum TaskResult { + /// A completed projection, or an empty split. + Array(Option), + /// An error that occurred before a row limit reserved rows. + Recoverable(VortexError), + /// An error that occurred after a row limit reserved rows. + Terminal(VortexError), +} + +pub(crate) type TaskFuture = BoxFuture<'static, TaskResult>; /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this @@ -48,22 +63,36 @@ pub fn split_exec( let row_mask = read_mask.mask().clone(); let Some(filter) = ctx.filter.as_ref() else { + let limited = row_limit.is_some(); let row_mask = if let Some(limit) = row_limit { limit.limit(row_mask) } else { row_mask }; if row_mask.all_false() { - return Ok(async { Ok(None) }.boxed()); + return Ok(async { TaskResult::Array(None) }.boxed()); } // With no filter, limit the selection before constructing projection work. - let projection = ctx.reader.projection_evaluation( + let projection = match ctx.reader.projection_evaluation( &row_range, &ctx.projection, MaskFuture::ready(row_mask), - )?; - return Ok(async move { projection.await.map(Some) }.boxed()); + ) { + Ok(projection) => projection, + Err(err) if limited => return Ok(async move { TaskResult::Terminal(err) }.boxed()), + Err(err) => return Err(err), + }; + return Ok( + async move { + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(err) if limited => TaskResult::Terminal(err), + Err(err) => TaskResult::Recoverable(err), + } + } + .boxed(), + ); }; let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); @@ -75,31 +104,46 @@ pub fn split_exec( ctx.reader .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; let array_fut = async move { - let mask = filter_mask.await?; + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(err) => return TaskResult::Recoverable(err), + }; if mask.all_false() { - return Ok(None); + return TaskResult::Array(None); } - projection.await.map(Some) + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(err) => TaskResult::Recoverable(err), + } }; return Ok(array_fut.boxed()); }; let array_fut = async move { - let mask = filter_mask.await?; + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(err) => return TaskResult::Recoverable(err), + }; // A filter error above returns before reserving any rows. Once filtering has succeeded, // reserve only the matching rows and construct projection work for that limited mask. let mask = row_limit.limit(mask); if mask.all_false() { - return Ok(None); + return TaskResult::Array(None); } - let projection = ctx.reader.projection_evaluation( + let projection = match ctx.reader.projection_evaluation( &row_range, &ctx.projection, MaskFuture::ready(mask), - )?; - projection.await.map(Some) + ) { + Ok(projection) => projection, + Err(err) => return TaskResult::Terminal(err), + }; + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(err) => TaskResult::Terminal(err), + } }; Ok(array_fut.boxed()) @@ -139,11 +183,22 @@ pub fn project_split( ctx: Arc, row_range: Range, mask: Mask, -) -> VortexResult { - let projection = - ctx.reader - .projection_evaluation(&row_range, &ctx.projection, MaskFuture::ready(mask))?; - Ok(async move { projection.await.map(Some) }.boxed()) +) -> TaskFuture { + async move { + let projection = match ctx.reader.projection_evaluation( + &row_range, + &ctx.projection, + MaskFuture::ready(mask), + ) { + Ok(projection) => projection, + Err(err) => return TaskResult::Terminal(err), + }; + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(err) => TaskResult::Terminal(err), + } + } + .boxed() } /// Build the filtered mask for a split. From 789297e84cd5193c3dac90c79fa967f9fbdb643e Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 11:45:07 +0100 Subject: [PATCH 12/17] Harden limited scan task scheduling Signed-off-by: Adam Gutglick --- vortex-datafusion/src/v2/source.rs | 11 +- vortex-layout/src/scan/repeated_scan.rs | 279 +++++++++++------ vortex-layout/src/scan/scan_builder.rs | 182 ++++++++++- vortex-layout/src/scan/tasks.rs | 391 +++++++++++++++++++----- 4 files changed, 673 insertions(+), 190 deletions(-) diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 155ea463af9..4f8438db810 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -729,11 +729,11 @@ fn estimate_to_df_precision(est: &Precision) -> DFPrecision { #[cfg(test)] mod tests { - use std::sync::Arc; - use std::time::Duration; use std::pin::Pin; + use std::sync::Arc; use std::task::Context; use std::task::Poll; + use std::time::Duration; use futures::Stream; use futures::TryStreamExt; @@ -845,9 +845,10 @@ mod tests { let streams = futures::stream::iter([Ok::<_, VortexError>(first), Ok::<_, VortexError>(second)]); let flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); - let chunks = tokio::time::timeout(Duration::from_secs(1), flattened.try_collect::>()) - .await - .map_err(|_| vortex_err!("ordered flatten did not warm the later partition"))??; + let chunks = + tokio::time::timeout(Duration::from_secs(1), flattened.try_collect::>()) + .await + .map_err(|_| vortex_err!("ordered flatten did not warm the later partition"))??; assert_eq!(collect_i32(chunks)?, [0, 100]); Ok(()) diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 9f4e5b7875d..9a2cdcc7f27 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -4,11 +4,15 @@ use std::cmp; use std::iter; use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; -use futures::FutureExt; use futures::Stream; use futures::StreamExt; +use futures::future; use futures::stream::BoxStream; use itertools::Either; use itertools::Itertools; @@ -22,6 +26,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; +use vortex_io::runtime::Handle; use vortex_io::runtime::Task; use vortex_io::session::RuntimeSessionExt; use vortex_scan::selection::Selection; @@ -65,6 +70,155 @@ pub struct RepeatedScan { dtype: DType, } +/// A source of split tasks that has not yet applied task concurrency or output error handling. +#[must_use = "task streams must be scheduled"] +struct TaskStream { + inner: BoxStream<'static, Task>, +} + +impl TaskStream { + fn new(stream: impl Stream> + Send + 'static) -> Self { + Self { + inner: stream.boxed(), + } + } + + fn eager(handle: Handle, tasks: Vec) -> Self { + Self::new(futures::stream::iter(tasks).map(move |task| handle.spawn(task))) + } + + fn unordered_filtered_limit( + handle: Handle, + split_ranges: Vec>, + selection: Selection, + ctx: Arc, + row_limit: RowLimit, + ) -> Self { + let take_limit = row_limit.clone(); + Self::new( + futures::stream::iter(split_ranges) + .take_while(move |_| future::ready(!take_limit.is_exhausted())) + .filter_map(move |range| { + // Build the row mask and split task synchronously so the I/O system sees the + // split's ranges as soon as `buffer_unordered` pulls it, without cloning + // `selection`. + let row_mask = selection.row_mask(&range); + let task = (!row_mask.mask().all_false()).then(|| { + split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) + .unwrap_or_else(TaskFuture::terminal) + }); + future::ready(task.map(|task| handle.spawn(task))) + }), + ) + } + + fn ordered_filtered_limit( + handle: Handle, + split_ranges: Vec>, + selection: Selection, + ctx: Arc, + row_limit: RowLimit, + concurrency: usize, + ) -> Self { + // Stage one: evaluate the filter for each split concurrently while preserving split + // order. + let filter_tasks = { + let filter_ctx = Arc::clone(&ctx); + let filter_handle = handle.clone(); + let take_limit = row_limit.clone(); + futures::stream::iter(split_ranges) + .take_while(move |_| future::ready(!take_limit.is_exhausted())) + .filter_map(move |range| { + // Build the row mask synchronously so the I/O system sees the split's ranges + // as soon as `buffered` pulls it, without cloning `selection`. + let row_mask = selection.row_mask(&range); + let task = (!row_mask.mask().all_false()).then(|| { + filter_handle.spawn(filter_split(Arc::clone(&filter_ctx), row_mask)) + }); + future::ready(task) + }) + .buffered(concurrency) + }; + + // Seam + stage two: reserve in split order (this map runs sequentially as the ordered + // stage-one stream is consumed), then spawn projection work for the reserved mask. The + // second gate drops results that were already filtering when an earlier split exhausted + // the budget, including any later filter errors. + let stage_two_limit = row_limit.clone(); + Self::new( + filter_tasks + .take_while(move |_| future::ready(!stage_two_limit.is_exhausted())) + .map(move |result| { + let task = match result { + Ok((row_range, mask)) => { + let mask = row_limit.limit(mask); + if mask.all_false() { + TaskFuture::empty() + } else { + project_split(Arc::clone(&ctx), row_range, mask) + } + } + Err(error) => TaskFuture::recoverable(error), + }; + handle.spawn(task) + }), + ) + } +} + +impl Stream for TaskStream { + type Item = Task; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_next(cx) + } +} + +/// A buffered task stream that exposes arrays and applies the task error policy. +pub(crate) struct ScheduledTaskStream { + tasks: Option>, +} + +impl ScheduledTaskStream { + fn new(tasks: TaskStream, ordered: bool, concurrency: usize) -> Self { + let tasks = if ordered { + tasks.buffered(concurrency).boxed() + } else { + tasks.buffer_unordered(concurrency).boxed() + }; + Self { tasks: Some(tasks) } + } +} + +impl Stream for ScheduledTaskStream { + type Item = VortexResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let Some(tasks) = self.tasks.as_mut() else { + return Poll::Ready(None); + }; + let result = ready!(tasks.as_mut().poll_next(cx)); + let Some(result) = result else { + self.tasks = None; + return Poll::Ready(None); + }; + match result { + TaskResult::Array(Some(array)) => return Poll::Ready(Some(Ok(array))), + TaskResult::Array(None) => {} + TaskResult::Recoverable(error) => return Poll::Ready(Some(Err(error))), + TaskResult::Terminal(error) => { + // Drop queued and in-flight task handles before yielding the error. Their + // `Drop` implementations abort work that can no longer contribute to this + // limited scan. + self.tasks = None; + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + impl RepeatedScan { pub fn dtype(&self) -> &DType { &self.dtype @@ -220,7 +374,7 @@ impl RepeatedScan { pub(crate) fn execute_stream( &self, row_range: Option>, - ) -> VortexResult> + Send + 'static> { + ) -> VortexResult { let num_workers = get_available_parallelism().unwrap_or(1); let row_limit = self .row_limit @@ -243,33 +397,24 @@ impl RepeatedScan { return Ok(self.ordered_filtered_limit_stream(row_range, row_limit, concurrency)); } - let ctx = self.task_context(); - let selection = self.selection.clone(); - let task_limit = row_limit.clone(); - let tasks = futures::stream::iter(self.split_ranges(row_range)) - .take_while(move |_| futures::future::ready(!task_limit.is_exhausted())) - .filter_map(move |range| { - // Build the row mask and split task synchronously so the IO system sees the - // split's ranges as soon as `buffer_unordered` pulls it, without cloning - // `selection`. - let row_mask = selection.row_mask(&range); - let spawned = (!row_mask.mask().all_false()).then(|| { - let task = split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) - .unwrap_or_else(|err| async move { TaskResult::Terminal(err) }.boxed()); - handle.spawn(task) - }); - async move { spawned } - }); - - return Ok(schedule(tasks, false, concurrency)); + let tasks = TaskStream::unordered_filtered_limit( + handle, + self.split_ranges(row_range), + self.selection.clone(), + self.task_context(), + row_limit, + ); + return Ok(ScheduledTaskStream::new(tasks, false, concurrency)); } // No filter (or no limit): build every task eagerly so the IO system sees all split // ranges up front. A no-filter limit is applied to each selection mask inside `execute`. - let tasks = futures::stream::iter(self.execute(row_range, row_limit)?) - .map(move |task| handle.spawn(task)); + let tasks = TaskStream::eager(handle, self.execute(row_range, row_limit)?); - Ok(schedule(tasks, self.ordered, concurrency)) + Ok({ + let ordered = self.ordered; + ScheduledTaskStream::new(tasks, ordered, concurrency) + }) } /// Ordered filtered scan with a shared row limit. @@ -287,83 +432,17 @@ impl RepeatedScan { row_range: Option>, row_limit: RowLimit, concurrency: usize, - ) -> BoxStream<'static, VortexResult> { - let ctx = self.task_context(); - let handle = self.session.handle(); - - // Stage one: evaluate the filter for each split concurrently while preserving split order. - let filter_tasks = { - let ctx = Arc::clone(&ctx); - let handle = handle.clone(); - let selection = self.selection.clone(); - let take_limit = row_limit.clone(); - futures::stream::iter(self.split_ranges(row_range)) - .take_while(move |_| futures::future::ready(!take_limit.is_exhausted())) - .filter_map(move |range| { - // Build the row mask synchronously so the IO system sees the split's ranges as - // soon as `buffered` pulls it, without cloning `selection`. - let row_mask = selection.row_mask(&range); - let spawned = (!row_mask.mask().all_false()) - .then(|| handle.spawn(filter_split(Arc::clone(&ctx), row_mask))); - async move { spawned } - }) - .buffered(concurrency) - }; - - // Seam + stage two: reserve in split order (this map runs sequentially as the ordered - // stage-one stream is consumed), then spawn projection work for the reserved mask. The - // second gate drops results that were already filtering when an earlier split exhausted - // the budget, including any later filter errors. - let stage_two_limit = row_limit.clone(); - let tasks = filter_tasks - .take_while(move |_| futures::future::ready(!stage_two_limit.is_exhausted())) - .map(move |result| { - let task: TaskFuture = match result { - Ok((row_range, mask)) => { - let mask = row_limit.limit(mask); - if mask.all_false() { - async { TaskResult::Array(None) }.boxed() - } else { - project_split(Arc::clone(&ctx), row_range, mask) - } - } - Err(err) => async move { TaskResult::Recoverable(err) }.boxed(), - }; - handle.spawn(task) - }); - schedule(tasks, true, concurrency) - } -} - -/// Spawn-buffer a stream of split tasks, preserving recoverable filter errors and stopping after -/// a projection error that occurred after a row limit reservation. -fn schedule( - tasks: S, - ordered: bool, - concurrency: usize, -) -> BoxStream<'static, VortexResult> -where - S: Stream> + Send + 'static, -{ - let mut stream = if ordered { - tasks.buffered(concurrency).boxed() - } else { - tasks.buffer_unordered(concurrency).boxed() - }; - async_stream::stream! { - while let Some(result) = stream.next().await { - match result { - TaskResult::Array(Some(array)) => yield Ok(array), - TaskResult::Array(None) => {} - TaskResult::Recoverable(err) => yield Err(err), - TaskResult::Terminal(err) => { - yield Err(err); - return; - } - } - } + ) -> ScheduledTaskStream { + let tasks = TaskStream::ordered_filtered_limit( + self.session.handle(), + self.split_ranges(row_range), + self.selection.clone(), + self.task_context(), + row_limit, + concurrency, + ); + ScheduledTaskStream::new(tasks, true, concurrency) } - .boxed() } fn intersect_ranges(left: Option<&Range>, right: Option>) -> Option> { diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index d92d42d54a4..316f9d1cfc0 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -406,6 +406,7 @@ mod test { use std::time::Duration; use futures::Stream; + use futures::channel::oneshot; use futures::task::noop_waker_ref; use parking_lot::Mutex; use vortex_array::ArrayRef; @@ -982,24 +983,32 @@ mod test { } } + #[derive(Debug, Clone, Copy)] + enum FirstSplitFailure { + Filter, + Projection, + } + #[derive(Debug)] - struct ErrorThenMatchLayoutReader { + struct FirstSplitFailureLayoutReader { name: Arc, dtype: DType, projection_masks: Arc>>, + failure: FirstSplitFailure, } - impl ErrorThenMatchLayoutReader { - fn new(projection_masks: Arc>>) -> Self { + impl FirstSplitFailureLayoutReader { + fn new(projection_masks: Arc>>, failure: FirstSplitFailure) -> Self { Self { - name: Arc::from("error-then-match"), + name: Arc::from("first-split-failure"), dtype: DType::Primitive(PType::I32, Nullability::NonNullable), projection_masks, + failure, } } } - impl LayoutReader for ErrorThenMatchLayoutReader { + impl LayoutReader for FirstSplitFailureLayoutReader { fn name(&self) -> &Arc { &self.name } @@ -1039,7 +1048,7 @@ mod test { _expr: &Expression, mask: MaskFuture, ) -> VortexResult { - if row_range.start == 0 { + if row_range.start == 0 && matches!(self.failure, FirstSplitFailure::Filter) { let len = mask.len(); return Ok(MaskFuture::new(len, async move { Err(vortex_err!("first split filter failed")) @@ -1056,10 +1065,14 @@ mod test { ) -> VortexResult { let row_range = row_range.clone(); let projection_masks = Arc::clone(&self.projection_masks); + let failure = self.failure; Ok(Box::pin(async move { let mask = mask.await?; projection_masks.lock().push(mask.true_count()); + if row_range.start == 0 && matches!(failure, FirstSplitFailure::Projection) { + return Err(vortex_err!("first split projection failed")); + } let start = i32::try_from(row_range.start) .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; let end = i32::try_from(row_range.end) @@ -1075,6 +1088,114 @@ mod test { } } + struct MatchThenErrorLayoutReader { + name: Arc, + dtype: DType, + release_first_filter: Mutex>>, + wait_for_later_filter: Mutex>>, + } + + impl MatchThenErrorLayoutReader { + fn new() -> Self { + let (release_first_filter, wait_for_later_filter) = oneshot::channel(); + Self { + name: Arc::from("match-then-error"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + release_first_filter: Mutex::new(Some(release_first_filter)), + wait_for_later_filter: Mutex::new(Some(wait_for_later_filter)), + } + } + } + + impl LayoutReader for MatchThenErrorLayoutReader { + fn name(&self) -> &Arc { + &self.name + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 2 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + let row_range = split_range.row_range(); + splits.push(split_range.row_offset() + row_range.start + 1); + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let len = mask.len(); + if row_range.start == 0 { + let wait_for_later_filter = self + .wait_for_later_filter + .lock() + .take() + .ok_or_else(|| vortex_err!("first split filter was evaluated twice"))?; + return Ok(MaskFuture::new(len, async move { + wait_for_later_filter + .await + .map_err(|_| vortex_err!("later split filter was cancelled"))?; + mask.await + })); + } + + self.release_first_filter + .lock() + .take() + .ok_or_else(|| vortex_err!("later split filter was evaluated twice"))? + .send(()) + .map_err(|_| vortex_err!("first split filter was cancelled"))?; + Ok(MaskFuture::new(len, async move { + Err(vortex_err!("later split filter failed")) + })) + } + + fn projection_evaluation( + &self, + row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + let row_range = row_range.clone(); + Ok(Box::pin(async move { + let start = i32::try_from(row_range.start) + .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; + let end = i32::try_from(row_range.end) + .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; + PrimitiveArray::from_iter(start..end) + .into_array() + .filter(mask.await?) + })) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + fn collect_scan_values(iter: I) -> VortexResult> where I: IntoIterator>, @@ -1141,14 +1262,34 @@ mod test { Ok(()) } + #[test] + fn ordered_filtered_limit_drops_prefetched_filter_errors_after_the_budget_is_full() + -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new(MatchThenErrorLayoutReader::new()); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(1) + .into_stream()?; + + // The first split waits until the later split has already produced its error. Once the + // first matching row fills the limit, that speculative error must be dropped. + let values = collect_scan_values(runtime.block_on_stream(stream))?; + + assert_eq!(values, [0]); + Ok(()) + } + #[test] fn filter_errors_are_stream_items_and_do_not_consume_the_limit() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let projection_masks = Arc::new(Mutex::new(Vec::new())); - let reader = Arc::new(ErrorThenMatchLayoutReader::new(Arc::clone( - &projection_masks, - ))); + let reader = Arc::new(FirstSplitFailureLayoutReader::new( + Arc::clone(&projection_masks), + FirstSplitFailure::Filter, + )); let stream = ScanBuilder::new(session, reader) .with_filter(root()) .with_limit(1) @@ -1170,6 +1311,29 @@ mod test { Ok(()) } + #[test] + fn projection_error_after_reservation_terminates_the_limited_scan() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let projection_masks = Arc::new(Mutex::new(Vec::new())); + let reader = Arc::new(FirstSplitFailureLayoutReader::new( + Arc::clone(&projection_masks), + FirstSplitFailure::Projection, + )); + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + // A budget of two leaves room for the second matching split. Continuing after the + // first projection failure would therefore yield a second stream item. + .with_limit(2) + .into_stream()?; + let mut iter = runtime.block_on_stream(stream); + + assert!(matches!(iter.next(), Some(Err(_)))); + assert!(iter.next().is_none()); + assert!(projection_masks.lock().contains(&1)); + Ok(()) + } + #[test] fn projection_errors_are_stream_items() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index a830930fb35..a11baedb219 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -3,9 +3,13 @@ //! Split scanning task implementation. +use std::future::Future; use std::ops::BitAnd; use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; use bit_vec::BitVec; use futures::FutureExt; @@ -13,12 +17,13 @@ use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::expr::Expression; -use vortex_error::VortexExpect; use vortex_error::VortexError; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_scan::row_mask::RowMask; +use crate::ArrayFuture; use crate::LayoutReader; use crate::scan::filter::FilterExpr; use crate::scan::limit::RowLimit; @@ -37,7 +42,125 @@ pub(crate) enum TaskResult { Terminal(VortexError), } -pub(crate) type TaskFuture = BoxFuture<'static, TaskResult>; +/// A future that executes one split and classifies any failure by whether it happened before or +/// after a row-limit reservation. +#[must_use = "split tasks must be scheduled or awaited"] +pub(crate) struct TaskFuture { + inner: BoxFuture<'static, TaskResult>, +} + +#[derive(Clone, Copy)] +enum ErrorDisposition { + Recoverable, + Terminal, +} + +impl ErrorDisposition { + fn result(self, error: VortexError) -> TaskResult { + match self { + Self::Recoverable => TaskResult::Recoverable(error), + Self::Terminal => TaskResult::Terminal(error), + } + } +} + +impl TaskFuture { + fn new(future: impl Future + Send + 'static) -> Self { + Self { + inner: future.boxed(), + } + } + + fn ready(result: TaskResult) -> Self { + Self::new(futures::future::ready(result)) + } + + pub(crate) fn empty() -> Self { + Self::ready(TaskResult::Array(None)) + } + + pub(crate) fn recoverable(error: VortexError) -> Self { + Self::ready(TaskResult::Recoverable(error)) + } + + pub(crate) fn terminal(error: VortexError) -> Self { + Self::ready(TaskResult::Terminal(error)) + } + + fn projection(projection: ArrayFuture, errors: ErrorDisposition) -> Self { + Self::new(async move { + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => errors.result(error), + } + }) + } + + fn filtered_projection(filter_mask: MaskFuture, projection: ArrayFuture) -> Self { + Self::new(async move { + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(error) => return TaskResult::Recoverable(error), + }; + if mask.all_false() { + return TaskResult::Array(None); + } + + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => TaskResult::Recoverable(error), + } + }) + } + + fn limited_filtered_projection( + ctx: Arc, + row_range: Range, + filter_mask: MaskFuture, + row_limit: RowLimit, + ) -> Self { + Self::new(async move { + let mask = match filter_mask.await { + Ok(mask) => mask, + Err(error) => return TaskResult::Recoverable(error), + }; + // A filter error above returns before reserving any rows. Once filtering has + // succeeded, reserve only the matching rows and defer projection construction until + // this task is polled by the runtime. + let mask = row_limit.limit(mask); + if mask.all_false() { + return TaskResult::Array(None); + } + + Self::deferred_projection(ctx, row_range, mask).await + }) + } + + fn deferred_projection(ctx: Arc, row_range: Range, mask: Mask) -> Self { + Self::new(async move { + let projection = match ctx.reader.projection_evaluation( + &row_range, + &ctx.projection, + MaskFuture::ready(mask), + ) { + Ok(projection) => projection, + Err(error) => return TaskResult::Terminal(error), + }; + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => TaskResult::Terminal(error), + } + }) + } +} + +impl Future for TaskFuture { + type Output = TaskResult; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.inner.as_mut().poll(cx) + } +} /// Logic for executing a single split reading task. /// N.B. read_mask should be evaluated against all_false() before calling this @@ -70,7 +193,7 @@ pub fn split_exec( row_mask }; if row_mask.all_false() { - return Ok(async { TaskResult::Array(None) }.boxed()); + return Ok(TaskFuture::empty()); } // With no filter, limit the selection before constructing projection work. @@ -80,19 +203,15 @@ pub fn split_exec( MaskFuture::ready(row_mask), ) { Ok(projection) => projection, - Err(err) if limited => return Ok(async move { TaskResult::Terminal(err) }.boxed()), + Err(err) if limited => return Ok(TaskFuture::terminal(err)), Err(err) => return Err(err), }; - return Ok( - async move { - match projection.await { - Ok(array) => TaskResult::Array(Some(array)), - Err(err) if limited => TaskResult::Terminal(err), - Err(err) => TaskResult::Recoverable(err), - } - } - .boxed(), - ); + let errors = if limited { + ErrorDisposition::Terminal + } else { + ErrorDisposition::Recoverable + }; + return Ok(TaskFuture::projection(projection, errors)); }; let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); @@ -103,50 +222,15 @@ pub fn split_exec( let projection = ctx.reader .projection_evaluation(&row_range, &ctx.projection, filter_mask.clone())?; - let array_fut = async move { - let mask = match filter_mask.await { - Ok(mask) => mask, - Err(err) => return TaskResult::Recoverable(err), - }; - if mask.all_false() { - return TaskResult::Array(None); - } - - match projection.await { - Ok(array) => TaskResult::Array(Some(array)), - Err(err) => TaskResult::Recoverable(err), - } - }; - return Ok(array_fut.boxed()); - }; - - let array_fut = async move { - let mask = match filter_mask.await { - Ok(mask) => mask, - Err(err) => return TaskResult::Recoverable(err), - }; - // A filter error above returns before reserving any rows. Once filtering has succeeded, - // reserve only the matching rows and construct projection work for that limited mask. - let mask = row_limit.limit(mask); - if mask.all_false() { - return TaskResult::Array(None); - } - - let projection = match ctx.reader.projection_evaluation( - &row_range, - &ctx.projection, - MaskFuture::ready(mask), - ) { - Ok(projection) => projection, - Err(err) => return TaskResult::Terminal(err), - }; - match projection.await { - Ok(array) => TaskResult::Array(Some(array)), - Err(err) => TaskResult::Terminal(err), - } + return Ok(TaskFuture::filtered_projection(filter_mask, projection)); }; - Ok(array_fut.boxed()) + Ok(TaskFuture::limited_filtered_projection( + ctx, + row_range, + filter_mask, + row_limit, + )) } /// Evaluate the filter for a split and return the matching mask together with its row range. @@ -179,26 +263,8 @@ pub fn filter_split( /// /// This is the second stage of the ordered filtered-limit pipeline, run after the caller has /// reserved rows against the limit in split order (see [`filter_split`]). -pub fn project_split( - ctx: Arc, - row_range: Range, - mask: Mask, -) -> TaskFuture { - async move { - let projection = match ctx.reader.projection_evaluation( - &row_range, - &ctx.projection, - MaskFuture::ready(mask), - ) { - Ok(projection) => projection, - Err(err) => return TaskResult::Terminal(err), - }; - match projection.await { - Ok(array) => TaskResult::Array(Some(array)), - Err(err) => TaskResult::Terminal(err), - } - } - .boxed() +pub fn project_split(ctx: Arc, row_range: Range, mask: Mask) -> TaskFuture { + TaskFuture::deferred_projection(ctx, row_range, mask) } /// Build the filtered mask for a split. @@ -286,3 +352,176 @@ pub struct TaskContext { /// The projection expression to apply to gather the scanned rows. pub projection: Expression, } + +#[cfg(test)] +mod tests { + use std::ops::Range; + use std::sync::Arc; + use std::sync::mpsc; + use std::thread; + use std::time::Duration; + + use parking_lot::Mutex; + use vortex_array::IntoArray; + use vortex_array::MaskFuture; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::FieldMask; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::Expression; + use vortex_array::expr::root; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::TaskContext; + use super::TaskResult; + use super::project_split; + use crate::ArrayFuture; + use crate::LayoutReader; + use crate::RowSplits; + use crate::SplitRange; + + struct BlockingProjectionReader { + name: Arc, + dtype: DType, + started: mpsc::Sender<()>, + gate: Arc>, + } + + impl LayoutReader for BlockingProjectionReader { + fn name(&self) -> &Arc { + &self.name + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + 1 + } + + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + splits.push(split_range.root_row_range().end); + Ok(()) + } + + fn pruning_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: Mask, + ) -> VortexResult { + Ok(MaskFuture::ready(mask)) + } + + fn filter_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + mask: MaskFuture, + ) -> VortexResult { + Ok(mask) + } + + fn projection_evaluation( + &self, + _row_range: &Range, + _expr: &Expression, + _mask: MaskFuture, + ) -> VortexResult { + self.started + .send(()) + .map_err(|_| vortex_err!("test projection-start receiver dropped"))?; + let _guard = self.gate.lock(); + let array = PrimitiveArray::from_iter([0_i32]).into_array(); + Ok(Box::pin(async move { Ok(array) })) + } + } + + #[test] + fn project_split_defers_projection_construction_until_task_poll() -> VortexResult<()> { + let gate = Arc::new(Mutex::new(())); + let guard = gate.lock(); + let (started_send, started_recv) = mpsc::channel(); + let reader: Arc = Arc::new(BlockingProjectionReader { + name: Arc::from("blocking-projection"), + dtype: DType::Primitive(PType::I32, Nullability::NonNullable), + started: started_send, + gate: Arc::clone(&gate), + }); + let ctx = Arc::new(TaskContext { + filter: None, + reader, + projection: root(), + }); + + let (task_send, task_recv) = mpsc::channel(); + let construction = thread::spawn(move || { + let task = project_split(ctx, 0..1, Mask::new_true(1)); + drop(task_send.send(task)); + }); + + let task = match task_recv.recv_timeout(Duration::from_secs(1)) { + Ok(task) => task, + Err(_) => { + drop(guard); + drop(construction.join()); + return Err(vortex_err!( + "constructing a projection task blocked before it was polled" + )); + } + }; + if construction.join().is_err() { + return Err(vortex_err!("projection-task construction panicked")); + } + match started_recv.try_recv() { + Err(mpsc::TryRecvError::Empty) => {} + Ok(()) => { + return Err(vortex_err!( + "projection construction started before the task was polled" + )); + } + Err(mpsc::TryRecvError::Disconnected) => { + return Err(vortex_err!( + "projection-start sender disconnected unexpectedly" + )); + } + } + + let (result_send, result_recv) = mpsc::channel(); + let execution = thread::spawn(move || { + let result = futures::executor::block_on(task); + drop(result_send.send(result)); + }); + + if started_recv.recv_timeout(Duration::from_secs(1)).is_err() { + drop(guard); + drop(execution.join()); + return Err(vortex_err!( + "projection construction did not start when polled" + )); + } + drop(guard); + let result = result_recv + .recv_timeout(Duration::from_secs(1)) + .map_err(|_| vortex_err!("projection task did not complete"))?; + if execution.join().is_err() { + return Err(vortex_err!("projection task panicked")); + } + + assert!(matches!(result, TaskResult::Array(Some(_)))); + Ok(()) + } +} From 9ef4b5252fee380b461d764cee124b011739c250 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 11:58:06 +0100 Subject: [PATCH 13/17] clippy Signed-off-by: Adam Gutglick --- vortex-datafusion/src/v2/source.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index 4f8438db810..cd3aa68ec51 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -98,6 +98,7 @@ use futures::StreamExt; use futures::TryStreamExt; use futures::future::try_join_all; use futures::stream::BoxStream; +use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use vortex::array::ArrayRef; use vortex::array::VortexSessionExecute; @@ -688,7 +689,7 @@ where .map(move |stream_result| { let handle = handle.clone(); let receiver = stream_result.map(|mut stream| { - let (tx, rx) = tokio::sync::mpsc::channel(CHUNK_BUFFER_CAPACITY); + let (tx, rx) = mpsc::channel(CHUNK_BUFFER_CAPACITY); handle .spawn(async move { loop { @@ -738,6 +739,7 @@ mod tests { use futures::Stream; use futures::TryStreamExt; use tokio::sync::Notify; + use tokio::sync::oneshot; use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; @@ -773,7 +775,7 @@ mod tests { Ok(values) } - struct DropNotifier(Option>); + struct DropNotifier(Option>); impl Drop for DropNotifier { fn drop(&mut self) { @@ -858,7 +860,7 @@ mod tests { async fn ordered_flatten_cancels_a_pending_drain_when_dropped() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let started = Arc::new(Notify::new()); - let (dropped_send, dropped_recv) = tokio::sync::oneshot::channel(); + let (dropped_send, dropped_recv) = oneshot::channel(); let pending = ArrayStreamAdapter::new( dtype, DropNotifyingPendingStream { From 32128d2f0f74df14b6af75d0a7c8467a18e79ba4 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 14:36:35 +0100 Subject: [PATCH 14/17] another pass Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/opener.rs | 21 +++- vortex-datafusion/src/v2/source.rs | 128 ++++++++++++++++++--- vortex-layout/src/scan/arrow.rs | 5 +- vortex-layout/src/scan/limit.rs | 73 ++++++++++++ vortex-layout/src/scan/repeated_scan.rs | 16 ++- vortex-layout/src/scan/scan_builder.rs | 53 +++++++++ vortex-layout/src/scan/tasks.rs | 18 +-- 7 files changed, 282 insertions(+), 32 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 50009b10cd5..3988b036b2b 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -1014,13 +1014,30 @@ mod tests { write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; let file = PartitionedFile::new(file_path.to_string(), data_size); let table_schema = TableSchema::from_file_schema(batch.schema()); - let filter = logical2physical(&col("a").gt(lit(0_i32)), table_schema.table_schema()); + // `a > 3` excludes the first three rows, so a limit applied *before* filtering would take + // rows [1, 2, 3] and filter them all out (yielding nothing), whereas a limit applied + // *after* filtering yields the first three matching rows [4, 5, 6]. Asserting the values + // (not just the count) is what makes this test able to detect a pre-filter regression. + let filter = logical2physical(&col("a").gt(lit(3_i32)), table_schema.table_schema()); let mut opener = make_opener(object_store, table_schema, Some(filter)); opener.limit = Some(3); let batches = opener.open(file)?.await?.try_collect::>().await?; - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("projected column should be Int32") + .values() + .to_vec() + }) + .collect::>(); + + assert_eq!(values, [4, 5, 6]); Ok(()) } diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index cd3aa68ec51..d6db39066bc 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -106,8 +106,10 @@ use vortex::array::stream::SendableArrayStream; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::Nullability; +use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_err; use vortex::expr::Expression; use vortex::expr::and as vx_and; use vortex::expr::get_item; @@ -116,6 +118,7 @@ use vortex::expr::root; use vortex::expr::stats::Precision; use vortex::expr::transform::replace; use vortex::io::runtime::Handle; +use vortex::io::runtime::JoinOutcome; use vortex::io::session::RuntimeSessionExt; use vortex::scan::DataSourceRef; use vortex::scan::ScanRequest; @@ -690,24 +693,38 @@ where let handle = handle.clone(); let receiver = stream_result.map(|mut stream| { let (tx, rx) = mpsc::channel(CHUNK_BUFFER_CAPACITY); - handle - .spawn(async move { - loop { - tokio::select! { - _ = tx.closed() => break, - item = stream.next() => { - let Some(item) = item else { - break; - }; - if tx.send(item).await.is_err() { - break; - } + // Keep the drain task's handle (do not `detach` it) so that a panic inside the + // source stream is surfaced as a terminal error. A detached task's panic would be + // discarded; `tx` would drop, the receiver would end cleanly, and `try_flatten` + // would treat the truncated partition as a normal completion — silently returning + // incomplete results as success. + let drain = handle.spawn(async move { + loop { + tokio::select! { + _ = tx.closed() => break, + item = stream.next() => { + let Some(item) = item else { + break; + }; + if tx.send(item).await.is_err() { + break; } } } - }) - .detach(); - ReceiverStream::new(rx) + } + }); + // Once the receiver is drained, join the drain task and turn a panic into a + // terminal error item. A clean completion or a benign runtime abort ends the + // partition without appending an extra item. + let finalizer = futures::stream::once(async move { + let mut drain = drain; + match futures::future::poll_fn(|cx| drain.poll_join(cx)).await { + JoinOutcome::Completed(()) | JoinOutcome::Aborted => None, + JoinOutcome::Panicked(payload) => Some(Err(panic_to_vortex_err(payload))), + } + }) + .filter_map(|item| async move { item }); + ReceiverStream::new(rx).chain(finalizer) }); futures::future::ready(receiver) }) @@ -716,6 +733,17 @@ where .boxed() } +/// Convert a caught panic payload from an ordered drain task into a terminal [`VortexError`], so +/// the scan surfaces the failure instead of silently truncating its output. +fn panic_to_vortex_err(payload: Box) -> VortexError { + let message = payload + .downcast_ref::<&'static str>() + .map(|s| (*s).to_string()) + .or_else(|| payload.downcast_ref::().cloned()) + .unwrap_or_else(|| "unknown panic".to_string()); + vortex_err!("ordered scan partition drain task panicked: {message}") +} + /// Convert a Vortex [`Option`] to a DataFusion /// [`DataFusionPrecision`]. /// @@ -886,4 +914,74 @@ mod tests { dropped.map_err(|_| vortex_err!("pending source dropped without notifying the test"))?; Ok(()) } + + struct PanickingStream; + + impl Stream for PanickingStream { + type Item = VortexResult; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + panic!("boom") + } + } + + #[tokio::test] + async fn ordered_flatten_surfaces_a_partition_panic_as_an_error() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let panicking = ArrayStreamAdapter::new(dtype, PanickingStream).boxed(); + + let streams = futures::stream::iter([Ok::<_, VortexError>(panicking)]); + let flattened = flatten_scan_streams(streams, true, 1, TokioRuntime::current()); + let result = flattened.try_collect::>().await; + + assert!( + result.is_err(), + "a panicking ordered partition must surface as an error, not silent truncation" + ); + Ok(()) + } + + #[tokio::test] + async fn ordered_flatten_cancels_all_warmed_drains_when_dropped() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let started = Arc::new(Notify::new()); + let mut dropped_recvs = Vec::new(); + let mut streams = Vec::new(); + for _ in 0..3 { + let (dropped_send, dropped_recv) = oneshot::channel(); + dropped_recvs.push(dropped_recv); + let pending = ArrayStreamAdapter::new( + dtype.clone(), + DropNotifyingPendingStream { + started: Arc::clone(&started), + _drop_notifier: DropNotifier(Some(dropped_send)), + }, + ) + .boxed(); + streams.push(Ok::<_, VortexError>(pending)); + } + + // A concurrency of 4 yields a lookahead of 4, so all three partitions warm (spawn their + // drains) concurrently before the first one is drained. + let streams = futures::stream::iter(streams); + let mut flattened = flatten_scan_streams(streams, true, 4, TokioRuntime::current()); + let mut next = Box::pin(futures::StreamExt::next(&mut flattened)); + tokio::select! { + _ = started.notified() => {} + _ = &mut next => return Err(vortex_err!("pending drain unexpectedly produced a chunk")), + } + drop(next); + drop(flattened); + + for dropped_recv in dropped_recvs { + let dropped = tokio::time::timeout(Duration::from_secs(1), dropped_recv) + .await + .map_err(|_| { + vortex_err!("dropping the flattened stream did not cancel every warmed drain") + })?; + dropped + .map_err(|_| vortex_err!("pending source dropped without notifying the test"))?; + } + Ok(()) + } } diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index f62bf4e920f..223c5962441 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -19,6 +19,7 @@ use vortex_arrow::ArrowSessionExt; use vortex_error::VortexResult; use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; +use vortex_utils::parallelism::get_available_parallelism; use crate::scan::scan_builder::ScanBuilder; @@ -54,9 +55,7 @@ impl ScanBuilder { let struct_field = Arc::new(Field::new_struct("", schema.fields().clone(), false)); let session = self.session().clone(); let handle = session.handle(); - let concurrency = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(1); + let concurrency = get_available_parallelism().unwrap_or(1); let stream = self .into_stream()? diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs index 60dcf921432..7242e002c38 100644 --- a/vortex-layout/src/scan/limit.rs +++ b/vortex-layout/src/scan/limit.rs @@ -50,3 +50,76 @@ impl RowLimit { } } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; + use std::thread; + + use vortex_mask::Mask; + + use super::RowLimit; + + #[test] + fn reserve_grants_up_to_the_remaining_budget() { + let limit = RowLimit::new(5); + assert_eq!(limit.reserve(3), 3); + assert!(!limit.is_exhausted()); + // Only two rows remain, so a larger request saturates at what is left. + assert_eq!(limit.reserve(10), 2); + assert!(limit.is_exhausted()); + // Once exhausted, further requests grant nothing. + assert_eq!(limit.reserve(1), 0); + } + + #[test] + fn limit_keeps_the_earliest_granted_rows() { + let limit = RowLimit::new(2); + // Rows 0, 2, 3, 5 are selected; only the first two survive the budget of 2. + let mask = Mask::from_iter([true, false, true, true, false, true]); + let limited = limit.limit(mask); + + assert_eq!(limited.true_count(), 2); + assert!(limited.value(0)); + assert!(limited.value(2)); + assert!(!limited.value(3)); + assert!(!limited.value(5)); + assert!(limit.is_exhausted()); + } + + #[test] + fn concurrent_reservations_never_exceed_the_budget() { + const THREADS: usize = 8; + const PER_THREAD: u64 = 10_000; + const LIMIT: u64 = 25_000; + + let limit = RowLimit::new(LIMIT); + let granted_total = Arc::new(AtomicU64::new(0)); + let barrier = Arc::new(Barrier::new(THREADS)); + + thread::scope(|scope| { + for _ in 0..THREADS { + let limit = limit.clone(); + let granted_total = Arc::clone(&granted_total); + let barrier = Arc::clone(&barrier); + scope.spawn(move || { + // Start all threads together to maximize contention on the atomic. + barrier.wait(); + let mut local = 0; + for _ in 0..PER_THREAD { + local += limit.reserve(1); + } + granted_total.fetch_add(local, Ordering::Relaxed); + }); + } + }); + + // Total requested (THREADS * PER_THREAD = 80_000) exceeds the budget, so exactly the + // budget is granted across all threads — no double-grant, over-grant, or lost reservation. + assert_eq!(granted_total.load(Ordering::Relaxed), LIMIT); + assert!(limit.is_exhausted()); + } +} diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 9a2cdcc7f27..7c3ea7973b4 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -63,6 +63,11 @@ pub struct RepeatedScan { /// The number of splits to make progress on concurrently **per-thread**. concurrency: usize, /// Maximal number of rows to read (after filtering). + /// + /// When no shared [`RowLimit`] (`row_limit`) is set, this limit is applied independently to + /// each `execute_*` call: a `RepeatedScan` executed over several row ranges may therefore + /// return up to `limit` rows *per call*. Supply a shared `row_limit` to cap the total across + /// calls and sibling partitions. limit: Option, /// An optional row limit shared with sibling external partitions. row_limit: Option, @@ -104,8 +109,11 @@ impl TaskStream { // `selection`. let row_mask = selection.row_mask(&range); let task = (!row_mask.mask().all_false()).then(|| { + // A synchronous split-construction failure happens before any row is + // reserved, so it is recoverable (yielded as a stream error) rather than + // terminal (which would abort the whole scan). See the `TaskResult` docs. split_exec(Arc::clone(&ctx), row_mask, Some(row_limit.clone())) - .unwrap_or_else(TaskFuture::terminal) + .unwrap_or_else(TaskFuture::recoverable) }); future::ready(task.map(|task| handle.spawn(task))) }), @@ -411,10 +419,8 @@ impl RepeatedScan { // ranges up front. A no-filter limit is applied to each selection mask inside `execute`. let tasks = TaskStream::eager(handle, self.execute(row_range, row_limit)?); - Ok({ - let ordered = self.ordered; - ScheduledTaskStream::new(tasks, ordered, concurrency) - }) + let ordered = self.ordered; + Ok(ScheduledTaskStream::new(tasks, ordered, concurrency)) } /// Ordered filtered scan with a shared row limit. diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 316f9d1cfc0..a90274befa8 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -1372,6 +1372,59 @@ mod test { Ok(()) } + #[test] + fn filtered_limit_zero_produces_no_rows() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new(FilteringLayoutReader::new(8, |_| true)); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(0) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + assert!(values.is_empty()); + Ok(()) + } + + #[test] + fn filtered_limit_exceeding_matches_returns_all_matches() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + // Only odd rows match, so there are four matching rows (1, 3, 5, 7). + let reader = Arc::new(FilteringLayoutReader::new(8, |row| row % 2 == 1)); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(100) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + // The scan terminates cleanly at end of input rather than hanging on the unfilled budget. + assert_eq!(values, [1, 3, 5, 7]); + Ok(()) + } + + #[test] + fn filtered_limit_over_empty_input_produces_no_rows() -> VortexResult<()> { + let runtime = SingleThreadRuntime::default(); + let session = session_with_handle(runtime.handle()); + let reader = Arc::new(FilteringLayoutReader::new(0, |_| true)); + + let stream = ScanBuilder::new(session, reader) + .with_filter(root()) + .with_limit(3) + .into_stream()?; + let values = collect_scan_values(runtime.block_on_stream(stream))?; + drain_runtime(&runtime); + + assert!(values.is_empty()); + Ok(()) + } + #[derive(Debug)] struct BlockingSplitsLayoutReader { name: Arc, diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index a11baedb219..034a80eed37 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -177,7 +177,7 @@ impl Future for TaskFuture { /// The final mask is limited before it is given to the reader to perform a filtered projection /// over the split data, yielding the projected array (or `None` when the split selects no rows). /// Limiting before projection prevents decode work for rows that the scan cannot return. -pub fn split_exec( +pub(crate) fn split_exec( ctx: Arc, read_mask: RowMask, row_limit: Option, @@ -240,7 +240,7 @@ pub fn split_exec( /// system can prefetch while earlier splits are still reserving), but it neither reserves against /// the limit nor projects. The caller reserves the returned mask in split order and then projects /// it via [`project_split`], which keeps ordered `LIMIT` semantics without serializing I/O. -pub fn filter_split( +pub(crate) fn filter_split( ctx: Arc, read_mask: RowMask, ) -> BoxFuture<'static, VortexResult<(Range, Mask)>> { @@ -263,7 +263,11 @@ pub fn filter_split( /// /// This is the second stage of the ordered filtered-limit pipeline, run after the caller has /// reserved rows against the limit in split order (see [`filter_split`]). -pub fn project_split(ctx: Arc, row_range: Range, mask: Mask) -> TaskFuture { +pub(crate) fn project_split( + ctx: Arc, + row_range: Range, + mask: Mask, +) -> TaskFuture { TaskFuture::deferred_projection(ctx, row_range, mask) } @@ -344,13 +348,13 @@ fn build_filter_mask( /// Information needed to execute a single split task. /// /// Row selection is evaluated before creating a split task so it's not included -pub struct TaskContext { +pub(crate) struct TaskContext { /// The shared filter expression. - pub filter: Option>, + pub(crate) filter: Option>, /// The layout reader. - pub reader: Arc, + pub(crate) reader: Arc, /// The projection expression to apply to gather the scanned rows. - pub projection: Expression, + pub(crate) projection: Expression, } #[cfg(test)] From 39c8c04f5b5b88f153386698b73a22d1b30ca64e Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 15:17:06 +0100 Subject: [PATCH 15/17] clear up tests Signed-off-by: Adam Gutglick --- vortex-datafusion/src/v2/source.rs | 31 -- vortex-layout/src/scan/layout.rs | 192 ++++--------- vortex-layout/src/scan/scan_builder.rs | 382 +++++++------------------ 3 files changed, 167 insertions(+), 438 deletions(-) diff --git a/vortex-datafusion/src/v2/source.rs b/vortex-datafusion/src/v2/source.rs index d6db39066bc..c802a1c0502 100644 --- a/vortex-datafusion/src/v2/source.rs +++ b/vortex-datafusion/src/v2/source.rs @@ -884,37 +884,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn ordered_flatten_cancels_a_pending_drain_when_dropped() -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - let started = Arc::new(Notify::new()); - let (dropped_send, dropped_recv) = oneshot::channel(); - let pending = ArrayStreamAdapter::new( - dtype, - DropNotifyingPendingStream { - started: Arc::clone(&started), - _drop_notifier: DropNotifier(Some(dropped_send)), - }, - ) - .boxed(); - - let streams = futures::stream::iter([Ok::<_, VortexError>(pending)]); - let mut flattened = flatten_scan_streams(streams, true, 1, TokioRuntime::current()); - let mut next = Box::pin(futures::StreamExt::next(&mut flattened)); - tokio::select! { - _ = started.notified() => {} - _ = &mut next => return Err(vortex_err!("pending drain unexpectedly produced a chunk")), - } - drop(next); - drop(flattened); - - let dropped = tokio::time::timeout(Duration::from_secs(1), dropped_recv) - .await - .map_err(|_| vortex_err!("dropping the flattened stream did not cancel its drain"))?; - dropped.map_err(|_| vortex_err!("pending source dropped without notifying the test"))?; - Ok(()) - } - struct PanickingStream; impl Stream for PanickingStream { diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index b018b6bd98c..8d0feb95e37 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -437,12 +437,20 @@ mod tests { use crate::SplitRange; use crate::scan::test::session_with_handle; + /// A configurable [`LayoutReader`] test double. Splits come from the data source's + /// `with_split_max_row_count`; the filter passes rows through (optionally delaying the first + /// split to exercise concurrent prefetch), and projection masks/ranges plus evaluated filter + /// ranges are recorded for assertions. #[derive(Debug)] struct TestLayoutReader { name: Arc, dtype: DType, row_count: u64, + split_size: Option, + delay_first_filter: bool, projection_masks: Option>>>, + projection_ranges: Option>>>>, + filter_ranges: Arc>>>, } impl TestLayoutReader { @@ -451,110 +459,45 @@ mod tests { name: Arc::from("test"), dtype: DType::Primitive(PType::I32, Nullability::NonNullable), row_count, + split_size: None, + delay_first_filter: false, projection_masks: None, + projection_ranges: None, + filter_ranges: Arc::new(Mutex::new(Vec::new())), } } - fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { - self.projection_masks = Some(projection_masks); + fn with_split_size(mut self, split_size: u64) -> Self { + self.split_size = Some(split_size); self } - } - impl LayoutReader for TestLayoutReader { - fn name(&self) -> &Arc { - &self.name - } - - fn as_any(&self) -> &dyn std::any::Any { + fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = Some(projection_masks); self } - fn dtype(&self) -> &DType { - &self.dtype - } - - fn row_count(&self) -> u64 { - self.row_count - } - - fn register_splits( - &self, - _field_mask: &[FieldMask], - split_range: &SplitRange, - splits: &mut RowSplits, - ) -> VortexResult<()> { - splits.push(split_range.root_row_range().end); - Ok(()) - } - - fn pruning_evaluation( - &self, - _row_range: &Range, - _expr: &Expression, - mask: Mask, - ) -> VortexResult { - Ok(MaskFuture::ready(mask)) - } - - fn filter_evaluation( - &self, - _row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - Ok(mask) - } - - fn projection_evaluation( - &self, - row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - let row_range = row_range.clone(); - let projection_masks = self.projection_masks.clone(); - - Ok(Box::pin(async move { - let mask = mask.await?; - if let Some(projection_masks) = projection_masks { - projection_masks.lock().push(mask.true_count()); - } - let start = i32::try_from(row_range.start)?; - let end = i32::try_from(row_range.end)?; - PrimitiveArray::from_iter(start..end) - .into_array() - .filter(mask) - })) + fn with_projection_ranges( + mut self, + projection_ranges: Arc>>>, + ) -> Self { + self.projection_ranges = Some(projection_ranges); + self } - } - #[derive(Debug)] - struct DelayedFirstSplitReader { - name: Arc, - dtype: DType, - projection_ranges: Arc>>>, - filter_ranges: Arc>>>, - } - - impl DelayedFirstSplitReader { - fn new(projection_ranges: Arc>>>) -> Self { - Self { - name: Arc::from("delayed-first-split"), - dtype: DType::Primitive(PType::I32, Nullability::NonNullable), - projection_ranges, - filter_ranges: Arc::new(Mutex::new(Vec::new())), - } + fn with_delayed_first_filter(mut self) -> Self { + self.delay_first_filter = true; + self } - /// Row ranges whose filter was actually evaluated (recorded when the split's filter - /// future runs, not when it is merely scheduled). + /// Row ranges whose filter was actually evaluated (recorded when the filter future runs, + /// not when it is merely scheduled). fn filter_ranges(&self) -> Arc>>> { Arc::clone(&self.filter_ranges) } } - impl LayoutReader for DelayedFirstSplitReader { + impl LayoutReader for TestLayoutReader { fn name(&self) -> &Arc { &self.name } @@ -568,7 +511,7 @@ mod tests { } fn row_count(&self) -> u64 { - 4 + self.row_count } fn register_splits( @@ -578,7 +521,13 @@ mod tests { splits: &mut RowSplits, ) -> VortexResult<()> { let row_range = split_range.row_range(); - splits.push(split_range.row_offset() + row_range.start + 2); + if let Some(size) = self.split_size { + let mut boundary = row_range.start + size; + while boundary < row_range.end { + splits.push(split_range.row_offset() + boundary); + boundary += size; + } + } splits.push(split_range.root_row_range().end); Ok(()) } @@ -599,11 +548,9 @@ mod tests { mask: MaskFuture, ) -> VortexResult { self.filter_ranges.lock().push(row_range.clone()); - let delay = row_range.start == 0; - let len = mask.len(); - - Ok(MaskFuture::new(len, async move { - if delay { + if self.delay_first_filter && row_range.start == 0 { + let len = mask.len(); + return Ok(MaskFuture::new(len, async move { let mut yielded = false; futures::future::poll_fn(move |cx| { if yielded { @@ -615,9 +562,10 @@ mod tests { } }) .await; - } - mask.await - })) + mask.await + })); + } + Ok(mask) } fn projection_evaluation( @@ -627,15 +575,22 @@ mod tests { mask: MaskFuture, ) -> VortexResult { let row_range = row_range.clone(); - let projection_ranges = Arc::clone(&self.projection_ranges); + let projection_masks = self.projection_masks.clone(); + let projection_ranges = self.projection_ranges.clone(); Ok(Box::pin(async move { - projection_ranges.lock().push(row_range.clone()); + let mask = mask.await?; + if let Some(projection_masks) = projection_masks { + projection_masks.lock().push(mask.true_count()); + } + if let Some(projection_ranges) = projection_ranges { + projection_ranges.lock().push(row_range.clone()); + } let start = i32::try_from(row_range.start)?; let end = i32::try_from(row_range.end)?; PrimitiveArray::from_iter(start..end) .into_array() - .filter(mask.await?) + .filter(mask) })) } } @@ -669,48 +624,15 @@ mod tests { Ok(()) } - #[test] - fn ordered_filtered_limit_waits_for_the_earlier_split() -> VortexResult<()> { - let runtime = SingleThreadRuntime::default(); - let session = session_with_handle(runtime.handle()); - let projection_ranges = Arc::new(Mutex::new(Vec::new())); - let source = LayoutReaderDataSource::new( - Arc::new(DelayedFirstSplitReader::new(Arc::clone(&projection_ranges))), - session, - ) - .with_split_max_row_count(2); - - let scan = runtime.block_on(source.scan(ScanRequest { - filter: Some(root()), - limit: Some(1), - ordered: true, - ..Default::default() - }))?; - let partitions = runtime.block_on(scan.partitions().try_collect::>())?; - assert_eq!(partitions.len(), 1); - - let mut ctx = array_session().create_execution_ctx(); - let mut values = Vec::new(); - for partition in partitions { - for chunk in runtime.block_on_stream(partition.execute()?) { - let primitive = chunk?.execute::(&mut ctx)?; - values.extend(primitive.into_buffer::()); - } - } - - assert_eq!(values, [0]); - let projection_ranges = projection_ranges.lock(); - assert_eq!(projection_ranges.len(), 1); - assert_eq!(projection_ranges[0], 0..2); - Ok(()) - } - #[test] fn ordered_filtered_limit_evaluates_later_split_filter_concurrently() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let projection_ranges = Arc::new(Mutex::new(Vec::new())); - let reader = DelayedFirstSplitReader::new(Arc::clone(&projection_ranges)); + let reader = TestLayoutReader::new(4) + .with_split_size(2) + .with_projection_ranges(Arc::clone(&projection_ranges)) + .with_delayed_first_filter(); let filter_ranges = reader.filter_ranges(); let source = LayoutReaderDataSource::new(Arc::new(reader), session).with_split_max_row_count(2); diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index a90274befa8..6b952a601c5 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -409,6 +409,7 @@ mod test { use futures::channel::oneshot; use futures::task::noop_waker_ref; use parking_lot::Mutex; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; @@ -786,132 +787,71 @@ mod test { Ok(()) } - #[derive(Debug)] - struct FilteringLayoutReader { + /// A configurable [`LayoutReader`] test double that replaces several near-identical mocks. + /// + /// `split_size` controls the split layout (`None` is a single split), `keep_row` filters rows, + /// and the `fail_*` flags inject filter/projection errors. Every projection records its mask's + /// true-count into `projection_masks`, letting tests assert the limit is applied before + /// projection. + struct MockLayoutReader { name: Arc, dtype: DType, row_count: u64, + split_size: Option, keep_row: fn(u64) -> bool, + fail_first_filter: bool, + fail_first_projection: bool, + fail_projection: bool, + projection_masks: Arc>>, } - impl FilteringLayoutReader { - fn new(row_count: u64, keep_row: fn(u64) -> bool) -> Self { + impl MockLayoutReader { + fn new(row_count: u64) -> Self { Self { - name: Arc::from("filtering"), + name: Arc::from("mock"), dtype: DType::Primitive(PType::I32, Nullability::NonNullable), row_count, - keep_row, - } - } - } - - impl LayoutReader for FilteringLayoutReader { - fn name(&self) -> &Arc { - &self.name - } - - fn dtype(&self) -> &DType { - &self.dtype - } - - fn row_count(&self) -> u64 { - self.row_count - } - - fn register_splits( - &self, - _field_mask: &[FieldMask], - split_range: &SplitRange, - splits: &mut RowSplits, - ) -> VortexResult<()> { - let row_range = split_range.row_range(); - for split in ((row_range.start + 2)..row_range.end).step_by(2) { - splits.push(split_range.row_offset() + split); + split_size: None, + keep_row: |_| true, + fail_first_filter: false, + fail_first_projection: false, + fail_projection: false, + projection_masks: Arc::new(Mutex::new(Vec::new())), } - splits.push(split_range.root_row_range().end); - Ok(()) } - fn pruning_evaluation( - &self, - _row_range: &Range, - _expr: &Expression, - mask: Mask, - ) -> VortexResult { - Ok(MaskFuture::ready(mask)) + fn with_split_size(mut self, split_size: u64) -> Self { + self.split_size = Some(split_size); + self } - fn filter_evaluation( - &self, - row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - let row_range = row_range.clone(); - let keep_row = self.keep_row; - let row_count = usize::try_from(row_range.end - row_range.start) - .map_err(|_| vortex_err!("row range must fit in usize"))?; - - Ok(MaskFuture::new(row_count, async move { - let input_mask = mask.await?; - let filtered = (row_range.start..row_range.end) - .enumerate() - .map(|(idx, row)| input_mask.value(idx) && keep_row(row)); - Ok(Mask::from_iter(filtered)) - })) + fn with_keep_row(mut self, keep_row: fn(u64) -> bool) -> Self { + self.keep_row = keep_row; + self } - fn projection_evaluation( - &self, - row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - let row_range = row_range.clone(); - - Ok(Box::pin(async move { - let start = i32::try_from(row_range.start) - .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; - let end = i32::try_from(row_range.end) - .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; - - let array = PrimitiveArray::from_iter(start..end).into_array(); - array.filter(mask.await?) - })) + fn with_projection_masks(mut self, projection_masks: Arc>>) -> Self { + self.projection_masks = projection_masks; + self } - fn as_any(&self) -> &dyn std::any::Any { + fn with_fail_first_filter(mut self) -> Self { + self.fail_first_filter = true; self } - } - - #[derive(Debug)] - struct ProjectionMaskLayoutReader { - name: Arc, - dtype: DType, - row_count: u64, - projection_masks: Arc>>, - projection_error: bool, - } - impl ProjectionMaskLayoutReader { - fn new(row_count: u64, projection_masks: Arc>>) -> Self { - Self { - name: Arc::from("projection-mask"), - dtype: DType::Primitive(PType::I32, Nullability::NonNullable), - row_count, - projection_masks, - projection_error: false, - } + fn with_fail_first_projection(mut self) -> Self { + self.fail_first_projection = true; + self } fn with_projection_error(mut self) -> Self { - self.projection_error = true; + self.fail_projection = true; self } } - impl LayoutReader for ProjectionMaskLayoutReader { + impl LayoutReader for MockLayoutReader { fn name(&self) -> &Arc { &self.name } @@ -930,105 +870,14 @@ mod test { split_range: &SplitRange, splits: &mut RowSplits, ) -> VortexResult<()> { - splits.push(split_range.root_row_range().end); - Ok(()) - } - - fn pruning_evaluation( - &self, - _row_range: &Range, - _expr: &Expression, - mask: Mask, - ) -> VortexResult { - Ok(MaskFuture::ready(mask)) - } - - fn filter_evaluation( - &self, - _row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - Ok(mask) - } - - fn projection_evaluation( - &self, - row_range: &Range, - _expr: &Expression, - mask: MaskFuture, - ) -> VortexResult { - let row_range = row_range.clone(); - let projection_masks = Arc::clone(&self.projection_masks); - let projection_error = self.projection_error; - - Ok(Box::pin(async move { - let mask = mask.await?; - projection_masks.lock().push(mask.true_count()); - if projection_error { - return Err(vortex_err!("projection failed")); + let row_range = split_range.row_range(); + if let Some(size) = self.split_size { + let mut boundary = row_range.start + size; + while boundary < row_range.end { + splits.push(split_range.row_offset() + boundary); + boundary += size; } - let start = i32::try_from(row_range.start) - .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; - let end = i32::try_from(row_range.end) - .map_err(|_| vortex_err!("row_range.end must fit in i32"))?; - PrimitiveArray::from_iter(start..end) - .into_array() - .filter(mask) - })) - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - } - - #[derive(Debug, Clone, Copy)] - enum FirstSplitFailure { - Filter, - Projection, - } - - #[derive(Debug)] - struct FirstSplitFailureLayoutReader { - name: Arc, - dtype: DType, - projection_masks: Arc>>, - failure: FirstSplitFailure, - } - - impl FirstSplitFailureLayoutReader { - fn new(projection_masks: Arc>>, failure: FirstSplitFailure) -> Self { - Self { - name: Arc::from("first-split-failure"), - dtype: DType::Primitive(PType::I32, Nullability::NonNullable), - projection_masks, - failure, } - } - } - - impl LayoutReader for FirstSplitFailureLayoutReader { - fn name(&self) -> &Arc { - &self.name - } - - fn dtype(&self) -> &DType { - &self.dtype - } - - fn row_count(&self) -> u64 { - 2 - } - - fn register_splits( - &self, - _field_mask: &[FieldMask], - split_range: &SplitRange, - splits: &mut RowSplits, - ) -> VortexResult<()> { - let row_range = split_range.row_range(); - splits.push(split_range.row_offset() + row_range.start + 1); splits.push(split_range.root_row_range().end); Ok(()) } @@ -1048,13 +897,26 @@ mod test { _expr: &Expression, mask: MaskFuture, ) -> VortexResult { - if row_range.start == 0 && matches!(self.failure, FirstSplitFailure::Filter) { + if self.fail_first_filter && row_range.start == 0 { let len = mask.len(); return Ok(MaskFuture::new(len, async move { Err(vortex_err!("first split filter failed")) })); } - Ok(mask) + + let row_range = row_range.clone(); + let keep_row = self.keep_row; + let row_count = usize::try_from(row_range.end - row_range.start) + .map_err(|_| vortex_err!("row range must fit in usize"))?; + + Ok(MaskFuture::new(row_count, async move { + let input_mask = mask.await?; + Ok(Mask::from_iter( + (row_range.start..row_range.end) + .enumerate() + .map(|(idx, row)| input_mask.value(idx) && keep_row(row)), + )) + })) } fn projection_evaluation( @@ -1065,13 +927,13 @@ mod test { ) -> VortexResult { let row_range = row_range.clone(); let projection_masks = Arc::clone(&self.projection_masks); - let failure = self.failure; + let fail = self.fail_projection || (self.fail_first_projection && row_range.start == 0); Ok(Box::pin(async move { let mask = mask.await?; projection_masks.lock().push(mask.true_count()); - if row_range.start == 0 && matches!(failure, FirstSplitFailure::Projection) { - return Err(vortex_err!("first split projection failed")); + if fail { + return Err(vortex_err!("projection failed")); } let start = i32::try_from(row_range.start) .map_err(|_| vortex_err!("row_range.start must fit in i32"))?; @@ -1224,20 +1086,41 @@ mod test { } } - #[test] - fn into_stream_limits_filtered_results() -> VortexResult<()> { + fn keep_all(_: u64) -> bool { + true + } + + fn keep_odd(row: u64) -> bool { + row % 2 == 1 + } + + #[rstest] + #[case::limit_below_matches(8, keep_all, 3, &[0, 1, 2])] + #[case::limit_zero(8, keep_all, 0, &[])] + #[case::limit_exceeds_matches(8, keep_odd, 100, &[1, 3, 5, 7])] + #[case::empty_input(0, keep_all, 3, &[])] + fn filtered_limit_yields_expected_rows( + #[case] row_count: u64, + #[case] keep_row: fn(u64) -> bool, + #[case] limit: u64, + #[case] expected: &[i32], + ) -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); - let reader = Arc::new(FilteringLayoutReader::new(8, |_| true)); + let reader = Arc::new( + MockLayoutReader::new(row_count) + .with_split_size(2) + .with_keep_row(keep_row), + ); let stream = ScanBuilder::new(session, reader) .with_filter(root()) - .with_limit(3) + .with_limit(limit) .into_stream()?; let values = collect_scan_values(runtime.block_on_stream(stream))?; drain_runtime(&runtime); - assert_eq!(values, [0, 1, 2]); + assert_eq!(values.as_slice(), expected); Ok(()) } @@ -1246,10 +1129,9 @@ mod test { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let projection_masks = Arc::new(Mutex::new(Vec::new())); - let reader = Arc::new(ProjectionMaskLayoutReader::new( - 100_000, - Arc::clone(&projection_masks), - )); + let reader = Arc::new( + MockLayoutReader::new(100_000).with_projection_masks(Arc::clone(&projection_masks)), + ); let stream = ScanBuilder::new(session, reader) .with_filter(root()) @@ -1286,10 +1168,12 @@ mod test { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let projection_masks = Arc::new(Mutex::new(Vec::new())); - let reader = Arc::new(FirstSplitFailureLayoutReader::new( - Arc::clone(&projection_masks), - FirstSplitFailure::Filter, - )); + let reader = Arc::new( + MockLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_filter(), + ); let stream = ScanBuilder::new(session, reader) .with_filter(root()) .with_limit(1) @@ -1316,10 +1200,12 @@ mod test { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); let projection_masks = Arc::new(Mutex::new(Vec::new())); - let reader = Arc::new(FirstSplitFailureLayoutReader::new( - Arc::clone(&projection_masks), - FirstSplitFailure::Projection, - )); + let reader = Arc::new( + MockLayoutReader::new(2) + .with_split_size(1) + .with_projection_masks(Arc::clone(&projection_masks)) + .with_fail_first_projection(), + ); let stream = ScanBuilder::new(session, reader) .with_filter(root()) // A budget of two leaves room for the second matching split. Continuing after the @@ -1340,7 +1226,8 @@ mod test { let session = session_with_handle(runtime.handle()); let projection_masks = Arc::new(Mutex::new(Vec::new())); let reader = Arc::new( - ProjectionMaskLayoutReader::new(1, Arc::clone(&projection_masks)) + MockLayoutReader::new(1) + .with_projection_masks(Arc::clone(&projection_masks)) .with_projection_error(), ); let stream = ScanBuilder::new(session, reader) @@ -1359,7 +1246,11 @@ mod test { fn prepared_scan_limits_filtered_results() -> VortexResult<()> { let runtime = SingleThreadRuntime::default(); let session = session_with_handle(runtime.handle()); - let reader = Arc::new(FilteringLayoutReader::new(8, |row| row % 2 == 1)); + let reader = Arc::new( + MockLayoutReader::new(8) + .with_split_size(2) + .with_keep_row(keep_odd), + ); let scan = ScanBuilder::new(session, reader) .with_filter(root()) @@ -1372,59 +1263,6 @@ mod test { Ok(()) } - #[test] - fn filtered_limit_zero_produces_no_rows() -> VortexResult<()> { - let runtime = SingleThreadRuntime::default(); - let session = session_with_handle(runtime.handle()); - let reader = Arc::new(FilteringLayoutReader::new(8, |_| true)); - - let stream = ScanBuilder::new(session, reader) - .with_filter(root()) - .with_limit(0) - .into_stream()?; - let values = collect_scan_values(runtime.block_on_stream(stream))?; - drain_runtime(&runtime); - - assert!(values.is_empty()); - Ok(()) - } - - #[test] - fn filtered_limit_exceeding_matches_returns_all_matches() -> VortexResult<()> { - let runtime = SingleThreadRuntime::default(); - let session = session_with_handle(runtime.handle()); - // Only odd rows match, so there are four matching rows (1, 3, 5, 7). - let reader = Arc::new(FilteringLayoutReader::new(8, |row| row % 2 == 1)); - - let stream = ScanBuilder::new(session, reader) - .with_filter(root()) - .with_limit(100) - .into_stream()?; - let values = collect_scan_values(runtime.block_on_stream(stream))?; - drain_runtime(&runtime); - - // The scan terminates cleanly at end of input rather than hanging on the unfilled budget. - assert_eq!(values, [1, 3, 5, 7]); - Ok(()) - } - - #[test] - fn filtered_limit_over_empty_input_produces_no_rows() -> VortexResult<()> { - let runtime = SingleThreadRuntime::default(); - let session = session_with_handle(runtime.handle()); - let reader = Arc::new(FilteringLayoutReader::new(0, |_| true)); - - let stream = ScanBuilder::new(session, reader) - .with_filter(root()) - .with_limit(3) - .into_stream()?; - let values = collect_scan_values(runtime.block_on_stream(stream))?; - drain_runtime(&runtime); - - assert!(values.is_empty()); - Ok(()) - } - #[derive(Debug)] struct BlockingSplitsLayoutReader { name: Arc, From f2a82f41114fbc545369582b589b6eb2b497efda Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 15:50:22 +0100 Subject: [PATCH 16/17] save some ptrs Signed-off-by: Adam Gutglick --- vortex-layout/src/scan/tasks.rs | 65 ++++++++++++++------------------- 1 file changed, 27 insertions(+), 38 deletions(-) diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 034a80eed37..7a987555291 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -49,21 +49,6 @@ pub(crate) struct TaskFuture { inner: BoxFuture<'static, TaskResult>, } -#[derive(Clone, Copy)] -enum ErrorDisposition { - Recoverable, - Terminal, -} - -impl ErrorDisposition { - fn result(self, error: VortexError) -> TaskResult { - match self { - Self::Recoverable => TaskResult::Recoverable(error), - Self::Terminal => TaskResult::Terminal(error), - } - } -} - impl TaskFuture { fn new(future: impl Future + Send + 'static) -> Self { Self { @@ -87,11 +72,12 @@ impl TaskFuture { Self::ready(TaskResult::Terminal(error)) } - fn projection(projection: ArrayFuture, errors: ErrorDisposition) -> Self { + fn projection(projection: ArrayFuture, terminal: bool) -> Self { Self::new(async move { match projection.await { Ok(array) => TaskResult::Array(Some(array)), - Err(error) => errors.result(error), + Err(error) if terminal => TaskResult::Terminal(error), + Err(error) => TaskResult::Recoverable(error), } }) } @@ -132,25 +118,33 @@ impl TaskFuture { return TaskResult::Array(None); } - Self::deferred_projection(ctx, row_range, mask).await + run_deferred_projection(ctx, row_range, mask).await }) } fn deferred_projection(ctx: Arc, row_range: Range, mask: Mask) -> Self { - Self::new(async move { - let projection = match ctx.reader.projection_evaluation( - &row_range, - &ctx.projection, - MaskFuture::ready(mask), - ) { - Ok(projection) => projection, - Err(error) => return TaskResult::Terminal(error), - }; - match projection.await { - Ok(array) => TaskResult::Array(Some(array)), - Err(error) => TaskResult::Terminal(error), - } - }) + Self::new(run_deferred_projection(ctx, row_range, mask)) + } +} + +/// Construct and run the projection for an already-reserved mask, classifying any failure as +/// terminal (rows have been reserved and cannot be released back to a concurrent limit). +async fn run_deferred_projection( + ctx: Arc, + row_range: Range, + mask: Mask, +) -> TaskResult { + let projection = + match ctx + .reader + .projection_evaluation(&row_range, &ctx.projection, MaskFuture::ready(mask)) + { + Ok(projection) => projection, + Err(error) => return TaskResult::Terminal(error), + }; + match projection.await { + Ok(array) => TaskResult::Array(Some(array)), + Err(error) => TaskResult::Terminal(error), } } @@ -206,12 +200,7 @@ pub(crate) fn split_exec( Err(err) if limited => return Ok(TaskFuture::terminal(err)), Err(err) => return Err(err), }; - let errors = if limited { - ErrorDisposition::Terminal - } else { - ErrorDisposition::Recoverable - }; - return Ok(TaskFuture::projection(projection, errors)); + return Ok(TaskFuture::projection(projection, limited)); }; let filter_mask = build_filter_mask(&ctx.reader, filter, &row_range, row_mask); From bf8ba8a8b1f195959743d9a28a2e5027c6ac37c2 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 17 Jul 2026 16:39:33 +0100 Subject: [PATCH 17/17] try this out Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/opener.rs | 32 +++++++++------------- 1 file changed, 13 insertions(+), 19 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 3988b036b2b..f142bf486cd 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -46,7 +46,6 @@ use vortex::error::VortexError; use vortex::error::VortexExpect; 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::layout::scan::split_by::SplitBy; @@ -429,9 +428,7 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } - let stream_target_field = - Arc::new(Field::new_struct("", stream_schema.fields().clone(), false)); - let handle = session.handle(); + let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); let file_location = file.object_meta.location.clone(); let stream = scan_builder .with_metrics_registry(metrics_registry) @@ -440,24 +437,21 @@ impl FileOpener for VortexOpener { .with_ordered(has_output_ordering) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? + // Convert to Arrow inline on the polling thread: DataFusion sources are expected + // to do their CPU work inside `poll_next`, and spawning this onto the blocking + // pool oversubscribes the CPU. .map(move |chunk| { - let session = session.clone(); - let stream_target_field = Arc::clone(&stream_target_field); - let handle = handle.clone(); - handle.spawn_blocking(move || { - let mut ctx = session.create_execution_ctx(); - chunk.and_then(|chunk| { - let arrow_session = ctx.session().clone(); - let arrow = arrow_session.arrow().execute_arrow( - chunk, - Some(stream_target_field.as_ref()), - &mut ctx, - )?; - Ok(RecordBatch::from(arrow.as_struct().clone())) - }) + let mut ctx = session.create_execution_ctx(); + chunk.and_then(|chunk| { + let arrow_session = ctx.session().clone(); + let arrow = arrow_session.arrow().execute_arrow( + chunk, + Some(&stream_target_field), + &mut ctx, + )?; + Ok(RecordBatch::from(arrow.as_struct().clone())) }) }) - .buffered(2) .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) .map(move |batch| { let batch = if projector.projection().as_ref().is_empty() {