@@ -46,6 +46,17 @@ pub unsafe trait FromBytes: Sized {
4646 fn from_le_bytes ( bs : Self :: Buffer ) -> Self ;
4747}
4848
49+ /// Types that can be decoded from bitpacked representations.
50+ ///
51+ /// This is implemented for primitive types and bool that can be
52+ /// directly converted from a u64 value. Types like Int96, ByteArray,
53+ /// and FixedLenByteArray that cannot be represented in 64 bits do not
54+ /// implement this trait.
55+ pub trait FromBitpacked : FromBytes {
56+ /// Convert directly from a u64 value by truncation, avoiding byte slice copies.
57+ fn from_u64 ( v : u64 ) -> Self ;
58+ }
59+
4960macro_rules! from_le_bytes {
5061 ( $( $ty: ty) ,* ) => {
5162 $(
@@ -60,11 +71,55 @@ macro_rules! from_le_bytes {
6071 <$ty>:: from_le_bytes( bs)
6172 }
6273 }
74+ impl FromBitpacked for $ty {
75+ #[ inline]
76+ fn from_u64( v: u64 ) -> Self {
77+ v as Self
78+ }
79+ }
6380 ) *
6481 } ;
6582}
6683
67- from_le_bytes ! { u8 , u16 , u32 , u64 , i8 , i16 , i32 , i64 , f32 , f64 }
84+ from_le_bytes ! { u8 , u16 , u32 , u64 , i8 , i16 , i32 , i64 }
85+
86+ // SAFETY: all bit patterns are valid for f32 and f64.
87+ unsafe impl FromBytes for f32 {
88+ const BIT_CAPACITY : usize = 32 ;
89+ type Buffer = [ u8 ; 4 ] ;
90+ fn try_from_le_slice ( b : & [ u8 ] ) -> Result < Self > {
91+ Ok ( Self :: from_le_bytes ( array_from_slice ( b) ?) )
92+ }
93+ fn from_le_bytes ( bs : Self :: Buffer ) -> Self {
94+ f32:: from_le_bytes ( bs)
95+ }
96+ }
97+
98+ impl FromBitpacked for f32 {
99+ #[ inline]
100+ fn from_u64 ( v : u64 ) -> Self {
101+ f32:: from_bits ( v as u32 )
102+ }
103+ }
104+
105+ // SAFETY: all bit patterns are valid for f64.
106+ unsafe impl FromBytes for f64 {
107+ const BIT_CAPACITY : usize = 64 ;
108+ type Buffer = [ u8 ; 8 ] ;
109+ fn try_from_le_slice ( b : & [ u8 ] ) -> Result < Self > {
110+ Ok ( Self :: from_le_bytes ( array_from_slice ( b) ?) )
111+ }
112+ fn from_le_bytes ( bs : Self :: Buffer ) -> Self {
113+ f64:: from_le_bytes ( bs)
114+ }
115+ }
116+
117+ impl FromBitpacked for f64 {
118+ #[ inline]
119+ fn from_u64 ( v : u64 ) -> Self {
120+ f64:: from_bits ( v)
121+ }
122+ }
68123
69124// SAFETY: the 0000000x bit pattern is always valid for `bool`.
70125unsafe impl FromBytes for bool {
@@ -79,6 +134,13 @@ unsafe impl FromBytes for bool {
79134 }
80135}
81136
137+ impl FromBitpacked for bool {
138+ #[ inline]
139+ fn from_u64 ( v : u64 ) -> Self {
140+ v != 0
141+ }
142+ }
143+
82144// SAFETY: BIT_CAPACITY is 0.
83145unsafe impl FromBytes for Int96 {
84146 const BIT_CAPACITY : usize = 0 ;
@@ -139,7 +201,7 @@ pub(crate) fn read_num_bytes<T>(size: usize, src: &[u8]) -> T
139201where
140202 T : FromBytes ,
141203{
142- assert ! ( size <= src. len( ) ) ;
204+ debug_assert ! ( size <= src. len( ) ) ;
143205 let mut buffer = <T as FromBytes >:: Buffer :: default ( ) ;
144206 buffer. as_mut ( ) [ ..size] . copy_from_slice ( & src[ ..size] ) ;
145207 <T >:: from_le_bytes ( buffer)
@@ -413,9 +475,9 @@ impl BitReader {
413475 /// Reads a value of type `T` and of size `num_bits`.
414476///
415477/// Returns `None` if there's not enough data available. `Some` otherwise.
416- pub fn get_value < T : FromBytes > ( & mut self , num_bits : usize ) -> Option < T > {
417- assert ! ( num_bits <= 64 ) ;
418- assert ! ( num_bits <= size_of:: <T >( ) * 8 ) ;
478+ pub fn get_value < T : FromBitpacked > ( & mut self , num_bits : usize ) -> Option < T > {
479+ debug_assert ! ( num_bits <= 64 ) ;
480+ debug_assert ! ( num_bits <= size_of:: <T >( ) * 8 ) ;
419481
420482 if self . byte_offset * 8 + self . bit_offset + num_bits > self . buffer . len ( ) * 8 {
421483 return None ;
@@ -445,8 +507,7 @@ impl BitReader {
445507 }
446508 }
447509
448- // TODO: better to avoid copying here
449- T :: try_from_le_slice ( v. as_bytes ( ) ) . ok ( )
510+ Some ( T :: from_u64 ( v) )
450511 }
451512
452513 /// Read multiple values from their packed representation where each element is represented
@@ -457,8 +518,8 @@ impl BitReader {
457518/// This function panics if
458519/// - `num_bits` is larger than the bit-capacity of `T`
459520///
460- pub fn get_batch < T : FromBytes > ( & mut self , batch : & mut [ T ] , num_bits : usize ) -> usize {
461- assert ! ( num_bits <= size_of:: <T >( ) * 8 ) ;
521+ pub fn get_batch < T : FromBitpacked > ( & mut self , batch : & mut [ T ] , num_bits : usize ) -> usize {
522+ debug_assert ! ( num_bits <= size_of:: <T >( ) * 8 ) ;
462523
463524 let mut values_to_read = batch. len ( ) ;
464525 let needed_bits = num_bits * values_to_read;
@@ -602,7 +663,7 @@ impl BitReader {
602663///
603664/// Return the number of values skipped (up to num_values)
604665pub fn skip ( & mut self , num_values : usize , num_bits : usize ) -> usize {
605- assert ! ( num_bits <= 64 ) ;
666+ debug_assert ! ( num_bits <= 64 ) ;
606667
607668 let needed_bits = num_bits * num_values;
608669 let remaining_bits = ( self . buffer . len ( ) - self . byte_offset ) * 8 - self . bit_offset ;
@@ -1024,7 +1085,7 @@ mod tests {
10241085
10251086 fn test_get_batch_helper < T > ( total : usize , num_bits : usize )
10261087 where
1027- T : FromBytes + Default + Clone + Debug + Eq ,
1088+ T : FromBitpacked + Default + Clone + Debug + Eq ,
10281089 {
10291090 assert ! ( num_bits <= 64 ) ;
10301091 let num_bytes = ceil ( num_bits, 8 ) ;
0 commit comments