diff --git a/encodings/fastlanes/src/for/compare.rs b/encodings/fastlanes/src/for/compare.rs new file mode 100644 index 0000000000..f0f215915d --- /dev/null +++ b/encodings/fastlanes/src/for/compare.rs @@ -0,0 +1,44 @@ +use vortex_array::compute::{compare, CompareFn, Operator}; +use vortex_array::ArrayData; +use vortex_error::VortexResult; + +use crate::{decompress, FoRArray, FoREncoding}; + +impl CompareFn for FoREncoding { + fn compare( + &self, + lhs: &FoRArray, + rhs: &ArrayData, + operator: Operator, + ) -> VortexResult> { + // this is cheap + let owned_lhs = lhs.clone(); + let decompressed_lhs = decompress(owned_lhs)?; + compare(decompressed_lhs, rhs, operator).map(Some) + } +} + +#[cfg(test)] +mod tests { + use vortex_array::array::PrimitiveArray; + use vortex_array::compute::{compare, Operator}; + use vortex_array::validity::Validity; + use vortex_array::IntoArrayVariant; + + use crate::for_compress; + + #[test] + fn test_for_compare() { + let lhs = PrimitiveArray::from_vec(vec![1i32, 2, 3, 4, 5], Validity::AllValid); + let lhs = for_compress(&lhs).unwrap(); + let rhs = PrimitiveArray::from_vec(vec![1i32, 2, 9, 4, 5], Validity::AllValid); + assert_eq!( + compare(lhs, rhs, Operator::Eq) + .unwrap() + .into_bool() + .unwrap() + .boolean_buffer(), + vec![true, true, false, true, true].into() + ); + } +} diff --git a/encodings/fastlanes/src/for/compute.rs b/encodings/fastlanes/src/for/compute.rs index 5324c58224..77b2472193 100644 --- a/encodings/fastlanes/src/for/compute.rs +++ b/encodings/fastlanes/src/for/compute.rs @@ -2,8 +2,8 @@ use std::ops::AddAssign; use num_traits::{CheckedShl, CheckedShr, WrappingAdd, WrappingSub}; use vortex_array::compute::{ - filter, scalar_at, search_sorted, slice, take, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, - SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions, + filter, scalar_at, search_sorted, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, + ScalarAtFn, SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions, }; use vortex_array::variants::PrimitiveArrayTrait; use vortex_array::{ArrayDType, ArrayData, IntoArrayData}; @@ -14,6 +14,10 @@ use vortex_scalar::{PValue, Scalar}; use crate::{FoRArray, FoREncoding}; impl ComputeVTable for FoREncoding { + fn compare_fn(&self) -> Option<&dyn CompareFn> { + Some(self) + } + fn filter_fn(&self) -> Option<&dyn FilterFn> { Some(self) } diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs index 6292b4638c..d2a04260aa 100644 --- a/encodings/fastlanes/src/for/mod.rs +++ b/encodings/fastlanes/src/for/mod.rs @@ -14,6 +14,7 @@ use vortex_dtype::DType; use vortex_error::{vortex_bail, VortexExpect as _, VortexResult}; use vortex_scalar::{Scalar, ScalarValue}; +mod compare; mod compress; mod compute;