Skip to content

Commit f9cf88c

Browse files
committed
Teach zoned pruning to lower StatFn
Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 9700eec commit f9cf88c

12 files changed

Lines changed: 930 additions & 70 deletions

File tree

vortex-array/public-api.lock

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,54 @@ pub type vortex_array::aggregate_fn::combined::CombinedOptions<T> = vortex_array
174174

175175
pub mod vortex_array::aggregate_fn::fns
176176

177+
pub mod vortex_array::aggregate_fn::fns::all_nan
178+
179+
pub struct vortex_array::aggregate_fn::fns::all_nan::AllNan
180+
181+
impl core::clone::Clone for vortex_array::aggregate_fn::fns::all_nan::AllNan
182+
183+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::clone(&self) -> vortex_array::aggregate_fn::fns::all_nan::AllNan
184+
185+
impl core::fmt::Debug for vortex_array::aggregate_fn::fns::all_nan::AllNan
186+
187+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
188+
189+
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_nan::AllNan
190+
191+
pub type vortex_array::aggregate_fn::fns::all_nan::AllNan::Options = vortex_array::aggregate_fn::EmptyOptions
192+
193+
pub type vortex_array::aggregate_fn::fns::all_nan::AllNan::Partial = bool
194+
195+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::accumulate(&self, &mut Self::Partial, &vortex_array::Columnar, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
196+
197+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::coerce_args(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
198+
199+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::combine_partials(&self, &mut Self::Partial, vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
200+
201+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
202+
203+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::empty_partial(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
204+
205+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::finalize(&self, vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
206+
207+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::finalize_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
208+
209+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
210+
211+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::is_saturated(&self, &Self::Partial) -> bool
212+
213+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::partial_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
214+
215+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::reset(&self, &mut Self::Partial)
216+
217+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::return_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
218+
219+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::serialize(&self, &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
220+
221+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::to_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
222+
223+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
224+
177225
pub mod vortex_array::aggregate_fn::fns::all_non_distinct
178226

179227
pub struct vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct
@@ -226,6 +274,54 @@ pub struct vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinctPart
226274

227275
pub fn vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct(&vortex_array::ArrayRef, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
228276

277+
pub mod vortex_array::aggregate_fn::fns::all_non_nan
278+
279+
pub struct vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
280+
281+
impl core::clone::Clone for vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
282+
283+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::clone(&self) -> vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
284+
285+
impl core::fmt::Debug for vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
286+
287+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
288+
289+
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
290+
291+
pub type vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::Options = vortex_array::aggregate_fn::EmptyOptions
292+
293+
pub type vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::Partial = bool
294+
295+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::accumulate(&self, &mut Self::Partial, &vortex_array::Columnar, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
296+
297+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::coerce_args(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
298+
299+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::combine_partials(&self, &mut Self::Partial, vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
300+
301+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
302+
303+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::empty_partial(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
304+
305+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::finalize(&self, vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
306+
307+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::finalize_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
308+
309+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
310+
311+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::is_saturated(&self, &Self::Partial) -> bool
312+
313+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::partial_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
314+
315+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::reset(&self, &mut Self::Partial)
316+
317+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::return_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
318+
319+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::serialize(&self, &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
320+
321+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::to_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
322+
323+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
324+
229325
pub mod vortex_array::aggregate_fn::fns::all_non_null
230326

231327
pub struct vortex_array::aggregate_fn::fns::all_non_null::AllNonNull
@@ -1334,6 +1430,42 @@ pub fn vortex_array::aggregate_fn::AggregateFnVTable::to_scalar(&self, &Self::Pa
13341430

13351431
pub fn vortex_array::aggregate_fn::AggregateFnVTable::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
13361432

1433+
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_nan::AllNan
1434+
1435+
pub type vortex_array::aggregate_fn::fns::all_nan::AllNan::Options = vortex_array::aggregate_fn::EmptyOptions
1436+
1437+
pub type vortex_array::aggregate_fn::fns::all_nan::AllNan::Partial = bool
1438+
1439+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::accumulate(&self, &mut Self::Partial, &vortex_array::Columnar, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
1440+
1441+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::coerce_args(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
1442+
1443+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::combine_partials(&self, &mut Self::Partial, vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
1444+
1445+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
1446+
1447+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::empty_partial(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
1448+
1449+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::finalize(&self, vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
1450+
1451+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::finalize_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
1452+
1453+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
1454+
1455+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::is_saturated(&self, &Self::Partial) -> bool
1456+
1457+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::partial_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
1458+
1459+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::reset(&self, &mut Self::Partial)
1460+
1461+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::return_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
1462+
1463+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::serialize(&self, &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
1464+
1465+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::to_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
1466+
1467+
pub fn vortex_array::aggregate_fn::fns::all_nan::AllNan::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
1468+
13371469
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct
13381470

13391471
pub type vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct::Options = vortex_array::aggregate_fn::EmptyOptions
@@ -1370,6 +1502,42 @@ pub fn vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct::to_sca
13701502

13711503
pub fn vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
13721504

1505+
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan
1506+
1507+
pub type vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::Options = vortex_array::aggregate_fn::EmptyOptions
1508+
1509+
pub type vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::Partial = bool
1510+
1511+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::accumulate(&self, &mut Self::Partial, &vortex_array::Columnar, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<()>
1512+
1513+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::coerce_args(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::dtype::DType>
1514+
1515+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::combine_partials(&self, &mut Self::Partial, vortex_array::scalar::Scalar) -> vortex_error::VortexResult<()>
1516+
1517+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::deserialize(&self, &[u8], &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Options>
1518+
1519+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::empty_partial(&self, &Self::Options, &vortex_array::dtype::DType) -> vortex_error::VortexResult<Self::Partial>
1520+
1521+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::finalize(&self, vortex_array::ArrayRef) -> vortex_error::VortexResult<vortex_array::ArrayRef>
1522+
1523+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::finalize_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
1524+
1525+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::id(&self) -> vortex_array::aggregate_fn::AggregateFnId
1526+
1527+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::is_saturated(&self, &Self::Partial) -> bool
1528+
1529+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::partial_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
1530+
1531+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::reset(&self, &mut Self::Partial)
1532+
1533+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::return_dtype(&self, &Self::Options, &vortex_array::dtype::DType) -> core::option::Option<vortex_array::dtype::DType>
1534+
1535+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::serialize(&self, &Self::Options) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
1536+
1537+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::to_scalar(&self, &Self::Partial) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
1538+
1539+
pub fn vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan::try_accumulate(&self, &mut Self::Partial, &vortex_array::ArrayRef, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<bool>
1540+
13731541
impl vortex_array::aggregate_fn::AggregateFnVTable for vortex_array::aggregate_fn::fns::all_non_null::AllNonNull
13741542

13751543
pub type vortex_array::aggregate_fn::fns::all_non_null::AllNonNull::Options = vortex_array::aggregate_fn::EmptyOptions
@@ -20106,6 +20274,10 @@ pub fn vortex_array::scalar_fn::fns::stat::StatOptions::hash<__H: core::hash::Ha
2010620274

2010720275
impl core::marker::StructuralPartialEq for vortex_array::scalar_fn::fns::stat::StatOptions
2010820276

20277+
pub fn vortex_array::stats::expr::all_nan(vortex_array::expr::Expression) -> vortex_array::expr::Expression
20278+
20279+
pub fn vortex_array::stats::expr::all_non_nan(vortex_array::expr::Expression) -> vortex_array::expr::Expression
20280+
2010920281
pub fn vortex_array::stats::expr::all_non_null(vortex_array::expr::Expression) -> vortex_array::expr::Expression
2011020282

2011120283
pub fn vortex_array::stats::expr::all_null(vortex_array::expr::Expression) -> vortex_array::expr::Expression
@@ -20384,6 +20556,10 @@ pub fn vortex_array::stats::TypedStatsSetRef<'_, '_>::len(&self) -> usize
2038420556

2038520557
pub const vortex_array::stats::PRUNING_STATS: &[vortex_array::expr::stats::Stat]
2038620558

20559+
pub fn vortex_array::stats::all_nan(vortex_array::expr::Expression) -> vortex_array::expr::Expression
20560+
20561+
pub fn vortex_array::stats::all_non_nan(vortex_array::expr::Expression) -> vortex_array::expr::Expression
20562+
2038720563
pub fn vortex_array::stats::all_non_null(vortex_array::expr::Expression) -> vortex_array::expr::Expression
2038820564

2038920565
pub fn vortex_array::stats::all_null(vortex_array::expr::Expression) -> vortex_array::expr::Expression
Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexResult;
5+
6+
use crate::ArrayRef;
7+
use crate::Columnar;
8+
use crate::ExecutionCtx;
9+
use crate::IntoArray;
10+
use crate::aggregate_fn::AggregateFnId;
11+
use crate::aggregate_fn::AggregateFnVTable;
12+
use crate::aggregate_fn::EmptyOptions;
13+
use crate::aggregate_fn::fns::nan_count::nan_count;
14+
use crate::dtype::DType;
15+
use crate::dtype::Nullability;
16+
use crate::scalar::Scalar;
17+
18+
/// Compute whether every value in an array is NaN.
19+
#[derive(Clone, Debug)]
20+
pub struct AllNan;
21+
22+
impl AggregateFnVTable for AllNan {
23+
type Options = EmptyOptions;
24+
type Partial = bool;
25+
26+
fn id(&self) -> AggregateFnId {
27+
AggregateFnId::new("vortex.all_nan")
28+
}
29+
30+
fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
31+
Ok(None)
32+
}
33+
34+
fn return_dtype(&self, _options: &Self::Options, _input_dtype: &DType) -> Option<DType> {
35+
Some(DType::Bool(Nullability::NonNullable))
36+
}
37+
38+
fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
39+
self.return_dtype(options, input_dtype)
40+
}
41+
42+
fn empty_partial(
43+
&self,
44+
_options: &Self::Options,
45+
_input_dtype: &DType,
46+
) -> VortexResult<Self::Partial> {
47+
Ok(true)
48+
}
49+
50+
fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> {
51+
*partial &= bool::try_from(&other)?;
52+
Ok(())
53+
}
54+
55+
fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
56+
Ok(Scalar::bool(*partial, Nullability::NonNullable))
57+
}
58+
59+
fn reset(&self, partial: &mut Self::Partial) {
60+
*partial = true;
61+
}
62+
63+
fn is_saturated(&self, partial: &Self::Partial) -> bool {
64+
!*partial
65+
}
66+
67+
fn try_accumulate(
68+
&self,
69+
state: &mut Self::Partial,
70+
batch: &ArrayRef,
71+
ctx: &mut ExecutionCtx,
72+
) -> VortexResult<bool> {
73+
*state &= nan_count(batch, ctx)? == batch.len();
74+
Ok(true)
75+
}
76+
77+
fn accumulate(
78+
&self,
79+
partial: &mut Self::Partial,
80+
batch: &Columnar,
81+
ctx: &mut ExecutionCtx,
82+
) -> VortexResult<()> {
83+
let array = match batch {
84+
Columnar::Constant(c) => c.clone().into_array(),
85+
Columnar::Canonical(c) => c.clone().into_array(),
86+
};
87+
*partial &= nan_count(&array, ctx)? == array.len();
88+
Ok(())
89+
}
90+
91+
fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
92+
Ok(partials)
93+
}
94+
95+
fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
96+
self.to_scalar(partial)
97+
}
98+
}
99+
100+
#[cfg(test)]
101+
mod tests {
102+
use vortex_error::VortexResult;
103+
104+
use crate::IntoArray;
105+
use crate::LEGACY_SESSION;
106+
use crate::VortexSessionExecute;
107+
use crate::aggregate_fn::Accumulator;
108+
use crate::aggregate_fn::DynAccumulator;
109+
use crate::aggregate_fn::EmptyOptions;
110+
use crate::aggregate_fn::fns::all_nan::AllNan;
111+
use crate::arrays::PrimitiveArray;
112+
use crate::dtype::DType;
113+
use crate::dtype::Nullability;
114+
use crate::dtype::PType;
115+
116+
#[test]
117+
fn all_nan_aggregate_fn() -> VortexResult<()> {
118+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
119+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
120+
let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
121+
122+
let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN)]).into_array();
123+
acc.accumulate(&batch, &mut ctx)?;
124+
125+
assert!(bool::try_from(&acc.finish()?)?);
126+
Ok(())
127+
}
128+
129+
#[test]
130+
fn all_nan_false_with_non_nan() -> VortexResult<()> {
131+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
132+
let dtype = DType::Primitive(PType::F32, Nullability::Nullable);
133+
let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
134+
135+
let batch = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32)]).into_array();
136+
acc.accumulate(&batch, &mut ctx)?;
137+
138+
assert!(!bool::try_from(&acc.finish()?)?);
139+
Ok(())
140+
}
141+
142+
#[test]
143+
fn all_nan_false_for_non_float_values() -> VortexResult<()> {
144+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
145+
let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
146+
let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?;
147+
148+
let batch = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array();
149+
acc.accumulate(&batch, &mut ctx)?;
150+
151+
assert!(!bool::try_from(&acc.finish()?)?);
152+
Ok(())
153+
}
154+
}

0 commit comments

Comments
 (0)