vortex_datafusion/persistent/access_plan.rs
1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex::layout::scan::scan_builder::ScanBuilder;
5use vortex::scan::selection::Selection;
6
7/// Additional Vortex-specific scan constraints attached to a
8/// [`PartitionedFile`].
9///
10/// `VortexAccessPlan` is the hook to use when an external index or planner
11/// already knows that only part of a file needs to be scanned. The plan is
12/// attached as `extensions` on `PartitionedFile`, and the internal
13/// `VortexOpener` applies it before building the Vortex scan.
14///
15/// The current access plan surface is intentionally small: it lets callers
16/// provide a [`Selection`] that narrows the rows considered by the scan.
17///
18/// # Example
19///
20/// ```no_run
21/// # use std::sync::Arc;
22/// # use datafusion_datasource::PartitionedFile;
23/// # use vortex::scan::selection::Selection;
24/// use vortex_datafusion::VortexAccessPlan;
25///
26/// # let selection: Selection = todo!();
27/// let file = PartitionedFile::new("metrics.vortex", 1024).with_extensions(Arc::new(
28/// VortexAccessPlan::default().with_selection(selection),
29/// ));
30/// # let _ = file;
31/// ```
32///
33/// This is a low-level integration point for systems building their own access
34/// paths on top of DataFusion. For a conceptually similar Parquet example, see
35/// DataFusion's
36/// [`parquet_advanced_index`].
37///
38/// [`PartitionedFile`]: datafusion_datasource::PartitionedFile
39/// [`parquet_advanced_index`]: https://github.com/apache/datafusion/blob/47df535d2cd5aac5ad5a92bdc837f38e05ea0f0f/datafusion-examples/examples/data_io/parquet_advanced_index.rs
40#[derive(Default)]
41pub struct VortexAccessPlan {
42 selection: Option<Selection>,
43}
44
45impl VortexAccessPlan {
46 /// Sets the row [`Selection`] to apply when the file is opened.
47 pub fn with_selection(mut self, selection: Selection) -> Self {
48 self.selection = Some(selection);
49 self
50 }
51}
52
53impl VortexAccessPlan {
54 /// Returns the selection, if one was set.
55 pub fn selection(&self) -> Option<&Selection> {
56 self.selection.as_ref()
57 }
58
59 /// Applies this access plan to a [`ScanBuilder`].
60 ///
61 /// This is used internally by the file opener after it has translated a
62 /// `PartitionedFile` into a Vortex scan.
63 pub fn apply_to_builder<A>(&self, mut scan_builder: ScanBuilder<A>) -> ScanBuilder<A>
64 where
65 A: 'static + Send,
66 {
67 let Self { selection } = self;
68
69 if let Some(selection) = selection {
70 scan_builder = scan_builder.with_selection(selection.clone());
71 }
72
73 scan_builder
74 }
75}