@@ -35,7 +35,25 @@ use crate::bitpack_decompress::count_exceptions;
3535use crate :: bitpacking:: array:: ChunkWidths ;
3636use crate :: bitpacking:: array:: chunk_packed_bytes;
3737
38- /// Encode with caller-supplied chunk widths, gathering exceptions for values that do not fit.
38+ /// Choose a cost-model width for each chunk, then pack values and gather exceptions.
39+ pub fn bitpack_to_best_chunk_widths (
40+ array : & PrimitiveArray ,
41+ ctx : & mut ExecutionCtx ,
42+ ) -> VortexResult < BitPackedArray > {
43+ let plan = chunk_width_plan ( array. as_view ( ) , ctx) ?;
44+ bitpack_encode_planned ( array, plan, ctx)
45+ }
46+
47+ /// The cost-model-optimal bit width of every 1024-element chunk of `array`.
48+ pub fn best_chunk_widths (
49+ array : ArrayView < ' _ , Primitive > ,
50+ ctx : & mut ExecutionCtx ,
51+ ) -> VortexResult < ChunkWidths > {
52+ Ok ( chunk_width_plan ( array, ctx) ?. widths )
53+ }
54+
55+ /// Bit-pack `array` at the given per-chunk widths, gathering values that do not fit their chunk's
56+ /// width into patches.
3957pub fn bitpack_encode_with_widths (
4058 array : & PrimitiveArray ,
4159 widths : ChunkWidths ,
@@ -58,7 +76,7 @@ pub fn bitpack_encode_with_widths(
5876/// Bit-pack `array` at the single best global width chosen by [`find_best_bit_width`].
5977///
6078/// Every chunk shares that width, so the result serializes under the original
61- /// `fastlanes.bitpacked` format.
79+ /// `fastlanes.bitpacked` format. See [`bitpack_to_best_chunk_widths`] for per-chunk widths.
6280pub fn bitpack_to_best_bit_width (
6381 array : & PrimitiveArray ,
6482 ctx : & mut ExecutionCtx ,
@@ -171,6 +189,76 @@ struct ChunkWidthPlan {
171189 num_exceptions : Option < usize > ,
172190}
173191
192+ fn chunk_width_plan (
193+ array : ArrayView < ' _ , Primitive > ,
194+ ctx : & mut ExecutionCtx ,
195+ ) -> VortexResult < ChunkWidthPlan > {
196+ match_each_integer_ptype ! ( array. ptype( ) , |P | {
197+ chunk_width_plan_typed:: <P >( array, ctx)
198+ } )
199+ }
200+
201+ fn chunk_width_plan_typed < T : NativePType + PrimInt > (
202+ array : ArrayView < ' _ , Primitive > ,
203+ ctx : & mut ExecutionCtx ,
204+ ) -> VortexResult < ChunkWidthPlan > {
205+ let bytes_per_exception = bytes_per_exception ( T :: PTYPE ) ;
206+ let values = array. as_slice :: < T > ( ) ;
207+ let num_chunks = values. len ( ) . div_ceil ( FL_CHUNK_SIZE ) ;
208+ let bit_width: fn ( T ) -> usize =
209+ |v : T | ( 8 * size_of :: < T > ( ) ) - ( PrimInt :: leading_zeros ( v) as usize ) ;
210+
211+ let mut widths = BufferMut :: < u8 > :: with_capacity ( num_chunks) ;
212+ let mut num_exceptions = 0usize ;
213+ let mut histogram = vec ! [ 0usize ; size_of:: <T >( ) * 8 + 1 ] ;
214+
215+ // Score one chunk's histogram and reset it for the next chunk.
216+ let mut finish_chunk = |histogram : & mut [ usize ] | -> u8 {
217+ let best = best_chunk_width ( histogram, bytes_per_exception) ;
218+ num_exceptions += count_exceptions ( best, histogram) ;
219+ histogram. fill ( 0 ) ;
220+ best
221+ } ;
222+
223+ match array
224+ . validity ( ) ?
225+ . execute_mask ( array. as_ref ( ) . len ( ) , ctx) ?
226+ . bit_buffer ( )
227+ {
228+ AllOr :: All => {
229+ for chunk in values. chunks ( FL_CHUNK_SIZE ) {
230+ for v in chunk {
231+ histogram[ bit_width ( * v) ] += 1 ;
232+ }
233+ widths. push ( finish_chunk ( & mut histogram) ) ;
234+ }
235+ }
236+ AllOr :: None => {
237+ for _ in 0 ..num_chunks {
238+ widths. push ( 0 ) ;
239+ }
240+ }
241+ AllOr :: Some ( buffer) => {
242+ let mut valid = buffer. iter ( ) ;
243+ for chunk in values. chunks ( FL_CHUNK_SIZE ) {
244+ for v in chunk {
245+ if valid. next ( ) . unwrap_or ( true ) {
246+ histogram[ bit_width ( * v) ] += 1 ;
247+ } else {
248+ histogram[ 0 ] += 1 ;
249+ }
250+ }
251+ widths. push ( finish_chunk ( & mut histogram) ) ;
252+ }
253+ }
254+ }
255+
256+ Ok ( ChunkWidthPlan {
257+ widths : ChunkWidths :: new ( widths. freeze ( ) ) ,
258+ num_exceptions : Some ( num_exceptions) ,
259+ } )
260+ }
261+
174262fn bitpack_encode_planned (
175263 array : & PrimitiveArray ,
176264 plan : ChunkWidthPlan ,
@@ -383,6 +471,26 @@ where
383471 }
384472}
385473
474+ /// The width minimising one chunk's cost: its packed block plus the exceptions left behind.
475+ ///
476+ /// A chunk always occupies a whole `128 * width` byte block, so a partial trailing chunk is
477+ /// charged for its padding.
478+ fn best_chunk_width ( bit_width_freq : & [ usize ] , bytes_per_exception : usize ) -> u8 {
479+ let len: usize = bit_width_freq. iter ( ) . sum ( ) ;
480+ let mut num_packed = 0 ;
481+ let mut best_cost = usize:: MAX ;
482+ let mut best_width = 0 ;
483+ for ( bit_width, freq) in bit_width_freq. iter ( ) . enumerate ( ) {
484+ num_packed += * freq;
485+ let cost = chunk_packed_bytes ( bit_width as u8 ) + ( len - num_packed) * bytes_per_exception;
486+ if cost < best_cost {
487+ best_cost = cost;
488+ best_width = bit_width;
489+ }
490+ }
491+ best_width as u8
492+ }
493+
386494pub fn bit_width_histogram (
387495 array : ArrayView < ' _ , Primitive > ,
388496 ctx : & mut ExecutionCtx ,
@@ -536,6 +644,19 @@ mod tests {
536644 session
537645 } ) ;
538646
647+ #[ test]
648+ fn test_best_chunk_width ( ) {
649+ // 1000 3-bit values and 24 10-bit values in a u16 chunk: 3 bits plus 24 exceptions
650+ // (384 + 24 * 6 bytes) beats 10 bits for everything (1280 bytes).
651+ let mut freq = vec ! [ 0usize ; 17 ] ;
652+ freq[ 3 ] = 1000 ;
653+ freq[ 10 ] = 24 ;
654+ assert_eq ! ( best_chunk_width( & freq, bytes_per_exception( PType :: U16 ) ) , 3 ) ;
655+ // Make the exceptions expensive enough and the wide width wins.
656+ freq[ 10 ] = 200 ;
657+ assert_eq ! ( best_chunk_width( & freq, bytes_per_exception( PType :: U16 ) ) , 10 ) ;
658+ }
659+
539660 #[ test]
540661 fn null_patches ( ) {
541662 let mut ctx = SESSION . create_execution_ctx ( ) ;
0 commit comments