From aec7150540113f92a591f2b953c4879d6c66b075 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 17 Jun 2026 11:42:20 -0400 Subject: [PATCH 1/3] Inline Kleene boolean kernels Signed-off-by: Nicholas Gates --- encodings/bytebool/src/compute.rs | 82 ++ encodings/bytebool/src/kernel.rs | 2 + .../src/arrays/bool/compute/boolean.rs | 58 ++ vortex-array/src/arrays/bool/compute/mod.rs | 1 + vortex-array/src/arrays/bool/vtable/kernel.rs | 2 + .../src/scalar_fn/fns/binary/boolean.rs | 717 +++++++++++++++++- vortex-array/src/scalar_fn/fns/binary/mod.rs | 8 +- 7 files changed, 843 insertions(+), 27 deletions(-) create mode 100644 vortex-array/src/arrays/bool/compute/boolean.rs diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index e81b59dc9a4..89662f44e6e 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -11,12 +11,18 @@ use vortex_array::arrays::dict::TakeExecute; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::binary::BooleanKernel; +use vortex_array::scalar_fn::fns::binary::boolean_buffer_scalar; +use vortex_array::scalar_fn::fns::binary::boolean_buffers; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_array::scalar_fn::fns::mask::MaskReduce; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::ByteBool; @@ -106,9 +112,59 @@ impl TakeExecute for ByteBool { } } +impl BooleanKernel for ByteBool { + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + let lhs_values = truthy_bit_buffer(lhs); + + if let Some(rhs) = rhs.as_opt::() { + let rhs = rhs + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))?; + return boolean_buffer_scalar( + lhs_values, + lhs.validity()?, + &rhs, + operator, + nullability, + ctx, + ) + .map(Some); + } + + let Some(rhs) = rhs.as_opt::() else { + return Ok(None); + }; + + boolean_buffers( + lhs_values, + lhs.validity()?, + truthy_bit_buffer(rhs), + rhs.validity()?, + operator, + nullability, + ctx, + ) + .map(Some) + } +} + +fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer { + BitBuffer::from_iter(array.truthy_bytes().iter().map(|byte| *byte != 0)) +} + #[cfg(test)] mod tests { use rstest::rstest; + use vortex_array::LEGACY_SESSION; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::cast::test_cast_conformance; @@ -119,6 +175,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::vortex_err; use super::*; use crate::ByteBoolArray; @@ -184,6 +241,31 @@ mod tests { assert_arrays_eq!(arr, expected.into_array()); } + #[test] + fn test_boolean_kernel_kleene() -> VortexResult<()> { + let lhs = bb_opt(vec![Some(false), Some(true), None, Some(false), None]); + let rhs = bb_opt(vec![None, None, Some(true), Some(false), None]).into_array(); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + let and_result = + ::boolean(lhs.as_view(), &rhs, Operator::And, &mut ctx)? + .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean AND"))?; + assert_arrays_eq!( + and_result, + BoolArray::from_iter([Some(false), None, None, Some(false), None]) + ); + + let or_result = + ::boolean(lhs.as_view(), &rhs, Operator::Or, &mut ctx)? + .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean OR"))?; + assert_arrays_eq!( + or_result, + BoolArray::from_iter([None, Some(true), Some(true), Some(false), None]) + ); + + Ok(()) + } + #[test] fn test_mask_byte_bool() { test_mask_conformance(&bb(vec![true, false, true, true, false]).into_array()); diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index ddf9679b3d5..7801316445e 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -3,11 +3,13 @@ use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::kernel::ParentKernelSet; +use vortex_array::scalar_fn::fns::binary::BooleanExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::ByteBool; pub(crate) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ + ParentKernelSet::lift(&BooleanExecuteAdaptor(ByteBool)), ParentKernelSet::lift(&CastExecuteAdaptor(ByteBool)), ParentKernelSet::lift(&TakeExecuteAdaptor(ByteBool)), ]); diff --git a/vortex-array/src/arrays/bool/compute/boolean.rs b/vortex-array/src/arrays/bool/compute/boolean.rs new file mode 100644 index 00000000000..438833d84bb --- /dev/null +++ b/vortex-array/src/arrays/bool/compute/boolean.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::array::ArrayView; +use crate::arrays::Bool; +use crate::arrays::Constant; +use crate::arrays::bool::BoolArrayExt; +use crate::scalar_fn::fns::binary::BooleanKernel; +use crate::scalar_fn::fns::binary::boolean_buffer_scalar; +use crate::scalar_fn::fns::binary::boolean_buffers; +use crate::scalar_fn::fns::operators::Operator; + +impl BooleanKernel for Bool { + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + + if let Some(rhs) = rhs.as_opt::() { + let rhs = rhs + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))?; + return boolean_buffer_scalar( + lhs.to_bit_buffer(), + lhs.validity()?, + &rhs, + operator, + nullability, + ctx, + ) + .map(Some); + } + + let Some(rhs) = rhs.as_opt::() else { + return Ok(None); + }; + + boolean_buffers( + lhs.to_bit_buffer(), + lhs.validity()?, + rhs.to_bit_buffer(), + rhs.validity()?, + operator, + nullability, + ctx, + ) + .map(Some) + } +} diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs index 9b71578cd84..f15817ffc57 100644 --- a/vortex-array/src/arrays/bool/compute/mod.rs +++ b/vortex-array/src/arrays/bool/compute/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod boolean; mod cast; mod fill_null; pub(crate) mod filter; diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index 2fe2d974a8f..cbcf1727ebc 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -4,11 +4,13 @@ use crate::arrays::Bool; use crate::arrays::dict::TakeExecuteAdaptor; use crate::kernel::ParentKernelSet; +use crate::scalar_fn::fns::binary::BooleanExecuteAdaptor; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; pub(super) const PARENT_KERNELS: ParentKernelSet = ParentKernelSet::new(&[ + ParentKernelSet::lift(&BooleanExecuteAdaptor(Bool)), ParentKernelSet::lift(&CastExecuteAdaptor(Bool)), ParentKernelSet::lift(&FillNullExecuteAdaptor(Bool)), ParentKernelSet::lift(&TakeExecuteAdaptor(Bool)), diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 218aba806c9..974ba2d1594 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -1,21 +1,98 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::iter::repeat_n; + use arrow_array::cast::AsArray; +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::Bool; +use crate::arrays::BoolArray; use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::ScalarFn; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::arrow::ArrowSessionExt; use crate::arrow::FromArrowArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::executor::ExecutionCtx; +use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; +use crate::scalar::BoolScalar; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +/// Trait for encoding-specific Kleene boolean kernels that operate in encoded space. +/// +/// Implementations receive the encoded array as the left operand. `rhs` may be any boolean array +/// encoding or a constant; implementations should return `Ok(None)` when they cannot handle that +/// operand without falling back to ordinary execution. +pub trait BooleanKernel: VTable { + /// Execute `lhs rhs` using Kleene boolean semantics. + fn boolean( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +/// Adaptor that bridges [`BooleanKernel`] implementations to [`ExecuteParentKernel`]. +/// +/// When a `ScalarFnArray(Binary, And|Or)` wraps a child implementing [`BooleanKernel`], this +/// adaptor extracts the other operand and delegates to the encoding-specific kernel. +#[derive(Default, Debug)] +pub struct BooleanExecuteAdaptor(pub V); + +impl ExecuteParentKernel for BooleanExecuteAdaptor +where + V: BooleanKernel, +{ + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ScalarFnArrayView<'_, Binary>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let op = *parent.options; + if !is_boolean_operator(op) { + return Ok(None); + } + + let Some(scalar_fn_array) = parent.as_opt::() else { + return Ok(None); + }; + let other = match child_idx { + 0 => scalar_fn_array.get_child(1), + 1 => scalar_fn_array.get_child(0), + _ => return Ok(None), + }; + + if let Some(result) = constant_boolean(array.array(), other, op)? { + return Ok(Some(result)); + } + + V::boolean(array, other, op, ctx) + } +} /// Point-wise Kleene logical _and_ between two Boolean arrays. #[deprecated(note = "Use `ArrayBuiltins::binary` instead")] @@ -32,16 +109,35 @@ pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult { /// Execute a Kleene boolean operation between two arrays. /// /// This is the entry point for boolean operations from the binary expression. -/// Handles constant-constant directly, otherwise falls back to Arrow. +/// Handles constants and canonical boolean arrays directly, otherwise falls back to Arrow. pub(crate) fn execute_boolean( lhs: &ArrayRef, rhs: &ArrayRef, op: Operator, ctx: &mut ExecutionCtx, ) -> VortexResult { + let nullable = boolean_nullability(lhs, rhs); + + if lhs.is_empty() { + return Ok(Canonical::empty(&DType::Bool(nullable)).into_array()); + } + if let Some(result) = constant_boolean(lhs, rhs, op)? { return Ok(result); } + + if let Some(lhs) = lhs.as_opt::() + && let Some(result) = ::boolean(lhs, rhs, op, ctx)? + { + return Ok(result); + } + + if let Some(rhs) = rhs.as_opt::() + && let Some(result) = ::boolean(rhs, lhs, op, ctx)? + { + return Ok(result); + } + arrow_execute_boolean(lhs.clone(), rhs.clone(), op, ctx) } @@ -52,19 +148,21 @@ fn arrow_execute_boolean( op: Operator, ctx: &mut ExecutionCtx, ) -> VortexResult { - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + let nullable = boolean_nullability(&lhs, &rhs); let session = ctx.session().clone(); let lhs = session .arrow() .execute_arrow(lhs, None, ctx)? - .as_boolean() + .as_boolean_opt() + .ok_or_else(|| vortex_err!("expected lhs to be boolean"))? .clone(); let rhs = session .arrow() .execute_arrow(rhs, None, ctx)? - .as_boolean() + .as_boolean_opt() + .ok_or_else(|| vortex_err!("expected rhs to be boolean"))? .clone(); let array = match op { @@ -73,63 +171,632 @@ fn arrow_execute_boolean( other => return Err(vortex_err!("Not a boolean operator: {other}")), }; - ArrayRef::from_arrow(&array, nullable) + ArrayRef::from_arrow(&array, nullable == Nullability::Nullable) } -/// Constant-folds a boolean operation between two constant arrays. +/// Handles boolean operations where at least one operand is a constant array. fn constant_boolean( lhs: &ArrayRef, rhs: &ArrayRef, op: Operator, ) -> VortexResult> { - let (Some(lhs), Some(rhs)) = (lhs.as_opt::(), rhs.as_opt::()) else { - return Ok(None); - }; + let nullable = boolean_nullability(lhs, rhs); - let length = lhs.len(); - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); - let lhs_val = lhs.scalar().as_bool().value(); - let rhs_val = rhs - .scalar() - .as_bool_opt() - .ok_or_else(|| vortex_err!("expected rhs to be boolean"))? - .value(); + match (lhs.as_opt::(), rhs.as_opt::()) { + (Some(lhs), Some(rhs)) => { + let result = boolean_scalar_scalar( + bool_scalar_value(lhs.scalar())?, + bool_scalar_value(rhs.scalar())?, + op, + )?; + + Ok(Some(constant_bool_result(result, lhs.len(), nullable))) + } + (Some(lhs), None) => constant_array_boolean(lhs.scalar(), rhs, op, nullable), + (None, Some(rhs)) => constant_array_boolean(rhs.scalar(), lhs, op, nullable), + (None, None) => Ok(None), + } +} + +fn constant_array_boolean( + constant: &Scalar, + array: &ArrayRef, + op: Operator, + nullability: Nullability, +) -> VortexResult> { + match (op, bool_scalar_value(constant)?) { + (Operator::And, Some(false)) => Ok(Some(constant_bool_result( + Some(false), + array.len(), + nullability, + ))), + (Operator::And, Some(true)) => Ok(Some(cast_bool_nullability(array, nullability)?)), + (Operator::Or, Some(true)) => Ok(Some(constant_bool_result( + Some(true), + array.len(), + nullability, + ))), + (Operator::Or, Some(false)) => Ok(Some(cast_bool_nullability(array, nullability)?)), + (Operator::And | Operator::Or, None) => Ok(None), + (other, _) => Err(vortex_err!("Not a boolean operator: {other}")), + } +} - let result = match op { - Operator::And => match (lhs_val, rhs_val) { +fn boolean_scalar_scalar( + lhs: Option, + rhs: Option, + op: Operator, +) -> VortexResult> { + Ok(match op { + Operator::And => match (lhs, rhs) { (Some(false), _) | (_, Some(false)) => Some(false), (None, _) | (_, None) => None, (Some(l), Some(r)) => Some(l & r), }, - Operator::Or => match (lhs_val, rhs_val) { + Operator::Or => match (lhs, rhs) { (Some(true), _) | (_, Some(true)) => Some(true), (None, _) | (_, None) => None, (Some(l), Some(r)) => Some(l | r), }, other => return Err(vortex_err!("Not a boolean operator: {other}")), + }) +} + +fn bool_scalar_value(scalar: &Scalar) -> VortexResult> { + Ok(scalar + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))? + .value()) +} + +/// Execute a Kleene boolean operation from boolean value bitmaps and validity values. +pub fn boolean_buffers( + lhs_values: BitBuffer, + lhs_validity: Validity, + rhs_values: BitBuffer, + rhs_validity: Validity, + operator: Operator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs_values.len(); + debug_assert_eq!(rhs_values.len(), len); + + if lhs_validity.definitely_no_nulls() && rhs_validity.definitely_no_nulls() { + let values = match operator { + Operator::And => lhs_values & &rhs_values, + Operator::Or => lhs_values | &rhs_values, + other => return Err(vortex_err!("Not a boolean operator: {other}")), + }; + return Ok(BoolArray::try_new(values, all_validity(nullability))?.into_array()); + } + + let lhs_valid = lhs_validity.execute_mask(len, ctx)?; + let rhs_valid = rhs_validity.execute_mask(len, ctx)?; + fused_boolean_buffers( + len, + &lhs_values, + &lhs_valid, + &rhs_values, + &rhs_valid, + operator, + nullability, + ) +} + +/// Execute a Kleene boolean operation between boolean value bits and a scalar. +pub fn boolean_buffer_scalar( + values: BitBuffer, + validity: Validity, + scalar: &BoolScalar<'_>, + operator: Operator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_value = scalar.value(); + let len = values.len(); + let result = match (operator, scalar_value) { + (Operator::And, Some(false)) => { + return Ok(constant_bool_result(Some(false), len, nullability)); + } + (Operator::And, Some(true)) => { + return Ok( + BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(), + ); + } + (Operator::Or, Some(true)) => { + return Ok(constant_bool_result(Some(true), len, nullability)); + } + (Operator::Or, Some(false)) => { + return Ok( + BoolArray::try_new(values, validity.union_nullability(nullability))?.into_array(), + ); + } + (Operator::And, None) => { + let valid = validity + .execute_mask(len, ctx)? + .bitand_not(&Mask::from_buffer(values)); + BoolArray::try_new( + BitBuffer::new_unset(len), + mask_to_validity(valid, nullability), + )? + } + (Operator::Or, None) => { + let valid = validity.execute_mask(len, ctx)? & &Mask::from_buffer(values); + BoolArray::try_new( + BitBuffer::new_set(len), + mask_to_validity(valid, nullability), + )? + } + (other, _) => return Err(vortex_err!("Not a boolean operator: {other}")), }; - let scalar = result - .map(|b| Scalar::bool(b, nullable.into())) - .unwrap_or_else(|| Scalar::null(DType::Bool(nullable.into()))); + Ok(result.into_array()) +} + +fn fused_boolean_buffers( + len: usize, + lhs_values: &BitBuffer, + lhs_validity: &Mask, + rhs_values: &BitBuffer, + rhs_validity: &Mask, + operator: Operator, + nullability: Nullability, +) -> VortexResult { + if let Some(result) = fused_boolean_buffers_aligned( + len, + lhs_values, + lhs_validity, + rhs_values, + rhs_validity, + operator, + nullability, + )? { + return Ok(result); + } - Ok(Some(ConstantArray::new(scalar, length).into_array())) + let n_words = len.div_ceil(64); + + macro_rules! fuse { + ($lhs_valid_words:expr, $rhs_valid_words:expr) => { + fused_boolean_words( + len, + lhs_values.chunks().iter_padded(), + rhs_values.chunks().iter_padded(), + $lhs_valid_words, + $rhs_valid_words, + operator, + nullability, + ) + }; + } + + match (lhs_validity.bit_buffer(), rhs_validity.bit_buffer()) { + (AllOr::All, AllOr::All) => { + fuse!(repeat_n(u64::MAX, n_words), repeat_n(u64::MAX, n_words)) + } + (AllOr::All, AllOr::None) => { + fuse!(repeat_n(u64::MAX, n_words), repeat_n(0, n_words)) + } + (AllOr::All, AllOr::Some(rhs_validity)) => fuse!( + repeat_n(u64::MAX, n_words), + rhs_validity.chunks().iter_padded() + ), + (AllOr::None, AllOr::All) => { + fuse!(repeat_n(0, n_words), repeat_n(u64::MAX, n_words)) + } + (AllOr::None, AllOr::None) => { + fuse!(repeat_n(0, n_words), repeat_n(0, n_words)) + } + (AllOr::None, AllOr::Some(rhs_validity)) => { + fuse!(repeat_n(0, n_words), rhs_validity.chunks().iter_padded()) + } + (AllOr::Some(lhs_validity), AllOr::All) => fuse!( + lhs_validity.chunks().iter_padded(), + repeat_n(u64::MAX, n_words) + ), + (AllOr::Some(lhs_validity), AllOr::None) => { + fuse!(lhs_validity.chunks().iter_padded(), repeat_n(0, n_words)) + } + (AllOr::Some(lhs_validity), AllOr::Some(rhs_validity)) => fuse!( + lhs_validity.chunks().iter_padded(), + rhs_validity.chunks().iter_padded() + ), + } +} + +#[derive(Clone, Copy)] +enum WordSource<'a> { + Fill(u64), + Bytes(&'a [u8]), +} + +impl WordSource<'_> { + #[inline] + fn word_at(self, byte_offset: usize, len: usize) -> u64 { + match self { + Self::Fill(word) => word, + Self::Bytes(bytes) => read_u64_le(&bytes[byte_offset..byte_offset + len]), + } + } +} + +fn fused_boolean_buffers_aligned( + len: usize, + lhs_values: &BitBuffer, + lhs_validity: &Mask, + rhs_values: &BitBuffer, + rhs_validity: &Mask, + operator: Operator, + nullability: Nullability, +) -> VortexResult> { + let Some(lhs_values) = word_source_from_bit_buffer(lhs_values, len) else { + return Ok(None); + }; + let Some(rhs_values) = word_source_from_bit_buffer(rhs_values, len) else { + return Ok(None); + }; + let Some(lhs_validity) = word_source_from_mask(lhs_validity, len) else { + return Ok(None); + }; + let Some(rhs_validity) = word_source_from_mask(rhs_validity, len) else { + return Ok(None); + }; + + Ok(Some(fused_boolean_word_sources( + len, + lhs_values, + rhs_values, + lhs_validity, + rhs_validity, + operator, + nullability, + )?)) +} + +fn word_source_from_bit_buffer(buffer: &BitBuffer, len: usize) -> Option> { + if !buffer.offset().is_multiple_of(8) { + return None; + } + + let n_bytes = len.div_ceil(8); + let start = buffer.offset() / 8; + let end = start + n_bytes; + Some(WordSource::Bytes(&buffer.inner().as_slice()[start..end])) +} + +fn word_source_from_mask(mask: &Mask, len: usize) -> Option> { + match mask.bit_buffer() { + AllOr::All => Some(WordSource::Fill(u64::MAX)), + AllOr::None => Some(WordSource::Fill(0)), + AllOr::Some(buffer) => word_source_from_bit_buffer(buffer, len), + } +} + +fn read_u64_le(bytes: &[u8]) -> u64 { + debug_assert!(bytes.len() <= 8); + let mut buf = [0; 8]; + buf[..bytes.len()].copy_from_slice(bytes); + u64::from_le_bytes(buf) +} + +fn fused_boolean_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + operator: Operator, + nullability: Nullability, +) -> VortexResult { + let n_bytes = len.div_ceil(8); + let n_words = n_bytes.div_ceil(8); + let full_bytes = n_bytes - n_bytes % 8; + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + match operator { + Operator::And => { + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs & rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), + ); + } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs & rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), + ); + } + } + } + Operator::Or => { + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs | rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), + ); + } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs | rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), + ); + } + } + } + other => return Err(vortex_err!("Not a boolean operator: {other}")), + } + + finish_fused_boolean_words(len, n_bytes, values, validity, nullability) +} + +fn finish_fused_boolean_words( + len: usize, + n_bytes: usize, + values: BufferMut, + validity: BufferMut, + nullability: Nullability, +) -> VortexResult { + let mut values = values.into_byte_buffer(); + values.truncate(n_bytes); + let mut validity = validity.into_byte_buffer(); + validity.truncate(n_bytes); + Ok(BoolArray::try_new( + BitBuffer::new(values.freeze(), len), + mask_to_validity( + Mask::from_buffer(BitBuffer::new(validity.freeze(), len)), + nullability, + ), + )? + .into_array()) +} + +fn fused_boolean_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + operator: Operator, + nullability: Nullability, +) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + let n_words = len.div_ceil(64); + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + match operator { + Operator::And => { + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop is capped at `n_words`. + unsafe { + values.push_unchecked(lhs & rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), + ); + } + } + } + Operator::Or => { + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop is capped at `n_words`. + unsafe { + values.push_unchecked(lhs | rhs); + validity.push_unchecked( + (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), + ); + } + } + } + other => return Err(vortex_err!("Not a boolean operator: {other}")), + } + + let values = BitBuffer::new(values.freeze().into_byte_buffer(), len); + let validity = BitBuffer::new(validity.freeze().into_byte_buffer(), len); + + Ok(BoolArray::try_new( + values, + mask_to_validity(Mask::from_buffer(validity), nullability), + )? + .into_array()) +} + +fn constant_bool_result(value: Option, len: usize, nullability: Nullability) -> ArrayRef { + let scalar = value + .map(|b| Scalar::bool(b, nullability)) + .unwrap_or_else(|| Scalar::null(DType::Bool(nullability))); + + ConstantArray::new(scalar, len).into_array() +} + +fn cast_bool_nullability(array: &ArrayRef, nullability: Nullability) -> VortexResult { + let dtype = DType::Bool(nullability); + if array.dtype() == &dtype { + Ok(array.clone()) + } else { + array.cast(dtype) + } +} + +fn boolean_nullability(lhs: &ArrayRef, rhs: &ArrayRef) -> Nullability { + lhs.dtype().nullability() | rhs.dtype().nullability() +} + +fn all_validity(nullability: Nullability) -> Validity { + match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::AllValid, + } +} + +fn mask_to_validity(mask: Mask, nullability: Nullability) -> Validity { + match nullability { + Nullability::NonNullable => { + debug_assert!(mask.all_true()); + Validity::NonNullable + } + Nullability::Nullable => Validity::from(mask.into_bit_buffer()), + } +} + +#[inline] +fn is_boolean_operator(operator: Operator) -> bool { + matches!(operator, Operator::And | Operator::Or) } #[cfg(test)] mod tests { use rstest::rstest; + use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::BoolArray; + use crate::arrays::ConstantArray; + use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; + #[test] + fn test_kleene_truth_table() -> VortexResult<()> { + let lhs = BoolArray::from_iter([ + Some(true), + Some(true), + Some(true), + Some(false), + Some(false), + Some(false), + None, + None, + None, + ]) + .into_array(); + let rhs = BoolArray::from_iter([ + Some(true), + Some(false), + None, + Some(true), + Some(false), + None, + Some(true), + Some(false), + None, + ]) + .into_array(); + + assert_arrays_eq!( + lhs.binary(rhs.clone(), Operator::And)?, + BoolArray::from_iter([ + Some(true), + Some(false), + None, + Some(false), + Some(false), + Some(false), + None, + Some(false), + None, + ]) + ); + + assert_arrays_eq!( + lhs.binary(rhs, Operator::Or)?, + BoolArray::from_iter([ + Some(true), + Some(true), + Some(true), + Some(true), + Some(false), + None, + Some(true), + None, + None, + ]) + ); + + Ok(()) + } + + #[test] + fn test_null_constant_kleene() -> VortexResult<()> { + let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); + let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) + .into_array(); + + assert_arrays_eq!( + lhs.binary(null.clone(), Operator::And)?, + BoolArray::from_iter([Some(false), None, None]) + ); + assert_arrays_eq!( + lhs.binary(null, Operator::Or)?, + BoolArray::from_iter([None, Some(true), None]) + ); + + Ok(()) + } + #[rstest] #[case( BoolArray::from_iter([Some(true), Some(true), Some(false), Some(false)]).into_array(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 1c860cb75b5..531ee22f5ce 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -37,8 +37,12 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; -pub(crate) mod boolean; -pub(crate) use boolean::*; +pub mod boolean; +pub use boolean::BooleanExecuteAdaptor; +pub use boolean::BooleanKernel; +pub use boolean::boolean_buffer_scalar; +pub use boolean::boolean_buffers; +pub(crate) use boolean::execute_boolean; mod compare; pub use compare::*; mod numeric; From 65cedb57a225c478929077f8a5b90ec5283f6789 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 10:01:25 -0400 Subject: [PATCH 2/3] Address Kleene boolean review feedback Signed-off-by: Nicholas Gates --- .../src/scalar_fn/fns/binary/boolean.rs | 63 ++++++------------- vortex-buffer/src/bit/buf.rs | 30 +++++++++ vortex-buffer/src/bit/mod.rs | 20 ++++++ vortex-buffer/src/bit/ops.rs | 12 +--- 4 files changed, 69 insertions(+), 56 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 974ba2d1594..940ef57d6f0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -6,6 +6,7 @@ use std::iter::repeat_n; use arrow_array::cast::AsArray; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; +use vortex_buffer::read_u64_le; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::AllOr; @@ -37,11 +38,14 @@ use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; -/// Trait for encoding-specific Kleene boolean kernels that operate in encoded space. +/// Trait for encoding-specific boolean kernels that operate in encoded space. /// /// Implementations receive the encoded array as the left operand. `rhs` may be any boolean array /// encoding or a constant; implementations should return `Ok(None)` when they cannot handle that /// operand without falling back to ordinary execution. +/// +/// Vortex's boolean [`Operator::And`] and [`Operator::Or`] variants use Kleene semantics; there is +/// no separate two-valued boolean operator path to dispatch here. pub trait BooleanKernel: VTable { /// Execute `lhs rhs` using Kleene boolean semantics. fn boolean( @@ -268,7 +272,7 @@ pub fn boolean_buffers( Operator::Or => lhs_values | &rhs_values, other => return Err(vortex_err!("Not a boolean operator: {other}")), }; - return Ok(BoolArray::try_new(values, all_validity(nullability))?.into_array()); + return Ok(BoolArray::try_new(values, Validity::from(nullability))?.into_array()); } let lhs_valid = lhs_validity.execute_mask(len, ctx)?; @@ -318,14 +322,14 @@ pub fn boolean_buffer_scalar( .bitand_not(&Mask::from_buffer(values)); BoolArray::try_new( BitBuffer::new_unset(len), - mask_to_validity(valid, nullability), + Validity::from_mask(valid, nullability), )? } (Operator::Or, None) => { let valid = validity.execute_mask(len, ctx)? & &Mask::from_buffer(values); BoolArray::try_new( BitBuffer::new_set(len), - mask_to_validity(valid, nullability), + Validity::from_mask(valid, nullability), )? } (other, _) => return Err(vortex_err!("Not a boolean operator: {other}")), @@ -430,16 +434,16 @@ fn fused_boolean_buffers_aligned( operator: Operator, nullability: Nullability, ) -> VortexResult> { - let Some(lhs_values) = word_source_from_bit_buffer(lhs_values, len) else { + let Some(lhs_values) = word_source_from_bit_buffer(lhs_values) else { return Ok(None); }; - let Some(rhs_values) = word_source_from_bit_buffer(rhs_values, len) else { + let Some(rhs_values) = word_source_from_bit_buffer(rhs_values) else { return Ok(None); }; - let Some(lhs_validity) = word_source_from_mask(lhs_validity, len) else { + let Some(lhs_validity) = word_source_from_mask(lhs_validity) else { return Ok(None); }; - let Some(rhs_validity) = word_source_from_mask(rhs_validity, len) else { + let Some(rhs_validity) = word_source_from_mask(rhs_validity) else { return Ok(None); }; @@ -454,32 +458,18 @@ fn fused_boolean_buffers_aligned( )?)) } -fn word_source_from_bit_buffer(buffer: &BitBuffer, len: usize) -> Option> { - if !buffer.offset().is_multiple_of(8) { - return None; - } - - let n_bytes = len.div_ceil(8); - let start = buffer.offset() / 8; - let end = start + n_bytes; - Some(WordSource::Bytes(&buffer.inner().as_slice()[start..end])) +fn word_source_from_bit_buffer(buffer: &BitBuffer) -> Option> { + buffer.byte_aligned_bytes().map(WordSource::Bytes) } -fn word_source_from_mask(mask: &Mask, len: usize) -> Option> { +fn word_source_from_mask(mask: &Mask) -> Option> { match mask.bit_buffer() { AllOr::All => Some(WordSource::Fill(u64::MAX)), AllOr::None => Some(WordSource::Fill(0)), - AllOr::Some(buffer) => word_source_from_bit_buffer(buffer, len), + AllOr::Some(buffer) => word_source_from_bit_buffer(buffer), } } -fn read_u64_le(bytes: &[u8]) -> u64 { - debug_assert!(bytes.len() <= 8); - let mut buf = [0; 8]; - buf[..bytes.len()].copy_from_slice(bytes); - u64::from_le_bytes(buf) -} - fn fused_boolean_word_sources( len: usize, lhs_words: WordSource<'_>, @@ -581,7 +571,7 @@ fn finish_fused_boolean_words( validity.truncate(n_bytes); Ok(BoolArray::try_new( BitBuffer::new(values.freeze(), len), - mask_to_validity( + Validity::from_mask( Mask::from_buffer(BitBuffer::new(validity.freeze(), len)), nullability, ), @@ -651,7 +641,7 @@ where Ok(BoolArray::try_new( values, - mask_to_validity(Mask::from_buffer(validity), nullability), + Validity::from_mask(Mask::from_buffer(validity), nullability), )? .into_array()) } @@ -677,23 +667,6 @@ fn boolean_nullability(lhs: &ArrayRef, rhs: &ArrayRef) -> Nullability { lhs.dtype().nullability() | rhs.dtype().nullability() } -fn all_validity(nullability: Nullability) -> Validity { - match nullability { - Nullability::NonNullable => Validity::NonNullable, - Nullability::Nullable => Validity::AllValid, - } -} - -fn mask_to_validity(mask: Mask, nullability: Nullability) -> Validity { - match nullability { - Nullability::NonNullable => { - debug_assert!(mask.all_true()); - Validity::NonNullable - } - Nullability::Nullable => Validity::from(mask.into_bit_buffer()), - } -} - #[inline] fn is_boolean_operator(operator: Operator) -> bool { matches!(operator, Operator::And | Operator::Or) diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 61b69b3f372..a8bd9e1929c 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -272,6 +272,23 @@ impl BitBuffer { &self.buffer } + /// Return the backing bytes for this bit buffer when its logical offset is byte-aligned. + /// + /// The returned slice contains exactly `self.len().div_ceil(8)` bytes. Bits past the logical + /// length in the final byte are outside the buffer's logical range and should be ignored by + /// callers. + #[inline] + pub fn byte_aligned_bytes(&self) -> Option<&[u8]> { + if !self.offset.is_multiple_of(8) { + return None; + } + + let n_bytes = self.len.div_ceil(8); + let start = self.offset / 8; + let end = start + n_bytes; + Some(&self.buffer.as_slice()[start..end]) + } + /// Retrieve the value at the given index. /// /// Panics if the index is out of bounds. @@ -692,6 +709,19 @@ mod tests { assert_eq!(sliced.offset(), 2); } + #[test] + fn test_byte_aligned_bytes() { + let bytes: ByteBuffer = buffer![0b1010_0101u8, 0b0000_0011]; + let buf = BitBuffer::new(bytes.clone(), 10); + assert_eq!(buf.byte_aligned_bytes(), Some(bytes.as_slice())); + + let byte_sliced = buf.slice(8..10); + assert_eq!(byte_sliced.byte_aligned_bytes(), Some(&[0b0000_0011][..])); + + let bit_sliced = buf.slice(1..9); + assert!(bit_sliced.byte_aligned_bytes().is_none()); + } + #[test] fn test_from_indices_dense_crosses_words() { let len = 130; diff --git a/vortex-buffer/src/bit/mod.rs b/vortex-buffer/src/bit/mod.rs index 41bb5797266..a39c93e3f25 100644 --- a/vortex-buffer/src/bit/mod.rs +++ b/vortex-buffer/src/bit/mod.rs @@ -74,6 +74,19 @@ where } } +/// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 +/// bytes are supplied. +/// +/// This preserves Vortex's least-significant-bit-first bitmap numbering on little- and big-endian +/// targets. For a full 8-byte slice it lowers to a single word load. +#[inline] +pub fn read_u64_le(bytes: &[u8]) -> u64 { + debug_assert!(bytes.len() <= 8); + let mut buf = [0u8; 8]; + buf[..bytes.len()].copy_from_slice(bytes); + u64::from_le_bytes(buf) +} + /// Splice a packed word `w` (whose bits above the highest valid bit are zero) into /// `words` at the given bit position. /// @@ -178,6 +191,7 @@ pub unsafe fn unset_bit_unchecked(buf: *mut u8, index: usize) { mod tests { use super::collect_bool_word; use super::pack_bools_into_words; + use super::read_u64_le; #[test] fn collect_bool_word_packs_lsb_first() { @@ -190,6 +204,12 @@ mod tests { assert_eq!(collect_bool_word(0, |_| true), 0); } + #[test] + fn read_u64_le_zero_pads_tail() { + assert_eq!(read_u64_le(&[0x34, 0x12]), 0x1234); + assert_eq!(read_u64_le(&[0xff; 8]), u64::MAX); + } + #[test] #[should_panic(expected = "cannot pack 65 bits into a u64 word")] fn collect_bool_word_rejects_too_many_bits() { diff --git a/vortex-buffer/src/bit/ops.rs b/vortex-buffer/src/bit/ops.rs index 5006fbcd4cc..ec922b54bc2 100644 --- a/vortex-buffer/src/bit/ops.rs +++ b/vortex-buffer/src/bit/ops.rs @@ -7,17 +7,7 @@ use crate::BitBuffer; use crate::BitBufferMut; use crate::BufferMut; use crate::ByteBufferMut; - -/// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 are -/// supplied. Using [`u64::from_le_bytes`] keeps the bit-numbering identical on little- and -/// big-endian targets; for a full 8-byte slice it lowers to a single word load. -#[inline] -fn read_u64_le(bytes: &[u8]) -> u64 { - debug_assert!(bytes.len() <= 8); - let mut buf = [0u8; 8]; - buf[..bytes.len()].copy_from_slice(bytes); - u64::from_le_bytes(buf) -} +use crate::read_u64_le; trait BitWordTarget { fn byte_len(&self) -> usize; From 225076f795c85f82dd4b009c238bd1c40c5e7845 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 18 Jun 2026 11:21:23 -0400 Subject: [PATCH 3/3] Address Kleene boolean PR comments Signed-off-by: Nicholas Gates --- encodings/bytebool/src/compute.rs | 14 +- .../src/arrays/bool/compute/boolean.rs | 8 +- .../src/scalar_fn/fns/binary/boolean.rs | 308 +++++++++++------- vortex-array/src/scalar_fn/fns/binary/mod.rs | 4 +- 4 files changed, 207 insertions(+), 127 deletions(-) diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 89662f44e6e..fb72fc4e549 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -6,14 +6,15 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::match_each_integer_ptype; use vortex_array::scalar_fn::fns::binary::BooleanKernel; -use vortex_array::scalar_fn::fns::binary::boolean_buffer_scalar; -use vortex_array::scalar_fn::fns::binary::boolean_buffers; +use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; +use vortex_array::scalar_fn::fns::binary::kleene_boolean_buffers; use vortex_array::scalar_fn::fns::cast::CastKernel; use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_array::scalar_fn::fns::mask::MaskReduce; @@ -122,12 +123,12 @@ impl BooleanKernel for ByteBool { let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let lhs_values = truthy_bit_buffer(lhs); - if let Some(rhs) = rhs.as_opt::() { + if let Some(rhs) = rhs.as_opt::() { let rhs = rhs .scalar() .as_bool_opt() .ok_or_else(|| vortex_err!("expected boolean scalar"))?; - return boolean_buffer_scalar( + return kleene_boolean_buffer_scalar( lhs_values, lhs.validity()?, &rhs, @@ -142,7 +143,7 @@ impl BooleanKernel for ByteBool { return Ok(None); }; - boolean_buffers( + kleene_boolean_buffers( lhs_values, lhs.validity()?, truthy_bit_buffer(rhs), @@ -156,7 +157,8 @@ impl BooleanKernel for ByteBool { } fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer { - BitBuffer::from_iter(array.truthy_bytes().iter().map(|byte| *byte != 0)) + let bytes = array.truthy_bytes(); + BitBuffer::collect_bool(bytes.len(), |idx| bytes[idx] != 0) } #[cfg(test)] diff --git a/vortex-array/src/arrays/bool/compute/boolean.rs b/vortex-array/src/arrays/bool/compute/boolean.rs index 438833d84bb..28d1bbb7f52 100644 --- a/vortex-array/src/arrays/bool/compute/boolean.rs +++ b/vortex-array/src/arrays/bool/compute/boolean.rs @@ -11,8 +11,8 @@ use crate::arrays::Bool; use crate::arrays::Constant; use crate::arrays::bool::BoolArrayExt; use crate::scalar_fn::fns::binary::BooleanKernel; -use crate::scalar_fn::fns::binary::boolean_buffer_scalar; -use crate::scalar_fn::fns::binary::boolean_buffers; +use crate::scalar_fn::fns::binary::kleene_boolean_buffer_scalar; +use crate::scalar_fn::fns::binary::kleene_boolean_buffers; use crate::scalar_fn::fns::operators::Operator; impl BooleanKernel for Bool { @@ -29,7 +29,7 @@ impl BooleanKernel for Bool { .scalar() .as_bool_opt() .ok_or_else(|| vortex_err!("expected boolean scalar"))?; - return boolean_buffer_scalar( + return kleene_boolean_buffer_scalar( lhs.to_bit_buffer(), lhs.validity()?, &rhs, @@ -44,7 +44,7 @@ impl BooleanKernel for Bool { return Ok(None); }; - boolean_buffers( + kleene_boolean_buffers( lhs.to_bit_buffer(), lhs.validity()?, rhs.to_bit_buffer(), diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 940ef57d6f0..40e86deb839 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -8,6 +8,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_buffer::read_u64_le; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -172,7 +173,7 @@ fn arrow_execute_boolean( let array = match op { Operator::And => arrow_arith::boolean::and_kleene(&lhs, &rhs)?, Operator::Or => arrow_arith::boolean::or_kleene(&lhs, &rhs)?, - other => return Err(vortex_err!("Not a boolean operator: {other}")), + other => vortex_bail!("Not a boolean operator: {other}"), }; ArrayRef::from_arrow(&array, nullable == Nullability::Nullable) @@ -222,7 +223,7 @@ fn constant_array_boolean( ))), (Operator::Or, Some(false)) => Ok(Some(cast_bool_nullability(array, nullability)?)), (Operator::And | Operator::Or, None) => Ok(None), - (other, _) => Err(vortex_err!("Not a boolean operator: {other}")), + (other, _) => vortex_bail!("Not a boolean operator: {other}"), } } @@ -242,7 +243,7 @@ fn boolean_scalar_scalar( (None, _) | (_, None) => None, (Some(l), Some(r)) => Some(l | r), }, - other => return Err(vortex_err!("Not a boolean operator: {other}")), + other => vortex_bail!("Not a boolean operator: {other}"), }) } @@ -254,7 +255,7 @@ fn bool_scalar_value(scalar: &Scalar) -> VortexResult> { } /// Execute a Kleene boolean operation from boolean value bitmaps and validity values. -pub fn boolean_buffers( +pub fn kleene_boolean_buffers( lhs_values: BitBuffer, lhs_validity: Validity, rhs_values: BitBuffer, @@ -270,7 +271,7 @@ pub fn boolean_buffers( let values = match operator { Operator::And => lhs_values & &rhs_values, Operator::Or => lhs_values | &rhs_values, - other => return Err(vortex_err!("Not a boolean operator: {other}")), + other => vortex_bail!("Not a boolean operator: {other}"), }; return Ok(BoolArray::try_new(values, Validity::from(nullability))?.into_array()); } @@ -289,7 +290,7 @@ pub fn boolean_buffers( } /// Execute a Kleene boolean operation between boolean value bits and a scalar. -pub fn boolean_buffer_scalar( +pub fn kleene_boolean_buffer_scalar( values: BitBuffer, validity: Validity, scalar: &BoolScalar<'_>, @@ -332,7 +333,7 @@ pub fn boolean_buffer_scalar( Validity::from_mask(valid, nullability), )? } - (other, _) => return Err(vortex_err!("Not a boolean operator: {other}")), + (other, _) => vortex_bail!("Not a boolean operator: {other}"), }; Ok(result.into_array()) @@ -478,6 +479,35 @@ fn fused_boolean_word_sources( rhs_valid_words: WordSource<'_>, operator: Operator, nullability: Nullability, +) -> VortexResult { + match operator { + Operator::And => fused_boolean_and_word_sources( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + Operator::Or => fused_boolean_or_word_sources( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + other => vortex_bail!("Not a boolean operator: {other}"), + } +} + +fn fused_boolean_and_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + nullability: Nullability, ) -> VortexResult { let n_bytes = len.div_ceil(8); let n_words = n_bytes.div_ceil(8); @@ -485,74 +515,81 @@ fn fused_boolean_word_sources( let mut values = BufferMut::::with_capacity(n_words); let mut validity = BufferMut::::with_capacity(n_words); - match operator { - Operator::And => { - for byte_offset in (0..full_bytes).step_by(8) { - let lhs = lhs_words.word_at(byte_offset, 8); - let rhs = rhs_words.word_at(byte_offset, 8); - let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); - let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); - - // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this - // loop plus the optional tail push emits at most `n_words` words. - unsafe { - values.push_unchecked(lhs & rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), - ); - } - } - - if full_bytes != n_bytes { - let tail_len = n_bytes - full_bytes; - let lhs = lhs_words.word_at(full_bytes, tail_len); - let rhs = rhs_words.word_at(full_bytes, tail_len); - let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); - let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); - - // SAFETY: see the loop safety comment above. - unsafe { - values.push_unchecked(lhs & rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), - ); - } - } + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); } - Operator::Or => { - for byte_offset in (0..full_bytes).step_by(8) { - let lhs = lhs_words.word_at(byte_offset, 8); - let rhs = rhs_words.word_at(byte_offset, 8); - let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); - let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); - - // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this - // loop plus the optional tail push emits at most `n_words` words. - unsafe { - values.push_unchecked(lhs | rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), - ); - } - } - - if full_bytes != n_bytes { - let tail_len = n_bytes - full_bytes; - let lhs = lhs_words.word_at(full_bytes, tail_len); - let rhs = rhs_words.word_at(full_bytes, tail_len); - let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); - let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); - - // SAFETY: see the loop safety comment above. - unsafe { - values.push_unchecked(lhs | rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), - ); - } - } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); + } + } + + finish_fused_boolean_words(len, n_bytes, values, validity, nullability) +} + +fn fused_boolean_or_word_sources( + len: usize, + lhs_words: WordSource<'_>, + rhs_words: WordSource<'_>, + lhs_valid_words: WordSource<'_>, + rhs_valid_words: WordSource<'_>, + nullability: Nullability, +) -> VortexResult { + let n_bytes = len.div_ceil(8); + let n_words = n_bytes.div_ceil(8); + let full_bytes = n_bytes - n_bytes % 8; + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for byte_offset in (0..full_bytes).step_by(8) { + let lhs = lhs_words.word_at(byte_offset, 8); + let rhs = rhs_words.word_at(byte_offset, 8); + let lhs_valid = lhs_valid_words.word_at(byte_offset, 8); + let rhs_valid = rhs_valid_words.word_at(byte_offset, 8); + + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this + // loop plus the optional tail push emits at most `n_words` words. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); + } + } + + if full_bytes != n_bytes { + let tail_len = n_bytes - full_bytes; + let lhs = lhs_words.word_at(full_bytes, tail_len); + let rhs = rhs_words.word_at(full_bytes, tail_len); + let lhs_valid = lhs_valid_words.word_at(full_bytes, tail_len); + let rhs_valid = rhs_valid_words.word_at(full_bytes, tail_len); + + // SAFETY: see the loop safety comment above. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); } - other => return Err(vortex_err!("Not a boolean operator: {other}")), } finish_fused_boolean_words(len, n_bytes, values, validity, nullability) @@ -588,6 +625,41 @@ fn fused_boolean_words( operator: Operator, nullability: Nullability, ) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + match operator { + Operator::And => fused_boolean_and_words( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + Operator::Or => fused_boolean_or_words( + len, + lhs_words, + rhs_words, + lhs_valid_words, + rhs_valid_words, + nullability, + ), + other => vortex_bail!("Not a boolean operator: {other}"), + } +} + +fn fused_boolean_and_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + nullability: Nullability, +) -> VortexResult where L: Iterator, R: Iterator, @@ -598,52 +670,58 @@ where let mut values = BufferMut::::with_capacity(n_words); let mut validity = BufferMut::::with_capacity(n_words); - match operator { - Operator::And => { - for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words - .zip(rhs_words) - .zip(lhs_valid_words) - .zip(rhs_valid_words) - .take(n_words) - { - // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this - // loop is capped at `n_words`. - unsafe { - values.push_unchecked(lhs & rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs), - ); - } - } - } - Operator::Or => { - for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words - .zip(rhs_words) - .zip(lhs_valid_words) - .zip(rhs_valid_words) - .take(n_words) - { - // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this - // loop is capped at `n_words`. - unsafe { - values.push_unchecked(lhs | rhs); - validity.push_unchecked( - (lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs), - ); - } - } + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is + // capped at `n_words`. + unsafe { + values.push_unchecked(lhs & rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & !lhs) | (rhs_valid & !rhs)); } - other => return Err(vortex_err!("Not a boolean operator: {other}")), } - let values = BitBuffer::new(values.freeze().into_byte_buffer(), len); - let validity = BitBuffer::new(validity.freeze().into_byte_buffer(), len); + finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability) +} - Ok(BoolArray::try_new( - values, - Validity::from_mask(Mask::from_buffer(validity), nullability), - )? - .into_array()) +fn fused_boolean_or_words( + len: usize, + lhs_words: L, + rhs_words: R, + lhs_valid_words: LV, + rhs_valid_words: RV, + nullability: Nullability, +) -> VortexResult +where + L: Iterator, + R: Iterator, + LV: Iterator, + RV: Iterator, +{ + let n_words = len.div_ceil(64); + let mut values = BufferMut::::with_capacity(n_words); + let mut validity = BufferMut::::with_capacity(n_words); + + for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words + .zip(rhs_words) + .zip(lhs_valid_words) + .zip(rhs_valid_words) + .take(n_words) + { + // SAFETY: both buffers were allocated with exactly `n_words` capacity, and this loop is + // capped at `n_words`. + unsafe { + values.push_unchecked(lhs | rhs); + validity + .push_unchecked((lhs_valid & rhs_valid) | (lhs_valid & lhs) | (rhs_valid & rhs)); + } + } + + finish_fused_boolean_words(len, len.div_ceil(8), values, validity, nullability) } fn constant_bool_result(value: Option, len: usize, nullability: Nullability) -> ArrayRef { diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 531ee22f5ce..bbb392a3752 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -40,9 +40,9 @@ use crate::scalar_fn::fns::operators::Operator; pub mod boolean; pub use boolean::BooleanExecuteAdaptor; pub use boolean::BooleanKernel; -pub use boolean::boolean_buffer_scalar; -pub use boolean::boolean_buffers; pub(crate) use boolean::execute_boolean; +pub use boolean::kleene_boolean_buffer_scalar; +pub use boolean::kleene_boolean_buffers; mod compare; pub use compare::*; mod numeric;