diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index ebe6ff2e96a..049837b56df 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; @@ -78,15 +79,17 @@ 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()) + }) } }) - .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 b1c999656f6..7d14c9aa264 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -42,6 +42,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; @@ -396,9 +397,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); } @@ -406,51 +405,39 @@ impl FileOpener for VortexOpener { scan_builder = scan_builder.with_concurrency(concurrency); } - let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false); + 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) .with_some_filter(filter) .with_ordered(has_output_ordering) - .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())) - }) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_ok(move |rb| { - // We try and slice the stream into respecting datafusion's configured batch size. - stream::iter( - (0..rb.num_rows().div_ceil(batch_size * 2)) - .flat_map(move |block_idx| { - let offset = block_idx * batch_size * 2; - - // If we have less than two batches worth of rows left, we keep them together as a single batch. - if rb.num_rows() - offset < 2 * batch_size { - let length = rb.num_rows() - offset; - [Some(rb.slice(offset, length)), None].into_iter() - } else { - let first = rb.slice(offset, batch_size); - let second = rb.slice(offset + batch_size, batch_size); - [Some(first), Some(second)].into_iter() - } - }) - .flatten() - .map(Ok), - ) + .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())) + }) + }) }) - .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_ok(move |rb| { + stream::iter(split_record_batch(rb, batch_size).into_iter().map(Ok)) }) + .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e)) .try_flatten() .map(move |batch| { if projector.projection().as_ref().is_empty() { @@ -562,6 +549,33 @@ 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 split_record_batch(rb: RecordBatch, batch_size: usize) -> Vec { + assert!(batch_size > 0, "batch size must be positive"); + + let mut batches = Vec::new(); + let mut offset = 0; + + while offset < rb.num_rows() { + let remaining = rb.num_rows() - offset; + if remaining < 2 * batch_size { + batches.push(rb.slice(offset, remaining)); + break; + } + + batches.push(rb.slice(offset, batch_size)); + batches.push(rb.slice(offset + batch_size, batch_size)); + offset += batch_size * 2; + } + + batches +} + +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::sync::Arc; @@ -589,6 +603,7 @@ mod tests { use datafusion_expr::Operator; use datafusion_physical_expr::expressions as df_expr; use datafusion_physical_expr::projection::ProjectionExpr; + use futures::TryStreamExt; use insta::assert_snapshot; use itertools::Itertools; use object_store::ObjectStore; @@ -766,6 +781,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-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-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 ea7f8cbeda5..a038cd85b60 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1277,7 +1277,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/layout.rs b/vortex-layout/src/scan/layout.rs index 0b1f150c67e..d2145fab29f 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -40,6 +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::scan_builder::ScanBuilder; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. @@ -157,6 +160,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 +169,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 +191,7 @@ struct LayoutReaderScan { projection: Expression, filter: Option, limit: Option, + shared_limit: Option, ordered: bool, selection: Selection, metrics_registry: Option>, @@ -251,6 +258,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 +286,7 @@ struct LayoutReaderSplit { projection: Expression, filter: Option, limit: Option, + shared_limit: Option, ordered: bool, row_range: Range, selection: Selection, @@ -312,6 +321,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 +334,13 @@ 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 = 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, @@ -391,3 +407,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 new file mode 100644 index 00000000000..855c00e6ed2 --- /dev/null +++ b/vortex-layout/src/scan/limit.rs @@ -0,0 +1,145 @@ +// 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; + +/// A row limit shared by the streams executing one scan's independent partitions. +/// +/// 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); + +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, + } + } + } + + 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/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/multi.rs b/vortex-layout/src/scan/multi.rs index cb516d2500d..d78695884d7 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -57,6 +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::scan_builder::ScanBuilder; /// Default concurrency for opening deferred readers. @@ -300,10 +303,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 +326,7 @@ struct MultiLayoutScan { session: VortexSession, dtype: DType, request: ScanRequest, + shared_limit: Option, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -345,6 +352,7 @@ impl DataSourceScan for MultiLayoutScan { session, dtype: _, request, + shared_limit, ready, deferred, handle, @@ -399,7 +407,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 +430,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 +476,7 @@ fn reader_partition( row_range: Some(row_range), ..request }, + shared_limit, index: partition_idx, }) as PartitionRef) }) @@ -475,6 +491,7 @@ struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, request: ScanRequest, + shared_limit: Option, index: usize, } @@ -510,6 +527,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 +541,13 @@ impl Partition for MultiLayoutPartition { } let dtype = builder.dtype()?; - let stream = builder.into_stream()?; + 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, diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 681f33639bc..186b5c07406 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -6,8 +6,10 @@ 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 futures::stream::BoxStream; use itertools::Either; use itertools::Itertools; use vortex_array::ArrayRef; @@ -20,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; @@ -27,15 +30,18 @@ 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::splits::Splits; use crate::scan::tasks::TaskContext; +use crate::scan::tasks::TaskFuture; use crate::scan::tasks::split_exec; /// 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 +55,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) + /// 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 +85,7 @@ 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 +101,6 @@ impl RepeatedScan { selection: Selection, splits: Splits, concurrency: usize, - map_fn: Arc VortexResult + Send + Sync>, limit: Option, dtype: DType, ) -> Self { @@ -113,16 +114,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 +132,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,66 +151,126 @@ 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; - let mut tasks = Vec::new(); - let ctx = Arc::new(TaskContext { + 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(), - mapper: Arc::clone(&self.map_fn), - }); + }) + } + + pub(crate) fn execute(&self, row_range: Option>) -> VortexResult> { + let ctx = self.task_context(); + + let mut limit = self.limit; + 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; + } - for range in ranges { 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) } - pub fn execute_stream( + pub(crate) fn execute_stream( &self, row_range: Option>, - ) -> VortexResult> + Send + 'static + use> { - use futures::StreamExt; + ) -> VortexResult> + Send + 'static> { let num_workers = get_available_parallelism().unwrap_or(1); let concurrency = self.concurrency * num_workers; let handle = self.session.handle(); - let stream = + // 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 = self.task_context(); + let selection = self.selection.clone(); + 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(schedule(tasks, self.ordered, concurrency, self.limit)); + } + + // 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(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(); - Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) + match limit { + Some(limit) => LimitedStream::new(stream, RowBudget::Local(limit)).boxed(), + None => stream, } } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..76da28aab5f 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; @@ -45,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: /// @@ -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,19 +67,15 @@ 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) + /// 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. 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, } @@ -128,9 +117,7 @@ impl ScanBuilder { runtime.block_on_stream(stream), )) } -} -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 +224,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 +278,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 +294,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 +387,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 +674,167 @@ 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], + 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); + } + 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 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 ctx = array_session().create_execution_ctx(); + let mut values = Vec::new(); + for chunk in iter { + let primitive = chunk?.execute::(&mut ctx)?; + 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 = 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 = 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..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 @@ -34,12 +34,12 @@ pub type TaskFuture = BoxFuture<'static, VortexResult>; /// 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>, +/// 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(); @@ -136,15 +136,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 +151,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>, } 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, }