diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index 862437ebcdc..9fa5cecc6ba 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -13,7 +13,6 @@ mod varbin; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_mask::Mask; use self::bool::check_bool_constant; use self::decimal::check_decimal_constant; @@ -73,7 +72,7 @@ fn arrays_value_equal(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vor // Compare values element-wise. Result is null where both inputs are null, // true/false where both are valid. let eq_result = a.binary(b.clone(), Operator::Eq)?; - let eq_result = eq_result.fill_null(false)?.execute::(ctx)?; + let eq_result = eq_result.null_as_false().execute(ctx)?; Ok(eq_result.true_count() == valid_count) } diff --git a/vortex-array/src/mask.rs b/vortex-array/src/mask.rs index 09032aabbd0..d8ba622f347 100644 --- a/vortex-array/src/mask.rs +++ b/vortex-array/src/mask.rs @@ -13,6 +13,7 @@ use crate::arrays::BoolArray; use crate::columnar::Columnar; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::validity::Validity; impl Executable for Mask { /// Executes a boolean array into a [`Mask`]. @@ -41,6 +42,60 @@ impl Executable for Mask { } } +/// An adapter that coerces null elements of a boolean array to `false` before executing it into a +/// [`Mask`]. Created by [`ArrayRef::null_as_false`]. +/// +/// Use for filter and pruning predicates over nullable data, where SQL semantics treat `NULL` as +/// not matching. +/// +/// Prefer `array.null_as_false().execute(ctx)` over `array.fill_null(false)?.execute::(ctx)`: +/// `fill_null` on a lazy `ScalarFn` array (e.g. the result of `apply()`) is currently +/// slow because its `validity()` executes the predicate expression. +pub struct NullAsFalse(ArrayRef); + +impl ArrayRef { + /// Returns an adapter that treats null elements of this boolean array as `false` when executed + /// into a [`Mask`]. See [`NullAsFalse`]. + pub fn null_as_false(self) -> NullAsFalse { + NullAsFalse(self) + } +} + +impl NullAsFalse { + /// Executes the boolean array into a [`Mask`], coercing null elements to `false`. + /// + /// Canonicalizes the (possibly lazy) array exactly once and folds validity into the value bits + /// with a single `AND` that reuses the value buffer when it is uniquely owned. + pub fn execute(self, ctx: &mut ExecutionCtx) -> VortexResult { + let array = self.0; + if !matches!(array.dtype(), DType::Bool(_)) { + vortex_bail!("Mask array must have boolean dtype, not {}", array.dtype()); + } + // Non-nullable input needs no coercion; defer to the strict `Mask` execution. + if !array.dtype().is_nullable() { + return array.execute::(ctx); + } + + let len = array.len(); + Ok(match array.execute::(ctx)? { + Columnar::Constant(c) => Mask::new(len, c.scalar().as_bool().value().unwrap_or(false)), + Columnar::Canonical(c) => { + let bool = c.into_array().execute::(ctx)?; + match bool.as_ref().validity()? { + Validity::NonNullable | Validity::AllValid => { + Mask::from_buffer(bool.into_bit_buffer()) + } + Validity::AllInvalid => Mask::new_false(len), + Validity::Array(v) => { + let validity_bits = v.execute::(ctx)?.into_bit_buffer(); + Mask::from_buffer(bool.into_bit_buffer() & &validity_bits) + } + } + } + }) + } +} + #[cfg(test)] mod tests { use vortex_error::VortexResult; @@ -91,4 +146,39 @@ mod tests { assert_eq!(mask, Mask::new_false(4)); Ok(()) } + + #[test] + fn null_as_false_non_nullable() -> VortexResult<()> { + let array = BoolArray::from_iter([true, false, true]).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, true])); + Ok(()) + } + + #[test] + fn null_as_false_treats_null_as_false() -> VortexResult<()> { + let array = BoolArray::from_iter([Some(true), None, Some(false), None]).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::from_iter([true, false, false, false])); + Ok(()) + } + + #[test] + fn null_as_false_null_constant() -> VortexResult<()> { + let array = + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); + let mask = array.null_as_false().execute(&mut ctx())?; + assert_eq!(mask, Mask::new_false(4)); + Ok(()) + } + + #[test] + fn null_as_false_matches_fill_null_then_mask() -> VortexResult<()> { + let array = + BoolArray::from_iter([Some(true), None, Some(false), Some(true), None]).into_array(); + let via_fill_null = array.fill_null(false)?.execute::(&mut ctx())?; + let via_coerce = array.null_as_false().execute(&mut ctx())?; + assert_eq!(via_coerce, via_fill_null); + Ok(()) + } } diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 9e6111fc371..ac16d858d2d 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -22,7 +22,6 @@ use vortex::array::MaskFuture; use vortex::array::ProstMetadata; use vortex::array::VortexSessionExecute; use vortex::array::arrays::Constant; -use vortex::array::builtins::ArrayBuiltins; use vortex::array::expr::Expression; use vortex::array::expr::stats::Precision; use vortex::array::expr::stats::Stat; @@ -332,12 +331,12 @@ impl LayoutReader for CudaFlatReader { let array = array.apply(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { let array = array.apply(&expr)?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.bitand(&array_mask) }; diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 26749e42f46..685444a61e5 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -16,7 +16,6 @@ use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DictArray; use vortex_array::arrays::SharedArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; @@ -259,10 +258,7 @@ impl LayoutReader for DictReader { let mask = mask.await?; let mut ctx = session.create_execution_ctx(); - let dict_mask = values - .take(codes)? - .fill_null(false)? - .execute::(&mut ctx)?; + let dict_mask = values.take(codes)?.null_as_false().execute(&mut ctx)?; Ok(mask.bitand(&dict_mask)) })) diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index e03ffc20e1a..6881d3c814a 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -11,7 +11,6 @@ use tracing::trace; use vortex_array::ArrayRef; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::Expression; @@ -157,14 +156,14 @@ impl LayoutReader for FlatReader { let array = array.apply(&expr)?; let array = array.filter(mask.clone())?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.intersect_by_rank(&array_mask) } else { // Run over the full array, with a simpler bitand at the end. let array = array.apply(&expr)?; let mut ctx = session.create_execution_ctx(); - let array_mask = array.fill_null(false)?.execute::(&mut ctx)?; + let array_mask = array.null_as_false().execute(&mut ctx)?; mask.bitand(&array_mask) }; diff --git a/vortex-layout/src/layouts/partitioned.rs b/vortex-layout/src/layouts/partitioned.rs index cd927bbc074..76327a7ee5f 100644 --- a/vortex-layout/src/layouts/partitioned.rs +++ b/vortex-layout/src/layouts/partitioned.rs @@ -11,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; @@ -19,7 +18,6 @@ use vortex_array::expr::transform::PartitionedExpr; use vortex_array::validity::Validity; use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_mask::Mask; use vortex_session::VortexSession; use crate::ArrayFuture; @@ -93,8 +91,8 @@ impl PartitionedExprEval

for PartitionedExpr

{ let mut ctx = session.create_execution_ctx(); let root_mask = root_scope .apply(&self.root)? - .fill_null(false)? - .execute::(&mut ctx)?; + .null_as_false() + .execute(&mut ctx)?; let mask = mask.bitand(&root_mask); diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 2347d05b0bd..fb826a1b401 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -19,7 +19,6 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; @@ -296,10 +295,7 @@ fn row_idx_mask_future( let array = idx_array(row_offset, &row_range).into_array(); let mut ctx = session.create_execution_ctx(); - let result_mask = array - .apply(&expr)? - .fill_null(false)? - .execute::(&mut ctx)?; + let result_mask = array.apply(&expr)?.null_as_false().execute(&mut ctx)?; Ok(result_mask.bitand(&mask.await?)) }) diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index d7a5b7cd769..789a6a74e8d 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -17,7 +17,6 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; @@ -124,12 +123,12 @@ impl ZoneMap { let applied = self.array.clone().into_array().apply(&predicate)?; if !contains_row_count(&applied) { - return applied.fill_null(false)?.execute::(&mut ctx); + return applied.null_as_false().execute(&mut ctx); } let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?; let substituted = substitute_row_count(applied, &row_count_array)?; - substituted.fill_null(false)?.execute::(&mut ctx) + substituted.null_as_false().execute(&mut ctx) } fn lower_stats(&self, predicate: Expression) -> VortexResult {