Adamg/remove map fn limit

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::<Canonical>(&mut ctx)?;
-                        Ok(canonical.into_array())
+                        result.and_then(|a| {
+                            let canonical = a.execute::<Canonical>(&mut ctx)?;
+                            Ok(canonical.into_array())
+                        })
                     }
                 })
-                .into_array_stream()?
                 .try_collect()
                 .await?;
             chunks.extend(file_chunks);

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

@@ -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,61 +397,47 @@ 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);
             }
 
             if let Some(concurrency) = scan_concurrency {
                 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<u64>, 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<RecordBatch> {
+    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<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());
+        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::<Vec<_>>().await?;
+        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
+
+        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;
@@ -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<ScanBuilder<ArrayRef>> {
+    pub fn scan(&self) -> VortexResult<ScanBuilder> {
         Ok(ScanBuilder::new(
             self.session.clone(),
             self.layout_reader()?,

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<ArrayRef>) -> VortexResult<ScanBuilder<ArrayRef>>,
+    f: impl Fn(ScanBuilder) -> VortexResult<ScanBuilder>,
 ) -> VortexResult<ArrayRef> {
     let mut writer = vec![];
     SESSION

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<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 +36,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 })
@@ -53,11 +54,11 @@ impl ScanBuilder<ArrayRef> {
         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)