vortex_file/v2/
file_stats_reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A [`LayoutReader`] decorator that performs file-level stats pruning.
5//!
6//! If file-level statistics prove that a filter expression cannot match any rows in the file,
7//! [`FileStatsLayoutReader`] short-circuits [`pruning_evaluation`](LayoutReader::pruning_evaluation)
8//! by returning an all-false mask — avoiding all downstream I/O.
9
10use std::collections::BTreeSet;
11use std::ops::Range;
12use std::sync::Arc;
13
14use vortex_array::Canonical;
15use vortex_array::IntoArray;
16use vortex_array::MaskFuture;
17use vortex_array::VortexSessionExecute;
18use vortex_array::arrays::NullArray;
19use vortex_array::dtype::DType;
20use vortex_array::dtype::FieldMask;
21use vortex_array::dtype::FieldPath;
22use vortex_array::dtype::StructFields;
23use vortex_array::expr::Expression;
24use vortex_array::expr::StatsCatalog;
25use vortex_array::expr::lit;
26use vortex_array::expr::stats::Stat;
27use vortex_array::scalar::Scalar;
28use vortex_array::scalar_fn::fns::literal::Literal;
29use vortex_error::VortexResult;
30use vortex_layout::ArrayFuture;
31use vortex_layout::LayoutReader;
32use vortex_layout::LayoutReaderRef;
33use vortex_mask::Mask;
34use vortex_session::VortexSession;
35use vortex_utils::aliases::dash_map::DashMap;
36
37use crate::FileStatistics;
38
39/// A [`LayoutReader`] decorator that prunes entire files based on file-level statistics.
40///
41/// This reader wraps an inner `LayoutReader` and intercepts `pruning_evaluation` calls.
42/// When file-level stats prove that a filter expression is false for the entire file,
43/// it returns an all-false mask immediately — avoiding all downstream I/O.
44///
45/// Pruning results are cached per-expression since file-level stats are global
46/// (the result is the same regardless of which row range is requested).
47pub struct FileStatsLayoutReader {
48    child: LayoutReaderRef,
49    file_stats: FileStatistics,
50    struct_fields: StructFields,
51    session: VortexSession,
52    prune_cache: DashMap<Expression, bool>,
53}
54
55impl FileStatsLayoutReader {
56    /// Creates a new `FileStatsLayoutReader` wrapping the given child reader.
57    ///
58    /// The `struct_fields` are derived from the child reader's dtype. If the dtype is not a
59    /// struct, the available stats will be empty and no pruning will occur.
60    ///
61    /// Pre-computes the set of available stat field paths from the struct fields and file stats.
62    pub fn new(child: LayoutReaderRef, file_stats: FileStatistics, session: VortexSession) -> Self {
63        let struct_fields = child
64            .dtype()
65            .as_struct_fields_opt()
66            .cloned()
67            .unwrap_or_default();
68
69        Self {
70            child,
71            file_stats,
72            struct_fields,
73            session,
74            prune_cache: Default::default(),
75        }
76    }
77
78    /// Evaluates whether the file can be fully pruned for the given expression.
79    ///
80    /// Returns `true` if file-level stats prove no rows can match, `false` otherwise.
81    fn evaluate_file_stats(&self, expr: &Expression) -> VortexResult<bool> {
82        let Some(pruning_expr) = expr.stat_falsification(self) else {
83            // If there is no pruning expression, we can't prune.
84            return Ok(false);
85        };
86
87        // Given how we implemented the StatsCatalog, we know the expression must be all literals.
88        // We can therefore optimize with a null scope since there are no field references that
89        // need to be resolved.
90        let simplified = pruning_expr.optimize_recursive(&DType::Null)?;
91        if let Some(result) = simplified.as_opt::<Literal>() {
92            // Can prune if the result is non-nullable and true
93            return Ok(result.as_bool().value() == Some(true));
94        }
95
96        // Sometimes expressions don't implement constant folding to literals... In this case,
97        // we just execute the expression over a null array.
98        let pruning = NullArray::new(1).into_array().apply(&pruning_expr)?;
99
100        let mut ctx = self.session.create_execution_ctx();
101        let result = pruning
102            .execute::<Canonical>(&mut ctx)?
103            .into_bool()
104            .into_array()
105            .scalar_at(0)?;
106
107        Ok(result.as_bool().value() == Some(true))
108    }
109}
110
111/// Implements [`StatsCatalog`] to provide file-level stats to expressions during pruning evaluation.
112impl StatsCatalog for FileStatsLayoutReader {
113    fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<Expression> {
114        // FileStats currently only holds top-level field statistics.
115        if field_path.parts().len() != 1 {
116            return None;
117        }
118
119        let field_name = field_path.parts()[0].as_name()?;
120        let field_idx = self.struct_fields.find(field_name)?;
121        let field_stats = self.file_stats.stats_sets().get(field_idx)?;
122
123        let stat_value = field_stats.get(stat)?.as_exact()?;
124        let field_dtype = self.struct_fields.field_by_index(field_idx)?;
125        let stat_scalar = Scalar::try_new(field_dtype, Some(stat_value)).ok()?;
126
127        Some(lit(stat_scalar))
128    }
129}
130
131impl LayoutReader for FileStatsLayoutReader {
132    fn name(&self) -> &Arc<str> {
133        self.child.name()
134    }
135
136    fn dtype(&self) -> &DType {
137        self.child.dtype()
138    }
139
140    fn row_count(&self) -> u64 {
141        self.child.row_count()
142    }
143
144    fn register_splits(
145        &self,
146        field_mask: &[FieldMask],
147        row_range: &Range<u64>,
148        splits: &mut BTreeSet<u64>,
149    ) -> VortexResult<()> {
150        self.child.register_splits(field_mask, row_range, splits)
151    }
152
153    fn pruning_evaluation(
154        &self,
155        row_range: &Range<u64>,
156        expr: &Expression,
157        mask: Mask,
158    ) -> VortexResult<MaskFuture> {
159        // Check cache first with read-only lock.
160        if let Some(pruned) = self.prune_cache.get(expr) {
161            if *pruned {
162                return Ok(MaskFuture::ready(Mask::new_false(mask.len())));
163            }
164            return self.child.pruning_evaluation(row_range, expr, mask);
165        }
166
167        // Evaluate and cache.
168        let pruned = self.evaluate_file_stats(expr)?;
169        self.prune_cache.insert(expr.clone(), pruned);
170
171        if pruned {
172            Ok(MaskFuture::ready(Mask::new_false(mask.len())))
173        } else {
174            self.child.pruning_evaluation(row_range, expr, mask)
175        }
176    }
177
178    fn filter_evaluation(
179        &self,
180        row_range: &Range<u64>,
181        expr: &Expression,
182        mask: MaskFuture,
183    ) -> VortexResult<MaskFuture> {
184        self.child.filter_evaluation(row_range, expr, mask)
185    }
186
187    fn projection_evaluation(
188        &self,
189        row_range: &Range<u64>,
190        expr: &Expression,
191        mask: MaskFuture,
192    ) -> VortexResult<ArrayFuture> {
193        self.child.projection_evaluation(row_range, expr, mask)
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use std::sync::Arc;
200    use std::sync::LazyLock;
201
202    use vortex_array::ArrayContext;
203    use vortex_array::IntoArray as _;
204    use vortex_array::arrays::StructArray;
205    use vortex_array::dtype::DType;
206    use vortex_array::dtype::Nullability;
207    use vortex_array::dtype::PType;
208    use vortex_array::expr::get_item;
209    use vortex_array::expr::gt;
210    use vortex_array::expr::lit;
211    use vortex_array::expr::root;
212    use vortex_array::expr::stats::Precision;
213    use vortex_array::expr::stats::Stat;
214    use vortex_array::scalar::ScalarValue;
215    use vortex_array::scalar_fn::session::ScalarFnSession;
216    use vortex_array::session::ArraySession;
217    use vortex_array::stats::StatsSet;
218    use vortex_buffer::buffer;
219    use vortex_error::VortexResult;
220    use vortex_io::runtime::single::block_on;
221    use vortex_io::session::RuntimeSession;
222    use vortex_io::session::RuntimeSessionExt;
223    use vortex_layout::LayoutReader;
224    use vortex_layout::LayoutStrategy;
225    use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
226    use vortex_layout::layouts::table::TableStrategy;
227    use vortex_layout::segments::TestSegments;
228    use vortex_layout::sequence::SequenceId;
229    use vortex_layout::sequence::SequentialArrayStreamExt;
230    use vortex_layout::session::LayoutSession;
231    use vortex_mask::Mask;
232    use vortex_session::VortexSession;
233
234    use super::*;
235
236    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
237        VortexSession::empty()
238            .with::<ArraySession>()
239            .with::<LayoutSession>()
240            .with::<ScalarFnSession>()
241            .with::<RuntimeSession>()
242    });
243
244    fn test_file_stats(min: i32, max: i32) -> FileStatistics {
245        let mut stats = StatsSet::default();
246        stats.set(Stat::Min, Precision::exact(ScalarValue::from(min)));
247        stats.set(Stat::Max, Precision::exact(ScalarValue::from(max)));
248        FileStatistics::new(
249            Arc::from([stats]),
250            Arc::from([DType::Primitive(PType::I32, Nullability::NonNullable)]),
251        )
252    }
253
254    #[test]
255    fn pruning_when_filter_out_of_range() -> VortexResult<()> {
256        block_on(|handle| async {
257            let session = SESSION.clone().with_handle(handle);
258            let ctx = ArrayContext::empty();
259            let segments = Arc::new(TestSegments::default());
260            let (ptr, eof) = SequenceId::root().split();
261            let struct_array = StructArray::from_fields(
262                [("col", buffer![1i32, 2, 3, 4, 5].into_array())].as_slice(),
263            )?;
264            let strategy = TableStrategy::new(
265                Arc::new(FlatLayoutStrategy::default()),
266                Arc::new(FlatLayoutStrategy::default()),
267            );
268            let layout = strategy
269                .write_stream(
270                    ctx,
271                    Arc::<TestSegments>::clone(&segments),
272                    struct_array.into_array().to_array_stream().sequenced(ptr),
273                    eof,
274                    &session,
275                )
276                .await?;
277
278            let child = layout.new_reader("".into(), segments, &SESSION)?;
279
280            let reader =
281                FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone());
282
283            // col > 200 should be prunable since max is 100.
284            let expr = gt(get_item("col", root()), lit(200i32));
285            let mask = Mask::new_true(5);
286            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
287            assert_eq!(result, Mask::new_false(5));
288
289            Ok(())
290        })
291    }
292
293    #[test]
294    fn no_pruning_when_filter_in_range() -> VortexResult<()> {
295        block_on(|handle| async {
296            let session = SESSION.clone().with_handle(handle);
297            let ctx = ArrayContext::empty();
298            let segments = Arc::new(TestSegments::default());
299            let (ptr, eof) = SequenceId::root().split();
300            let struct_array = StructArray::from_fields(
301                [("col", buffer![1i32, 2, 3, 4, 5].into_array())].as_slice(),
302            )?;
303            let strategy = TableStrategy::new(
304                Arc::new(FlatLayoutStrategy::default()),
305                Arc::new(FlatLayoutStrategy::default()),
306            );
307            let layout = strategy
308                .write_stream(
309                    ctx,
310                    Arc::<TestSegments>::clone(&segments),
311                    struct_array.into_array().to_array_stream().sequenced(ptr),
312                    eof,
313                    &session,
314                )
315                .await?;
316
317            let child = layout.new_reader("".into(), segments, &SESSION)?;
318
319            let reader =
320                FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone());
321
322            // col > 50 should NOT be prunable since max is 100 (some rows could match).
323            let expr = gt(get_item("col", root()), lit(50i32));
324            let mask = Mask::new_true(5);
325            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
326            // Should delegate to child, which returns the mask unchanged (struct reader doesn't prune).
327            assert_eq!(result, Mask::new_true(5));
328
329            Ok(())
330        })
331    }
332}
vortex_file/v2/file_stats_reader.rs

vortex_file/v2/
file_stats_reader.rs
