Shared limit pushdown

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-datafusion/src/persistent/opener.rs

@@ -397,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);
             }
 
@@ -783,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-layout/src/scan/layout.rs

@@ -40,6 +40,8 @@ use vortex_scan::selection::Selection;
 use vortex_session::VortexSession;
 
 use crate::LayoutReaderRef;
+use crate::scan::limit::SharedRowLimit;
+use crate::scan::limit::limit_array_stream_shared;
 use crate::scan::scan_builder::ScanBuilder;
 
 /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`].
@@ -157,13 +159,16 @@ 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(),
             dtype,
             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 +190,7 @@ struct LayoutReaderScan {
     projection: Expression,
     filter: Option<Expression>,
     limit: Option<u64>,
+    shared_limit: Option<SharedRowLimit>,
     ordered: bool,
     selection: Selection,
     metrics_registry: Option<Arc<dyn MetricsRegistry>>,
@@ -251,6 +257,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 +285,7 @@ struct LayoutReaderSplit {
     projection: Expression,
     filter: Option<Expression>,
     limit: Option<u64>,
+    shared_limit: Option<SharedRowLimit>,
     ordered: bool,
     row_range: Range<u64>,
     selection: Selection,
@@ -312,6 +320,7 @@ impl Partition for LayoutReaderSplit {
     }
 
     fn execute(self: Box<Self>) -> VortexResult<SendableArrayStream> {
+        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 +333,7 @@ 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 = limit_array_stream_shared(builder.into_stream()?, shared_limit);
 
         Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new(
             dtype, stream,
@@ -391,3 +400,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<str>,
+        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<str> {
+            &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<u64>,
+            _expr: &Expression,
+            mask: Mask,
+        ) -> VortexResult<MaskFuture> {
+            Ok(MaskFuture::ready(mask))
+        }
+
+        fn filter_evaluation(
+            &self,
+            _row_range: &Range<u64>,
+            _expr: &Expression,
+            mask: MaskFuture,
+        ) -> VortexResult<MaskFuture> {
+            Ok(mask)
+        }
+
+        fn projection_evaluation(
+            &self,
+            row_range: &Range<u64>,
+            _expr: &Expression,
+            mask: MaskFuture,
+        ) -> VortexResult<ArrayFuture> {
+            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::<Vec<_>>())?;
+
+        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::<PrimitiveArray>(&mut ctx)?;
+                values.extend(primitive.into_buffer::<i32>());
+            }
+        }
+
+        assert_eq!(values, [0, 1, 2]);
+        Ok(())
+    }
+}

vortex-layout/src/scan/limit.rs

@@ -2,6 +2,9 @@
 // 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;
 
@@ -25,6 +28,45 @@ where
     }
 }
 
+/// A row limit shared by streams that execute independent scan partitions.
+#[derive(Clone)]
+pub(crate) struct SharedRowLimit(Arc<AtomicU64>);
+
+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,
+            }
+        }
+    }
+}
+
+pub(crate) fn limit_array_stream_shared<S>(
+    stream: S,
+    limit: Option<SharedRowLimit>,
+) -> BoxStream<'static, VortexResult<ArrayRef>>
+where
+    S: Stream<Item = VortexResult<ArrayRef>> + Send + 'static,
+{
+    match limit {
+        Some(limit) => SharedRowLimitedStream::new(stream.boxed(), limit).boxed(),
+        None => stream.boxed(),
+    }
+}
+
 struct RowLimitedStream {
     inner: BoxStream<'static, VortexResult<ArrayRef>>,
     remaining: u64,
@@ -72,3 +114,54 @@ impl Stream for RowLimitedStream {
         }
     }
 }
+
+struct SharedRowLimitedStream {
+    inner: BoxStream<'static, VortexResult<ArrayRef>>,
+    limit: SharedRowLimit,
+}
+
+impl SharedRowLimitedStream {
+    fn new(inner: BoxStream<'static, VortexResult<ArrayRef>>, limit: SharedRowLimit) -> Self {
+        Self { inner, limit }
+    }
+
+    fn abort_pending(&mut self) {
+        let inner = std::mem::replace(&mut self.inner, stream::empty().boxed());
+        drop(inner);
+    }
+}
+
+impl Stream for SharedRowLimitedStream {
+    type Item = VortexResult<ArrayRef>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.inner.as_mut().poll_next(cx) {
+            Poll::Ready(Some(Ok(chunk))) => {
+                let chunk_len = chunk.len() as u64;
+                let (reserved, exhausted) = self.limit.reserve(chunk_len);
+
+                if exhausted {
+                    self.abort_pending();
+                }
+
+                if reserved == 0 {
+                    if exhausted {
+                        Poll::Ready(None)
+                    } else {
+                        Poll::Ready(Some(Ok(chunk)))
+                    }
+                } else if reserved == chunk_len {
+                    Poll::Ready(Some(Ok(chunk)))
+                } else {
+                    let limit = match usize::try_from(reserved) {
+                        Ok(limit) => limit,
+                        Err(_) => unreachable!("reserved rows cannot exceed the current chunk"),
+                    };
+                    self.abort_pending();
+                    Poll::Ready(Some(chunk.slice(0..limit)))
+                }
+            }
+            other => other,
+        }
+    }
+}