Skip to content

Commit e09432d

Browse files
gatesngatesn
committed
Add aggregate satisfaction relation
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 4dec695 commit e09432d

8 files changed

Lines changed: 301 additions & 0 deletions

File tree

‎vortex-array/public-api.lock‎

Lines changed: 132 additions & 0 deletions
Large diffs are not rendered by default.

‎vortex-array/src/aggregate_fn/erased.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use vortex_utils::debug_with::DebugWith;
1616

1717
use crate::aggregate_fn::AccumulatorRef;
1818
use crate::aggregate_fn::AggregateFnId;
19+
use crate::aggregate_fn::AggregateFnSatisfaction;
1920
use crate::aggregate_fn::AggregateFnVTable;
2021
use crate::aggregate_fn::GroupedAccumulatorRef;
2122
use crate::aggregate_fn::options::AggregateFnOptions;
@@ -74,6 +75,11 @@ impl AggregateFnRef {
7475
AggregateFnOptions { inner: &*self.0 }
7576
}
7677

78+
/// Return whether this stored aggregate can satisfy `requested`.
79+
pub fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction {
80+
self.0.can_satisfy(requested)
81+
}
82+
7783
/// Coerce the input type for this aggregate function.
7884
pub fn coerce_args(&self, input_dtype: &DType) -> VortexResult<DType> {
7985
self.0.coerce_args(input_dtype)

‎vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@ use crate::Columnar;
1717
use crate::ExecutionCtx;
1818
use crate::IntoArray;
1919
use crate::aggregate_fn::AggregateFnId;
20+
use crate::aggregate_fn::AggregateFnRef;
21+
use crate::aggregate_fn::AggregateFnSatisfaction;
2022
use crate::aggregate_fn::AggregateFnVTable;
2123
use crate::aggregate_fn::EmptyOptions;
24+
use crate::aggregate_fn::fns::max::Max;
2225
use crate::aggregate_fn::fns::min_max::MinMax;
2326
use crate::aggregate_fn::fns::min_max::min_max;
2427
use crate::dtype::DType;
@@ -115,6 +118,25 @@ impl AggregateFnVTable for BoundedMax {
115118
supported_dtype(options, input_dtype).map(DType::as_nullable)
116119
}
117120

121+
fn can_satisfy(
122+
&self,
123+
options: &Self::Options,
124+
requested: &AggregateFnRef,
125+
) -> AggregateFnSatisfaction {
126+
if requested
127+
.as_opt::<Self>()
128+
.is_some_and(|other| other == options)
129+
{
130+
return AggregateFnSatisfaction::Exact;
131+
}
132+
133+
if requested.is::<Self>() || requested.is::<Max>() {
134+
AggregateFnSatisfaction::Approximate
135+
} else {
136+
AggregateFnSatisfaction::No
137+
}
138+
}
139+
118140
fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
119141
self.return_dtype(options, input_dtype)
120142
}
@@ -219,10 +241,15 @@ mod tests {
219241
use crate::LEGACY_SESSION;
220242
use crate::VortexSessionExecute;
221243
use crate::aggregate_fn::Accumulator;
244+
use crate::aggregate_fn::AggregateFnSatisfaction;
222245
use crate::aggregate_fn::AggregateFnVTable;
246+
use crate::aggregate_fn::AggregateFnVTableExt;
223247
use crate::aggregate_fn::DynAccumulator;
248+
use crate::aggregate_fn::EmptyOptions;
224249
use crate::aggregate_fn::fns::bounded_max::BoundedMax;
225250
use crate::aggregate_fn::fns::bounded_max::BoundedMaxOptions;
251+
use crate::aggregate_fn::fns::max::Max;
252+
use crate::aggregate_fn::fns::min::Min;
226253
use crate::arrays::PrimitiveArray;
227254
use crate::arrays::VarBinViewArray;
228255
use crate::dtype::Nullability;
@@ -333,6 +360,33 @@ mod tests {
333360
Ok(())
334361
}
335362

363+
#[test]
364+
fn bounded_max_satisfies_max_bounds() {
365+
let stored = BoundedMax.bind(BoundedMaxOptions {
366+
max_bytes: max_bytes(5),
367+
});
368+
let same = BoundedMax.bind(BoundedMaxOptions {
369+
max_bytes: max_bytes(5),
370+
});
371+
let other_bounded = BoundedMax.bind(BoundedMaxOptions {
372+
max_bytes: max_bytes(6),
373+
});
374+
375+
assert_eq!(stored.can_satisfy(&same), AggregateFnSatisfaction::Exact);
376+
assert_eq!(
377+
stored.can_satisfy(&other_bounded),
378+
AggregateFnSatisfaction::Approximate
379+
);
380+
assert_eq!(
381+
stored.can_satisfy(&Max.bind(EmptyOptions)),
382+
AggregateFnSatisfaction::Approximate
383+
);
384+
assert_eq!(
385+
stored.can_satisfy(&Min.bind(EmptyOptions)),
386+
AggregateFnSatisfaction::No
387+
);
388+
}
389+
336390
#[test]
337391
fn bounded_max_options_round_trip() -> VortexResult<()> {
338392
let options = BoundedMaxOptions {

‎vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@ use crate::Columnar;
1717
use crate::ExecutionCtx;
1818
use crate::IntoArray;
1919
use crate::aggregate_fn::AggregateFnId;
20+
use crate::aggregate_fn::AggregateFnRef;
21+
use crate::aggregate_fn::AggregateFnSatisfaction;
2022
use crate::aggregate_fn::AggregateFnVTable;
2123
use crate::aggregate_fn::EmptyOptions;
24+
use crate::aggregate_fn::fns::min::Min;
2225
use crate::aggregate_fn::fns::min_max::MinMax;
2326
use crate::aggregate_fn::fns::min_max::min_max;
2427
use crate::dtype::DType;
@@ -103,6 +106,25 @@ impl AggregateFnVTable for BoundedMin {
103106
supported_dtype(options, input_dtype).map(DType::as_nullable)
104107
}
105108

109+
fn can_satisfy(
110+
&self,
111+
options: &Self::Options,
112+
requested: &AggregateFnRef,
113+
) -> AggregateFnSatisfaction {
114+
if requested
115+
.as_opt::<Self>()
116+
.is_some_and(|other| other == options)
117+
{
118+
return AggregateFnSatisfaction::Exact;
119+
}
120+
121+
if requested.is::<Self>() || requested.is::<Min>() {
122+
AggregateFnSatisfaction::Approximate
123+
} else {
124+
AggregateFnSatisfaction::No
125+
}
126+
}
127+
106128
fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
107129
self.return_dtype(options, input_dtype)
108130
}
@@ -206,10 +228,15 @@ mod tests {
206228
use crate::LEGACY_SESSION;
207229
use crate::VortexSessionExecute;
208230
use crate::aggregate_fn::Accumulator;
231+
use crate::aggregate_fn::AggregateFnSatisfaction;
209232
use crate::aggregate_fn::AggregateFnVTable;
233+
use crate::aggregate_fn::AggregateFnVTableExt;
210234
use crate::aggregate_fn::DynAccumulator;
235+
use crate::aggregate_fn::EmptyOptions;
211236
use crate::aggregate_fn::fns::bounded_min::BoundedMin;
212237
use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions;
238+
use crate::aggregate_fn::fns::max::Max;
239+
use crate::aggregate_fn::fns::min::Min;
213240
use crate::arrays::PrimitiveArray;
214241
use crate::arrays::VarBinViewArray;
215242
use crate::dtype::Nullability;
@@ -263,6 +290,33 @@ mod tests {
263290
Ok(())
264291
}
265292

293+
#[test]
294+
fn bounded_min_satisfies_min_bounds() {
295+
let stored = BoundedMin.bind(BoundedMinOptions {
296+
max_bytes: max_bytes(5),
297+
});
298+
let same = BoundedMin.bind(BoundedMinOptions {
299+
max_bytes: max_bytes(5),
300+
});
301+
let other_bounded = BoundedMin.bind(BoundedMinOptions {
302+
max_bytes: max_bytes(6),
303+
});
304+
305+
assert_eq!(stored.can_satisfy(&same), AggregateFnSatisfaction::Exact);
306+
assert_eq!(
307+
stored.can_satisfy(&other_bounded),
308+
AggregateFnSatisfaction::Approximate
309+
);
310+
assert_eq!(
311+
stored.can_satisfy(&Min.bind(EmptyOptions)),
312+
AggregateFnSatisfaction::Approximate
313+
);
314+
assert_eq!(
315+
stored.can_satisfy(&Max.bind(EmptyOptions)),
316+
AggregateFnSatisfaction::No
317+
);
318+
}
319+
266320
#[test]
267321
fn bounded_min_options_round_trip() -> VortexResult<()> {
268322
let options = BoundedMinOptions {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ use vortex_session::registry::Id;
1111
mod accumulator;
1212
pub use accumulator::*;
1313

14+
mod satisfaction;
15+
pub use satisfaction::*;
16+
1417
mod accumulator_grouped;
1518
pub use accumulator_grouped::*;
1619

‎vortex-array/src/aggregate_fn/satisfaction.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
/// Whether a stored aggregate function can satisfy a requested aggregate function.
5+
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
6+
pub enum AggregateFnSatisfaction {
7+
/// The stored aggregate cannot satisfy the requested aggregate.
8+
#[default]
9+
No,
10+
/// The stored aggregate can satisfy the request as an approximate bound.
11+
Approximate,
12+
/// The stored aggregate exactly satisfies the request.
13+
Exact,
14+
}
15+
16+
impl AggregateFnSatisfaction {
17+
/// Returns whether the stored aggregate can satisfy the requested aggregate.
18+
pub fn is_satisfied(self) -> bool {
19+
!matches!(self, Self::No)
20+
}
21+
22+
/// Returns whether the stored aggregate exactly satisfies the requested aggregate.
23+
pub fn is_exact(self) -> bool {
24+
matches!(self, Self::Exact)
25+
}
26+
}

‎vortex-array/src/aggregate_fn/typed.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ use crate::aggregate_fn::Accumulator;
2424
use crate::aggregate_fn::AccumulatorRef;
2525
use crate::aggregate_fn::AggregateFnId;
2626
use crate::aggregate_fn::AggregateFnRef;
27+
use crate::aggregate_fn::AggregateFnSatisfaction;
2728
use crate::aggregate_fn::AggregateFnVTable;
2829
use crate::aggregate_fn::GroupedAccumulator;
2930
use crate::aggregate_fn::GroupedAccumulatorRef;
@@ -39,6 +40,7 @@ pub(super) trait DynAggregateFn: 'static + Send + Sync + super::sealed::Sealed {
3940
fn options_any(&self) -> &dyn Any;
4041

4142
fn coerce_args(&self, input_dtype: &DType) -> VortexResult<DType>;
43+
fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction;
4244
fn return_dtype(&self, input_dtype: &DType) -> Option<DType>;
4345
fn state_dtype(&self, input_dtype: &DType) -> Option<DType>;
4446
fn accumulator(&self, input_dtype: &DType) -> VortexResult<AccumulatorRef>;
@@ -80,6 +82,10 @@ impl<V: AggregateFnVTable> DynAggregateFn for AggregateFnInner<V> {
8082
V::coerce_args(&self.vtable, &self.options, input_dtype)
8183
}
8284

85+
fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction {
86+
V::can_satisfy(&self.vtable, &self.options, requested)
87+
}
88+
8389
fn return_dtype(&self, input_dtype: &DType) -> Option<DType> {
8490
V::return_dtype(&self.vtable, &self.options, input_dtype)
8591
}

‎vortex-array/src/aggregate_fn/vtable.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use crate::ExecutionCtx;
1717
use crate::aggregate_fn::AggregateFn;
1818
use crate::aggregate_fn::AggregateFnId;
1919
use crate::aggregate_fn::AggregateFnRef;
20+
use crate::aggregate_fn::AggregateFnSatisfaction;
2021
use crate::dtype::DType;
2122
use crate::scalar::Scalar;
2223

@@ -66,6 +67,25 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
6667
Ok(input_dtype.clone())
6768
}
6869

70+
/// Return whether this stored aggregate can satisfy `requested`.
71+
///
72+
/// The default implementation only treats exactly equal aggregate functions as satisfying the
73+
/// request. Approximate pruning aggregates can override this to expose looser-but-sound bounds.
74+
fn can_satisfy(
75+
&self,
76+
options: &Self::Options,
77+
requested: &AggregateFnRef,
78+
) -> AggregateFnSatisfaction {
79+
if requested
80+
.as_opt::<Self>()
81+
.is_some_and(|other| other == options)
82+
{
83+
AggregateFnSatisfaction::Exact
84+
} else {
85+
AggregateFnSatisfaction::No
86+
}
87+
}
88+
6989
/// The return [`DType`] of the aggregate.
7090
///
7191
/// Returns `None` if the aggregate function cannot be applied to the input dtype.

0 commit comments

Comments
 (0)