Skip to content

Commit

Permalink
chore: prefer into_buffer_mut() to into_buffer().into_mut()
Browse files Browse the repository at this point in the history
  • Loading branch information
danking committed Jan 2, 2025
1 parent 14f51e8 commit 922aee3
Showing 1 changed file with 13 additions and 19 deletions.
32 changes: 13 additions & 19 deletions encodings/zigzag/src/compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use vortex_array::array::PrimitiveArray;
use vortex_array::validity::Validity;
use vortex_array::variants::PrimitiveArrayTrait;
use vortex_array::IntoArrayData;
use vortex_buffer::Buffer;
use vortex_buffer::BufferMut;
use vortex_dtype::{NativePType, PType};
use vortex_error::{vortex_bail, VortexResult};
use zigzag::ZigZag as ExternalZigZag;
Expand All @@ -12,10 +12,10 @@ use crate::ZigZagArray;
pub fn zigzag_encode(parray: PrimitiveArray) -> VortexResult<ZigZagArray> {
let validity = parray.validity();
let encoded = match parray.ptype() {
PType::I8 => zigzag_encode_primitive::<i8>(parray.into_buffer(), validity),
PType::I16 => zigzag_encode_primitive::<i16>(parray.into_buffer(), validity),
PType::I32 => zigzag_encode_primitive::<i32>(parray.into_buffer(), validity),
PType::I64 => zigzag_encode_primitive::<i64>(parray.into_buffer(), validity),
PType::I8 => zigzag_encode_primitive::<i8>(parray.into_buffer_mut(), validity),
PType::I16 => zigzag_encode_primitive::<i16>(parray.into_buffer_mut(), validity),
PType::I32 => zigzag_encode_primitive::<i32>(parray.into_buffer_mut(), validity),
PType::I64 => zigzag_encode_primitive::<i64>(parray.into_buffer_mut(), validity),
_ => vortex_bail!(
"ZigZag can only encode signed integers, got {}",
parray.ptype()
Expand All @@ -25,25 +25,22 @@ pub fn zigzag_encode(parray: PrimitiveArray) -> VortexResult<ZigZagArray> {
}

fn zigzag_encode_primitive<T: ExternalZigZag + NativePType>(
values: Buffer<T>,
values: BufferMut<T>,
validity: Validity,
) -> PrimitiveArray
where
<T as ExternalZigZag>::UInt: NativePType,
{
PrimitiveArray::new(
values.into_mut().map_each(|v| T::encode(v)).freeze(),
validity,
)
PrimitiveArray::new(values.map_each(|v| T::encode(v)).freeze(), validity)
}

pub fn zigzag_decode(parray: PrimitiveArray) -> VortexResult<PrimitiveArray> {
let validity = parray.validity();
let decoded = match parray.ptype() {
PType::U8 => zigzag_decode_primitive::<i8>(parray.into_buffer(), validity),
PType::U16 => zigzag_decode_primitive::<i16>(parray.into_buffer(), validity),
PType::U32 => zigzag_decode_primitive::<i32>(parray.into_buffer(), validity),
PType::U64 => zigzag_decode_primitive::<i64>(parray.into_buffer(), validity),
PType::U8 => zigzag_decode_primitive::<i8>(parray.into_buffer_mut(), validity),
PType::U16 => zigzag_decode_primitive::<i16>(parray.into_buffer_mut(), validity),
PType::U32 => zigzag_decode_primitive::<i32>(parray.into_buffer_mut(), validity),
PType::U64 => zigzag_decode_primitive::<i64>(parray.into_buffer_mut(), validity),
_ => vortex_bail!(
"ZigZag can only decode unsigned integers, got {}",
parray.ptype()
Expand All @@ -53,16 +50,13 @@ pub fn zigzag_decode(parray: PrimitiveArray) -> VortexResult<PrimitiveArray> {
}

fn zigzag_decode_primitive<T: ExternalZigZag + NativePType>(
values: Buffer<T::UInt>,
values: BufferMut<T::UInt>,
validity: Validity,
) -> PrimitiveArray
where
<T as ExternalZigZag>::UInt: NativePType,
{
PrimitiveArray::new(
values.into_mut().map_each(|v| T::decode(v)).freeze(),
validity,
)
PrimitiveArray::new(values.map_each(|v| T::decode(v)).freeze(), validity)
}

#[cfg(test)]
Expand Down

0 comments on commit 922aee3

Please sign in to comment.