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::Columnar;
16use vortex_array::IntoArray;
17use vortex_array::VortexSessionExecute;
18use vortex_array::arrays::ConstantArray;
19use vortex_array::dtype::DType;
20use vortex_array::dtype::Field;
21use vortex_array::dtype::FieldMask;
22use vortex_array::dtype::FieldPath;
23use vortex_array::dtype::FieldPathSet;
24use vortex_array::expr::Expression;
25use vortex_array::expr::pruning::checked_pruning_expr;
26use vortex_array::scalar_fn::internal::row_count::substitute_row_count;
27use vortex_error::VortexResult;
28use vortex_layout::LayoutReader;
29use vortex_layout::scan::layout::LayoutReaderDataSource;
30use vortex_layout::scan::scan_builder::ScanBuilder;
31use vortex_layout::scan::split_by::SplitBy;
32use vortex_layout::segments::SegmentSource;
33use vortex_scan::DataSourceRef;
34use vortex_session::VortexSession;
35use vortex_utils::aliases::hash_map::HashMap;
36
37use crate::FileStatistics;
38use crate::footer::Footer;
39use crate::pruning::extract_relevant_file_stats_as_struct_row;
40use crate::v2::FileStatsLayoutReader;
41
42/// Represents a Vortex file, providing access to its metadata and content.
43///
44/// A `VortexFile` is created by opening a Vortex file using [`VortexOpenOptions`](crate::VortexOpenOptions).
45/// It provides methods for accessing file metadata (such as row count, data type, and statistics)
46/// and for initiating scans to read the file's contents.
47#[derive(Clone)]
48pub struct VortexFile {
49    /// The footer of the Vortex file, containing metadata and layout information.
50    footer: Footer,
51    /// The segment source used to read segments from this file.
52    segment_source: Arc<dyn SegmentSource>,
53    /// The Vortex session used to open this file.
54    session: VortexSession,
55    /// None id LayoutReader caching is turned off
56    layout_reader_cache: Option<OnceLock<Arc<dyn LayoutReader>>>,
57}
58
59fn layout_reader(
60    segment_source: Arc<dyn SegmentSource>,
61    footer: &Footer,
62    session: &VortexSession,
63) -> VortexResult<Arc<dyn LayoutReader>> {
64    let root_reader = footer
65        .layout()
66        // TODO(ngates): we may want to allow the user pass in a name here?
67        .new_reader("".into(), segment_source, session, &Default::default())?;
68
69    Ok(if let Some(stats) = footer.statistics().cloned() {
70        Arc::new(FileStatsLayoutReader::new(
71            root_reader,
72            stats,
73            session.clone(),
74        ))
75    } else {
76        root_reader
77    })
78}
79
80impl VortexFile {
81    /// Creates a new `VortexFile` from the given footer, segment source, and session.
82    pub fn new(
83        footer: Footer,
84        segment_source: Arc<dyn SegmentSource>,
85        session: VortexSession,
86    ) -> Self {
87        Self {
88            footer,
89            segment_source,
90            session,
91            layout_reader_cache: None,
92        }
93    }
94
95    /// Enable layout reader caching
96    pub fn with_caching(self) -> Self {
97        Self {
98            footer: self.footer,
99            segment_source: self.segment_source,
100            session: self.session,
101            layout_reader_cache: Some(OnceLock::new()),
102        }
103    }
104
105    /// Returns a reference to the file's footer, which contains metadata and layout information.
106    pub fn footer(&self) -> &Footer {
107        &self.footer
108    }
109
110    /// Returns the number of rows in the file.
111    pub fn row_count(&self) -> u64 {
112        self.footer.row_count()
113    }
114
115    /// Returns the data type of the file's contents.
116    pub fn dtype(&self) -> &DType {
117        self.footer.dtype()
118    }
119
120    /// Returns the file's statistics, if available.
121    ///
122    /// Statistics can be used for query optimization and data exploration.
123    pub fn file_stats(&self) -> Option<&FileStatistics> {
124        self.footer.statistics()
125    }
126
127    /// Create a new segment source for reading from the file.
128    ///
129    /// This may spawn a background I/O driver that will exit when the returned segment source
130    /// is dropped.
131    pub fn segment_source(&self) -> Arc<dyn SegmentSource> {
132        Arc::clone(&self.segment_source)
133    }
134
135    /// Returns a reference to the Vortex session used to open this file.
136    pub fn session(&self) -> &VortexSession {
137        &self.session
138    }
139
140    /// Create a new layout reader for the file.
141    ///
142    /// Wraps the root layout in a [`FileStatsLayoutReader`] if file stats are available.
143    pub fn layout_reader(&self) -> VortexResult<Arc<dyn LayoutReader>> {
144        match &self.layout_reader_cache {
145            None => layout_reader(
146                Arc::clone(&self.segment_source),
147                &self.footer,
148                &self.session,
149            ),
150            Some(reader) => {
151                // get_or_try_init is unstable
152                if let Some(val) = reader.get() {
153                    Ok(Arc::clone(val))
154                } else {
155                    let inner = layout_reader(
156                        Arc::clone(&self.segment_source),
157                        &self.footer,
158                        &self.session,
159                    )?;
160                    Ok(if let Err(val) = reader.set(Arc::clone(&inner)) {
161                        val
162                    } else {
163                        inner
164                    })
165                }
166            }
167        }
168    }
169
170    /// Create a [`DataSource`](vortex_scan::DataSource) from this file for scanning.
171    ///
172    /// Wraps the file's layout reader with [`FileStatsLayoutReader`] (when file-level
173    /// statistics are available) and [`LayoutReaderDataSource`].
174    pub fn data_source(&self) -> VortexResult<DataSourceRef> {
175        let reader = self.layout_reader()?;
176
177        Ok(Arc::new(LayoutReaderDataSource::new(
178            reader,
179            self.session.clone(),
180        )))
181    }
182
183    /// Initiate a scan of the file, returning a builder for configuring the scan.
184    pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
185        Ok(ScanBuilder::new(
186            self.session.clone(),
187            self.layout_reader()?,
188        ))
189    }
190
191    /// Returns `true` if file-level statistics prove the expression cannot
192    /// match any rows in this file.
193    ///
194    /// Row-count-aware pruning predicates are evaluated with the file's total
195    /// row count as their scope.
196    pub fn can_prune(&self, filter: &Expression) -> VortexResult<bool> {
197        let Some((stats, fields)) = self
198            .footer
199            .statistics()
200            .zip(self.footer.dtype().as_struct_fields_opt())
201        else {
202            return Ok(false);
203        };
204
205        let set = FieldPathSet::from_iter(
206            fields
207                .names()
208                .iter()
209                .zip(stats.stats_sets().iter())
210                .flat_map(|(name, stats)| {
211                    stats.iter().map(|(stat, _)| {
212                        FieldPath::from_iter([
213                            Field::Name(name.clone()),
214                            Field::Name(stat.name().into()),
215                        ])
216                    })
217                }),
218        );
219
220        let Some((predicate, required_stats)) = checked_pruning_expr(filter, &set) else {
221            return Ok(false);
222        };
223
224        let required_file_stats = HashMap::from_iter(
225            required_stats
226                .map()
227                .iter()
228                .map(|(path, stats)| (path.clone(), stats.clone())),
229        );
230
231        let Some(file_stats) = extract_relevant_file_stats_as_struct_row(
232            &required_file_stats,
233            stats.stats_sets(),
234            fields,
235        )?
236        else {
237            return Ok(false);
238        };
239
240        // Apply the predicate, then substitute any row_count placeholders in the resulting array
241        // tree with a ConstantArray carrying the file-level row count.
242        let applied = file_stats.apply(&predicate)?;
243        let row_count_replacement =
244            ConstantArray::new(self.footer.row_count(), applied.len()).into_array();
245        let applied = substitute_row_count(applied, &row_count_replacement)?;
246
247        let mut ctx = self.session.create_execution_ctx();
248        Ok(match applied.execute::<Columnar>(&mut ctx)? {
249            Columnar::Constant(s) => s.scalar().as_bool().value() == Some(true),
250            Columnar::Canonical(c) => {
251                c.into_array()
252                    .execute_scalar(0, &mut ctx)?
253                    .as_bool()
254                    .value()
255                    == Some(true)
256            }
257        })
258    }
259
260    pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
261        let reader = self.layout_reader()?;
262        Ok(SplitBy::Layout
263            .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
264            .into_iter()
265            .tuple_windows()
266            .map(|(start, end)| start..end)
267            .collect())
268    }
269}
vortex_file/file.rs

vortex_file/
file.rs
