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

vortex_layout/layouts/
file_stats.rs
