Skip to content

Commit 60c53a1

Browse files
committed
feat: select bitpacked widths independently per chunk
Signed-off-by: "Matt Katz" <mhkatz97@gmail.com>
1 parent 899c596 commit 60c53a1

3 files changed

Lines changed: 416 additions & 19 deletions

File tree

‎encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs‎

Lines changed: 123 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,25 @@ use crate::bitpack_decompress::count_exceptions;
3535
use crate::bitpacking::array::ChunkWidths;
3636
use 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.
3957
pub 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.
6280
pub 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+
174262
fn 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+
386494
pub 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

Comments
 (0)