Skip to content

Commit 9f9fc0c

Browse files
mhk197mhk197
committed
feat(btrblocks): pick per-chunk bit widths when fastlanes.bitpacked_v2 is permitted
BitPackingScheme chooses a width per 1024-element chunk when the compressor may use the v2 format and one global width otherwise, so an old edition keeps writing the original format and a newer one gets per-chunk widths from the same scheme. The width table child is re-encoded through the cascade. The scheme returns the original array when half or more of the values would be patches. A few wide values otherwise pack at width 0 with nearly every value patched, which beats raw storage by a couple of buffer bytes and loses them back in footer metadata. No shipped edition permits the format yet; the writer test declares its own. Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent 711ca0c commit 9f9fc0c

10 files changed

Lines changed: 311 additions & 66 deletions

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎vortex-btrblocks/src/schemes/integer/bitpacking.rs‎

Lines changed: 86 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,28 @@ use vortex_compressor::scheme::DeferredEstimate;
1717
use vortex_compressor::scheme::EstimateVerdict;
1818
use vortex_error::VortexResult;
1919
use vortex_fastlanes::BitPacked;
20+
use vortex_fastlanes::BitPackedArray;
21+
use vortex_fastlanes::BitPackedArrayExt;
22+
use vortex_fastlanes::BitPackedArraySlotsExt;
23+
use vortex_fastlanes::BitPackedSlots;
2024
use vortex_fastlanes::bitpack_compress::bit_width_histogram;
2125
use vortex_fastlanes::bitpack_compress::bitpack_encode;
26+
use vortex_fastlanes::bitpack_compress::bitpack_to_best_chunk_widths;
2227
use vortex_fastlanes::bitpack_compress::find_best_bit_width;
28+
use vortex_fastlanes::bitpacked_v2_id;
2329

2430
use crate::ArrayAndStats;
2531
use crate::CascadingCompressor;
2632
use crate::CompressorContext;
2733
use crate::Scheme;
34+
use crate::SchemeExt;
2835
use crate::compress_patches;
2936

3037
/// BitPacking encoding for non-negative integers.
38+
///
39+
/// Every 1024-element chunk gets its own bit width when the compressor may use the
40+
/// `fastlanes.bitpacked_v2` format. Otherwise every chunk shares one width, which is the original
41+
/// `fastlanes.bitpacked` format.
3142
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
3243
pub struct BitPackingScheme;
3344

@@ -66,31 +77,51 @@ impl Scheme for BitPackingScheme {
6677

6778
fn compress(
6879
&self,
69-
_compressor: &CascadingCompressor,
80+
compressor: &CascadingCompressor,
7081
data: &ArrayAndStats,
71-
_compress_ctx: CompressorContext,
82+
compress_ctx: CompressorContext,
7283
exec_ctx: &mut ExecutionCtx,
7384
) -> VortexResult<ArrayRef> {
7485
let primitive_array = data.array_as_primitive();
86+
let full_width = primitive_array.ptype().bit_width();
7587

76-
let histogram = bit_width_histogram(primitive_array, exec_ctx)?;
77-
let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
88+
// Per-chunk widths are the newer format of BitPacked. Produce them whenever the writer
89+
// may serialize them; otherwise every chunk shares one width, the original format.
90+
let packed = if compressor.permits_format(bitpacked_v2_id()) {
91+
bitpack_to_best_chunk_widths(&data.array_as_primitive().into_owned(), exec_ctx)?
92+
} else {
93+
let histogram = bit_width_histogram(primitive_array, exec_ctx)?;
94+
let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
95+
if bw as usize == full_width {
96+
return Ok(primitive_array.array().clone());
97+
}
98+
bitpack_encode(
99+
&data.array_as_primitive().into_owned(),
100+
bw,
101+
Some(&histogram),
102+
exec_ctx,
103+
)?
104+
};
78105

79-
// If best bw is determined to be the current bit-width, return the original array.
80-
if bw as usize == primitive_array.ptype().bit_width() {
106+
// If every chunk needs the full bit-width, return the original array.
107+
if packed.chunk_widths().uniform_width().map(usize::from) == Some(full_width) {
108+
return Ok(primitive_array.array().clone());
109+
}
110+
// Mostly patches means the values were never really packed: the array is sparse in all
111+
// but name, and beats raw storage only by the bytes of the few values that did fit.
112+
if packed
113+
.patches()
114+
.is_some_and(|p| p.num_patches() * 2 >= packed.len())
115+
{
81116
return Ok(primitive_array.array().clone());
82117
}
83-
84-
// Otherwise we can bitpack the array.
85-
let primitive_array = primitive_array.into_owned();
86-
let packed = bitpack_encode(&primitive_array, bw, Some(&histogram), exec_ctx)?;
87118

88119
let packed_stats = packed.statistics().to_owned();
89120
let ptype = packed.dtype().as_ptype();
90121
let mut parts = BitPacked::into_parts(packed);
122+
let patches = parts.patches.take();
91123

92-
let array = if use_experimental_patches() {
93-
let patches = parts.patches.take();
124+
if use_experimental_patches() {
94125
// Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray.
95126
let array = BitPacked::try_new(
96127
parts.packed,
@@ -100,36 +131,53 @@ impl Scheme for BitPackingScheme {
100131
parts.widths,
101132
parts.len,
102133
parts.offset,
103-
)?
104-
.into_array();
105-
106-
match patches {
134+
)?;
135+
let array =
136+
compress_width_table(compressor, array, &compress_ctx, exec_ctx)?.into_array();
137+
return Ok(match patches {
107138
None => array,
108139
Some(p) => Patched::from_array_and_patches(array, &p, exec_ctx)?
109140
.with_stats_set(packed_stats)
110141
.into_array(),
111-
}
112-
} else {
113-
// Compress patches and place back into BitPackedArray.
114-
let patches = parts
115-
.patches
116-
.take()
117-
.map(|p| compress_patches(p, exec_ctx))
118-
.transpose()?;
119-
parts.patches = patches;
120-
BitPacked::try_new(
121-
parts.packed,
122-
ptype,
123-
parts.validity,
124-
parts.patches,
125-
parts.widths,
126-
parts.len,
127-
parts.offset,
128-
)?
129-
.with_stats_set(packed_stats)
130-
.into_array()
131-
};
142+
});
143+
}
132144

133-
Ok(array)
145+
// Compress patches and place back into BitPackedArray.
146+
let patches = patches.map(|p| compress_patches(p, exec_ctx)).transpose()?;
147+
let array = BitPacked::try_new(
148+
parts.packed,
149+
ptype,
150+
parts.validity,
151+
patches,
152+
parts.widths,
153+
parts.len,
154+
parts.offset,
155+
)?;
156+
Ok(
157+
compress_width_table(compressor, array, &compress_ctx, exec_ctx)?
158+
.with_stats_set(packed_stats)
159+
.into_array(),
160+
)
134161
}
135162
}
163+
164+
/// Re-encode the width table child through the compressor. Arrays whose chunks share one width
165+
/// have no table.
166+
fn compress_width_table(
167+
compressor: &CascadingCompressor,
168+
packed: BitPackedArray,
169+
compress_ctx: &CompressorContext,
170+
exec_ctx: &mut ExecutionCtx,
171+
) -> VortexResult<BitPackedArray> {
172+
let Some(table) = packed.width_table().cloned() else {
173+
return Ok(packed);
174+
};
175+
let compressed = compressor.compress_child(
176+
&table,
177+
compress_ctx,
178+
BitPackingScheme.id(),
179+
BitPackedSlots::WIDTH_TABLE,
180+
exec_ctx,
181+
)?;
182+
BitPacked::with_width_table(packed, compressed)
183+
}

‎vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use vortex_array::IntoArray;
1313
use vortex_array::VortexSessionExecute;
1414
use vortex_array::arrays::Constant;
1515
use vortex_array::arrays::Dict;
16+
use vortex_array::arrays::Primitive;
1617
use vortex_array::arrays::PrimitiveArray;
1718
use vortex_array::expr::stats::Precision;
1819
use vortex_array::expr::stats::Stat;
@@ -273,3 +274,19 @@ fn has_nested_delta(array: &vortex_array::ArrayRef, under_delta: bool) -> bool {
273274
.iter()
274275
.any(|child| has_nested_delta(child, under_delta || is_delta))
275276
}
277+
278+
/// A few wide values pack at width 0 with nearly every value patched, which beats raw storage
279+
/// only by the bytes of the values that did fit. Such an array stays primitive.
280+
#[test]
281+
fn test_mostly_patched_stays_primitive() -> VortexResult<()> {
282+
let values: Vec<i64> = vec![0, 1 << 40, 1 << 41, 1 << 42, 1 << 43, 1 << 44];
283+
let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable);
284+
let btr = BtrBlocksCompressor::default();
285+
let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?;
286+
assert!(
287+
compressed.is::<Primitive>(),
288+
"got {}",
289+
compressed.encoding_id()
290+
);
291+
Ok(())
292+
}

‎vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs
33
expression: rendered
44
---
55
input: u64, len=16384, nbytes=131072
6-
root: fastlanes.for(u64, len=16384) nbytes=49152
6+
root: fastlanes.for(u64, len=16384) nbytes=47248
77
metadata: reference: 1700000001036u64
8-
encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49152
9-
metadata: bit_width: 24, offset: 0
8+
encoded: fastlanes.bitpacked(u64, len=16384) nbytes=47248
9+
metadata: bit_widths: 16 chunks, max 24, offset: 0
10+
width_table: vortex.primitive(u8, len=16) nbytes=16
11+
metadata: ptype: u8

‎vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs
33
expression: rendered
44
---
55
input: list(i32), len=4066, nbytes=81804
6-
root: vortex.list(list(i32), len=4066) nbytes=11146
6+
root: vortex.list(list(i32), len=4066) nbytes=10944
77
metadata:
88
elements: vortex.runend(i32, len=16384) nbytes=3968
99
metadata: offset: 0
@@ -15,11 +15,13 @@ root: vortex.list(list(i32), len=4066) nbytes=11146
1515
metadata: reference: -49931i32
1616
encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176
1717
metadata: bit_width: 17, offset: 0
18-
offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178
19-
metadata: bit_width: 14, offset: 0
20-
patch_indices: vortex.primitive(u16, len=1) nbytes=2
18+
offsets: fastlanes.bitpacked(u16, len=4067) nbytes=6976
19+
metadata: bit_widths: 4 chunks, max 14, offset: 0
20+
patch_indices: vortex.primitive(u16, len=14) nbytes=28
21+
metadata: ptype: u16
22+
patch_values: vortex.primitive(u16, len=14) nbytes=28
2123
metadata: ptype: u16
22-
patch_values: vortex.constant(u16, len=1) nbytes=4
23-
metadata: scalar: 16384u16
2424
patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4
2525
metadata: ptype: u8
26+
width_table: vortex.primitive(u8, len=4) nbytes=4
27+
metadata: ptype: u8

‎vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,21 @@ source: vortex-btrblocks/tests/golden.rs
33
expression: rendered
44
---
55
input: utf8, len=16384, nbytes=653785
6-
root: vortex.fsst(utf8, len=16384) nbytes=151382
6+
root: vortex.fsst(utf8, len=16384) nbytes=147198
77
metadata: len: 16384, nsymbols: 223
88
uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154
99
metadata: fill_value: 24u8
1010
patch_indices: vortex.primitive(u16, len=1575) nbytes=3150
1111
metadata: ptype: u16
1212
patch_values: vortex.constant(u8, len=1575) nbytes=2
1313
metadata: scalar: 23u8
14-
codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992
15-
metadata: bit_width: 17, offset: 0
14+
codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=32808
15+
metadata: bit_widths: 17 chunks, max 17, offset: 0
16+
patch_indices: vortex.primitive(u16, len=1) nbytes=2
17+
metadata: ptype: u16
18+
patch_values: vortex.constant(u32, len=1) nbytes=4
19+
metadata: scalar: 109229u32
20+
patch_chunk_offsets: vortex.primitive(u8, len=17) nbytes=17
21+
metadata: ptype: u8
22+
width_table: vortex.primitive(u8, len=17) nbytes=17
23+
metadata: ptype: u8

‎vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs
33
expression: rendered
44
---
55
input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072
6-
root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67584
6+
root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=65680
77
metadata:
8-
storage: fastlanes.for(i64, len=16384) nbytes=67584
8+
storage: fastlanes.for(i64, len=16384) nbytes=65680
99
metadata: reference: 1700000000891673i64
10-
encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67584
11-
metadata: bit_width: 33, offset: 0
10+
encoded: fastlanes.bitpacked(i64, len=16384) nbytes=65680
11+
metadata: bit_widths: 16 chunks, max 33, offset: 0
12+
width_table: vortex.primitive(u8, len=16) nbytes=16
13+
metadata: ptype: u8

‎vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,31 +3,35 @@ source: vortex-btrblocks/tests/golden.rs
33
expression: rendered
44
---
55
input: utf8, len=16384, nbytes=653785
6-
root: vortex.fsst(utf8, len=16384) nbytes=121766
6+
root: vortex.fsst(utf8, len=16384) nbytes=121333
77
metadata: len: 16384, nsymbols: 223
8-
uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=1802
8+
uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=1736
99
metadata: fill_value: 24u8
10-
patch_indices: fastlanes.delta(u16, len=1575) nbytes=1798
10+
patch_indices: fastlanes.delta(u16, len=1575) nbytes=1732
1111
metadata: offset: 0
1212
bases: vortex.primitive(u16, len=128) nbytes=256
1313
metadata: ptype: u16
14-
deltas: fastlanes.bitpacked(u16, len=2048) nbytes=1542
15-
metadata: bit_width: 6, offset: 0
16-
patch_indices: vortex.primitive(u16, len=1) nbytes=2
14+
deltas: fastlanes.bitpacked(u16, len=2048) nbytes=1476
15+
metadata: bit_widths: 2 chunks, max 6, offset: 0
16+
patch_indices: vortex.primitive(u16, len=16) nbytes=32
17+
metadata: ptype: u16
18+
patch_values: vortex.primitive(u16, len=16) nbytes=32
1719
metadata: ptype: u16
18-
patch_values: vortex.constant(u16, len=1) nbytes=2
19-
metadata: scalar: 67u16
2020
patch_chunk_offsets: vortex.primitive(u8, len=2) nbytes=2
2121
metadata: ptype: u8
22+
width_table: vortex.primitive(u8, len=2) nbytes=2
23+
metadata: ptype: u8
2224
patch_values: vortex.constant(u8, len=1575) nbytes=2
2325
metadata: scalar: 23u8
24-
codes_offsets: fastlanes.delta(u32, len=16385) nbytes=8728
26+
codes_offsets: fastlanes.delta(u32, len=16385) nbytes=8361
2527
metadata: offset: 0
2628
bases: vortex.primitive(u32, len=544) nbytes=2176
2729
metadata: ptype: u32
28-
deltas: vortex.dict(u32, len=17408) nbytes=6552
30+
deltas: vortex.dict(u32, len=17408) nbytes=6185
2931
metadata: all_values_referenced: true
30-
codes: fastlanes.bitpacked(u8, len=17408) nbytes=6528
31-
metadata: bit_width: 3, offset: 0
32+
codes: fastlanes.bitpacked(u8, len=17408) nbytes=6161
33+
metadata: bit_widths: 17 chunks, max 3, offset: 0
34+
width_table: vortex.primitive(u8, len=17) nbytes=17
35+
metadata: ptype: u8
3236
values: vortex.primitive(u32, len=6) nbytes=24
3337
metadata: ptype: u32

‎vortex/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ vortex-zigzag = { workspace = true }
5757
vortex-zstd = { workspace = true, optional = true }
5858

5959
[dev-dependencies]
60+
parking_lot = { workspace = true }
6061
anyhow = { workspace = true }
6162
arrow-array = { workspace = true }
6263
divan = { workspace = true }

0 commit comments

Comments
 (0)