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

vortex_layout/layouts/
file_stats.rs
