perf(encoding): avoid inline bitpacking input copy

rust/lance-encoding/Cargo.toml

@@ -29,6 +29,7 @@ hex = "0.4.3"
 itertools.workspace = true
 log.workspace = true
 num-traits.workspace = true
+object_store.workspace = true
 prost.workspace = true
 hyperloglogplus.workspace = true
 rand.workspace = true

rust/lance-encoding/benches/decoder.rs

@@ -1,14 +1,28 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, hint::black_box, sync::Arc};
 
 use arrow_array::{RecordBatch, UInt32Array};
+#[cfg(feature = "bitpacking")]
+use arrow_buffer::ArrowNativeType;
 use arrow_schema::{DataType, Field, Schema, TimeUnit};
 use arrow_select::take::take;
+#[cfg(feature = "bitpacking")]
+use bytemuck::Pod;
 use criterion::{Criterion, criterion_group, criterion_main};
 use futures::StreamExt;
+#[cfg(feature = "bitpacking")]
+use lance_bitpacking::BitPacking;
 use lance_core::cache::LanceCache;
 use lance_datagen::ArrayGeneratorExt;
+#[cfg(feature = "bitpacking")]
+use lance_encoding::buffer::LanceBuffer;
+#[cfg(feature = "bitpacking")]
+use lance_encoding::compression::BlockDecompressor;
+#[cfg(feature = "bitpacking")]
+use lance_encoding::data::{BlockInfo, DataBlock, FixedWidthDataBlock};
+#[cfg(feature = "bitpacking")]
+use lance_encoding::encodings::physical::bitpacking::InlineBitpacking;
 use lance_encoding::{
     decoder::{
         DecodeBatchScheduler, DecoderConfig, DecoderPlugins, FilterExpression, create_decode_stream,
@@ -49,6 +63,9 @@ const PRIMITIVE_TYPES: &[DataType] = &[
 // schema doesn't yet parse them in the context of a fixed size list.
 const PRIMITIVE_TYPES_FOR_FSL: &[DataType] = &[DataType::Int8, DataType::Float32];
 
+#[cfg(feature = "bitpacking")]
+const BITPACKING_ELEMS_PER_CHUNK: usize = 1024;
+
 fn bench_decode(c: &mut Criterion) {
     let rt = tokio::runtime::Runtime::new().unwrap();
     let mut group = c.benchmark_group("decode_primitive");
@@ -563,20 +580,156 @@ fn bench_decode_compressed_parallel(c: &mut Criterion) {
     }
 }
 
+#[cfg(feature = "bitpacking")]
+fn make_inline_bitpacking_chunk<T>(bit_width: usize) -> LanceBuffer
+where
+    T: ArrowNativeType + BitPacking + Pod,
+{
+    let value_range = 1_usize << bit_width;
+    let values: Vec<T> = (0..BITPACKING_ELEMS_PER_CHUNK)
+        .map(|i| T::from_usize((i * 31 + 7) % value_range).unwrap())
+        .collect();
+    let packed_words = BITPACKING_ELEMS_PER_CHUNK * bit_width / (std::mem::size_of::<T>() * 8);
+
+    let mut chunk = Vec::with_capacity(1 + packed_words);
+    chunk.push(T::from_usize(bit_width).unwrap());
+    let payload_start = chunk.len();
+    chunk.resize(payload_start + packed_words, T::from_usize(0).unwrap());
+    unsafe {
+        BitPacking::unchecked_pack(bit_width, &values, &mut chunk[payload_start..]);
+    }
+
+    LanceBuffer::reinterpret_vec(chunk)
+}
+
+#[cfg(feature = "bitpacking")]
+fn read_little_endian_header<T>(bytes: &[u8]) -> usize {
+    bytes[..std::mem::size_of::<T>()]
+        .iter()
+        .enumerate()
+        .fold(0_u64, |value, (idx, byte)| {
+            value | ((*byte as u64) << (idx * 8))
+        }) as usize
+}
+
+#[cfg(feature = "bitpacking")]
+fn legacy_copy_unchunk<T>(data: LanceBuffer, num_values: u64) -> DataBlock
+where
+    T: ArrowNativeType + BitPacking + Pod,
+{
+    assert!(data.len() >= std::mem::size_of::<T>());
+    assert!(num_values <= BITPACKING_ELEMS_PER_CHUNK as u64);
+
+    let chunk_in_u8 = data.to_vec();
+    let bit_width_value = read_little_endian_header::<T>(&chunk_in_u8);
+    let chunk = bytemuck::cast_slice(&chunk_in_u8[std::mem::size_of::<T>()..]);
+    assert!(std::mem::size_of_val(chunk) == bit_width_value * BITPACKING_ELEMS_PER_CHUNK / 8);
+
+    let mut decompressed = vec![T::from_usize(0).unwrap(); BITPACKING_ELEMS_PER_CHUNK];
+    unsafe {
+        BitPacking::unchecked_unpack(bit_width_value, chunk, &mut decompressed);
+    }
+
+    decompressed.truncate(num_values as usize);
+    DataBlock::FixedWidth(FixedWidthDataBlock {
+        data: LanceBuffer::reinterpret_vec(decompressed),
+        bits_per_value: (std::mem::size_of::<T>() * 8) as u64,
+        num_values,
+        block_info: BlockInfo::new(),
+    })
+}
+
+#[cfg(feature = "bitpacking")]
+fn typed_view_unchunk(buffer: LanceBuffer, uncompressed_bits: u64, num_values: u64) -> DataBlock {
+    InlineBitpacking::new(uncompressed_bits)
+        .decompress(buffer, num_values)
+        .unwrap()
+}
+
+#[cfg(feature = "bitpacking")]
+fn assert_same_fixed_width_payloads(legacy: &DataBlock, typed_view: &DataBlock) {
+    let legacy = legacy.as_fixed_width_ref().unwrap();
+    let typed_view = typed_view.as_fixed_width_ref().unwrap();
+
+    assert_eq!(legacy.num_values, typed_view.num_values);
+    assert_eq!(legacy.bits_per_value, typed_view.bits_per_value);
+    assert_eq!(legacy.data.as_ref(), typed_view.data.as_ref());
+}
+
+#[cfg(feature = "bitpacking")]
+fn bench_inline_bitpacking_case<T>(
+    group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
+    name: &str,
+    bit_width: usize,
+) where
+    T: ArrowNativeType + BitPacking + Pod,
+{
+    let buffer = make_inline_bitpacking_chunk::<T>(bit_width);
+    let compressed_bytes = buffer.len() as u64;
+    let uncompressed_bits = (std::mem::size_of::<T>() * 8) as u64;
+    group.throughput(criterion::Throughput::Bytes(compressed_bytes));
+
+    let legacy = legacy_copy_unchunk::<T>(buffer.clone(), BITPACKING_ELEMS_PER_CHUNK as u64);
+    let typed_view = typed_view_unchunk(
+        buffer.clone(),
+        uncompressed_bits,
+        BITPACKING_ELEMS_PER_CHUNK as u64,
+    );
+    assert_same_fixed_width_payloads(&legacy, &typed_view);
+
+    group.bench_function(format!("{name}/legacy_copy/compressed_bytes"), |b| {
+        b.iter(|| {
+            let decoded = legacy_copy_unchunk::<T>(
+                black_box(buffer.clone()),
+                black_box(BITPACKING_ELEMS_PER_CHUNK as u64),
+            );
+            let fixed = decoded.as_fixed_width().unwrap();
+            black_box(fixed.data.as_ref());
+        })
+    });
+
+    group.bench_function(format!("{name}/typed_view/compressed_bytes"), |b| {
+        b.iter(|| {
+            let decoded = typed_view_unchunk(
+                black_box(buffer.clone()),
+                black_box(uncompressed_bits),
+                black_box(BITPACKING_ELEMS_PER_CHUNK as u64),
+            );
+            let fixed = decoded.as_fixed_width().unwrap();
+            black_box(fixed.data.as_ref());
+        })
+    });
+}
+
+#[cfg(feature = "bitpacking")]
+fn bench_decode_inline_bitpacking_unchunk(c: &mut Criterion) {
+    let mut group = c.benchmark_group("decode_inline_bitpacking_unchunk");
+    bench_inline_bitpacking_case::<u32>(&mut group, "u32_bw12_1024", 12);
+    bench_inline_bitpacking_case::<u64>(&mut group, "u64_bw23_1024", 23);
+    group.finish();
+}
+
+#[cfg(not(feature = "bitpacking"))]
+fn bench_decode_inline_bitpacking_unchunk(c: &mut Criterion) {
+    let mut group = c.benchmark_group("decode_inline_bitpacking_unchunk");
+    group.bench_function("bitpacking_feature_disabled", |b| b.iter(|| black_box(())));
+    group.finish();
+}
+
 #[cfg(target_os = "linux")]
 criterion_group!(
     name=benches;
     config = Criterion::default().significance_level(0.1).sample_size(10)
         .with_profiler(lance_testing::pprof::PProfProfiler::new(100, lance_testing::pprof::Output::Flamegraph(None)));
     targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct,
                 bench_decode_str_with_fixed_size_binary_encoding, bench_decode_compressed,
-                bench_decode_compressed_parallel);
+                bench_decode_compressed_parallel, bench_decode_inline_bitpacking_unchunk);
 
 // Non-linux version does not support pprof.
 #[cfg(not(target_os = "linux"))]
 criterion_group!(
     name=benches;
     config = Criterion::default().significance_level(0.1).sample_size(10);
     targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct,
-                bench_decode_compressed, bench_decode_compressed_parallel);
+                bench_decode_compressed, bench_decode_compressed_parallel, bench_decode_inline_bitpacking_unchunk);
 criterion_main!(benches);