Skip to content

Commit c04ebf9

Browse files
gatesngatesn
committed
Use aggregate descriptors for zoned stats
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 0473953 commit c04ebf9

10 files changed

Lines changed: 689 additions & 237 deletions

File tree

‎vortex-array/src/stats/rewrite/builtins.rs‎

Lines changed: 43 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -136,12 +136,18 @@ impl StatsRewriteRule for BoundedBinaryStatsRewrite {
136136
fn falsify(
137137
&self,
138138
expr: &Expression,
139-
_ctx: &StatsRewriteCtx<'_>,
139+
ctx: &StatsRewriteCtx<'_>,
140140
) -> VortexResult<Option<Expression>> {
141141
let operator = expr.as_::<Binary>();
142142
let lhs = expr.child(0);
143143
let rhs = expr.child(1);
144144

145+
if !supports_bounded_rewrite(&ctx.return_dtype(lhs)?)
146+
|| !supports_bounded_rewrite(&ctx.return_dtype(rhs)?)
147+
{
148+
return Ok(None);
149+
}
150+
145151
Ok(match operator {
146152
Operator::Eq => {
147153
let left = gt(bounded_min(lhs), bounded_max(rhs));
@@ -432,6 +438,10 @@ fn has_nans(dtype: &DType) -> bool {
432438
matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float())
433439
}
434440

441+
fn supports_bounded_rewrite(dtype: &DType) -> bool {
442+
matches!(dtype, DType::Binary(_) | DType::Utf8(_))
443+
}
444+
435445
fn stat_expr(expr: &Expression, stat: Stat) -> Option<Expression> {
436446
if let Some(literal) = literal_stat(expr, stat) {
437447
return Some(literal);
@@ -570,6 +580,8 @@ mod tests {
570580
StructFields::from_iter([
571581
("a", DType::Primitive(PType::I32, Nullability::NonNullable)),
572582
("b", DType::Primitive(PType::I32, Nullability::NonNullable)),
583+
("s", DType::Utf8(Nullability::NonNullable)),
584+
("t", DType::Utf8(Nullability::NonNullable)),
573585
]),
574586
Nullability::NonNullable,
575587
)
@@ -592,29 +604,38 @@ mod tests {
592604
let expr = gt(col("a"), lit(10));
593605
assert_eq!(
594606
falsify(&expr)?,
595-
Some(or(
596-
and(
597-
nan_free(col("a")),
598-
lt_eq(stat(col("a"), Stat::Max), lit(10)),
599-
),
600-
lt_eq(bounded_max(&col("a")), lit(10)),
607+
Some(and(
608+
nan_free(col("a")),
609+
lt_eq(stat(col("a"), Stat::Max), lit(10)),
601610
))
602611
);
603612

604613
let expr = eq(col("a"), col("b"));
614+
assert_eq!(
615+
falsify(&expr)?,
616+
Some(and(
617+
and(nan_free(col("a")), nan_free(col("b"))),
618+
or(
619+
gt(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)),
620+
gt(stat(col("b"), Stat::Min), stat(col("a"), Stat::Max)),
621+
),
622+
))
623+
);
624+
625+
let expr = eq(col("s"), col("t"));
605626
assert_eq!(
606627
falsify(&expr)?,
607628
Some(or(
608629
and(
609-
and(nan_free(col("a")), nan_free(col("b"))),
630+
and(nan_free(col("s")), nan_free(col("t"))),
610631
or(
611-
gt(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)),
612-
gt(stat(col("b"), Stat::Min), stat(col("a"), Stat::Max)),
632+
gt(stat(col("s"), Stat::Min), stat(col("t"), Stat::Max)),
633+
gt(stat(col("t"), Stat::Min), stat(col("s"), Stat::Max)),
613634
),
614635
),
615636
or(
616-
gt(bounded_min(&col("a")), bounded_max(&col("b"))),
617-
gt(bounded_min(&col("b")), bounded_max(&col("a"))),
637+
gt(bounded_min(&col("s")), bounded_max(&col("t"))),
638+
gt(bounded_min(&col("t")), bounded_max(&col("s"))),
618639
),
619640
))
620641
);
@@ -627,19 +648,13 @@ mod tests {
627648
assert_eq!(
628649
falsify(&expr)?,
629650
Some(or(
630-
or(
631-
and(
632-
nan_free(col("a")),
633-
lt_eq(stat(col("a"), Stat::Max), lit(10)),
634-
),
635-
lt_eq(bounded_max(&col("a")), lit(10)),
651+
and(
652+
nan_free(col("a")),
653+
lt_eq(stat(col("a"), Stat::Max), lit(10)),
636654
),
637-
or(
638-
and(
639-
nan_free(col("a")),
640-
gt_eq(stat(col("a"), Stat::Min), lit(50)),
641-
),
642-
gt_eq(bounded_min(&col("a")), lit(50)),
655+
and(
656+
nan_free(col("a")),
657+
gt_eq(stat(col("a"), Stat::Min), lit(50)),
643658
),
644659
))
645660
);
@@ -661,14 +676,8 @@ mod tests {
661676
assert_eq!(
662677
falsify(&expr)?,
663678
Some(or(
664-
or(
665-
and(nan_free(col("a")), gt(lit(10), stat(col("a"), Stat::Max))),
666-
gt(lit(10), bounded_max(&col("a"))),
667-
),
668-
or(
669-
and(nan_free(col("a")), gt(stat(col("a"), Stat::Min), lit(50))),
670-
gt(bounded_min(&col("a")), lit(50)),
671-
),
679+
and(nan_free(col("a")), gt(lit(10), stat(col("a"), Stat::Max))),
680+
and(nan_free(col("a")), gt(stat(col("a"), Stat::Min), lit(50))),
672681
))
673682
);
674683
Ok(())
@@ -762,14 +771,8 @@ mod tests {
762771
assert_eq!(
763772
falsify(&expr)?,
764773
Some(or(
765-
or(
766-
gt(cast(stat(col("a"), Stat::Min), dtype.clone()), lit(42i64)),
767-
gt(lit(42i64), cast(stat(col("a"), Stat::Max), dtype.clone())),
768-
),
769-
or(
770-
gt(cast(bounded_min(&col("a")), dtype.clone()), lit(42i64)),
771-
gt(lit(42i64), cast(bounded_max(&col("a")), dtype)),
772-
),
774+
gt(cast(stat(col("a"), Stat::Min), dtype.clone()), lit(42i64)),
775+
gt(lit(42i64), cast(stat(col("a"), Stat::Max), dtype)),
773776
))
774777
);
775778
Ok(())

‎vortex-layout/public-api.lock‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -766,13 +766,9 @@ pub mod vortex_layout::layouts::zoned::writer
766766

767767
pub struct vortex_layout::layouts::zoned::writer::ZonedLayoutOptions
768768

769-
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::block_size: usize
770-
771-
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::concurrency: usize
769+
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::aggregate_fns: alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>
772770

773-
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::max_variable_length_statistics_size: usize
774-
775-
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::stats: alloc::sync::Arc<[vortex_array::expr::stats::Stat]>
771+
pub vortex_layout::layouts::zoned::writer::ZonedLayoutOptions::block_size: usize
776772

777773
impl core::default::Default for vortex_layout::layouts::zoned::writer::ZonedLayoutOptions
778774

@@ -800,7 +796,7 @@ pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::dtype_for_stats_table(&
800796

801797
pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::prune(&self, &vortex_array::expr::expression::Expression, &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_mask::Mask>
802798

803-
pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::try_new(vortex_array::dtype::DType, vortex_array::arrays::struct_::vtable::StructArray, alloc::sync::Arc<[vortex_array::expr::stats::Stat]>, u64, u64) -> vortex_error::VortexResult<Self>
799+
pub fn vortex_layout::layouts::zoned::zone_map::ZoneMap::try_new(vortex_array::dtype::DType, vortex_array::arrays::struct_::vtable::StructArray, alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>, u64, u64) -> vortex_error::VortexResult<Self>
804800

805801
impl core::clone::Clone for vortex_layout::layouts::zoned::zone_map::ZoneMap
806802

@@ -848,11 +844,11 @@ pub struct vortex_layout::layouts::zoned::ZonedLayout
848844

849845
impl vortex_layout::layouts::zoned::ZonedLayout
850846

851-
pub fn vortex_layout::layouts::zoned::ZonedLayout::new(vortex_layout::LayoutRef, vortex_layout::LayoutRef, usize, alloc::sync::Arc<[vortex_array::expr::stats::Stat]>) -> Self
847+
pub fn vortex_layout::layouts::zoned::ZonedLayout::new(vortex_layout::LayoutRef, vortex_layout::LayoutRef, usize, alloc::sync::Arc<[vortex_array::aggregate_fn::erased::AggregateFnRef]>) -> Self
852848

853849
pub fn vortex_layout::layouts::zoned::ZonedLayout::nzones(&self) -> usize
854850

855-
pub fn vortex_layout::layouts::zoned::ZonedLayout::present_stats(&self) -> &alloc::sync::Arc<[vortex_array::expr::stats::Stat]>
851+
pub fn vortex_layout::layouts::zoned::ZonedLayout::present_aggregates(&self) -> &alloc::sync::Arc<[alloc::string::String]>
856852

857853
pub fn vortex_layout::layouts::zoned::ZonedLayout::zone_len(&self) -> usize
858854

‎vortex-layout/src/layouts/zoned/builder.rs‎

Lines changed: 129 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use vortex_array::ExecutionCtx;
1212
use vortex_array::IntoArray;
1313
use vortex_array::LEGACY_SESSION;
1414
use vortex_array::VortexSessionExecute;
15+
use vortex_array::aggregate_fn::AggregateFnRef;
1516
use vortex_array::aggregate_fn::fns::sum::sum;
1617
use vortex_array::arrays::ConstantArray;
1718
use vortex_array::arrays::StructArray;
@@ -25,7 +26,6 @@ use vortex_array::dtype::Nullability;
2526
use vortex_array::dtype::PType;
2627
use vortex_array::expr::stats::Precision;
2728
use vortex_array::expr::stats::Stat;
28-
use vortex_array::expr::stats::StatsProvider;
2929
use vortex_array::scalar::Scalar;
3030
use vortex_array::scalar::ScalarTruncation;
3131
use vortex_array::scalar::lower_bound;
@@ -35,9 +35,12 @@ use vortex_array::validity::Validity;
3535
use vortex_buffer::BufferString;
3636
use vortex_buffer::ByteBuffer;
3737
use vortex_error::VortexResult;
38+
use vortex_error::vortex_ensure_eq;
3839

3940
use crate::layouts::zoned::schema::MAX_IS_TRUNCATED;
4041
use crate::layouts::zoned::schema::MIN_IS_TRUNCATED;
42+
use crate::layouts::zoned::schema::aggregate_descriptor;
43+
use crate::layouts::zoned::schema::aggregate_state_dtype;
4144

4245
/// Accumulates write-time statistics for each logical zone.
4346
pub struct StatsAccumulator {
@@ -74,18 +77,6 @@ impl StatsAccumulator {
7477
}
7578
}
7679

77-
pub fn push_chunk_without_compute(&mut self, array: &ArrayRef) -> VortexResult<()> {
78-
for builder in &mut self.builders {
79-
if let Precision::Exact(value) = array.statistics().get(builder.stat()) {
80-
builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?;
81-
} else {
82-
builder.append_null();
83-
}
84-
}
85-
self.length += 1;
86-
Ok(())
87-
}
88-
8980
pub fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
9081
for builder in &mut self.builders {
9182
if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? {
@@ -176,6 +167,102 @@ fn supports_file_stats(dtype: &DType) -> bool {
176167
!matches!(dtype, DType::Variant(_))
177168
}
178169

170+
/// Accumulates aggregate-function partials for each logical zone.
171+
pub(crate) struct AggregateStatsAccumulator {
172+
builders: Vec<AggregateStatsArrayBuilder>,
173+
length: usize,
174+
}
175+
176+
impl AggregateStatsAccumulator {
177+
pub(crate) fn new(dtype: &DType, aggregate_fns: &[AggregateFnRef]) -> Self {
178+
let builders = aggregate_fns
179+
.iter()
180+
.filter_map(|aggregate_fn| {
181+
aggregate_state_dtype(dtype, aggregate_fn).map(|partial_dtype| {
182+
AggregateStatsArrayBuilder::new(
183+
aggregate_fn.clone(),
184+
&partial_dtype.as_nullable(),
185+
1024,
186+
)
187+
})
188+
})
189+
.collect::<Vec<_>>();
190+
191+
Self {
192+
builders,
193+
length: 0,
194+
}
195+
}
196+
197+
pub(crate) fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
198+
self.builders
199+
.iter()
200+
.map(|builder| builder.aggregate_fn.clone())
201+
.collect::<Vec<_>>()
202+
.into()
203+
}
204+
205+
pub(crate) fn push_partials(&mut self, partials: Vec<Scalar>) -> VortexResult<()> {
206+
vortex_ensure_eq!(
207+
partials.len(),
208+
self.builders.len(),
209+
"aggregate partial count must match zone stats builder count"
210+
);
211+
212+
for (builder, value) in self.builders.iter_mut().zip_eq(partials) {
213+
builder.append_scalar(value)?;
214+
}
215+
self.length += 1;
216+
Ok(())
217+
}
218+
219+
pub(crate) fn as_array(
220+
&mut self,
221+
) -> VortexResult<Option<(StructArray, Arc<[AggregateFnRef]>)>> {
222+
let mut names = Vec::new();
223+
let mut fields = Vec::new();
224+
let mut aggregate_fns = Vec::new();
225+
226+
for builder in self
227+
.builders
228+
.iter_mut()
229+
.sorted_unstable_by(|lhs, rhs| lhs.descriptor.cmp(&rhs.descriptor))
230+
{
231+
let values = builder.finish();
232+
233+
if values.all_invalid()? {
234+
continue;
235+
}
236+
237+
aggregate_fns.push(builder.aggregate_fn.clone());
238+
names.extend(values.names);
239+
fields.extend(values.arrays);
240+
}
241+
242+
if names.is_empty() {
243+
return Ok(None);
244+
}
245+
246+
let array = StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable)?;
247+
Ok(Some((array, aggregate_fns.into())))
248+
}
249+
}
250+
251+
pub(crate) fn aggregate_partials(
252+
array: &ArrayRef,
253+
aggregate_fns: &[AggregateFnRef],
254+
ctx: &mut ExecutionCtx,
255+
) -> VortexResult<Vec<Scalar>> {
256+
aggregate_fns
257+
.iter()
258+
.map(|aggregate_fn| {
259+
let mut accumulator = aggregate_fn.accumulator(array.dtype())?;
260+
accumulator.accumulate(array, ctx)?;
261+
accumulator.partial_scalar()
262+
})
263+
.collect()
264+
}
265+
179266
fn stats_builder_with_capacity(
180267
stat: Stat,
181268
dtype: &DType,
@@ -214,6 +301,35 @@ fn stats_builder_with_capacity(
214301
}
215302
}
216303

304+
struct AggregateStatsArrayBuilder {
305+
aggregate_fn: AggregateFnRef,
306+
descriptor: String,
307+
dtype: DType,
308+
builder: Box<dyn ArrayBuilder>,
309+
}
310+
311+
impl AggregateStatsArrayBuilder {
312+
fn new(aggregate_fn: AggregateFnRef, dtype: &DType, capacity: usize) -> Self {
313+
Self {
314+
descriptor: aggregate_descriptor(&aggregate_fn),
315+
aggregate_fn,
316+
dtype: dtype.clone(),
317+
builder: builder_with_capacity(dtype, capacity),
318+
}
319+
}
320+
321+
fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> {
322+
self.builder.append_scalar(&value.cast(&self.dtype)?)
323+
}
324+
325+
fn finish(&mut self) -> NamedArrays {
326+
NamedArrays {
327+
names: vec![self.descriptor.clone().into()],
328+
arrays: vec![self.builder.finish()],
329+
}
330+
}
331+
}
332+
217333
/// Arrays with their associated names, reduced version of a `StructArray`.
218334
struct NamedArrays {
219335
names: Vec<FieldName>,

0 commit comments

Comments
 (0)