bitapcked scalar_at

spiraldb · Apr 2, 2024 · 9dcc6c0 · 9dcc6c0
1 parent 936015d
commit 9dcc6c0
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/fastlanez/src/bitpack.rs b/fastlanez/src/bitpack.rs
@@ -1,4 +1,4 @@
-use std::mem::{MaybeUninit, size_of};
+use std::mem::{size_of, MaybeUninit};
 
 use arrayref::{array_mut_ref, array_ref};
 use num_traits::{PrimInt, Unsigned};

diff --git a/fastlanez/src/lib.rs b/fastlanez/src/lib.rs
@@ -18,4 +18,4 @@ pub struct Pred<const B: bool>;
 
 pub trait Satisfied {}
 
-impl Satisfied for Pred<true> {}
+impl Satisfied for Pred<true> {}
diff --git a/vortex-fastlanes/src/bitpacking/compress.rs b/vortex-fastlanes/src/bitpacking/compress.rs
@@ -1,25 +1,25 @@
 use arrayref::array_ref;
 
 use fastlanez::TryBitPack;
-use vortex::array::{Array, ArrayRef};
 use vortex::array::downcast::DowncastArrayBuiltin;
-use vortex::array::IntoArray;
 use vortex::array::primitive::PrimitiveArray;
 use vortex::array::sparse::SparseArray;
+use vortex::array::IntoArray;
+use vortex::array::{Array, ArrayRef};
 use vortex::compress::{CompressConfig, CompressCtx, EncodingCompression};
 use vortex::compute::cast::cast;
 use vortex::compute::flatten::flatten_primitive;
 use vortex::compute::patch::patch;
 use vortex::match_each_integer_ptype;
-use vortex::ptype::{NativePType, PType};
 use vortex::ptype::PType::{I16, I32, I64, I8, U16, U32, U64, U8};
+use vortex::ptype::{NativePType, PType};
 use vortex::scalar::{ListScalarVec, Scalar};
 use vortex::stats::Stat;
 use vortex::validity::ArrayValidity;
 use vortex_error::{vortex_bail, vortex_err, VortexResult};
 
-use crate::{BitPackedArray, BitPackedEncoding};
 use crate::downcast::DowncastFastlanes;
+use crate::{BitPackedArray, BitPackedEncoding};
 
 impl EncodingCompression for BitPackedEncoding {
     fn cost(&self) -> u8 {

diff --git a/vortex-fastlanes/src/bitpacking/compute.rs b/vortex-fastlanes/src/bitpacking/compute.rs
@@ -4,9 +4,9 @@ use vortex::array::primitive::PrimitiveArray;
 use vortex::array::{Array, ArrayRef};
 use vortex::compute::as_contiguous::as_contiguous;
 use vortex::compute::flatten::{flatten_primitive, FlattenFn, FlattenedArray};
+use vortex::compute::scalar_at::ScalarAtFn;
 use vortex::compute::take::{take, TakeFn};
 use vortex::compute::ArrayCompute;
-use vortex::compute::scalar_at::ScalarAtFn;
 use vortex::match_each_integer_ptype;
 use vortex::scalar::Scalar;
 use vortex_error::VortexResult;
@@ -20,6 +20,10 @@ impl ArrayCompute for BitPackedArray {
         Some(self)
     }
 
+    fn scalar_at(&self) -> Option<&dyn ScalarAtFn> {
+        Some(self)
+    }
+
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
@@ -42,11 +46,11 @@ impl TakeFn for BitPackedArray {
         let prim_indices = flatten_primitive(indices)?;
         // Group indices into 1024 chunks and relativise them to the beginning of each chunk
         let relative_indices: Vec<(usize, Vec<u16>)> = match_each_integer_ptype!(prim_indices.ptype(), |$P| {
-            let groupped_indices = prim_indices
+            let grouped_indices = prim_indices
                 .typed_data::<$P>()
                 .iter()
                 .group_by(|idx| (**idx / 1024) as usize);
-            groupped_indices
+            grouped_indices
                 .into_iter()
                 .map(|(k, g)| (k, g.map(|idx| (*idx % 1024) as u16).collect()))
                 .collect()