diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index b1c999656f6..c77265614c0 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -86,7 +86,7 @@ pub(crate) struct VortexOpener { /// A hint for the desired row count of record batches returned from the scan. pub batch_size: usize, /// If provided, the scan will not return more than this many rows. - pub limit: Option, + pub limit: Option, /// A metrics object for tracking performance of the scan. pub metrics_registry: Arc, /// A shared cache of file readers. @@ -399,6 +399,11 @@ impl FileOpener for VortexOpener { if let Some(limit) = limit && filter.is_none() { + // ScanBuilder cannot combine a filter and a limit. PrunableStream applies every + // limit after filtering; this pushdown avoids reading unnecessary rows when no + // filter is present. + let limit = u64::try_from(limit) + .map_err(|_| exec_datafusion_err!("Vortex scan limit exceeds u64"))?; scan_builder = scan_builder.with_limit(limit); } @@ -461,8 +466,8 @@ impl FileOpener for VortexOpener { }) .boxed(); - if let Some(file_pruner) = file_pruner { - Ok(PrunableStream::new(file_pruner, stream).boxed()) + if file_pruner.is_some() || limit.is_some() { + Ok(PrunableStream::new(file_pruner, limit, stream).boxed()) } else { Ok(stream) } @@ -766,6 +771,42 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_applies_limit_after_filter() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "filtered/file.vortex"; + let batch = record_batch!(("a", Int32, vec![0, 1, 2, 3, 4]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let table_schema = TableSchema::from_file_schema(batch.schema()); + let filter = logical2physical(&col("a").gt(lit(1)), table_schema.table_schema()); + + let mut opener = make_opener(object_store, table_schema, Some(filter)); + opener.limit = Some(2); + + let batches = opener + .open(PartitionedFile::new(file_path.to_string(), data_size))? + .await? + .try_collect::>() + .await?; + + assert_snapshot!(pretty_format_batches_with_options( + &batches, + &FormatOptions::new().with_types_info(true), + )? + .to_string(), @r" + +-------+ + | a | + | Int32 | + +-------+ + | 2 | + | 3 | + +-------+ + "); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 469759527f6..47070cfc2bf 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -359,7 +359,7 @@ impl VortexSource { expr_adapter_factory, table_schema: self.table_schema.clone(), batch_size, - limit: base_config.limit.map(|l| l as u64), + limit: base_config.limit, metrics_registry: Arc::clone(&self.vx_metrics_registry), layout_readers: Arc::clone(&self.layout_readers), natural_split_ranges: Arc::clone(&self.natural_split_ranges), diff --git a/vortex-datafusion/src/persistent/stream.rs b/vortex-datafusion/src/persistent/stream.rs index af2fbc8693e..9a40c2fd20e 100644 --- a/vortex-datafusion/src/persistent/stream.rs +++ b/vortex-datafusion/src/persistent/stream.rs @@ -12,18 +12,25 @@ use futures::Stream; use futures::StreamExt; use futures::stream::BoxStream; -/// Utility to end a stream early if its backing [`PartitionedFile`] can be pruned away by an updated dynamic expression. +/// Utility to end a stream early when it reaches its limit or its backing +/// [`PartitionedFile`] can be pruned by an updated dynamic expression. /// /// [`PartitionedFile`]: datafusion_datasource::PartitionedFile pub(crate) struct PrunableStream { - file_pruner: FilePruner, + file_pruner: Option, + remaining: Option, stream: BoxStream<'static, DFResult>, } impl PrunableStream { - pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult>) -> Self { + pub fn new( + file_pruner: Option, + limit: Option, + stream: BoxStream<'static, DFResult>, + ) -> Self { Self { file_pruner, + remaining: limit, stream, } } @@ -33,10 +40,28 @@ impl Stream for PrunableStream { type Item = DFResult; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - if self.as_mut().file_pruner.should_prune()? { + if self.remaining == Some(0) { + Poll::Ready(None) + } else if let Some(file_pruner) = self.file_pruner.as_mut() + && file_pruner.should_prune()? + { Poll::Ready(None) } else { - self.stream.poll_next_unpin(cx) + match self.stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => match &mut self.remaining { + Some(remaining) if batch.num_rows() > *remaining => { + let batch = batch.slice(0, *remaining); + *remaining = 0; + Poll::Ready(Some(Ok(batch))) + } + Some(remaining) => { + *remaining -= batch.num_rows(); + Poll::Ready(Some(Ok(batch))) + } + None => Poll::Ready(Some(Ok(batch))), + }, + poll => poll, + } } } }