Skip to content

Commit d963a44

Browse files
gatesngatesn
committed
Fix NaN pruning aggregate support
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent abdf067 commit d963a44

5 files changed

Lines changed: 93 additions & 65 deletions

File tree

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

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ impl AggregateFnVTable for AllNan {
3838
Ok(None)
3939
}
4040

41-
fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
42-
Some(DType::Bool(Nullability::NonNullable))
41+
fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
42+
has_nans(input_dtype).then_some(DType::Bool(Nullability::Nullable))
4343
}
4444

4545
fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -49,9 +49,9 @@ impl AggregateFnVTable for AllNan {
4949
fn empty_partial(
5050
&self,
5151
_options: &Self::Options,
52-
input_dtype: &DType,
52+
_input_dtype: &DType,
5353
) -> VortexResult<Self::Partial> {
54-
Ok(has_nans(input_dtype))
54+
Ok(true)
5555
}
5656

5757
fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
@@ -60,7 +60,7 @@ impl AggregateFnVTable for AllNan {
6060
}
6161

6262
fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
63-
Ok(Scalar::bool(*partial, Nullability::NonNullable))
63+
Ok(Scalar::bool(*partial, Nullability::Nullable))
6464
}
6565

6666
fn reset(&self, partial: &mut Self::Partial) {
@@ -161,24 +161,31 @@ mod tests {
161161
}
162162

163163
#[test]
164-
fn all_nan_false_for_non_float_values() -> VortexResult<()> {
165-
let mut ctx = LEGACY_SESSION.create_execution_ctx();
164+
fn all_nan_unsupported_for_non_float_values() -> VortexResult<()> {
166165
let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
166+
assert!(Accumulator::try_new(AllNan, EmptyOptions, dtype).is_err());
167+
Ok(())
168+
}
169+
170+
#[test]
171+
fn all_nan_false_with_null() -> VortexResult<()> {
172+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
173+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
167174
let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
168175

169-
let batch = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array();
176+
let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), None]).into_array();
170177
acc.accumulate(&batch, &mut ctx)?;
171178

172179
assert!(!bool::try_from(&acc.finish()?)?);
173180
Ok(())
174181
}
175182

176183
#[test]
177-
fn all_nan_false_for_empty_non_float_values() -> VortexResult<()> {
178-
let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
184+
fn all_nan_true_for_empty_float_values() -> VortexResult<()> {
185+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
179186
let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
180187

181-
assert!(!bool::try_from(&acc.finish()?)?);
188+
assert!(bool::try_from(&acc.finish()?)?);
182189
Ok(())
183190
}
184191
}

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

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ impl AggregateFnVTable for AllNonNan {
3838
Ok(None)
3939
}
4040

41-
fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
42-
Some(DType::Bool(Nullability::NonNullable))
41+
fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
42+
has_nans(input_dtype).then_some(DType::Bool(Nullability::Nullable))
4343
}
4444

4545
fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
@@ -60,7 +60,7 @@ impl AggregateFnVTable for AllNonNan {
6060
}
6161

6262
fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
63-
Ok(Scalar::bool(*partial, Nullability::NonNullable))
63+
Ok(Scalar::bool(*partial, Nullability::Nullable))
6464
}
6565

6666
fn reset(&self, partial: &mut Self::Partial) {
@@ -104,6 +104,10 @@ impl AggregateFnVTable for AllNonNan {
104104
}
105105
}
106106

107+
fn has_nans(dtype: &DType) -> bool {
108+
matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float())
109+
}
110+
107111
#[cfg(test)]
108112
mod tests {
109113
use vortex_error::VortexResult;
@@ -147,12 +151,28 @@ mod tests {
147151
}
148152

149153
#[test]
150-
fn all_non_nan_true_for_non_float() -> VortexResult<()> {
151-
let mut ctx = LEGACY_SESSION.create_execution_ctx();
154+
fn all_non_nan_unsupported_for_non_float() -> VortexResult<()> {
152155
let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
156+
assert!(Accumulator::try_new(AllNonNan, EmptyOptions, dtype).is_err());
157+
Ok(())
158+
}
159+
160+
#[test]
161+
fn all_non_nan_true_for_empty_float() -> VortexResult<()> {
162+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
163+
let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
164+
165+
assert!(bool::try_from(&acc.finish()?)?);
166+
Ok(())
167+
}
168+
169+
#[test]
170+
fn all_non_nan_true_with_nulls() -> VortexResult<()> {
171+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
172+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
153173
let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?;
154174

155-
let batch = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array();
175+
let batch = PrimitiveArray::from_option_iter([Some(1.0f32), None]).into_array();
156176
acc.accumulate(&batch, &mut ctx)?;
157177

158178
assert!(bool::try_from(&acc.finish()?)?);

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

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use std::fmt::Formatter;
66

77
use vortex_buffer::BufferString;
88
use vortex_buffer::ByteBuffer;
9+
use vortex_error::VortexExpect;
910
use vortex_error::VortexResult;
1011
use vortex_error::vortex_ensure;
1112
use vortex_session::VortexSession;
@@ -65,7 +66,7 @@ impl BoundedMaxPartial {
6566
self.state = match std::mem::replace(&mut self.state, BoundedMaxState::Empty) {
6667
BoundedMaxState::Empty => BoundedMaxState::Value(max),
6768
BoundedMaxState::Value(current) => BoundedMaxState::Value(
68-
partial_max(max, current).expect("incomparable bounded max scalars"),
69+
partial_max(max, current).vortex_expect("incomparable bounded max scalars"),
6970
),
7071
BoundedMaxState::Unknown => BoundedMaxState::Unknown,
7172
};
@@ -187,6 +188,25 @@ fn supported_dtype<'a>(options: &BoundedMaxOptions, input_dtype: &'a DType) -> O
187188
.map(|_| input_dtype)
188189
}
189190

191+
fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult<Option<Scalar>> {
192+
let nullability = value.dtype().nullability();
193+
match value.dtype() {
194+
DType::Utf8(_) => {
195+
Ok(
196+
upper_bound(BufferString::from_scalar(value)?, max_bytes, nullability)
197+
.map(|(bound, _)| bound),
198+
)
199+
}
200+
DType::Binary(_) => {
201+
Ok(
202+
upper_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability)
203+
.map(|(bound, _)| bound),
204+
)
205+
}
206+
_ => Ok(Some(value)),
207+
}
208+
}
209+
190210
#[cfg(test)]
191211
mod tests {
192212
use vortex_buffer::buffer;
@@ -270,22 +290,3 @@ mod tests {
270290
Ok(())
271291
}
272292
}
273-
274-
fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult<Option<Scalar>> {
275-
let nullability = value.dtype().nullability();
276-
match value.dtype() {
277-
DType::Utf8(_) => {
278-
Ok(
279-
upper_bound(BufferString::from_scalar(value)?, max_bytes, nullability)
280-
.map(|(bound, _)| bound),
281-
)
282-
}
283-
DType::Binary(_) => {
284-
Ok(
285-
upper_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability)
286-
.map(|(bound, _)| bound),
287-
)
288-
}
289-
_ => Ok(Some(value)),
290-
}
291-
}

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

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use std::fmt::Formatter;
66

77
use vortex_buffer::BufferString;
88
use vortex_buffer::ByteBuffer;
9+
use vortex_error::VortexExpect;
910
use vortex_error::VortexResult;
1011
use vortex_error::vortex_ensure;
1112
use vortex_session::VortexSession;
@@ -56,7 +57,9 @@ impl BoundedMinPartial {
5657
}
5758

5859
self.min = Some(match self.min.take() {
59-
Some(current) => partial_min(min, current).expect("incomparable bounded min scalars"),
60+
Some(current) => {
61+
partial_min(min, current).vortex_expect("incomparable bounded min scalars")
62+
}
6063
None => min,
6164
});
6265
}
@@ -172,6 +175,25 @@ fn supported_dtype<'a>(options: &BoundedMinOptions, input_dtype: &'a DType) -> O
172175
.map(|_| input_dtype)
173176
}
174177

178+
fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult<Option<Scalar>> {
179+
let nullability = value.dtype().nullability();
180+
match value.dtype() {
181+
DType::Utf8(_) => {
182+
Ok(
183+
lower_bound(BufferString::from_scalar(value)?, max_bytes, nullability)
184+
.map(|(bound, _)| bound),
185+
)
186+
}
187+
DType::Binary(_) => {
188+
Ok(
189+
lower_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability)
190+
.map(|(bound, _)| bound),
191+
)
192+
}
193+
_ => Ok(Some(value)),
194+
}
195+
}
196+
175197
#[cfg(test)]
176198
mod tests {
177199
use vortex_buffer::buffer;
@@ -243,22 +265,3 @@ mod tests {
243265
Ok(())
244266
}
245267
}
246-
247-
fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult<Option<Scalar>> {
248-
let nullability = value.dtype().nullability();
249-
match value.dtype() {
250-
DType::Utf8(_) => {
251-
Ok(
252-
lower_bound(BufferString::from_scalar(value)?, max_bytes, nullability)
253-
.map(|(bound, _)| bound),
254-
)
255-
}
256-
DType::Binary(_) => {
257-
Ok(
258-
lower_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability)
259-
.map(|(bound, _)| bound),
260-
)
261-
}
262-
_ => Ok(Some(value)),
263-
}
264-
}

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

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -322,18 +322,15 @@ mod tests {
322322
}
323323

324324
#[test]
325-
fn stat_expr_reads_all_nan_false_for_empty_non_float() -> VortexResult<()> {
325+
fn stat_expr_rejects_all_nan_for_non_float() -> VortexResult<()> {
326326
let array = PrimitiveArray::empty::<i32>(Nullability::NonNullable).into_array();
327+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
327328

328329
let result = array
329-
.apply(&super::all_nan(root()))?
330-
.execute::<Canonical>(&mut LEGACY_SESSION.create_execution_ctx())?
331-
.into_array();
332-
333-
let expected =
334-
ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 0).into_array();
335-
assert_arrays_eq!(result, expected);
330+
.apply(&super::all_nan(root()))
331+
.and_then(|array| array.execute::<Canonical>(&mut ctx));
336332

333+
assert!(result.is_err());
337334
Ok(())
338335
}
339336

0 commit comments

Comments
 (0)