Adamg/remove map fn limit

benchmarks/string-bench/src/serialized.rs

@@ -7,10 +7,8 @@
 //! * **write** runs the full default write pipeline — repartition into row
 //!   blocks, zoned statistics, dictionary probe, coalesce, compress with the
 //!   forced string scheme and its children, layout, serialize into a buffer;
-//! * **read** opens that buffer and runs the scan, decoding each row split to
-//!   canonical `VarBinViewArray` form inside its own scan task and dropping it
-//!   before the next split runs — the shape production uses, where
-//!   `into_record_batch_stream` fuses the Arrow conversion into the split task.
+//! * **read** opens that buffer and runs the scan, decoding each emitted chunk to
+//!   canonical `VarBinViewArray` form and dropping it before polling the next chunk.
 //!
 //! Both run on a current-thread runtime, so these are single-threaded CPU costs
 //! and exclude physical I/O.
@@ -201,35 +199,20 @@ async fn write_serialized_file(
 }
 
 /// Time one complete read of a serialized Vortex buffer: open the file, then run
-/// the scan with the canonical decode fused into each row split's task, dropping
-/// each decoded chunk before the next split runs.
-///
-/// The splits are awaited one at a time rather than through
-/// `ScanBuilder::into_array_stream`, which spawns
-/// `concurrency * available_parallelism()` of them at once. On a current-thread
-/// runtime that read-ahead buys no parallelism; it only holds that many chunks in
-/// memory and makes the result depend on the host's core count. Awaiting one at a
-/// time keeps the per-split work identical to production while making the
-/// measurement machine-independent.
+/// the scan and decode each emitted chunk before polling the next one.
 async fn read_serialized_buffer(session: &VortexSession, data: Bytes) -> Result<Duration> {
     let decode_session = session.clone();
 
     let start = Instant::now();
     let file = session.open_options().open_buffer(data)?;
-    let splits = file
-        .scan()?
-        .map(move |chunk: ArrayRef| {
-            let mut ctx = decode_session.create_execution_ctx();
-            chunk.execute::<VarBinViewArray>(&mut ctx)
-        })
-        .build()?;
+    let mut chunks = file.scan()?.into_stream()?;
 
     let mut rows = 0usize;
-    for split in splits {
-        if let Some(canonical) = split.await? {
-            rows += canonical.len();
-            drop(black_box(canonical));
-        }
+    while let Some(chunk) = chunks.try_next().await? {
+        let mut ctx = decode_session.create_execution_ctx();
+        let canonical = chunk.execute::<VarBinViewArray>(&mut ctx)?;
+        rows += canonical.len();
+        drop(black_box(canonical));
     }
 
     black_box(rows);

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 vortex::array::ArrayRef;
 use vortex::array::Canonical;
@@ -72,15 +73,17 @@ impl Dataset for TPCHLCommentChunked {
         let chunks: Vec<_> = file
             .scan()?
             .with_projection(projection)
+            .into_array_stream()?
             .map({
                 let ctx = ctx.clone();
                 move |a| {
                     let mut ctx = ctx.clone();
-                    let canonical = a.execute::<Canonical>(&mut ctx)?;
-                    Ok(canonical.into_array())
+                    a.and_then(|a| {
+                        let canonical = a.execute::<Canonical>(&mut ctx)?;
+                        Ok(canonical.into_array())
+                    })
                 }
             })
-            .into_array_stream()?
             .try_collect()
             .await?;
 

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<A>(&self, mut scan_builder: ScanBuilder<A>) -> ScanBuilder<A>
-    where
-        A: 'static + Send,
-    {
+    pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder {
         let Self { selection } = self;
 
         if let Some(selection) = selection {

vortex-datafusion/src/persistent/opener.rs

@@ -397,9 +397,7 @@ impl FileOpener for VortexOpener {
                 .transpose()
                 .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?;
 
-            if let Some(limit) = limit
-                && filter.is_none()
-            {
+            if let Some(limit) = limit {
                 scan_builder = scan_builder.with_limit(limit);
             }
 
@@ -445,28 +443,29 @@ impl FileOpener for VortexOpener {
             }
 
             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)
                 .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_err(move |e: VortexError| {
-                    DataFusionError::External(Box::new(e.with_context(format!(
-                        "Failed to read Vortex file: {}",
-                        file.object_meta.location
-                    ))))
+                // 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 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()))
+                    })
                 })
-                .map(move |batch| {
+                .map_err(move |e: VortexError| vortex_file_read_error(&file_location, e))
+                .map(move |batch| -> DFResult<RecordBatch> {
                     let batch = if projector.projection().as_ref().is_empty() {
                         batch
                     } else {
@@ -549,10 +548,10 @@ impl NaturalSplits {
 }
 
 /// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use.
-fn natural_splits_for_file<A: 'static + Send>(
+fn natural_splits_for_file(
     natural_splits: &DashMap<Path, Arc<NaturalSplits>>,
     path: &Path,
-    scan_builder: &ScanBuilder<A>,
+    scan_builder: &ScanBuilder,
     total_size: u64,
 ) -> DFResult<Arc<NaturalSplits>> {
     if let Some(splits) = natural_splits.get(path) {
@@ -574,8 +573,8 @@ fn natural_splits_for_file<A: 'static + Send>(
 
 /// Walk the layout tree to compute the file's full natural split boundaries for the fields
 /// referenced by the scan's projection and filter.
-fn compute_natural_splits<A: 'static + Send>(
-    scan_builder: &ScanBuilder<A>,
+fn compute_natural_splits(
+    scan_builder: &ScanBuilder,
     total_size: u64,
 ) -> DFResult<Arc<NaturalSplits>> {
     let row_boundaries = scan_builder
@@ -636,6 +635,12 @@ fn split_midpoint_to_byte(split_range: &Range<u64>, 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;
@@ -668,6 +673,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;
@@ -1086,6 +1092,48 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_open_applies_limit_after_filtering() -> anyhow::Result<()> {
+        let object_store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        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());
+        // `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::<Vec<_>>().await?;
+        let values = batches
+            .iter()
+            .flat_map(|batch| {
+                batch
+                    .column(0)
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .expect("projected column should be Int32")
+                    .values()
+                    .to_vec()
+            })
+            .collect::<Vec<i32>>();
+
+        assert_eq!(values, [4, 5, 6]);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_open_empty_file() -> anyhow::Result<()> {
         use futures::TryStreamExt;

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<RecordBatch>>,
+    stream: Option<BoxStream<'static, DFResult<RecordBatch>>>,
 }
 
 impl PrunableStream {
     pub fn new(file_pruner: FilePruner, stream: BoxStream<'static, DFResult<RecordBatch>>) -> 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<Option<Self::Item>> {
         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),
+            }
         }
     }
 }

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;
@@ -214,7 +213,7 @@ impl VortexFile {
 
     /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and
     /// execution options.
-    pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
+    pub fn scan(&self) -> VortexResult<ScanBuilder> {
         Ok(ScanBuilder::new(
             self.session.clone(),
             self.layout_reader()?,

vortex-file/src/tests.rs

@@ -1328,7 +1328,7 @@ async fn write_nullable_top_level_struct() {
 
 async fn round_trip(
     array: &ArrayRef,
-    f: impl FnOnce(ScanBuilder<ArrayRef>) -> VortexResult<ScanBuilder<ArrayRef>>,
+    f: impl FnOnce(ScanBuilder) -> VortexResult<ScanBuilder>,
 ) -> VortexResult<ArrayRef> {
     let mut writer = vec![];
     SESSION

vortex-layout/src/scan/arrow.rs

@@ -10,17 +10,20 @@ 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;
 use vortex_array::VortexSessionExecute;
 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;
 
-impl ScanBuilder<ArrayRef> {
+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 +38,11 @@ impl ScanBuilder<ArrayRef> {
         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 })
@@ -49,15 +52,29 @@ impl ScanBuilder<ArrayRef> {
         self,
         schema: SchemaRef,
     ) -> VortexResult<impl Stream<Item = Result<RecordBatch, ArrowError>> + 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 = get_available_parallelism().unwrap_or(1);
 
         let stream = self
+            .into_stream()?
             .map(move |chunk| {
-                let mut ctx = session.create_execution_ctx();
-                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
+                }
             })
-            .into_stream()?
+            .buffered(concurrency)
             .map_err(|e| ArrowError::ExternalError(Box::new(e)));
 
         Ok(stream)