For decompression works on owned arrays

encodings/fastlanes/benches/for_decode.rs

@@ -55,36 +55,56 @@ fn values<T: NativePType + TryFrom<usize>>(len: usize) -> Buffer<T> {
         .collect()
 }
 
-fn for_array<T: NativePType + TryFrom<usize>>(len: usize, bitpacked: bool) -> ArrayRef {
+fn bitpacked_for_array<T: NativePType + TryFrom<usize>>(len: usize) -> ArrayRef {
     let mut ctx = SESSION.create_execution_ctx();
     let array = PrimitiveArray::new(values::<T>(len), Validity::NonNullable);
     let for_array = FoR::encode_chunked(array, &mut ctx).unwrap();
-    if !bitpacked {
-        return for_array.into_array();
-    }
     let packed = BitPacked::encode(for_array.encoded(), BIT_WIDTH, &mut ctx).unwrap();
     FoR::try_new_chunked(packed.into_array(), for_array.references().clone(), 0)
         .unwrap()
         .into_array()
 }
 
-fn run<T: NativePType + TryFrom<usize>>(bencher: Bencher, bytes: usize, bitpacked: bool) {
+/// Each iteration decodes a uniquely owned FoR array over a uniquely owned buffer, so the
+/// references are added in place instead of into a copy.
+#[vortex_bench_support::cpu_features]
+#[divan::bench(types = [i64], args = INPUT_BYTES)]
+fn decode_chunked<T: NativePType + TryFrom<usize>>(bencher: Bencher, bytes: usize) {
     let len = bytes / size_of::<T>();
-    let array = for_array::<T>(len, bitpacked);
+    let mut ctx = SESSION.create_execution_ctx();
+    let array = PrimitiveArray::new(values::<T>(len), Validity::NonNullable);
+    let for_array = FoR::encode_chunked(array, &mut ctx).unwrap();
+    let encoded = for_array
+        .encoded()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)
+        .unwrap();
+    let references = for_array.references().clone();
+
     bencher
         .counter(ItemsCount::new(len))
-        .with_inputs(|| (&array, SESSION.create_execution_ctx()))
-        .bench_refs(|(array, ctx)| (*array).clone().execute::<PrimitiveArray>(ctx).unwrap());
-}
+        .with_inputs(|| {
+            let encoded = PrimitiveArray::new(
+                Buffer::<T>::copy_from(encoded.as_slice::<T>()),
+                Validity::NonNullable,
+            );
+            let array = FoR::try_new_chunked(encoded.into_array(), references.clone(), 0)
+                .unwrap()
+                .into_array();
 
-#[vortex_bench_support::cpu_features]
-#[divan::bench(types = [i64], args = INPUT_BYTES)]
-fn decode_chunked<T: NativePType + TryFrom<usize>>(bencher: Bencher, bytes: usize) {
-    run::<T>(bencher, bytes, false);
+            (array, SESSION.create_execution_ctx())
+        })
+        .bench_values(|(array, mut ctx)| array.execute::<PrimitiveArray>(&mut ctx).unwrap());
 }
 
 #[vortex_bench_support::cpu_features]
 #[divan::bench(types = [u32, i64], args = INPUT_BYTES)]
 fn decode_bitpacked_chunked<T: NativePType + TryFrom<usize>>(bencher: Bencher, bytes: usize) {
-    run::<T>(bencher, bytes, true);
+    let len = bytes / size_of::<T>();
+    let array = bitpacked_for_array::<T>(len);
+
+    bencher
+        .counter(ItemsCount::new(len))
+        .with_inputs(|| (&array, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, ctx)| (*array).clone().execute::<PrimitiveArray>(ctx).unwrap());
 }

encodings/fastlanes/src/for_/array/for_compress.rs

@@ -390,7 +390,7 @@ mod test {
         let expected_unsigned = PrimitiveArray::from_iter(unsigned);
         assert_eq!(encoded.as_slice::<u8>(), expected_unsigned.as_slice::<u8>());
 
-        let decompressed = decompress(&compressed, &mut ctx)?;
+        let decompressed = decompress(compressed.clone(), &mut ctx)?;
         array
             .as_slice::<i8>()
             .iter()

encodings/fastlanes/src/for_/array/for_decompress.rs

@@ -11,8 +11,10 @@ use itertools::Itertools;
 use num_traits::AsPrimitive;
 use num_traits::PrimInt;
 use num_traits::WrappingAdd;
+use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
 use vortex_array::arrays::Primitive;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::builders::PrimitiveBuilder;
@@ -31,83 +33,84 @@ use vortex_compute::lane_kernels::IndexedSourceExt;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
-use vortex_error::vortex_err;
 
 use crate::BitPacked;
 use crate::BitPackedArrayExt;
 use crate::BitWidthsView;
 use crate::FL_CHUNK_SIZE;
 use crate::FoRArray;
-use crate::for_::array::FoRArrayExt;
-use crate::for_::array::FoRArraySlotsExt;
+use crate::for_::array::FoRArrayOwnedExt;
+use crate::for_::array::FoRParts;
 use crate::unpack_iter::for_each_packed_chunk;
 
-pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
-    match array.constant_reference() {
-        Some(reference) => decompress_one_ref(array, &reference, ctx),
-        None => decompress_many_refs(array, ctx),
+pub fn decompress(array: FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
+    let FoRParts {
+        encoded,
+        references,
+        offset,
+    } = array.into_parts();
+
+    match references.as_constant() {
+        Some(reference) => decompress_one_ref(encoded, &reference, ctx),
+        None => decompress_many_refs(encoded, references.downcast::<Primitive>(), offset, ctx),
     }
 }
 
 /// Decompress an array whose chunks all share `reference`.
 fn decompress_one_ref(
-    array: &FoRArray,
+    encoded: ArrayRef,
     reference: &Scalar,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
     // Try to do fused unpack.
-    if let Some(bp) = array.encoded().as_opt::<BitPacked>() {
-        return fused_decompress(array, bp, ctx);
+    match encoded.try_downcast::<BitPacked>() {
+        Ok(bp) => fused_decompress(bp.as_view(), reference, ctx),
+        Err(encoded) => add_reference(encoded.downcast::<Primitive>(), reference, ctx),
     }
-
-    add_reference(array, reference, ctx)
 }
 
 /// Unpack a BitPacked child and add the constant reference in one pass.
 pub(crate) fn fused_decompress(
-    for_: &FoRArray,
     bp: ArrayView<'_, BitPacked>,
+    reference: &Scalar,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    match_each_integer_ptype!(for_.ptype(), |T| {
-        fused_decompress_typed::<T>(for_, bp, ctx)
+    match_each_integer_ptype!(bp.dtype().as_ptype(), |T| {
+        fused_decompress_typed::<T>(bp, reference, ctx)
     })
 }
 
 fn fused_decompress_typed<
     T: PhysicalPType<Physical: FoR + BitPacking> + AsPrimitive<T::Physical> + WrappingAdd,
 >(
-    for_: &FoRArray,
     bp: ArrayView<'_, BitPacked>,
+    reference: &Scalar,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    let ref_ = for_
-        .constant_reference()
-        .ok_or_else(|| vortex_err!("fused FoR decompression requires a constant reference"))?
+    let ref_ = reference
         .as_primitive()
         .as_::<T>()
         .vortex_expect("cannot be null");
 
-    fused_unpack(for_, bp, |_| ref_, ctx)
+    fused_unpack(bp, |_| ref_, ctx)
 }
 
-/// Decode `encoded`, then add `reference` to every value.
+/// Add `reference` to every value of `encoded`.
 fn add_reference(
-    array: &FoRArray,
+    encoded: PrimitiveArray,
     reference: &Scalar,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    match_each_integer_ptype!(array.ptype(), |T| {
-        add_reference_typed::<T>(array, reference, ctx)
+    match_each_integer_ptype!(encoded.ptype(), |T| {
+        add_reference_typed::<T>(encoded, reference, ctx)
     })
 }
 
 fn add_reference_typed<T: NativePType + WrappingAdd + PrimInt>(
-    array: &FoRArray,
+    encoded: PrimitiveArray,
     reference: &Scalar,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    let encoded = array.encoded().as_::<Primitive>().into_owned();
     let min = reference
         .as_primitive()
         .typed_value::<T>()
@@ -152,36 +155,50 @@ fn decompress_primitive<T: NativePType + WrappingAdd + PrimInt>(
 }
 
 /// Decompress an array whose chunks have different references.
-fn decompress_many_refs(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
+fn decompress_many_refs(
+    encoded: ArrayRef,
+    references: PrimitiveArray,
+    offset: u16,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
     // Try to do fused unpack. BitPacked chunks line up with FoR chunks when the offsets match.
-    if let Some(bp) = array.encoded().as_opt::<BitPacked>()
-        && bp.offset() == array.offset()
-    {
-        return fused_decompress_many_refs(array, bp, ctx);
-    }
+    let encoded = match encoded.try_downcast::<BitPacked>() {
+        Ok(bp) if bp.offset() == offset => {
+            return fused_decompress_many_refs(bp.as_view(), &references, ctx);
+        }
+        Ok(bp) => bp.into_array(),
+        Err(encoded) => encoded,
+    };
 
-    add_references(array, ctx)
+    add_references(encoded.downcast::<Primitive>(), &references, offset, ctx)
 }
 
-/// Decode `encoded`, then add each chunk's reference in place.
-fn add_references(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
-    match_each_integer_ptype!(array.ptype(), |T| { add_references_typed::<T>(array, ctx) })
+/// Add each chunk's reference to `encoded` in place.
+fn add_references(
+    encoded: PrimitiveArray,
+    references: &PrimitiveArray,
+    offset: u16,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
+    match_each_integer_ptype!(encoded.ptype(), |T| {
+        add_references_typed::<T>(encoded, references, offset, ctx)
+    })
 }
 
 fn add_references_typed<T: NativePType + WrappingAdd + PrimInt>(
-    array: &FoRArray,
+    encoded: PrimitiveArray,
+    references: &PrimitiveArray,
+    offset: u16,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    let encoded = array.encoded().as_::<Primitive>().into_owned();
     if encoded.is_empty() {
         return Ok(encoded);
     }
     let validity = encoded.validity()?;
-    let references = array.references().as_::<Primitive>().into_owned();
     let references = references.as_slice::<T>();
 
     // The first chunk may be partial when the array was sliced.
-    let first_len = (FL_CHUNK_SIZE - usize::from(array.offset())).min(array.len());
+    let first_len = (FL_CHUNK_SIZE - usize::from(offset)).min(encoded.len());
     let values = match encoded.into_buffer::<T>().try_into_mut() {
         // Try to add references in place if we hold only strong reference.
         Ok(mut values) => {
@@ -240,25 +257,24 @@ fn chunks_mut<V>(values: &mut [V], first_len: usize) -> impl Iterator<Item = &mu
 
 /// Unpack each BitPacked chunk and add its chunk's reference in one pass.
 fn fused_decompress_many_refs(
-    for_: &FoRArray,
     bp: ArrayView<'_, BitPacked>,
+    references: &PrimitiveArray,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    match_each_integer_ptype!(for_.ptype(), |T| {
-        fused_decompress_many_refs_typed::<T>(for_, bp, ctx)
+    match_each_integer_ptype!(bp.dtype().as_ptype(), |T| {
+        fused_decompress_many_refs_typed::<T>(bp, references, ctx)
     })
 }
 
 fn fused_decompress_many_refs_typed<
     T: PhysicalPType<Physical: FoR + BitPacking> + AsPrimitive<T::Physical> + WrappingAdd,
 >(
-    for_: &FoRArray,
     bp: ArrayView<'_, BitPacked>,
+    references: &PrimitiveArray,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
-    let references = for_.references().as_::<Primitive>().into_owned();
     let references = references.as_slice::<T>();
-    fused_unpack(for_, bp, |chunk| references[chunk], ctx)
+    fused_unpack(bp, |chunk| references[chunk], ctx)
 }
 
 /// Unpack each BitPacked chunk and add its reference in one pass.
@@ -268,14 +284,13 @@ fn fused_decompress_many_refs_typed<
 fn fused_unpack<
     T: PhysicalPType<Physical: FoR + BitPacking> + AsPrimitive<T::Physical> + WrappingAdd,
 >(
-    for_: &FoRArray,
     bp: ArrayView<'_, BitPacked>,
     chunk_reference: impl Fn(usize) -> T,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
     let len = bp.len();
     let mut builder =
-        PrimitiveBuilder::<T>::with_capacity_in(for_.dtype().nullability(), len, ctx.allocator());
+        PrimitiveBuilder::<T>::with_capacity_in(bp.dtype().nullability(), len, ctx.allocator());
     let mut uninit_range = builder.uninit_range(len);
     unsafe {
         // Append a dense null Mask.
@@ -442,8 +457,7 @@ mod tests {
         let expect = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7 + 10));
         let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7));
         let bp = BitPackedData::encode(&array.into_array(), 2, &mut ctx)?;
-        let compressed = FoR::try_new(bp.clone().into_array(), 10u32.into())?;
-        let decompressed = fused_decompress(&compressed, bp.as_view(), &mut ctx)?;
+        let decompressed = fused_decompress(bp.as_view(), &10u32.into(), &mut ctx)?;
         assert_arrays_eq!(decompressed, expect, &mut ctx);
         Ok(())
     }
@@ -454,8 +468,7 @@ mod tests {
         let expect = PrimitiveArray::from_iter((0i64..1024).map(|x| x % 7 - 1_000_000));
         let array = PrimitiveArray::from_iter((0i64..1024).map(|x| x % 7));
         let bp = BitPackedData::encode(&array.into_array(), 2, &mut ctx)?;
-        let compressed = FoR::try_new(bp.clone().into_array(), (-1_000_000i64).into())?;
-        let decompressed = fused_decompress(&compressed, bp.as_view(), &mut ctx)?;
+        let decompressed = fused_decompress(bp.as_view(), &(-1_000_000i64).into(), &mut ctx)?;
         assert_arrays_eq!(decompressed, expect, &mut ctx);
         Ok(())
     }

encodings/fastlanes/src/for_/array/mod.rs

@@ -4,6 +4,7 @@
 use std::fmt::Display;
 use std::fmt::Formatter;
 
+use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::TypedArrayRef;
 use vortex_array::array_slots;
@@ -60,6 +61,41 @@ pub trait FoRArrayExt: FoRArraySlotsExt {
 
 impl<T: TypedArrayRef<crate::FoR>> FoRArrayExt for T {}
 
+/// The owned children and offset of a FoR array.
+pub(crate) struct FoRParts {
+    pub(crate) encoded: ArrayRef,
+    pub(crate) references: ArrayRef,
+    pub(crate) offset: u16,
+}
+
+pub(crate) trait FoRArrayOwnedExt {
+    /// Take the children of the array without a clone when this handle is the unique owner.
+    fn into_parts(self) -> FoRParts;
+}
+
+impl FoRArrayOwnedExt for Array<crate::FoR> {
+    fn into_parts(self) -> FoRParts {
+        match self.try_into_parts() {
+            Ok(parts) => {
+                let slots = FoRSlots::from_slots(parts.slots);
+                FoRParts {
+                    encoded: slots.encoded,
+                    references: slots.references,
+                    offset: parts.data.offset,
+                }
+            }
+            Err(array) => {
+                let view = FoRSlotsView::from_slots(array.as_ref().slots());
+                FoRParts {
+                    encoded: view.encoded.clone(),
+                    references: view.references.clone(),
+                    offset: array.offset(),
+                }
+            }
+        }
+    }
+}
+
 impl Display for FoRData {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "offset: {}", self.offset)

encodings/fastlanes/src/for_/vtable/mod.rs

@@ -161,7 +161,7 @@ impl VTable for FoR {
         } else {
             require_child!(array, array.encoded(), FoRSlots::ENCODED => Primitive)
         };
-        Ok(ExecutionResult::done(decompress(&array, ctx)?.into_array()))
+        Ok(ExecutionResult::done(decompress(array, ctx)?.into_array()))
     }
 }