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

vortex_file/v2/
file_stats_reader.rs
