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

vortex_file/
file.rs
