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

vortex_file/v2/
file_stats_reader.rs
