1use arrow_buffer::{ArrowNativeType, ScalarBuffer};
58use log::trace;
59
60use crate::buffer::LanceBuffer;
61use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor};
62use crate::data::DataBlock;
63use crate::data::{BlockInfo, FixedWidthDataBlock};
64use crate::encodings::logical::primitive::miniblock::{
65 MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed,
66 MiniBlockCompressionContext, MiniBlockCompressor,
67};
68use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor};
69use crate::format::ProtobufUtils21;
70use crate::format::pb21::CompressiveEncoding;
71
72use lance_core::{Error, Result};
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub(crate) enum RunLengthWidth {
77 U8,
79 U16,
81 U32,
83}
84
85impl RunLengthWidth {
86 pub(crate) fn from_bits(bits_per_value: u64) -> Option<Self> {
87 match bits_per_value {
88 8 => Some(Self::U8),
89 16 => Some(Self::U16),
90 32 => Some(Self::U32),
91 _ => None,
92 }
93 }
94
95 pub(crate) fn bits_per_value(self) -> u64 {
96 match self {
97 Self::U8 => 8,
98 Self::U16 => 16,
99 Self::U32 => 32,
100 }
101 }
102
103 fn bytes_per_value(self) -> usize {
104 match self {
105 Self::U8 => 1,
106 Self::U16 => 2,
107 Self::U32 => 4,
108 }
109 }
110
111 fn max_run_length(self) -> u64 {
112 match self {
113 Self::U8 => u8::MAX as u64,
114 Self::U16 => u16::MAX as u64,
115 Self::U32 => u32::MAX as u64,
116 }
117 }
118
119 fn write_length(self, length: u64, dst: &mut Vec<u8>) {
120 match self {
121 Self::U8 => dst.push(length as u8),
122 Self::U16 => dst.extend_from_slice(&(length as u16).to_le_bytes()),
123 Self::U32 => dst.extend_from_slice(&(length as u32).to_le_bytes()),
124 }
125 }
126
127 fn read_length(self, bytes: &[u8]) -> u64 {
128 match self {
129 Self::U8 => bytes[0] as u64,
130 Self::U16 => u16::from_le_bytes([bytes[0], bytes[1]]) as u64,
131 Self::U32 => u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as u64,
132 }
133 }
134}
135
136const RUN_LENGTH_WIDTHS: [RunLengthWidth; 3] =
137 [RunLengthWidth::U8, RunLengthWidth::U16, RunLengthWidth::U32];
138
139pub(crate) fn select_run_length_width_from_entries(
141 entries: &[u64],
142 bits_per_value: u64,
143) -> Result<(RunLengthWidth, u128)> {
144 if entries.len() != RUN_LENGTH_WIDTHS.len() {
145 return Err(Error::invalid_input_source(
146 format!(
147 "RLE run length entry statistics must have {} values, got {}",
148 RUN_LENGTH_WIDTHS.len(),
149 entries.len()
150 )
151 .into(),
152 ));
153 }
154
155 if !matches!(bits_per_value, 8 | 16 | 32 | 64) {
156 return Err(Error::invalid_input_source(
157 format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}")
158 .into(),
159 ));
160 }
161
162 let mut best_width = RUN_LENGTH_WIDTHS[0];
163 let mut best_cost = rle_encoded_size_from_entries(entries[0], bits_per_value, best_width);
164 for (&width, &entry_count) in RUN_LENGTH_WIDTHS.iter().zip(entries.iter()).skip(1) {
165 let cost = rle_encoded_size_from_entries(entry_count, bits_per_value, width);
166 if cost < best_cost {
167 best_width = width;
168 best_cost = cost;
169 }
170 }
171
172 Ok((best_width, best_cost))
173}
174
175pub(crate) fn rle_encoded_size_from_entries(
176 entry_count: u64,
177 bits_per_value: u64,
178 run_length_width: RunLengthWidth,
179) -> u128 {
180 let bytes_per_value = (bits_per_value / 8) as u128;
181 let bytes_per_length = run_length_width.bytes_per_value() as u128;
182 (entry_count as u128) * (bytes_per_value + bytes_per_length)
183}
184
185pub(crate) fn run_length_width_index(run_length_width: RunLengthWidth) -> usize {
186 match run_length_width {
187 RunLengthWidth::U8 => 0,
188 RunLengthWidth::U16 => 1,
189 RunLengthWidth::U32 => 2,
190 }
191}
192
193pub(crate) fn select_run_length_width(
194 data: &LanceBuffer,
195 num_values: u64,
196 bits_per_value: u64,
197 max_segment_values: Option<u64>,
198) -> Result<(RunLengthWidth, u128)> {
199 let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?;
200 select_run_length_width_from_entries(&entries, bits_per_value)
201}
202
203pub(crate) fn rle_encoded_size(
204 data: &LanceBuffer,
205 num_values: u64,
206 bits_per_value: u64,
207 max_segment_values: Option<u64>,
208 run_length_width: RunLengthWidth,
209) -> Result<u128> {
210 let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?;
211 let width_idx = run_length_width_index(run_length_width);
212 Ok(rle_encoded_size_from_entries(
213 entries[width_idx],
214 bits_per_value,
215 run_length_width,
216 ))
217}
218
219fn collect_run_length_entries(
220 data: &LanceBuffer,
221 num_values: u64,
222 bits_per_value: u64,
223 max_segment_values: Option<u64>,
224) -> Result<[u64; 3]> {
225 let num_values = usize::try_from(num_values).map_err(|_| {
226 Error::invalid_input_source(
227 format!("RLE num_values does not fit in usize: {num_values}").into(),
228 )
229 })?;
230
231 macro_rules! collect_entries {
232 ($ty:ty) => {{
233 let type_size = std::mem::size_of::<$ty>();
234 let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| {
235 Error::invalid_input_source(
236 format!(
237 "RLE input byte length overflow: {num_values} values of {type_size} bytes"
238 )
239 .into(),
240 )
241 })?;
242 if data.len() != expected_bytes {
243 return Err(Error::invalid_input_source(
244 format!(
245 "RLE input data size mismatch: {} bytes for {} values of {} bytes",
246 data.len(),
247 num_values,
248 type_size
249 )
250 .into(),
251 ));
252 }
253 let values = data.borrow_to_typed_slice::<$ty>();
254 let values = values.get(..num_values).ok_or_else(|| {
255 Error::invalid_input_source(
256 format!(
257 "RLE data has {} values but {} were expected",
258 values.len(),
259 num_values
260 )
261 .into(),
262 )
263 })?;
264 Ok(collect_run_length_entries_from_slice(
265 values,
266 max_segment_values,
267 ))
268 }};
269 }
270
271 match bits_per_value {
272 8 => collect_entries!(u8),
273 16 => collect_entries!(u16),
274 32 => collect_entries!(u32),
275 64 => collect_entries!(u64),
276 _ => Err(Error::invalid_input_source(
277 format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}")
278 .into(),
279 )),
280 }
281}
282
283fn collect_run_length_entries_from_slice<T: PartialEq + Copy>(
284 values: &[T],
285 max_segment_values: Option<u64>,
286) -> [u64; 3] {
287 if values.is_empty() {
288 return [0; 3];
289 }
290
291 let mut entries = [0u64; 3];
292 let mut prev = values[0];
293 let mut current_length = 1u64;
294
295 for &value in &values[1..] {
296 if value != prev {
297 accumulate_run_length_entries(current_length, max_segment_values, &mut entries);
298 prev = value;
299 current_length = 1;
300 } else {
301 current_length += 1;
302 }
303 }
304 accumulate_run_length_entries(current_length, max_segment_values, &mut entries);
305
306 entries
307}
308
309pub(crate) fn accumulate_run_length_entries(
310 run_length: u64,
311 max_segment_values: Option<u64>,
312 entries: &mut [u64; 3],
313) {
314 let max_segment_values = max_segment_values.unwrap_or(run_length).max(1);
315 let mut remaining = run_length;
316 while remaining > 0 {
317 let segment = remaining.min(max_segment_values);
318 for (idx, width) in RUN_LENGTH_WIDTHS.iter().enumerate() {
319 let entry_count = segment.div_ceil(width.max_run_length());
320 entries[idx] = entries[idx].saturating_add(entry_count);
321 }
322 remaining -= segment;
323 }
324}
325
326#[derive(Debug)]
328pub struct RleEncoder {
329 run_length_width: RunLengthWidth,
330 values_compression: Option<CompressionConfig>,
331 run_lengths_compression: Option<CompressionConfig>,
332 use_child_bitpacking: bool,
333}
334
335#[derive(Clone)]
336struct RleChildCandidate {
337 encoding: CompressiveEncoding,
338 data: LanceBuffer,
339 chunk_sizes: Vec<u32>,
340 size: usize,
341 requires_num_values: bool,
342}
343
344impl Default for RleEncoder {
345 fn default() -> Self {
346 Self::new()
347 }
348}
349
350impl RleEncoder {
351 pub fn new() -> Self {
352 Self {
353 run_length_width: RunLengthWidth::U8,
354 values_compression: None,
355 run_lengths_compression: None,
356 use_child_bitpacking: false,
357 }
358 }
359
360 pub(crate) fn with_run_length_width(run_length_width: RunLengthWidth) -> Self {
361 Self {
362 run_length_width,
363 values_compression: None,
364 run_lengths_compression: None,
365 use_child_bitpacking: false,
366 }
367 }
368
369 pub(crate) fn with_child_encoding(
370 run_length_width: RunLengthWidth,
371 values_compression: Option<CompressionConfig>,
372 run_lengths_compression: Option<CompressionConfig>,
373 use_child_bitpacking: bool,
374 ) -> Self {
375 Self {
376 run_length_width,
377 values_compression,
378 run_lengths_compression,
379 use_child_bitpacking,
380 }
381 }
382
383 fn encode_data(
384 &self,
385 data: &LanceBuffer,
386 num_values: u64,
387 bits_per_value: u64,
388 ) -> Result<(Vec<LanceBuffer>, Vec<MiniBlockChunk>)> {
389 if num_values == 0 {
390 return Ok((Vec::new(), Vec::new()));
391 }
392
393 let num_values = usize::try_from(num_values).map_err(|_| {
394 Error::invalid_input_source(
395 format!("RLE num_values does not fit in usize: {num_values}").into(),
396 )
397 })?;
398 let bytes_per_value = (bits_per_value / 8) as usize;
399 let bytes_per_length = self.run_length_width.bytes_per_value();
400
401 let estimated_runs = num_values / 10;
404 let mut all_values = Vec::with_capacity(estimated_runs * bytes_per_value);
405 let mut all_lengths = Vec::with_capacity(estimated_runs * bytes_per_length);
406 let mut chunks = Vec::new();
407
408 let mut offset = 0usize;
409 let mut values_remaining = num_values;
410
411 while values_remaining > 0 {
412 let values_start = all_values.len();
413 let lengths_start = all_lengths.len();
414
415 let (_num_runs, values_processed, is_last_chunk) = match bits_per_value {
416 8 => self.encode_chunk_rolling::<u8>(
417 data,
418 offset,
419 values_remaining,
420 &mut all_values,
421 &mut all_lengths,
422 ),
423 16 => self.encode_chunk_rolling::<u16>(
424 data,
425 offset,
426 values_remaining,
427 &mut all_values,
428 &mut all_lengths,
429 ),
430 32 => self.encode_chunk_rolling::<u32>(
431 data,
432 offset,
433 values_remaining,
434 &mut all_values,
435 &mut all_lengths,
436 ),
437 64 => self.encode_chunk_rolling::<u64>(
438 data,
439 offset,
440 values_remaining,
441 &mut all_values,
442 &mut all_lengths,
443 ),
444 _ => {
445 return Err(Error::invalid_input_source(
446 format!(
447 "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}"
448 )
449 .into(),
450 ));
451 }
452 };
453
454 if values_processed == 0 {
455 return Err(Error::internal(format!(
458 "RLE encoder made no progress: values_remaining={values_remaining}, \
459 offset={offset}, data_len={}, bits_per_value={bits_per_value}, \
460 max_miniblock_values={}",
461 data.len(),
462 *MAX_MINIBLOCK_VALUES
463 )));
464 }
465
466 let log_num_values = if is_last_chunk {
467 0
468 } else {
469 assert!(
470 values_processed.is_power_of_two(),
471 "Non-last chunk must have power-of-2 values"
472 );
473 values_processed.ilog2() as u8
474 };
475
476 let values_size = all_values.len() - values_start;
477 let lengths_size = all_lengths.len() - lengths_start;
478
479 let chunk = MiniBlockChunk {
480 buffer_sizes: vec![values_size as u32, lengths_size as u32],
481 log_num_values,
482 };
483
484 chunks.push(chunk);
485
486 offset += values_processed;
487 values_remaining -= values_processed;
488 }
489
490 Ok((
492 vec![
493 LanceBuffer::from(all_values),
494 LanceBuffer::from(all_lengths),
495 ],
496 chunks,
497 ))
498 }
499
500 fn encode_block_data(
501 &self,
502 data: &LanceBuffer,
503 num_values: u64,
504 bits_per_value: u64,
505 ) -> Result<Vec<LanceBuffer>> {
506 match bits_per_value {
507 8 => self.encode_block_data_generic::<u8>(data, num_values),
508 16 => self.encode_block_data_generic::<u16>(data, num_values),
509 32 => self.encode_block_data_generic::<u32>(data, num_values),
510 64 => self.encode_block_data_generic::<u64>(data, num_values),
511 _ => Err(Error::invalid_input_source(
512 format!(
513 "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}"
514 )
515 .into(),
516 )),
517 }
518 }
519
520 fn encode_block_data_generic<T>(
521 &self,
522 data: &LanceBuffer,
523 num_values: u64,
524 ) -> Result<Vec<LanceBuffer>>
525 where
526 T: bytemuck::Pod + PartialEq + Copy + ArrowNativeType,
527 {
528 let num_values = usize::try_from(num_values).map_err(|_| {
529 Error::invalid_input_source(
530 format!("RLE num_values does not fit in usize: {num_values}").into(),
531 )
532 })?;
533 let type_size = std::mem::size_of::<T>();
534 let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| {
535 Error::invalid_input_source(
536 format!("RLE input byte length overflow: {num_values} values of {type_size} bytes")
537 .into(),
538 )
539 })?;
540 if data.len() != expected_bytes {
541 return Err(Error::invalid_input_source(
542 format!(
543 "RLE input data size mismatch: {} bytes for {} values of {} bytes",
544 data.len(),
545 num_values,
546 type_size
547 )
548 .into(),
549 ));
550 }
551 if num_values == 0 {
552 return Ok(vec![LanceBuffer::empty(), LanceBuffer::empty()]);
553 }
554
555 let values_ref = data.borrow_to_typed_slice::<T>();
556 let values = values_ref.as_ref();
557 let estimated_runs = num_values / 10;
558 let mut all_values = Vec::with_capacity(estimated_runs * type_size);
559 let mut all_lengths =
560 Vec::with_capacity(estimated_runs * self.run_length_width.bytes_per_value());
561 self.encode_values(values, &mut all_values, &mut all_lengths);
562 Ok(vec![
563 LanceBuffer::from(all_values),
564 LanceBuffer::from(all_lengths),
565 ])
566 }
567
568 fn encode_chunk_rolling<T>(
570 &self,
571 data: &LanceBuffer,
572 offset: usize,
573 values_remaining: usize,
574 all_values: &mut Vec<u8>,
575 all_lengths: &mut Vec<u8>,
576 ) -> (usize, usize, bool)
577 where
578 T: bytemuck::Pod + PartialEq + Copy + std::fmt::Debug + ArrowNativeType,
579 {
580 let type_size = std::mem::size_of::<T>();
581 let chunk_start = offset * type_size;
582 let max_by_count = *MAX_MINIBLOCK_VALUES as usize;
583 let max_values = values_remaining.min(max_by_count);
584 let chunk_end = chunk_start + max_values * type_size;
585
586 if chunk_start >= data.len() {
587 return (0, 0, false);
588 }
589
590 let chunk_len = chunk_end.min(data.len()) - chunk_start;
591 let chunk_buffer = data.slice_with_length(chunk_start, chunk_len);
592 let typed_data_ref = chunk_buffer.borrow_to_typed_slice::<T>();
593 let typed_data: &[T] = typed_data_ref.as_ref();
594 let max_values = max_values.min(typed_data.len());
595
596 if typed_data.is_empty() {
597 return (0, 0, false);
598 }
599
600 let values_start = all_values.len();
601 let all_remaining_values_fit = values_remaining <= max_by_count;
602 let encoded_size = self.encoded_size(&typed_data[..max_values]);
603 let (values_to_encode, is_last_chunk) = if all_remaining_values_fit
604 && encoded_size <= MAX_MINIBLOCK_BYTES as usize
605 {
606 (max_values, true)
607 } else if let Some(values_to_encode) = self.largest_power_of_two_prefix::<T>(typed_data) {
608 (values_to_encode, false)
609 } else {
610 return (0, 0, false);
611 };
612
613 self.encode_values(&typed_data[..values_to_encode], all_values, all_lengths);
614
615 let num_runs = (all_values.len() - values_start) / type_size;
616 (num_runs, values_to_encode, is_last_chunk)
617 }
618
619 fn largest_power_of_two_prefix<T>(&self, values: &[T]) -> Option<usize>
620 where
621 T: bytemuck::Pod + PartialEq + Copy,
622 {
623 let max_prefix = values.len().min(*MAX_MINIBLOCK_VALUES as usize);
624 let mut prefix = 1usize << max_prefix.ilog2();
625 while prefix > 1 {
626 if self.encoded_size(&values[..prefix]) <= MAX_MINIBLOCK_BYTES as usize {
627 return Some(prefix);
628 }
629 prefix >>= 1;
630 }
631 None
632 }
633
634 fn encoded_size<T>(&self, values: &[T]) -> usize
635 where
636 T: bytemuck::Pod + PartialEq + Copy,
637 {
638 if values.is_empty() {
639 return 0;
640 }
641
642 let mut current_value = values[0];
643 let mut current_length = 1u64;
644 let mut encoded_size = 0usize;
645
646 for &value in values.iter().skip(1) {
647 if value == current_value {
648 current_length += 1;
649 } else {
650 encoded_size += self.run_size::<T>(current_length);
651 current_value = value;
652 current_length = 1;
653 }
654 }
655 encoded_size += self.run_size::<T>(current_length);
656 encoded_size
657 }
658
659 fn run_size<T>(&self, length: u64) -> usize
660 where
661 T: bytemuck::Pod,
662 {
663 let type_size = std::mem::size_of::<T>();
664 let run_chunks = length.div_ceil(self.run_length_width.max_run_length()) as usize;
665 run_chunks * (type_size + self.run_length_width.bytes_per_value())
666 }
667
668 fn encode_values<T>(&self, values: &[T], all_values: &mut Vec<u8>, all_lengths: &mut Vec<u8>)
669 where
670 T: bytemuck::Pod + PartialEq + Copy,
671 {
672 if values.is_empty() {
673 return;
674 }
675
676 let mut current_value = values[0];
677 let mut current_length = 1u64;
678
679 for &value in values.iter().skip(1) {
680 if value == current_value {
681 current_length += 1;
682 } else {
683 self.add_run(¤t_value, current_length, all_values, all_lengths);
684 current_value = value;
685 current_length = 1;
686 }
687 }
688 self.add_run(¤t_value, current_length, all_values, all_lengths);
689 }
690
691 fn add_run<T>(
692 &self,
693 value: &T,
694 length: u64,
695 all_values: &mut Vec<u8>,
696 all_lengths: &mut Vec<u8>,
697 ) -> usize
698 where
699 T: bytemuck::Pod,
700 {
701 let value_bytes = bytemuck::bytes_of(value);
702 let type_size = std::mem::size_of::<T>();
703 let max_run_length = self.run_length_width.max_run_length();
704 let num_full_chunks = (length / max_run_length) as usize;
705 let remainder = length % max_run_length;
706
707 let total_chunks = num_full_chunks + if remainder > 0 { 1 } else { 0 };
708 all_values.reserve(total_chunks * type_size);
709 all_lengths.reserve(total_chunks * self.run_length_width.bytes_per_value());
710
711 for _ in 0..num_full_chunks {
712 all_values.extend_from_slice(value_bytes);
713 self.run_length_width
714 .write_length(max_run_length, all_lengths);
715 }
716
717 if remainder > 0 {
718 all_values.extend_from_slice(value_bytes);
719 self.run_length_width.write_length(remainder, all_lengths);
720 }
721
722 total_chunks * (type_size + self.run_length_width.bytes_per_value())
723 }
724
725 fn flat_child_candidate(
726 buffers: &[LanceBuffer],
727 chunks: &[MiniBlockChunk],
728 buffer_index: usize,
729 bits_per_value: u64,
730 ) -> RleChildCandidate {
731 RleChildCandidate {
732 encoding: ProtobufUtils21::flat(bits_per_value, None),
733 data: buffers[buffer_index].clone(),
734 chunk_sizes: chunks
735 .iter()
736 .map(|chunk| chunk.buffer_sizes[buffer_index])
737 .collect(),
738 size: buffers[buffer_index].len(),
739 requires_num_values: false,
740 }
741 }
742
743 fn general_child_candidate(
744 buffers: &[LanceBuffer],
745 chunks: &[MiniBlockChunk],
746 buffer_index: usize,
747 bits_per_value: u64,
748 compression: CompressionConfig,
749 ) -> Result<Option<RleChildCandidate>> {
750 if buffers.is_empty() || buffers[buffer_index].is_empty() {
751 return Ok(None);
752 };
753
754 let compressor = GeneralBufferCompressor::get_compressor(compression)?;
755 let original = &buffers[buffer_index];
756 let mut compressed = Vec::new();
757 let mut offset = 0usize;
758 let mut total_original_size = 0usize;
759 let mut compressed_sizes = Vec::with_capacity(chunks.len());
760
761 for chunk in chunks.iter() {
762 let chunk_size = chunk.buffer_sizes[buffer_index] as usize;
763 let end = offset.checked_add(chunk_size).ok_or_else(|| {
764 Error::invalid_input_source("RLE child buffer offset overflow".into())
765 })?;
766 if end > original.len() {
767 return Err(Error::invalid_input_source(
768 format!(
769 "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}",
770 buffer_index,
771 end,
772 original.len()
773 )
774 .into(),
775 ));
776 }
777
778 let start = compressed.len();
779 compressor.compress(&original.as_ref()[offset..end], &mut compressed)?;
780 let compressed_size = compressed.len() - start;
781 let compressed_size = u32::try_from(compressed_size).map_err(|_| {
782 Error::invalid_input_source(
783 format!(
784 "RLE child buffer {} compressed chunk is too large: {} bytes",
785 buffer_index, compressed_size
786 )
787 .into(),
788 )
789 })?;
790 compressed_sizes.push(compressed_size);
791 total_original_size += chunk_size;
792 offset = end;
793 }
794
795 if compressed.len() >= total_original_size {
796 return Ok(None);
797 }
798
799 let encoding =
800 ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(bits_per_value, None))?;
801 Ok(Some(
802 RleChildCandidate {
803 encoding,
804 data: LanceBuffer::from(compressed),
805 chunk_sizes: compressed_sizes,
806 size: 0,
807 requires_num_values: false,
808 }
809 .with_size_from_data(),
810 ))
811 }
812
813 #[cfg(feature = "bitpacking")]
814 fn bitpacked_child_candidate(
815 buffers: &[LanceBuffer],
816 chunks: &[MiniBlockChunk],
817 buffer_index: usize,
818 bits_per_value: u64,
819 ) -> Result<Option<RleChildCandidate>> {
820 let original = &buffers[buffer_index];
821 if original.is_empty() {
822 return Ok(None);
823 }
824 let packed_bits = Self::required_bits(original, bits_per_value)?;
825 if packed_bits >= bits_per_value {
826 return Ok(None);
827 }
828
829 let compressor = crate::encodings::physical::bitpacking::OutOfLineBitpacking::new(
830 packed_bits,
831 bits_per_value,
832 );
833 let mut packed = Vec::new();
834 let mut offset = 0usize;
835 let mut packed_sizes = Vec::with_capacity(chunks.len());
836 let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
837 Error::invalid_input_source(
838 format!("RLE child bit width is too large: {bits_per_value}").into(),
839 )
840 })?;
841
842 for chunk in chunks {
843 let chunk_size = chunk.buffer_sizes[buffer_index] as usize;
844 let end = offset.checked_add(chunk_size).ok_or_else(|| {
845 Error::invalid_input_source("RLE child buffer offset overflow".into())
846 })?;
847 if end > original.len() {
848 return Err(Error::invalid_input_source(
849 format!(
850 "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}",
851 buffer_index,
852 end,
853 original.len()
854 )
855 .into(),
856 ));
857 }
858 if bytes_per_value == 0 || !chunk_size.is_multiple_of(bytes_per_value) {
859 return Err(Error::invalid_input_source(
860 format!(
861 "RLE child buffer {} chunk has invalid size {} for {} bits per value",
862 buffer_index, chunk_size, bits_per_value
863 )
864 .into(),
865 ));
866 }
867
868 let child_values = (chunk_size / bytes_per_value) as u64;
869 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
870 bits_per_value,
871 data: original.slice_with_length(offset, chunk_size),
872 num_values: child_values,
873 block_info: BlockInfo::default(),
874 });
875 let chunk_packed = BlockCompressor::compress(&compressor, block)?;
876 let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| {
877 Error::invalid_input_source(
878 format!(
879 "RLE child buffer {} bitpacked chunk is too large: {} bytes",
880 buffer_index,
881 chunk_packed.len()
882 )
883 .into(),
884 )
885 })?;
886 packed_sizes.push(packed_size);
887 packed.extend_from_slice(chunk_packed.as_ref());
888 offset = end;
889 }
890
891 if packed.len() >= original.len() {
892 return Ok(None);
893 }
894
895 Ok(Some(
896 RleChildCandidate {
897 encoding: ProtobufUtils21::out_of_line_bitpacking(
898 bits_per_value,
899 ProtobufUtils21::flat(packed_bits, None),
900 ),
901 data: LanceBuffer::from(packed),
902 chunk_sizes: packed_sizes,
903 size: 0,
904 requires_num_values: true,
905 }
906 .with_size_from_data(),
907 ))
908 }
909
910 #[cfg(feature = "bitpacking")]
911 fn required_bits(buffer: &LanceBuffer, bits_per_value: u64) -> Result<u64> {
912 let max_value = match bits_per_value {
913 8 => buffer.as_ref().iter().map(|value| *value as u64).max(),
914 16 => buffer
915 .as_ref()
916 .chunks_exact(2)
917 .map(|value| u16::from_le_bytes(value.try_into().unwrap()) as u64)
918 .max(),
919 32 => buffer
920 .as_ref()
921 .chunks_exact(4)
922 .map(|value| u32::from_le_bytes(value.try_into().unwrap()) as u64)
923 .max(),
924 64 => buffer
925 .as_ref()
926 .chunks_exact(8)
927 .map(|value| u64::from_le_bytes(value.try_into().unwrap()))
928 .max(),
929 _ => {
930 return Err(Error::invalid_input_source(
931 format!(
932 "RLE child bitpacking only supports 8, 16, 32, or 64-bit values, got {bits_per_value}"
933 )
934 .into(),
935 ));
936 }
937 }
938 .unwrap_or(0);
939 Ok((u64::BITS - max_value.leading_zeros()).max(1) as u64)
940 }
941
942 fn child_candidates(
943 buffers: &[LanceBuffer],
944 chunks: &[MiniBlockChunk],
945 buffer_index: usize,
946 bits_per_value: u64,
947 compression: Option<CompressionConfig>,
948 use_child_bitpacking: bool,
949 ) -> Result<Vec<RleChildCandidate>> {
950 #[cfg(not(feature = "bitpacking"))]
951 let _ = use_child_bitpacking;
952 let mut candidates = vec![Self::flat_child_candidate(
953 buffers,
954 chunks,
955 buffer_index,
956 bits_per_value,
957 )];
958 if let Some(compression) = compression
959 && let Some(candidate) = Self::general_child_candidate(
960 buffers,
961 chunks,
962 buffer_index,
963 bits_per_value,
964 compression,
965 )?
966 {
967 candidates.push(candidate);
968 }
969 #[cfg(feature = "bitpacking")]
970 {
971 if use_child_bitpacking
972 && let Some(candidate) =
973 Self::bitpacked_child_candidate(buffers, chunks, buffer_index, bits_per_value)?
974 {
975 candidates.push(candidate);
976 }
977 }
978 Ok(candidates)
979 }
980
981 fn select_child_candidates(
982 values: Vec<RleChildCandidate>,
983 run_lengths: Vec<RleChildCandidate>,
984 ) -> (RleChildCandidate, RleChildCandidate) {
985 let mut best: Option<(usize, usize, usize)> = None;
986 for (value_idx, value) in values.iter().enumerate() {
987 for (length_idx, length) in run_lengths.iter().enumerate() {
988 if value.requires_num_values && length.requires_num_values {
989 continue;
990 }
991 let size = value.size + length.size;
992 if best.is_none_or(|(_, _, best_size)| size < best_size) {
993 best = Some((value_idx, length_idx, size));
994 }
995 }
996 }
997 let (value_idx, length_idx, _) =
998 best.expect("flat RLE child candidates should always be selectable");
999 (values[value_idx].clone(), run_lengths[length_idx].clone())
1000 }
1001
1002 pub(crate) fn selected_payload_size(&self, data: &FixedWidthDataBlock) -> Result<u128> {
1003 let (all_buffers, chunks) =
1004 self.encode_data(&data.data, data.num_values, data.bits_per_value)?;
1005 if all_buffers.is_empty() {
1006 return Ok(0);
1007 }
1008
1009 let values_candidates = Self::child_candidates(
1010 &all_buffers,
1011 &chunks,
1012 0,
1013 data.bits_per_value,
1014 self.values_compression,
1015 self.use_child_bitpacking,
1016 )?;
1017 let run_lengths_candidates = Self::child_candidates(
1018 &all_buffers,
1019 &chunks,
1020 1,
1021 self.run_length_width.bits_per_value(),
1022 self.run_lengths_compression,
1023 self.use_child_bitpacking,
1024 )?;
1025 let (values, run_lengths) =
1026 Self::select_child_candidates(values_candidates, run_lengths_candidates);
1027 Ok((values.size as u128).saturating_add(run_lengths.size as u128))
1028 }
1029}
1030
1031impl RleChildCandidate {
1032 fn with_size_from_data(mut self) -> Self {
1033 self.size = self.data.len();
1034 self
1035 }
1036}
1037
1038impl MiniBlockCompressor for RleEncoder {
1039 fn compress(
1040 &self,
1041 _context: MiniBlockCompressionContext,
1042 data: DataBlock,
1043 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1044 match data {
1045 DataBlock::FixedWidth(fixed_width) => {
1046 let num_values = fixed_width.num_values;
1047 let bits_per_value = fixed_width.bits_per_value;
1048
1049 let (all_buffers, chunks) =
1050 self.encode_data(&fixed_width.data, num_values, bits_per_value)?;
1051 if all_buffers.is_empty() {
1052 let compressed = MiniBlockCompressed {
1053 data: all_buffers,
1054 chunks,
1055 num_values,
1056 };
1057 let encoding = ProtobufUtils21::rle(
1058 ProtobufUtils21::flat(bits_per_value, None),
1059 ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None),
1060 );
1061 return Ok((compressed, encoding));
1062 }
1063
1064 let values_candidates = Self::child_candidates(
1065 &all_buffers,
1066 &chunks,
1067 0,
1068 bits_per_value,
1069 self.values_compression,
1070 self.use_child_bitpacking,
1071 )?;
1072 let run_lengths_candidates = Self::child_candidates(
1073 &all_buffers,
1074 &chunks,
1075 1,
1076 self.run_length_width.bits_per_value(),
1077 self.run_lengths_compression,
1078 self.use_child_bitpacking,
1079 )?;
1080 let (values, run_lengths) =
1081 Self::select_child_candidates(values_candidates, run_lengths_candidates);
1082 let chunks = chunks
1083 .into_iter()
1084 .enumerate()
1085 .map(|(idx, chunk)| MiniBlockChunk {
1086 buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]],
1087 log_num_values: chunk.log_num_values,
1088 })
1089 .collect();
1090
1091 let compressed = MiniBlockCompressed {
1092 data: vec![values.data, run_lengths.data],
1093 chunks,
1094 num_values,
1095 };
1096
1097 let encoding = ProtobufUtils21::rle(values.encoding, run_lengths.encoding);
1098
1099 Ok((compressed, encoding))
1100 }
1101 _ => Err(Error::invalid_input_source(
1102 "RLE encoding only supports FixedWidth data blocks".into(),
1103 )),
1104 }
1105 }
1106}
1107
1108impl BlockCompressor for RleEncoder {
1109 fn compress(&self, data: DataBlock) -> Result<LanceBuffer> {
1111 match data {
1112 DataBlock::FixedWidth(fixed_width) => {
1113 let num_values = fixed_width.num_values;
1114 let bits_per_value = fixed_width.bits_per_value;
1115
1116 let all_buffers =
1117 self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?;
1118
1119 let values_size = all_buffers[0].len() as u64;
1120
1121 let mut combined = Vec::new();
1122 combined.extend_from_slice(&values_size.to_le_bytes());
1123 combined.extend_from_slice(&all_buffers[0]);
1124 combined.extend_from_slice(&all_buffers[1]);
1125 Ok(LanceBuffer::from(combined))
1126 }
1127 _ => Err(Error::invalid_input_source(
1128 "RLE encoding only supports FixedWidth data blocks".into(),
1129 )),
1130 }
1131 }
1132}
1133
1134#[derive(Debug)]
1136pub struct RleDecompressor {
1137 bits_per_value: u64,
1138 run_length_width: RunLengthWidth,
1139 values: RleChildDecompressor,
1140 run_lengths: RleChildDecompressor,
1141}
1142
1143#[derive(Debug)]
1144pub(crate) struct RleChildDecompressor {
1145 bits_per_value: u64,
1146 inner: RleChildDecompressorInner,
1147}
1148
1149#[derive(Debug)]
1150enum RleChildDecompressorInner {
1151 Flat,
1152 Block {
1153 decompressor: Box<dyn BlockDecompressor>,
1154 requires_num_values: bool,
1155 },
1156}
1157
1158impl RleChildDecompressor {
1159 pub(crate) fn flat(bits_per_value: u64) -> Self {
1160 Self {
1161 bits_per_value,
1162 inner: RleChildDecompressorInner::Flat,
1163 }
1164 }
1165
1166 pub(crate) fn block(
1167 bits_per_value: u64,
1168 decompressor: Box<dyn BlockDecompressor>,
1169 requires_num_values: bool,
1170 ) -> Self {
1171 Self {
1172 bits_per_value,
1173 inner: RleChildDecompressorInner::Block {
1174 decompressor,
1175 requires_num_values,
1176 },
1177 }
1178 }
1179
1180 pub(crate) fn bits_per_value(&self) -> u64 {
1181 self.bits_per_value
1182 }
1183
1184 pub(crate) fn requires_num_values(&self) -> bool {
1185 match &self.inner {
1186 RleChildDecompressorInner::Flat => false,
1187 RleChildDecompressorInner::Block {
1188 requires_num_values,
1189 ..
1190 } => *requires_num_values,
1191 }
1192 }
1193
1194 pub(crate) fn is_identity(&self) -> bool {
1195 matches!(self.inner, RleChildDecompressorInner::Flat)
1196 }
1197
1198 fn decode(
1199 &self,
1200 data: LanceBuffer,
1201 num_values: Option<u64>,
1202 label: &str,
1203 ) -> Result<LanceBuffer> {
1204 match &self.inner {
1205 RleChildDecompressorInner::Flat => Ok(data),
1206 RleChildDecompressorInner::Block {
1207 decompressor,
1208 requires_num_values,
1209 } => {
1210 let num_values = if *requires_num_values {
1211 num_values.ok_or_else(|| {
1212 Error::invalid_input_source(
1213 format!("RLE {label} child compression requires the run count").into(),
1214 )
1215 })?
1216 } else {
1217 num_values.unwrap_or(0)
1218 };
1219 let decoded = decompressor.decompress(data, num_values)?;
1220 self.extract_fixed_width(decoded, num_values, label)
1221 }
1222 }
1223 }
1224
1225 fn extract_fixed_width(
1226 &self,
1227 data: DataBlock,
1228 expected_num_values: u64,
1229 label: &str,
1230 ) -> Result<LanceBuffer> {
1231 match data {
1232 DataBlock::FixedWidth(block) => {
1233 if block.bits_per_value != self.bits_per_value {
1234 return Err(Error::invalid_input_source(
1235 format!(
1236 "RLE {label} child decoded {}-bit values, expected {}",
1237 block.bits_per_value, self.bits_per_value
1238 )
1239 .into(),
1240 ));
1241 }
1242 if expected_num_values != 0 && block.num_values != expected_num_values {
1243 return Err(Error::invalid_input_source(
1244 format!(
1245 "RLE {label} child decoded {} values, expected {}",
1246 block.num_values, expected_num_values
1247 )
1248 .into(),
1249 ));
1250 }
1251 Ok(block.data)
1252 }
1253 _ => Err(Error::invalid_input_source(
1254 format!("RLE {label} child decoded to a non fixed-width block").into(),
1255 )),
1256 }
1257 }
1258}
1259
1260impl RleDecompressor {
1261 pub fn new(bits_per_value: u64) -> Self {
1262 Self {
1263 bits_per_value,
1264 run_length_width: RunLengthWidth::U8,
1265 values: RleChildDecompressor::flat(bits_per_value),
1266 run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()),
1267 }
1268 }
1269
1270 pub(crate) fn with_run_length_width(
1271 bits_per_value: u64,
1272 run_length_width: RunLengthWidth,
1273 ) -> Self {
1274 Self {
1275 bits_per_value,
1276 run_length_width,
1277 values: RleChildDecompressor::flat(bits_per_value),
1278 run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()),
1279 }
1280 }
1281
1282 pub(crate) fn with_child_decompressors(
1283 bits_per_value: u64,
1284 run_length_width: RunLengthWidth,
1285 values: RleChildDecompressor,
1286 run_lengths: RleChildDecompressor,
1287 ) -> Self {
1288 Self {
1289 bits_per_value,
1290 run_length_width,
1291 values,
1292 run_lengths,
1293 }
1294 }
1295
1296 fn decode_data(
1297 &self,
1298 data: Vec<LanceBuffer>,
1299 num_values: u64,
1300 clamp_overflow: bool,
1301 ) -> Result<DataBlock> {
1302 if num_values == 0 {
1303 return Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1304 bits_per_value: self.bits_per_value,
1305 data: LanceBuffer::from(vec![]),
1306 num_values: 0,
1307 block_info: BlockInfo::default(),
1308 }));
1309 }
1310
1311 if data.len() != 2 {
1312 return Err(Error::invalid_input_source(
1313 format!(
1314 "RLE decompressor expects exactly 2 buffers, got {}",
1315 data.len()
1316 )
1317 .into(),
1318 ));
1319 }
1320
1321 let mut data_iter = data.into_iter();
1322 let values_buffer = data_iter.next().unwrap();
1323 let lengths_buffer = data_iter.next().unwrap();
1324 let (values_buffer, lengths_buffer) =
1325 self.decode_child_buffers(values_buffer, lengths_buffer)?;
1326
1327 let decoded_data = match self.bits_per_value {
1328 8 => self.decode_generic::<u8>(
1329 &values_buffer,
1330 &lengths_buffer,
1331 num_values,
1332 clamp_overflow,
1333 )?,
1334 16 => self.decode_generic::<u16>(
1335 &values_buffer,
1336 &lengths_buffer,
1337 num_values,
1338 clamp_overflow,
1339 )?,
1340 32 => self.decode_generic::<u32>(
1341 &values_buffer,
1342 &lengths_buffer,
1343 num_values,
1344 clamp_overflow,
1345 )?,
1346 64 => self.decode_generic::<u64>(
1347 &values_buffer,
1348 &lengths_buffer,
1349 num_values,
1350 clamp_overflow,
1351 )?,
1352 _ => {
1353 return Err(Error::invalid_input_source(
1354 format!(
1355 "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}",
1356 self.bits_per_value
1357 )
1358 .into(),
1359 ));
1360 }
1361 };
1362
1363 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1364 bits_per_value: self.bits_per_value,
1365 data: decoded_data,
1366 num_values,
1367 block_info: BlockInfo::default(),
1368 }))
1369 }
1370
1371 fn decode_child_buffers(
1372 &self,
1373 values_buffer: LanceBuffer,
1374 lengths_buffer: LanceBuffer,
1375 ) -> Result<(LanceBuffer, LanceBuffer)> {
1376 let values_requires_num_runs = self.values.requires_num_values();
1377 let lengths_requires_num_runs = self.run_lengths.requires_num_values();
1378 if values_requires_num_runs && lengths_requires_num_runs {
1379 return Err(Error::invalid_input_source(
1380 "RLE values and run lengths child compression both require the run count".into(),
1381 ));
1382 }
1383
1384 if values_requires_num_runs {
1385 let lengths_buffer = self
1386 .run_lengths
1387 .decode(lengths_buffer, None, "run lengths")?;
1388 let num_runs = Self::num_child_values(
1389 &lengths_buffer,
1390 self.run_lengths.bits_per_value(),
1391 "run lengths",
1392 )?;
1393 let values_buffer = self
1394 .values
1395 .decode(values_buffer, Some(num_runs), "values")?;
1396 Ok((values_buffer, lengths_buffer))
1397 } else if lengths_requires_num_runs {
1398 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1399 let num_runs =
1400 Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?;
1401 let lengths_buffer =
1402 self.run_lengths
1403 .decode(lengths_buffer, Some(num_runs), "run lengths")?;
1404 Ok((values_buffer, lengths_buffer))
1405 } else {
1406 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1407 let lengths_buffer = self
1408 .run_lengths
1409 .decode(lengths_buffer, None, "run lengths")?;
1410 Ok((values_buffer, lengths_buffer))
1411 }
1412 }
1413
1414 fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result<u64> {
1415 let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
1416 Error::invalid_input_source(
1417 format!("RLE {label} child bit width is too large: {bits_per_value}").into(),
1418 )
1419 })?;
1420 if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) {
1421 return Err(Error::invalid_input_source(
1422 format!(
1423 "RLE {label} child decoded to {} bytes, not divisible by {}",
1424 buffer.len(),
1425 bytes_per_value
1426 )
1427 .into(),
1428 ));
1429 }
1430 Ok((buffer.len() / bytes_per_value) as u64)
1431 }
1432
1433 fn decode_generic<T>(
1434 &self,
1435 values_buffer: &LanceBuffer,
1436 lengths_buffer: &LanceBuffer,
1437 num_values: u64,
1438 clamp_overflow: bool,
1439 ) -> Result<LanceBuffer>
1440 where
1441 T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType,
1442 {
1443 let type_size = std::mem::size_of::<T>();
1444 let length_size = self.run_length_width.bytes_per_value();
1445
1446 if values_buffer.is_empty() || lengths_buffer.is_empty() {
1447 if num_values == 0 {
1448 return Ok(LanceBuffer::empty());
1449 } else {
1450 return Err(Error::invalid_input_source(
1451 format!("Empty buffers but expected {} values", num_values).into(),
1452 ));
1453 }
1454 }
1455
1456 if !values_buffer.len().is_multiple_of(type_size)
1457 || !lengths_buffer.len().is_multiple_of(length_size)
1458 {
1459 return Err(Error::invalid_input_source(format!(
1460 "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})",
1461 std::any::type_name::<T>(),
1462 values_buffer.len(),
1463 type_size,
1464 lengths_buffer.len(),
1465 length_size
1466 )
1467 .into()));
1468 }
1469
1470 let num_runs = values_buffer.len() / type_size;
1471 let num_length_entries = lengths_buffer.len() / length_size;
1472 if num_runs != num_length_entries {
1473 return Err(Error::invalid_input_source(
1474 format!(
1475 "Inconsistent RLE buffers: {} runs but {} length entries",
1476 num_runs, num_length_entries
1477 )
1478 .into(),
1479 ));
1480 }
1481
1482 let values_ref = values_buffer.borrow_to_typed_slice::<T>();
1483 let values: &[T] = values_ref.as_ref();
1484 let lengths = lengths_buffer.as_ref();
1485
1486 let expected_value_count = usize::try_from(num_values).map_err(|_| {
1487 Error::invalid_input_source(
1488 format!("RLE num_values does not fit in usize: {num_values}").into(),
1489 )
1490 })?;
1491 let mut decoded: Vec<T> = Vec::new();
1498 decoded
1499 .try_reserve_exact(expected_value_count)
1500 .map_err(|_| {
1501 Error::invalid_input_source(
1502 format!("RLE decoding cannot allocate {expected_value_count} values").into(),
1503 )
1504 })?;
1505 for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) {
1506 let length = self.run_length_width.read_length(length_bytes);
1507 if length == 0 {
1508 return Err(Error::invalid_input_source(
1509 "RLE decoding encountered a zero run length".into(),
1510 ));
1511 }
1512 let length = usize::try_from(length).map_err(|_| {
1513 Error::invalid_input_source(
1514 format!("RLE run length does not fit in usize: {length}").into(),
1515 )
1516 })?;
1517 let remaining = expected_value_count - decoded.len();
1518 if length > remaining {
1519 if !clamp_overflow {
1520 return Err(Error::invalid_input_source(
1521 format!(
1522 "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1523 decoded.len() + length,
1524 expected_value_count
1525 )
1526 .into(),
1527 ));
1528 }
1529 decoded.resize(expected_value_count, *value);
1530 break;
1531 }
1532 decoded.resize(decoded.len() + length, *value);
1533 }
1534
1535 if decoded.len() != expected_value_count {
1536 return Err(Error::invalid_input_source(
1537 format!(
1538 "RLE decoding produced {} values, expected {}",
1539 decoded.len(),
1540 expected_value_count
1541 )
1542 .into(),
1543 ));
1544 }
1545
1546 trace!(
1547 "RLE decoded {} {} values",
1548 num_values,
1549 std::any::type_name::<T>()
1550 );
1551 Ok(LanceBuffer::reinterpret_vec(decoded))
1552 }
1553}
1554
1555impl MiniBlockDecompressor for RleDecompressor {
1556 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
1557 self.decode_data(data, num_values, true)
1558 }
1559
1560 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
1561 num_values
1562 .checked_mul(self.bits_per_value)
1563 .map(|bits| bits.div_ceil(8))
1564 }
1565}
1566
1567impl BlockDecompressor for RleDecompressor {
1568 fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result<DataBlock> {
1569 let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
1570 self.decode_data(vec![values_buffer, lengths_buffer], num_values, false)
1571 }
1572}
1573
1574fn parse_rle_block_frame(data: &LanceBuffer) -> Result<(LanceBuffer, LanceBuffer)> {
1577 if data.len() < 8 {
1579 return Err(Error::invalid_input_source(
1580 format!("Insufficient data size: {}", data.len()).into(),
1581 ));
1582 }
1583
1584 let values_size_bytes: [u8; 8] = data[..8].try_into().expect("slice length already checked");
1585 let values_size: usize = u64::from_le_bytes(values_size_bytes)
1586 .try_into()
1587 .map_err(|_| {
1588 Error::invalid_input_source(
1589 format!(
1590 "Invalid values buffer size: {}",
1591 u64::from_le_bytes(values_size_bytes)
1592 )
1593 .into(),
1594 )
1595 })?;
1596
1597 let values_start: usize = 8;
1599 let lengths_start = values_start
1600 .checked_add(values_size)
1601 .ok_or_else(|| Error::invalid_input_source("Invalid RLE values buffer size".into()))?;
1602
1603 if data.len() < lengths_start {
1604 return Err(Error::invalid_input_source(
1605 format!("Insufficient data size: {}", data.len()).into(),
1606 ));
1607 }
1608
1609 let values_buffer = data.slice_with_length(values_start, values_size);
1610 let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start);
1611 Ok((values_buffer, lengths_buffer))
1612}
1613
1614#[derive(Clone, Debug)]
1615enum RleRunLengths {
1616 U8(ScalarBuffer<u8>),
1617 U16(ScalarBuffer<u16>),
1618 U32(ScalarBuffer<u32>),
1619}
1620
1621impl RleRunLengths {
1622 fn try_new(buffer: LanceBuffer, width: RunLengthWidth) -> Result<Self> {
1623 let width_bytes = width.bytes_per_value();
1624 if !buffer.len().is_multiple_of(width_bytes) {
1625 return Err(Error::invalid_input_source(
1626 format!(
1627 "Invalid RLE run lengths buffer: {} bytes (not divisible by {})",
1628 buffer.len(),
1629 width_bytes
1630 )
1631 .into(),
1632 ));
1633 }
1634 Ok(match width {
1635 RunLengthWidth::U8 => Self::U8(buffer.borrow_to_typed_slice()),
1636 RunLengthWidth::U16 => Self::U16(buffer.borrow_to_typed_slice()),
1637 RunLengthWidth::U32 => Self::U32(buffer.borrow_to_typed_slice()),
1638 })
1639 }
1640
1641 fn len(&self) -> usize {
1642 match self {
1643 Self::U8(lengths) => lengths.len(),
1644 Self::U16(lengths) => lengths.len(),
1645 Self::U32(lengths) => lengths.len(),
1646 }
1647 }
1648
1649 fn get(&self, index: usize) -> usize {
1650 match self {
1651 Self::U8(lengths) => lengths[index] as usize,
1652 Self::U16(lengths) => lengths[index] as usize,
1653 Self::U32(lengths) => lengths[index] as usize,
1654 }
1655 }
1656
1657 fn owned_size(&self) -> usize {
1658 match self {
1659 Self::U8(lengths) => std::mem::size_of_val(lengths.as_ref()),
1660 Self::U16(lengths) => std::mem::size_of_val(lengths.as_ref()),
1661 Self::U32(lengths) => std::mem::size_of_val(lengths.as_ref()),
1662 }
1663 }
1664
1665 fn into_owned(self) -> Self {
1666 match self {
1667 Self::U8(lengths) => Self::U8(ScalarBuffer::from(lengths.as_ref().to_vec())),
1668 Self::U16(lengths) => Self::U16(ScalarBuffer::from(lengths.as_ref().to_vec())),
1669 Self::U32(lengths) => Self::U32(ScalarBuffer::from(lengths.as_ref().to_vec())),
1670 }
1671 }
1672
1673 fn deep_size(&self) -> usize {
1674 match self {
1675 Self::U8(lengths) => lengths.inner().capacity(),
1676 Self::U16(lengths) => lengths.inner().capacity(),
1677 Self::U32(lengths) => lengths.inner().capacity(),
1678 }
1679 }
1680}
1681
1682#[derive(Clone, Debug)]
1688pub(crate) struct RleRuns {
1689 values: ScalarBuffer<u16>,
1690 lengths: RleRunLengths,
1691 num_values: usize,
1692 coalesced_runs: usize,
1693}
1694
1695impl RleRuns {
1696 fn try_new(
1697 values_buffer: LanceBuffer,
1698 lengths_buffer: LanceBuffer,
1699 run_length_width: RunLengthWidth,
1700 num_values: u64,
1701 ) -> Result<Self> {
1702 let num_values = usize::try_from(num_values).map_err(|_| {
1703 Error::invalid_input_source(
1704 format!("RLE num_values does not fit in usize: {num_values}").into(),
1705 )
1706 })?;
1707 let type_size = std::mem::size_of::<u16>();
1708 if !values_buffer.len().is_multiple_of(type_size) {
1709 return Err(Error::invalid_input_source(
1710 format!(
1711 "Invalid RLE u16 values buffer: {} bytes (not divisible by {})",
1712 values_buffer.len(),
1713 type_size
1714 )
1715 .into(),
1716 ));
1717 }
1718
1719 let values = values_buffer.borrow_to_typed_slice::<u16>();
1720 let lengths = RleRunLengths::try_new(lengths_buffer, run_length_width)?;
1721 if values.len() != lengths.len() {
1722 return Err(Error::invalid_input_source(
1723 format!(
1724 "Inconsistent RLE buffers: {} runs but {} length entries",
1725 values.len(),
1726 lengths.len()
1727 )
1728 .into(),
1729 ));
1730 }
1731 if values.is_empty() && num_values != 0 {
1732 return Err(Error::invalid_input_source(
1733 format!("Empty RLE buffers but expected {num_values} values").into(),
1734 ));
1735 }
1736
1737 let mut decoded_values = 0usize;
1738 let mut coalesced_runs = 0usize;
1739 let mut previous_value = None;
1740 for run in 0..values.len() {
1741 let length = lengths.get(run);
1742 if length == 0 {
1743 return Err(Error::invalid_input_source(
1744 "RLE decoding encountered a zero run length".into(),
1745 ));
1746 }
1747 decoded_values = decoded_values.checked_add(length).ok_or_else(|| {
1748 Error::invalid_input_source("RLE run length sum overflowed usize".into())
1749 })?;
1750 if decoded_values > num_values {
1751 return Err(Error::invalid_input_source(
1752 format!(
1753 "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1754 decoded_values, num_values
1755 )
1756 .into(),
1757 ));
1758 }
1759 if previous_value != Some(values[run]) {
1760 coalesced_runs += 1;
1761 previous_value = Some(values[run]);
1762 }
1763 }
1764 if decoded_values != num_values {
1765 return Err(Error::invalid_input_source(
1766 format!(
1767 "RLE decoding produced {} values, expected {}",
1768 decoded_values, num_values
1769 )
1770 .into(),
1771 ));
1772 }
1773
1774 Ok(Self {
1775 values,
1776 lengths,
1777 num_values,
1778 coalesced_runs,
1779 })
1780 }
1781
1782 pub(crate) fn num_values(&self) -> usize {
1783 self.num_values
1784 }
1785
1786 pub(crate) fn num_runs(&self) -> usize {
1787 self.values.len()
1788 }
1789
1790 pub(crate) fn coalesced_runs(&self) -> usize {
1791 self.coalesced_runs
1792 }
1793
1794 pub(crate) fn owned_size(&self) -> usize {
1795 std::mem::size_of_val(self.values.as_ref()) + self.lengths.owned_size()
1796 }
1797
1798 pub(crate) fn into_owned(self) -> Self {
1799 Self {
1800 values: ScalarBuffer::from(self.values.as_ref().to_vec()),
1801 lengths: self.lengths.into_owned(),
1802 num_values: self.num_values,
1803 coalesced_runs: self.coalesced_runs,
1804 }
1805 }
1806
1807 pub(crate) fn deep_size(&self) -> usize {
1808 self.values.inner().capacity() + self.lengths.deep_size()
1809 }
1810
1811 pub(crate) fn value(&self, run: usize) -> u16 {
1812 self.values[run]
1813 }
1814
1815 pub(crate) fn length(&self, run: usize) -> usize {
1816 self.lengths.get(run)
1817 }
1818
1819 pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = (u16, usize)> + '_ {
1820 (0..self.num_runs()).map(|run| (self.value(run), self.length(run)))
1821 }
1822}
1823
1824impl RleDecompressor {
1825 pub(crate) fn decode_u16_runs(&self, data: LanceBuffer, num_values: u64) -> Result<RleRuns> {
1827 if self.bits_per_value != 16 {
1828 return Err(Error::invalid_input_source(
1829 format!(
1830 "RLE level values must be 16 bits, got {}",
1831 self.bits_per_value
1832 )
1833 .into(),
1834 ));
1835 }
1836 let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
1837 let (values_buffer, lengths_buffer) =
1838 self.decode_child_buffers(values_buffer, lengths_buffer)?;
1839 RleRuns::try_new(
1840 values_buffer,
1841 lengths_buffer,
1842 self.run_length_width,
1843 num_values,
1844 )
1845 }
1846}
1847
1848#[cfg(test)]
1849mod tests {
1850 use std::sync::Arc;
1851
1852 use super::*;
1853 use crate::compression::{
1854 DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor,
1855 };
1856 use crate::data::DataBlock;
1857 use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES;
1858 use crate::encodings::physical::block::{CompressionConfig, CompressionScheme};
1859 use crate::{
1860 buffer::LanceBuffer,
1861 compression::{BlockCompressor, BlockDecompressor},
1862 };
1863 use arrow_array::Int32Array;
1864 use rstest::rstest;
1865
1866 fn compress_miniblock(
1867 compressor: &dyn MiniBlockCompressor,
1868 data: DataBlock,
1869 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1870 compressor.compress(MiniBlockCompressionContext::new(0, true, true), data)
1871 }
1872
1873 fn expand_u16_runs(runs: &RleRuns) -> Vec<u16> {
1874 let mut expanded = Vec::with_capacity(runs.num_values());
1875 for (value, length) in runs.iter() {
1876 expanded.extend(std::iter::repeat_n(value, length));
1877 }
1878 expanded
1879 }
1880
1881 #[test]
1882 fn decode_u16_runs_matches_eager() {
1883 let mut levels: Vec<u16> = vec![1u16; 1000];
1885 levels.extend(std::iter::repeat_n(0u16, 500));
1886 levels.extend(std::iter::repeat_n(2u16, 300));
1887 let num_values = levels.len() as u64;
1888
1889 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
1890 data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
1891 bits_per_value: 16,
1892 num_values,
1893 block_info: BlockInfo::new(),
1894 });
1895 let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
1896
1897 let eager =
1898 BlockDecompressor::decompress(&RleDecompressor::new(16), frame.clone(), num_values)
1899 .unwrap();
1900 let DataBlock::FixedWidth(eager) = eager else {
1901 panic!("expected fixed-width block");
1902 };
1903 assert_eq!(
1904 eager.data.borrow_to_typed_slice::<u16>().as_ref(),
1905 levels.as_slice()
1906 );
1907
1908 let runs = RleDecompressor::new(16)
1910 .decode_u16_runs(frame, num_values)
1911 .unwrap();
1912 assert_eq!(runs.num_values(), num_values as usize);
1913 assert_eq!(runs.num_runs(), 8);
1916 assert_eq!(runs.coalesced_runs(), 3);
1917 assert_eq!(expand_u16_runs(&runs), levels);
1918 }
1919
1920 #[test]
1921 fn decode_u16_runs_empty() {
1922 let mut empty_frame = Vec::new();
1923 empty_frame.extend_from_slice(&0u64.to_le_bytes());
1924 let runs = RleDecompressor::new(16)
1925 .decode_u16_runs(LanceBuffer::from(empty_frame), 0)
1926 .unwrap();
1927 assert_eq!(runs.num_values(), 0);
1928 assert_eq!(runs.num_runs(), 0);
1929 assert_eq!(runs.coalesced_runs(), 0);
1930
1931 let error = RleDecompressor::new(16)
1932 .decode_u16_runs(LanceBuffer::empty(), 0)
1933 .unwrap_err();
1934 assert!(error.to_string().contains("Insufficient data size: 0"));
1935 }
1936
1937 #[rstest]
1938 #[case::zero(0, 1, "zero run length")]
1939 #[case::underflow(1, 2, "produced 1 values, expected 2")]
1940 #[case::overflow(2, 1, "overflowed expected value count")]
1941 #[case::nonempty_for_empty(1, 0, "overflowed expected value count")]
1942 fn decode_u16_runs_rejects_invalid_coverage(
1943 #[case] run_length: u8,
1944 #[case] num_values: u64,
1945 #[case] expected_message: &str,
1946 ) {
1947 let mut frame = Vec::new();
1948 frame.extend_from_slice(&2u64.to_le_bytes());
1949 frame.extend_from_slice(&7u16.to_le_bytes());
1950 frame.push(run_length);
1951
1952 let error = RleDecompressor::new(16)
1953 .decode_u16_runs(LanceBuffer::from(frame), num_values)
1954 .unwrap_err();
1955 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
1956 assert!(error.to_string().contains(expected_message));
1957 }
1958
1959 #[test]
1960 #[cfg(any(feature = "lz4", feature = "zstd"))]
1961 fn decode_u16_runs_supports_compressed_values_child() {
1962 let levels: Vec<u16> = (0..1024)
1963 .flat_map(|run| std::iter::repeat_n((run % 8) as u16, 4))
1964 .collect();
1965 let num_values = levels.len() as u64;
1966 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
1967 data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
1968 bits_per_value: 16,
1969 num_values,
1970 block_info: BlockInfo::new(),
1971 });
1972 let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
1973 let (values, lengths) = parse_rle_block_frame(&frame).unwrap();
1974
1975 let compression = test_general_compression();
1976 let compressor = GeneralBufferCompressor::get_compressor(compression).unwrap();
1977 let mut compressed_values = Vec::new();
1978 compressor
1979 .compress(values.as_ref(), &mut compressed_values)
1980 .unwrap();
1981 let mut compressed_frame = Vec::new();
1982 compressed_frame.extend_from_slice(&(compressed_values.len() as u64).to_le_bytes());
1983 compressed_frame.extend_from_slice(&compressed_values);
1984 compressed_frame.extend_from_slice(lengths.as_ref());
1985
1986 let encoding = ProtobufUtils21::rle(
1987 ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(16, None)).unwrap(),
1988 ProtobufUtils21::flat(8, None),
1989 );
1990 let decompressor = create_rle_decompressor(
1991 expect_rle(&encoding),
1992 &DefaultDecompressionStrategy::default(),
1993 )
1994 .unwrap();
1995 let runs = decompressor
1996 .decode_u16_runs(LanceBuffer::from(compressed_frame), num_values)
1997 .unwrap();
1998 assert_eq!(expand_u16_runs(&runs), levels);
1999 }
2000
2001 #[test]
2002 fn decode_u16_runs_counts_coalesced_runs() {
2003 let num_values = 5000u64;
2007 let constant: Vec<u16> = vec![7u16; num_values as usize];
2008 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2009 data: LanceBuffer::reinterpret_slice(Arc::from(constant)),
2010 bits_per_value: 16,
2011 num_values,
2012 block_info: BlockInfo::new(),
2013 });
2014 let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
2015 let runs = RleDecompressor::new(16)
2016 .decode_u16_runs(frame, num_values)
2017 .unwrap();
2018 assert_eq!(
2019 runs.num_runs(),
2020 num_values.div_ceil(u8::MAX as u64) as usize
2021 );
2022 assert_eq!(runs.coalesced_runs(), 1);
2023 assert_eq!(expand_u16_runs(&runs), vec![7u16; num_values as usize]);
2024
2025 let alternating: Vec<u16> = (0..200u16).map(|i| i % 2).collect();
2028 let n = alternating.len() as u64;
2029 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2030 data: LanceBuffer::reinterpret_slice(Arc::from(alternating.clone())),
2031 bits_per_value: 16,
2032 num_values: n,
2033 block_info: BlockInfo::new(),
2034 });
2035 let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
2036 let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap();
2037 assert_eq!(
2038 runs.coalesced_runs() as u64,
2039 n,
2040 "no two adjacent values are equal"
2041 );
2042 assert_eq!(expand_u16_runs(&runs), alternating);
2043 }
2044
2045 #[test]
2048 fn test_basic_miniblock_rle_encoding() {
2049 let encoder = RleEncoder::new();
2050
2051 let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]);
2053 let data_block = DataBlock::from_array(array);
2054
2055 let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap();
2056
2057 assert_eq!(compressed.num_values, 9);
2058 assert_eq!(compressed.chunks.len(), 1);
2059
2060 let values_buffer = &compressed.data[0];
2062 let lengths_buffer = &compressed.data[1];
2063 assert_eq!(values_buffer.len(), 12); assert_eq!(lengths_buffer.len(), 3); }
2066
2067 #[test]
2068 fn test_long_run_splitting() {
2069 let encoder = RleEncoder::new();
2070
2071 let mut data = vec![42i32; 1000]; data.extend(&[100i32; 300]); let array = Int32Array::from(data);
2076 let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2077
2078 let lengths_buffer = &compressed.data[1];
2080 assert_eq!(lengths_buffer.len(), 6);
2081 }
2082
2083 #[test]
2084 fn test_rle_v2_u16_miniblock_encoding() {
2085 let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16);
2086
2087 let data = vec![42i32; 1000];
2088 let array = Int32Array::from(data);
2089 let (compressed, encoding) =
2090 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2091
2092 assert_eq!(compressed.data[0].len(), 4);
2093 assert_eq!(compressed.data[1].len(), 2);
2094 assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes());
2095
2096 let rle = match encoding.compression.as_ref().unwrap() {
2097 crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
2098 other => panic!("expected RLE encoding, got {other:?}"),
2099 };
2100 let run_lengths = rle.run_lengths.as_ref().unwrap();
2101 let flat = match run_lengths.compression.as_ref().unwrap() {
2102 crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat,
2103 other => panic!("expected flat run lengths, got {other:?}"),
2104 };
2105 assert_eq!(flat.bits_per_value, 16);
2106
2107 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2108 let decompressed = MiniBlockDecompressor::decompress(
2109 &decompressor,
2110 compressed.data,
2111 compressed.num_values,
2112 )
2113 .unwrap();
2114 match decompressed {
2115 DataBlock::FixedWidth(block) => {
2116 let values = block.data.borrow_to_typed_slice::<i32>();
2117 assert_eq!(values.as_ref(), vec![42i32; 1000]);
2118 }
2119 _ => panic!("Expected FixedWidth block"),
2120 }
2121 }
2122
2123 #[test]
2124 #[cfg(any(feature = "lz4", feature = "zstd"))]
2125 fn test_rle_miniblock_compressed_values_child() {
2126 let compression = test_general_compression();
2127 let encoder =
2128 RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false);
2129 let array = Int32Array::from(repeating_runs(1024, 4));
2130 let (compressed, encoding) =
2131 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2132
2133 let rle = expect_rle(&encoding);
2134 assert!(matches!(
2135 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2136 crate::format::pb21::compressive_encoding::Compression::General(_)
2137 ));
2138 assert!(matches!(
2139 rle.run_lengths
2140 .as_ref()
2141 .unwrap()
2142 .compression
2143 .as_ref()
2144 .unwrap(),
2145 crate::format::pb21::compressive_encoding::Compression::Flat(_)
2146 ));
2147
2148 let decompressor = DefaultDecompressionStrategy::default()
2149 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2150 .unwrap();
2151 let decoded =
2152 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2153 .unwrap();
2154 assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4));
2155 }
2156
2157 #[test]
2158 #[cfg(any(feature = "lz4", feature = "zstd"))]
2159 fn test_rle_miniblock_compressed_run_lengths_child() {
2160 let compression = test_general_compression();
2161 let encoder =
2162 RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false);
2163 let expected = repeating_runs(1024, 4);
2164 let (compressed, encoding) = compress_miniblock(
2165 &encoder,
2166 DataBlock::from_array(Int32Array::from(expected.clone())),
2167 )
2168 .unwrap();
2169
2170 let rle = expect_rle(&encoding);
2171 assert!(matches!(
2172 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2173 crate::format::pb21::compressive_encoding::Compression::Flat(_)
2174 ));
2175 assert!(matches!(
2176 rle.run_lengths
2177 .as_ref()
2178 .unwrap()
2179 .compression
2180 .as_ref()
2181 .unwrap(),
2182 crate::format::pb21::compressive_encoding::Compression::General(_)
2183 ));
2184
2185 let decompressor = DefaultDecompressionStrategy::default()
2186 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2187 .unwrap();
2188 let decoded =
2189 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2190 .unwrap();
2191 assert_decoded_i32_eq(decoded, &expected);
2192 }
2193
2194 #[test]
2195 #[cfg(feature = "bitpacking")]
2196 fn test_rle_miniblock_bitpacked_run_lengths_child() {
2197 use crate::encodings::physical::bitpacking::OutOfLineBitpacking;
2198
2199 let expected = repeating_runs(1024, 4);
2200 let (compressed, _) = compress_miniblock(
2201 &RleEncoder::new(),
2202 DataBlock::from_array(Int32Array::from(expected.clone())),
2203 )
2204 .unwrap();
2205 let run_lengths = compressed.data[1].clone();
2206 let num_runs = run_lengths.len() as u64;
2207 let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2208 bits_per_value: 8,
2209 data: run_lengths,
2210 num_values: num_runs,
2211 block_info: BlockInfo::default(),
2212 });
2213 let bitpacked_run_lengths =
2214 BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap();
2215 let encoding = ProtobufUtils21::rle(
2216 ProtobufUtils21::flat(32, None),
2217 ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
2218 );
2219
2220 let decompressor = DefaultDecompressionStrategy::default()
2221 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2222 .unwrap();
2223 let decoded = MiniBlockDecompressor::decompress(
2224 decompressor.as_ref(),
2225 vec![compressed.data[0].clone(), bitpacked_run_lengths],
2226 expected.len() as u64,
2227 )
2228 .unwrap();
2229 assert_decoded_i32_eq(decoded, &expected);
2230 }
2231
2232 #[test]
2233 #[cfg(feature = "bitpacking")]
2234 fn test_rle_rejects_two_count_dependent_child_encodings() {
2235 let encoding = ProtobufUtils21::rle(
2236 ProtobufUtils21::out_of_line_bitpacking(32, ProtobufUtils21::flat(3, None)),
2237 ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
2238 );
2239
2240 let err = DefaultDecompressionStrategy::default()
2241 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2242 .unwrap_err();
2243 assert!(
2244 err.to_string()
2245 .contains("cannot both require the run count")
2246 );
2247 }
2248
2249 #[cfg(any(feature = "lz4", feature = "zstd"))]
2250 fn test_general_compression() -> CompressionConfig {
2251 if cfg!(feature = "zstd") {
2252 CompressionConfig::new(CompressionScheme::Zstd, Some(3))
2253 } else {
2254 CompressionConfig::new(CompressionScheme::Lz4, None)
2255 }
2256 }
2257
2258 fn repeating_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2259 let mut values = Vec::with_capacity(num_runs * run_length);
2260 for run in 0..num_runs {
2261 values.extend(std::iter::repeat_n((run % 8) as i32, run_length));
2262 }
2263 values
2264 }
2265
2266 fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle {
2267 match encoding.compression.as_ref().unwrap() {
2268 crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
2269 other => panic!("expected RLE encoding, got {other:?}"),
2270 }
2271 }
2272
2273 fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) {
2274 match decoded {
2275 DataBlock::FixedWidth(block) => {
2276 let values = block.data.borrow_to_typed_slice::<i32>();
2277 assert_eq!(values.as_ref(), expected);
2278 }
2279 _ => panic!("Expected FixedWidth block"),
2280 }
2281 }
2282
2283 #[test]
2284 #[cfg(any(feature = "lz4", feature = "zstd"))]
2285 fn test_rle_miniblock_compressed_children_multiple_chunks() {
2286 let compression = test_general_compression();
2287 let encoder = RleEncoder::with_child_encoding(
2288 RunLengthWidth::U8,
2289 Some(compression),
2290 Some(compression),
2291 false,
2292 );
2293 let expected = repeating_runs(8192, 4);
2294 let (compressed, encoding) = compress_miniblock(
2295 &encoder,
2296 DataBlock::from_array(Int32Array::from(expected.clone())),
2297 )
2298 .unwrap();
2299
2300 assert!(compressed.chunks.len() > 1);
2301 let rle = expect_rle(&encoding);
2302 assert!(matches!(
2303 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2304 crate::format::pb21::compressive_encoding::Compression::General(_)
2305 ));
2306 assert!(matches!(
2307 rle.run_lengths
2308 .as_ref()
2309 .unwrap()
2310 .compression
2311 .as_ref()
2312 .unwrap(),
2313 crate::format::pb21::compressive_encoding::Compression::General(_)
2314 ));
2315
2316 let decoded = decompress_i32_chunks(&compressed, &encoding);
2317 assert_eq!(decoded, expected);
2318 }
2319
2320 #[test]
2321 #[cfg(feature = "bitpacking")]
2322 fn test_rle_miniblock_bitpacks_values_child_when_smaller() {
2323 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
2324 let expected = monotonic_runs(2048, 4);
2325 let (compressed, encoding) = compress_miniblock(
2326 &encoder,
2327 DataBlock::from_array(Int32Array::from(expected.clone())),
2328 )
2329 .unwrap();
2330
2331 let rle = expect_rle(&encoding);
2332 assert!(matches!(
2333 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2334 crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
2335 ));
2336 assert!(matches!(
2337 rle.run_lengths
2338 .as_ref()
2339 .unwrap()
2340 .compression
2341 .as_ref()
2342 .unwrap(),
2343 crate::format::pb21::compressive_encoding::Compression::Flat(_)
2344 ));
2345
2346 let decoded = decompress_i32_chunks(&compressed, &encoding);
2347 assert_eq!(decoded, expected);
2348 }
2349
2350 #[test]
2351 #[cfg(feature = "bitpacking")]
2352 fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() {
2353 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
2354 let expected = high_entropy_runs(2048, 4);
2355 let (compressed, encoding) = compress_miniblock(
2356 &encoder,
2357 DataBlock::from_array(Int32Array::from(expected.clone())),
2358 )
2359 .unwrap();
2360
2361 let rle = expect_rle(&encoding);
2362 assert!(matches!(
2363 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2364 crate::format::pb21::compressive_encoding::Compression::Flat(_)
2365 ));
2366 assert!(matches!(
2367 rle.run_lengths
2368 .as_ref()
2369 .unwrap()
2370 .compression
2371 .as_ref()
2372 .unwrap(),
2373 crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
2374 ));
2375
2376 let decoded = decompress_i32_chunks(&compressed, &encoding);
2377 assert_eq!(decoded, expected);
2378 }
2379
2380 fn decompress_i32_chunks(
2381 compressed: &MiniBlockCompressed,
2382 encoding: &CompressiveEncoding,
2383 ) -> Vec<i32> {
2384 let strategy = DefaultDecompressionStrategy::default();
2385 let decompressor = strategy
2386 .create_miniblock_decompressor(encoding, &strategy)
2387 .unwrap();
2388 let mut offsets = vec![0usize; compressed.data.len()];
2389 let mut values_processed = 0u64;
2390 let mut decoded_values = Vec::new();
2391
2392 for chunk in &compressed.chunks {
2393 let chunk_values = chunk.num_values(values_processed, compressed.num_values);
2394 let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len());
2395 for (idx, size) in chunk.buffer_sizes.iter().enumerate() {
2396 let size = *size as usize;
2397 chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size));
2398 offsets[idx] += size;
2399 }
2400
2401 let decoded = decompressor
2402 .decompress(chunk_buffers, chunk_values)
2403 .unwrap();
2404 match decoded {
2405 DataBlock::FixedWidth(block) => {
2406 let values = block.data.borrow_to_typed_slice::<i32>();
2407 decoded_values.extend_from_slice(values.as_ref());
2408 }
2409 _ => panic!("Expected FixedWidth block"),
2410 }
2411 values_processed += chunk_values;
2412 }
2413
2414 assert_eq!(values_processed, compressed.num_values);
2415 decoded_values
2416 }
2417
2418 #[cfg(feature = "bitpacking")]
2419 fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2420 let mut values = Vec::with_capacity(num_runs * run_length);
2421 for run in 0..num_runs {
2422 values.extend(std::iter::repeat_n(run as i32, run_length));
2423 }
2424 values
2425 }
2426
2427 #[cfg(feature = "bitpacking")]
2428 fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2429 let mut values = Vec::with_capacity(num_runs * run_length);
2430 let mut state = 7u64;
2431 for _ in 0..num_runs {
2432 state = state
2433 .wrapping_mul(6364136223846793005)
2434 .wrapping_add(1442695040888963407);
2435 values.extend(std::iter::repeat_n((state >> 32) as i32, run_length));
2436 }
2437 values
2438 }
2439
2440 #[test]
2441 fn test_select_run_length_width_prefers_u16_for_long_runs() {
2442 let mut entries = [0u64; 3];
2443 accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries);
2444 let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap();
2445 assert_eq!(width, RunLengthWidth::U16);
2446 }
2447
2448 #[test]
2451 fn test_round_trip_all_types() {
2452 test_round_trip_helper(vec![42u8, 42, 42, 100, 100, 255, 255, 255, 255], 8);
2454
2455 test_round_trip_helper(vec![1000u16, 1000, 2000, 2000, 2000, 3000], 16);
2457
2458 test_round_trip_helper(vec![100i32, 100, 100, -200, -200, 300, 300, 300, 300], 32);
2460
2461 test_round_trip_helper(vec![1_000_000_000u64; 5], 64);
2463 }
2464
2465 fn test_round_trip_helper<T>(data: Vec<T>, bits_per_value: u64)
2466 where
2467 T: bytemuck::Pod + PartialEq + std::fmt::Debug,
2468 {
2469 let encoder = RleEncoder::new();
2470 let bytes: Vec<u8> = data
2471 .iter()
2472 .flat_map(|v| bytemuck::bytes_of(v))
2473 .copied()
2474 .collect();
2475
2476 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2477 bits_per_value,
2478 data: LanceBuffer::from(bytes),
2479 num_values: data.len() as u64,
2480 block_info: BlockInfo::default(),
2481 });
2482
2483 let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
2484 let decompressor = RleDecompressor::new(bits_per_value);
2485 let decompressed = MiniBlockDecompressor::decompress(
2486 &decompressor,
2487 compressed.data,
2488 compressed.num_values,
2489 )
2490 .unwrap();
2491
2492 match decompressed {
2493 DataBlock::FixedWidth(ref block) => {
2494 assert_eq!(block.data.len(), data.len() * std::mem::size_of::<T>());
2496 }
2497 _ => panic!("Expected FixedWidth block"),
2498 }
2499 }
2500
2501 #[test]
2504 fn test_power_of_two_chunking() {
2505 let encoder = RleEncoder::new();
2506
2507 let test_sizes = vec![1000, 2500, 5000, 10000];
2509
2510 for size in test_sizes {
2511 let data: Vec<i32> = (0..size)
2512 .map(|i| i / 50) .collect();
2514
2515 let array = Int32Array::from(data);
2516 let (compressed, _) =
2517 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2518
2519 for (i, chunk) in compressed.chunks.iter().enumerate() {
2521 if i < compressed.chunks.len() - 1 {
2522 assert!(chunk.log_num_values > 0);
2523 let chunk_values = 1u64 << chunk.log_num_values;
2524 assert!(chunk_values.is_power_of_two());
2525 assert!(chunk_values <= *MAX_MINIBLOCK_VALUES);
2526 } else {
2527 assert_eq!(chunk.log_num_values, 0);
2528 }
2529 }
2530 }
2531 }
2532
2533 #[rstest]
2534 #[case::u8_lengths(RunLengthWidth::U8)]
2535 #[case::u16_lengths(RunLengthWidth::U16)]
2536 #[case::u32_lengths(RunLengthWidth::U32)]
2537 fn test_miniblock_chunk_counts_match_encoded_runs(#[case] run_length_width: RunLengthWidth) {
2538 let levels = (0..4098)
2540 .map(|index| if index % 3 == 0 { 1u16 } else { 0u16 })
2541 .collect::<Vec<_>>();
2542 let num_values = levels.len() as u64;
2543 let encoder = RleEncoder::with_run_length_width(run_length_width);
2544 let (buffers, chunks) = encoder
2545 .encode_data(
2546 &LanceBuffer::reinterpret_vec(levels),
2547 num_values,
2548 u16::BITS as u64,
2549 )
2550 .unwrap();
2551
2552 assert_eq!(buffers.len(), 2);
2553 let bytes_per_length = run_length_width.bytes_per_value();
2554 let mut values_offset = 0usize;
2555 let mut lengths_offset = 0usize;
2556 let mut values_processed = 0u64;
2557
2558 for chunk in &chunks {
2559 let values_size = chunk.buffer_sizes[0] as usize;
2560 let lengths_size = chunk.buffer_sizes[1] as usize;
2561 let lengths_end = lengths_offset + lengths_size;
2562 let chunk_lengths = &buffers[1].as_ref()[lengths_offset..lengths_end];
2563 let length_chunks = chunk_lengths.chunks_exact(bytes_per_length);
2564 assert!(length_chunks.remainder().is_empty());
2565 let num_runs = length_chunks.len();
2566 let encoded_values = length_chunks
2567 .map(|bytes| run_length_width.read_length(bytes))
2568 .sum::<u64>();
2569 let declared_values = chunk.num_values(values_processed, num_values);
2570
2571 assert_eq!(values_size, num_runs * size_of::<u16>());
2572 assert_eq!(encoded_values, declared_values);
2573
2574 values_offset += values_size;
2575 lengths_offset = lengths_end;
2576 values_processed += declared_values;
2577 }
2578
2579 assert_eq!(values_processed, num_values);
2580 assert_eq!(values_offset, buffers[0].len());
2581 assert_eq!(lengths_offset, buffers[1].len());
2582 }
2583
2584 #[test]
2587 fn test_encoder_rejects_zero_progress() {
2588 let error = RleEncoder::new()
2589 .encode_data(&LanceBuffer::empty(), 1, u16::BITS as u64)
2590 .unwrap_err();
2591
2592 assert!(
2593 matches!(&error, Error::Internal { .. }),
2594 "expected internal error, got: {error:?}"
2595 );
2596 assert!(error.to_string().contains("made no progress"));
2597 assert!(error.to_string().contains("values_remaining=1"));
2598 }
2599
2600 #[test]
2601 fn test_invalid_buffer_count() {
2602 let decompressor = RleDecompressor::new(32);
2603 let result = MiniBlockDecompressor::decompress(
2604 &decompressor,
2605 vec![LanceBuffer::from(vec![1, 2, 3, 4])],
2606 10,
2607 );
2608 assert!(result.is_err());
2609 assert!(
2610 result
2611 .unwrap_err()
2612 .to_string()
2613 .contains("expects exactly 2 buffers")
2614 );
2615 }
2616
2617 #[test]
2618 fn test_buffer_consistency() {
2619 let decompressor = RleDecompressor::new(32);
2620 let values = LanceBuffer::from(vec![1, 0, 0, 0]); let lengths = LanceBuffer::from(vec![5, 10]); let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 15);
2623 assert!(result.is_err());
2624 assert!(
2625 result
2626 .unwrap_err()
2627 .to_string()
2628 .contains("Inconsistent RLE buffers")
2629 );
2630 }
2631
2632 #[test]
2633 fn test_u16_length_buffer_must_be_aligned() {
2634 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2635 let values = LanceBuffer::from(vec![1, 0, 0, 0]);
2636 let lengths = LanceBuffer::from(vec![5]);
2637 let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5);
2638 assert!(matches!(&result, Err(Error::InvalidInput { .. })));
2639 assert!(
2640 result
2641 .unwrap_err()
2642 .to_string()
2643 .contains("not divisible by 2")
2644 );
2645 }
2646
2647 #[test]
2648 fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() {
2649 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2650 let value = LanceBuffer::from(1i32.to_le_bytes().to_vec());
2651
2652 let underflow = MiniBlockDecompressor::decompress(
2653 &decompressor,
2654 vec![
2655 value.clone(),
2656 LanceBuffer::from(4u16.to_le_bytes().to_vec()),
2657 ],
2658 5,
2659 )
2660 .unwrap_err();
2661 assert!(underflow.to_string().contains("produced 4 values"));
2662
2663 let overflow = MiniBlockDecompressor::decompress(
2664 &decompressor,
2665 vec![
2666 value.clone(),
2667 LanceBuffer::from(6u16.to_le_bytes().to_vec()),
2668 ],
2669 5,
2670 )
2671 .unwrap();
2672 match overflow {
2673 DataBlock::FixedWidth(block) => {
2674 assert_eq!(block.num_values, 5);
2675 let decoded = block.data.borrow_to_typed_slice::<i32>();
2676 assert_eq!(decoded.as_ref(), &[1i32; 5]);
2677 }
2678 _ => panic!("Expected FixedWidth block"),
2679 }
2680
2681 let zero = MiniBlockDecompressor::decompress(
2682 &decompressor,
2683 vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())],
2684 5,
2685 )
2686 .unwrap_err();
2687 assert!(zero.to_string().contains("zero run length"));
2688 }
2689
2690 #[test]
2691 fn test_block_rle_rejects_overflow() {
2692 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2695 let values = 1i32.to_le_bytes();
2696 let lengths = 6u16.to_le_bytes();
2697 let mut payload = Vec::new();
2698 payload.extend_from_slice(&(values.len() as u64).to_le_bytes());
2699 payload.extend_from_slice(&values);
2700 payload.extend_from_slice(&lengths);
2701
2702 let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5)
2703 .unwrap_err();
2704 assert!(matches!(&error, Error::InvalidInput { .. }));
2705 assert!(
2706 error
2707 .to_string()
2708 .contains("overflowed expected value count")
2709 );
2710 }
2711
2712 #[test]
2713 fn test_rle_truncates_legacy_chunk_boundary_overflow() {
2714 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2718 let mut values = Vec::new();
2719 values.extend_from_slice(&7i32.to_le_bytes());
2720 values.extend_from_slice(&8i32.to_le_bytes());
2721 let mut lengths = Vec::new();
2722 lengths.extend_from_slice(&2000u16.to_le_bytes());
2723 lengths.extend_from_slice(&80u16.to_le_bytes());
2724
2725 let decoded = MiniBlockDecompressor::decompress(
2726 &decompressor,
2727 vec![LanceBuffer::from(values), LanceBuffer::from(lengths)],
2728 2048,
2729 )
2730 .unwrap();
2731 match decoded {
2732 DataBlock::FixedWidth(block) => {
2733 assert_eq!(block.num_values, 2048);
2734 let decoded = block.data.borrow_to_typed_slice::<i32>();
2735 let decoded = decoded.as_ref();
2736 assert_eq!(decoded.len(), 2048);
2737 assert!(decoded[..2000].iter().all(|&v| v == 7));
2738 assert!(decoded[2000..].iter().all(|&v| v == 8));
2739 }
2740 _ => panic!("Expected FixedWidth block"),
2741 }
2742 }
2743
2744 #[test]
2745 fn test_empty_data_handling() {
2746 let encoder = RleEncoder::new();
2747
2748 let empty_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2750 bits_per_value: 32,
2751 data: LanceBuffer::from(vec![]),
2752 num_values: 0,
2753 block_info: BlockInfo::default(),
2754 });
2755
2756 let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap();
2757 assert_eq!(compressed.num_values, 0);
2758 assert!(compressed.data.is_empty());
2759
2760 let decompressor = RleDecompressor::new(32);
2762 let decompressed = MiniBlockDecompressor::decompress(&decompressor, vec![], 0).unwrap();
2763
2764 match decompressed {
2765 DataBlock::FixedWidth(ref block) => {
2766 assert_eq!(block.num_values, 0);
2767 assert_eq!(block.data.len(), 0);
2768 }
2769 _ => panic!("Expected FixedWidth block"),
2770 }
2771 }
2772
2773 #[test]
2776 fn test_multi_chunk_round_trip() {
2777 let encoder = RleEncoder::new();
2778
2779 let mut data = Vec::new();
2781
2782 data.extend(vec![999i32; 2000]);
2784 data.extend(0..1000);
2786 data.extend(vec![777i32; 2000]);
2788
2789 let array = Int32Array::from(data.clone());
2790 let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2791
2792 let mut reconstructed = Vec::new();
2794 let mut values_offset = 0usize;
2795 let mut lengths_offset = 0usize;
2796 let mut values_processed = 0u64;
2797
2798 assert_eq!(compressed.data.len(), 2);
2800 let global_values = &compressed.data[0];
2801 let global_lengths = &compressed.data[1];
2802
2803 for chunk in &compressed.chunks {
2804 let chunk_values = if chunk.log_num_values > 0 {
2805 1u64 << chunk.log_num_values
2806 } else {
2807 compressed.num_values - values_processed
2808 };
2809
2810 let values_size = chunk.buffer_sizes[0] as usize;
2812 let lengths_size = chunk.buffer_sizes[1] as usize;
2813
2814 let chunk_values_buffer = global_values.slice_with_length(values_offset, values_size);
2815 let chunk_lengths_buffer =
2816 global_lengths.slice_with_length(lengths_offset, lengths_size);
2817
2818 let decompressor = RleDecompressor::new(32);
2819 let chunk_data = MiniBlockDecompressor::decompress(
2820 &decompressor,
2821 vec![chunk_values_buffer, chunk_lengths_buffer],
2822 chunk_values,
2823 )
2824 .unwrap();
2825
2826 values_offset += values_size;
2827 lengths_offset += lengths_size;
2828 values_processed += chunk_values;
2829
2830 match chunk_data {
2831 DataBlock::FixedWidth(ref block) => {
2832 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2833 reconstructed.extend_from_slice(values);
2834 }
2835 _ => panic!("Expected FixedWidth block"),
2836 }
2837 }
2838
2839 assert_eq!(reconstructed, data);
2840 }
2841
2842 #[test]
2843 fn test_1024_boundary_conditions() {
2844 let encoder = RleEncoder::new();
2847 let decompressor = RleDecompressor::new(32);
2848
2849 let test_cases = [
2850 ("runs_of_2", {
2851 let mut data = Vec::new();
2852 for i in 0..512 {
2853 data.push(i);
2854 data.push(i);
2855 }
2856 data
2857 }),
2858 ("single_run_1024", vec![42i32; 1024]),
2859 ("alternating_values", {
2860 let mut data = Vec::new();
2861 for i in 0..1024 {
2862 data.push(i % 2);
2863 }
2864 data
2865 }),
2866 ("run_boundary_255s", {
2867 let mut data = Vec::new();
2868 data.extend(vec![1i32; 255]);
2869 data.extend(vec![2i32; 255]);
2870 data.extend(vec![3i32; 255]);
2871 data.extend(vec![4i32; 255]);
2872 data.extend(vec![5i32; 4]);
2873 data
2874 }),
2875 ("unique_values_1024", (0..1024).collect::<Vec<_>>()),
2876 ("unique_plus_duplicate", {
2877 let mut data = Vec::new();
2879 for i in 0..1023 {
2880 data.push(i);
2881 }
2882 data.push(1022i32); data
2884 }),
2885 ("bug_4092_pattern", {
2886 let mut data = Vec::new();
2888 for i in 0..1022 {
2889 data.push(i);
2890 }
2891 data.push(999999i32);
2892 data.push(999999i32);
2893 data
2894 }),
2895 ];
2896
2897 for (test_name, data) in test_cases.iter() {
2898 assert_eq!(data.len(), 1024, "Test case {} has wrong length", test_name);
2899
2900 let array = Int32Array::from(data.clone());
2902 let (compressed, _) =
2903 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2904
2905 match MiniBlockDecompressor::decompress(
2907 &decompressor,
2908 compressed.data,
2909 compressed.num_values,
2910 ) {
2911 Ok(decompressed) => match decompressed {
2912 DataBlock::FixedWidth(ref block) => {
2913 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2914 assert_eq!(
2915 values.len(),
2916 1024,
2917 "Test case {} got {} values, expected 1024",
2918 test_name,
2919 values.len()
2920 );
2921 assert_eq!(
2922 block.data.len(),
2923 4096,
2924 "Test case {} got {} bytes, expected 4096",
2925 test_name,
2926 block.data.len()
2927 );
2928 assert_eq!(values, &data[..], "Test case {} data mismatch", test_name);
2929 }
2930 _ => panic!("Test case {} expected FixedWidth block", test_name),
2931 },
2932 Err(e) => {
2933 if e.to_string().contains("4092") {
2934 panic!("Test case {} found bug 4092: {}", test_name, e);
2935 }
2936 panic!("Test case {} failed with error: {}", test_name, e);
2937 }
2938 }
2939 }
2940 }
2941
2942 #[test]
2943 fn test_low_repetition_50pct_bug() {
2944 let encoder = RleEncoder::new();
2947
2948 let num_values = 1_048_576; let mut data = Vec::with_capacity(num_values);
2951 let mut value = 0i32;
2952 let mut rng = 12345u64; for _ in 0..num_values {
2955 data.push(value);
2956 rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
2958 if (rng >> 16) & 1 == 1 {
2960 value += 1;
2961 }
2962 }
2963
2964 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
2965
2966 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2967 bits_per_value: 32,
2968 data: LanceBuffer::from(bytes),
2969 num_values: num_values as u64,
2970 block_info: BlockInfo::default(),
2971 });
2972
2973 let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
2974
2975 for (i, chunk) in compressed.chunks.iter().take(5).enumerate() {
2977 let _chunk_values = if chunk.log_num_values > 0 {
2978 1 << chunk.log_num_values
2979 } else {
2980 let prev_total: usize = compressed.chunks[..i]
2982 .iter()
2983 .map(|c| 1usize << c.log_num_values)
2984 .sum();
2985 num_values - prev_total
2986 };
2987 }
2988
2989 let decompressor = RleDecompressor::new(32);
2991 match MiniBlockDecompressor::decompress(
2992 &decompressor,
2993 compressed.data,
2994 compressed.num_values,
2995 ) {
2996 Ok(decompressed) => match decompressed {
2997 DataBlock::FixedWidth(ref block) => {
2998 assert_eq!(
2999 block.data.len(),
3000 num_values * 4,
3001 "Expected {} bytes but got {}",
3002 num_values * 4,
3003 block.data.len()
3004 );
3005 }
3006 _ => panic!("Expected FixedWidth block"),
3007 },
3008 Err(e) => {
3009 if e.to_string().contains("4092") {
3010 panic!("Bug reproduced! {}", e);
3011 } else {
3012 panic!("Unexpected error: {}", e);
3013 }
3014 }
3015 }
3016 }
3017
3018 #[test_log::test(tokio::test)]
3021 async fn test_rle_encoding_verification() {
3022 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
3023 use arrow_array::{Array, Int32Array};
3024 use lance_datagen::{ArrayGenerator, RowCount};
3025 use std::collections::HashMap;
3026 use std::sync::Arc;
3027
3028 let test_cases = TestCases::default()
3029 .with_expected_encoding("rle")
3030 .with_structural_encodings();
3031
3032 let mut metadata_explicit = HashMap::new();
3035 metadata_explicit.insert(
3036 "lance-encoding:rle-threshold".to_string(),
3037 "0.8".to_string(),
3038 );
3039 metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
3040
3041 let mut generator = RleDataGenerator::new(vec![
3042 i32::MIN,
3043 i32::MIN,
3044 i32::MIN,
3045 i32::MIN,
3046 i32::MIN + 1,
3047 i32::MIN + 1,
3048 i32::MIN + 1,
3049 i32::MIN + 1,
3050 i32::MIN + 2,
3051 i32::MIN + 2,
3052 i32::MIN + 2,
3053 i32::MIN + 2,
3054 ]);
3055 let data_explicit = generator.generate_default(RowCount::from(10000)).unwrap();
3056 check_round_trip_encoding_of_data(vec![data_explicit], &test_cases, metadata_explicit)
3057 .await;
3058
3059 let mut metadata = HashMap::new();
3065 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3066
3067 let mut values = vec![i32::MIN; 8000]; values.extend(
3069 [
3070 i32::MIN + 1,
3071 i32::MIN + 2,
3072 i32::MIN + 3,
3073 i32::MIN + 4,
3074 i32::MIN + 5,
3075 ]
3076 .repeat(400),
3077 ); let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3079 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3080
3081 #[cfg(any(feature = "lz4", feature = "zstd"))]
3082 {
3083 let mut metadata = HashMap::new();
3084 metadata.insert(
3085 "lance-encoding:rle-threshold".to_string(),
3086 "0.8".to_string(),
3087 );
3088 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3089 metadata.insert(
3090 "lance-encoding:compression".to_string(),
3091 if cfg!(feature = "zstd") {
3092 "zstd".to_string()
3093 } else {
3094 "lz4".to_string()
3095 },
3096 );
3097 let mut values = Vec::with_capacity(2048 * 4);
3098 for run in 0..2048 {
3099 values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4));
3100 }
3101 let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3102 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3103 }
3104 }
3105
3106 #[derive(Debug)]
3108 struct RleDataGenerator {
3109 pattern: Vec<i32>,
3110 idx: usize,
3111 }
3112
3113 impl RleDataGenerator {
3114 fn new(pattern: Vec<i32>) -> Self {
3115 Self { pattern, idx: 0 }
3116 }
3117 }
3118
3119 impl lance_datagen::ArrayGenerator for RleDataGenerator {
3120 fn generate(
3121 &mut self,
3122 _length: lance_datagen::RowCount,
3123 _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
3124 ) -> std::result::Result<std::sync::Arc<dyn arrow_array::Array>, arrow_schema::ArrowError>
3125 {
3126 use arrow_array::Int32Array;
3127 use std::sync::Arc;
3128
3129 let mut values = Vec::new();
3131 for _ in 0..10000 {
3132 values.push(self.pattern[self.idx]);
3133 self.idx = (self.idx + 1) % self.pattern.len();
3134 }
3135 Ok(Arc::new(Int32Array::from(values)))
3136 }
3137
3138 fn data_type(&self) -> &arrow_schema::DataType {
3139 &arrow_schema::DataType::Int32
3140 }
3141
3142 fn element_size_bytes(&self) -> Option<lance_datagen::ByteCount> {
3143 Some(lance_datagen::ByteCount::from(4))
3144 }
3145 }
3146
3147 #[test]
3149 fn test_block_decompressor_rejects_overflowing_values_size() {
3150 let decompressor = RleDecompressor::new(32);
3151
3152 let mut data = Vec::new();
3153 data.extend_from_slice(&u64::MAX.to_le_bytes());
3154 let result = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(data), 1);
3155 assert!(result.is_err());
3156 assert!(
3157 result
3158 .unwrap_err()
3159 .to_string()
3160 .contains("Invalid RLE values buffer size")
3161 );
3162 }
3163
3164 #[test]
3165 fn test_block_decompressor_too_small() {
3166 let decompressor = RleDecompressor::new(32);
3167 let result =
3168 BlockDecompressor::decompress(&decompressor, LanceBuffer::from(vec![1, 2, 3]), 10);
3169 assert!(result.is_err());
3170 assert!(
3171 result
3172 .unwrap_err()
3173 .to_string()
3174 .contains("Insufficient data size: 3")
3175 );
3176 }
3177
3178 #[test]
3179 fn test_block_compressor_header_format() {
3180 let encoder = RleEncoder::new();
3181
3182 let data = vec![1i32, 1, 1];
3183 let array = Int32Array::from(data);
3184 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
3185
3186 assert!(compressed.len() >= 8);
3188 let values_size_bytes: [u8; 8] = compressed.as_ref()[..8].try_into().unwrap();
3189 let values_size = u64::from_le_bytes(values_size_bytes);
3190
3191 assert_eq!(values_size, 4);
3193
3194 assert_eq!(compressed.len(), 13);
3196 }
3197
3198 #[test]
3199 fn test_block_compressor_round_trip() {
3200 let encoder = RleEncoder::new();
3201 let decompressor = RleDecompressor::new(32);
3202
3203 let data = vec![1i32, 1, 1, 2, 2, 3, 3, 3, 3];
3205 let array = Int32Array::from(data.clone());
3206 let data_block = DataBlock::from_array(array);
3207
3208 let compressed = BlockCompressor::compress(&encoder, data_block).unwrap();
3209 let decompressed =
3210 BlockDecompressor::decompress(&decompressor, compressed, data.len() as u64).unwrap();
3211
3212 match decompressed {
3213 DataBlock::FixedWidth(block) => {
3214 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3215 assert_eq!(values, &data[..]);
3216 }
3217 _ => panic!("Expected FixedWidth block"),
3218 }
3219 }
3220
3221 #[test]
3222 fn test_block_compressor_large_data() {
3223 let encoder = RleEncoder::new();
3224 let decompressor = RleDecompressor::new(32);
3225
3226 let mut data = Vec::new();
3229 data.extend(vec![999i32; 3000]); data.extend(vec![777i32; 3000]); data.extend(vec![555i32; 4000]); let total_values = data.len();
3234 assert_eq!(total_values, 10000);
3235
3236 let array = Int32Array::from(data.clone());
3237 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
3238 let decompressed =
3239 BlockDecompressor::decompress(&decompressor, compressed, total_values as u64).unwrap();
3240
3241 match decompressed {
3242 DataBlock::FixedWidth(block) => {
3243 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3244 assert_eq!(values.len(), total_values);
3245 assert_eq!(values, &data[..]);
3246 }
3247 _ => panic!("Expected FixedWidth block"),
3248 }
3249 }
3250}