claude CR stuff

AdamGS · AdamGS · commit 1ca3a3480526 · 2026-07-15T13:47:57.000+01:00
Signed-off-by: Adam Gutglick &lt;adam@spiraldb.com&gt;
diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs
@@ -40,8 +40,9 @@ use vortex_scan::selection::Selection;
 use vortex_session::VortexSession;
 
 use crate::LayoutReaderRef;
+use crate::scan::limit::LimitedStream;
+use crate::scan::limit::RowBudget;
 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`].
@@ -333,7 +334,11 @@ 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 = limit_array_stream_shared(builder.into_stream()?, shared_limit);
+        let stream = builder.into_stream()?.boxed();
+        let stream = match shared_limit {
+            Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(),
+            None => stream,
+        };
 
         Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new(
             dtype, stream,
diff --git a/vortex-layout/src/scan/limit.rs b/vortex-layout/src/scan/limit.rs
@@ -13,22 +13,15 @@ use futures::StreamExt;
 use futures::stream;
 use futures::stream::BoxStream;
 use vortex_array::ArrayRef;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
-pub(crate) fn limit_array_stream<S>(
-    stream: S,
-    limit: Option<u64>,
-) -> BoxStream<'static, VortexResult<ArrayRef>>
-where
-    S: Stream<Item = VortexResult<ArrayRef>> + Send + 'static,
-{
-    match limit {
-        Some(limit) => RowLimitedStream::new(stream.boxed(), limit).boxed(),
-        None => stream.boxed(),
-    }
-}
-
 /// A row limit shared by streams that execute independent scan partitions.
+///
+/// The shared budget gives "at most `limit` rows in total" semantics without any ordering
+/// guarantee across the streams that share it: whichever stream produces a chunk first claims
+/// the budget first. It is therefore only correct for consumers that treat the combined output
+/// as unordered (e.g. a bare `LIMIT n`), not for order-preserving cross-partition consumption.
 #[derive(Clone)]
 pub(crate) struct SharedRowLimit(Arc<AtomicU64>);
 
@@ -52,99 +45,87 @@ impl SharedRowLimit {
             }
         }
     }
-}
 
-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(),
+    fn is_exhausted(&self) -> bool {
+        self.0.load(Ordering::Relaxed) == 0
     }
 }
 
-struct RowLimitedStream {
-    inner: BoxStream<'static, VortexResult<ArrayRef>>,
-    remaining: u64,
-}
-
-impl RowLimitedStream {
-    fn new(inner: BoxStream<'static, VortexResult<ArrayRef>>, remaining: u64) -> Self {
-        Self { inner, remaining }
-    }
-
-    fn abort_pending(&mut self) {
-        let inner = std::mem::replace(&mut self.inner, stream::empty().boxed());
-        drop(inner);
-    }
+/// The remaining rows a [`LimitedStream`] may emit, either privately or shared across streams.
+pub(crate) enum RowBudget {
+    /// A budget private to a single stream.
+    Local(u64),
+    /// A budget shared across streams executing independent scan partitions.
+    Shared(SharedRowLimit),
 }
 
-impl Stream for RowLimitedStream {
-    type Item = VortexResult<ArrayRef>;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        if self.remaining == 0 {
-            return Poll::Ready(None);
+impl RowBudget {
+    /// Reserve up to `requested` rows from the budget.
+    ///
+    /// Returns `(reserved, exhausted)` where `reserved <= requested` and `exhausted` is true
+    /// once the budget has reached zero.
+    fn reserve(&mut self, requested: u64) -> (u64, bool) {
+        match self {
+            RowBudget::Local(remaining) => {
+                let reserved = (*remaining).min(requested);
+                *remaining -= reserved;
+                (reserved, *remaining == 0)
+            }
+            RowBudget::Shared(shared) => shared.reserve(requested),
         }
+    }
 
-        match self.inner.as_mut().poll_next(cx) {
-            Poll::Ready(Some(Ok(chunk))) => {
-                let chunk_len = chunk.len() as u64;
-                if chunk_len <= self.remaining {
-                    self.remaining -= chunk_len;
-                    if self.remaining == 0 {
-                        self.abort_pending();
-                    }
-                    Poll::Ready(Some(Ok(chunk)))
-                } else {
-                    let limit = match usize::try_from(self.remaining) {
-                        Ok(limit) => limit,
-                        Err(_) => unreachable!("remaining rows cannot exceed the current chunk"),
-                    };
-                    self.remaining = 0;
-                    self.abort_pending();
-                    Poll::Ready(Some(chunk.slice(0..limit)))
-                }
-            }
-            other => other,
+    fn is_exhausted(&self) -> bool {
+        match self {
+            RowBudget::Local(remaining) => *remaining == 0,
+            RowBudget::Shared(shared) => shared.is_exhausted(),
         }
     }
 }
 
-struct SharedRowLimitedStream {
+/// Wraps a stream, emitting chunks until its [`RowBudget`] is exhausted, then terminating.
+pub(crate) struct LimitedStream {
     inner: BoxStream<'static, VortexResult<ArrayRef>>,
-    limit: SharedRowLimit,
+    budget: RowBudget,
 }
 
-impl SharedRowLimitedStream {
-    fn new(inner: BoxStream<'static, VortexResult<ArrayRef>>, limit: SharedRowLimit) -> Self {
-        Self { inner, limit }
+impl LimitedStream {
+    pub(crate) fn new(
+        inner: BoxStream<'static, VortexResult<ArrayRef>>,
+        budget: RowBudget,
+    ) -> Self {
+        Self { inner, budget }
     }
 
+    /// Drop the inner stream so no further work (including spawned split tasks) is polled.
     fn abort_pending(&mut self) {
-        let inner = std::mem::replace(&mut self.inner, stream::empty().boxed());
-        drop(inner);
+        self.inner = stream::empty().boxed();
     }
 }
 
-impl Stream for SharedRowLimitedStream {
+impl Stream for LimitedStream {
     type Item = VortexResult<ArrayRef>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        // Avoid reading a chunk we have no budget for. For a shared budget this also stops a
+        // partition whose siblings already consumed the limit.
+        if self.budget.is_exhausted() {
+            self.abort_pending();
+            return Poll::Ready(None);
+        }
+
         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);
+                let (reserved, exhausted) = self.budget.reserve(chunk_len);
 
                 if exhausted {
                     self.abort_pending();
                 }
 
                 if reserved == 0 {
+                    // Either the budget was already exhausted (stop), or this is an empty chunk
+                    // while the budget still has room (pass it through).
                     if exhausted {
                         Poll::Ready(None)
                     } else {
@@ -153,11 +134,8 @@ impl Stream for SharedRowLimitedStream {
                 } 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();
+                    let limit = usize::try_from(reserved)
+                        .vortex_expect("reserved rows are bounded by the chunk length");
                     Poll::Ready(Some(chunk.slice(0..limit)))
                 }
             }
diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs
@@ -57,8 +57,9 @@ use vortex_session::VortexSession;
 use vortex_utils::parallelism::get_available_parallelism;
 
 use crate::LayoutReaderRef;
+use crate::scan::limit::LimitedStream;
+use crate::scan::limit::RowBudget;
 use crate::scan::limit::SharedRowLimit;
-use crate::scan::limit::limit_array_stream_shared;
 use crate::scan::scan_builder::ScanBuilder;
 
 /// Default concurrency for opening deferred readers.
@@ -540,7 +541,11 @@ impl Partition for MultiLayoutPartition {
         }
 
         let dtype = builder.dtype()?;
-        let stream = limit_array_stream_shared(builder.into_stream()?, shared_limit);
+        let stream = builder.into_stream()?.boxed();
+        let stream = match shared_limit {
+            Some(limit) => LimitedStream::new(stream, RowBudget::Shared(limit)).boxed(),
+            None => stream,
+        };
 
         Ok(ArrayStreamExt::boxed(ArrayStreamAdapter::new(
             dtype, stream,
diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs
@@ -9,6 +9,7 @@ use std::sync::Arc;
 use futures::FutureExt;
 use futures::Stream;
 use futures::StreamExt;
+use futures::stream::BoxStream;
 use itertools::Either;
 use itertools::Itertools;
 use vortex_array::ArrayRef;
@@ -21,14 +22,16 @@ use vortex_array::stream::ArrayStreamAdapter;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_io::runtime::BlockingRuntime;
+use vortex_io::runtime::Task;
 use vortex_io::session::RuntimeSessionExt;
 use vortex_scan::selection::Selection;
 use vortex_session::VortexSession;
 use vortex_utils::parallelism::get_available_parallelism;
 
 use crate::LayoutReaderRef;
 use crate::scan::filter::FilterExpr;
-use crate::scan::limit::limit_array_stream;
+use crate::scan::limit::LimitedStream;
+use crate::scan::limit::RowBudget;
 use crate::scan::splits::Splits;
 use crate::scan::tasks::TaskContext;
 use crate::scan::tasks::TaskFuture;
@@ -172,15 +175,19 @@ impl RepeatedScan {
         }
     }
 
+    fn task_context(&self) -> Arc<TaskContext> {
+        Arc::new(TaskContext {
+            filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))),
+            reader: Arc::clone(&self.layout_reader),
+            projection: self.projection.clone(),
+        })
+    }
+
     pub(crate) fn execute(
         &self,
         row_range: Option<Range<u64>>,
     ) -> VortexResult<Vec<TaskFuture<Option<ArrayRef>>>> {
-        let ctx = Arc::new(TaskContext {
-            filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))),
-            reader: Arc::clone(&self.layout_reader),
-            projection: self.projection.clone(),
-        });
+        let ctx = self.task_context();
 
         let mut limit = self.limit;
         let mut tasks = Vec::new();
@@ -209,57 +216,64 @@ impl RepeatedScan {
         row_range: Option<Range<u64>>,
     ) -> VortexResult<impl Stream<Item = VortexResult<ArrayRef>> + Send + 'static> {
         let num_workers = get_available_parallelism().unwrap_or(1);
+        let concurrency = self.concurrency * num_workers;
         let handle = self.session.handle();
 
+        // With both a filter and a limit we cannot know each split's output row count ahead of
+        // time, so split tasks are built lazily as the stream is polled. `buffered`'s read-ahead
+        // (bounded by `concurrency`) registers IO for splits eagerly, but only as far as the
+        // limit requires: `limit_array_stream` drops the inner stream once the limit is reached,
+        // capping over-read at `concurrency` splits.
         if self.filter.is_some() && self.limit.is_some() {
-            let ctx = Arc::new(TaskContext {
-                filter: self.filter.clone().map(|f| Arc::new(FilterExpr::new(f))),
-                reader: Arc::clone(&self.layout_reader),
-                projection: self.projection.clone(),
-            });
+            let ctx = self.task_context();
             let selection = self.selection.clone();
-            let ordered = self.ordered;
-            let limit = self.limit;
-            let stream = futures::stream::iter(self.split_ranges(row_range)).map(move |range| {
-                let ctx = Arc::clone(&ctx);
-                let handle = handle.clone();
-                let selection = selection.clone();
-                async move {
+            let tasks =
+                futures::stream::iter(self.split_ranges(row_range)).filter_map(move |range| {
+                    // Build the row mask and split task synchronously so the IO system sees the
+                    // split's ranges as soon as `buffered` pulls it, without cloning `selection`.
                     let row_mask = selection.row_mask(&range);
-                    if row_mask.mask().all_false() {
-                        return Ok(None);
-                    }
+                    let spawned = (!row_mask.mask().all_false()).then(|| {
+                        let task = split_exec(Arc::clone(&ctx), row_mask, None)
+                            .unwrap_or_else(|err| async move { Err(err) }.boxed());
+                        handle.spawn(task)
+                    });
+                    async move { spawned }
+                });
 
-                    let task = split_exec(ctx, row_mask, None)?;
-                    handle.spawn(task).await
-                }
-                .boxed()
-            });
-            let stream = if ordered {
-                stream.buffered(1).boxed()
-            } else {
-                stream.buffer_unordered(1).boxed()
-            };
-
-            return Ok(limit_array_stream(
-                stream.filter_map(|chunk| async move { chunk.transpose() }),
-                limit,
-            ));
+            return Ok(schedule(tasks, self.ordered, concurrency, self.limit));
         }
 
-        let stream =
+        // No filter (or no limit): build every task eagerly so the IO system sees all split
+        // ranges up front. A no-filter limit is applied exactly per split inside `execute`.
+        let tasks =
             futures::stream::iter(self.execute(row_range)?).map(move |task| handle.spawn(task));
-        let concurrency = self.concurrency * num_workers;
-        let stream = if self.ordered {
-            stream.buffered(concurrency).boxed()
-        } else {
-            stream.buffer_unordered(concurrency).boxed()
-        };
 
-        Ok(limit_array_stream(
-            stream.filter_map(|chunk| async move { chunk.transpose() }),
-            self.limit,
-        ))
+        Ok(schedule(tasks, self.ordered, concurrency, self.limit))
+    }
+}
+
+/// Spawn-buffer a stream of split tasks, transposing empty splits away and applying `limit`.
+fn schedule<S>(
+    tasks: S,
+    ordered: bool,
+    concurrency: usize,
+    limit: Option<u64>,
+) -> BoxStream<'static, VortexResult<ArrayRef>>
+where
+    S: Stream<Item = Task<VortexResult<Option<ArrayRef>>>> + Send + 'static,
+{
+    let stream = if ordered {
+        tasks.buffered(concurrency).boxed()
+    } else {
+        tasks.buffer_unordered(concurrency).boxed()
+    };
+    let stream = stream
+        .filter_map(|chunk| async move { chunk.transpose() })
+        .boxed();
+
+    match limit {
+        Some(limit) => LimitedStream::new(stream, RowBudget::Local(limit)).boxed(),
+        None => stream,
     }
 }
 
