1use arrow_buffer::ArrowNativeType;
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 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(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1040 match data {
1041 DataBlock::FixedWidth(fixed_width) => {
1042 let num_values = fixed_width.num_values;
1043 let bits_per_value = fixed_width.bits_per_value;
1044
1045 let (all_buffers, chunks) =
1046 self.encode_data(&fixed_width.data, num_values, bits_per_value)?;
1047 if all_buffers.is_empty() {
1048 let compressed = MiniBlockCompressed {
1049 data: all_buffers,
1050 chunks,
1051 num_values,
1052 };
1053 let encoding = ProtobufUtils21::rle(
1054 ProtobufUtils21::flat(bits_per_value, None),
1055 ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None),
1056 );
1057 return Ok((compressed, encoding));
1058 }
1059
1060 let values_candidates = Self::child_candidates(
1061 &all_buffers,
1062 &chunks,
1063 0,
1064 bits_per_value,
1065 self.values_compression,
1066 self.use_child_bitpacking,
1067 )?;
1068 let run_lengths_candidates = Self::child_candidates(
1069 &all_buffers,
1070 &chunks,
1071 1,
1072 self.run_length_width.bits_per_value(),
1073 self.run_lengths_compression,
1074 self.use_child_bitpacking,
1075 )?;
1076 let (values, run_lengths) =
1077 Self::select_child_candidates(values_candidates, run_lengths_candidates);
1078 let chunks = chunks
1079 .into_iter()
1080 .enumerate()
1081 .map(|(idx, chunk)| MiniBlockChunk {
1082 buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]],
1083 log_num_values: chunk.log_num_values,
1084 })
1085 .collect();
1086
1087 let compressed = MiniBlockCompressed {
1088 data: vec![values.data, run_lengths.data],
1089 chunks,
1090 num_values,
1091 };
1092
1093 let encoding = ProtobufUtils21::rle(values.encoding, run_lengths.encoding);
1094
1095 Ok((compressed, encoding))
1096 }
1097 _ => Err(Error::invalid_input_source(
1098 "RLE encoding only supports FixedWidth data blocks".into(),
1099 )),
1100 }
1101 }
1102}
1103
1104impl BlockCompressor for RleEncoder {
1105 fn compress(&self, data: DataBlock) -> Result<LanceBuffer> {
1107 match data {
1108 DataBlock::FixedWidth(fixed_width) => {
1109 let num_values = fixed_width.num_values;
1110 let bits_per_value = fixed_width.bits_per_value;
1111
1112 let all_buffers =
1113 self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?;
1114
1115 let values_size = all_buffers[0].len() as u64;
1116
1117 let mut combined = Vec::new();
1118 combined.extend_from_slice(&values_size.to_le_bytes());
1119 combined.extend_from_slice(&all_buffers[0]);
1120 combined.extend_from_slice(&all_buffers[1]);
1121 Ok(LanceBuffer::from(combined))
1122 }
1123 _ => Err(Error::invalid_input_source(
1124 "RLE encoding only supports FixedWidth data blocks".into(),
1125 )),
1126 }
1127 }
1128}
1129
1130#[derive(Debug)]
1132pub struct RleDecompressor {
1133 bits_per_value: u64,
1134 run_length_width: RunLengthWidth,
1135 values: RleChildDecompressor,
1136 run_lengths: RleChildDecompressor,
1137}
1138
1139#[derive(Debug)]
1140pub(crate) struct RleChildDecompressor {
1141 bits_per_value: u64,
1142 inner: RleChildDecompressorInner,
1143}
1144
1145#[derive(Debug)]
1146enum RleChildDecompressorInner {
1147 Flat,
1148 Block {
1149 decompressor: Box<dyn BlockDecompressor>,
1150 requires_num_values: bool,
1151 },
1152}
1153
1154impl RleChildDecompressor {
1155 pub(crate) fn flat(bits_per_value: u64) -> Self {
1156 Self {
1157 bits_per_value,
1158 inner: RleChildDecompressorInner::Flat,
1159 }
1160 }
1161
1162 pub(crate) fn block(
1163 bits_per_value: u64,
1164 decompressor: Box<dyn BlockDecompressor>,
1165 requires_num_values: bool,
1166 ) -> Self {
1167 Self {
1168 bits_per_value,
1169 inner: RleChildDecompressorInner::Block {
1170 decompressor,
1171 requires_num_values,
1172 },
1173 }
1174 }
1175
1176 pub(crate) fn bits_per_value(&self) -> u64 {
1177 self.bits_per_value
1178 }
1179
1180 pub(crate) fn requires_num_values(&self) -> bool {
1181 match &self.inner {
1182 RleChildDecompressorInner::Flat => false,
1183 RleChildDecompressorInner::Block {
1184 requires_num_values,
1185 ..
1186 } => *requires_num_values,
1187 }
1188 }
1189
1190 pub(crate) fn is_identity(&self) -> bool {
1191 matches!(self.inner, RleChildDecompressorInner::Flat)
1192 }
1193
1194 fn decode(
1195 &self,
1196 data: LanceBuffer,
1197 num_values: Option<u64>,
1198 label: &str,
1199 ) -> Result<LanceBuffer> {
1200 match &self.inner {
1201 RleChildDecompressorInner::Flat => Ok(data),
1202 RleChildDecompressorInner::Block {
1203 decompressor,
1204 requires_num_values,
1205 } => {
1206 let num_values = if *requires_num_values {
1207 num_values.ok_or_else(|| {
1208 Error::invalid_input_source(
1209 format!("RLE {label} child compression requires the run count").into(),
1210 )
1211 })?
1212 } else {
1213 num_values.unwrap_or(0)
1214 };
1215 let decoded = decompressor.decompress(data, num_values)?;
1216 self.extract_fixed_width(decoded, num_values, label)
1217 }
1218 }
1219 }
1220
1221 fn extract_fixed_width(
1222 &self,
1223 data: DataBlock,
1224 expected_num_values: u64,
1225 label: &str,
1226 ) -> Result<LanceBuffer> {
1227 match data {
1228 DataBlock::FixedWidth(block) => {
1229 if block.bits_per_value != self.bits_per_value {
1230 return Err(Error::invalid_input_source(
1231 format!(
1232 "RLE {label} child decoded {}-bit values, expected {}",
1233 block.bits_per_value, self.bits_per_value
1234 )
1235 .into(),
1236 ));
1237 }
1238 if expected_num_values != 0 && block.num_values != expected_num_values {
1239 return Err(Error::invalid_input_source(
1240 format!(
1241 "RLE {label} child decoded {} values, expected {}",
1242 block.num_values, expected_num_values
1243 )
1244 .into(),
1245 ));
1246 }
1247 Ok(block.data)
1248 }
1249 _ => Err(Error::invalid_input_source(
1250 format!("RLE {label} child decoded to a non fixed-width block").into(),
1251 )),
1252 }
1253 }
1254}
1255
1256impl RleDecompressor {
1257 pub fn new(bits_per_value: u64) -> Self {
1258 Self {
1259 bits_per_value,
1260 run_length_width: RunLengthWidth::U8,
1261 values: RleChildDecompressor::flat(bits_per_value),
1262 run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()),
1263 }
1264 }
1265
1266 pub(crate) fn with_run_length_width(
1267 bits_per_value: u64,
1268 run_length_width: RunLengthWidth,
1269 ) -> Self {
1270 Self {
1271 bits_per_value,
1272 run_length_width,
1273 values: RleChildDecompressor::flat(bits_per_value),
1274 run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()),
1275 }
1276 }
1277
1278 pub(crate) fn with_child_decompressors(
1279 bits_per_value: u64,
1280 run_length_width: RunLengthWidth,
1281 values: RleChildDecompressor,
1282 run_lengths: RleChildDecompressor,
1283 ) -> Self {
1284 Self {
1285 bits_per_value,
1286 run_length_width,
1287 values,
1288 run_lengths,
1289 }
1290 }
1291
1292 fn decode_data(
1293 &self,
1294 data: Vec<LanceBuffer>,
1295 num_values: u64,
1296 clamp_overflow: bool,
1297 ) -> Result<DataBlock> {
1298 if num_values == 0 {
1299 return Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1300 bits_per_value: self.bits_per_value,
1301 data: LanceBuffer::from(vec![]),
1302 num_values: 0,
1303 block_info: BlockInfo::default(),
1304 }));
1305 }
1306
1307 if data.len() != 2 {
1308 return Err(Error::invalid_input_source(
1309 format!(
1310 "RLE decompressor expects exactly 2 buffers, got {}",
1311 data.len()
1312 )
1313 .into(),
1314 ));
1315 }
1316
1317 let mut data_iter = data.into_iter();
1318 let values_buffer = data_iter.next().unwrap();
1319 let lengths_buffer = data_iter.next().unwrap();
1320 let (values_buffer, lengths_buffer) =
1321 self.decode_child_buffers(values_buffer, lengths_buffer)?;
1322
1323 let decoded_data = match self.bits_per_value {
1324 8 => self.decode_generic::<u8>(
1325 &values_buffer,
1326 &lengths_buffer,
1327 num_values,
1328 clamp_overflow,
1329 )?,
1330 16 => self.decode_generic::<u16>(
1331 &values_buffer,
1332 &lengths_buffer,
1333 num_values,
1334 clamp_overflow,
1335 )?,
1336 32 => self.decode_generic::<u32>(
1337 &values_buffer,
1338 &lengths_buffer,
1339 num_values,
1340 clamp_overflow,
1341 )?,
1342 64 => self.decode_generic::<u64>(
1343 &values_buffer,
1344 &lengths_buffer,
1345 num_values,
1346 clamp_overflow,
1347 )?,
1348 _ => {
1349 return Err(Error::invalid_input_source(
1350 format!(
1351 "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}",
1352 self.bits_per_value
1353 )
1354 .into(),
1355 ));
1356 }
1357 };
1358
1359 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1360 bits_per_value: self.bits_per_value,
1361 data: decoded_data,
1362 num_values,
1363 block_info: BlockInfo::default(),
1364 }))
1365 }
1366
1367 fn decode_child_buffers(
1368 &self,
1369 values_buffer: LanceBuffer,
1370 lengths_buffer: LanceBuffer,
1371 ) -> Result<(LanceBuffer, LanceBuffer)> {
1372 let values_requires_num_runs = self.values.requires_num_values();
1373 let lengths_requires_num_runs = self.run_lengths.requires_num_values();
1374 if values_requires_num_runs && lengths_requires_num_runs {
1375 return Err(Error::invalid_input_source(
1376 "RLE values and run lengths child compression both require the run count".into(),
1377 ));
1378 }
1379
1380 if values_requires_num_runs {
1381 let lengths_buffer = self
1382 .run_lengths
1383 .decode(lengths_buffer, None, "run lengths")?;
1384 let num_runs = Self::num_child_values(
1385 &lengths_buffer,
1386 self.run_lengths.bits_per_value(),
1387 "run lengths",
1388 )?;
1389 let values_buffer = self
1390 .values
1391 .decode(values_buffer, Some(num_runs), "values")?;
1392 Ok((values_buffer, lengths_buffer))
1393 } else if lengths_requires_num_runs {
1394 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1395 let num_runs =
1396 Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?;
1397 let lengths_buffer =
1398 self.run_lengths
1399 .decode(lengths_buffer, Some(num_runs), "run lengths")?;
1400 Ok((values_buffer, lengths_buffer))
1401 } else {
1402 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1403 let lengths_buffer = self
1404 .run_lengths
1405 .decode(lengths_buffer, None, "run lengths")?;
1406 Ok((values_buffer, lengths_buffer))
1407 }
1408 }
1409
1410 fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result<u64> {
1411 let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
1412 Error::invalid_input_source(
1413 format!("RLE {label} child bit width is too large: {bits_per_value}").into(),
1414 )
1415 })?;
1416 if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) {
1417 return Err(Error::invalid_input_source(
1418 format!(
1419 "RLE {label} child decoded to {} bytes, not divisible by {}",
1420 buffer.len(),
1421 bytes_per_value
1422 )
1423 .into(),
1424 ));
1425 }
1426 Ok((buffer.len() / bytes_per_value) as u64)
1427 }
1428
1429 fn decode_generic<T>(
1430 &self,
1431 values_buffer: &LanceBuffer,
1432 lengths_buffer: &LanceBuffer,
1433 num_values: u64,
1434 clamp_overflow: bool,
1435 ) -> Result<LanceBuffer>
1436 where
1437 T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType,
1438 {
1439 let type_size = std::mem::size_of::<T>();
1440 let length_size = self.run_length_width.bytes_per_value();
1441
1442 if values_buffer.is_empty() || lengths_buffer.is_empty() {
1443 if num_values == 0 {
1444 return Ok(LanceBuffer::empty());
1445 } else {
1446 return Err(Error::invalid_input_source(
1447 format!("Empty buffers but expected {} values", num_values).into(),
1448 ));
1449 }
1450 }
1451
1452 if !values_buffer.len().is_multiple_of(type_size)
1453 || !lengths_buffer.len().is_multiple_of(length_size)
1454 {
1455 return Err(Error::invalid_input_source(format!(
1456 "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})",
1457 std::any::type_name::<T>(),
1458 values_buffer.len(),
1459 type_size,
1460 lengths_buffer.len(),
1461 length_size
1462 )
1463 .into()));
1464 }
1465
1466 let num_runs = values_buffer.len() / type_size;
1467 let num_length_entries = lengths_buffer.len() / length_size;
1468 if num_runs != num_length_entries {
1469 return Err(Error::invalid_input_source(
1470 format!(
1471 "Inconsistent RLE buffers: {} runs but {} length entries",
1472 num_runs, num_length_entries
1473 )
1474 .into(),
1475 ));
1476 }
1477
1478 let values_ref = values_buffer.borrow_to_typed_slice::<T>();
1479 let values: &[T] = values_ref.as_ref();
1480 let lengths = lengths_buffer.as_ref();
1481
1482 let expected_value_count = usize::try_from(num_values).map_err(|_| {
1483 Error::invalid_input_source(
1484 format!("RLE num_values does not fit in usize: {num_values}").into(),
1485 )
1486 })?;
1487 let mut decoded: Vec<T> = Vec::new();
1494 decoded
1495 .try_reserve_exact(expected_value_count)
1496 .map_err(|_| {
1497 Error::invalid_input_source(
1498 format!("RLE decoding cannot allocate {expected_value_count} values").into(),
1499 )
1500 })?;
1501 for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) {
1502 let length = self.run_length_width.read_length(length_bytes);
1503 if length == 0 {
1504 return Err(Error::invalid_input_source(
1505 "RLE decoding encountered a zero run length".into(),
1506 ));
1507 }
1508 let length = usize::try_from(length).map_err(|_| {
1509 Error::invalid_input_source(
1510 format!("RLE run length does not fit in usize: {length}").into(),
1511 )
1512 })?;
1513 let remaining = expected_value_count - decoded.len();
1514 if length > remaining {
1515 if !clamp_overflow {
1516 return Err(Error::invalid_input_source(
1517 format!(
1518 "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1519 decoded.len() + length,
1520 expected_value_count
1521 )
1522 .into(),
1523 ));
1524 }
1525 decoded.resize(expected_value_count, *value);
1526 break;
1527 }
1528 decoded.resize(decoded.len() + length, *value);
1529 }
1530
1531 if decoded.len() != expected_value_count {
1532 return Err(Error::invalid_input_source(
1533 format!(
1534 "RLE decoding produced {} values, expected {}",
1535 decoded.len(),
1536 expected_value_count
1537 )
1538 .into(),
1539 ));
1540 }
1541
1542 trace!(
1543 "RLE decoded {} {} values",
1544 num_values,
1545 std::any::type_name::<T>()
1546 );
1547 Ok(LanceBuffer::reinterpret_vec(decoded))
1548 }
1549}
1550
1551impl MiniBlockDecompressor for RleDecompressor {
1552 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
1553 self.decode_data(data, num_values, true)
1554 }
1555}
1556
1557impl BlockDecompressor for RleDecompressor {
1558 fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result<DataBlock> {
1559 if data.len() < 8 {
1561 return Err(Error::invalid_input_source(
1562 format!("Insufficient data size: {}", data.len()).into(),
1563 ));
1564 }
1565
1566 let values_size_bytes: [u8; 8] =
1567 data[..8].try_into().expect("slice length already checked");
1568 let values_size: u64 = u64::from_le_bytes(values_size_bytes);
1569
1570 let values_start: usize = 8;
1572 let values_size: usize = values_size.try_into().map_err(|_| {
1573 Error::invalid_input_source(
1574 format!("Invalid values buffer size: {}", values_size).into(),
1575 )
1576 })?;
1577 let lengths_start = values_start
1578 .checked_add(values_size)
1579 .ok_or_else(|| Error::invalid_input_source("Invalid RLE values buffer size".into()))?;
1580
1581 if data.len() < lengths_start {
1582 return Err(Error::invalid_input_source(
1583 format!("Insufficient data size: {}", data.len()).into(),
1584 ));
1585 }
1586
1587 let values_buffer = data.slice_with_length(values_start, values_size);
1588 let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start);
1589
1590 self.decode_data(vec![values_buffer, lengths_buffer], num_values, false)
1591 }
1592}
1593
1594#[cfg(test)]
1595mod tests {
1596 use super::*;
1597 use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy};
1598 use crate::data::DataBlock;
1599 use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES;
1600 use crate::encodings::physical::block::{CompressionConfig, CompressionScheme};
1601 use crate::{
1602 buffer::LanceBuffer,
1603 compression::{BlockCompressor, BlockDecompressor},
1604 };
1605 use arrow_array::Int32Array;
1606 use rstest::rstest;
1607
1608 #[test]
1611 fn test_basic_miniblock_rle_encoding() {
1612 let encoder = RleEncoder::new();
1613
1614 let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]);
1616 let data_block = DataBlock::from_array(array);
1617
1618 let (compressed, _) = MiniBlockCompressor::compress(&encoder, data_block).unwrap();
1619
1620 assert_eq!(compressed.num_values, 9);
1621 assert_eq!(compressed.chunks.len(), 1);
1622
1623 let values_buffer = &compressed.data[0];
1625 let lengths_buffer = &compressed.data[1];
1626 assert_eq!(values_buffer.len(), 12); assert_eq!(lengths_buffer.len(), 3); }
1629
1630 #[test]
1631 fn test_long_run_splitting() {
1632 let encoder = RleEncoder::new();
1633
1634 let mut data = vec![42i32; 1000]; data.extend(&[100i32; 300]); let array = Int32Array::from(data);
1639 let (compressed, _) =
1640 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
1641
1642 let lengths_buffer = &compressed.data[1];
1644 assert_eq!(lengths_buffer.len(), 6);
1645 }
1646
1647 #[test]
1648 fn test_rle_v2_u16_miniblock_encoding() {
1649 let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16);
1650
1651 let data = vec![42i32; 1000];
1652 let array = Int32Array::from(data);
1653 let (compressed, encoding) =
1654 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
1655
1656 assert_eq!(compressed.data[0].len(), 4);
1657 assert_eq!(compressed.data[1].len(), 2);
1658 assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes());
1659
1660 let rle = match encoding.compression.as_ref().unwrap() {
1661 crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
1662 other => panic!("expected RLE encoding, got {other:?}"),
1663 };
1664 let run_lengths = rle.run_lengths.as_ref().unwrap();
1665 let flat = match run_lengths.compression.as_ref().unwrap() {
1666 crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat,
1667 other => panic!("expected flat run lengths, got {other:?}"),
1668 };
1669 assert_eq!(flat.bits_per_value, 16);
1670
1671 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
1672 let decompressed = MiniBlockDecompressor::decompress(
1673 &decompressor,
1674 compressed.data,
1675 compressed.num_values,
1676 )
1677 .unwrap();
1678 match decompressed {
1679 DataBlock::FixedWidth(block) => {
1680 let values = block.data.borrow_to_typed_slice::<i32>();
1681 assert_eq!(values.as_ref(), vec![42i32; 1000]);
1682 }
1683 _ => panic!("Expected FixedWidth block"),
1684 }
1685 }
1686
1687 #[test]
1688 #[cfg(any(feature = "lz4", feature = "zstd"))]
1689 fn test_rle_miniblock_compressed_values_child() {
1690 let compression = test_general_compression();
1691 let encoder =
1692 RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false);
1693 let array = Int32Array::from(repeating_runs(1024, 4));
1694 let (compressed, encoding) =
1695 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
1696
1697 let rle = expect_rle(&encoding);
1698 assert!(matches!(
1699 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
1700 crate::format::pb21::compressive_encoding::Compression::General(_)
1701 ));
1702 assert!(matches!(
1703 rle.run_lengths
1704 .as_ref()
1705 .unwrap()
1706 .compression
1707 .as_ref()
1708 .unwrap(),
1709 crate::format::pb21::compressive_encoding::Compression::Flat(_)
1710 ));
1711
1712 let decompressor = DefaultDecompressionStrategy::default()
1713 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
1714 .unwrap();
1715 let decoded =
1716 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
1717 .unwrap();
1718 assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4));
1719 }
1720
1721 #[test]
1722 #[cfg(any(feature = "lz4", feature = "zstd"))]
1723 fn test_rle_miniblock_compressed_run_lengths_child() {
1724 let compression = test_general_compression();
1725 let encoder =
1726 RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false);
1727 let expected = repeating_runs(1024, 4);
1728 let (compressed, encoding) = MiniBlockCompressor::compress(
1729 &encoder,
1730 DataBlock::from_array(Int32Array::from(expected.clone())),
1731 )
1732 .unwrap();
1733
1734 let rle = expect_rle(&encoding);
1735 assert!(matches!(
1736 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
1737 crate::format::pb21::compressive_encoding::Compression::Flat(_)
1738 ));
1739 assert!(matches!(
1740 rle.run_lengths
1741 .as_ref()
1742 .unwrap()
1743 .compression
1744 .as_ref()
1745 .unwrap(),
1746 crate::format::pb21::compressive_encoding::Compression::General(_)
1747 ));
1748
1749 let decompressor = DefaultDecompressionStrategy::default()
1750 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
1751 .unwrap();
1752 let decoded =
1753 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
1754 .unwrap();
1755 assert_decoded_i32_eq(decoded, &expected);
1756 }
1757
1758 #[test]
1759 #[cfg(feature = "bitpacking")]
1760 fn test_rle_miniblock_bitpacked_run_lengths_child() {
1761 use crate::encodings::physical::bitpacking::OutOfLineBitpacking;
1762
1763 let expected = repeating_runs(1024, 4);
1764 let (compressed, _) = MiniBlockCompressor::compress(
1765 &RleEncoder::new(),
1766 DataBlock::from_array(Int32Array::from(expected.clone())),
1767 )
1768 .unwrap();
1769 let run_lengths = compressed.data[1].clone();
1770 let num_runs = run_lengths.len() as u64;
1771 let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock {
1772 bits_per_value: 8,
1773 data: run_lengths,
1774 num_values: num_runs,
1775 block_info: BlockInfo::default(),
1776 });
1777 let bitpacked_run_lengths =
1778 BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap();
1779 let encoding = ProtobufUtils21::rle(
1780 ProtobufUtils21::flat(32, None),
1781 ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
1782 );
1783
1784 let decompressor = DefaultDecompressionStrategy::default()
1785 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
1786 .unwrap();
1787 let decoded = MiniBlockDecompressor::decompress(
1788 decompressor.as_ref(),
1789 vec![compressed.data[0].clone(), bitpacked_run_lengths],
1790 expected.len() as u64,
1791 )
1792 .unwrap();
1793 assert_decoded_i32_eq(decoded, &expected);
1794 }
1795
1796 #[test]
1797 #[cfg(feature = "bitpacking")]
1798 fn test_rle_rejects_two_count_dependent_child_encodings() {
1799 let encoding = ProtobufUtils21::rle(
1800 ProtobufUtils21::out_of_line_bitpacking(32, ProtobufUtils21::flat(3, None)),
1801 ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
1802 );
1803
1804 let err = DefaultDecompressionStrategy::default()
1805 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
1806 .unwrap_err();
1807 assert!(
1808 err.to_string()
1809 .contains("cannot both require the run count")
1810 );
1811 }
1812
1813 #[cfg(any(feature = "lz4", feature = "zstd"))]
1814 fn test_general_compression() -> CompressionConfig {
1815 if cfg!(feature = "zstd") {
1816 CompressionConfig::new(CompressionScheme::Zstd, Some(3))
1817 } else {
1818 CompressionConfig::new(CompressionScheme::Lz4, None)
1819 }
1820 }
1821
1822 fn repeating_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
1823 let mut values = Vec::with_capacity(num_runs * run_length);
1824 for run in 0..num_runs {
1825 values.extend(std::iter::repeat_n((run % 8) as i32, run_length));
1826 }
1827 values
1828 }
1829
1830 fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle {
1831 match encoding.compression.as_ref().unwrap() {
1832 crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
1833 other => panic!("expected RLE encoding, got {other:?}"),
1834 }
1835 }
1836
1837 fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) {
1838 match decoded {
1839 DataBlock::FixedWidth(block) => {
1840 let values = block.data.borrow_to_typed_slice::<i32>();
1841 assert_eq!(values.as_ref(), expected);
1842 }
1843 _ => panic!("Expected FixedWidth block"),
1844 }
1845 }
1846
1847 #[test]
1848 #[cfg(any(feature = "lz4", feature = "zstd"))]
1849 fn test_rle_miniblock_compressed_children_multiple_chunks() {
1850 let compression = test_general_compression();
1851 let encoder = RleEncoder::with_child_encoding(
1852 RunLengthWidth::U8,
1853 Some(compression),
1854 Some(compression),
1855 false,
1856 );
1857 let expected = repeating_runs(8192, 4);
1858 let (compressed, encoding) = MiniBlockCompressor::compress(
1859 &encoder,
1860 DataBlock::from_array(Int32Array::from(expected.clone())),
1861 )
1862 .unwrap();
1863
1864 assert!(compressed.chunks.len() > 1);
1865 let rle = expect_rle(&encoding);
1866 assert!(matches!(
1867 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
1868 crate::format::pb21::compressive_encoding::Compression::General(_)
1869 ));
1870 assert!(matches!(
1871 rle.run_lengths
1872 .as_ref()
1873 .unwrap()
1874 .compression
1875 .as_ref()
1876 .unwrap(),
1877 crate::format::pb21::compressive_encoding::Compression::General(_)
1878 ));
1879
1880 let decoded = decompress_i32_chunks(&compressed, &encoding);
1881 assert_eq!(decoded, expected);
1882 }
1883
1884 #[test]
1885 #[cfg(feature = "bitpacking")]
1886 fn test_rle_miniblock_bitpacks_values_child_when_smaller() {
1887 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
1888 let expected = monotonic_runs(2048, 4);
1889 let (compressed, encoding) = MiniBlockCompressor::compress(
1890 &encoder,
1891 DataBlock::from_array(Int32Array::from(expected.clone())),
1892 )
1893 .unwrap();
1894
1895 let rle = expect_rle(&encoding);
1896 assert!(matches!(
1897 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
1898 crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
1899 ));
1900 assert!(matches!(
1901 rle.run_lengths
1902 .as_ref()
1903 .unwrap()
1904 .compression
1905 .as_ref()
1906 .unwrap(),
1907 crate::format::pb21::compressive_encoding::Compression::Flat(_)
1908 ));
1909
1910 let decoded = decompress_i32_chunks(&compressed, &encoding);
1911 assert_eq!(decoded, expected);
1912 }
1913
1914 #[test]
1915 #[cfg(feature = "bitpacking")]
1916 fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() {
1917 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
1918 let expected = high_entropy_runs(2048, 4);
1919 let (compressed, encoding) = MiniBlockCompressor::compress(
1920 &encoder,
1921 DataBlock::from_array(Int32Array::from(expected.clone())),
1922 )
1923 .unwrap();
1924
1925 let rle = expect_rle(&encoding);
1926 assert!(matches!(
1927 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
1928 crate::format::pb21::compressive_encoding::Compression::Flat(_)
1929 ));
1930 assert!(matches!(
1931 rle.run_lengths
1932 .as_ref()
1933 .unwrap()
1934 .compression
1935 .as_ref()
1936 .unwrap(),
1937 crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
1938 ));
1939
1940 let decoded = decompress_i32_chunks(&compressed, &encoding);
1941 assert_eq!(decoded, expected);
1942 }
1943
1944 fn decompress_i32_chunks(
1945 compressed: &MiniBlockCompressed,
1946 encoding: &CompressiveEncoding,
1947 ) -> Vec<i32> {
1948 let strategy = DefaultDecompressionStrategy::default();
1949 let decompressor = strategy
1950 .create_miniblock_decompressor(encoding, &strategy)
1951 .unwrap();
1952 let mut offsets = vec![0usize; compressed.data.len()];
1953 let mut values_processed = 0u64;
1954 let mut decoded_values = Vec::new();
1955
1956 for chunk in &compressed.chunks {
1957 let chunk_values = chunk.num_values(values_processed, compressed.num_values);
1958 let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len());
1959 for (idx, size) in chunk.buffer_sizes.iter().enumerate() {
1960 let size = *size as usize;
1961 chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size));
1962 offsets[idx] += size;
1963 }
1964
1965 let decoded = decompressor
1966 .decompress(chunk_buffers, chunk_values)
1967 .unwrap();
1968 match decoded {
1969 DataBlock::FixedWidth(block) => {
1970 let values = block.data.borrow_to_typed_slice::<i32>();
1971 decoded_values.extend_from_slice(values.as_ref());
1972 }
1973 _ => panic!("Expected FixedWidth block"),
1974 }
1975 values_processed += chunk_values;
1976 }
1977
1978 assert_eq!(values_processed, compressed.num_values);
1979 decoded_values
1980 }
1981
1982 #[cfg(feature = "bitpacking")]
1983 fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
1984 let mut values = Vec::with_capacity(num_runs * run_length);
1985 for run in 0..num_runs {
1986 values.extend(std::iter::repeat_n(run as i32, run_length));
1987 }
1988 values
1989 }
1990
1991 #[cfg(feature = "bitpacking")]
1992 fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
1993 let mut values = Vec::with_capacity(num_runs * run_length);
1994 let mut state = 7u64;
1995 for _ in 0..num_runs {
1996 state = state
1997 .wrapping_mul(6364136223846793005)
1998 .wrapping_add(1442695040888963407);
1999 values.extend(std::iter::repeat_n((state >> 32) as i32, run_length));
2000 }
2001 values
2002 }
2003
2004 #[test]
2005 fn test_select_run_length_width_prefers_u16_for_long_runs() {
2006 let mut entries = [0u64; 3];
2007 accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries);
2008 let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap();
2009 assert_eq!(width, RunLengthWidth::U16);
2010 }
2011
2012 #[test]
2015 fn test_round_trip_all_types() {
2016 test_round_trip_helper(vec![42u8, 42, 42, 100, 100, 255, 255, 255, 255], 8);
2018
2019 test_round_trip_helper(vec![1000u16, 1000, 2000, 2000, 2000, 3000], 16);
2021
2022 test_round_trip_helper(vec![100i32, 100, 100, -200, -200, 300, 300, 300, 300], 32);
2024
2025 test_round_trip_helper(vec![1_000_000_000u64; 5], 64);
2027 }
2028
2029 fn test_round_trip_helper<T>(data: Vec<T>, bits_per_value: u64)
2030 where
2031 T: bytemuck::Pod + PartialEq + std::fmt::Debug,
2032 {
2033 let encoder = RleEncoder::new();
2034 let bytes: Vec<u8> = data
2035 .iter()
2036 .flat_map(|v| bytemuck::bytes_of(v))
2037 .copied()
2038 .collect();
2039
2040 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2041 bits_per_value,
2042 data: LanceBuffer::from(bytes),
2043 num_values: data.len() as u64,
2044 block_info: BlockInfo::default(),
2045 });
2046
2047 let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap();
2048 let decompressor = RleDecompressor::new(bits_per_value);
2049 let decompressed = MiniBlockDecompressor::decompress(
2050 &decompressor,
2051 compressed.data,
2052 compressed.num_values,
2053 )
2054 .unwrap();
2055
2056 match decompressed {
2057 DataBlock::FixedWidth(ref block) => {
2058 assert_eq!(block.data.len(), data.len() * std::mem::size_of::<T>());
2060 }
2061 _ => panic!("Expected FixedWidth block"),
2062 }
2063 }
2064
2065 #[test]
2068 fn test_power_of_two_chunking() {
2069 let encoder = RleEncoder::new();
2070
2071 let test_sizes = vec![1000, 2500, 5000, 10000];
2073
2074 for size in test_sizes {
2075 let data: Vec<i32> = (0..size)
2076 .map(|i| i / 50) .collect();
2078
2079 let array = Int32Array::from(data);
2080 let (compressed, _) =
2081 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
2082
2083 for (i, chunk) in compressed.chunks.iter().enumerate() {
2085 if i < compressed.chunks.len() - 1 {
2086 assert!(chunk.log_num_values > 0);
2087 let chunk_values = 1u64 << chunk.log_num_values;
2088 assert!(chunk_values.is_power_of_two());
2089 assert!(chunk_values <= *MAX_MINIBLOCK_VALUES);
2090 } else {
2091 assert_eq!(chunk.log_num_values, 0);
2092 }
2093 }
2094 }
2095 }
2096
2097 #[rstest]
2098 #[case::u8_lengths(RunLengthWidth::U8)]
2099 #[case::u16_lengths(RunLengthWidth::U16)]
2100 #[case::u32_lengths(RunLengthWidth::U32)]
2101 fn test_miniblock_chunk_counts_match_encoded_runs(#[case] run_length_width: RunLengthWidth) {
2102 let levels = (0..4098)
2104 .map(|index| if index % 3 == 0 { 1u16 } else { 0u16 })
2105 .collect::<Vec<_>>();
2106 let num_values = levels.len() as u64;
2107 let encoder = RleEncoder::with_run_length_width(run_length_width);
2108 let (buffers, chunks) = encoder
2109 .encode_data(
2110 &LanceBuffer::reinterpret_vec(levels),
2111 num_values,
2112 u16::BITS as u64,
2113 )
2114 .unwrap();
2115
2116 assert_eq!(buffers.len(), 2);
2117 let bytes_per_length = run_length_width.bytes_per_value();
2118 let mut values_offset = 0usize;
2119 let mut lengths_offset = 0usize;
2120 let mut values_processed = 0u64;
2121
2122 for chunk in &chunks {
2123 let values_size = chunk.buffer_sizes[0] as usize;
2124 let lengths_size = chunk.buffer_sizes[1] as usize;
2125 let lengths_end = lengths_offset + lengths_size;
2126 let chunk_lengths = &buffers[1].as_ref()[lengths_offset..lengths_end];
2127 let length_chunks = chunk_lengths.chunks_exact(bytes_per_length);
2128 assert!(length_chunks.remainder().is_empty());
2129 let num_runs = length_chunks.len();
2130 let encoded_values = length_chunks
2131 .map(|bytes| run_length_width.read_length(bytes))
2132 .sum::<u64>();
2133 let declared_values = chunk.num_values(values_processed, num_values);
2134
2135 assert_eq!(values_size, num_runs * size_of::<u16>());
2136 assert_eq!(encoded_values, declared_values);
2137
2138 values_offset += values_size;
2139 lengths_offset = lengths_end;
2140 values_processed += declared_values;
2141 }
2142
2143 assert_eq!(values_processed, num_values);
2144 assert_eq!(values_offset, buffers[0].len());
2145 assert_eq!(lengths_offset, buffers[1].len());
2146 }
2147
2148 #[test]
2151 fn test_encoder_rejects_zero_progress() {
2152 let error = RleEncoder::new()
2153 .encode_data(&LanceBuffer::empty(), 1, u16::BITS as u64)
2154 .unwrap_err();
2155
2156 assert!(
2157 matches!(&error, Error::Internal { .. }),
2158 "expected internal error, got: {error:?}"
2159 );
2160 assert!(error.to_string().contains("made no progress"));
2161 assert!(error.to_string().contains("values_remaining=1"));
2162 }
2163
2164 #[test]
2165 fn test_invalid_buffer_count() {
2166 let decompressor = RleDecompressor::new(32);
2167 let result = MiniBlockDecompressor::decompress(
2168 &decompressor,
2169 vec![LanceBuffer::from(vec![1, 2, 3, 4])],
2170 10,
2171 );
2172 assert!(result.is_err());
2173 assert!(
2174 result
2175 .unwrap_err()
2176 .to_string()
2177 .contains("expects exactly 2 buffers")
2178 );
2179 }
2180
2181 #[test]
2182 fn test_buffer_consistency() {
2183 let decompressor = RleDecompressor::new(32);
2184 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);
2187 assert!(result.is_err());
2188 assert!(
2189 result
2190 .unwrap_err()
2191 .to_string()
2192 .contains("Inconsistent RLE buffers")
2193 );
2194 }
2195
2196 #[test]
2197 fn test_u16_length_buffer_must_be_aligned() {
2198 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2199 let values = LanceBuffer::from(vec![1, 0, 0, 0]);
2200 let lengths = LanceBuffer::from(vec![5]);
2201 let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5);
2202 assert!(matches!(&result, Err(Error::InvalidInput { .. })));
2203 assert!(
2204 result
2205 .unwrap_err()
2206 .to_string()
2207 .contains("not divisible by 2")
2208 );
2209 }
2210
2211 #[test]
2212 fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() {
2213 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2214 let value = LanceBuffer::from(1i32.to_le_bytes().to_vec());
2215
2216 let underflow = MiniBlockDecompressor::decompress(
2217 &decompressor,
2218 vec![
2219 value.clone(),
2220 LanceBuffer::from(4u16.to_le_bytes().to_vec()),
2221 ],
2222 5,
2223 )
2224 .unwrap_err();
2225 assert!(underflow.to_string().contains("produced 4 values"));
2226
2227 let overflow = MiniBlockDecompressor::decompress(
2228 &decompressor,
2229 vec![
2230 value.clone(),
2231 LanceBuffer::from(6u16.to_le_bytes().to_vec()),
2232 ],
2233 5,
2234 )
2235 .unwrap();
2236 match overflow {
2237 DataBlock::FixedWidth(block) => {
2238 assert_eq!(block.num_values, 5);
2239 let decoded = block.data.borrow_to_typed_slice::<i32>();
2240 assert_eq!(decoded.as_ref(), &[1i32; 5]);
2241 }
2242 _ => panic!("Expected FixedWidth block"),
2243 }
2244
2245 let zero = MiniBlockDecompressor::decompress(
2246 &decompressor,
2247 vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())],
2248 5,
2249 )
2250 .unwrap_err();
2251 assert!(zero.to_string().contains("zero run length"));
2252 }
2253
2254 #[test]
2255 fn test_block_rle_rejects_overflow() {
2256 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2259 let values = 1i32.to_le_bytes();
2260 let lengths = 6u16.to_le_bytes();
2261 let mut payload = Vec::new();
2262 payload.extend_from_slice(&(values.len() as u64).to_le_bytes());
2263 payload.extend_from_slice(&values);
2264 payload.extend_from_slice(&lengths);
2265
2266 let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5)
2267 .unwrap_err();
2268 assert!(matches!(&error, Error::InvalidInput { .. }));
2269 assert!(
2270 error
2271 .to_string()
2272 .contains("overflowed expected value count")
2273 );
2274 }
2275
2276 #[test]
2277 fn test_rle_truncates_legacy_chunk_boundary_overflow() {
2278 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2282 let mut values = Vec::new();
2283 values.extend_from_slice(&7i32.to_le_bytes());
2284 values.extend_from_slice(&8i32.to_le_bytes());
2285 let mut lengths = Vec::new();
2286 lengths.extend_from_slice(&2000u16.to_le_bytes());
2287 lengths.extend_from_slice(&80u16.to_le_bytes());
2288
2289 let decoded = MiniBlockDecompressor::decompress(
2290 &decompressor,
2291 vec![LanceBuffer::from(values), LanceBuffer::from(lengths)],
2292 2048,
2293 )
2294 .unwrap();
2295 match decoded {
2296 DataBlock::FixedWidth(block) => {
2297 assert_eq!(block.num_values, 2048);
2298 let decoded = block.data.borrow_to_typed_slice::<i32>();
2299 let decoded = decoded.as_ref();
2300 assert_eq!(decoded.len(), 2048);
2301 assert!(decoded[..2000].iter().all(|&v| v == 7));
2302 assert!(decoded[2000..].iter().all(|&v| v == 8));
2303 }
2304 _ => panic!("Expected FixedWidth block"),
2305 }
2306 }
2307
2308 #[test]
2309 fn test_empty_data_handling() {
2310 let encoder = RleEncoder::new();
2311
2312 let empty_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2314 bits_per_value: 32,
2315 data: LanceBuffer::from(vec![]),
2316 num_values: 0,
2317 block_info: BlockInfo::default(),
2318 });
2319
2320 let (compressed, _) = MiniBlockCompressor::compress(&encoder, empty_block).unwrap();
2321 assert_eq!(compressed.num_values, 0);
2322 assert!(compressed.data.is_empty());
2323
2324 let decompressor = RleDecompressor::new(32);
2326 let decompressed = MiniBlockDecompressor::decompress(&decompressor, vec![], 0).unwrap();
2327
2328 match decompressed {
2329 DataBlock::FixedWidth(ref block) => {
2330 assert_eq!(block.num_values, 0);
2331 assert_eq!(block.data.len(), 0);
2332 }
2333 _ => panic!("Expected FixedWidth block"),
2334 }
2335 }
2336
2337 #[test]
2340 fn test_multi_chunk_round_trip() {
2341 let encoder = RleEncoder::new();
2342
2343 let mut data = Vec::new();
2345
2346 data.extend(vec![999i32; 2000]);
2348 data.extend(0..1000);
2350 data.extend(vec![777i32; 2000]);
2352
2353 let array = Int32Array::from(data.clone());
2354 let (compressed, _) =
2355 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
2356
2357 let mut reconstructed = Vec::new();
2359 let mut values_offset = 0usize;
2360 let mut lengths_offset = 0usize;
2361 let mut values_processed = 0u64;
2362
2363 assert_eq!(compressed.data.len(), 2);
2365 let global_values = &compressed.data[0];
2366 let global_lengths = &compressed.data[1];
2367
2368 for chunk in &compressed.chunks {
2369 let chunk_values = if chunk.log_num_values > 0 {
2370 1u64 << chunk.log_num_values
2371 } else {
2372 compressed.num_values - values_processed
2373 };
2374
2375 let values_size = chunk.buffer_sizes[0] as usize;
2377 let lengths_size = chunk.buffer_sizes[1] as usize;
2378
2379 let chunk_values_buffer = global_values.slice_with_length(values_offset, values_size);
2380 let chunk_lengths_buffer =
2381 global_lengths.slice_with_length(lengths_offset, lengths_size);
2382
2383 let decompressor = RleDecompressor::new(32);
2384 let chunk_data = MiniBlockDecompressor::decompress(
2385 &decompressor,
2386 vec![chunk_values_buffer, chunk_lengths_buffer],
2387 chunk_values,
2388 )
2389 .unwrap();
2390
2391 values_offset += values_size;
2392 lengths_offset += lengths_size;
2393 values_processed += chunk_values;
2394
2395 match chunk_data {
2396 DataBlock::FixedWidth(ref block) => {
2397 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2398 reconstructed.extend_from_slice(values);
2399 }
2400 _ => panic!("Expected FixedWidth block"),
2401 }
2402 }
2403
2404 assert_eq!(reconstructed, data);
2405 }
2406
2407 #[test]
2408 fn test_1024_boundary_conditions() {
2409 let encoder = RleEncoder::new();
2412 let decompressor = RleDecompressor::new(32);
2413
2414 let test_cases = [
2415 ("runs_of_2", {
2416 let mut data = Vec::new();
2417 for i in 0..512 {
2418 data.push(i);
2419 data.push(i);
2420 }
2421 data
2422 }),
2423 ("single_run_1024", vec![42i32; 1024]),
2424 ("alternating_values", {
2425 let mut data = Vec::new();
2426 for i in 0..1024 {
2427 data.push(i % 2);
2428 }
2429 data
2430 }),
2431 ("run_boundary_255s", {
2432 let mut data = Vec::new();
2433 data.extend(vec![1i32; 255]);
2434 data.extend(vec![2i32; 255]);
2435 data.extend(vec![3i32; 255]);
2436 data.extend(vec![4i32; 255]);
2437 data.extend(vec![5i32; 4]);
2438 data
2439 }),
2440 ("unique_values_1024", (0..1024).collect::<Vec<_>>()),
2441 ("unique_plus_duplicate", {
2442 let mut data = Vec::new();
2444 for i in 0..1023 {
2445 data.push(i);
2446 }
2447 data.push(1022i32); data
2449 }),
2450 ("bug_4092_pattern", {
2451 let mut data = Vec::new();
2453 for i in 0..1022 {
2454 data.push(i);
2455 }
2456 data.push(999999i32);
2457 data.push(999999i32);
2458 data
2459 }),
2460 ];
2461
2462 for (test_name, data) in test_cases.iter() {
2463 assert_eq!(data.len(), 1024, "Test case {} has wrong length", test_name);
2464
2465 let array = Int32Array::from(data.clone());
2467 let (compressed, _) =
2468 MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
2469
2470 match MiniBlockDecompressor::decompress(
2472 &decompressor,
2473 compressed.data,
2474 compressed.num_values,
2475 ) {
2476 Ok(decompressed) => match decompressed {
2477 DataBlock::FixedWidth(ref block) => {
2478 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2479 assert_eq!(
2480 values.len(),
2481 1024,
2482 "Test case {} got {} values, expected 1024",
2483 test_name,
2484 values.len()
2485 );
2486 assert_eq!(
2487 block.data.len(),
2488 4096,
2489 "Test case {} got {} bytes, expected 4096",
2490 test_name,
2491 block.data.len()
2492 );
2493 assert_eq!(values, &data[..], "Test case {} data mismatch", test_name);
2494 }
2495 _ => panic!("Test case {} expected FixedWidth block", test_name),
2496 },
2497 Err(e) => {
2498 if e.to_string().contains("4092") {
2499 panic!("Test case {} found bug 4092: {}", test_name, e);
2500 }
2501 panic!("Test case {} failed with error: {}", test_name, e);
2502 }
2503 }
2504 }
2505 }
2506
2507 #[test]
2508 fn test_low_repetition_50pct_bug() {
2509 let encoder = RleEncoder::new();
2512
2513 let num_values = 1_048_576; let mut data = Vec::with_capacity(num_values);
2516 let mut value = 0i32;
2517 let mut rng = 12345u64; for _ in 0..num_values {
2520 data.push(value);
2521 rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
2523 if (rng >> 16) & 1 == 1 {
2525 value += 1;
2526 }
2527 }
2528
2529 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
2530
2531 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2532 bits_per_value: 32,
2533 data: LanceBuffer::from(bytes),
2534 num_values: num_values as u64,
2535 block_info: BlockInfo::default(),
2536 });
2537
2538 let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap();
2539
2540 for (i, chunk) in compressed.chunks.iter().take(5).enumerate() {
2542 let _chunk_values = if chunk.log_num_values > 0 {
2543 1 << chunk.log_num_values
2544 } else {
2545 let prev_total: usize = compressed.chunks[..i]
2547 .iter()
2548 .map(|c| 1usize << c.log_num_values)
2549 .sum();
2550 num_values - prev_total
2551 };
2552 }
2553
2554 let decompressor = RleDecompressor::new(32);
2556 match MiniBlockDecompressor::decompress(
2557 &decompressor,
2558 compressed.data,
2559 compressed.num_values,
2560 ) {
2561 Ok(decompressed) => match decompressed {
2562 DataBlock::FixedWidth(ref block) => {
2563 assert_eq!(
2564 block.data.len(),
2565 num_values * 4,
2566 "Expected {} bytes but got {}",
2567 num_values * 4,
2568 block.data.len()
2569 );
2570 }
2571 _ => panic!("Expected FixedWidth block"),
2572 },
2573 Err(e) => {
2574 if e.to_string().contains("4092") {
2575 panic!("Bug reproduced! {}", e);
2576 } else {
2577 panic!("Unexpected error: {}", e);
2578 }
2579 }
2580 }
2581 }
2582
2583 #[test_log::test(tokio::test)]
2586 async fn test_rle_encoding_verification() {
2587 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
2588 use crate::version::LanceFileVersion;
2589 use arrow_array::{Array, Int32Array};
2590 use lance_datagen::{ArrayGenerator, RowCount};
2591 use std::collections::HashMap;
2592 use std::sync::Arc;
2593
2594 let test_cases = TestCases::default()
2595 .with_expected_encoding("rle")
2596 .with_min_file_version(LanceFileVersion::V2_1);
2597
2598 let mut metadata_explicit = HashMap::new();
2601 metadata_explicit.insert(
2602 "lance-encoding:rle-threshold".to_string(),
2603 "0.8".to_string(),
2604 );
2605 metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
2606
2607 let mut generator = RleDataGenerator::new(vec![
2608 i32::MIN,
2609 i32::MIN,
2610 i32::MIN,
2611 i32::MIN,
2612 i32::MIN + 1,
2613 i32::MIN + 1,
2614 i32::MIN + 1,
2615 i32::MIN + 1,
2616 i32::MIN + 2,
2617 i32::MIN + 2,
2618 i32::MIN + 2,
2619 i32::MIN + 2,
2620 ]);
2621 let data_explicit = generator.generate_default(RowCount::from(10000)).unwrap();
2622 check_round_trip_encoding_of_data(vec![data_explicit], &test_cases, metadata_explicit)
2623 .await;
2624
2625 let mut metadata = HashMap::new();
2631 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
2632
2633 let mut values = vec![i32::MIN; 8000]; values.extend(
2635 [
2636 i32::MIN + 1,
2637 i32::MIN + 2,
2638 i32::MIN + 3,
2639 i32::MIN + 4,
2640 i32::MIN + 5,
2641 ]
2642 .repeat(400),
2643 ); let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
2645 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
2646
2647 #[cfg(any(feature = "lz4", feature = "zstd"))]
2648 {
2649 let mut metadata = HashMap::new();
2650 metadata.insert(
2651 "lance-encoding:rle-threshold".to_string(),
2652 "0.8".to_string(),
2653 );
2654 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
2655 metadata.insert(
2656 "lance-encoding:compression".to_string(),
2657 if cfg!(feature = "zstd") {
2658 "zstd".to_string()
2659 } else {
2660 "lz4".to_string()
2661 },
2662 );
2663 let mut values = Vec::with_capacity(2048 * 4);
2664 for run in 0..2048 {
2665 values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4));
2666 }
2667 let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
2668 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
2669 }
2670 }
2671
2672 #[derive(Debug)]
2674 struct RleDataGenerator {
2675 pattern: Vec<i32>,
2676 idx: usize,
2677 }
2678
2679 impl RleDataGenerator {
2680 fn new(pattern: Vec<i32>) -> Self {
2681 Self { pattern, idx: 0 }
2682 }
2683 }
2684
2685 impl lance_datagen::ArrayGenerator for RleDataGenerator {
2686 fn generate(
2687 &mut self,
2688 _length: lance_datagen::RowCount,
2689 _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
2690 ) -> std::result::Result<std::sync::Arc<dyn arrow_array::Array>, arrow_schema::ArrowError>
2691 {
2692 use arrow_array::Int32Array;
2693 use std::sync::Arc;
2694
2695 let mut values = Vec::new();
2697 for _ in 0..10000 {
2698 values.push(self.pattern[self.idx]);
2699 self.idx = (self.idx + 1) % self.pattern.len();
2700 }
2701 Ok(Arc::new(Int32Array::from(values)))
2702 }
2703
2704 fn data_type(&self) -> &arrow_schema::DataType {
2705 &arrow_schema::DataType::Int32
2706 }
2707
2708 fn element_size_bytes(&self) -> Option<lance_datagen::ByteCount> {
2709 Some(lance_datagen::ByteCount::from(4))
2710 }
2711 }
2712
2713 #[test]
2715 fn test_block_decompressor_rejects_overflowing_values_size() {
2716 let decompressor = RleDecompressor::new(32);
2717
2718 let mut data = Vec::new();
2719 data.extend_from_slice(&u64::MAX.to_le_bytes());
2720 let result = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(data), 1);
2721 assert!(result.is_err());
2722 assert!(
2723 result
2724 .unwrap_err()
2725 .to_string()
2726 .contains("Invalid RLE values buffer size")
2727 );
2728 }
2729
2730 #[test]
2731 fn test_block_decompressor_too_small() {
2732 let decompressor = RleDecompressor::new(32);
2733 let result =
2734 BlockDecompressor::decompress(&decompressor, LanceBuffer::from(vec![1, 2, 3]), 10);
2735 assert!(result.is_err());
2736 assert!(
2737 result
2738 .unwrap_err()
2739 .to_string()
2740 .contains("Insufficient data size: 3")
2741 );
2742 }
2743
2744 #[test]
2745 fn test_block_compressor_header_format() {
2746 let encoder = RleEncoder::new();
2747
2748 let data = vec![1i32, 1, 1];
2749 let array = Int32Array::from(data);
2750 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
2751
2752 assert!(compressed.len() >= 8);
2754 let values_size_bytes: [u8; 8] = compressed.as_ref()[..8].try_into().unwrap();
2755 let values_size = u64::from_le_bytes(values_size_bytes);
2756
2757 assert_eq!(values_size, 4);
2759
2760 assert_eq!(compressed.len(), 13);
2762 }
2763
2764 #[test]
2765 fn test_block_compressor_round_trip() {
2766 let encoder = RleEncoder::new();
2767 let decompressor = RleDecompressor::new(32);
2768
2769 let data = vec![1i32, 1, 1, 2, 2, 3, 3, 3, 3];
2771 let array = Int32Array::from(data.clone());
2772 let data_block = DataBlock::from_array(array);
2773
2774 let compressed = BlockCompressor::compress(&encoder, data_block).unwrap();
2775 let decompressed =
2776 BlockDecompressor::decompress(&decompressor, compressed, data.len() as u64).unwrap();
2777
2778 match decompressed {
2779 DataBlock::FixedWidth(block) => {
2780 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2781 assert_eq!(values, &data[..]);
2782 }
2783 _ => panic!("Expected FixedWidth block"),
2784 }
2785 }
2786
2787 #[test]
2788 fn test_block_compressor_large_data() {
2789 let encoder = RleEncoder::new();
2790 let decompressor = RleDecompressor::new(32);
2791
2792 let mut data = Vec::new();
2795 data.extend(vec![999i32; 3000]); data.extend(vec![777i32; 3000]); data.extend(vec![555i32; 4000]); let total_values = data.len();
2800 assert_eq!(total_values, 10000);
2801
2802 let array = Int32Array::from(data.clone());
2803 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
2804 let decompressed =
2805 BlockDecompressor::decompress(&decompressor, compressed, total_values as u64).unwrap();
2806
2807 match decompressed {
2808 DataBlock::FixedWidth(block) => {
2809 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2810 assert_eq!(values.len(), total_values);
2811 assert_eq!(values, &data[..]);
2812 }
2813 _ => panic!("Expected FixedWidth block"),
2814 }
2815 }
2816}