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::default(),
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_bail;
530    use vortex_error::vortex_err;
531    use vortex_io::runtime::BlockingRuntime;
532    use vortex_io::runtime::single::SingleThreadRuntime;
533    use vortex_mask::Mask;
534
535    use super::ScanBuilder;
536    use super::referenced_field_masks;
537    use crate::ArrayFuture;
538    use crate::LayoutReader;
539    use crate::RowSplits;
540    use crate::SplitRange;
541    use crate::scan::test::SCAN_SESSION;
542    use crate::scan::test::session_with_handle;
543
544    fn nested_dtype() -> DType {
545        DType::Struct(
546            StructFields::from_iter([
547                (
548                    "a",
549                    DType::Struct(
550                        StructFields::from_iter([
551                            ("1", DType::Primitive(PType::I32, Nullability::NonNullable)),
552                            ("2", DType::Primitive(PType::I32, Nullability::NonNullable)),
553                        ]),
554                        Nullability::NonNullable,
555                    ),
556                ),
557                ("b", DType::Primitive(PType::I32, Nullability::NonNullable)),
558            ]),
559            Nullability::NonNullable,
560        )
561    }
562
563    #[test]
564    fn bound_setters_preserve_identity() -> VortexResult<()> {
565        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
566        let projection = eq(root(), lit(1_i32)).bind(&dtype)?;
567        let filter = eq(root(), lit(2_i32)).bind(&dtype)?;
568        let expected_projection = ExactBoundExpr(projection.clone());
569        let expected_filter = ExactBoundExpr(filter.clone());
570        let reader = Arc::new(CountingLayoutReader::new(Arc::new(AtomicUsize::new(0))));
571
572        let builder = ScanBuilder::new(SCAN_SESSION.clone(), reader)
573            .with_projection(projection)
574            .with_filter(filter);
575
576        assert_eq!(ExactBoundExpr(builder.projection), expected_projection);
577        assert_eq!(builder.filter.map(ExactBoundExpr), Some(expected_filter));
578        Ok(())
579    }
580
581    #[test]
582    fn root_projection_produces_all_mask() -> VortexResult<()> {
583        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
584        let projection = root().bind(&dtype)?;
585
586        assert_eq!(referenced_field_masks(&projection, None)?, [FieldMask::All]);
587        Ok(())
588    }
589
590    #[test]
591    fn nested_projection_preserves_field_path_in_split_mask() -> VortexResult<()> {
592        let dtype = nested_dtype();
593        let projection = get_item("1", get_item("a", root())).bind(&dtype)?;
594        let filter = eq(get_item("2", get_item("a", root())), lit(0_i32)).bind(&dtype)?;
595
596        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
597
598        assert_eq!(field_masks.len(), 2);
599        assert!(field_masks.contains(&FieldMask::Prefix(FieldPath::from_name("a").push("1"))));
600        assert!(field_masks.contains(&FieldMask::Prefix(FieldPath::from_name("a").push("2"))));
601        Ok(())
602    }
603
604    #[test]
605    fn filter_path_covers_nested_projection_path() -> VortexResult<()> {
606        let dtype = nested_dtype();
607        let projection = get_item("1", get_item("a", root())).bind(&dtype)?;
608        let filter = is_not_null(get_item("a", root())).bind(&dtype)?;
609
610        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
611
612        assert_eq!(field_masks, [FieldMask::Prefix(FieldPath::from_name("a"))]);
613        Ok(())
614    }
615
616    #[test]
617    fn parent_projection_path_covers_nested_filter_path() -> VortexResult<()> {
618        let dtype = nested_dtype();
619        let projection = get_item("a", root()).bind(&dtype)?;
620        let filter = is_not_null(get_item("1", get_item("a", root()))).bind(&dtype)?;
621
622        let field_masks = referenced_field_masks(&projection, Some(&filter))?;
623
624        assert_eq!(field_masks, [FieldMask::Prefix(FieldPath::from_name("a"))]);
625        Ok(())
626    }
627
628    #[derive(Debug)]
629    struct CountingLayoutReader {
630        name: Arc<str>,
631        dtype: DType,
632        row_count: u64,
633        register_splits_calls: Arc<AtomicUsize>,
634    }
635
636    impl CountingLayoutReader {
637        fn new(register_splits_calls: Arc<AtomicUsize>) -> Self {
638            Self {
639                name: Arc::from("counting"),
640                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
641                row_count: 1,
642                register_splits_calls,
643            }
644        }
645    }
646
647    impl LayoutReader for CountingLayoutReader {
648        fn name(&self) -> &Arc<str> {
649            &self.name
650        }
651
652        fn dtype(&self) -> &DType {
653            &self.dtype
654        }
655
656        fn row_count(&self) -> u64 {
657            self.row_count
658        }
659
660        fn register_splits(
661            &self,
662            _field_mask: &[FieldMask],
663            split_range: &SplitRange,
664            splits: &mut RowSplits,
665        ) -> VortexResult<()> {
666            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
667            splits.push(split_range.root_row_range().end);
668            Ok(())
669        }
670
671        fn pruning_evaluation(
672            &self,
673            _row_range: &Range<u64>,
674            _expr: &BoundExpression,
675            _mask: Mask,
676        ) -> VortexResult<MaskFuture> {
677            vortex_bail!("not needed for this test");
678        }
679
680        fn filter_evaluation(
681            &self,
682            _row_range: &Range<u64>,
683            _expr: &BoundExpression,
684            _mask: MaskFuture,
685        ) -> VortexResult<MaskFuture> {
686            vortex_bail!("not needed for this test");
687        }
688
689        fn projection_evaluation(
690            &self,
691            _row_range: &Range<u64>,
692            _expr: &BoundExpression,
693            _mask: MaskFuture,
694        ) -> VortexResult<ArrayFuture> {
695            Ok(Box::pin(async move {
696                unreachable!("scan should not be polled in this test")
697            }))
698        }
699
700        fn as_any(&self) -> &dyn std::any::Any {
701            self
702        }
703    }
704
705    #[test]
706    fn into_stream_is_lazy() {
707        let calls = Arc::new(AtomicUsize::new(0));
708        let reader = Arc::new(CountingLayoutReader::new(Arc::clone(&calls)));
709
710        let session = SCAN_SESSION.clone();
711
712        let _stream = ScanBuilder::new(session, reader).into_stream().unwrap();
713
714        assert_eq!(calls.load(Ordering::Relaxed), 0);
715    }
716
717    #[derive(Debug)]
718    struct SplittingLayoutReader {
719        name: Arc<str>,
720        dtype: DType,
721        row_count: u64,
722        register_splits_calls: Arc<AtomicUsize>,
723    }
724
725    impl SplittingLayoutReader {
726        fn new(register_splits_calls: Arc<AtomicUsize>) -> Self {
727            Self {
728                name: Arc::from("splitting"),
729                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
730                row_count: 4,
731                register_splits_calls,
732            }
733        }
734    }
735
736    impl LayoutReader for SplittingLayoutReader {
737        fn name(&self) -> &Arc<str> {
738            &self.name
739        }
740
741        fn dtype(&self) -> &DType {
742            &self.dtype
743        }
744
745        fn row_count(&self) -> u64 {
746            self.row_count
747        }
748
749        fn register_splits(
750            &self,
751            _field_mask: &[FieldMask],
752            split_range: &SplitRange,
753            splits: &mut RowSplits,
754        ) -> VortexResult<()> {
755            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
756            for split in (split_range.row_range().start + 1)..=split_range.row_range().end {
757                splits.push(split_range.row_offset() + split);
758            }
759            Ok(())
760        }
761
762        fn pruning_evaluation(
763            &self,
764            _row_range: &Range<u64>,
765            _expr: &BoundExpression,
766            mask: Mask,
767        ) -> VortexResult<MaskFuture> {
768            Ok(MaskFuture::ready(mask))
769        }
770
771        fn filter_evaluation(
772            &self,
773            _row_range: &Range<u64>,
774            _expr: &BoundExpression,
775            mask: MaskFuture,
776        ) -> VortexResult<MaskFuture> {
777            Ok(mask)
778        }
779
780        fn projection_evaluation(
781            &self,
782            row_range: &Range<u64>,
783            _expr: &BoundExpression,
784            _mask: MaskFuture,
785        ) -> VortexResult<ArrayFuture> {
786            let start = usize::try_from(row_range.start)
787                .map_err(|_| vortex_err!("row_range.start must fit in usize"))?;
788            let end = usize::try_from(row_range.end)
789                .map_err(|_| vortex_err!("row_range.end must fit in usize"))?;
790
791            let values: VortexResult<Vec<i32>> = (start..end)
792                .map(|v| i32::try_from(v).map_err(|_| vortex_err!("split value must fit in i32")))
793                .collect();
794
795            let array = PrimitiveArray::from_iter(values?).into_array();
796            Ok(Box::pin(async move { Ok(array) }))
797        }
798
799        fn as_any(&self) -> &dyn std::any::Any {
800            self
801        }
802    }
803
804    #[test]
805    fn into_stream_executes_after_prepare() -> VortexResult<()> {
806        let mut ctx = array_session().create_execution_ctx();
807        let calls = Arc::new(AtomicUsize::new(0));
808        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
809
810        let runtime = SingleThreadRuntime::default();
811        let session = session_with_handle(runtime.handle());
812
813        let stream = ScanBuilder::new(session, reader).into_stream()?;
814        let mut iter = runtime.block_on_stream(stream);
815
816        let mut values = Vec::new();
817        for chunk in &mut iter {
818            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
819            values.push(prim.into_buffer::<i32>()[0]);
820        }
821
822        assert_eq!(calls.load(Ordering::Relaxed), 1);
823        assert_eq!(values.as_ref(), [0, 1, 2, 3]);
824
825        Ok(())
826    }
827
828    #[test]
829    fn supplied_natural_splits_skip_layout_walk() -> VortexResult<()> {
830        let mut ctx = array_session().create_execution_ctx();
831        let calls = Arc::new(AtomicUsize::new(0));
832        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
833
834        let runtime = SingleThreadRuntime::default();
835        let session = session_with_handle(runtime.handle());
836
837        let stream = ScanBuilder::new(session, reader)
838            .with_natural_splits(vec![0u64, 2, 4].into())
839            .with_row_range(1..4)
840            .into_stream()?;
841        let mut iter = runtime.block_on_stream(stream);
842
843        let mut chunks = Vec::new();
844        for chunk in &mut iter {
845            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
846            chunks.push(prim.into_buffer::<i32>().to_vec());
847        }
848
849        assert_eq!(calls.load(Ordering::Relaxed), 0);
850        // Supplied full-file boundaries [0, 2, 4] clamped to rows 1..4.
851        assert_eq!(chunks, [vec![1], vec![2, 3]]);
852
853        Ok(())
854    }
855
856    #[test]
857    fn full_file_splits_ignore_row_range() -> VortexResult<()> {
858        let calls = Arc::new(AtomicUsize::new(0));
859        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
860
861        let splits = ScanBuilder::new(SCAN_SESSION.clone(), reader)
862            .with_row_range(1..3)
863            .full_file_splits()?;
864
865        assert_eq!(splits, [0, 1, 2, 3, 4]);
866        Ok(())
867    }
868
869    #[derive(Debug)]
870    struct BlockingSplitsLayoutReader {
871        name: Arc<str>,
872        dtype: DType,
873        row_count: u64,
874        register_splits_calls: Arc<AtomicUsize>,
875        gate: Arc<Mutex<()>>,
876    }
877
878    impl BlockingSplitsLayoutReader {
879        fn new(gate: Arc<Mutex<()>>, register_splits_calls: Arc<AtomicUsize>) -> Self {
880            Self {
881                name: Arc::from("blocking-splits"),
882                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
883                row_count: 1,
884                register_splits_calls,
885                gate,
886            }
887        }
888    }
889
890    impl LayoutReader for BlockingSplitsLayoutReader {
891        fn name(&self) -> &Arc<str> {
892            &self.name
893        }
894
895        fn dtype(&self) -> &DType {
896            &self.dtype
897        }
898
899        fn row_count(&self) -> u64 {
900            self.row_count
901        }
902
903        fn register_splits(
904            &self,
905            _field_mask: &[FieldMask],
906            split_range: &SplitRange,
907            splits: &mut RowSplits,
908        ) -> VortexResult<()> {
909            self.register_splits_calls.fetch_add(1, Ordering::Relaxed);
910            let _guard = self.gate.lock();
911            splits.push(split_range.root_row_range().end);
912            Ok(())
913        }
914
915        fn pruning_evaluation(
916            &self,
917            _row_range: &Range<u64>,
918            _expr: &BoundExpression,
919            _mask: Mask,
920        ) -> VortexResult<MaskFuture> {
921            vortex_bail!("not needed for this test");
922        }
923
924        fn filter_evaluation(
925            &self,
926            _row_range: &Range<u64>,
927            _expr: &BoundExpression,
928            _mask: MaskFuture,
929        ) -> VortexResult<MaskFuture> {
930            vortex_bail!("not needed for this test");
931        }
932
933        fn projection_evaluation(
934            &self,
935            _row_range: &Range<u64>,
936            _expr: &BoundExpression,
937            _mask: MaskFuture,
938        ) -> VortexResult<ArrayFuture> {
939            Ok(Box::pin(async move {
940                unreachable!("scan should not be polled in this test")
941            }))
942        }
943
944        fn as_any(&self) -> &dyn std::any::Any {
945            self
946        }
947    }
948
949    #[test]
950    fn into_stream_first_poll_does_not_block() {
951        let gate = Arc::new(Mutex::new(()));
952        let guard = gate.lock();
953
954        let calls = Arc::new(AtomicUsize::new(0));
955        let reader = Arc::new(BlockingSplitsLayoutReader::new(
956            Arc::clone(&gate),
957            Arc::clone(&calls),
958        ));
959
960        let runtime = SingleThreadRuntime::default();
961        let session = session_with_handle(runtime.handle());
962
963        let mut stream = ScanBuilder::new(session, reader).into_stream().unwrap();
964
965        let (send, recv) = mpsc::channel::<bool>();
966        let join = std::thread::spawn(move || {
967            let waker = noop_waker_ref();
968            let mut cx = Context::from_waker(waker);
969            let poll = Pin::new(&mut stream).poll_next(&mut cx);
970            let _ = send.send(matches!(poll, Poll::Pending));
971        });
972
973        let polled_pending = recv.recv_timeout(Duration::from_secs(1)).ok();
974
975        // Always release the gate and join the thread so failures don't hang the test process.
976        drop(guard);
977        drop(join.join());
978
979        let polled_pending = polled_pending.expect("poll_next blocked; expected quick return");
980        assert!(
981            polled_pending,
982            "expected Poll::Pending while prepare is blocked"
983        );
984        assert_eq!(calls.load(Ordering::Relaxed), 0);
985
986        drop(runtime);
987    }
988
989    #[test]
990    fn into_stream_with_row_range() -> VortexResult<()> {
991        let mut ctx = array_session().create_execution_ctx();
992        let calls = Arc::new(AtomicUsize::new(0));
993        let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls)));
994
995        let runtime = SingleThreadRuntime::default();
996        let session = session_with_handle(runtime.handle());
997
998        let stream = ScanBuilder::new(session, reader)
999            .with_row_range(1..3)
1000            .into_stream()?;
1001        let mut iter = runtime.block_on_stream(stream);
1002
1003        let mut values = Vec::new();
1004        for chunk in &mut iter {
1005            let prim = chunk?.execute::<PrimitiveArray>(&mut ctx)?;
1006            values.extend(prim.into_buffer::<i32>().iter().copied());
1007        }
1008
1009        assert_eq!(calls.load(Ordering::Relaxed), 1);
1010        assert_eq!(values.as_ref(), [1, 2]);
1011
1012        Ok(())
1013    }
1014}
vortex_layout/scan/scan_builder.rs

vortex_layout/scan/
scan_builder.rs
