Skip to content

Reduce constant mask for MaskedArray and Mask scalarfn #9038

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mhk197 wants to merge 4 commits into
base: develop
Choose a base branch
from
Open

Reduce constant mask for MaskedArray and Mask scalarfn #9038

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2529e20
all true constant mask reduces to nullability cast
mhk197 Jul 28, 2026
9761085
remove list enablement
mhk197 Jul 29, 2026
d2ca64d
fix
mhk197 Jul 29, 2026
a1bb95a
change test
mhk197 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions vortex-array/src/arrays/masked/vtable/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -193,13 +195,33 @@ impl VTable for Masked {
))
}

fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {
// 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);
};
Comment on lines +200 to +202

@joseph-isaacs joseph-isaacs Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there is array.validity() method from the array_slots macro that is generated to get the validity, would be nice to use it


let Some(constant) = mask_child.as_opt::<Constant>() 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),

@joseph-isaacs joseph-isaacs Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do none mean here? This is not allowed. This scalar must be non-nullable!

@mhk197 mhk197 Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
let validity = Validity::Array(
     ConstantArray::new(
          Scalar::null(DType::Bool(Nullability::Nullable)),
          child.len(),
      )
      .into_array(),
 );
let masked = MaskedArray::try_new(child, validity)?.into_array();

This succeeds currently

}
}

fn reduce_parent(
array: ArrayView<'_, Self>,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>> {
PARENT_RULES.evaluate(array, parent, child_idx)
}

fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
MaskedSlots::NAMES[idx].to_string()
}
Expand All @@ -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;
Expand Down Expand Up @@ -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::<Masked>());
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::<Constant>()
.is_some_and(|constant| { constant.scalar().is_null() })
);
Ok(())
}
}
34 changes: 20 additions & 14 deletions vortex-array/src/scalar_fn/fns/mask/kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,19 +125,20 @@ mod tests {
use crate::arrays::Primitive;
use crate::arrays::PrimitiveArray;
use crate::arrays::ScalarFn;
use crate::arrays::scalar_fn::ScalarFnArrayExt;
use crate::arrays::scalar_fn::ScalarFnFactoryExt;
use crate::assert_arrays_eq;
use crate::dtype::Nullability;
use crate::executor::VortexSessionExecute;
use crate::optimizer::ArrayOptimizer;
use crate::scalar::Scalar;
use crate::scalar_fn::EmptyOptions;
use crate::scalar_fn::fns::literal::Literal;
use crate::scalar_fn::fns::mask::Mask as MaskExpr;

/// A constant Boolean mask child must take the metadata-only reduction path (pushing into the
/// input encoding) rather than surviving as a `ScalarFn` wrapper that falls through to
/// execution. Asserting the optimized encoding makes this fail before the adaptor accepts
/// `Constant` masks, not just verifying values that could pass through the execution fallback.
/// A constant Boolean mask must take the metadata-only self-reduction path rather than survive
/// until execution. An all-true mask becomes the nullable input, while an all-false mask
/// becomes a typed null literal.
#[rstest]
#[case(true)]
#[case(false)]
Expand All @@ -156,16 +157,21 @@ mod tests {
);

let optimized = masked.optimize()?;
assert!(
!optimized.is::<ScalarFn>(),
"constant mask should not fall through to execution, got {}",
optimized.encoding_id()
);
assert!(
optimized.is::<Primitive>(),
"constant mask should reduce into the Primitive input, got {}",
optimized.encoding_id()
);
if mask_value {
assert!(
optimized.is::<Primitive>(),
"constant true mask should reduce into the Primitive input, got {}",
optimized.encoding_id()
);
} else {
assert!(
optimized
.as_opt::<ScalarFn>()
.is_some_and(|array| { array.scalar_fn().is::<Literal>() }),
"constant false mask should reduce to a null literal, got {}",
optimized.encoding_id()
);
}

let mut ctx = crate::array_session().create_execution_ctx();
let expected = if mask_value {
Expand Down
117 changes: 116 additions & 1 deletion vortex-array/src/scalar_fn/fns/mask/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<Option<ReduceNodeRef>> {
let mask = node.child(1);
// Expression literals are handled by `simplify`; this path recognizes array metadata.
let Some(mask) = mask.as_any().downcast_ref::<ArrayRef>() else {
return Ok(None);
};
let Some(constant) = mask.as_opt::<Constant>() 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),
}
}
Comment on lines +110 to +138

@joseph-isaacs joseph-isaacs Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how does this happen?

@mhk197 mhk197 Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if we apply .mask and the encoding has no MaskReduce i believe. Also this runs before parent reduce kernels


fn simplify(
&self,
_options: &Self::Options,
Expand Down Expand Up @@ -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() {
Expand All @@ -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::<ScalarFn>());

let optimized = masked.optimize()?;
assert!(
optimized.is::<Filter>(),
"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_::<ScalarFn>();

assert!(
optimized.scalar_fn().is::<Cast>(),
"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_::<ScalarFn>();
let scalar = optimized
.scalar_fn()
.as_opt::<Literal>()
.vortex_expect("expected null literal");

assert!(scalar.is_null());
assert_eq!(scalar.dtype(), &DType::Primitive(PType::I32, Nullable));

Ok(())
}
}
Loading