From 877fe11201ef552450ee9521df01f6b312df9e8b Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 11 Sep 2024 17:03:01 +0100 Subject: [PATCH 1/2] . --- .../src/array/primitive/compute/compare.rs | 23 +++++++--------- vortex-array/src/compute/compare.rs | 27 ++++++++++++++++--- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/array/primitive/compute/compare.rs b/vortex-array/src/array/primitive/compute/compare.rs index 840f3395968..48633ebc215 100644 --- a/vortex-array/src/array/primitive/compute/compare.rs +++ b/vortex-array/src/array/primitive/compute/compare.rs @@ -1,5 +1,5 @@ use arrow_buffer::bit_util::ceil; -use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer}; +use arrow_buffer::{BooleanBuffer, MutableBuffer}; use vortex_dtype::{match_each_native_ptype, NativePType}; use vortex_error::{VortexExpect, VortexResult}; use vortex_scalar::PrimitiveScalar; @@ -7,7 +7,6 @@ use vortex_scalar::PrimitiveScalar; use crate::array::primitive::PrimitiveArray; use crate::array::{BoolArray, ConstantArray}; use crate::compute::{MaybeCompareFn, Operator}; -use crate::validity::ArrayValidity; use crate::{Array, IntoArray}; impl MaybeCompareFn for PrimitiveArray { @@ -42,23 +41,19 @@ fn primitive_const_compare( other: ConstantArray, operator: Operator, ) -> VortexResult { - let mut builder = BooleanBufferBuilder::new(this.len()); let primitive_scalar = PrimitiveScalar::try_from(other.scalar()).vortex_expect("Expected a primitive scalar"); - match_each_native_ptype!(this.ptype(), |$T| { + let buffer = match_each_native_ptype!(this.ptype(), |$T| { let op_fn = operator.to_fn::<$T>(); let typed_value = primitive_scalar.typed_value::<$T>().unwrap(); - for v in this.maybe_null_slice::<$T>() { - builder.append(op_fn(*v, typed_value)); - } + let slice = this.maybe_null_slice::<$T>(); + BooleanBuffer::collect_bool(this.len(), |idx| { + op_fn(unsafe { *slice.get_unchecked(idx) }, typed_value) + }) }); - let validity = this - .validity() - .and(other.logical_validity().into_validity())? - .into_nullable(); - Ok(BoolArray::try_new(builder.finish(), validity)?.into_array()) + Ok(BoolArray::try_new(buffer, this.validity().into_nullable())?.into_array()) } fn apply_predicate bool>( @@ -78,7 +73,7 @@ fn apply_predicate bool>( let mut packed_block = 0_u64; for bit_idx in 0..BLOCK_SIZE { let idx = bit_idx + block * BLOCK_SIZE; - let r = f(lhs[idx], rhs[idx]); + let r = unsafe { f(*lhs.get_unchecked(idx), *rhs.get_unchecked(idx)) }; packed_block |= (r as u64) << bit_idx; } @@ -91,7 +86,7 @@ fn apply_predicate bool>( let mut packed_block = 0_u64; for bit_idx in 0..reminder { let idx = bit_idx + block_count * BLOCK_SIZE; - let r = f(lhs[idx], rhs[idx]); + let r = unsafe { f(*lhs.get_unchecked(idx), *rhs.get_unchecked(idx)) }; packed_block |= (r as u64) << bit_idx; } diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs index b0a5354ed4d..f98927f28e9 100644 --- a/vortex-array/src/compute/compare.rs +++ b/vortex-array/src/compute/compare.rs @@ -89,12 +89,31 @@ pub fn compare(left: &Array, right: &Array, operator: Operator) -> VortexResult< vortex_bail!("Compare operations only support arrays of the same type"); } - if ConstantArray::try_from(left).is_ok() { - let scalar = scalar_at(left, 0)?; - let left_const = ConstantArray::new(scalar, left.len()).into_array(); - return compare(right, &left_const, operator.swap()); + use crate::stats::ArrayStatistics; + use crate::stats::Stat; + + let l_is_const = ConstantArray::try_from(left).is_ok()|| left.statistics().get_as::(Stat::IsConstant).unwrap_or_default(); + let r_is_const = ConstantArray::try_from(right).is_ok()|| right.statistics().get_as::(Stat::IsConstant).unwrap_or_default(); + + match (l_is_const, r_is_const) { + (true, false) => { + let scalar = scalar_at(left, 0)?; + let left_const = ConstantArray::new(scalar, left.len()).into_array(); + return compare(right, &left_const, operator.swap()); + } + (true, true) => { + let original_len = left.len(); + let left = scalar_at(left, 0)?; + let right = scalar_at(right, 0)?; + let r = scalar_cmp(&left, &right, operator); + let arr = ConstantArray::new(r, original_len); + return Ok(arr.into_array()); + } + (false, true) | (false, false) => {} } + + if let Some(selection) = left.with_dyn(|lhs| lhs.compare(right, operator)) { return selection; } From 3b3774a89d99bfd95c99ae18336293ca5ead55fa Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Thu, 12 Sep 2024 10:20:14 +0100 Subject: [PATCH 2/2] . --- vortex-array/src/compute/compare.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs index f98927f28e9..5ab56391165 100644 --- a/vortex-array/src/compute/compare.rs +++ b/vortex-array/src/compute/compare.rs @@ -89,11 +89,18 @@ pub fn compare(left: &Array, right: &Array, operator: Operator) -> VortexResult< vortex_bail!("Compare operations only support arrays of the same type"); } - use crate::stats::ArrayStatistics; - use crate::stats::Stat; - - let l_is_const = ConstantArray::try_from(left).is_ok()|| left.statistics().get_as::(Stat::IsConstant).unwrap_or_default(); - let r_is_const = ConstantArray::try_from(right).is_ok()|| right.statistics().get_as::(Stat::IsConstant).unwrap_or_default(); + use crate::stats::{ArrayStatistics, Stat}; + + let l_is_const = ConstantArray::try_from(left).is_ok() + || left + .statistics() + .get_as::(Stat::IsConstant) + .unwrap_or_default(); + let r_is_const = ConstantArray::try_from(right).is_ok() + || right + .statistics() + .get_as::(Stat::IsConstant) + .unwrap_or_default(); match (l_is_const, r_is_const) { (true, false) => { @@ -112,8 +119,6 @@ pub fn compare(left: &Array, right: &Array, operator: Operator) -> VortexResult< (false, true) | (false, false) => {} } - - if let Some(selection) = left.with_dyn(|lhs| lhs.compare(right, operator)) { return selection; }