vortex_layout/layouts/
file_stats.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::future;
5use std::marker::PhantomData;
6use std::sync::Arc;
7
8use futures::StreamExt;
9use itertools::Itertools;
10use parking_lot::Mutex;
11use vortex_array::ArrayRef;
12use vortex_array::ExecutionCtx;
13use vortex_array::VortexSessionExecute;
14use vortex_array::aggregate_fn::fns::sum::sum;
15use vortex_array::arrays::StructArray;
16use vortex_array::arrays::struct_::StructArrayExt;
17use vortex_array::builders::ArrayBuilder;
18use vortex_array::builders::BoolBuilder;
19use vortex_array::builders::builder_with_capacity_in;
20use vortex_array::dtype::DType;
21use vortex_array::dtype::FieldName;
22use vortex_array::dtype::Nullability;
23use vortex_array::dtype::PType;
24use vortex_array::expr::stats::Precision;
25use vortex_array::expr::stats::Stat;
26use vortex_array::scalar::Scalar;
27use vortex_array::scalar::ScalarTruncation;
28use vortex_array::scalar::lower_bound;
29use vortex_array::scalar::upper_bound;
30use vortex_array::stats::StatsSet;
31use vortex_array::validity::Validity;
32use vortex_buffer::BufferString;
33use vortex_buffer::ByteBuffer;
34use vortex_error::VortexExpect;
35use vortex_error::VortexResult;
36use vortex_error::vortex_panic;
37use vortex_session::VortexSession;
38
39use crate::layouts::zoned::MAX_IS_TRUNCATED;
40use crate::layouts::zoned::MIN_IS_TRUNCATED;
41use crate::sequence::SendableSequentialStream;
42use crate::sequence::SequenceId;
43use crate::sequence::SequentialStreamAdapter;
44use crate::sequence::SequentialStreamExt;
45
46pub fn accumulate_stats(
47    stream: SendableSequentialStream,
48    stats: Arc<[Stat]>,
49    max_variable_length_statistics_size: usize,
50    session: &VortexSession,
51) -> (FileStatsAccumulator, SendableSequentialStream) {
52    let accumulator = FileStatsAccumulator::new(
53        stream.dtype(),
54        stats,
55        max_variable_length_statistics_size,
56        session,
57    );
58    let stream = SequentialStreamAdapter::new(
59        stream.dtype().clone(),
60        stream.scan(accumulator.clone(), |acc, item| {
61            future::ready(Some(acc.process(item)))
62        }),
63    )
64    .sendable();
65    (accumulator, stream)
66}
67
68/// Accumulates write-time statistics for a single file column.
69struct StatsAccumulator {
70    builders: Vec<Box<dyn StatsArrayBuilder>>,
71    length: usize,
72}
73
74impl StatsAccumulator {
75    fn new(dtype: &DType, stats: &[Stat], max_variable_length_statistics_size: usize) -> Self {
76        if !supports_file_stats(dtype) {
77            return Self {
78                builders: Vec::new(),
79                length: 0,
80            };
81        }
82
83        let builders = stats
84            .iter()
85            .filter_map(|&stat| {
86                stat.dtype(dtype).map(|stat_dtype| {
87                    stats_builder_with_capacity(
88                        stat,
89                        &stat_dtype.as_nullable(),
90                        1024,
91                        max_variable_length_statistics_size,
92                    )
93                })
94            })
95            .collect::<Vec<_>>();
96
97        Self {
98            builders,
99            length: 0,
100        }
101    }
102
103    fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
104        for builder in &mut self.builders {
105            if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? {
106                builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?;
107            } else {
108                builder.append_null();
109            }
110        }
111        self.length += 1;
112        Ok(())
113    }
114
115    fn as_array(&mut self, ctx: &mut ExecutionCtx) -> VortexResult<Option<StructArray>> {
116        let mut names = Vec::new();
117        let mut fields = Vec::new();
118
119        for builder in self
120            .builders
121            .iter_mut()
122            // We sort the stats so the DType is deterministic based on which stats are present.
123            .sorted_unstable_by_key(|builder| builder.stat())
124        {
125            let values = builder.finish();
126
127            // We drop any all-null stats columns.
128            if values.all_invalid(ctx)? {
129                continue;
130            }
131
132            names.extend(values.names);
133            fields.extend(values.arrays);
134        }
135
136        if names.is_empty() {
137            return Ok(None);
138        }
139
140        StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable).map(Some)
141    }
142
143    /// Returns an aggregated stats set for the table.
144    fn as_stats_set(&mut self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult<StatsSet> {
145        let mut stats_set = StatsSet::default();
146        let Some(stats_table) = self.as_array(ctx)? else {
147            return Ok(stats_set);
148        };
149
150        for &stat in stats {
151            let Some(values) = stats_table.unmasked_field_by_name_opt(stat.name()) else {
152                continue;
153            };
154
155            match stat {
156                Stat::Max if is_varlen_dtype(values.dtype()) && !values.all_valid(ctx)? => {
157                    // A null truncated varlen max can mean either an empty chunk or no finite
158                    // upper bound, so aggregating by skipping nulls would be unsound.
159                    continue;
160                }
161                Stat::Min | Stat::Max | Stat::Sum => {
162                    if let Some(s) = values.statistics().compute_stat(stat, ctx)?
163                        && let Some(v) = s.into_value()
164                    {
165                        let precision = if stat_was_truncated(&stats_table, stat, ctx)? {
166                            Precision::inexact(v)
167                        } else {
168                            Precision::exact(v)
169                        };
170                        stats_set.set(stat, precision)
171                    }
172                }
173                Stat::NullCount | Stat::NaNCount | Stat::UncompressedSizeInBytes => {
174                    if let Some(sum_value) = sum(values, ctx)?
175                        .cast(&DType::Primitive(PType::U64, Nullability::Nullable))?
176                        .into_value()
177                    {
178                        stats_set.set(stat, Precision::exact(sum_value));
179                    }
180                }
181                Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted => {}
182            }
183        }
184        Ok(stats_set)
185    }
186}
187
188fn stat_was_truncated(
189    stats_table: &StructArray,
190    stat: Stat,
191    ctx: &mut ExecutionCtx,
192) -> VortexResult<bool> {
193    let field_name = match stat {
194        Stat::Min => MIN_IS_TRUNCATED,
195        Stat::Max => MAX_IS_TRUNCATED,
196        _ => return Ok(false),
197    };
198    let Some(is_truncated) = stats_table.unmasked_field_by_name_opt(field_name) else {
199        return Ok(false);
200    };
201
202    Ok(is_truncated
203        .statistics()
204        .compute_stat(Stat::Max, ctx)?
205        .is_some_and(|max| max.as_bool().value() == Some(true)))
206}
207
208fn supports_file_stats(dtype: &DType) -> bool {
209    !matches!(dtype, DType::Variant(_))
210}
211
212fn is_varlen_dtype(dtype: &DType) -> bool {
213    matches!(dtype, DType::Utf8(_) | DType::Binary(_))
214}
215
216fn stats_builder_with_capacity(
217    stat: Stat,
218    dtype: &DType,
219    capacity: usize,
220    max_length: usize,
221) -> Box<dyn StatsArrayBuilder> {
222    let values_builder = builder_with_capacity_in(
223        dtype,
224        capacity,
225        vortex_buffer::BufferAllocatorRef::static_ref(),
226    );
227    match stat {
228        Stat::Max => match dtype {
229            DType::Utf8(_) => Box::new(TruncatedMaxBinaryStatsBuilder::<BufferString>::new(
230                values_builder,
231                BoolBuilder::with_capacity_in(
232                    Nullability::NonNullable,
233                    capacity,
234                    vortex_buffer::BufferAllocatorRef::static_ref(),
235                ),
236                max_length,
237            )),
238            DType::Binary(_) => Box::new(TruncatedMaxBinaryStatsBuilder::<ByteBuffer>::new(
239                values_builder,
240                BoolBuilder::with_capacity_in(
241                    Nullability::NonNullable,
242                    capacity,
243                    vortex_buffer::BufferAllocatorRef::static_ref(),
244                ),
245                max_length,
246            )),
247            _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)),
248        },
249        Stat::Min => match dtype {
250            DType::Utf8(_) => Box::new(TruncatedMinBinaryStatsBuilder::<BufferString>::new(
251                values_builder,
252                BoolBuilder::with_capacity_in(
253                    Nullability::NonNullable,
254                    capacity,
255                    vortex_buffer::BufferAllocatorRef::static_ref(),
256                ),
257                max_length,
258            )),
259            DType::Binary(_) => Box::new(TruncatedMinBinaryStatsBuilder::<ByteBuffer>::new(
260                values_builder,
261                BoolBuilder::with_capacity_in(
262                    Nullability::NonNullable,
263                    capacity,
264                    vortex_buffer::BufferAllocatorRef::static_ref(),
265                ),
266                max_length,
267            )),
268            _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)),
269        },
270        _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)),
271    }
272}
273
274/// Arrays with their associated names, reduced version of a `StructArray`.
275struct NamedArrays {
276    names: Vec<FieldName>,
277    arrays: Vec<ArrayRef>,
278}
279
280impl NamedArrays {
281    fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
282        self.arrays[0].all_invalid(ctx)
283    }
284}
285
286trait StatsArrayBuilder: Send {
287    fn stat(&self) -> Stat;
288
289    fn append_scalar(&mut self, value: Scalar) -> VortexResult<()>;
290
291    fn append_null(&mut self);
292
293    fn finish(&mut self) -> NamedArrays;
294}
295
296struct StatNameArrayBuilder {
297    stat: Stat,
298    builder: Box<dyn ArrayBuilder>,
299}
300
301impl StatNameArrayBuilder {
302    fn new(stat: Stat, builder: Box<dyn ArrayBuilder>) -> Self {
303        Self { stat, builder }
304    }
305}
306
307impl StatsArrayBuilder for StatNameArrayBuilder {
308    fn stat(&self) -> Stat {
309        self.stat
310    }
311
312    fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> {
313        self.builder.append_scalar(&value)
314    }
315
316    fn append_null(&mut self) {
317        self.builder.append_null()
318    }
319
320    fn finish(&mut self) -> NamedArrays {
321        NamedArrays {
322            names: vec![self.stat.name().into()],
323            arrays: vec![self.builder.finish()],
324        }
325    }
326}
327
328struct TruncatedMaxBinaryStatsBuilder<T: ScalarTruncation> {
329    values: Box<dyn ArrayBuilder>,
330    is_truncated: BoolBuilder,
331    max_value_length: usize,
332    _marker: PhantomData<T>,
333}
334
335impl<T: ScalarTruncation> TruncatedMaxBinaryStatsBuilder<T> {
336    fn new(
337        values: Box<dyn ArrayBuilder>,
338        is_truncated: BoolBuilder,
339        max_value_length: usize,
340    ) -> Self {
341        Self {
342            values,
343            is_truncated,
344            max_value_length,
345            _marker: PhantomData,
346        }
347    }
348}
349
350struct TruncatedMinBinaryStatsBuilder<T: ScalarTruncation> {
351    values: Box<dyn ArrayBuilder>,
352    is_truncated: BoolBuilder,
353    max_value_length: usize,
354    _marker: PhantomData<T>,
355}
356
357impl<T: ScalarTruncation> TruncatedMinBinaryStatsBuilder<T> {
358    fn new(
359        values: Box<dyn ArrayBuilder>,
360        is_truncated: BoolBuilder,
361        max_value_length: usize,
362    ) -> Self {
363        Self {
364            values,
365            is_truncated,
366            max_value_length,
367            _marker: PhantomData,
368        }
369    }
370}
371
372impl<T: ScalarTruncation> StatsArrayBuilder for TruncatedMaxBinaryStatsBuilder<T> {
373    fn stat(&self) -> Stat {
374        Stat::Max
375    }
376
377    fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> {
378        let nullability = value.dtype().nullability();
379        if let Some((upper_bound, truncated)) =
380            upper_bound(T::from_scalar(value)?, self.max_value_length, nullability)
381        {
382            self.values.append_scalar(&upper_bound)?;
383            self.is_truncated.append_value(truncated);
384        } else {
385            self.append_null()
386        }
387        Ok(())
388    }
389
390    fn append_null(&mut self) {
391        ArrayBuilder::append_null(self.values.as_mut());
392        self.is_truncated.append_value(false);
393    }
394
395    fn finish(&mut self) -> NamedArrays {
396        NamedArrays {
397            names: vec![Stat::Max.name().into(), MAX_IS_TRUNCATED.into()],
398            arrays: vec![
399                ArrayBuilder::finish(self.values.as_mut()),
400                ArrayBuilder::finish(&mut self.is_truncated),
401            ],
402        }
403    }
404}
405
406impl<T: ScalarTruncation> StatsArrayBuilder for TruncatedMinBinaryStatsBuilder<T> {
407    fn stat(&self) -> Stat {
408        Stat::Min
409    }
410
411    fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> {
412        let nullability = value.dtype().nullability();
413        if let Some((lower_bound, truncated)) =
414            lower_bound(T::from_scalar(value)?, self.max_value_length, nullability)
415        {
416            self.values.append_scalar(&lower_bound)?;
417            self.is_truncated.append_value(truncated);
418        } else {
419            self.append_null()
420        }
421        Ok(())
422    }
423
424    fn append_null(&mut self) {
425        ArrayBuilder::append_null(self.values.as_mut());
426        self.is_truncated.append_value(false);
427    }
428
429    fn finish(&mut self) -> NamedArrays {
430        NamedArrays {
431            names: vec![Stat::Min.name().into(), MIN_IS_TRUNCATED.into()],
432            arrays: vec![
433                ArrayBuilder::finish(self.values.as_mut()),
434                ArrayBuilder::finish(&mut self.is_truncated),
435            ],
436        }
437    }
438}
439
440/// An array stream processor that computes aggregate statistics for all fields.
441///
442/// Note: for now this only collects top-level struct fields.
443#[derive(Clone)]
444pub struct FileStatsAccumulator {
445    stats: Arc<[Stat]>,
446    accumulators: Arc<Mutex<Vec<StatsAccumulator>>>,
447    ctx: Arc<Mutex<ExecutionCtx>>,
448}
449
450impl FileStatsAccumulator {
451    fn new(
452        dtype: &DType,
453        stats: Arc<[Stat]>,
454        max_variable_length_statistics_size: usize,
455        session: &VortexSession,
456    ) -> Self {
457        let accumulators = Arc::new(Mutex::new(match dtype.as_struct_fields_opt() {
458            Some(struct_dtype) => {
459                if dtype.nullability() == Nullability::Nullable {
460                    // top level dtype could be nullable, but we don't support it yet
461                    vortex_panic!(
462                        "FileStatsAccumulator temporarily does not support nullable top-level structs, got: {}. Use Validity::NonNullable",
463                        dtype
464                    );
465                }
466
467                struct_dtype
468                    .fields()
469                    .map(|field_dtype| {
470                        StatsAccumulator::new(
471                            &field_dtype,
472                            &stats,
473                            max_variable_length_statistics_size,
474                        )
475                    })
476                    .collect()
477            }
478            None => [StatsAccumulator::new(
479                dtype,
480                &stats,
481                max_variable_length_statistics_size,
482            )]
483            .into(),
484        }));
485
486        Self {
487            stats,
488            accumulators,
489            ctx: Arc::new(Mutex::new(session.create_execution_ctx())),
490        }
491    }
492
493    fn process(
494        &self,
495        chunk: VortexResult<(SequenceId, ArrayRef)>,
496    ) -> VortexResult<(SequenceId, ArrayRef)> {
497        let (sequence_id, chunk) = chunk?;
498        let mut ctx = self.ctx.lock();
499        if chunk.dtype().is_struct() {
500            let struct_chunk = chunk.clone().execute::<StructArray>(&mut ctx)?;
501            for (acc, field) in self
502                .accumulators
503                .lock()
504                .iter_mut()
505                .zip_eq(struct_chunk.iter_unmasked_fields())
506            {
507                acc.push_chunk(field, &mut ctx)?;
508            }
509        } else {
510            self.accumulators.lock()[0].push_chunk(&chunk, &mut ctx)?;
511        }
512        Ok((sequence_id, chunk))
513    }
514
515    pub fn stats_sets(&self) -> Vec<StatsSet> {
516        let mut ctx = self.ctx.lock();
517        self.accumulators
518            .lock()
519            .iter_mut()
520            .map(|acc| {
521                acc.as_stats_set(&self.stats, &mut ctx)
522                    .vortex_expect("as_stats_table should not fail")
523            })
524            .collect()
525    }
526}
527
528#[cfg(test)]
529mod tests {
530    use rstest::rstest;
531    use vortex_array::IntoArray;
532    use vortex_array::array_session;
533    use vortex_array::arrays::BoolArray;
534    use vortex_array::arrays::bool::BoolArrayExt;
535    use vortex_array::builders::VarBinViewBuilder;
536    use vortex_buffer::BitBuffer;
537    use vortex_buffer::buffer;
538
539    use super::*;
540
541    #[rstest]
542    #[case(DType::Utf8(Nullability::NonNullable))]
543    #[case(DType::Binary(Nullability::NonNullable))]
544    fn truncates_accumulated_stats(#[case] dtype: DType) {
545        let mut ctx = array_session().create_execution_ctx();
546        let mut builder = VarBinViewBuilder::with_capacity_in(
547            dtype.clone(),
548            2,
549            vortex_buffer::BufferAllocatorRef::statically_allocated(),
550        );
551        builder.append_value("Value to be truncated");
552        builder.append_value("untruncated");
553        let mut builder2 = VarBinViewBuilder::with_capacity_in(
554            dtype,
555            2,
556            vortex_buffer::BufferAllocatorRef::statically_allocated(),
557        );
558        builder2.append_value("Another");
559        builder2.append_value("wait a minute");
560        let mut acc =
561            StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12);
562        acc.push_chunk(&builder.finish(), &mut ctx)
563            .vortex_expect("push_chunk should succeed for test data");
564        acc.push_chunk(&builder2.finish(), &mut ctx)
565            .vortex_expect("push_chunk should succeed for test data");
566        let stats_table = acc
567            .as_array(&mut ctx)
568            .unwrap()
569            .expect("Must have stats table");
570        assert_eq!(
571            stats_table.names().as_ref(),
572            &[
573                Stat::Max.name(),
574                MAX_IS_TRUNCATED,
575                Stat::Min.name(),
576                MIN_IS_TRUNCATED,
577            ]
578        );
579        let field1_bool = stats_table
580            .unmasked_field(1)
581            .clone()
582            .execute::<BoolArray>(&mut ctx)
583            .unwrap();
584        assert_eq!(
585            field1_bool.to_bit_buffer(),
586            BitBuffer::from(vec![false, true])
587        );
588        let field3_bool = stats_table
589            .unmasked_field(3)
590            .clone()
591            .execute::<BoolArray>(&mut ctx)
592            .unwrap();
593        assert_eq!(
594            field3_bool.to_bit_buffer(),
595            BitBuffer::from(vec![true, false])
596        );
597    }
598
599    #[rstest]
600    #[case(DType::Utf8(Nullability::NonNullable))]
601    #[case(DType::Binary(Nullability::NonNullable))]
602    fn truncated_accumulated_stats_are_inexact(#[case] dtype: DType) {
603        let mut ctx = array_session().create_execution_ctx();
604        let mut builder = VarBinViewBuilder::with_capacity_in(
605            dtype,
606            2,
607            vortex_buffer::BufferAllocatorRef::statically_allocated(),
608        );
609        builder.append_value("Value to be truncated");
610        builder.append_value("Another truncated value");
611        let mut acc = StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min], 12);
612        acc.push_chunk(&builder.finish(), &mut ctx)
613            .vortex_expect("push_chunk should succeed for test data");
614
615        let stats = acc
616            .as_stats_set(&[Stat::Max, Stat::Min], &mut ctx)
617            .vortex_expect("as_stats_set should succeed for test data");
618
619        assert!(matches!(stats.get(Stat::Min), Precision::Inexact(_)));
620        assert!(matches!(stats.get(Stat::Max), Precision::Inexact(_)));
621    }
622
623    #[test]
624    fn fixed_width_stats_omit_is_truncated_columns() {
625        let mut ctx = array_session().create_execution_ctx();
626        let array = buffer![0, 1, 2].into_array();
627        let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12);
628        acc.push_chunk(&array, &mut ctx)
629            .vortex_expect("push_chunk should succeed for test array");
630        let stats_table = acc
631            .as_array(&mut ctx)
632            .unwrap()
633            .expect("Must have stats table");
634        assert_eq!(
635            stats_table.names().as_ref(),
636            &[Stat::Max.name(), Stat::Min.name(), Stat::Sum.name()]
637        );
638    }
639}
vortex_layout/layouts/file_stats.rs

vortex_layout/layouts/
file_stats.rs
