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::ops::Range;
11use std::sync::Arc;
12
13use vortex_array::MaskFuture;
14use vortex_array::dtype::DType;
15use vortex_array::dtype::FieldMask;
16use vortex_array::dtype::StructFields;
17use vortex_array::expr::Expression;
18use vortex_error::VortexResult;
19use vortex_layout::ArrayFuture;
20use vortex_layout::LayoutReader;
21use vortex_layout::LayoutReaderRef;
22use vortex_layout::RowSplits;
23use vortex_layout::SplitRange;
24use vortex_mask::Mask;
25use vortex_session::VortexSession;
26use vortex_utils::aliases::dash_map::DashMap;
27
28use crate::FileStatistics;
29use crate::pruning::can_prune_file_stats;
30
31/// A [`LayoutReader`] decorator that prunes entire files based on file-level statistics.
32///
33/// This reader wraps an inner `LayoutReader` and intercepts `pruning_evaluation` calls.
34/// When file-level stats prove that a filter expression is false for the entire file,
35/// it returns an all-false mask immediately — avoiding all downstream I/O.
36///
37/// Pruning results are cached per-expression since file-level stats are global
38/// (the result is the same regardless of which row range is requested).
39pub struct FileStatsLayoutReader {
40    child: LayoutReaderRef,
41    file_stats: FileStatistics,
42    struct_fields: StructFields,
43    session: VortexSession,
44    prune_cache: DashMap<Expression, bool>,
45}
46
47impl FileStatsLayoutReader {
48    /// Creates a new `FileStatsLayoutReader` wrapping the given child reader.
49    ///
50    /// The `struct_fields` are derived from the child reader's dtype. If the dtype is not a
51    /// struct, the available stats will be empty and no pruning will occur.
52    ///
53    /// Pre-computes the set of available stat field paths from the struct fields and file stats.
54    pub fn new(child: LayoutReaderRef, file_stats: FileStatistics, session: VortexSession) -> Self {
55        let struct_fields = child
56            .dtype()
57            .as_struct_fields_opt()
58            .cloned()
59            .unwrap_or_default();
60
61        Self {
62            child,
63            file_stats,
64            struct_fields,
65            session,
66            prune_cache: Default::default(),
67        }
68    }
69
70    /// Evaluates whether file-level statistics prove `expr` cannot match.
71    ///
72    /// Row-count placeholders are resolved against the full file row count,
73    /// independent of the requested row range.
74    fn evaluate_file_stats(&self, expr: &Expression) -> VortexResult<bool> {
75        can_prune_file_stats(
76            expr,
77            self.child.dtype(),
78            self.child.row_count(),
79            &self.file_stats,
80            &self.struct_fields,
81            &self.session,
82        )
83    }
84
85    /// Returns the file-level statistics used by this reader.
86    pub fn file_stats(&self) -> &FileStatistics {
87        &self.file_stats
88    }
89}
90
91impl LayoutReader for FileStatsLayoutReader {
92    fn name(&self) -> &Arc<str> {
93        self.child.name()
94    }
95
96    fn dtype(&self) -> &DType {
97        self.child.dtype()
98    }
99
100    fn row_count(&self) -> u64 {
101        self.child.row_count()
102    }
103
104    fn register_splits(
105        &self,
106        field_mask: &[FieldMask],
107        split_range: &SplitRange,
108        splits: &mut RowSplits,
109    ) -> VortexResult<()> {
110        self.child.register_splits(field_mask, split_range, splits)
111    }
112
113    fn pruning_evaluation(
114        &self,
115        row_range: &Range<u64>,
116        expr: &Expression,
117        mask: Mask,
118    ) -> VortexResult<MaskFuture> {
119        // Check cache first with read-only lock.
120        if let Some(pruned) = self.prune_cache.get(expr) {
121            if *pruned {
122                return Ok(MaskFuture::ready(Mask::new_false(mask.len())));
123            }
124            return self.child.pruning_evaluation(row_range, expr, mask);
125        }
126
127        // Evaluate and cache.
128        let pruned = self.evaluate_file_stats(expr)?;
129        self.prune_cache.insert(expr.clone(), pruned);
130
131        if pruned {
132            Ok(MaskFuture::ready(Mask::new_false(mask.len())))
133        } else {
134            self.child.pruning_evaluation(row_range, expr, mask)
135        }
136    }
137
138    fn filter_evaluation(
139        &self,
140        row_range: &Range<u64>,
141        expr: &Expression,
142        mask: MaskFuture,
143    ) -> VortexResult<MaskFuture> {
144        self.child.filter_evaluation(row_range, expr, mask)
145    }
146
147    fn projection_evaluation(
148        &self,
149        row_range: &Range<u64>,
150        expr: &Expression,
151        mask: MaskFuture,
152    ) -> VortexResult<ArrayFuture> {
153        self.child.projection_evaluation(row_range, expr, mask)
154    }
155
156    fn as_any(&self) -> &dyn std::any::Any {
157        self
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use std::sync::Arc;
164    use std::sync::LazyLock;
165
166    use vortex_array::ArrayContext;
167    use vortex_array::IntoArray as _;
168    use vortex_array::arrays::PrimitiveArray;
169    use vortex_array::arrays::StructArray;
170    use vortex_array::arrays::datetime::TemporalData;
171    use vortex_array::dtype::DType;
172    use vortex_array::dtype::Nullability;
173    use vortex_array::dtype::PType;
174    use vortex_array::expr::checked_add;
175    use vortex_array::expr::get_item;
176    use vortex_array::expr::gt;
177    use vortex_array::expr::is_not_null;
178    use vortex_array::expr::is_null;
179    use vortex_array::expr::lit;
180    use vortex_array::expr::root;
181    use vortex_array::expr::stats::Precision;
182    use vortex_array::expr::stats::Stat;
183    use vortex_array::extension::datetime::TimeUnit;
184    use vortex_array::scalar::ScalarValue;
185    use vortex_array::stats::StatsSet;
186    use vortex_buffer::buffer;
187    use vortex_error::VortexResult;
188    use vortex_io::runtime::single::block_on;
189    use vortex_io::session::RuntimeSession;
190    use vortex_io::session::RuntimeSessionExt;
191    use vortex_layout::LayoutReader;
192    use vortex_layout::LayoutStrategy;
193    use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
194    use vortex_layout::layouts::table::TableStrategy;
195    use vortex_layout::segments::SegmentSink;
196    use vortex_layout::segments::TestSegments;
197    use vortex_layout::sequence::SequenceId;
198    use vortex_layout::sequence::SequentialArrayStreamExt;
199    use vortex_layout::session::LayoutSession;
200    use vortex_mask::Mask;
201    use vortex_session::VortexSession;
202
203    use super::*;
204
205    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
206        vortex_array::array_session()
207            .with::<LayoutSession>()
208            .with::<RuntimeSession>()
209    });
210
211    fn test_file_stats(min: i32, max: i32) -> FileStatistics {
212        let mut stats = StatsSet::default();
213        stats.set(Stat::Min, Precision::exact(ScalarValue::from(min)));
214        stats.set(Stat::Max, Precision::exact(ScalarValue::from(max)));
215        FileStatistics::new(
216            Arc::from([stats]),
217            Arc::from([DType::Primitive(PType::I32, Nullability::NonNullable)]),
218        )
219    }
220
221    fn test_file_null_count_stats(null_count: u64) -> FileStatistics {
222        let mut stats = StatsSet::default();
223        stats.set(
224            Stat::NullCount,
225            Precision::exact(ScalarValue::from(null_count)),
226        );
227        FileStatistics::new(
228            Arc::from([stats]),
229            Arc::from([DType::Primitive(PType::I32, Nullability::Nullable)]),
230        )
231    }
232
233    #[test]
234    fn pruning_when_filter_out_of_range() -> VortexResult<()> {
235        block_on(|handle| async {
236            let session = SESSION.clone().with_handle(handle);
237            let ctx = ArrayContext::empty();
238            let segments = Arc::new(TestSegments::default());
239            let (ptr, eof) = SequenceId::root().split();
240            let struct_array = StructArray::from_fields(
241                [("col", buffer![1i32, 2, 3, 4, 5].into_array())].as_slice(),
242            )?;
243            let strategy = TableStrategy::new(
244                Arc::new(FlatLayoutStrategy::default()),
245                Arc::new(FlatLayoutStrategy::default()),
246            );
247            let layout = strategy
248                .write_stream(
249                    ctx,
250                    Arc::<TestSegments>::clone(&segments),
251                    struct_array.into_array().to_array_stream().sequenced(ptr),
252                    eof,
253                    &session,
254                )
255                .await?;
256
257            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
258
259            let reader =
260                FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone());
261
262            // col > 200 should be prunable since max is 100.
263            let expr = gt(get_item("col", root()), lit(200i32));
264            let mask = Mask::new_true(5);
265            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
266            assert_eq!(result, Mask::new_false(5));
267
268            Ok(())
269        })
270    }
271
272    #[test]
273    fn no_pruning_when_filter_in_range() -> VortexResult<()> {
274        block_on(|handle| async {
275            let session = SESSION.clone().with_handle(handle);
276            let ctx = ArrayContext::empty();
277            let segments = Arc::new(TestSegments::default());
278            let (ptr, eof) = SequenceId::root().split();
279            let struct_array = StructArray::from_fields(
280                [("col", buffer![1i32, 2, 3, 4, 5].into_array())].as_slice(),
281            )?;
282            let strategy = TableStrategy::new(
283                Arc::new(FlatLayoutStrategy::default()),
284                Arc::new(FlatLayoutStrategy::default()),
285            );
286            let layout = strategy
287                .write_stream(
288                    ctx,
289                    Arc::<TestSegments>::clone(&segments),
290                    struct_array.into_array().to_array_stream().sequenced(ptr),
291                    eof,
292                    &session,
293                )
294                .await?;
295
296            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
297
298            let reader =
299                FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone());
300
301            // col > 50 should NOT be prunable since max is 100 (some rows could match).
302            let expr = gt(get_item("col", root()), lit(50i32));
303            let mask = Mask::new_true(5);
304            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
305            // Should delegate to child, which returns the mask unchanged (struct reader doesn't prune).
306            assert_eq!(result, Mask::new_true(5));
307
308            Ok(())
309        })
310    }
311
312    #[test]
313    fn no_pruning_for_computed_expression_stats() -> VortexResult<()> {
314        block_on(|handle| async {
315            let session = SESSION.clone().with_handle(handle);
316            let ctx = ArrayContext::empty();
317            let segments = Arc::new(TestSegments::default());
318            let (ptr, eof) = SequenceId::root().split();
319            let struct_array =
320                StructArray::from_fields([("col", buffer![0i32, 100].into_array())].as_slice())?;
321            let strategy = TableStrategy::new(
322                Arc::new(FlatLayoutStrategy::default()),
323                Arc::new(FlatLayoutStrategy::default()),
324            );
325            let layout = strategy
326                .write_stream(
327                    ctx,
328                    Arc::<TestSegments>::clone(&segments),
329                    struct_array.into_array().to_array_stream().sequenced(ptr),
330                    eof,
331                    &session,
332                )
333                .await?;
334
335            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
336            let reader =
337                FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone());
338
339            let expr = gt(checked_add(get_item("col", root()), lit(5i32)), lit(102i32));
340            let mask = Mask::new_true(2);
341            let result = reader.pruning_evaluation(&(0..2), &expr, mask)?.await?;
342
343            assert_eq!(result, Mask::new_true(2));
344
345            Ok(())
346        })
347    }
348
349    /// Regression test: `IS NULL` on a nullable timestamp column must not fail with a
350    /// dtype mismatch. The bug was that `stats_ref` used the *field* dtype (timestamp)
351    /// for the `NullCount` stat scalar instead of the stat's own dtype (u64).
352    #[test]
353    fn is_null_pruning_on_nullable_timestamp_column() -> VortexResult<()> {
354        block_on(|handle| async {
355            let session = SESSION.clone().with_handle(handle);
356            let ctx = ArrayContext::empty();
357            let segments = Arc::new(TestSegments::default());
358            let (ptr, eof) = SequenceId::root().split();
359
360            // Build a struct with a nullable timestamp column containing some nulls.
361            let prim_array =
362                PrimitiveArray::from_option_iter([Some(1_000_000i64), None, Some(3_000_000)])
363                    .into_array();
364            let ts_data = TemporalData::new_timestamp(prim_array, TimeUnit::Microseconds, None);
365            let ts_dtype = ts_data.dtype().clone();
366            let ts_array = ts_data.into_array();
367
368            let struct_array = StructArray::from_fields([("deleted_at", ts_array)].as_slice())?;
369
370            let strategy = TableStrategy::new(
371                Arc::new(FlatLayoutStrategy::default()),
372                Arc::new(FlatLayoutStrategy::default()),
373            );
374            let layout = strategy
375                .write_stream(
376                    ctx,
377                    Arc::clone(&segments) as Arc<dyn SegmentSink>,
378                    struct_array.into_array().to_array_stream().sequenced(ptr),
379                    eof,
380                    &session,
381                )
382                .await?;
383
384            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
385
386            // File-level stats: 1 null in deleted_at.
387            let mut stats = StatsSet::default();
388            stats.set(Stat::NullCount, Precision::exact(ScalarValue::from(1u64)));
389            let file_stats = FileStatistics::new(Arc::from([stats]), Arc::from([ts_dtype]));
390
391            let reader = FileStatsLayoutReader::new(child, file_stats, SESSION.clone());
392
393            // `is_null(deleted_at)` — should NOT panic or error due to dtype mismatch.
394            let expr = is_null(get_item("deleted_at", root()));
395            let mask = Mask::new_true(3);
396            let result = reader.pruning_evaluation(&(0..3), &expr, mask)?.await?;
397            // null_count is 1 (non-zero), so is_null is not falsified => not pruned.
398            assert_eq!(result, Mask::new_true(3));
399
400            Ok(())
401        })
402    }
403
404    #[test]
405    fn pruning_is_not_null_when_file_is_all_null() -> VortexResult<()> {
406        block_on(|handle| async {
407            let session = SESSION.clone().with_handle(handle);
408            let ctx = ArrayContext::empty();
409            let segments = Arc::new(TestSegments::default());
410            let (ptr, eof) = SequenceId::root().split();
411            let struct_array = StructArray::from_fields(
412                [(
413                    "col",
414                    PrimitiveArray::from_option_iter([None::<i32>, None, None, None, None])
415                        .into_array(),
416                )]
417                .as_slice(),
418            )?;
419            let strategy = TableStrategy::new(
420                Arc::new(FlatLayoutStrategy::default()),
421                Arc::new(FlatLayoutStrategy::default()),
422            );
423            let layout = strategy
424                .write_stream(
425                    ctx,
426                    Arc::clone(&segments) as Arc<dyn SegmentSink>,
427                    struct_array.into_array().to_array_stream().sequenced(ptr),
428                    eof,
429                    &session,
430                )
431                .await?;
432
433            let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?;
434
435            let reader =
436                FileStatsLayoutReader::new(child, test_file_null_count_stats(5), SESSION.clone());
437
438            let expr = is_not_null(get_item("col", root()));
439            let mask = Mask::new_true(5);
440            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
441            assert_eq!(result, Mask::new_false(5));
442
443            Ok(())
444        })
445    }
446}
vortex_file/v2/file_stats_reader.rs

vortex_file/v2/
file_stats_reader.rs
