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

vortex_layout/scan/
scan_builder.rs
