Skip to content

Commit 4c41dfc

Browse files
gatesngatesn
committed
MinMaxAggregateFn
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent d1fc11f commit 4c41dfc

5 files changed

Lines changed: 153 additions & 13 deletions

File tree

‎encodings/runend/src/compute/min_max.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use vortex_array::ArrayRef;
55
use vortex_array::ExecutionCtx;
66
use vortex_array::aggregate_fn::AggregateFnRef;
77
use vortex_array::aggregate_fn::fns::min_max::MinMax;
8-
use vortex_array::aggregate_fn::fns::min_max::make_struct_dtype;
8+
use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype;
99
use vortex_array::aggregate_fn::fns::min_max::min_max;
1010
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
1111
use vortex_array::scalar::Scalar;
@@ -35,7 +35,7 @@ impl DynAggregateKernel for RunEndMinMaxKernel {
3535
return Ok(None);
3636
};
3737

38-
let struct_dtype = make_struct_dtype(batch.dtype());
38+
let struct_dtype = make_minmax_dtype(batch.dtype());
3939
match min_max(run_end.values(), ctx)? {
4040
Some(result) => Ok(Some(Scalar::struct_(
4141
struct_dtype,

‎encodings/sequence/src/compute/min_max.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use vortex_array::ArrayRef;
55
use vortex_array::ExecutionCtx;
66
use vortex_array::aggregate_fn::AggregateFnRef;
77
use vortex_array::aggregate_fn::fns::min_max::MinMax;
8-
use vortex_array::aggregate_fn::fns::min_max::make_struct_dtype;
8+
use vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype;
99
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
1010
use vortex_array::dtype::DType;
1111
use vortex_array::dtype::Nullability;
@@ -39,7 +39,7 @@ impl DynAggregateKernel for SequenceMinMaxKernel {
3939
return Ok(None);
4040
};
4141

42-
let struct_dtype = make_struct_dtype(batch.dtype());
42+
let struct_dtype = make_minmax_dtype(batch.dtype());
4343

4444
// Empty sequences shouldn't exist (try_new validates length), but handle gracefully.
4545
if seq.is_empty() {

‎vortex-array/public-api.lock‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ pub fn vortex_array::aggregate_fn::fns::min_max::MinMaxResult::fmt(&self, f: &mu
102102

103103
impl core::marker::StructuralPartialEq for vortex_array::aggregate_fn::fns::min_max::MinMaxResult
104104

105-
pub fn vortex_array::aggregate_fn::fns::min_max::make_struct_dtype(element_dtype: &vortex_array::dtype::DType) -> vortex_array::dtype::DType
105+
pub fn vortex_array::aggregate_fn::fns::min_max::make_minmax_dtype(element_dtype: &vortex_array::dtype::DType) -> vortex_array::dtype::DType
106106

107107
pub fn vortex_array::aggregate_fn::fns::min_max::min_max(array: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::aggregate_fn::fns::min_max::MinMaxResult>>
108108

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

Lines changed: 46 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ pub fn min_max(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Option<
6969

7070
// Short-circuit for unsupported dtypes.
7171
if MinMax.return_dtype(&EmptyOptions, array.dtype()).is_none() {
72-
vortex_bail!("min_max not supported for dtype {}", array.dtype());
72+
return Ok(None);
7373
}
7474

7575
// Compute using Accumulator<MinMax>.
@@ -155,7 +155,7 @@ impl MinMaxPartial {
155155
}
156156

157157
/// Creates the struct dtype `{min: T, max: T}` (nullable) used for min/max aggregate results.
158-
pub fn make_struct_dtype(element_dtype: &DType) -> DType {
158+
pub fn make_minmax_dtype(element_dtype: &DType) -> DType {
159159
DType::Struct(
160160
StructFields::new(
161161
NAMES.clone(),
@@ -178,8 +178,13 @@ impl AggregateFnVTable for MinMax {
178178

179179
fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
180180
match input_dtype {
181-
DType::Null => None,
182-
_ => Some(make_struct_dtype(input_dtype)),
181+
DType::Bool(_)
182+
| DType::Primitive(..)
183+
| DType::Decimal(..)
184+
| DType::Utf8(..)
185+
| DType::Binary(..)
186+
| DType::Extension(..) => Some(make_minmax_dtype(input_dtype)),
187+
_ => None,
183188
}
184189
}
185190

@@ -206,7 +211,7 @@ impl AggregateFnVTable for MinMax {
206211
}
207212

208213
fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
209-
let dtype = make_struct_dtype(&partial.element_dtype);
214+
let dtype = make_minmax_dtype(&partial.element_dtype);
210215
let result = match (partial.min.take(), partial.max.take()) {
211216
(Some(min), Some(max)) => Scalar::struct_(dtype, vec![min, max]),
212217
_ => Scalar::null(dtype),
@@ -435,7 +440,7 @@ mod tests {
435440
let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
436441
let mut state = MinMax.empty_partial(&EmptyOptions, &dtype)?;
437442

438-
let struct_dtype = crate::aggregate_fn::fns::min_max::make_struct_dtype(&dtype);
443+
let struct_dtype = crate::aggregate_fn::fns::min_max::make_minmax_dtype(&dtype);
439444
let scalar1 = Scalar::struct_(
440445
struct_dtype.clone(),
441446
vec![Scalar::from(5i32), Scalar::from(15i32)],
@@ -551,6 +556,30 @@ mod tests {
551556
})
552557
);
553558

559+
let result = min_max(
560+
&BoolArray::from_iter(vec![None, Some(true), Some(true)]).into_array(),
561+
&mut ctx,
562+
)?;
563+
assert_eq!(
564+
result,
565+
Some(MinMaxResult {
566+
min: Scalar::bool(true, Nullability::NonNullable),
567+
max: Scalar::bool(true, Nullability::NonNullable),
568+
})
569+
);
570+
571+
let result = min_max(
572+
&BoolArray::from_iter(vec![None, Some(true), Some(true), None]).into_array(),
573+
&mut ctx,
574+
)?;
575+
assert_eq!(
576+
result,
577+
Some(MinMaxResult {
578+
min: Scalar::bool(true, Nullability::NonNullable),
579+
max: Scalar::bool(true, Nullability::NonNullable),
580+
})
581+
);
582+
554583
let result = min_max(
555584
&BoolArray::from_iter(vec![Some(false), Some(false), None, None]).into_array(),
556585
&mut ctx,
@@ -564,4 +593,15 @@ mod tests {
564593
);
565594
Ok(())
566595
}
596+
597+
#[test]
598+
fn test_varbin_all_nulls() -> VortexResult<()> {
599+
let array = VarBinArray::from_iter(
600+
vec![Option::<&str>::None, None, None],
601+
DType::Utf8(Nullability::Nullable),
602+
);
603+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
604+
assert_eq!(min_max(&array.into_array(), &mut ctx)?, None);
605+
Ok(())
606+
}
567607
}

‎vortex-array/src/arrays/dict/compute/min_max.rs‎

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::ArrayRef;
88
use crate::ExecutionCtx;
99
use crate::aggregate_fn::AggregateFnRef;
1010
use crate::aggregate_fn::fns::min_max::MinMax;
11-
use crate::aggregate_fn::fns::min_max::make_struct_dtype;
11+
use crate::aggregate_fn::fns::min_max::make_minmax_dtype;
1212
use crate::aggregate_fn::fns::min_max::min_max;
1313
use crate::aggregate_fn::kernels::DynAggregateKernel;
1414
use crate::arrays::Dict;
@@ -36,7 +36,7 @@ impl DynAggregateKernel for DictMinMaxKernel {
3636
return Ok(None);
3737
};
3838

39-
let struct_dtype = make_struct_dtype(batch.dtype());
39+
let struct_dtype = make_minmax_dtype(batch.dtype());
4040

4141
let result = if dict.has_all_values_referenced() {
4242
// All values are referenced, compute min/max directly on the values array.
@@ -55,3 +55,103 @@ impl DynAggregateKernel for DictMinMaxKernel {
5555
}
5656
}
5757
}
58+
59+
#[cfg(test)]
60+
mod tests {
61+
use rstest::rstest;
62+
use vortex_buffer::buffer;
63+
use vortex_error::VortexResult;
64+
65+
use crate::ArrayRef;
66+
use crate::IntoArray;
67+
use crate::LEGACY_SESSION;
68+
use crate::VortexSessionExecute;
69+
use crate::aggregate_fn::fns::min_max::min_max;
70+
use crate::arrays::DictArray;
71+
use crate::arrays::PrimitiveArray;
72+
use crate::builders::dict::dict_encode;
73+
74+
fn assert_min_max(array: &ArrayRef, expected: Option<(i32, i32)>) -> VortexResult<()> {
75+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
76+
match (min_max(array, &mut ctx)?, expected) {
77+
(Some(result), Some((expected_min, expected_max))) => {
78+
assert_eq!(i32::try_from(&result.min)?, expected_min);
79+
assert_eq!(i32::try_from(&result.max)?, expected_max);
80+
}
81+
(None, None) => {}
82+
(got, expected) => panic!(
83+
"min_max mismatch: expected {expected:?}, got {:?}",
84+
got.as_ref().map(|r| (
85+
i32::try_from(&r.min.clone()).ok(),
86+
i32::try_from(&r.max.clone()).ok()
87+
))
88+
),
89+
}
90+
Ok(())
91+
}
92+
93+
#[rstest]
94+
#[case::covering(
95+
DictArray::try_new(
96+
buffer![0u32, 1, 2, 3, 0, 1].into_array(),
97+
buffer![10i32, 20, 30, 40].into_array(),
98+
).unwrap(),
99+
(10, 40)
100+
)]
101+
#[case::non_covering_duplicates(
102+
DictArray::try_new(
103+
buffer![1u32, 1, 1, 3, 3].into_array(),
104+
buffer![1i32, 2, 3, 4, 5].into_array(),
105+
).unwrap(),
106+
(2, 4)
107+
)]
108+
#[case::non_covering_gaps(
109+
DictArray::try_new(
110+
buffer![0u32, 2, 4].into_array(),
111+
buffer![1i32, 2, 3, 4, 5].into_array(),
112+
).unwrap(),
113+
(1, 5)
114+
)]
115+
#[case::single(dict_encode(&buffer![42i32].into_array()).unwrap(), (42, 42))]
116+
#[case::nullable_codes(
117+
DictArray::try_new(
118+
PrimitiveArray::from_option_iter([Some(0u32), None, Some(1), Some(2)]).into_array(),
119+
buffer![10i32, 20, 30].into_array(),
120+
).unwrap(),
121+
(10, 30)
122+
)]
123+
#[case::nullable_values(
124+
dict_encode(
125+
&PrimitiveArray::from_option_iter([Some(1i32), None, Some(2), Some(1), None]).into_array()
126+
).unwrap(),
127+
(1, 2)
128+
)]
129+
fn test_min_max(#[case] dict: DictArray, #[case] expected: (i32, i32)) -> VortexResult<()> {
130+
assert_min_max(&dict.into_array(), Some(expected))
131+
}
132+
133+
#[test]
134+
fn test_sliced_dict() -> VortexResult<()> {
135+
let reference = PrimitiveArray::from_iter([1, 5, 10, 50, 100]);
136+
let dict = dict_encode(&reference.into_array())?;
137+
let sliced = dict.slice(1..3)?;
138+
assert_min_max(&sliced, Some((5, 10)))
139+
}
140+
141+
#[rstest]
142+
#[case::empty(
143+
DictArray::try_new(
144+
PrimitiveArray::from_iter(Vec::<u32>::new()).into_array(),
145+
buffer![10i32, 20, 30].into_array(),
146+
).unwrap()
147+
)]
148+
#[case::all_null_codes(
149+
DictArray::try_new(
150+
PrimitiveArray::from_option_iter([Option::<u32>::None, None, None]).into_array(),
151+
buffer![10i32, 20, 30].into_array(),
152+
).unwrap()
153+
)]
154+
fn test_min_max_none(#[case] dict: DictArray) -> VortexResult<()> {
155+
assert_min_max(&dict.into_array(), None)
156+
}
157+
}

0 commit comments

Comments
 (0)