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

vortex_file/v2/
file_stats_reader.rs
