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