vortex_layout/scan/
scan_builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::Context;
8use std::task::Poll;
9use std::task::ready;
10
11use futures::Stream;
12use futures::StreamExt;
13use futures::future::BoxFuture;
14use futures::stream::BoxStream;
15use itertools::Itertools;
16use vortex_array::ArrayRef;
17use vortex_array::dtype::DType;
18use vortex_array::dtype::FieldMask;
19use vortex_array::expr::BoundExpression;
20use vortex_array::expr::analysis::referenced_field_paths;
21use vortex_array::iter::ArrayIterator;
22use vortex_array::iter::ArrayIteratorAdapter;
23use vortex_array::stats::StatsSet;
24use vortex_array::stream::ArrayStream;
25use vortex_array::stream::ArrayStreamAdapter;
26use vortex_error::VortexExpect;
27use vortex_error::VortexResult;
28use vortex_error::vortex_bail;
29use vortex_io::runtime::BlockingRuntime;
30use vortex_io::runtime::Handle;
31use vortex_io::runtime::Task;
32use vortex_io::session::RuntimeSessionExt;
33use vortex_metrics::MetricsRegistry;
34use vortex_scan::selection::Selection;
35use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
36use vortex_session::VortexSession;
37use vortex_utils::parallelism::get_available_parallelism;
38
39use crate::LayoutReader;
40use crate::LayoutReaderRef;
41use crate::layouts::row_idx::RowIdx;
42use crate::layouts::row_idx::RowIdxLayoutReader;
43use crate::scan::repeated_scan::RepeatedScan;
44use crate::scan::split_by::SplitBy;
45use crate::scan::splits::Splits;
46use crate::scan::splits::attempt_split_ranges;
47
48/// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs.
49///
50/// A scan has three independent row restriction mechanisms:
51///
52/// - [`with_row_range`](Self::with_row_range) selects a contiguous range before scanning.
53/// - [`with_selection`](Self::with_selection) applies a [`Selection`] inside that range.
54/// - [`with_filter`](Self::with_filter) evaluates an expression predicate during execution.
55///
56/// Projection and filter expressions must be bound against the reader dtype. Work is divided by
57/// the configured [`SplitBy`] strategy or by explicit selection ranges.
58pub struct ScanBuilder<A> {
59    session: VortexSession,
60    layout_reader: LayoutReaderRef,
61    projection: BoundExpression,
62    filter: Option<BoundExpression>,
63    /// Whether the scan needs to return splits in the order they appear in the file.
64    ordered: bool,
65    /// Optionally read a subset of the rows in the file.
66    row_range: Option<Range<u64>>,
67    /// The selection mask to apply to the selected row range.
68    // TODO(joe): replace this is usage of row_id selection, see
69    selection: Selection,
70    /// How to split the file for concurrent processing.
71    split_by: SplitBy,
72    /// Precomputed full-file natural split boundaries; when set, [`prepare`](Self::prepare)
73    /// uses them instead of walking the layout.
74    natural_splits: Option<Arc<[u64]>>,
75    /// The number of splits to make progress on concurrently **per-thread**.
76    concurrency: usize,
77    /// Function to apply to each [`ArrayRef`] within the spawned split tasks.
78    map_fn: Arc<dyn Fn(ArrayRef) -> VortexResult<A> + Send + Sync>,
79    metrics_registry: Option<Arc<dyn MetricsRegistry>>,
80    /// Should we try to prune the file (using stats) on open.
81    file_stats: Option<Arc<[StatsSet]>>,
82    /// Maximal number of rows to read (after filtering)
83    limit: Option<u64>,
84    /// The row-offset assigned to the first row of the file. Used by the `row_idx` expression,
85    /// but not by the scan [`Selection`] which remains relative.
86    row_offset: u64,
87}
88
89impl ScanBuilder<ArrayRef> {
90    /// Create a scan builder over `layout_reader` using `session` for runtime and execution state.
91    pub fn new(session: VortexSession, layout_reader: Arc<dyn LayoutReader>) -> Self {
92        let projection = BoundExpression::new_root(layout_reader.dtype().clone());
93        Self {
94            session,
95            layout_reader,
96            projection,
97            filter: None,
98            ordered: true,
99            row_range: None,
100            selection: Default::default(),
101            split_by: SplitBy::Layout,
102            natural_splits: None,
103            // We default to four tasks per worker thread, which allows for some I/O lookahead
104            // without too much impact on work-stealing.
105            concurrency: 4,
106            map_fn: Arc::new(Ok),
107            metrics_registry: None,
108            file_stats: None,
109            limit: None,
110            row_offset: 0,
111        }
112    }
113
114    /// Returns an [`ArrayStream`] with tasks spawned onto the session's runtime handle.
115    ///
116    /// See [`ScanBuilder::into_stream`] for more details.
117    pub fn into_array_stream(self) -> VortexResult<impl ArrayStream + Send + 'static> {
118        let dtype = self.dtype()?;
119        let stream = self.into_stream()?;
120        Ok(ArrayStreamAdapter::new(dtype, stream))
121    }
122
123    /// Returns an [`ArrayIterator`] using the given blocking runtime.
124    pub fn into_array_iter<B: BlockingRuntime>(
125        self,
126        runtime: &B,
127    ) -> VortexResult<impl ArrayIterator + 'static> {
128        let stream = self.into_array_stream()?;
129        let dtype = stream.dtype().clone();
130        Ok(ArrayIteratorAdapter::new(
131            dtype,
132            runtime.block_on_stream(stream),
133        ))
134    }
135}
136
137impl<A: 'static + Send> ScanBuilder<A> {
138    /// Add a filter expression bound against the reader dtype.
139    pub fn with_filter(mut self, filter: BoundExpression) -> Self {
140        self.filter = Some(filter);
141        self
142    }
143
144    /// Add or clear a filter expression bound against the reader dtype.
145    pub fn with_some_filter(mut self, filter: Option<BoundExpression>) -> Self {
146        self.filter = filter;
147        self
148    }
149
150    /// Set a projection expression bound against the reader dtype.
151    pub fn with_projection(mut self, projection: BoundExpression) -> Self {
152        self.projection = projection;
153        self
154    }
155
156    /// Returns whether output chunks are yielded in file order.
157    pub fn ordered(&self) -> bool {
158        self.ordered
159    }
160
161    /// Configure whether output chunks must be yielded in file order.
162    pub fn with_ordered(mut self, ordered: bool) -> Self {
163        self.ordered = ordered;
164        self
165    }
166
167    /// Restrict scanning to a contiguous row range.
168    pub fn with_row_range(mut self, row_range: Range<u64>) -> Self {
169        self.row_range = Some(row_range);
170        self
171    }
172
173    /// Apply a row selection to the selected row range.
174    pub fn with_selection(mut self, selection: Selection) -> Self {
175        self.selection = selection;
176        self
177    }
178
179    /// Select rows by strictly sorted absolute indices relative to the scan input.
180    pub fn with_row_indices(mut self, row_indices: StrictSortedBuffer<u64>) -> Self {
181        self.selection = Selection::IncludeByIndex(row_indices);
182        self
183    }
184
185    /// Set the root row offset used by row-index expressions.
186    pub fn with_row_offset(mut self, row_offset: u64) -> Self {
187        self.row_offset = row_offset;
188        self
189    }
190
191    /// Configure how natural scan work is split for concurrency.
192    pub fn with_split_by(mut self, split_by: SplitBy) -> Self {
193        self.split_by = split_by;
194        self
195    }
196
197    /// Supply precomputed full-file natural split boundaries (see
198    /// [`full_file_splits`](Self::full_file_splits)) so [`prepare`](Self::prepare) reuses them
199    /// instead of walking the layout. Callers translating external partitions into row ranges
200    /// can compute the boundaries once per file and share them across partitions.
201    ///
202    /// Takes precedence over [`with_split_by`](Self::with_split_by); boundaries outside the
203    /// scan's row range are clamped during execution. Boundaries must be strictly increasing.
204    pub fn with_natural_splits(mut self, boundaries: Arc<[u64]>) -> Self {
205        debug_assert!(
206            boundaries.windows(2).all(|w| w[0] < w[1]),
207            "natural split boundaries must be strictly increasing"
208        );
209        self.natural_splits = Some(boundaries);
210        self
211    }
212
213    /// Compute the full-file natural split boundaries for the fields referenced by this scan's
214    /// projection and filter, ignoring any configured row range.
215    ///
216    /// These are the boundaries [`prepare`](Self::prepare) derives for a whole-file scan; hand
217    /// them back via [`with_natural_splits`](Self::with_natural_splits) to skip the layout walk
218    /// in `prepare`.
219    pub fn full_file_splits(&self) -> VortexResult<Vec<u64>> {
220        let field_mask = referenced_field_masks(&self.projection, self.filter.as_ref())?;
221        self.split_by.splits(
222            self.layout_reader.as_ref(),
223            &(0..self.layout_reader.row_count()),
224            &field_mask,
225        )
226    }
227
228    /// Returns the per-worker row-split concurrency.
229    pub fn concurrency(&self) -> usize {
230        self.concurrency
231    }
232
233    /// The number of row splits to make progress on concurrently per-thread, must
234    /// be greater than 0.
235    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
236        assert!(concurrency > 0);
237        self.concurrency = concurrency;
238        self
239    }
240
241    /// Add or clear the metrics registry used by scan execution.
242    pub fn with_some_metrics_registry(mut self, metrics: Option<Arc<dyn MetricsRegistry>>) -> Self {
243        self.metrics_registry = metrics;
244        self
245    }
246
247    /// Set the metrics registry used by scan execution.
248    pub fn with_metrics_registry(mut self, metrics: Arc<dyn MetricsRegistry>) -> Self {
249        self.metrics_registry = Some(metrics);
250        self
251    }
252
253    /// Add or clear the maximum number of rows returned after filtering.
254    pub fn with_some_limit(mut self, limit: Option<u64>) -> Self {
255        self.limit = limit;
256        self
257    }
258
259    /// Set the maximum number of rows returned after filtering.
260    pub fn with_limit(mut self, limit: u64) -> Self {
261        self.limit = Some(limit);
262        self
263    }
264
265    /// The [`DType`] returned by the scan, after applying the projection.
266    pub fn dtype(&self) -> VortexResult<DType> {
267        Ok(self.projection.dtype().clone())
268    }
269
270    /// The session used by the scan.
271    pub fn session(&self) -> &VortexSession {
272        &self.session
273    }
274
275    /// Map each split of the scan. The function will be run on the spawned task.
276    pub fn map<B: 'static>(
277        self,
278        map_fn: impl Fn(A) -> VortexResult<B> + 'static + Send + Sync,
279    ) -> ScanBuilder<B> {
280        let old_map_fn = self.map_fn;
281        ScanBuilder {
282            session: self.session,
283            layout_reader: self.layout_reader,
284            projection: self.projection,
285            filter: self.filter,
286            ordered: self.ordered,
287            row_range: self.row_range,
288            selection: self.selection,
289            split_by: self.split_by,
290            natural_splits: self.natural_splits,
291            concurrency: self.concurrency,
292            metrics_registry: self.metrics_registry,
293            file_stats: self.file_stats,
294            limit: self.limit,
295            row_offset: self.row_offset,
296            map_fn: Arc::new(move |a| old_map_fn(a).and_then(&map_fn)),
297        }
298    }
299
300    /// Optimize expressions, compute split ranges, and return an executable repeated scan.
301    pub fn prepare(self) -> VortexResult<RepeatedScan<A>> {
302        let dtype = self.dtype()?;
303
304        if self.filter.is_some() && self.limit.is_some() {
305            vortex_bail!("Vortex doesn't support scans with both a filter and a limit")
306        }
307
308        // Spin up the root layout reader, and wrap it in a FilterLayoutReader to perform
309        // conjunction splitting if a filter is provided.
310        let mut layout_reader = self.layout_reader;
311
312        // Enrich the layout reader to support RowIdx expressions if scan uses #row_idx.
313        // Note that this is applied below the filter layout reader since it can perform
314        // better over individual conjunctions.
315        let mut found_row_idx = self.projection.contains::<RowIdx>()?;
316        if !found_row_idx && let Some(filter) = self.filter.as_ref() {
317            found_row_idx = filter.contains::<RowIdx>()?;
318        }
319        if found_row_idx {
320            layout_reader = Arc::new(RowIdxLayoutReader::new(
321                self.row_offset,
322                layout_reader,
323                self.session.clone(),
324            ));
325        }
326
327        let bound_projection = self.projection;
328        let bound_filter = self.filter;
329
330        // Compute the row splits of the scan.
331        let splits =
332            if let Some(ranges) = attempt_split_ranges(&self.selection, self.row_range.as_ref()) {
333                Splits::Ranges(ranges)
334            } else if let Some(boundaries) = self.natural_splits {
335                // Caller-supplied full-file boundaries; execution clamps them to the row range.
336                Splits::Natural(boundaries)
337            } else {
338                let field_mask = referenced_field_masks(&bound_projection, bound_filter.as_ref())?;
339                let split_range = self
340                    .row_range
341                    .clone()
342                    .unwrap_or_else(|| 0..layout_reader.row_count());
343                Splits::Natural(
344                    self.split_by
345                        .splits(layout_reader.as_ref(), &split_range, &field_mask)?
346                        .into(),
347                )
348            };
349
350        Ok(RepeatedScan::new(
351            self.session.clone(),
352            layout_reader,
353            bound_projection,
354            bound_filter,
355            self.ordered,
356            self.row_range,
357            self.selection,
358            splits,
359            self.concurrency,
360            self.map_fn,
361            self.limit,
362            dtype,
363        ))
364    }
365
366    /// Constructs a task per row split of the scan, returned as a vector of futures.
367    pub fn build(self) -> VortexResult<Vec<BoxFuture<'static, VortexResult<Option<A>>>>> {
368        // The ultimate short circuit
369        if self.limit.is_some_and(|l| l == 0) {
370            return Ok(vec![]);
371        }
372
373        self.prepare()?.execute(None)
374    }
375
376    /// Returns a [`Stream`] with tasks spawned onto the session's runtime handle.
377    pub fn into_stream(
378        self,
379    ) -> VortexResult<impl Stream<Item = VortexResult<A>> + Send + 'static + use<A>> {
380        Ok(LazyScanStream::new(self))
381    }
382
383    /// Returns an [`Iterator`] using the session's runtime.
384    pub fn into_iter<B: BlockingRuntime>(
385        self,
386        runtime: &B,
387    ) -> VortexResult<impl Iterator<Item = VortexResult<A>> + 'static> {
388        let stream = self.into_stream()?;
389        Ok(runtime.block_on_stream(stream))
390    }
391}
392
393enum LazyScanState<A: 'static + Send> {
394    Builder(Option<Box<ScanBuilder<A>>>),
395    Preparing(PreparingScan<A>),
396    Stream(BoxStream<'static, VortexResult<A>>),
397    Error(Option<vortex_error::VortexError>),
398}
399
400type PreparedScanTasks<A> = Vec<BoxFuture<'static, VortexResult<Option<A>>>>;
401
402struct PreparingScan<A: 'static + Send> {
403    ordered: bool,
404    concurrency: usize,
405    handle: Handle,
406    task: Task<VortexResult<PreparedScanTasks<A>>>,
407}
408
409struct LazyScanStream<A: 'static + Send> {
410    state: LazyScanState<A>,
411}
412
413impl<A: 'static + Send> LazyScanStream<A> {
414    fn new(builder: ScanBuilder<A>) -> Self {
415        Self {
416            state: LazyScanState::Builder(Some(Box::new(builder))),
417        }
418    }
419}
420
421impl<A: 'static + Send> Unpin for LazyScanStream<A> {}
422
423impl<A: 'static + Send> Stream for LazyScanStream<A> {
424    type Item = VortexResult<A>;
425
426    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
427        loop {
428            match &mut self.state {
429                LazyScanState::Builder(builder) => {
430                    let builder = builder.take().vortex_expect("polled after completion");
431                    let ordered = builder.ordered;
432                    let num_workers = get_available_parallelism().unwrap_or(1);
433                    let concurrency = builder.concurrency * num_workers;
434                    let handle = builder.session.handle();
435                    let task = handle
436                        .spawn_cpu(move || builder.prepare().and_then(|scan| scan.execute(None)));
437                    self.state = LazyScanState::Preparing(PreparingScan {
438                        ordered,
439                        concurrency,
440                        handle,
441                        task,
442                    });
443                }
444                LazyScanState::Preparing(preparing) => {
445                    match ready!(Pin::new(&mut preparing.task).poll(cx)) {
446                        Ok(tasks) => {
447                            let ordered = preparing.ordered;
448                            let concurrency = preparing.concurrency;
449                            let handle = preparing.handle.clone();
450                            let stream =
451                                futures::stream::iter(tasks).map(move |task| handle.spawn(task));
452                            let stream = if ordered {
453                                stream.buffered(concurrency).boxed()
454                            } else {
455                                stream.buffer_unordered(concurrency).boxed()
456                            };
457                            let stream = stream
458                                .filter_map(|chunk| async move { chunk.transpose() })
459                                .boxed();
460                            self.state = LazyScanState::Stream(stream);
461                        }
462                        Err(err) => self.state = LazyScanState::Error(Some(err)),
463                    }
464                }
465                LazyScanState::Stream(stream) => return stream.as_mut().poll_next(cx),
466                LazyScanState::Error(err) => return Poll::Ready(err.take().map(Err)),
467            }
468        }
469    }
470}
471
472/// Compute masks of field paths referenced by the projection and filter in the scan.
473///
474/// Projection and filter must be pre-simplified and bound against the scan dtype.
475pub fn referenced_field_masks(
476    projection: &BoundExpression,
477    filter: Option<&BoundExpression>,
478) -> VortexResult<Vec<FieldMask>> {
479    let mut field_paths = referenced_field_paths(projection)?;
480    if let Some(filter) = filter {
481        field_paths.extend(referenced_field_paths(filter)?);
482    }
483    Ok(field_paths
484        .into_iter()
485        .map(|path| {
486            if path.is_root() {
487                FieldMask::All
488            } else {
489                FieldMask::Prefix(path)
490            }
491        })
492        .collect_vec())
493}
494
495#[cfg(test)]
496mod test {
497    use std::ops::Range;
498    use std::pin::Pin;
499    use std::sync::Arc;
500    use std::sync::atomic::AtomicUsize;
501    use std::sync::atomic::Ordering;
502    use std::sync::mpsc;
503    use std::task::Context;
504    use std::task::Poll;
505    use std::time::Duration;
506
507    use futures::Stream;
508    use futures::task::noop_waker_ref;
509    use parking_lot::Mutex;
510    use vortex_array::IntoArray;
511    use vortex_array::MaskFuture;
512    use vortex_array::VortexSessionExecute;
513    use vortex_array::array_session;
514    use vortex_array::arrays::PrimitiveArray;
515    use vortex_array::dtype::DType;
516    use vortex_array::dtype::FieldMask;
517    use vortex_array::dtype::FieldPath;
518    use vortex_array::dtype::Nullability;
519    use vortex_array::dtype::PType;
520    use vortex_array::dtype::StructFields;
521    use vortex_array::expr::BoundExpression;
522    use vortex_array::expr::ExactBoundExpr;
523    use vortex_array::expr::eq;
524    use vortex_array::expr::get_item;
525    use vortex_array::expr::is_not_null;
526    use vortex_array::expr::lit;
527    use vortex_array::expr::root;
528    use vortex_error::VortexResult;
529    use vortex_error::vortex_err;
530    use vortex_io::runtime::BlockingRuntime;
531    use vortex_io::runtime::single::SingleThreadRuntime;
532    use vortex_mask::Mask;
533
534    use super::ScanBuilder;
535    use super::referenced_field_masks;
536    use crate::ArrayFuture;
537    use crate::LayoutReader;
538    use crate::RowSplits;
539    use crate::SplitRange;
540    use crate::scan::test::SCAN_SESSION;
541    use crate::scan::test::session_with_handle;
542
543    fn nested_dtype() -> DType {
544        DType::Struct(
545            StructFields::from_iter([
546                (
547                    "a",
548                    DType::Struct(
549                        StructFields::from_iter([
550                            ("1", DType::Primitive(PType::I32, Nullability::NonNullable)),
551                            ("2", DType::Primitive(PType::I32, Nullability::NonNullable)),
552                        ]),
553                        Nullability::NonNullable,
554                    ),
555                ),
556                ("b", DType::Primitive(PType::I32, Nullability::NonNullable)),
557            ]),
558            Nullability::NonNullable,
559        )
560    }
561
562    #[test]
563    fn bound_setters_preserve_identity() -> VortexResult<()> {
564        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
565        let projection = eq(root(), lit(1_i32)).bind(&dtype)?;
566        let filter = eq(root(), lit(2_i32)).bind(&dtype)?;
567        let expected_projection = ExactBoundExpr(projection.clone());
568        let expected_filter = ExactBoundExpr(filter.clone());
569        let reader = Arc::new(CountingLayoutReader::new(Arc::new(AtomicUsize::new(0))));
570
571        let builder = ScanBuilder::new(SCAN_SESSION.clone(), reader)
572            .with_projection(projection)
573            .with_filter(filter);
574
575        assert_eq!(ExactBoundExpr(builder.projection), expected_projection);
576        assert_eq!(builder.filter.map(ExactBoundExpr), Some(expected_filter));
577        Ok(())
578    }
579
580    #[test]
581    fn root_projection_produces_all_mask() -> VortexResult<()> {
582        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
583        let projection = root().bind(&dtype)?;
584
585        assert_eq!(referenced_field_masks(&projection, None)?, [FieldMask::All]);
586        Ok(())
587    }
588
589    #[test]
590    fn nested_projection_preserves_field_path_in_split_mask() -> VortexResult<()> {
591        let dtype = nested_dtype();
592        let projection = get_item("1", get_item("a", root())).bind(&dtype)?;
593        let filter = eq(get_item("2", get_item("a", root())), lit(0_i32)).bind(&dtype)?;
594
595        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
596
597        assert_eq!(field_masks.len(), 2);
598        assert!(field_masks.contains(&FieldMask::Prefix(FieldPath::from_name("a").push("1"))));
599        assert!(field_masks.contains(&FieldMask::Prefix(FieldPath::from_name("a").push("2"))));
600        Ok(())
601    }
602
603    #[test]
604    fn filter_path_covers_nested_projection_path() -> VortexResult<()> {
605        let dtype = nested_dtype();
606        let projection = get_item("1", get_item("a", root())).bind(&dtype)?;
607        let filter = is_not_null(get_item("a", root())).bind(&dtype)?;
608
609        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
610
611        assert_eq!(field_masks, [FieldMask::Prefix(FieldPath::from_name("a"))]);
612        Ok(())
613    }
614
615    #[test]
616    fn parent_projection_path_covers_nested_filter_path() -> VortexResult<()> {
617        let dtype = nested_dtype();
618        let projection = get_item("a", root()).bind(&dtype)?;
619        let filter = is_not_null(get_item("1", get_item("a", root()))).bind(&dtype)?;
620
621        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
622
623        assert_eq!(field_masks, [FieldMask::Prefix(FieldPath::from_name("a"))]);
624        Ok(())
625    }
626
627    #[derive(Debug)]
628    struct CountingLayoutReader {
629        name: Arc<str>,
630        dtype: DType,
631        row_count: u64,
632        register_splits_calls: Arc<AtomicUsize>,
633    }
634
635    impl CountingLayoutReader {
636        fn new(register_splits_calls: Arc<AtomicUsize>) -> Self {
637            Self {
638                name: Arc::from("counting"),
639                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
640                row_count: 1,
641                register_splits_calls,
642            }
643        }
644    }
645
646    impl LayoutReader for CountingLayoutReader {
647        fn name(&self) -> &Arc<str> {
648            &self.name
649        }
650
651        fn dtype(&self) -> &DType {
652            &self.dtype
653        }
654
655        fn row_count(&self) -> u64 {
656            self.row_count
657        }
658
659        fn register_splits(
660            &self,
661            _field_mask: &[FieldMask],
662            split_range: &SplitRange,
663            splits: &mut RowSplits,
664        ) -> VortexResult<()> {
665            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
666            splits.push(split_range.root_row_range().end);
667            Ok(())
668        }
669
670        fn pruning_evaluation(
671            &self,
672            _row_range: &Range<u64>,
673            _expr: &BoundExpression,
674            _mask: Mask,
675        ) -> VortexResult<MaskFuture> {
676            unimplemented!("not needed for this test");
677        }
678
679        fn filter_evaluation(
680            &self,
681            _row_range: &Range<u64>,
682            _expr: &BoundExpression,
683            _mask: MaskFuture,
684        ) -> VortexResult<MaskFuture> {
685            unimplemented!("not needed for this test");
686        }
687
688        fn projection_evaluation(
689            &self,
690            _row_range: &Range<u64>,
691            _expr: &BoundExpression,
692            _mask: MaskFuture,
693        ) -> VortexResult<ArrayFuture> {
694            Ok(Box::pin(async move {
695                unreachable!("scan should not be polled in this test")
696            }))
697        }
698
699        fn as_any(&self) -> &dyn std::any::Any {
700            self
701        }
702    }
703
704    #[test]
705    fn into_stream_is_lazy() {
706        let calls = Arc::new(AtomicUsize::new(0));
707        let reader = Arc::new(CountingLayoutReader::new(Arc::clone(&calls)));
708
709        let session = SCAN_SESSION.clone();
710
711        let _stream = ScanBuilder::new(session, reader).into_stream().unwrap();
712
713        assert_eq!(calls.load(Ordering::Relaxed), 0);
714    }
715
716    #[derive(Debug)]
717    struct SplittingLayoutReader {
718        name: Arc<str>,
719        dtype: DType,
720        row_count: u64,
721        register_splits_calls: Arc<AtomicUsize>,
722    }
723
724    impl SplittingLayoutReader {
725        fn new(register_splits_calls: Arc<AtomicUsize>) -> Self {
726            Self {
727                name: Arc::from("splitting"),
728                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
729                row_count: 4,
730                register_splits_calls,
731            }
732        }
733    }
734
735    impl LayoutReader for SplittingLayoutReader {
736        fn name(&self) -> &Arc<str> {
737            &self.name
738        }
739
740        fn dtype(&self) -> &DType {
741            &self.dtype
742        }
743
744        fn row_count(&self) -> u64 {
745            self.row_count
746        }
747
748        fn register_splits(
749            &self,
750            _field_mask: &[FieldMask],
751            split_range: &SplitRange,
752            splits: &mut RowSplits,
753        ) -> VortexResult<()> {
754            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
755            for split in (split_range.row_range().start + 1)..=split_range.row_range().end {
756                splits.push(split_range.row_offset() + split);
757            }
758            Ok(())
759        }
760
761        fn pruning_evaluation(
762            &self,
763            _row_range: &Range<u64>,
764            _expr: &BoundExpression,
765            mask: Mask,
766        ) -> VortexResult<MaskFuture> {
767            Ok(MaskFuture::ready(mask))
768        }
769
770        fn filter_evaluation(
771            &self,
772            _row_range: &Range<u64>,
773            _expr: &BoundExpression,
774            mask: MaskFuture,
775        ) -> VortexResult<MaskFuture> {
776            Ok(mask)
777        }
778
779        fn projection_evaluation(
780            &self,
781            row_range: &Range<u64>,
782            _expr: &BoundExpression,
783            _mask: MaskFuture,
784        ) -> VortexResult<ArrayFuture> {
785            let start = usize::try_from(row_range.start)
786                .map_err(|_| vortex_err!("row_range.start must fit in usize"))?;
787            let end = usize::try_from(row_range.end)
788                .map_err(|_| vortex_err!("row_range.end must fit in usize"))?;
789
790            let values: VortexResult<Vec<i32>> = (start..end)
791                .map(|v| i32::try_from(v).map_err(|_| vortex_err!("split value must fit in i32")))
792                .collect();
793
794            let array = PrimitiveArray::from_iter(values?).into_array();
795            Ok(Box::pin(async move { Ok(array) }))
796        }
797
798        fn as_any(&self) -> &dyn std::any::Any {
799            self
800        }
801    }
802
803    #[test]
804    fn into_stream_executes_after_prepare() -> VortexResult<()> {
805        let mut ctx = array_session().create_execution_ctx();
806        let calls = Arc::new(AtomicUsize::new(0));
807        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
808
809        let runtime = SingleThreadRuntime::default();
810        let session = session_with_handle(runtime.handle());
811
812        let stream = ScanBuilder::new(session, reader).into_stream()?;
813        let mut iter = runtime.block_on_stream(stream);
814
815        let mut values = Vec::new();
816        for chunk in &mut iter {
817            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
818            values.push(prim.into_buffer::<i32>()[0]);
819        }
820
821        assert_eq!(calls.load(Ordering::Relaxed), 1);
822        assert_eq!(values.as_ref(), [0, 1, 2, 3]);
823
824        Ok(())
825    }
826
827    #[test]
828    fn supplied_natural_splits_skip_layout_walk() -> VortexResult<()> {
829        let mut ctx = array_session().create_execution_ctx();
830        let calls = Arc::new(AtomicUsize::new(0));
831        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
832
833        let runtime = SingleThreadRuntime::default();
834        let session = session_with_handle(runtime.handle());
835
836        let stream = ScanBuilder::new(session, reader)
837            .with_natural_splits(vec![0u64, 2, 4].into())
838            .with_row_range(1..4)
839            .into_stream()?;
840        let mut iter = runtime.block_on_stream(stream);
841
842        let mut chunks = Vec::new();
843        for chunk in &mut iter {
844            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
845            chunks.push(prim.into_buffer::<i32>().to_vec());
846        }
847
848        assert_eq!(calls.load(Ordering::Relaxed), 0);
849        // Supplied full-file boundaries [0, 2, 4] clamped to rows 1..4.
850        assert_eq!(chunks, [vec![1], vec![2, 3]]);
851
852        Ok(())
853    }
854
855    #[test]
856    fn full_file_splits_ignore_row_range() -> VortexResult<()> {
857        let calls = Arc::new(AtomicUsize::new(0));
858        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
859
860        let splits = ScanBuilder::new(SCAN_SESSION.clone(), reader)
861            .with_row_range(1..3)
862            .full_file_splits()?;
863
864        assert_eq!(splits, [0, 1, 2, 3, 4]);
865        Ok(())
866    }
867
868    #[derive(Debug)]
869    struct BlockingSplitsLayoutReader {
870        name: Arc<str>,
871        dtype: DType,
872        row_count: u64,
873        register_splits_calls: Arc<AtomicUsize>,
874        gate: Arc<Mutex<()>>,
875    }
876
877    impl BlockingSplitsLayoutReader {
878        fn new(gate: Arc<Mutex<()>>, register_splits_calls: Arc<AtomicUsize>) -> Self {
879            Self {
880                name: Arc::from("blocking-splits"),
881                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
882                row_count: 1,
883                register_splits_calls,
884                gate,
885            }
886        }
887    }
888
889    impl LayoutReader for BlockingSplitsLayoutReader {
890        fn name(&self) -> &Arc<str> {
891            &self.name
892        }
893
894        fn dtype(&self) -> &DType {
895            &self.dtype
896        }
897
898        fn row_count(&self) -> u64 {
899            self.row_count
900        }
901
902        fn register_splits(
903            &self,
904            _field_mask: &[FieldMask],
905            split_range: &SplitRange,
906            splits: &mut RowSplits,
907        ) -> VortexResult<()> {
908            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
909            let _guard = self.gate.lock();
910            splits.push(split_range.root_row_range().end);
911            Ok(())
912        }
913
914        fn pruning_evaluation(
915            &self,
916            _row_range: &Range<u64>,
917            _expr: &BoundExpression,
918            _mask: Mask,
919        ) -> VortexResult<MaskFuture> {
920            unimplemented!("not needed for this test");
921        }
922
923        fn filter_evaluation(
924            &self,
925            _row_range: &Range<u64>,
926            _expr: &BoundExpression,
927            _mask: MaskFuture,
928        ) -> VortexResult<MaskFuture> {
929            unimplemented!("not needed for this test");
930        }
931
932        fn projection_evaluation(
933            &self,
934            _row_range: &Range<u64>,
935            _expr: &BoundExpression,
936            _mask: MaskFuture,
937        ) -> VortexResult<ArrayFuture> {
938            Ok(Box::pin(async move {
939                unreachable!("scan should not be polled in this test")
940            }))
941        }
942
943        fn as_any(&self) -> &dyn std::any::Any {
944            self
945        }
946    }
947
948    #[test]
949    fn into_stream_first_poll_does_not_block() {
950        let gate = Arc::new(Mutex::new(()));
951        let guard = gate.lock();
952
953        let calls = Arc::new(AtomicUsize::new(0));
954        let reader = Arc::new(BlockingSplitsLayoutReader::new(
955            Arc::clone(&gate),
956            Arc::clone(&calls),
957        ));
958
959        let runtime = SingleThreadRuntime::default();
960        let session = session_with_handle(runtime.handle());
961
962        let mut stream = ScanBuilder::new(session, reader).into_stream().unwrap();
963
964        let (send, recv) = mpsc::channel::<bool>();
965        let join = std::thread::spawn(move || {
966            let waker = noop_waker_ref();
967            let mut cx = Context::from_waker(waker);
968            let poll = Pin::new(&mut stream).poll_next(&mut cx);
969            let _ = send.send(matches!(poll, Poll::Pending));
970        });
971
972        let polled_pending = recv.recv_timeout(Duration::from_secs(1)).ok();
973
974        // Always release the gate and join the thread so failures don't hang the test process.
975        drop(guard);
976        drop(join.join());
977
978        let polled_pending = polled_pending.expect("poll_next blocked; expected quick return");
979        assert!(
980            polled_pending,
981            "expected Poll::Pending while prepare is blocked"
982        );
983        assert_eq!(calls.load(Ordering::Relaxed), 0);
984
985        drop(runtime);
986    }
987
988    #[test]
989    fn into_stream_with_row_range() -> VortexResult<()> {
990        let mut ctx = array_session().create_execution_ctx();
991        let calls = Arc::new(AtomicUsize::new(0));
992        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
993
994        let runtime = SingleThreadRuntime::default();
995        let session = session_with_handle(runtime.handle());
996
997        let stream = ScanBuilder::new(session, reader)
998            .with_row_range(1..3)
999            .into_stream()?;
1000        let mut iter = runtime.block_on_stream(stream);
1001
1002        let mut values = Vec::new();
1003        for chunk in &mut iter {
1004            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
1005            values.extend(prim.into_buffer::<i32>().iter().copied());
1006        }
1007
1008        assert_eq!(calls.load(Ordering::Relaxed), 1);
1009        assert_eq!(values.as_ref(), [1, 2]);
1010
1011        Ok(())
1012    }
1013}
vortex_layout/scan/scan_builder.rs

vortex_layout/scan/
scan_builder.rs
