Apply top-level limit in VortexSource

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<u64>,
+    pub limit: Option<usize>,
     /// A metrics object for tracking performance of the scan.
     pub metrics_registry: Arc<dyn MetricsRegistry>,
     /// 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<dyn ObjectStore>;
+        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::<Vec<_>>()
+            .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;

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),

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<FilePruner>,
+    remaining: Option<usize>,
     stream: BoxStream<'static, DFResult<RecordBatch>>,
 }
 
 impl PrunableStream {
-    pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult<RecordBatch>>) -> Self {
+    pub fn new(
+        file_pruner: Option<FilePruner>,
+        limit: Option<usize>,
+        stream: BoxStream<'static, DFResult<RecordBatch>>,
+    ) -> Self {
         Self {
             file_pruner,
+            remaining: limit,
             stream,
         }
     }
@@ -33,10 +40,28 @@ impl Stream for PrunableStream {
     type Item = DFResult<RecordBatch>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        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,
+            }
         }
     }
 }