vortex_layout/scan/v2/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! The planner-driven scan executor.
5//!
6//! A V2 scan lowers the file's layout once into a physical plan and runs every split of it
7//! through the planning protocol: the [`Run`](vortex_scan::planning::driver::Run) driver advances
8//! the split's planners and morsels, and every segment read is a protocol request served from the
9//! file's [`SegmentSource`]. The layout reader is never asked to evaluate anything.
10//!
11//! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`]
12//! directly and supports projection, row ranges, and selections. Unsupported layouts return an
13//! error; filters and limits are not part of this builder's API. Work is partitioned into fixed
14//! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries.
15//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high.
16
17mod io;
18mod split;
19#[cfg(test)]
20mod tests;
21
22use std::ops::Range;
23use std::sync::Arc;
24
25use arrow_array::RecordBatch;
26use arrow_schema::ArrowError;
27use arrow_schema::Field;
28use arrow_schema::SchemaRef;
29use futures::Stream;
30use futures::StreamExt;
31use futures::TryStreamExt;
32use futures::stream;
33use futures::stream::BoxStream;
34use vortex_array::ArrayRef;
35use vortex_array::VortexSessionExecute;
36use vortex_array::dtype::DType;
37use vortex_array::expr::BoundExpression;
38use vortex_error::VortexResult;
39use vortex_io::request::IoService;
40use vortex_io::session::RuntimeSessionExt;
41use vortex_scan::planning::planner::WorkScope;
42use vortex_scan::selection::Selection;
43use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
44use vortex_session::VortexSession;
45use vortex_utils::parallelism::get_available_parallelism;
46
47use crate::LayoutRef;
48use crate::plan::physical::lower;
49use crate::plan::physical::optimize;
50use crate::plan::physical::plan_row_idx_expression;
51use crate::scan::arrow::to_record_batch;
52use crate::scan::execution::ScanPlans;
53pub use crate::scan::execution::SegmentLocation;
54use crate::scan::execution::plan_selection;
55use crate::scan::v2::io::SegmentScanIo;
56use crate::scan::v2::io::segment_ranges;
57use crate::scan::v2::split::SplitPlan;
58use crate::segments::SegmentSource;
59
60const ROW_BLOCK_SIZE: usize = 8192;
61const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE;
62const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024;
63
64/// Limit estimated bytes per task without changing the writer-aligned row granularity.
65fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize {
66    if row_count <= MAX_ROWS_PER_CHUNK as u64 {
67        return MAX_ROWS_PER_CHUNK;
68    }
69    let bytes: u64 = locations
70        .iter()
71        .map(|location| u64::from(location.length))
72        .sum();
73    let bytes_per_row = bytes.div_ceil(row_count).max(1);
74    let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64;
75    blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE
76}
77
78/// Fixed boundaries in the file's row domain, clipped to the requested range.
79fn row_chunks(rows: Range<u64>, rows_per_chunk: usize) -> Vec<Range<u64>> {
80    if rows.is_empty() {
81        return Vec::new();
82    }
83    let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64;
84    (first..rows.end)
85        .step_by(rows_per_chunk)
86        .map(|start| {
87            start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end)
88        })
89        .collect()
90}
91
92/// The parts of a file the planner-driven executor reads through.
93#[derive(Clone)]
94pub struct ScanFile {
95    /// The root layout to scan.
96    pub layout: LayoutRef,
97    /// The location of every segment, indexed by segment id.
98    pub locations: Arc<[SegmentLocation]>,
99    /// Fetches segment bytes by id.
100    pub segments: Arc<dyn SegmentSource>,
101}
102
103/// Builds a projection scan on the planner-driven V2 executor.
104///
105/// This builder reads the file's layout and segments directly. Unsupported layouts return an
106/// error when the stream is polled. Filters and row limits are not supported.
107pub struct ScanBuilder<A = ArrayRef> {
108    session: VortexSession,
109    file: ScanFile,
110    projection: BoundExpression,
111    ordered: bool,
112    row_range: Option<Range<u64>>,
113    selection: Selection,
114    concurrency: usize,
115    map_fn: Arc<dyn Fn(ArrayRef) -> VortexResult<A> + Send + Sync>,
116    row_offset: u64,
117}
118
119impl ScanBuilder<ArrayRef> {
120    /// Creates a V2 scan over the file's layout and segment source.
121    pub fn new(session: VortexSession, file: ScanFile) -> Self {
122        Self {
123            projection: BoundExpression::new_root(file.layout.dtype().clone()),
124            session,
125            file,
126            ordered: true,
127            row_range: None,
128            selection: Selection::default(),
129            concurrency: 4,
130            map_fn: Arc::new(Ok),
131            row_offset: 0,
132        }
133    }
134
135    /// Converts each output batch to Arrow on the split's task.
136    pub fn into_record_batch_stream(
137        self,
138        schema: SchemaRef,
139    ) -> VortexResult<impl Stream<Item = Result<RecordBatch, ArrowError>> + Send + 'static> {
140        let struct_field = Field::new_struct("", schema.fields().clone(), false);
141        let session = self.session.clone();
142        Ok(self
143            .map(move |chunk| {
144                let mut ctx = session.create_execution_ctx();
145                to_record_batch(chunk, &struct_field, &mut ctx)
146            })
147            .into_stream()?
148            .map_err(|error| ArrowError::ExternalError(Box::new(error))))
149    }
150}
151
152impl<A: 'static + Send> ScanBuilder<A> {
153    /// The output dtype after projection, before any output mapping.
154    pub fn dtype(&self) -> &DType {
155        self.projection.dtype()
156    }
157
158    /// Sets a projection bound against the file's dtype.
159    pub fn with_projection(mut self, projection: BoundExpression) -> Self {
160        self.projection = projection;
161        self
162    }
163
164    /// Configures whether batches are yielded in file order.
165    pub fn with_ordered(mut self, ordered: bool) -> Self {
166        self.ordered = ordered;
167        self
168    }
169
170    /// Restricts scanning to a contiguous file-local row range.
171    pub fn with_row_range(mut self, row_range: Range<u64>) -> Self {
172        self.row_range = Some(row_range);
173        self
174    }
175
176    /// Applies a row selection within the scan's row range.
177    pub fn with_selection(mut self, selection: Selection) -> Self {
178        self.selection = selection;
179        self
180    }
181
182    /// Selects rows by strictly sorted file-local indices.
183    pub fn with_row_indices(self, row_indices: StrictSortedBuffer<u64>) -> Self {
184        self.with_selection(Selection::IncludeByIndex(row_indices))
185    }
186
187    /// Sets the root row offset used by row-index expressions.
188    pub fn with_row_offset(mut self, row_offset: u64) -> Self {
189        self.row_offset = row_offset;
190        self
191    }
192
193    /// Sets the number of active splits per worker thread. Must be greater than zero.
194    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
195        assert!(concurrency > 0);
196        self.concurrency = concurrency;
197        self
198    }
199
200    /// Maps each output batch on its split's task.
201    fn map<B: 'static>(
202        self,
203        map_fn: impl Fn(A) -> VortexResult<B> + 'static + Send + Sync,
204    ) -> ScanBuilder<B> {
205        let old_map_fn = self.map_fn;
206        ScanBuilder {
207            session: self.session,
208            file: self.file,
209            projection: self.projection,
210            ordered: self.ordered,
211            row_range: self.row_range,
212            selection: self.selection,
213            concurrency: self.concurrency,
214            row_offset: self.row_offset,
215            map_fn: Arc::new(move |array| old_map_fn(array).and_then(&map_fn)),
216        }
217    }
218
219    /// Runs the scan on the session's runtime, preparing it on the first poll.
220    /// Preparation errors are returned as stream items.
221    pub fn into_stream(self) -> VortexResult<BoxStream<'static, VortexResult<A>>> {
222        let handle = self.session.handle();
223        Ok(stream::once(async move {
224            handle
225                .spawn_cpu(move || PreparedScan::try_new(self)?.execute_stream())
226                .await
227        })
228        .try_flatten()
229        .boxed())
230    }
231}
232
233/// A scan planned over its file, ready to run its splits.
234struct PreparedScan<A: 'static + Send> {
235    session: VortexSession,
236    plans: ScanPlans,
237    /// Serves the splits' reads.
238    io: Arc<dyn IoService>,
239    /// The scan's splits, over the root row domain, in row order.
240    splits: Vec<Range<u64>>,
241    selection: Selection,
242    map_fn: Arc<dyn Fn(ArrayRef) -> VortexResult<A> + Send + Sync>,
243    ordered: bool,
244    concurrency: usize,
245}
246
247impl<A: 'static + Send> PreparedScan<A> {
248    /// Lowers the file's layout, plans the projection, and partitions the selected range into
249    /// fixed-size row chunks.
250    fn try_new(builder: ScanBuilder<A>) -> VortexResult<Self> {
251        let file = builder.file;
252        let root = lower(&file.layout)?;
253        let row_count = root.row_count();
254        // An identity projection needs no rewrites. Keep the lowered children lazy, including
255        // statistics that a projection-only scan never reads.
256        let identity = builder.projection.is_root() && builder.projection.dtype() == root.dtype();
257        let projection = if identity {
258            root
259        } else {
260            optimize(plan_row_idx_expression(builder.projection, root)?)?
261        };
262        let chunk_rows = if identity {
263            rows_per_chunk(row_count, &file.locations)
264        } else {
265            MAX_ROWS_PER_CHUNK
266        };
267        let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows);
268        let io = Arc::new(SegmentScanIo::new(
269            file.segments,
270            segment_ranges(&file.locations),
271        ));
272        Ok(Self {
273            session: builder.session.clone(),
274            plans: ScanPlans {
275                session: builder.session,
276                locations: file.locations,
277                projection,
278                row_offset: builder.row_offset,
279                values: Default::default(),
280            },
281            io,
282            splits,
283            selection: builder.selection,
284            map_fn: builder.map_fn,
285            ordered: builder.ordered,
286            concurrency: builder.concurrency,
287        })
288    }
289
290    /// The splits with selected rows, each ready to be admitted to a planning run.
291    fn split_plans(&self) -> Vec<SplitPlan> {
292        self.splits
293            .iter()
294            .filter_map(|rows| {
295                let (scope, root) = plan_selection(
296                    &self.plans,
297                    WorkScope {
298                        file_ordinal: 0,
299                        rows: rows.clone(),
300                    },
301                    &self.selection,
302                )?;
303                Some(SplitPlan { scope, root })
304            })
305            .collect()
306    }
307
308    /// Runs every split as a task on the session's runtime and yields their mapped batches.
309    ///
310    /// Like the default executor, up to `concurrency` splits per thread run at once, and their
311    /// batches are yielded in row order when the scan is ordered.
312    fn execute_stream(self) -> VortexResult<BoxStream<'static, VortexResult<A>>> {
313        let concurrency = self.concurrency * get_available_parallelism().unwrap_or(1);
314        let handle = self.session.handle();
315        let split_plans = self.split_plans();
316        let io = self.io;
317        let map_fn = self.map_fn;
318        // Spawned as the buffer pulls them, so at most `concurrency` splits read at once.
319        let tasks = stream::iter(split_plans).map(move |split| {
320            let io = Arc::clone(&io);
321            let map_fn = Arc::clone(&map_fn);
322            handle.spawn(async move {
323                split
324                    .run(io)
325                    .await?
326                    .into_iter()
327                    .map(|array| map_fn(array))
328                    .collect::<VortexResult<Vec<_>>>()
329            })
330        });
331        let batches = if self.ordered {
332            tasks.buffered(concurrency).boxed()
333        } else {
334            tasks.buffer_unordered(concurrency).boxed()
335        };
336        Ok(batches
337            .map_ok(|batches| stream::iter(batches.into_iter().map(Ok)))
338            .try_flatten()
339            .boxed())
340    }
341}
vortex_layout/scan/v2/mod.rs

vortex_layout/scan/v2/
mod.rs
