44use fastlanes:: BitPacking ;
55use itertools:: Itertools ;
66use num_traits:: PrimInt ;
7+ use num_traits:: Zero ;
78use vortex_array:: ArrayView ;
89use vortex_array:: ExecutionCtx ;
910use vortex_array:: IntoArray ;
@@ -14,10 +15,12 @@ use vortex_array::buffer::BufferHandle;
1415use vortex_array:: dtype:: IntegerPType ;
1516use vortex_array:: dtype:: NativePType ;
1617use vortex_array:: dtype:: PType ;
18+ use vortex_array:: dtype:: PhysicalPType ;
1719use vortex_array:: match_each_integer_ptype;
1820use vortex_array:: match_each_unsigned_integer_ptype;
1921use vortex_array:: patches:: Patches ;
2022use vortex_array:: validity:: Validity ;
23+ use vortex_buffer:: BitBuffer ;
2124use vortex_buffer:: Buffer ;
2225use vortex_buffer:: BufferMut ;
2326use vortex_buffer:: ByteBuffer ;
@@ -35,15 +38,222 @@ use crate::bitpack_decompress::count_exceptions;
3538use crate :: bitpacking:: array:: ChunkWidths ;
3639use crate :: bitpacking:: array:: chunk_packed_bytes;
3740
38- /// Choose a cost-model width for each chunk, then pack values and gather exceptions.
41+ /// Bit-pack an array choosing the cost-model-optimal width for every 1024-element chunk.
42+ ///
43+ /// Each chunk is charged for its packed block plus the exceptions left behind, so a chunk of small
44+ /// values stays narrow no matter how wide its neighbours are.
45+ ///
46+ /// Every chunk is processed in one go while it sits in L1: histogram, width choice, exception
47+ /// gathering and packing, so the values are streamed from memory once.
3948pub fn bitpack_to_best_chunk_widths (
4049 array : & PrimitiveArray ,
4150 ctx : & mut ExecutionCtx ,
51+ ) -> VortexResult < BitPackedArray > {
52+ ensure_non_negative ( array, ctx) ?;
53+ let validity = array. validity ( ) ?;
54+ let mask = validity. execute_mask ( array. len ( ) , ctx) ?;
55+ let patch_validity = match validity {
56+ Validity :: NonNullable => Validity :: NonNullable ,
57+ _ => Validity :: AllValid ,
58+ } ;
59+
60+ let len = array. len ( ) ;
61+ let ( widths, packed, patches) = match_each_integer_ptype ! ( array. ptype( ) , |T | {
62+ encode_chunks:: <T >( array. as_slice:: <T >( ) , & mask, patch_validity) ?
63+ } ) ;
64+
65+ let offsets = widths. offsets_array ( ) ;
66+ let bitpacked = BitPacked :: try_new (
67+ BufferHandle :: new_host ( packed) ,
68+ array. ptype ( ) ,
69+ validity,
70+ patches,
71+ widths. into_array ( ) ,
72+ offsets,
73+ len,
74+ 0 ,
75+ ) ?;
76+ bitpacked. statistics ( ) . inherit_from ( array. statistics ( ) ) ;
77+ Ok ( bitpacked)
78+ }
79+
80+ /// Multi-pass reference for [`bitpack_to_best_chunk_widths`]: one pass to choose widths, one to
81+ /// pack, and one to gather exceptions. Kept to check the fused encoder against.
82+ #[ cfg( test) ]
83+ pub ( crate ) fn bitpack_to_best_chunk_widths_multipass (
84+ array : & PrimitiveArray ,
85+ ctx : & mut ExecutionCtx ,
4286) -> VortexResult < BitPackedArray > {
4387 let plan = chunk_width_plan ( array. as_view ( ) , ctx) ?;
4488 bitpack_encode_planned ( array, plan, ctx)
4589}
4690
91+ /// Histogram, choose a width, gather exceptions and pack, one 1024-element chunk at a time.
92+ ///
93+ /// Patch indices use the narrowest unsigned type that can address every value.
94+ fn encode_chunks < T > (
95+ values : & [ T ] ,
96+ mask : & Mask ,
97+ patch_validity : Validity ,
98+ ) -> VortexResult < ( ChunkWidths , ByteBuffer , Option < Patches > ) >
99+ where
100+ T : NativePType + PrimInt + PhysicalPType ,
101+ T :: Physical : BitPacking + NativePType ,
102+ {
103+ let len = values. len ( ) ;
104+ if len < u8:: MAX as usize {
105+ encode_chunks_indexed :: < T , u8 > ( values, mask, patch_validity)
106+ } else if len < u16:: MAX as usize {
107+ encode_chunks_indexed :: < T , u16 > ( values, mask, patch_validity)
108+ } else if len < u32:: MAX as usize {
109+ encode_chunks_indexed :: < T , u32 > ( values, mask, patch_validity)
110+ } else {
111+ encode_chunks_indexed :: < T , u64 > ( values, mask, patch_validity)
112+ }
113+ }
114+
115+ fn encode_chunks_indexed < T , P > (
116+ values : & [ T ] ,
117+ mask : & Mask ,
118+ patch_validity : Validity ,
119+ ) -> VortexResult < ( ChunkWidths , ByteBuffer , Option < Patches > ) >
120+ where
121+ T : NativePType + PrimInt + PhysicalPType ,
122+ T :: Physical : BitPacking + NativePType ,
123+ P : IntegerPType ,
124+ {
125+ let bits = T :: PTYPE . bit_width ( ) ;
126+ let bytes_per_exception = bytes_per_exception ( T :: PTYPE ) ;
127+ let num_chunks = values. len ( ) . div_ceil ( FL_CHUNK_SIZE ) ;
128+
129+ let mut widths = BufferMut :: < u8 > :: with_capacity ( num_chunks) ;
130+ let mut chunk_offsets = BufferMut :: < u64 > :: with_capacity ( num_chunks) ;
131+ let mut indices = BufferMut :: < P > :: empty ( ) ;
132+ let mut patch_values = BufferMut :: < T > :: empty ( ) ;
133+
134+ let validity = match mask. bit_buffer ( ) {
135+ AllOr :: All => None ,
136+ AllOr :: Some ( bits) => Some ( bits) ,
137+ // Every value is null: every chunk is zero-width and there is nothing to pack.
138+ AllOr :: None => {
139+ widths. extend_trusted ( std:: iter:: repeat_n ( 0u8 , num_chunks) ) ;
140+ chunk_offsets. extend_trusted ( std:: iter:: repeat_n ( 0u64 , num_chunks) ) ;
141+ return Ok ( ( ChunkWidths :: new ( widths. freeze ( ) ) , ByteBuffer :: empty ( ) , None ) ) ;
142+ }
143+ } ;
144+
145+ // Every chunk packs into a whole block, so the padded size bounds the output; a short trailing
146+ // chunk can pack to more than its raw size. The buffer is shrunk to its exact size at the end.
147+ let mut packed = BufferMut :: < T :: Physical > :: with_capacity ( num_chunks * FL_CHUNK_SIZE ) ;
148+ let mut histogram = vec ! [ 0usize ; bits + 1 ] ;
149+ // Zero-padded copy of the trailing partial chunk.
150+ let mut padded = [ T :: Physical :: zero ( ) ; FL_CHUNK_SIZE ] ;
151+
152+ for ( chunk_idx, chunk) in values. chunks ( FL_CHUNK_SIZE ) . enumerate ( ) {
153+ let base = chunk_idx * FL_CHUNK_SIZE ;
154+ let chunk_validity = validity. map ( |v| v. slice ( base..base + chunk. len ( ) ) ) ;
155+
156+ histogram. fill ( 0 ) ;
157+ for_each_valid_width ( chunk, chunk_validity. as_ref ( ) , |_, _, width| {
158+ histogram[ width] += 1 ;
159+ } ) ;
160+
161+ let bit_width = best_chunk_width ( & histogram, bytes_per_exception) ;
162+ widths. push ( bit_width) ;
163+ chunk_offsets. push ( patch_values. len ( ) as u64 ) ;
164+
165+ // The chunk is still in L1, so a second walk over it is cheaper than remembering widths.
166+ if count_exceptions ( bit_width, & histogram) > 0 {
167+ for_each_valid_width ( chunk, chunk_validity. as_ref ( ) , |i, value, width| {
168+ if width > bit_width as usize {
169+ indices. push ( P :: from ( base + i) . vortex_expect ( "cast index from usize" ) ) ;
170+ patch_values. push ( value) ;
171+ }
172+ } ) ;
173+ }
174+
175+ if bit_width > 0 {
176+ let input: & [ T :: Physical ] = if chunk. len ( ) == FL_CHUNK_SIZE {
177+ as_physical ( chunk)
178+ } else {
179+ padded[ ..chunk. len ( ) ] . copy_from_slice ( as_physical ( chunk) ) ;
180+ & padded
181+ } ;
182+ let packed_len = chunk_packed_bytes ( bit_width) / size_of :: < T :: Physical > ( ) ;
183+ let start = packed. len ( ) ;
184+ // SAFETY: `input` holds exactly 1024 values and the output window is exactly one
185+ // packed block at `bit_width`, within the raw-size capacity reserved above.
186+ unsafe {
187+ packed. set_len ( start + packed_len) ;
188+ BitPacking :: unchecked_pack ( bit_width as usize , input, & mut packed[ start..] ) ;
189+ }
190+ }
191+ }
192+
193+ let packed = if packed. len ( ) < packed. capacity ( ) {
194+ let mut exact = BufferMut :: < T :: Physical > :: with_capacity ( packed. len ( ) ) ;
195+ exact. extend_from_slice ( & packed) ;
196+ exact. freeze ( )
197+ } else {
198+ packed. freeze ( )
199+ } ;
200+
201+ let patches = if indices. is_empty ( ) {
202+ None
203+ } else {
204+ Some ( Patches :: new (
205+ values. len ( ) ,
206+ 0 ,
207+ indices. into_array ( ) ,
208+ PrimitiveArray :: new ( patch_values, patch_validity) . into_array ( ) ,
209+ Some ( chunk_offsets. into_array ( ) ) ,
210+ ) ?)
211+ } ;
212+
213+ Ok ( (
214+ ChunkWidths :: new ( widths. freeze ( ) ) ,
215+ packed. into_byte_buffer ( ) ,
216+ patches,
217+ ) )
218+ }
219+
220+ /// Call `f(index, value, bit_width)` for every value of `chunk`; nulls report a width of zero.
221+ #[ inline]
222+ fn for_each_valid_width < T : NativePType + PrimInt > (
223+ chunk : & [ T ] ,
224+ validity : Option < & BitBuffer > ,
225+ mut f : impl FnMut ( usize , T , usize ) ,
226+ ) {
227+ let bits = T :: PTYPE . bit_width ( ) ;
228+ match validity {
229+ None => {
230+ for ( i, & v) in chunk. iter ( ) . enumerate ( ) {
231+ f ( i, v, bits - PrimInt :: leading_zeros ( v) as usize ) ;
232+ }
233+ }
234+ Some ( validity) => {
235+ for ( ( i, & v) , valid) in chunk. iter ( ) . enumerate ( ) . zip ( validity. iter ( ) ) {
236+ let width = if valid {
237+ bits - PrimInt :: leading_zeros ( v) as usize
238+ } else {
239+ 0
240+ } ;
241+ f ( i, v, width) ;
242+ }
243+ }
244+ }
245+ }
246+
247+ /// View signed or unsigned values as their unsigned physical twin, which FastLanes packs.
248+ fn as_physical < T : PhysicalPType > ( values : & [ T ] ) -> & [ T :: Physical ] {
249+ const {
250+ assert ! ( size_of:: <T >( ) == size_of:: <T :: Physical >( ) ) ;
251+ assert ! ( align_of:: <T >( ) == align_of:: <T :: Physical >( ) ) ;
252+ }
253+ // SAFETY: `Physical` is the same-width unsigned integer, so the layouts match exactly.
254+ unsafe { std:: slice:: from_raw_parts ( values. as_ptr ( ) . cast ( ) , values. len ( ) ) }
255+ }
256+
47257/// The cost-model-optimal bit width of every 1024-element chunk of `array`.
48258pub fn best_chunk_widths (
49259 array : ArrayView < ' _ , Primitive > ,
0 commit comments