vortex_file/
file.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! This module defines the [`VortexFile`] struct, which represents a Vortex file on disk or in memory.
5//!
6//! The `VortexFile` provides methods for accessing file metadata, creating segment sources for reading
7//! data from the file, and initiating scans to read the file's contents into memory as Vortex arrays.
8
9use std::ops::Range;
10use std::sync::Arc;
11use std::sync::OnceLock;
12
13use itertools::Itertools;
14use vortex_array::ArrayRef;
15use vortex_array::dtype::DType;
16use vortex_array::dtype::FieldMask;
17use vortex_array::expr::Expression;
18use vortex_buffer::ByteBuffer;
19use vortex_error::VortexResult;
20use vortex_layout::LayoutReader;
21use vortex_layout::scan::layout::LayoutReaderDataSource;
22use vortex_layout::scan::scan_builder::ScanBuilder;
23use vortex_layout::scan::split_by::SplitBy;
24use vortex_layout::scan::v2::ScanCache;
25use vortex_layout::scan::v2::ScanFile;
26use vortex_layout::scan::v2::SegmentLocation;
27use vortex_layout::segments::SegmentSource;
28use vortex_scan::DataSourceRef;
29use vortex_session::VortexSession;
30use vortex_utils::aliases::hash_map::HashMap;
31
32use crate::FileStatistics;
33use crate::footer::Footer;
34use crate::pruning::can_prune_file_stats;
35use crate::v2::FileStatsLayoutReader;
36
37/// Represents a Vortex file, providing access to its metadata and content.
38///
39/// A `VortexFile` is created by opening a Vortex file using [`VortexOpenOptions`](crate::VortexOpenOptions).
40/// It provides methods for accessing file metadata (such as row count, data type, and statistics)
41/// and for initiating scans to read the file's contents.
42#[derive(Clone)]
43pub struct VortexFile {
44    /// The footer of the Vortex file, containing metadata and layout information.
45    footer: Footer,
46    /// The segment source used to read segments from this file.
47    segment_source: Arc<dyn SegmentSource>,
48    /// The Vortex session used to open this file.
49    session: VortexSession,
50    /// User-defined metadata values resolved for this file open.
51    metadata: Arc<HashMap<String, ByteBuffer>>,
52    /// None id LayoutReader caching is turned off
53    layout_reader_cache: Option<OnceLock<Arc<dyn LayoutReader>>>,
54    /// None if V2 scan caching is turned off
55    scan_cache: Option<FileScanCache>,
56}
57
58/// What repeated V2 scans of a file share: its segment locations, read from the footer once, and
59/// the plan state derived from them.
60#[derive(Clone)]
61struct FileScanCache {
62    locations: Arc<[SegmentLocation]>,
63    plans: ScanCache,
64}
65
66fn segment_locations(footer: &Footer) -> Arc<[SegmentLocation]> {
67    footer
68        .segment_specs_with_metadata()
69        .iter()
70        .map(|spec| SegmentLocation {
71            offset: spec.offset,
72            length: spec.length,
73            alignment: spec.alignment,
74        })
75        .collect()
76}
77
78fn layout_reader(
79    segment_source: Arc<dyn SegmentSource>,
80    footer: &Footer,
81    session: &VortexSession,
82) -> VortexResult<Arc<dyn LayoutReader>> {
83    let root_reader = footer
84        .layout()
85        // TODO(ngates): we may want to allow the user pass in a name here?
86        .new_reader("".into(), segment_source, session, &Default::default())?;
87
88    Ok(if let Some(stats) = footer.statistics().cloned() {
89        Arc::new(FileStatsLayoutReader::new(
90            root_reader,
91            stats,
92            session.clone(),
93        ))
94    } else {
95        root_reader
96    })
97}
98
99impl VortexFile {
100    /// Creates a new `VortexFile` from the given footer, segment source, and session.
101    pub fn new(
102        footer: Footer,
103        segment_source: Arc<dyn SegmentSource>,
104        session: VortexSession,
105    ) -> Self {
106        Self {
107            footer,
108            segment_source,
109            session,
110            metadata: Arc::new(HashMap::new()),
111            layout_reader_cache: None,
112            scan_cache: None,
113        }
114    }
115
116    pub(crate) fn with_metadata(mut self, metadata: Arc<HashMap<String, ByteBuffer>>) -> Self {
117        self.metadata = metadata;
118        self
119    }
120
121    /// Enable layout reader and V2 scan caching.
122    ///
123    /// Repeated calls to [`layout_reader`](Self::layout_reader), [`scan`](Self::scan), and
124    /// [`data_source`](Self::data_source) will share the same reader tree. Repeated V2 scans of
125    /// [`scan_file`](Self::scan_file) will share its segment locations and a [`ScanCache`].
126    pub fn with_caching(self) -> Self {
127        let scan_cache = FileScanCache {
128            locations: segment_locations(&self.footer),
129            plans: ScanCache::default(),
130        };
131        Self {
132            footer: self.footer,
133            segment_source: self.segment_source,
134            session: self.session,
135            metadata: self.metadata,
136            layout_reader_cache: Some(OnceLock::new()),
137            scan_cache: Some(scan_cache),
138        }
139    }
140
141    /// Returns a reference to the file's footer, which contains metadata and layout information.
142    pub fn footer(&self) -> &Footer {
143        &self.footer
144    }
145
146    /// Returns the number of rows in the file.
147    pub fn row_count(&self) -> u64 {
148        self.footer.row_count()
149    }
150
151    /// Returns the data type of the file's contents.
152    pub fn dtype(&self) -> &DType {
153        self.footer.dtype()
154    }
155
156    /// Returns the file's statistics, if available.
157    ///
158    /// Statistics can be used for query optimization and data exploration.
159    pub fn file_stats(&self) -> Option<&FileStatistics> {
160        self.footer.statistics()
161    }
162
163    /// Returns the user-defined metadata segments loaded for this file.
164    ///
165    /// Metadata is only loaded when requested during open. Iteration order is unspecified.
166    pub fn metadata_segments(&self) -> impl Iterator<Item = (&str, &ByteBuffer)> {
167        self.metadata
168            .iter()
169            .map(|(key, metadata)| (key.as_str(), metadata))
170    }
171
172    /// Returns the loaded user-defined metadata segment for the given key.
173    ///
174    /// Returns `None` when the key is absent or metadata was not loaded.
175    pub fn metadata_segment(&self, key: &str) -> Option<&ByteBuffer> {
176        self.metadata.get(key)
177    }
178
179    /// Create a new segment source for reading from the file.
180    ///
181    /// This may spawn a background I/O driver that will exit when the returned segment source
182    /// is dropped.
183    pub fn segment_source(&self) -> Arc<dyn SegmentSource> {
184        Arc::clone(&self.segment_source)
185    }
186
187    /// Replace the segment source used by this file.
188    ///
189    /// Any cached layout reader or V2 scan state is cleared so that subsequent scans read from
190    /// the replacement source.
191    pub fn with_segment_source(mut self, segment_source: Arc<dyn SegmentSource>) -> Self {
192        self.segment_source = segment_source;
193        if self.layout_reader_cache.is_some() {
194            self.layout_reader_cache = Some(OnceLock::new());
195        }
196        if let Some(cache) = &mut self.scan_cache {
197            cache.plans = ScanCache::default();
198        }
199        self
200    }
201
202    /// Returns a reference to the Vortex session used to open this file.
203    pub fn session(&self) -> &VortexSession {
204        &self.session
205    }
206
207    /// Create a new layout reader for the file.
208    ///
209    /// Wraps the root layout in a [`FileStatsLayoutReader`] if file stats are available.
210    pub fn layout_reader(&self) -> VortexResult<Arc<dyn LayoutReader>> {
211        match &self.layout_reader_cache {
212            None => layout_reader(
213                Arc::clone(&self.segment_source),
214                &self.footer,
215                &self.session,
216            ),
217            Some(reader) => {
218                // get_or_try_init is unstable
219                if let Some(val) = reader.get() {
220                    Ok(Arc::clone(val))
221                } else {
222                    let inner = layout_reader(
223                        Arc::clone(&self.segment_source),
224                        &self.footer,
225                        &self.session,
226                    )?;
227                    Ok(if let Err(val) = reader.set(Arc::clone(&inner)) {
228                        val
229                    } else {
230                        inner
231                    })
232                }
233            }
234        }
235    }
236
237    /// Create a [`DataSource`](vortex_scan::DataSource) from this file for scanning.
238    ///
239    /// Wraps the file's layout reader with [`FileStatsLayoutReader`] (when file-level
240    /// statistics are available) and [`LayoutReaderDataSource`].
241    pub fn data_source(&self) -> VortexResult<DataSourceRef> {
242        let reader = self.layout_reader()?;
243
244        Ok(Arc::new(LayoutReaderDataSource::new(
245            reader,
246            self.session.clone(),
247        )))
248    }
249
250    /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and
251    /// execution options.
252    pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
253        Ok(ScanBuilder::new(
254            self.session.clone(),
255            self.layout_reader()?,
256        ))
257    }
258
259    /// The parts of this file the planner-driven scan executor reads through: its root layout,
260    /// where each of its segments lives, and its segment source.
261    ///
262    /// Pass this to [`vortex_layout::scan::v2::ScanBuilder::new`] to build a V2 scan explicitly.
263    /// With [caching](Self::with_caching), every returned [`ScanFile`] shares one [`ScanCache`].
264    pub fn scan_file(&self) -> ScanFile {
265        let (locations, cache) = match &self.scan_cache {
266            Some(cache) => (Arc::clone(&cache.locations), Some(cache.plans.clone())),
267            None => (segment_locations(&self.footer), None),
268        };
269        ScanFile {
270            layout: Arc::clone(self.footer.layout()),
271            locations,
272            segments: Arc::clone(&self.segment_source),
273            cache,
274        }
275    }
276
277    /// Returns `true` if file-level statistics prove the expression cannot
278    /// match any rows in this file.
279    ///
280    /// Row-count-aware pruning predicates are evaluated with the file's total
281    /// row count as their scope.
282    pub fn can_prune(&self, filter: &Expression) -> VortexResult<bool> {
283        let Some((stats, fields)) = self
284            .footer
285            .statistics()
286            .zip(self.footer.dtype().as_struct_fields_opt())
287        else {
288            return Ok(false);
289        };
290
291        can_prune_file_stats(
292            &filter.bind(self.footer.dtype())?,
293            self.footer.row_count(),
294            stats,
295            fields,
296            &self.session,
297        )
298    }
299
300    /// Return the file's natural row splits as root-coordinate ranges.
301    ///
302    /// These are the ranges that the default [`SplitBy`] would use for an all-fields scan.
303    pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
304        let reader = self.layout_reader()?;
305        Ok(SplitBy::default()
306            .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
307            .into_iter()
308            .tuple_windows()
309            .map(|(start, end)| start..end)
310            .collect())
311    }
312}
vortex_file/file.rs

vortex_file/
file.rs
