diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 8b83d0f8128..2dd08c55832 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -4,7 +4,6 @@ use std::fmt::Display; use std::fmt::Formatter; -use arrow_array::BooleanArray; use smallvec::smallvec; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMeta; @@ -338,14 +337,16 @@ impl FromIterator for BoolArray { impl FromIterator> for BoolArray { fn from_iter>>(iter: I) -> Self { - let (buffer, nulls) = BooleanArray::from_iter(iter).into_parts(); + let iter = iter.into_iter(); + let capacity = iter.size_hint().0; + let mut bits = BitBufferMut::with_capacity(capacity); + let mut validity = BitBufferMut::with_capacity(capacity); + for value in iter { + bits.append(value.unwrap_or_default()); + validity.append(value.is_some()); + } - BoolArray::new( - BitBuffer::from(buffer), - nulls - .map(|n| Validity::from(BitBuffer::from(n.into_inner()))) - .unwrap_or(Validity::AllValid), - ) + BoolArray::new(bits.freeze(), Validity::from(validity.freeze())) } } diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index f35f628c332..52b0373ac49 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -83,19 +83,13 @@ pub(super) fn execute_filter_fast_paths( } /// Filter a canonical array by a mask, returning a new canonical array. -pub(super) fn execute_filter( - canonical: Canonical, - mask: &Arc, - ctx: &mut ExecutionCtx, -) -> Canonical { +pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Canonical { match canonical { Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())), Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)), Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)), Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)), - Canonical::VarBinView(a) => { - Canonical::VarBinView(varbinview::filter_varbinview(&a, mask, ctx)) - } + Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)), Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)), Canonical::FixedSizeList(a) => { Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask)) diff --git a/vortex-array/src/arrays/filter/execute/varbinview.rs b/vortex-array/src/arrays/filter/execute/varbinview.rs index 1f9df8f99cd..58a348b41b0 100644 --- a/vortex-array/src/arrays/filter/execute/varbinview.rs +++ b/vortex-array/src/arrays/filter/execute/varbinview.rs @@ -3,54 +3,32 @@ use std::sync::Arc; -use arrow_array::BooleanArray; -use vortex_error::VortexExpect; -use vortex_mask::Mask; +use vortex_buffer::Buffer; use vortex_mask::MaskValues; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::VarBinView; use crate::arrays::VarBinViewArray; -use crate::arrow::ArrowSessionExt; -use crate::arrow::FromArrowArray; - -pub fn filter_varbinview( - array: &VarBinViewArray, - mask: &Arc, - ctx: &mut ExecutionCtx, -) -> VarBinViewArray { - // Delegate to the Arrow implementation of filter over `VarBinView`. - arrow_filter_fn( - &array.clone().into_array(), - &Mask::Values(Arc::clone(mask)), - ctx, - ) - .vortex_expect("VarBinViewArray is Arrow-compatible and supports arrow_filter_fn") - .as_::() - .into_owned() -} - -fn arrow_filter_fn( - array: &ArrayRef, - mask: &Mask, - ctx: &mut ExecutionCtx, -) -> vortex_error::VortexResult { - let values = match &mask { - Mask::Values(values) => values, - Mask::AllTrue(_) | Mask::AllFalse(_) => unreachable!("check in filter invoke"), - }; - - let array_ref = ctx - .session() - .arrow() - .clone() - .execute_arrow(array.clone(), None, ctx)?; - let mask_array = BooleanArray::new(values.bit_buffer().clone().into(), None); - let filtered = arrow_select::filter::filter(array_ref.as_ref(), &mask_array)?; - - ArrayRef::from_arrow(filtered.as_ref(), array.dtype().is_nullable()) +use crate::arrays::filter::execute::buffer; +use crate::arrays::filter::execute::filter_validity; +use crate::arrays::varbinview::BinaryView; +use crate::arrays::varbinview::VarBinViewArrayExt; +use crate::buffer::BufferHandle; + +pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> VarBinViewArray { + let filtered_validity = filter_validity(array.varbinview_validity(), mask); + + let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()); + let filtered_views = buffer::filter_buffer(views, mask.as_ref()); + + // SAFETY: the filtered views are a subset of the original views and reference the same data + // buffers, and the validity is filtered by the same mask so lengths stay aligned. + unsafe { + VarBinViewArray::new_handle_unchecked( + BufferHandle::new_host(filtered_views.into_byte_buffer()), + Arc::clone(array.data_buffers()), + array.dtype().clone(), + filtered_validity, + ) + } } #[cfg(test)] diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 4f556eff61b..0f54a97c574 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -168,7 +168,7 @@ impl VTable for Filter { // TODO(joe): fix the ownership of AnyCanonical let child = Canonical::from(array.child().as_::()); Ok(ExecutionResult::done( - execute_filter(child, &mask_values, ctx).into_array(), + execute_filter(child, &mask_values).into_array(), )) } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 42876a92aca..f6fff972c9b 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -1,17 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use arrow_array::BinaryArray; -use arrow_array::LargeBinaryArray; -use arrow_array::LargeStringArray; -use arrow_array::StringArray; -use arrow_ord::cmp; -use arrow_schema::DataType; +use std::cmp::Ordering; + use vortex_buffer::BitBuffer; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; @@ -20,19 +15,15 @@ use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBin; -use crate::arrays::VarBinViewArray; use crate::arrays::varbin::VarBinArrayExt; -use crate::arrow::Datum; -use crate::arrow::from_arrow_columnar; -use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::match_each_integer_ptype; use crate::scalar_fn::fns::binary::CompareKernel; use crate::scalar_fn::fns::operators::CompareOperator; -use crate::scalar_fn::fns::operators::Operator; -// This implementation exists so we can have custom translation of RHS to arrow that's not the same as IntoCanonical +// This implementation exists so we can compare against a constant in encoded space, without +// canonicalizing the VarBin array to VarBinView. impl CompareKernel for VarBin { fn compare( lhs: ArrayView<'_, VarBin>, @@ -40,106 +31,66 @@ impl CompareKernel for VarBin { operator: CompareOperator, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(rhs_const) = rhs.as_constant() { - let nullable = lhs.dtype().is_nullable() || rhs_const.dtype().is_nullable(); - let len = lhs.len(); - - let rhs_is_empty = match rhs_const.dtype() { - DType::Binary(_) => rhs_const - .as_binary() - .is_empty() - .vortex_expect("RHS should not be null"), - DType::Utf8(_) => rhs_const - .as_utf8() - .is_empty() - .vortex_expect("RHS should not be null"), - _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"), - }; - - if rhs_is_empty { - let buffer = match operator { - CompareOperator::Gte => BitBuffer::new_set(len), /* Every possible value is >= "" */ - CompareOperator::Lt => BitBuffer::new_unset(len), // No value is < "" - CompareOperator::Eq | CompareOperator::Lte => { - let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(lhs_offsets.ptype(), |P| { - compare_offsets_to_empty::

(lhs_offsets, true) - }) - } - CompareOperator::NotEq | CompareOperator::Gt => { - let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; - match_each_integer_ptype!(lhs_offsets.ptype(), |P| { - compare_offsets_to_empty::

(lhs_offsets, false) - }) - } - }; + let Some(rhs_const) = rhs.as_constant() else { + return Ok(None); + }; - return Ok(Some( - BoolArray::new( - buffer, - lhs.validity()?.union_nullability(rhs.dtype().nullability()), - ) - .into_array(), - )); - } - - let lhs = Datum::try_new(lhs.array(), ctx)?; + let len = lhs.len(); - // The RHS scalar must match the LHS Arrow data type. VarBin with i64 offsets is - // converted to LargeBinary/LargeUtf8 (see `preferred_arrow_type`), and Arrow refuses to - // compare LargeBinary with Binary (or LargeUtf8 with Utf8). - let arrow_rhs: &dyn arrow_array::Datum = match (rhs_const.dtype(), lhs.data_type()) { - (DType::Utf8(_), DataType::LargeUtf8) => &rhs_const - .as_utf8() - .value() - .map(LargeStringArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(LargeStringArray::new_null(1))), - (DType::Utf8(_), _) => &rhs_const - .as_utf8() - .value() - .map(StringArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(StringArray::new_null(1))), - (DType::Binary(_), DataType::LargeBinary) => &rhs_const - .as_binary() - .value() - .map(LargeBinaryArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(LargeBinaryArray::new_null(1))), - (DType::Binary(_), _) => &rhs_const - .as_binary() - .value() - .map(BinaryArray::new_scalar) - .unwrap_or_else(|| arrow_array::Scalar::new(BinaryArray::new_null(1))), - _ => vortex_bail!( - "VarBin array RHS can only be Utf8 or Binary, given {}", - rhs_const.dtype() - ), - }; + // The compare adaptor resolves null constants before dispatching to this kernel, so + // the scalar always carries a value. + let rhs_bytes: &[u8] = match rhs_const.dtype() { + DType::Binary(_) => rhs_const + .as_binary() + .value() + .vortex_expect("RHS should not be null") + .as_slice(), + DType::Utf8(_) => rhs_const + .as_utf8() + .value() + .vortex_expect("RHS should not be null") + .as_str() + .as_bytes(), + _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"), + }; - let array = match operator { - CompareOperator::Eq => cmp::eq(&lhs, arrow_rhs), - CompareOperator::NotEq => cmp::neq(&lhs, arrow_rhs), - CompareOperator::Gt => cmp::gt(&lhs, arrow_rhs), - CompareOperator::Gte => cmp::gt_eq(&lhs, arrow_rhs), - CompareOperator::Lt => cmp::lt(&lhs, arrow_rhs), - CompareOperator::Lte => cmp::lt_eq(&lhs, arrow_rhs), + let buffer = if rhs_bytes.is_empty() { + // Comparisons against "" only need the value lengths, i.e. the offset deltas. + match operator { + CompareOperator::Gte => BitBuffer::new_set(len), /* Every possible value is >= "" */ + CompareOperator::Lt => BitBuffer::new_unset(len), // No value is < "" + CompareOperator::Eq | CompareOperator::Lte => { + let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_offsets_to_empty::

(lhs_offsets, true) + }) + } + CompareOperator::NotEq | CompareOperator::Gt => { + let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_offsets_to_empty::

(lhs_offsets, false) + }) + } } - .map_err(|err| vortex_err!("Failed to compare VarBin array: {}", err))?; - - Ok(Some(from_arrow_columnar(&array, len, nullable, ctx)?)) - } else if !rhs.is::() { - // NOTE: If the rhs is not a VarBin array it will be canonicalized to a VarBinView - // Arrow doesn't support comparing VarBin to VarBinView arrays, so we convert ourselves - // to VarBinView and re-invoke. - Ok(Some( - lhs.array() - .clone() - .execute::(ctx)? - .into_array() - .binary(rhs.clone(), Operator::from(operator))?, - )) } else { - Ok(None) - } + let lhs_offsets = lhs.offsets().clone().execute::(ctx)?; + match_each_integer_ptype!(lhs_offsets.ptype(), |P| { + compare_bytes_to_constant( + lhs_offsets.as_slice::

(), + lhs.bytes().as_slice(), + rhs_bytes, + operator, + ) + }) + }; + + Ok(Some( + BoolArray::new( + buffer, + lhs.validity()?.union_nullability(rhs.dtype().nullability()), + ) + .into_array(), + )) } } @@ -153,6 +104,71 @@ fn compare_offsets_to_empty(offsets: PrimitiveArray, eq: bool) }) } +/// Compare every value in a VarBin array against a constant, resolving values through the +/// offsets. Dispatches the operator outside the lane loop so each predicate inlines into its +/// own loop. +fn compare_bytes_to_constant( + offsets: &[P], + bytes: &[u8], + constant: &[u8], + operator: CompareOperator, +) -> BitBuffer { + match operator { + CompareOperator::Eq => { + collect_lane_bits(offsets, |start, end| value_eq(bytes, start, end, constant)) + } + CompareOperator::NotEq => { + collect_lane_bits(offsets, |start, end| !value_eq(bytes, start, end, constant)) + } + CompareOperator::Gt => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_gt() + }), + CompareOperator::Gte => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_ge() + }), + CompareOperator::Lt => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_lt() + }), + CompareOperator::Lte => collect_lane_bits(offsets, |start, end| { + value_cmp(bytes, start, end, constant).is_le() + }), + } +} + +/// Bit-pack `predicate(offsets[i], offsets[i + 1])` over each lane of a VarBin array. +fn collect_lane_bits( + offsets: &[P], + predicate: impl Fn(usize, usize) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(offsets.len() - 1, |idx| { + // SAFETY: `collect_bool` yields idx < offsets.len() - 1. + let start = unsafe { offsets.get_unchecked(idx) }.as_(); + let end = unsafe { offsets.get_unchecked(idx + 1) }.as_(); + predicate(start, end) + }) +} + +/// Whether `bytes[start..end]` equals `constant`, comparing lengths first so lanes of a +/// different length never touch the value bytes. +/// +/// Offsets at null positions are not validated, so an out-of-bounds or inverted range is +/// possible there; such lanes answer `false`, and validity masks them out of the result anyway. +#[inline(always)] +fn value_eq(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> bool { + // A lane can only match when its length equals the constant's, so lanes of a different + // length answer without touching the value bytes. An inverted garbage range (start > end) + // wraps to a huge value that never equals `constant.len()`. + end.wrapping_sub(start) == constant.len() + && bytes.get(start..end).is_some_and(|value| value == constant) +} + +/// Order `bytes[start..end]` against `constant`, treating the unvalidated garbage ranges that +/// can appear at null positions as empty; validity masks those lanes out of the result anyway. +#[inline(always)] +fn value_cmp(bytes: &[u8], start: usize, end: usize, constant: &[u8]) -> Ordering { + bytes.get(start..end).unwrap_or_default().cmp(constant) +} + #[cfg(test)] mod test { use vortex_buffer::BitBuffer; @@ -277,11 +293,9 @@ mod tests { ); } - /// Regression: a [`VarBinArray`] built with `i64` offsets is canonicalised to - /// Arrow `LargeUtf8` / `LargeBinary` by `preferred_arrow_type`. Without an explicit - /// branch in [`CompareKernel`], the constant RHS is wrapped in a `StringArray` / - /// `BinaryArray` and Arrow rejects the `LargeUtf8 == Utf8` mismatch. Triggering - /// this only requires `i64` offsets, not large data. + /// Regression: [`CompareKernel`] must handle every offset width; a `VarBinArray` built with + /// `i64` offsets once failed the constant comparison. Triggering this only requires `i64` + /// offsets, not large data. /// /// [`CompareKernel`]: super::CompareKernel #[test]