diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index c7e32a7ee3c..46c3e88689a 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -30,6 +30,7 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::validity_to_child; use crate::array::with_empty_buffers; +use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::masked::MaskedArrayExt; use crate::arrays::masked::MaskedArraySlotsExt; @@ -38,6 +39,7 @@ use crate::arrays::masked::array::MaskedSlots; use crate::arrays::masked::compute::rules::PARENT_RULES; use crate::arrays::masked::mask_validity_canonical; use crate::buffer::BufferHandle; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; @@ -193,6 +195,25 @@ impl VTable for Masked { )) } + fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { + // No validity child represents AllValid. + let Some(mask_child) = array.slots()[MaskedSlots::VALIDITY].as_ref() else { + return array.child().cast(array.dtype().clone()).map(Some); + }; + + let Some(constant) = mask_child.as_opt::() else { + return Ok(None); + }; + + match constant.scalar().as_bool().value() { + Some(true) => array.child().cast(array.dtype().clone()).map(Some), + Some(false) => Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(), + )), + None => Ok(None), + } + } + fn reduce_parent( array: ArrayView<'_, Self>, parent: &ArrayRef, @@ -200,6 +221,7 @@ impl VTable for Masked { ) -> VortexResult> { PARENT_RULES.evaluate(array, parent, child_idx) } + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { MaskedSlots::NAMES[idx].to_string() } @@ -217,10 +239,13 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::Constant; + use crate::arrays::ConstantArray; use crate::arrays::Masked; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; + use crate::optimizer::ArrayOptimizer; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; use crate::validity::Validity; @@ -307,4 +332,37 @@ mod tests { Ok(()) } + + #[rstest] + #[case::without_validity_child(Validity::AllValid)] + #[case::with_constant_true_child(Validity::Array( + ConstantArray::new(true, 3).into_array() + ))] + fn reduce_all_valid(#[case] validity: Validity) -> Result<(), VortexError> { + let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let masked = MaskedArray::try_new(child, validity)?.into_array(); + + let optimized = masked.optimize()?; + + assert!(!optimized.is::()); + assert_eq!(optimized.dtype().nullability(), Nullability::Nullable); + Ok(()) + } + + #[test] + fn reduce_all_invalid_preserves_dtype() -> Result<(), VortexError> { + let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let masked = MaskedArray::try_new(child, Validity::AllInvalid)?.into_array(); + let dtype = masked.dtype().clone(); + + let optimized = masked.optimize()?; + + assert_eq!(optimized.dtype(), &dtype); + assert!( + optimized + .as_opt::() + .is_some_and(|constant| { constant.scalar().is_null() }) + ); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/fns/mask/mod.rs index 4c081e99c30..61086f70af3 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mask/mod.rs @@ -29,9 +29,14 @@ use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ReduceCtx; +use crate::scalar_fn::ReduceNode; +use crate::scalar_fn::ReduceNodeRef; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::SimplifyCtx; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::literal::Literal; /// An expression that masks an input based on a boolean mask. @@ -98,6 +103,40 @@ impl ScalarFnVTable for Mask { execute_canonical(input, mask_array, ctx) } + fn reduce( + &self, + _options: &Self::Options, + node: &dyn ReduceNode, + ctx: &dyn ReduceCtx, + ) -> VortexResult> { + let mask = node.child(1); + // Expression literals are handled by `simplify`; this path recognizes array metadata. + let Some(mask) = mask.as_any().downcast_ref::() else { + return Ok(None); + }; + let Some(constant) = mask.as_opt::() else { + return Ok(None); + }; + + match constant.scalar().as_bool().value() { + Some(true) => { + let input = node.child(0); + let output_dtype = node.node_dtype()?; + if input.node_dtype()? == output_dtype { + return Ok(Some(input)); + } + + ctx.new_node(Cast.bind(output_dtype), &[input]).map(Some) + } + Some(false) => { + let output_dtype = node.node_dtype()?; + ctx.new_node(Literal.bind(Scalar::null(output_dtype)), &[]) + .map(Some) + } + None => Ok(None), + } + } + fn simplify( &self, _options: &Self::Options, @@ -179,15 +218,30 @@ fn execute_canonical( } #[cfg(test)] -mod test { +mod tests { use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_mask::Mask as VortexMask; + use crate::IntoArray; + use crate::arrays::ConstantArray; + use crate::arrays::Filter; + use crate::arrays::FilterArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::ScalarFn; + use crate::arrays::scalar_fn::ScalarFnArrayExt; + use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::dtype::PType; use crate::expr::lit; use crate::expr::mask; + use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::fns::cast::Cast; + use crate::scalar_fn::fns::literal::Literal; + use crate::scalar_fn::fns::mask::Mask; #[test] fn test_simplify() { @@ -208,4 +262,65 @@ mod test { let expected_null_expr = lit(Scalar::null(DType::Primitive(PType::U32, Nullable))); assert_eq!(&simplified_false, &expected_null_expr); } + + #[test] + fn constant_true_mask_self_reduces_without_input_mask_rule() -> VortexResult<()> { + let input = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); + let input = + FilterArray::new(input, VortexMask::from_iter([true, false, true])).into_array(); + let mask = ConstantArray::new(true, input.len()).into_array(); + + let masked = Mask.try_new_array(input.len(), EmptyOptions, [input, mask])?; + assert!(masked.is::()); + + let optimized = masked.optimize()?; + assert!( + optimized.is::(), + "expected the mask to self-reduce to its already-nullable Filter input, got {}", + optimized.encoding_id() + ); + assert_eq!(optimized.dtype(), &DType::Primitive(PType::I32, Nullable)); + + Ok(()) + } + + #[test] + fn constant_true_mask_self_reduces_to_nullable_cast() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let input = + FilterArray::new(input, VortexMask::from_iter([true, false, true])).into_array(); + let mask = ConstantArray::new(true, input.len()).into_array(); + + let masked = Mask.try_new_array(input.len(), EmptyOptions, [input, mask])?; + let optimized = masked.optimize()?; + let optimized = optimized.as_::(); + + assert!( + optimized.scalar_fn().is::(), + "expected the mask to self-reduce to a nullable cast, got {}", + optimized.scalar_fn().id() + ); + assert_eq!(optimized.dtype(), &DType::Primitive(PType::I32, Nullable)); + + Ok(()) + } + + #[test] + fn constant_false_mask_self_reduces_to_null_literal() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let mask = ConstantArray::new(false, input.len()).into_array(); + + let masked = Mask.try_new_array(input.len(), EmptyOptions, [input, mask])?; + let optimized = masked.optimize()?; + let optimized = optimized.as_::(); + let scalar = optimized + .scalar_fn() + .as_opt::() + .vortex_expect("expected null literal"); + + assert!(scalar.is_null()); + assert_eq!(scalar.dtype(), &DType::Primitive(PType::I32, Nullable)); + + Ok(()) + } }