Skip to content

Commit 9c91f1a

Browse files
thorfourthorfour
committed
replace footer usage of Stats enum with AggPartialRef
1 parent de12c75 commit 9c91f1a

6 files changed

Lines changed: 141 additions & 100 deletions

File tree

‎vortex-array/src/stats/mod.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ pub use array::*;
3131
pub use session::*;
3232
use vortex_error::VortexExpect;
3333

34+
use crate::aggregate_fn::AggregateFnRef;
3435
use crate::expr::stats::Stat;
3536

3637
/// Statistics that are used for pruning files (i.e., we want to ensure they are computed when compressing/writing).
@@ -43,6 +44,14 @@ pub const PRUNING_STATS: &[Stat] = &[
4344
Stat::NaNCount,
4445
];
4546

47+
/// [`PRUNING_STATS`], converted to their aggregate function form.
48+
pub fn pruning_aggregate_fns() -> Vec<AggregateFnRef> {
49+
PRUNING_STATS
50+
.iter()
51+
.filter_map(|stat| stat.aggregate_fn())
52+
.collect()
53+
}
54+
4655
pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec<u8> {
4756
let max_stat = u8::from(last::<Stat>().vortex_expect("last stat")) as usize + 1;
4857
// TODO(ngates): use vortex-buffer::BitBuffer

‎vortex-duckdb/src/copy.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -338,8 +338,9 @@ pub fn copy_to_initialize_global(
338338
#[cfg(test)]
339339
mod tests {
340340
use vortex::array::IntoArray;
341+
use vortex::array::aggregate_fn::AggregateFnRef;
341342
use vortex::array::arrays::StructArray;
342-
use vortex::array::stats::PRUNING_STATS;
343+
use vortex::array::stats::pruning_aggregate_fns;
343344
use vortex::array::validity::Validity;
344345
use vortex::buffer::ByteBufferMut;
345346
use vortex::buffer::buffer;
@@ -348,7 +349,7 @@ mod tests {
348349

349350
/// Writes a one-column file and returns its summary, with `file_statistics` controlling which
350351
/// statistics the footer carries (empty means none at all).
351-
fn write_summary(file_statistics: Vec<Stat>) -> WriteSummary {
352+
fn write_summary(file_statistics: Vec<AggregateFnRef>) -> WriteSummary {
352353
RUNTIME.block_on(async {
353354
let array = StructArray::from_fields(&[("i", buffer![1u32, 2, 3].into_array())])
354355
.unwrap()
@@ -365,7 +366,7 @@ mod tests {
365366

366367
#[test]
367368
fn column_stats_out_of_range_is_an_error() {
368-
let summary = write_summary(PRUNING_STATS.to_vec());
369+
let summary = write_summary(pruning_aggregate_fns());
369370
assert!(column_stats_from_summary(&summary, 0, &[]).is_ok());
370371
assert!(column_stats_from_summary(&summary, 1, &[]).is_err());
371372
}
@@ -393,7 +394,7 @@ mod tests {
393394
let mut buf = ByteBufferMut::empty();
394395
let mut writer = SESSION
395396
.write_options()
396-
.with_file_statistics(PRUNING_STATS.to_vec())
397+
.with_file_statistics(pruning_aggregate_fns())
397398
.writer(&mut buf, outer.dtype().clone());
398399
writer.push(outer).await.unwrap();
399400
writer.finish().await.unwrap()

‎vortex-duckdb/src/file_reader.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -339,7 +339,7 @@ pub fn footer_get_statistics(footer: &Footer, index: usize) -> Option<ColumnStat
339339
mod tests {
340340
use vortex::array::IntoArray;
341341
use vortex::array::arrays::StructArray;
342-
use vortex::array::stats::PRUNING_STATS;
342+
use vortex::array::stats::pruning_aggregate_fns;
343343
use vortex::array::validity::Validity;
344344
use vortex::buffer::ByteBufferMut;
345345
use vortex::buffer::buffer;
@@ -366,7 +366,7 @@ mod tests {
366366
let mut buf = ByteBufferMut::empty();
367367
let mut writer = SESSION
368368
.write_options()
369-
.with_file_statistics(PRUNING_STATS.to_vec())
369+
.with_file_statistics(pruning_aggregate_fns())
370370
.writer(&mut buf, outer.dtype().clone());
371371
writer.push(outer).await.unwrap();
372372
writer.finish().await.unwrap().footer().clone()

‎vortex-file/src/tests.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ use vortex_array::scalar::ScalarValue;
6464
use vortex_array::scalar_fn::ScalarFnVTableExt;
6565
use vortex_array::scalar_fn::fns::pack::Pack;
6666
use vortex_array::scalar_fn::fns::pack::PackOptions;
67-
use vortex_array::stats::PRUNING_STATS;
67+
use vortex_array::stats::pruning_aggregate_fns;
6868
use vortex_array::stream::ArrayStreamAdapter;
6969
use vortex_array::stream::ArrayStreamExt;
7070
use vortex_array::validity::Validity;
@@ -1322,7 +1322,7 @@ async fn write_nullable_top_level_struct() -> VortexResult<()> {
13221322
let mut buf = ByteBufferMut::empty();
13231323
let summary = SESSION
13241324
.write_options()
1325-
.with_file_statistics(PRUNING_STATS.to_vec())
1325+
.with_file_statistics(pruning_aggregate_fns())
13261326
.write(&mut buf, array.to_array_stream())
13271327
.await?;
13281328

@@ -1359,14 +1359,14 @@ async fn exclude_legacy_statistics_omits_legacy_but_keeps_nested() -> VortexResu
13591359
let mut buf_with_legacy = ByteBufferMut::empty();
13601360
SESSION
13611361
.write_options()
1362-
.with_file_statistics(PRUNING_STATS.to_vec())
1362+
.with_file_statistics(pruning_aggregate_fns())
13631363
.write(&mut buf_with_legacy, array.to_array_stream())
13641364
.await?;
13651365

13661366
let mut buf_without_legacy = ByteBufferMut::empty();
13671367
let summary = SESSION
13681368
.write_options()
1369-
.with_file_statistics(PRUNING_STATS.to_vec())
1369+
.with_file_statistics(pruning_aggregate_fns())
13701370
.exclude_legacy_statistics()
13711371
.write(&mut buf_without_legacy, array.to_array_stream())
13721372
.await?;
@@ -2089,7 +2089,7 @@ async fn test_writer_with_statistics() -> VortexResult<()> {
20892089
let mut buf = ByteBufferMut::empty();
20902090
let mut writer = SESSION
20912091
.write_options()
2092-
.with_file_statistics(PRUNING_STATS.to_vec())
2092+
.with_file_statistics(pruning_aggregate_fns())
20932093
.writer(&mut buf, array.dtype().clone());
20942094

20952095
writer.push(array).await?;

‎vortex-file/src/writer.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,13 @@ use itertools::Itertools;
1919
use vortex_array::ArrayContext;
2020
use vortex_array::ArrayId;
2121
use vortex_array::ArrayRef;
22+
use vortex_array::aggregate_fn::AggregateFnRef;
2223
use vortex_array::dtype::DType;
2324
use vortex_array::dtype::FieldPath;
24-
use vortex_array::expr::stats::Stat;
2525
use vortex_array::iter::ArrayIterator;
2626
use vortex_array::iter::ArrayIteratorExt;
2727
use vortex_array::session::ArraySessionExt;
28-
use vortex_array::stats::PRUNING_STATS;
28+
use vortex_array::stats::pruning_aggregate_fns;
2929
use vortex_array::stream::ArrayStream;
3030
use vortex_array::stream::ArrayStreamAdapter;
3131
use vortex_array::stream::ArrayStreamExt;
@@ -84,7 +84,7 @@ pub struct VortexWriteOptions {
8484
buffered_bytes: BufferedBytesTracker,
8585
exclude_dtype: bool,
8686
max_variable_length_statistics_size: usize,
87-
file_statistics: Vec<Stat>,
87+
file_statistics: Vec<AggregateFnRef>,
8888
write_legacy_statistics: bool,
8989
metadata: HashMap<String, ByteBuffer>,
9090
}
@@ -107,7 +107,7 @@ impl VortexWriteOptions {
107107
buffered_bytes: BufferedBytesTracker::new(),
108108
session,
109109
exclude_dtype: false,
110-
file_statistics: PRUNING_STATS.to_vec(),
110+
file_statistics: pruning_aggregate_fns(),
111111
write_legacy_statistics: true,
112112
max_variable_length_statistics_size: 64,
113113
metadata: HashMap::default(),
@@ -159,7 +159,7 @@ impl VortexWriteOptions {
159159
/// Configure which statistics to compute at the file level.
160160
///
161161
/// Pass an empty vector to omit file-level statistics.
162-
pub fn with_file_statistics(mut self, file_statistics: Vec<Stat>) -> Self {
162+
pub fn with_file_statistics(mut self, file_statistics: Vec<AggregateFnRef>) -> Self {
163163
self.file_statistics = file_statistics;
164164
self
165165
}

0 commit comments

Comments
 (0)