ALP Kernel

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

encodings/alp/src/alp/array.rs

@@ -17,7 +17,6 @@ use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
 use vortex_array::kernel::BindCtx;
 use vortex_array::kernel::KernelRef;
-use vortex_array::kernel::kernel;
 use vortex_array::patches::Patches;
 use vortex_array::patches::PatchesMetadata;
 use vortex_array::serde::ArrayChildren;
@@ -49,7 +48,8 @@ use crate::ALPFloat;
 use crate::alp::Exponents;
 use crate::alp::alp_encode;
 use crate::alp::decompress::decompress_into_array;
-use crate::alp::decompress::decompress_into_vector;
+use crate::alp::kernel::ALPKernel;
+use crate::alp::kernel::PatchKernels;
 use crate::match_each_alp_float_ptype;
 
 vtable!(ALP);
@@ -190,15 +190,15 @@ impl VTable for ALPVTable {
     fn bind_kernel(array: &ALPArray, ctx: &mut BindCtx) -> VortexResult<KernelRef> {
         let encoded = array.encoded().bind_kernel(ctx)?;
         let patches_kernels = if let Some(patches) = array.patches() {
-            Some((
-                patches.indices().bind_kernel(ctx)?,
-                patches.values().bind_kernel(ctx)?,
-                patches
+            Some(PatchKernels {
+                indices: patches.indices().bind_kernel(ctx)?,
+                values: patches.values().bind_kernel(ctx)?,
+                chunk_offsets: patches
                     .chunk_offsets()
                     .as_ref()
                     .map(|co| co.bind_kernel(ctx))
                     .transpose()?,
-            ))
+            })
         } else {
             None
         };
@@ -207,24 +207,12 @@ impl VTable for ALPVTable {
         let exponents = array.exponents();
 
         match_each_alp_float_ptype!(array.dtype().as_ptype(), |T| {
-            Ok(kernel(move || {
-                let encoded_vector = encoded.execute()?;
-                let patches_vectors = match patches_kernels {
-                    Some((idx_kernel, val_kernel, co_kernel)) => Some((
-                        idx_kernel.execute()?,
-                        val_kernel.execute()?,
-                        co_kernel.map(|k| k.execute()).transpose()?,
-                    )),
-                    None => None,
-                };
-
-                decompress_into_vector::<T>(
-                    encoded_vector,
-                    exponents,
-                    patches_vectors,
-                    patches_offset,
-                )
-            }))
+            Ok(Box::new(ALPKernel::<T>::new(
+                exponents,
+                encoded,
+                patches_kernels,
+                patches_offset,
+            )))
         })
     }
 }

encodings/alp/src/alp/decompress.rs

@@ -18,7 +18,9 @@ use vortex_error::VortexResult;
 use vortex_vector::Vector;
 use vortex_vector::VectorMutOps;
 use vortex_vector::VectorOps;
+use vortex_vector::primitive::PVector;
 use vortex_vector::primitive::PVectorMut;
+use vortex_vector::primitive::PrimitiveVector;
 
 use crate::ALPArray;
 use crate::ALPFloat;
@@ -53,23 +55,20 @@ pub fn decompress_into_array(array: ALPArray) -> PrimitiveArray {
 ///
 /// A `Vector` containing the decompressed floating-point values with all patches applied.
 pub fn decompress_into_vector<T: ALPFloat>(
-    encoded_vector: Vector,
+    encoded_vector: PVector<T::ALPInt>,
     exponents: Exponents,
-    patches_vectors: Option<(Vector, Vector, Option<Vector>)>,
+    patches_vectors: Option<(PrimitiveVector, PrimitiveVector, Option<PrimitiveVector>)>,
     patches_offset: usize,
 ) -> VortexResult<Vector> {
-    let encoded_primitive = encoded_vector.into_primitive().into_mut();
-    let (mut alp_buffer, mask) = T::ALPInt::downcast(encoded_primitive).into_parts();
+    // TODO(ngates): we should try_into_mut and fallback to decode _while_ copying.
+    let (mut alp_buffer, mask) = encoded_vector.into_mut().into_parts();
     <T>::decode_slice_inplace(alp_buffer.as_mut_slice(), exponents);
 
     // SAFETY: `Buffer<T::ALPInt> and `BufferMut<T>` have the same layout.
     let mut decoded_buffer: BufferMut<T> = unsafe { transmute(alp_buffer) };
 
     // Apply patches if they exist.
     if let Some((patches_indices, patches_values, _)) = patches_vectors {
-        let patches_indices = patches_indices.into_primitive();
-        let patches_values = patches_values.into_primitive();
-
         let values_buffer = T::downcast(patches_values.into_mut()).into_parts().0;
         let values_slice = values_buffer.as_slice();
         let decoded_slice = decoded_buffer.as_mut_slice();

encodings/alp/src/alp/kernel.rs

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::marker::PhantomData;
+
+use vortex_array::kernel::Kernel;
+use vortex_array::kernel::KernelRef;
+use vortex_array::kernel::PushDownResult;
+use vortex_dtype::PTypeDowncastExt;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
+
+use crate::ALPFloat;
+use crate::Exponents;
+use crate::alp::decompress::decompress_into_vector;
+
+#[derive(Debug)]
+pub(super) struct ALPKernel<A: ALPFloat> {
+    exponents: Exponents,
+    encoded: KernelRef,
+    patches: Option<PatchKernels>,
+    patches_offset: usize,
+    _phantom: PhantomData<A>,
+}
+
+impl<A: ALPFloat> ALPKernel<A> {
+    pub(super) fn new(
+        exponents: Exponents,
+        encoded: KernelRef,
+        patches: Option<PatchKernels>,
+        patches_offset: usize,
+    ) -> Self {
+        Self {
+            exponents,
+            encoded,
+            patches,
+            patches_offset,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct PatchKernels {
+    pub(super) indices: KernelRef,
+    pub(super) values: KernelRef,
+    pub(super) chunk_offsets: Option<KernelRef>,
+}
+
+impl<A: ALPFloat> Kernel for ALPKernel<A> {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        let encoded = self
+            .encoded
+            .execute()?
+            .into_primitive()
+            .downcast::<A::ALPInt>();
+
+        let patches_vectors = match self.patches {
+            None => None,
+            Some(PatchKernels {
+                indices,
+                values,
+                chunk_offsets,
+            }) => {
+                let indices = indices.execute()?.into_primitive();
+                let values = values.execute()?.into_primitive();
+                let chunk_offsets = match chunk_offsets {
+                    None => None,
+                    Some(co) => Some(co.execute()?.into_primitive()),
+                };
+                Some((indices, values, chunk_offsets))
+            }
+        };
+
+        decompress_into_vector::<A>(
+            encoded,
+            self.exponents,
+            patches_vectors,
+            self.patches_offset,
+        )
+    }
+
+    fn push_down_filter(self: Box<Self>, _selection: &Mask) -> VortexResult<PushDownResult> {
+        Ok(PushDownResult::NotPushed(self))
+    }
+}

encodings/alp/src/alp/mod.rs

@@ -17,6 +17,7 @@ mod array;
 mod compress;
 mod compute;
 mod decompress;
+mod kernel;
 mod ops;
 
 #[cfg(test)]