1use log::trace;
5
6use crate::{
7 Result,
8 buffer::LanceBuffer,
9 compression::{
10 BlockCompressor, BlockDecompressor, MiniBlockDecompressor, require_block_payload,
11 },
12 data::DataBlock,
13 encodings::{
14 logical::primitive::miniblock::{
15 MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor,
16 },
17 physical::block::{CompressionConfig, GeneralBufferCompressor},
18 },
19 format::{EncodingUtils, metadata::CompressiveEncoding},
20};
21use lance_core::Error;
22
23#[derive(Debug)]
25pub struct GeneralBlockCompressor {
26 child: Box<dyn BlockCompressor>,
27 compression: CompressionConfig,
28}
29
30impl GeneralBlockCompressor {
31 pub fn new(child: Box<dyn BlockCompressor>, compression: CompressionConfig) -> Self {
32 Self { child, compression }
33 }
34}
35
36impl BlockCompressor for GeneralBlockCompressor {
37 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
38 let (payload, child_encoding) = self.child.compress(data)?;
39 let payload = payload.ok_or_else(|| {
40 Error::invalid_input("General block compression requires a payload-bearing child")
41 })?;
42 let compressor = GeneralBufferCompressor::get_compressor(self.compression)?;
43 let mut compressed = Vec::new();
44 compressor.compress(&payload, &mut compressed)?;
45 Ok((
46 Some(LanceBuffer::from(compressed)),
47 EncodingUtils::wrapped(self.compression, child_encoding)?,
48 ))
49 }
50}
51
52#[derive(Debug)]
54pub(crate) struct FixedWidthGeneralBlockDecompressor {
55 child: Box<dyn BlockDecompressor>,
56 compression: CompressionConfig,
57 bits_per_value: u64,
58}
59
60impl FixedWidthGeneralBlockDecompressor {
61 pub(crate) fn new(
62 child: Box<dyn BlockDecompressor>,
63 compression: CompressionConfig,
64 bits_per_value: u64,
65 ) -> Self {
66 Self {
67 child,
68 compression,
69 bits_per_value,
70 }
71 }
72}
73
74impl BlockDecompressor for FixedWidthGeneralBlockDecompressor {
75 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
76 let data = require_block_payload(data, "General fixed-width block")?;
77 let expected_bytes = usize::try_from(num_values)
78 .ok()
79 .and_then(|num_values| {
80 num_values.checked_mul(usize::try_from(self.bits_per_value / 8).ok()?)
81 })
82 .ok_or_else(|| {
83 Error::invalid_input("General fixed-width output length overflows usize")
84 })?;
85 let compressor = GeneralBufferCompressor::get_compressor(self.compression)?;
86 let mut decompressed = Vec::new();
87 compressor.decompress_exact(&data, &mut decompressed, expected_bytes)?;
88 self.child
89 .decompress(Some(LanceBuffer::from(decompressed)), num_values)
90 }
91}
92
93#[derive(Debug)]
96pub struct GeneralMiniBlockCompressor {
97 inner: Box<dyn MiniBlockCompressor>,
98 compression: CompressionConfig,
99}
100
101impl GeneralMiniBlockCompressor {
102 pub fn new(inner: Box<dyn MiniBlockCompressor>, compression: CompressionConfig) -> Self {
103 Self { inner, compression }
104 }
105}
106
107const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024;
109
110use super::super::logical::primitive::miniblock::MiniBlockChunk;
111
112impl MiniBlockCompressor for GeneralMiniBlockCompressor {
113 fn compress(
114 &self,
115 context: MiniBlockCompressionContext,
116 page: DataBlock,
117 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
118 let (inner_compressed, inner_encoding) = self.inner.compress(context, page)?;
120
121 if inner_compressed.data.is_empty()
124 || inner_compressed.data[0].len() < MIN_BUFFER_SIZE_FOR_COMPRESSION
125 {
126 return Ok((inner_compressed, inner_encoding));
127 }
128
129 let first_buffer = &inner_compressed.data[0];
131 let mut compressed_first_buffer = Vec::new();
132 let mut new_chunks = Vec::with_capacity(inner_compressed.chunks.iter().len());
133 let mut offset = 0usize;
134 let mut total_original_size = 0usize;
135
136 let compressor = GeneralBufferCompressor::get_compressor(self.compression)?;
137
138 for chunk in &inner_compressed.chunks {
139 let chunk_first_buffer_size = chunk.buffer_sizes[0] as usize;
140
141 let chunk_data = &first_buffer.as_ref()[offset..offset + chunk_first_buffer_size];
142 total_original_size += chunk_first_buffer_size;
143
144 let compressed_start = compressed_first_buffer.len();
145 compressor.compress(chunk_data, &mut compressed_first_buffer)?;
146 let compressed_size = compressed_first_buffer.len() - compressed_start;
147
148 let mut new_buffer_sizes = chunk.buffer_sizes.clone();
150 new_buffer_sizes[0] = compressed_size as u32;
151
152 new_chunks.push(MiniBlockChunk {
153 buffer_sizes: new_buffer_sizes,
154 log_num_values: chunk.log_num_values,
155 });
156
157 offset += chunk_first_buffer_size;
158 }
159
160 let compressed_total_size = compressed_first_buffer.len();
162 if compressed_total_size >= total_original_size {
163 return Ok((inner_compressed, inner_encoding));
165 }
166
167 trace!(
168 "First buffer compressed from {} to {} bytes (ratio: {:.2})",
169 total_original_size,
170 compressed_total_size,
171 compressed_total_size as f32 / total_original_size as f32
172 );
173
174 let mut final_buffers = vec![LanceBuffer::from(compressed_first_buffer)];
176 final_buffers.extend(inner_compressed.data.into_iter().skip(1));
177
178 let compressed_result = MiniBlockCompressed {
179 data: final_buffers,
180 chunks: new_chunks,
181 num_values: inner_compressed.num_values,
182 };
183
184 let encoding = EncodingUtils::wrapped(self.compression, inner_encoding)?;
186 Ok((compressed_result, encoding))
187 }
188}
189
190#[derive(Debug)]
193pub struct GeneralMiniBlockDecompressor {
194 inner: Box<dyn MiniBlockDecompressor>,
195 compression: CompressionConfig,
196}
197
198impl GeneralMiniBlockDecompressor {
199 pub fn new(inner: Box<dyn MiniBlockDecompressor>, compression: CompressionConfig) -> Self {
200 Self { inner, compression }
201 }
202}
203
204impl MiniBlockDecompressor for GeneralMiniBlockDecompressor {
205 fn num_buffers(&self) -> usize {
206 self.inner.num_buffers()
207 }
208
209 fn decompress(&self, mut data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
210 if data.len() != self.num_buffers() || data.is_empty() {
211 return Err(Error::invalid_input(format!(
212 "General mini-block has {} buffers, expected {}",
213 data.len(),
214 self.num_buffers()
215 )));
216 }
217 let mut decompressed_buffer = Vec::new();
218
219 let decompressor = GeneralBufferCompressor::get_compressor(self.compression)?;
220 decompressor.decompress(&data[0], &mut decompressed_buffer)?;
221 data[0] = LanceBuffer::from(decompressed_buffer);
222
223 self.inner.decompress(data, num_values)
224 }
225
226 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
227 self.inner.decoded_size_bytes(num_values)
228 }
229}
230
231#[cfg(test)]
232mod tests {
233 use super::*;
234 use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy};
235 use crate::data::{BlockInfo, FixedWidthDataBlock};
236 use crate::encodings::physical::block::CompressionScheme;
237 use crate::encodings::physical::rle::RleEncoder;
238 use crate::encodings::physical::value::ValueEncoder;
239 use crate::format::metadata;
240 use crate::format::metadata::compressive_encoding::Compression;
241 use arrow_array::{Float64Array, Int32Array};
242
243 fn miniblock_context() -> MiniBlockCompressionContext {
244 MiniBlockCompressionContext::new(0, true, true)
245 }
246
247 #[derive(Debug)]
248 struct TestCase {
249 name: &'static str,
250 inner_encoder: Box<dyn MiniBlockCompressor>,
251 compression: CompressionConfig,
252 data: DataBlock,
253 expected_compressed: bool, min_compression_ratio: f32, }
256
257 fn create_test_cases() -> Vec<TestCase> {
258 vec![
259 TestCase {
261 name: "small_rle_data",
262 inner_encoder: Box::new(RleEncoder::new()),
263 compression: CompressionConfig {
264 scheme: CompressionScheme::Lz4,
265 level: None,
266 },
267 data: create_repeated_i32_block(vec![1, 1, 1, 1, 2, 2, 2, 2]),
268 expected_compressed: false,
269 min_compression_ratio: 1.0,
270 },
271 TestCase {
273 name: "large_rle_lz4",
274 inner_encoder: Box::new(RleEncoder::new()),
275 compression: CompressionConfig {
276 scheme: CompressionScheme::Lz4,
277 level: None,
278 },
279 data: create_pattern_i32_block(2048, |i| (i / 8) as i32),
280 expected_compressed: false, min_compression_ratio: 1.0,
282 },
283 TestCase {
285 name: "large_rle_zstd",
286 inner_encoder: Box::new(RleEncoder::new()),
287 compression: CompressionConfig {
288 scheme: CompressionScheme::Zstd,
289 level: Some(3),
290 },
291 data: create_pattern_i32_block(8192, |i| (i / 16) as i32),
292 expected_compressed: true, min_compression_ratio: 0.9, },
295 TestCase {
297 name: "sequential_value_lz4",
298 inner_encoder: Box::new(ValueEncoder {}),
299 compression: CompressionConfig {
300 scheme: CompressionScheme::Lz4,
301 level: None,
302 },
303 data: create_pattern_i32_block(1024, |i| i as i32),
304 expected_compressed: false, min_compression_ratio: 1.0,
306 },
307 TestCase {
309 name: "float_value_zstd",
310 inner_encoder: Box::new(ValueEncoder {}),
311 compression: CompressionConfig {
312 scheme: CompressionScheme::Zstd,
313 level: Some(3),
314 },
315 data: create_pattern_f64_block(1024, |i| i as f64 * 0.1),
316 expected_compressed: true,
317 min_compression_ratio: 0.9,
318 },
319 ]
320 }
321
322 fn create_repeated_i32_block(values: Vec<i32>) -> DataBlock {
323 let array = Int32Array::from(values);
324 DataBlock::from_array(array)
325 }
326
327 fn create_pattern_i32_block<F>(size: usize, pattern: F) -> DataBlock
328 where
329 F: Fn(usize) -> i32,
330 {
331 let values: Vec<i32> = (0..size).map(pattern).collect();
332 let array = Int32Array::from(values);
333 DataBlock::from_array(array)
334 }
335
336 fn create_pattern_f64_block<F>(size: usize, pattern: F) -> DataBlock
337 where
338 F: Fn(usize) -> f64,
339 {
340 let values: Vec<f64> = (0..size).map(pattern).collect();
341 let array = Float64Array::from(values);
342 DataBlock::from_array(array)
343 }
344
345 fn run_round_trip_test(test_case: TestCase) {
346 let compressor =
347 GeneralMiniBlockCompressor::new(test_case.inner_encoder, test_case.compression);
348
349 let (compressed, encoding) = compressor
351 .compress(miniblock_context(), test_case.data)
352 .unwrap();
353
354 match &encoding.compression {
356 Some(Compression::General(cm)) => {
357 assert!(
358 test_case.expected_compressed,
359 "{}: Expected compression to be applied",
360 test_case.name
361 );
362 assert_eq!(
363 CompressionScheme::try_from(cm.compression.as_ref().unwrap().scheme()).unwrap(),
364 test_case.compression.scheme
365 );
366 }
367 _ => {
368 if test_case.expected_compressed {
370 match &encoding.compression {
372 Some(Compression::Rle(_)) => {
373 }
375 Some(Compression::Flat(_)) => {
376 }
378 _ => {
379 panic!(
380 "{}: Expected GeneralMiniBlock but got {:?}",
381 test_case.name, encoding.compression
382 );
383 }
384 }
385 }
386 }
387 }
388
389 assert!(
391 !compressed.chunks.is_empty(),
392 "{}: No chunks created",
393 test_case.name
394 );
395
396 let decompressed_data = decompress_miniblock_chunks(&compressed, &encoding);
398
399 let bytes_per_value = if test_case.name.contains("float") {
403 8 } else {
405 4 };
407 let expected_bytes = compressed.num_values as usize * bytes_per_value;
408 assert_eq!(
409 expected_bytes,
410 decompressed_data.len(),
411 "{}: Data size mismatch",
412 test_case.name
413 );
414
415 if test_case.expected_compressed {
417 let compression_ratio = compressed.data[0].len() as f32 / expected_bytes as f32;
418 assert!(
419 compression_ratio <= test_case.min_compression_ratio,
420 "{}: Compression ratio {} > expected {}",
421 test_case.name,
422 compression_ratio,
423 test_case.min_compression_ratio
424 );
425 }
426 }
427
428 fn decompress_miniblock_chunks(
429 compressed: &MiniBlockCompressed,
430 encoding: &CompressiveEncoding,
431 ) -> Vec<u8> {
432 let mut decompressed_data = Vec::new();
433 let mut offsets = vec![0usize; compressed.data.len()]; let decompression_strategy = DefaultDecompressionStrategy::default();
435
436 for chunk in &compressed.chunks {
437 let chunk_values = if chunk.log_num_values > 0 {
438 1u64 << chunk.log_num_values
439 } else {
440 let decompressed_values =
442 decompressed_data.len() as u64 / get_bytes_per_value(compressed) as u64;
443 compressed.num_values.saturating_sub(decompressed_values)
444 };
445
446 let mut chunk_buffers = Vec::new();
448 for (i, &size) in chunk.buffer_sizes.iter().enumerate() {
449 if i < compressed.data.len() {
450 let buffer_data =
451 compressed.data[i].slice_with_length(offsets[i], size as usize);
452 chunk_buffers.push(buffer_data);
453 offsets[i] += size as usize;
454 }
455 }
456
457 let decompressor = decompression_strategy
459 .create_miniblock_decompressor(encoding, &decompression_strategy)
460 .unwrap();
461
462 let chunk_decompressed = decompressor
464 .decompress(chunk_buffers, chunk_values)
465 .unwrap();
466
467 match chunk_decompressed {
468 DataBlock::FixedWidth(ref block) => {
469 decompressed_data.extend_from_slice(block.data.as_ref());
470 }
471 _ => panic!("Expected FixedWidth block"),
472 }
473 }
474
475 decompressed_data
476 }
477
478 fn get_bytes_per_value(compressed: &MiniBlockCompressed) -> usize {
479 if compressed.num_values == 0 {
483 return 4; }
485
486 if compressed.num_values == 1024 {
488 return 8; }
490
491 4 }
493
494 #[test]
495 fn test_compressed_mini_block_table_driven() {
496 for test_case in create_test_cases() {
497 run_round_trip_test(test_case);
498 }
499 }
500
501 #[test]
502 fn test_compressed_mini_block_threshold() {
503 let small_test = TestCase {
505 name: "small_buffer_no_compression",
506 inner_encoder: Box::new(RleEncoder::new()),
507 compression: CompressionConfig {
508 scheme: CompressionScheme::Lz4,
509 level: None,
510 },
511 data: create_repeated_i32_block(vec![1, 1, 2, 2]),
512 expected_compressed: false,
513 min_compression_ratio: 1.0,
514 };
515 run_round_trip_test(small_test);
516 }
517
518 #[test]
519 fn test_compressed_mini_block_with_doubles() {
520 let test_case = TestCase {
524 name: "float_values_with_zstd",
525 inner_encoder: Box::new(ValueEncoder {}),
526 compression: CompressionConfig {
527 scheme: CompressionScheme::Zstd,
528 level: Some(3),
529 },
530 data: create_pattern_f64_block(1024, |i| (i / 10) as f64),
532 expected_compressed: true,
533 min_compression_ratio: 0.5, };
535
536 run_round_trip_test(test_case);
537 }
538
539 #[test]
540 fn test_compressed_mini_block_large_buffers() {
541 let values: Vec<i32> = (0..1024).collect();
544 let data = LanceBuffer::from_bytes(
545 bytemuck::cast_slice(&values).to_vec().into(),
546 std::mem::align_of::<i32>() as u64,
547 );
548 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
549 bits_per_value: 32,
550 data,
551 num_values: 1024,
552 block_info: BlockInfo::new(),
553 });
554
555 let inner = Box::new(ValueEncoder {});
557 let compression = CompressionConfig {
558 scheme: CompressionScheme::Zstd,
559 level: Some(3),
560 };
561 let compressor = GeneralMiniBlockCompressor::new(inner, compression);
562
563 let (compressed, encoding) = compressor.compress(miniblock_context(), block).unwrap();
565
566 match &encoding.compression {
568 Some(Compression::General(cm)) => {
569 assert!(cm.values.is_some());
570 assert_eq!(
571 cm.compression.as_ref().unwrap().scheme(),
572 metadata::CompressionScheme::CompressionAlgorithmZstd
573 );
574 assert_eq!(cm.compression.as_ref().unwrap().level, Some(3));
575
576 match &cm.values.as_ref().unwrap().compression {
578 Some(Compression::Flat(flat)) => {
579 assert_eq!(flat.bits_per_value, 32);
580 }
581 _ => panic!("Expected Flat inner encoding"),
582 }
583 }
584 _ => panic!("Expected GeneralMiniBlock encoding"),
585 }
586
587 assert_eq!(compressed.num_values, 1024);
588 assert_eq!(compressed.data.len(), 1);
590 }
591
592 #[test]
595 fn test_compressed_mini_block_rle_multiple_buffers() {
596 let data = create_repeated_i32_block(vec![1; 100]);
598 let compressor = GeneralMiniBlockCompressor::new(
599 Box::new(RleEncoder::new()),
600 CompressionConfig {
601 scheme: CompressionScheme::Lz4,
602 level: None,
603 },
604 );
605
606 let (compressed, _) = compressor.compress(miniblock_context(), data).unwrap();
607 assert_eq!(compressed.data.len(), 2);
609 }
610
611 #[test]
612 fn test_rle_with_general_miniblock_wrapper() {
613 let test_32 = TestCase {
621 name: "rle_32bit_with_general_wrapper",
622 inner_encoder: Box::new(RleEncoder::new()),
623 compression: CompressionConfig {
624 scheme: CompressionScheme::Lz4,
625 level: None,
626 },
627 data: create_repeated_i32_block(vec![1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]),
628 expected_compressed: false, min_compression_ratio: 1.0,
630 };
631
632 let compressor = GeneralMiniBlockCompressor::new(
635 Box::new(RleEncoder::new()),
636 CompressionConfig {
637 scheme: CompressionScheme::Lz4,
638 level: None,
639 },
640 );
641
642 let (_compressed, encoding) = compressor
643 .compress(miniblock_context(), test_32.data)
644 .unwrap();
645
646 match &encoding.compression {
648 Some(Compression::General(cm)) => {
649 match &cm.values.as_ref().unwrap().compression {
651 Some(Compression::Rle(rle)) => {
652 let Compression::Flat(values) =
653 rle.values.as_ref().unwrap().compression.as_ref().unwrap()
654 else {
655 panic!("Expected flat for RLE values")
656 };
657 let Compression::Flat(run_lengths) = rle
658 .run_lengths
659 .as_ref()
660 .unwrap()
661 .compression
662 .as_ref()
663 .unwrap()
664 else {
665 panic!("Expected flat for RLE run lengths")
666 };
667 assert_eq!(values.bits_per_value, 32);
668 assert_eq!(run_lengths.bits_per_value, 8);
669 }
670 _ => panic!("Expected RLE as inner encoding"),
671 }
672 assert_eq!(
674 cm.compression.as_ref().unwrap().scheme(),
675 metadata::CompressionScheme::CompressionAlgorithmLz4
676 );
677 }
678 Some(Compression::Rle(_)) => {
679 }
681 _ => panic!("Expected GeneralMiniBlock or Rle encoding"),
682 }
683
684 let values_64: Vec<i64> = vec![100i64; 50]
686 .into_iter()
687 .chain(vec![200i64; 50])
688 .chain(vec![300i64; 50])
689 .collect();
690 let array_64 = arrow_array::Int64Array::from(values_64);
691 let block_64 = DataBlock::from_array(array_64);
692
693 let compressor_64 = GeneralMiniBlockCompressor::new(
694 Box::new(RleEncoder::new()),
695 CompressionConfig {
696 scheme: CompressionScheme::Lz4,
697 level: None,
698 },
699 );
700
701 let (_compressed_64, encoding_64) = compressor_64
702 .compress(miniblock_context(), block_64)
703 .unwrap();
704
705 match &encoding_64.compression {
707 Some(Compression::General(cm)) => {
708 match &cm.values.as_ref().unwrap().compression {
710 Some(Compression::Rle(rle)) => {
711 let Compression::Flat(values) =
712 rle.values.as_ref().unwrap().compression.as_ref().unwrap()
713 else {
714 panic!("Expected flat for RLE values")
715 };
716 let Compression::Flat(run_lengths) = rle
717 .run_lengths
718 .as_ref()
719 .unwrap()
720 .compression
721 .as_ref()
722 .unwrap()
723 else {
724 panic!("Expected flat for RLE run lengths")
725 };
726 assert_eq!(values.bits_per_value, 64);
727 assert_eq!(run_lengths.bits_per_value, 8);
728 }
729 _ => panic!("Expected RLE as inner encoding for 64-bit"),
730 }
731 assert_eq!(
733 cm.compression.as_ref().unwrap().scheme(),
734 metadata::CompressionScheme::CompressionAlgorithmLz4
735 );
736 }
737 Some(Compression::Rle(_)) => {
738 }
740 _ => panic!("Expected GeneralMiniBlock or Rle encoding for 64-bit"),
741 }
742 }
743
744 #[test]
745 fn test_compressed_mini_block_empty_data() {
746 let empty_array = Int32Array::from(vec![] as Vec<i32>);
747 let empty_block = DataBlock::from_array(empty_array);
748
749 let compressor = GeneralMiniBlockCompressor::new(
750 Box::new(ValueEncoder {}),
751 CompressionConfig {
752 scheme: CompressionScheme::Lz4,
753 level: None,
754 },
755 );
756
757 let result = compressor.compress(miniblock_context(), empty_block);
758 match result {
759 Ok((compressed, _)) => {
760 assert_eq!(compressed.num_values, 0);
761 }
762 Err(_) => {
763 }
765 }
766 }
767}