diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index fb8bfe227aa..f860c12c040 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -34,6 +34,7 @@ use crate::scalar_fn::fns::dynamic::Rhs; use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::is_nan::IsNan; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; @@ -855,6 +856,28 @@ pub fn bound_is_null(child: BoundExpression) -> BoundExpression { .vortex_expect("is-null expressions are always well-typed") } +// ---- IsNan ---- + +/// Creates an expression that checks for NaN values. +/// +/// The expression is strict: null inputs produce null outputs, so the output nullability +/// follows the input. Only primitive float inputs are supported. +/// +/// ```rust +/// # use vortex_array::expr::{is_nan, root}; +/// let expr = is_nan(root()); +/// ``` +pub fn is_nan(child: Expression) -> Expression { + IsNan.new_expr(EmptyOptions, vec![child]) +} + +/// Creates a bound expression that checks for NaN values. +pub fn bound_is_nan(child: BoundExpression) -> BoundExpression { + IsNan + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-nan expressions are always well-typed") +} + // ---- IsNotNull ---- /// Creates an expression that checks for non-null values. @@ -1242,6 +1265,7 @@ pub mod bound { pub use super::bound_gt as gt; pub use super::bound_gt_eq as gt_eq; pub use super::bound_ilike as ilike; + pub use super::bound_is_nan as is_nan; pub use super::bound_is_not_null as is_not_null; pub use super::bound_is_null as is_null; pub use super::bound_like as like; diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 59fd21c46ee..10a2fea3edf 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -92,6 +92,7 @@ pub use exprs::get_item; pub use exprs::gt; pub use exprs::gt_eq; pub use exprs::ilike; +pub use exprs::is_nan; pub use exprs::is_not_null; pub use exprs::is_null; pub use exprs::is_root; diff --git a/vortex-array/src/scalar_fn/fns/is_nan.rs b/vortex-array/src/scalar_fn/fns/is_nan.rs new file mode 100755 index 00000000000..9f345634abe --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/is_nan.rs @@ -0,0 +1,251 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use crate::dtype::DType; +use crate::match_each_float_ptype; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; + +/// Expression that checks for NaN values. +/// +/// The function is strict: null inputs produce null outputs. Only primitive float inputs +/// are supported. +#[derive(Clone)] +pub struct IsNan; + +impl RowFn for IsNan { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.is_nan"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = match args { + [DType::Primitive(ptype, _)] if ptype.is_float() => *ptype, + _ => vortex_bail!("is_nan expects a single float input, got {args:?}"), + }; + match_each_float_ptype!(ptype, |T| { + visitor.visit_bool::<(T,), false>(|(value,)| value.is_nan()) + }) + } +} + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexExpect as _; + + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::get_item; + use crate::expr::is_nan; + use crate::expr::lit; + use crate::expr::root; + use crate::scalar::Scalar; + use crate::validity::Validity; + + #[test] + fn dtype() { + assert_eq!( + is_nan(root()) + .return_dtype(&DType::Primitive(PType::F32, Nullability::Nullable)) + .unwrap(), + DType::Bool(Nullability::Nullable) + ); + assert_eq!( + is_nan(root()) + .return_dtype(&DType::Primitive(PType::F64, Nullability::NonNullable)) + .unwrap(), + DType::Bool(Nullability::NonNullable) + ); + } + + #[test] + fn dtype_rejects_non_float() { + assert!( + is_nan(root()) + .return_dtype(&DType::Primitive(PType::I32, Nullability::NonNullable)) + .is_err() + ); + assert!( + is_nan(root()) + .return_dtype(&DType::Bool(Nullability::NonNullable)) + .is_err() + ); + } + + #[test] + fn replace_children() { + let expr = is_nan(root()); + expr.with_children([root()]) + .vortex_expect("operation should succeed in test"); + } + + #[test] + fn evaluate_floats() { + let test_array = PrimitiveArray::from_option_iter([ + Some(1.0f32), + Some(f32::NAN), + None, + Some(f32::NEG_INFINITY), + Some(-f32::NAN), + ]) + .into_array(); + let expected = [Some(false), Some(true), None, Some(false), Some(true)]; + + let result = test_array.clone().apply(&is_nan(root())).unwrap(); + + assert_eq!(result.len(), test_array.len()); + assert_eq!(result.dtype(), &DType::Bool(Nullability::Nullable)); + for (i, expected_value) in expected.iter().enumerate() { + let expected_scalar = match expected_value { + Some(value) => Scalar::bool(*value, Nullability::Nullable), + None => Scalar::null(DType::Bool(Nullability::Nullable)), + }; + assert_eq!( + result + .execute_scalar(i, &mut array_session().create_execution_ctx()) + .unwrap(), + expected_scalar + ); + } + } + + #[test] + fn evaluate_all_valid_floats() { + let test_array = PrimitiveArray::new( + buffer![1.0f64, f64::NAN, f64::INFINITY, -f64::NAN], + Validity::NonNullable, + ) + .into_array(); + let expected = [false, true, false, true]; + + let result = test_array.clone().apply(&is_nan(root())).unwrap(); + + assert_eq!(result.len(), test_array.len()); + assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable)); + for (i, expected_value) in expected.iter().enumerate() { + assert_eq!( + result + .execute_scalar(i, &mut array_session().create_execution_ctx()) + .unwrap(), + Scalar::bool(*expected_value, Nullability::NonNullable) + ); + } + } + + #[test] + fn evaluate_all_null_floats() { + let test_array = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); + + let result = test_array.clone().apply(&is_nan(root())).unwrap(); + + assert_eq!(result.len(), test_array.len()); + for i in 0..result.len() { + assert_eq!( + result + .execute_scalar(i, &mut array_session().create_execution_ctx()) + .unwrap(), + Scalar::null(DType::Bool(Nullability::Nullable)) + ); + } + } + + #[test] + fn evaluate_constant() { + let test_array = buffer![1.0f32, 2.0, 3.0].into_array(); + let cases = [ + (lit(f32::NAN), Some(true)), + (lit(1.0f32), Some(false)), + ( + lit(Scalar::null(DType::Primitive( + PType::F32, + Nullability::Nullable, + ))), + None, + ), + ]; + + for (expr_child, expected_value) in cases { + let result = test_array.clone().apply(&is_nan(expr_child)).unwrap(); + for i in 0..result.len() { + let expected_scalar = match expected_value { + Some(value) => Scalar::bool(value, Nullability::Nullable), + None => Scalar::null(DType::Bool(Nullability::Nullable)), + }; + assert_eq!( + result + .execute_scalar(i, &mut array_session().create_execution_ctx()) + .unwrap(), + expected_scalar + ); + } + } + } + + #[test] + fn evaluate_rejects_non_float() { + let test_array = buffer![1i32, 2, 3].into_array(); + assert!(test_array.apply(&is_nan(root())).is_err()); + } + + #[test] + fn evaluate_sliced() { + let test_array = buffer![1.0f32, f32::NAN, 2.0, f32::NAN, 3.0] + .into_array() + .slice(1..4) + .unwrap(); + let expected = [true, false, true]; + + let result = test_array.clone().apply(&is_nan(root())).unwrap(); + + assert_eq!(result.len(), test_array.len()); + for (i, expected_value) in expected.iter().enumerate() { + assert_eq!( + result + .execute_scalar(i, &mut array_session().create_execution_ctx()) + .unwrap(), + Scalar::bool(*expected_value, Nullability::NonNullable) + ); + } + } + + #[test] + fn test_display() { + let expr = is_nan(get_item("name", root())); + assert_eq!(expr.to_string(), "vortex.is_nan($.name)"); + + let expr2 = is_nan(root()); + assert_eq!(expr2.to_string(), "vortex.is_nan($)"); + } + + #[test] + fn test_is_nan_is_strict() { + assert!( + is_nan(root()) + .as_scalar() + .is_some_and(|f| f.signature().is_strict()) + ); + } +} diff --git a/vortex-array/src/scalar_fn/fns/mod.rs b/vortex-array/src/scalar_fn/fns/mod.rs index 087f540ce7b..ece71396faf 100644 --- a/vortex-array/src/scalar_fn/fns/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mod.rs @@ -10,6 +10,7 @@ pub mod dynamic; pub mod ext_storage; pub mod fill_null; pub mod get_item; +pub mod is_nan; pub mod is_not_null; pub mod is_null; pub mod like; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 211227858ca..4e85b4d078c 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -19,6 +19,7 @@ use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::is_nan::IsNan; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; @@ -70,6 +71,7 @@ impl Default for ScalarFnSession { this.register(ExtStorage); this.register(FillNull); this.register(GetItem); + this.register(IsNan); this.register(IsNotNull); this.register(IsNull); this.register(Like); diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 3cb5fdb06df..ed03ed573b0 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -9,6 +9,7 @@ use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; +use crate::aggregate_fn::fns::all_nan::AllNan; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; @@ -38,6 +39,7 @@ use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; +use crate::scalar_fn::fns::is_nan::IsNan; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::like::Like; @@ -63,6 +65,9 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(IsNotNullNullCountStatsRewrite); session.register_rewrite(IsNotNullAllNullStatsRewrite); session.register_rewrite(IsNotNullAllNonNullStatsRewrite); + session.register_rewrite(IsNanNaNCountStatsRewrite); + session.register_rewrite(IsNanAllNonNanStatsRewrite); + session.register_rewrite(IsNanAllNanStatsRewrite); session.register_rewrite(LikeStatsRewrite); session.register_rewrite(ListContainsNanCountStatsRewrite); session.register_rewrite(ListContainsAllNonNanStatsRewrite); @@ -327,6 +332,71 @@ impl StatsRewriteRule for IsNotNullAllNonNullStatsRewrite { } } +/// Rewrites `is_nan` using the `NaNCount` stat: the predicate is falsified when a zone +/// contains no NaN values, and satisfied when every row in the zone is NaN. +#[derive(Debug)] +struct IsNanNaNCountStatsRewrite; + +impl StatsRewriteRule for IsNanNaNCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(nan_count(expr.child(0), ctx).map(|nan_count| eq(nan_count, lit(0u64)))) + } + + fn satisfy( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(nan_count(expr.child(0), ctx).map(|nan_count| eq(nan_count, row_count()))) + } +} + +/// Falsifies `is_nan` when the `AllNonNan` pruning stat proves that no value in the zone +/// is NaN. +#[derive(Debug)] +struct IsNanAllNonNanStatsRewrite; + +impl StatsRewriteRule for IsNanAllNonNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + _ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(Some(all_non_nan(expr.child(0)))) + } +} + +/// Satisfies `is_nan` when the `AllNan` pruning stat proves that every value in the zone +/// is NaN. +#[derive(Debug)] +struct IsNanAllNanStatsRewrite; + +impl StatsRewriteRule for IsNanAllNanStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNan.id() + } + + fn satisfy( + &self, + expr: &BoundExpression, + _ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(Some(all_nan(expr.child(0)))) + } +} + #[derive(Debug)] struct LikeStatsRewrite; @@ -530,6 +600,10 @@ fn null_count(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option) -> Option { + stat_expr(expr, Stat::NaNCount, ctx) +} + fn all_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) } @@ -538,6 +612,14 @@ fn all_non_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } +fn all_nan(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNan.bind(AggregateEmptyOptions)) +} + +fn all_non_nan(expr: &BoundExpression) -> BoundExpression { + stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)) +} + enum NanCheck { NotNeeded, Check(BoundExpression), @@ -736,8 +818,10 @@ mod tests { use crate::expr::col; use crate::expr::dynamic; use crate::expr::eq; + use crate::expr::get_item; use crate::expr::gt; use crate::expr::gt_eq; + use crate::expr::is_nan; use crate::expr::is_not_null; use crate::expr::is_null; use crate::expr::like; @@ -817,6 +901,14 @@ mod tests { crate::stats::all_non_null(expr.clone()) } + fn all_nan(expr: &Expression) -> Expression { + crate::stats::all_nan(expr.clone()) + } + + fn all_non_nan(expr: &Expression) -> Expression { + crate::stats::all_non_nan(expr.clone()) + } + macro_rules! assert_rewrite_eq { ($actual:expr, $expected:expr) => { assert_eq!($actual, bind_expected($expected)?) @@ -992,6 +1084,56 @@ mod tests { Ok(()) } + #[test] + fn rewrites_nan_falsifiers() -> VortexResult<()> { + assert_rewrite_eq!( + falsify(&is_nan(col("f")))?, + Some(or( + eq(stat(col("f"), Stat::NaNCount), lit(0u64)), + all_non_nan(&col("f")), + )) + ); + + // Nullable floats prune on nan_count just the same. + assert_rewrite_eq!( + falsify(&is_nan(get_item("x", col("n"))))?, + Some(or( + eq(stat(get_item("x", col("n")), Stat::NaNCount), lit(0u64)), + all_non_nan(&get_item("x", col("n"))), + )) + ); + + // NaN literals fold through the nan_count stat. + assert_rewrite_eq!( + falsify(&is_nan(lit(f32::NAN)))?, + Some(or(eq(lit(1u64), lit(0u64)), all_non_nan(&lit(f32::NAN)))) + ); + Ok(()) + } + + #[test] + fn rewrites_nan_satisfiers() -> VortexResult<()> { + assert_rewrite_eq!( + satisfy(&is_nan(col("f")))?, + Some(or( + eq( + stat(col("f"), Stat::NaNCount), + RowCount.new_expr(EmptyOptions, []) + ), + all_nan(&col("f")), + )) + ); + + assert_rewrite_eq!( + satisfy(&is_nan(lit(f32::NAN)))?, + Some(or( + eq(lit(1u64), RowCount.new_expr(EmptyOptions, [])), + all_nan(&lit(f32::NAN)), + )) + ); + Ok(()) + } + #[test] fn rewrites_list_contains_falsifier() -> VortexResult<()> { let list = Scalar::list(