1use 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#[derive(Clone)]
47pub struct VortexFile {
48 pub(crate) footer: Footer,
50 pub(crate) segment_source: Arc<dyn SegmentSource>,
52 pub(crate) session: VortexSession,
54}
55
56impl VortexFile {
57 pub fn footer(&self) -> &Footer {
59 &self.footer
60 }
61
62 pub fn row_count(&self) -> u64 {
64 self.footer.row_count()
65 }
66
67 pub fn dtype(&self) -> &DType {
69 self.footer.dtype()
70 }
71
72 pub fn file_stats(&self) -> Option<&FileStatistics> {
76 self.footer.statistics()
77 }
78
79 pub fn segment_source(&self) -> Arc<dyn SegmentSource> {
84 Arc::clone(&self.segment_source)
85 }
86
87 pub fn layout_reader(&self) -> VortexResult<Arc<dyn LayoutReader>> {
89 let segment_source = self.segment_source();
90 self.footer
91 .layout()
92 .new_reader("".into(), segment_source, &self.session)
94 }
95
96 pub fn data_source(&self) -> VortexResult<DataSourceRef> {
101 let mut reader = self.layout_reader()?;
102 if let Some(stats) = self.file_stats().cloned() {
103 reader = Arc::new(FileStatsLayoutReader::new(
104 reader,
105 stats,
106 self.session.clone(),
107 ));
108 }
109 Ok(Arc::new(LayoutReaderDataSource::new(
110 reader,
111 self.session.clone(),
112 )))
113 }
114
115 pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
117 Ok(ScanBuilder::new(
118 self.session.clone(),
119 self.layout_reader()?,
120 ))
121 }
122
123 pub fn can_prune(&self, filter: &Expression) -> VortexResult<bool> {
129 let Some((stats, fields)) = self
130 .footer
131 .statistics()
132 .zip(self.footer.dtype().as_struct_fields_opt())
133 else {
134 return Ok(false);
135 };
136
137 let set = FieldPathSet::from_iter(
138 fields
139 .names()
140 .iter()
141 .zip(stats.stats_sets().iter())
142 .flat_map(|(name, stats)| {
143 stats.iter().map(|(stat, _)| {
144 FieldPath::from_iter([
145 Field::Name(name.clone()),
146 Field::Name(stat.name().into()),
147 ])
148 })
149 }),
150 );
151
152 let Some((predicate, required_stats)) = checked_pruning_expr(filter, &set) else {
153 return Ok(false);
154 };
155
156 let required_file_stats = HashMap::from_iter(
157 required_stats
158 .map()
159 .iter()
160 .map(|(path, stats)| (path.clone(), stats.clone())),
161 );
162
163 let Some(file_stats) = extract_relevant_file_stats_as_struct_row(
164 &required_file_stats,
165 stats.stats_sets(),
166 fields,
167 )?
168 else {
169 return Ok(false);
170 };
171
172 let applied = file_stats.apply(&predicate)?;
175 let row_count_replacement =
176 ConstantArray::new(self.footer.row_count(), applied.len()).into_array();
177 let applied = substitute_row_count(applied, &row_count_replacement)?;
178
179 let mut ctx = self.session.create_execution_ctx();
180 Ok(match applied.execute::<Columnar>(&mut ctx)? {
181 Columnar::Constant(s) => s.scalar().as_bool().value() == Some(true),
182 Columnar::Canonical(_) => false,
183 })
184 }
185
186 pub fn splits(&self) -> VortexResult<Vec<Range<u64>>> {
187 let reader = self.layout_reader()?;
188 Ok(SplitBy::Layout
189 .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])?
190 .into_iter()
191 .tuple_windows()
192 .map(|(start, end)| start..end)
193 .collect())
194 }
195}