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

vortex_layout/scan/v2/
mod.rs
