Skip to content

Commit

Permalink
Fix transmute hack
Browse files Browse the repository at this point in the history
  • Loading branch information
gatesn committed May 11, 2024
1 parent 7c47059 commit 6d1b362
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 26 deletions.
15 changes: 15 additions & 0 deletions vortex-dtype/src/ptype.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,21 @@ macro_rules! match_each_integer_ptype {
})
}

#[macro_export]
macro_rules! match_each_unsigned_integer_ptype {
($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
use $crate::PType;
match $self {
PType::U8 => __with__! { u8 },
PType::U16 => __with__! { u16 },
PType::U32 => __with__! { u32 },
PType::U64 => __with__! { u64 },
_ => panic!("Unsupported ptype {}", $self),
}
})
}

#[macro_export]
macro_rules! match_each_float_ptype {
($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
Expand Down
14 changes: 8 additions & 6 deletions vortex-fastlanes/src/bitpacking/compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,13 @@ use vortex::stats::ArrayStatistics;
use vortex::validity::Validity;
use vortex::{Array, ArrayDType, ArrayDef, ArrayTrait, IntoArray, OwnedArray, ToStatic};
use vortex_dtype::PType::U8;
use vortex_dtype::{match_each_integer_ptype, NativePType, PType};
use vortex_dtype::{
match_each_integer_ptype, match_each_unsigned_integer_ptype, NativePType, PType,
};
use vortex_error::{vortex_bail, vortex_err, VortexResult};
use vortex_scalar::Scalar;

use crate::{match_integers_by_width, BitPackedArray, BitPackedEncoding, OwnedBitPackedArray};
use crate::{BitPackedArray, BitPackedEncoding, OwnedBitPackedArray};

impl EncodingCompression for BitPackedEncoding {
fn cost(&self) -> u8 {
Expand Down Expand Up @@ -118,8 +120,8 @@ pub(crate) fn bitpack_encode(

pub(crate) fn bitpack(parray: &PrimitiveArray, bit_width: usize) -> VortexResult<OwnedArray> {
// We know the min is > 0, so it's safe to re-interpret signed integers as unsigned.
// TODO(ngates): we should implement this using a vortex cast to centralize this hack.
let bytes = match_integers_by_width!(parray.ptype(), |$P| {
let parray = parray.reinterpret_cast(parray.ptype().to_unsigned());
let bytes = match_each_unsigned_integer_ptype!(parray.ptype(), |$P| {
bitpack_primitive(parray.typed_data::<$P>(), bit_width)
});
Ok(PrimitiveArray::from(bytes).into_array())
Expand Down Expand Up @@ -185,7 +187,7 @@ pub fn unpack<'a>(array: BitPackedArray) -> VortexResult<PrimitiveArray<'a>> {
let encoded = cast(&array.packed(), U8.into())?.flatten_primitive()?;
let ptype: PType = array.dtype().try_into()?;

let mut unpacked = match_integers_by_width!(ptype, |$P| {
let mut unpacked = match_each_unsigned_integer_ptype!(ptype, |$P| {
PrimitiveArray::from_vec(
unpack_primitive::<$P>(encoded.typed_data::<u8>(), bit_width, offset, length),
array.validity(),
Expand Down Expand Up @@ -287,7 +289,7 @@ pub(crate) fn unpack_single(array: &BitPackedArray, index: usize) -> VortexResul
let ptype: PType = array.dtype().try_into()?;
let index_in_encoded = index + array.offset();

let scalar: Scalar = match_integers_by_width!(ptype, |$P| {
let scalar: Scalar = match_each_unsigned_integer_ptype!(ptype, |$P| {
unsafe {
unpack_single_primitive::<$P>(encoded.typed_data::<u8>(), bit_width, index_in_encoded).map(|v| v.into())
}
Expand Down
8 changes: 4 additions & 4 deletions vortex-fastlanes/src/bitpacking/compute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@ use vortex::compute::slice::{slice, SliceFn};
use vortex::compute::take::{take, TakeFn};
use vortex::compute::ArrayCompute;
use vortex::{Array, ArrayDType, ArrayTrait, IntoArray, OwnedArray};
use vortex_dtype::{match_each_integer_ptype, NativePType};
use vortex_error::{vortex_bail, vortex_err, VortexResult};
use vortex_dtype::{match_each_integer_ptype, match_each_unsigned_integer_ptype, NativePType};
use vortex_error::{vortex_err, VortexResult};
use vortex_scalar::Scalar;

use crate::bitpacking::compress::unpack_single;
use crate::{match_integers_by_width, unpack_single_primitive, BitPackedArray};
use crate::{unpack_single_primitive, BitPackedArray};

mod slice;

Expand Down Expand Up @@ -66,7 +66,7 @@ impl TakeFn for BitPackedArray<'_> {
}

let indices = indices.clone().flatten_primitive()?;
let taken = match_integers_by_width!(ptype, |$T| {
let taken = match_each_unsigned_integer_ptype!(ptype, |$T| {
PrimitiveArray::from_vec(take_primitive::<$T>(self, &indices)?, taken_validity)
});
Ok(taken.reinterpret_cast(ptype).into_array())
Expand Down
16 changes: 0 additions & 16 deletions vortex-fastlanes/src/bitpacking/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -174,22 +174,6 @@ impl ArrayTrait for BitPackedArray<'_> {
}
}

#[macro_export]
macro_rules! match_integers_by_width {
($self:expr, | $_:tt $enc:ident | $($body:tt)*) => ({
macro_rules! __with__ {( $_ $enc:ident ) => ( $($body)* )}
use vortex_dtype::PType;
use vortex_error::vortex_bail;
match $self {
PType::I8 | PType::U8 => __with__! { u8 },
PType::I16 | PType::U16 => __with__! { u16 },
PType::I32 | PType::U32 => __with__! { u32 },
PType::I64 | PType::U64 => __with__! { u64 },
_ => vortex_bail!("Unsupported ptype {}", $self),
}
})
}

#[cfg(test)]
mod test {
use vortex::array::primitive::PrimitiveArray;
Expand Down

0 comments on commit 6d1b362

Please sign in to comment.