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::segments::SegmentSource;
25use vortex_scan::DataSourceRef;
26use vortex_session::VortexSession;
27use vortex_utils::aliases::hash_map::HashMap;
28
29use crate::FileStatistics;
30use crate::footer::Footer;
31use crate::pruning::can_prune_file_stats;
32use crate::v2::FileStatsLayoutReader;
33
34/// Represents a Vortex file, providing access to its metadata and content.
35///
36/// A `VortexFile` is created by opening a Vortex file using [`VortexOpenOptions`](crate::VortexOpenOptions).
37/// It provides methods for accessing file metadata (such as row count, data type, and statistics)
38/// and for initiating scans to read the file's contents.
39#[derive(Clone)]
40pub struct VortexFile {
41    /// The footer of the Vortex file, containing metadata and layout information.
42    footer: Footer,
43    /// The segment source used to read segments from this file.
44    segment_source: Arc<dyn SegmentSource>,
45    /// The Vortex session used to open this file.
46    session: VortexSession,
47    /// User-defined metadata values resolved for this file open.
48    metadata: Arc<HashMap<String, ByteBuffer>>,
49    /// None id LayoutReader caching is turned off
50    layout_reader_cache: Option<OnceLock<Arc<dyn LayoutReader>>>,
51}
52
53fn layout_reader(
54    segment_source: Arc<dyn SegmentSource>,
55    footer: &Footer,
56    session: &VortexSession,
57) -> VortexResult<Arc<dyn LayoutReader>> {
58    let root_reader = footer
59        .layout()
60        // TODO(ngates): we may want to allow the user pass in a name here?
61        .new_reader("".into(), segment_source, session, &Default::default())?;
62
63    Ok(if let Some(stats) = footer.statistics().cloned() {
64        Arc::new(FileStatsLayoutReader::new(
65            root_reader,
66            stats,
67            session.clone(),
68        ))
69    } else {
70        root_reader
71    })
72}
73
74impl VortexFile {
75    /// Creates a new `VortexFile` from the given footer, segment source, and session.
76    pub fn new(
77        footer: Footer,
78        segment_source: Arc<dyn SegmentSource>,
79        session: VortexSession,
80    ) -> Self {
81        Self {
82            footer,
83            segment_source,
84            session,
85            metadata: Arc::new(HashMap::new()),
86            layout_reader_cache: None,
87        }
88    }
89
90    pub(crate) fn with_metadata(mut self, metadata: Arc<HashMap<String, ByteBuffer>>) -> Self {
91        self.metadata = metadata;
92        self
93    }
94
95    /// Enable layout reader caching.
96    ///
97    /// Repeated calls to [`layout_reader`](Self::layout_reader), [`scan`](Self::scan), and
98    /// [`data_source`](Self::data_source) will share the same reader tree.
99    pub fn with_caching(self) -> Self {
100        Self {
101            footer: self.footer,
102            segment_source: self.segment_source,
103            session: self.session,
104            metadata: self.metadata,
105            layout_reader_cache: Some(OnceLock::new()),
106        }
107    }
108
109    /// Returns a reference to the file's footer, which contains metadata and layout information.
110    pub fn footer(&self) -> &Footer {
111        &self.footer
112    }
113
114    /// Returns the number of rows in the file.
115    pub fn row_count(&self) -> u64 {
116        self.footer.row_count()
117    }
118
119    /// Returns the data type of the file's contents.
120    pub fn dtype(&self) -> &DType {
121        self.footer.dtype()
122    }
123
124    /// Returns the file's statistics, if available.
125    ///
126    /// Statistics can be used for query optimization and data exploration.
127    pub fn file_stats(&self) -> Option<&FileStatistics> {
128        self.footer.statistics()
129    }
130
131    /// Returns the user-defined metadata segments loaded for this file.
132    ///
133    /// Metadata is only loaded when requested during open. Iteration order is unspecified.
134    pub fn metadata_segments(&self) -> impl Iterator<Item = (&str, &ByteBuffer)> {
135        self.metadata
136            .iter()
137            .map(|(key, metadata)| (key.as_str(), metadata))
138    }
139
140    /// Returns the loaded user-defined metadata segment for the given key.
141    ///
142    /// Returns `None` when the key is absent or metadata was not loaded.
143    pub fn metadata_segment(&self, key: &str) -> Option<&ByteBuffer> {
144        self.metadata.get(key)
145    }
146
147    /// Create a new segment source for reading from the file.
148    ///
149    /// This may spawn a background I/O driver that will exit when the returned segment source
150    /// is dropped.
151    pub fn segment_source(&self) -> Arc<dyn SegmentSource> {
152        Arc::clone(&self.segment_source)
153    }
154
155    /// Replace the segment source used by this file.
156    ///
157    /// Any cached layout reader is cleared so that subsequent scans construct readers over the
158    /// replacement source.
159    pub fn with_segment_source(mut self, segment_source: Arc<dyn SegmentSource>) -> Self {
160        self.segment_source = segment_source;
161        if self.layout_reader_cache.is_some() {
162            self.layout_reader_cache = Some(OnceLock::new());
163        }
164        self
165    }
166
167    /// Returns a reference to the Vortex session used to open this file.
168    pub fn session(&self) -> &VortexSession {
169        &self.session
170    }
171
172    /// Create a new layout reader for the file.
173    ///
174    /// Wraps the root layout in a [`FileStatsLayoutReader`] if file stats are available.
175    pub fn layout_reader(&self) -> VortexResult<Arc<dyn LayoutReader>> {
176        match &self.layout_reader_cache {
177            None => layout_reader(
178                Arc::clone(&self.segment_source),
179                &self.footer,
180                &self.session,
181            ),
182            Some(reader) => {
183                // get_or_try_init is unstable
184                if let Some(val) = reader.get() {
185                    Ok(Arc::clone(val))
186                } else {
187                    let inner = layout_reader(
188                        Arc::clone(&self.segment_source),
189                        &self.footer,
190                        &self.session,
191                    )?;
192                    Ok(if let Err(val) = reader.set(Arc::clone(&inner)) {
193                        val
194                    } else {
195                        inner
196                    })
197                }
198            }
199        }
200    }
201
202    /// Create a [`DataSource`](vortex_scan::DataSource) from this file for scanning.
203    ///
204    /// Wraps the file's layout reader with [`FileStatsLayoutReader`] (when file-level
205    /// statistics are available) and [`LayoutReaderDataSource`].
206    pub fn data_source(&self) -> VortexResult<DataSourceRef> {
207        let reader = self.layout_reader()?;
208
209        Ok(Arc::new(LayoutReaderDataSource::new(
210            reader,
211            self.session.clone(),
212        )))
213    }
214
215    /// Initiate a scan of the file, returning a builder for projection, filtering, selection, and
216    /// execution options.
217    pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
218        Ok(ScanBuilder::new(
219            self.session.clone(),
220            self.layout_reader()?,
221        ))
222    }
223
224    /// Returns `true` if file-level statistics prove the expression cannot
225    /// match any rows in this file.
226    ///
227    /// Row-count-aware pruning predicates are evaluated with the file's total
228    /// row count as their scope.
229    pub fn can_prune(&self, filter: &Expression) -> VortexResult<bool> {
230        let Some((stats, fields)) = self
231            .footer
232            .statistics()
233            .zip(self.footer.dtype().as_struct_fields_opt())
234        else {
235            return Ok(false);
236        };
237
238        can_prune_file_stats(
239            filter,
240            self.footer.dtype(),
241            self.footer.row_count(),
242            stats,
243            fields,
244            &self.session,
245        )
246    }
247
248    /// Return the file's natural row splits as root-coordinate ranges.
249    ///
250    /// These are the ranges that [`SplitBy::Layout`] would use for an all-fields scan.
251    pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
252        let reader = self.layout_reader()?;
253        Ok(SplitBy::Layout
254            .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
255            .into_iter()
256            .tuple_windows()
257            .map(|(start, end)| start..end)
258            .collect())
259    }
260}
vortex_file/file.rs

vortex_file/
file.rs
