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