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

vortex_file/v2/
file_stats_reader.rs
