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::EncodingUtils;
72use crate::format::metadata::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 select_miniblock_run_length_width(
208 data: &FixedWidthDataBlock,
209 full_limit: u64,
210 composite_limit: u64,
211) -> Result<(RunLengthWidth, u128, u64)> {
212 fn collect<T: Copy + PartialEq>(
213 values: &[T],
214 full_limit: u64,
215 composite_limit: u64,
216 ) -> ([u64; 3], [u64; 3]) {
217 let mut full = [0; 3];
218 let mut composite = [0; 3];
219 let mut add_run = |length| {
220 let mut entries = [0; 3];
221 accumulate_run_length_entries(length, Some(full_limit), &mut entries);
222 for (total, count) in full.iter_mut().zip(entries) {
223 *total += count;
224 }
225 if length <= composite_limit {
226 for (total, count) in composite.iter_mut().zip(entries) {
227 *total += count;
228 }
229 } else {
230 accumulate_run_length_entries(length, Some(composite_limit), &mut composite);
231 }
232 };
233 if let Some((&first, rest)) = values.split_first() {
234 let mut previous = first;
235 let mut length = 1;
236 for &value in rest {
237 if value == previous {
238 length += 1;
239 } else {
240 add_run(length);
241 previous = value;
242 length = 1;
243 }
244 }
245 add_run(length);
246 }
247 (full, composite)
248 }
249 use super::checked_fixed_values;
250 let (full, composite) = match data.bits_per_value {
251 8 => collect(
252 &checked_fixed_values::<u8>(data, "RLE values")?,
253 full_limit,
254 composite_limit,
255 ),
256 16 => collect(
257 &checked_fixed_values::<u16>(data, "RLE values")?,
258 full_limit,
259 composite_limit,
260 ),
261 32 => collect(
262 &checked_fixed_values::<u32>(data, "RLE values")?,
263 full_limit,
264 composite_limit,
265 ),
266 64 => collect(
267 &checked_fixed_values::<u64>(data, "RLE values")?,
268 full_limit,
269 composite_limit,
270 ),
271 bits => {
272 return Err(Error::invalid_input(format!(
273 "Unsupported RLE value width: {bits}"
274 )));
275 }
276 };
277 let (width, bytes) = select_run_length_width_from_entries(&full, data.bits_per_value)?;
278 if width == RunLengthWidth::U8 || full_limit == composite_limit {
279 Ok((width, bytes, full_limit))
280 } else {
281 let (width, bytes) = select_run_length_width_from_entries(&composite, data.bits_per_value)?;
282 Ok((width, bytes, composite_limit))
283 }
284}
285
286pub(crate) fn rle_encoded_size(
287 data: &LanceBuffer,
288 num_values: u64,
289 bits_per_value: u64,
290 max_segment_values: Option<u64>,
291 run_length_width: RunLengthWidth,
292) -> Result<u128> {
293 let entries = collect_run_length_entries(data, num_values, bits_per_value, max_segment_values)?;
294 let width_idx = run_length_width_index(run_length_width);
295 Ok(rle_encoded_size_from_entries(
296 entries[width_idx],
297 bits_per_value,
298 run_length_width,
299 ))
300}
301
302fn collect_run_length_entries(
303 data: &LanceBuffer,
304 num_values: u64,
305 bits_per_value: u64,
306 max_segment_values: Option<u64>,
307) -> Result<[u64; 3]> {
308 let num_values = usize::try_from(num_values).map_err(|_| {
309 Error::invalid_input_source(
310 format!("RLE num_values does not fit in usize: {num_values}").into(),
311 )
312 })?;
313
314 macro_rules! collect_entries {
315 ($ty:ty) => {{
316 let type_size = std::mem::size_of::<$ty>();
317 let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| {
318 Error::invalid_input_source(
319 format!(
320 "RLE input byte length overflow: {num_values} values of {type_size} bytes"
321 )
322 .into(),
323 )
324 })?;
325 if data.len() != expected_bytes {
326 return Err(Error::invalid_input_source(
327 format!(
328 "RLE input data size mismatch: {} bytes for {} values of {} bytes",
329 data.len(),
330 num_values,
331 type_size
332 )
333 .into(),
334 ));
335 }
336 let values = data.borrow_to_typed_slice::<$ty>();
337 let values = values.get(..num_values).ok_or_else(|| {
338 Error::invalid_input_source(
339 format!(
340 "RLE data has {} values but {} were expected",
341 values.len(),
342 num_values
343 )
344 .into(),
345 )
346 })?;
347 Ok(collect_run_length_entries_from_slice(
348 values,
349 max_segment_values,
350 ))
351 }};
352 }
353
354 match bits_per_value {
355 8 => collect_entries!(u8),
356 16 => collect_entries!(u16),
357 32 => collect_entries!(u32),
358 64 => collect_entries!(u64),
359 _ => Err(Error::invalid_input_source(
360 format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}")
361 .into(),
362 )),
363 }
364}
365
366fn collect_run_length_entries_from_slice<T: PartialEq + Copy>(
367 values: &[T],
368 max_segment_values: Option<u64>,
369) -> [u64; 3] {
370 if values.is_empty() {
371 return [0; 3];
372 }
373
374 let mut entries = [0u64; 3];
375 let mut prev = values[0];
376 let mut current_length = 1u64;
377
378 for &value in &values[1..] {
379 if value != prev {
380 accumulate_run_length_entries(current_length, max_segment_values, &mut entries);
381 prev = value;
382 current_length = 1;
383 } else {
384 current_length += 1;
385 }
386 }
387 accumulate_run_length_entries(current_length, max_segment_values, &mut entries);
388
389 entries
390}
391
392pub(crate) fn accumulate_run_length_entries(
393 run_length: u64,
394 max_segment_values: Option<u64>,
395 entries: &mut [u64; 3],
396) {
397 let max_segment_values = max_segment_values.unwrap_or(run_length).max(1);
398 let mut remaining = run_length;
399 while remaining > 0 {
400 let segment = remaining.min(max_segment_values);
401 for (idx, width) in RUN_LENGTH_WIDTHS.iter().enumerate() {
402 let entry_count = segment.div_ceil(width.max_run_length());
403 entries[idx] = entries[idx].saturating_add(entry_count);
404 }
405 remaining -= segment;
406 }
407}
408
409#[derive(Debug)]
411pub struct RleEncoder {
412 run_length_width: RunLengthWidth,
413 max_miniblock_values: usize,
414 values_compression: Option<CompressionConfig>,
415 run_lengths_compression: Option<CompressionConfig>,
416 use_child_bitpacking: bool,
417}
418
419#[derive(Clone)]
420struct RleChildCandidate {
421 encoding: CompressiveEncoding,
422 data: LanceBuffer,
423 chunk_sizes: Vec<u32>,
424 size: usize,
425 requires_num_values: bool,
426}
427
428impl Default for RleEncoder {
429 fn default() -> Self {
430 Self::new()
431 }
432}
433
434impl RleEncoder {
435 pub fn new() -> Self {
436 Self {
437 run_length_width: RunLengthWidth::U8,
438 max_miniblock_values: *MAX_MINIBLOCK_VALUES as usize,
439 values_compression: None,
440 run_lengths_compression: None,
441 use_child_bitpacking: false,
442 }
443 }
444
445 pub(crate) fn with_run_length_width(run_length_width: RunLengthWidth) -> Self {
446 Self {
447 run_length_width,
448 max_miniblock_values: *MAX_MINIBLOCK_VALUES as usize,
449 values_compression: None,
450 run_lengths_compression: None,
451 use_child_bitpacking: false,
452 }
453 }
454
455 pub(crate) fn with_child_encoding(
456 run_length_width: RunLengthWidth,
457 values_compression: Option<CompressionConfig>,
458 run_lengths_compression: Option<CompressionConfig>,
459 use_child_bitpacking: bool,
460 ) -> Self {
461 Self {
462 run_length_width,
463 max_miniblock_values: *MAX_MINIBLOCK_VALUES as usize,
464 values_compression,
465 run_lengths_compression,
466 use_child_bitpacking,
467 }
468 }
469
470 pub(crate) fn with_max_miniblock_values(mut self, max_values: usize) -> Self {
471 self.max_miniblock_values = self.max_miniblock_values.min(max_values);
472 self
473 }
474
475 fn encode_data(
476 &self,
477 data: &LanceBuffer,
478 num_values: u64,
479 bits_per_value: u64,
480 ) -> Result<(Vec<LanceBuffer>, Vec<MiniBlockChunk>)> {
481 if num_values == 0 {
482 return Ok((Vec::new(), Vec::new()));
483 }
484
485 let num_values = usize::try_from(num_values).map_err(|_| {
486 Error::invalid_input_source(
487 format!("RLE num_values does not fit in usize: {num_values}").into(),
488 )
489 })?;
490 let bytes_per_value = (bits_per_value / 8) as usize;
491 let bytes_per_length = self.run_length_width.bytes_per_value();
492
493 let estimated_runs = num_values / 10;
496 let mut all_values = Vec::with_capacity(estimated_runs * bytes_per_value);
497 let mut all_lengths = Vec::with_capacity(estimated_runs * bytes_per_length);
498 let mut chunks = Vec::new();
499
500 let mut offset = 0usize;
501 let mut values_remaining = num_values;
502
503 while values_remaining > 0 {
504 let values_start = all_values.len();
505 let lengths_start = all_lengths.len();
506
507 let (_num_runs, values_processed, is_last_chunk) = match bits_per_value {
508 8 => self.encode_chunk_rolling::<u8>(
509 data,
510 offset,
511 values_remaining,
512 &mut all_values,
513 &mut all_lengths,
514 ),
515 16 => self.encode_chunk_rolling::<u16>(
516 data,
517 offset,
518 values_remaining,
519 &mut all_values,
520 &mut all_lengths,
521 ),
522 32 => self.encode_chunk_rolling::<u32>(
523 data,
524 offset,
525 values_remaining,
526 &mut all_values,
527 &mut all_lengths,
528 ),
529 64 => self.encode_chunk_rolling::<u64>(
530 data,
531 offset,
532 values_remaining,
533 &mut all_values,
534 &mut all_lengths,
535 ),
536 _ => {
537 return Err(Error::invalid_input_source(
538 format!(
539 "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}"
540 )
541 .into(),
542 ));
543 }
544 };
545
546 if values_processed == 0 {
547 return Err(Error::internal(format!(
550 "RLE encoder made no progress: values_remaining={values_remaining}, \
551 offset={offset}, data_len={}, bits_per_value={bits_per_value}, \
552 max_miniblock_values={}",
553 data.len(),
554 *MAX_MINIBLOCK_VALUES
555 )));
556 }
557
558 let log_num_values = if is_last_chunk {
559 0
560 } else {
561 assert!(
562 values_processed.is_power_of_two(),
563 "Non-last chunk must have power-of-2 values"
564 );
565 values_processed.ilog2() as u8
566 };
567
568 let values_size = all_values.len() - values_start;
569 let lengths_size = all_lengths.len() - lengths_start;
570
571 let chunk = MiniBlockChunk {
572 buffer_sizes: vec![values_size as u32, lengths_size as u32],
573 log_num_values,
574 };
575
576 chunks.push(chunk);
577
578 offset += values_processed;
579 values_remaining -= values_processed;
580 }
581
582 Ok((
584 vec![
585 LanceBuffer::from(all_values),
586 LanceBuffer::from(all_lengths),
587 ],
588 chunks,
589 ))
590 }
591
592 fn encode_block_data(
593 &self,
594 data: &LanceBuffer,
595 num_values: u64,
596 bits_per_value: u64,
597 ) -> Result<Vec<LanceBuffer>> {
598 match bits_per_value {
599 8 => self.encode_block_data_generic::<u8>(data, num_values),
600 16 => self.encode_block_data_generic::<u16>(data, num_values),
601 32 => self.encode_block_data_generic::<u32>(data, num_values),
602 64 => self.encode_block_data_generic::<u64>(data, num_values),
603 _ => Err(Error::invalid_input_source(
604 format!(
605 "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}"
606 )
607 .into(),
608 )),
609 }
610 }
611
612 fn encode_block_data_generic<T>(
613 &self,
614 data: &LanceBuffer,
615 num_values: u64,
616 ) -> Result<Vec<LanceBuffer>>
617 where
618 T: bytemuck::Pod + PartialEq + Copy + ArrowNativeType,
619 {
620 let num_values = usize::try_from(num_values).map_err(|_| {
621 Error::invalid_input_source(
622 format!("RLE num_values does not fit in usize: {num_values}").into(),
623 )
624 })?;
625 let type_size = std::mem::size_of::<T>();
626 let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| {
627 Error::invalid_input_source(
628 format!("RLE input byte length overflow: {num_values} values of {type_size} bytes")
629 .into(),
630 )
631 })?;
632 if data.len() != expected_bytes {
633 return Err(Error::invalid_input_source(
634 format!(
635 "RLE input data size mismatch: {} bytes for {} values of {} bytes",
636 data.len(),
637 num_values,
638 type_size
639 )
640 .into(),
641 ));
642 }
643 if num_values == 0 {
644 return Ok(vec![LanceBuffer::empty(), LanceBuffer::empty()]);
645 }
646
647 let values_ref = data.borrow_to_typed_slice::<T>();
648 let values = values_ref.as_ref();
649 let estimated_runs = num_values / 10;
650 let mut all_values = Vec::with_capacity(estimated_runs * type_size);
651 let mut all_lengths =
652 Vec::with_capacity(estimated_runs * self.run_length_width.bytes_per_value());
653 self.encode_values(values, &mut all_values, &mut all_lengths);
654 Ok(vec![
655 LanceBuffer::from(all_values),
656 LanceBuffer::from(all_lengths),
657 ])
658 }
659
660 fn encode_chunk_rolling<T>(
662 &self,
663 data: &LanceBuffer,
664 offset: usize,
665 values_remaining: usize,
666 all_values: &mut Vec<u8>,
667 all_lengths: &mut Vec<u8>,
668 ) -> (usize, usize, bool)
669 where
670 T: bytemuck::Pod + PartialEq + Copy + std::fmt::Debug + ArrowNativeType,
671 {
672 let type_size = std::mem::size_of::<T>();
673 let chunk_start = offset * type_size;
674 let max_by_count = self.max_miniblock_values;
675 let max_values = values_remaining.min(max_by_count);
676 let chunk_end = chunk_start + max_values * type_size;
677
678 if chunk_start >= data.len() {
679 return (0, 0, false);
680 }
681
682 let chunk_len = chunk_end.min(data.len()) - chunk_start;
683 let chunk_buffer = data.slice_with_length(chunk_start, chunk_len);
684 let typed_data_ref = chunk_buffer.borrow_to_typed_slice::<T>();
685 let typed_data: &[T] = typed_data_ref.as_ref();
686 let max_values = max_values.min(typed_data.len());
687
688 if typed_data.is_empty() {
689 return (0, 0, false);
690 }
691
692 let values_start = all_values.len();
693 let all_remaining_values_fit = values_remaining <= max_by_count;
694 let encoded_size = self.encoded_size(&typed_data[..max_values]);
695 let (values_to_encode, is_last_chunk) = if all_remaining_values_fit
696 && encoded_size <= MAX_MINIBLOCK_BYTES as usize
697 {
698 (max_values, true)
699 } else if let Some(values_to_encode) = self.largest_power_of_two_prefix::<T>(typed_data) {
700 (values_to_encode, false)
701 } else {
702 return (0, 0, false);
703 };
704
705 self.encode_values(&typed_data[..values_to_encode], all_values, all_lengths);
706
707 let num_runs = (all_values.len() - values_start) / type_size;
708 (num_runs, values_to_encode, is_last_chunk)
709 }
710
711 fn largest_power_of_two_prefix<T>(&self, values: &[T]) -> Option<usize>
712 where
713 T: bytemuck::Pod + PartialEq + Copy,
714 {
715 let max_prefix = values.len().min(self.max_miniblock_values);
716 let mut prefix = 1usize << max_prefix.ilog2();
717 while prefix > 1 {
718 if self.encoded_size(&values[..prefix]) <= MAX_MINIBLOCK_BYTES as usize {
719 return Some(prefix);
720 }
721 prefix >>= 1;
722 }
723 None
724 }
725
726 fn encoded_size<T>(&self, values: &[T]) -> usize
727 where
728 T: bytemuck::Pod + PartialEq + Copy,
729 {
730 if values.is_empty() {
731 return 0;
732 }
733
734 let mut current_value = values[0];
735 let mut current_length = 1u64;
736 let mut encoded_size = 0usize;
737
738 for &value in values.iter().skip(1) {
739 if value == current_value {
740 current_length += 1;
741 } else {
742 encoded_size += self.run_size::<T>(current_length);
743 current_value = value;
744 current_length = 1;
745 }
746 }
747 encoded_size += self.run_size::<T>(current_length);
748 encoded_size
749 }
750
751 fn run_size<T>(&self, length: u64) -> usize
752 where
753 T: bytemuck::Pod,
754 {
755 let type_size = std::mem::size_of::<T>();
756 let run_chunks = length.div_ceil(self.run_length_width.max_run_length()) as usize;
757 run_chunks * (type_size + self.run_length_width.bytes_per_value())
758 }
759
760 fn encode_values<T>(&self, values: &[T], all_values: &mut Vec<u8>, all_lengths: &mut Vec<u8>)
761 where
762 T: bytemuck::Pod + PartialEq + Copy,
763 {
764 if values.is_empty() {
765 return;
766 }
767
768 let mut current_value = values[0];
769 let mut current_length = 1u64;
770
771 for &value in values.iter().skip(1) {
772 if value == current_value {
773 current_length += 1;
774 } else {
775 self.add_run(¤t_value, current_length, all_values, all_lengths);
776 current_value = value;
777 current_length = 1;
778 }
779 }
780 self.add_run(¤t_value, current_length, all_values, all_lengths);
781 }
782
783 fn add_run<T>(
784 &self,
785 value: &T,
786 length: u64,
787 all_values: &mut Vec<u8>,
788 all_lengths: &mut Vec<u8>,
789 ) -> usize
790 where
791 T: bytemuck::Pod,
792 {
793 let value_bytes = bytemuck::bytes_of(value);
794 let type_size = std::mem::size_of::<T>();
795 let max_run_length = self.run_length_width.max_run_length();
796 let num_full_chunks = (length / max_run_length) as usize;
797 let remainder = length % max_run_length;
798
799 let total_chunks = num_full_chunks + if remainder > 0 { 1 } else { 0 };
800 all_values.reserve(total_chunks * type_size);
801 all_lengths.reserve(total_chunks * self.run_length_width.bytes_per_value());
802
803 for _ in 0..num_full_chunks {
804 all_values.extend_from_slice(value_bytes);
805 self.run_length_width
806 .write_length(max_run_length, all_lengths);
807 }
808
809 if remainder > 0 {
810 all_values.extend_from_slice(value_bytes);
811 self.run_length_width.write_length(remainder, all_lengths);
812 }
813
814 total_chunks * (type_size + self.run_length_width.bytes_per_value())
815 }
816
817 fn flat_child_candidate(
818 buffers: &[LanceBuffer],
819 chunks: &[MiniBlockChunk],
820 buffer_index: usize,
821 bits_per_value: u64,
822 ) -> RleChildCandidate {
823 RleChildCandidate {
824 encoding: EncodingUtils::flat(bits_per_value, None),
825 data: buffers[buffer_index].clone(),
826 chunk_sizes: chunks
827 .iter()
828 .map(|chunk| chunk.buffer_sizes[buffer_index])
829 .collect(),
830 size: buffers[buffer_index].len(),
831 requires_num_values: false,
832 }
833 }
834
835 fn general_child_candidate(
836 buffers: &[LanceBuffer],
837 chunks: &[MiniBlockChunk],
838 buffer_index: usize,
839 bits_per_value: u64,
840 compression: CompressionConfig,
841 ) -> Result<Option<RleChildCandidate>> {
842 if buffers.is_empty() || buffers[buffer_index].is_empty() {
843 return Ok(None);
844 };
845
846 let compressor = GeneralBufferCompressor::get_compressor(compression)?;
847 let original = &buffers[buffer_index];
848 let mut compressed = Vec::new();
849 let mut offset = 0usize;
850 let mut total_original_size = 0usize;
851 let mut compressed_sizes = Vec::with_capacity(chunks.len());
852
853 for chunk in chunks.iter() {
854 let chunk_size = chunk.buffer_sizes[buffer_index] as usize;
855 let end = offset.checked_add(chunk_size).ok_or_else(|| {
856 Error::invalid_input_source("RLE child buffer offset overflow".into())
857 })?;
858 if end > original.len() {
859 return Err(Error::invalid_input_source(
860 format!(
861 "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}",
862 buffer_index,
863 end,
864 original.len()
865 )
866 .into(),
867 ));
868 }
869
870 let start = compressed.len();
871 compressor.compress(&original.as_ref()[offset..end], &mut compressed)?;
872 let compressed_size = compressed.len() - start;
873 let compressed_size = u32::try_from(compressed_size).map_err(|_| {
874 Error::invalid_input_source(
875 format!(
876 "RLE child buffer {} compressed chunk is too large: {} bytes",
877 buffer_index, compressed_size
878 )
879 .into(),
880 )
881 })?;
882 compressed_sizes.push(compressed_size);
883 total_original_size += chunk_size;
884 offset = end;
885 }
886
887 if compressed.len() >= total_original_size {
888 return Ok(None);
889 }
890
891 let encoding =
892 EncodingUtils::wrapped(compression, EncodingUtils::flat(bits_per_value, None))?;
893 Ok(Some(
894 RleChildCandidate {
895 encoding,
896 data: LanceBuffer::from(compressed),
897 chunk_sizes: compressed_sizes,
898 size: 0,
899 requires_num_values: false,
900 }
901 .with_size_from_data(),
902 ))
903 }
904
905 #[cfg(feature = "bitpacking")]
906 fn bitpacked_child_candidate(
907 buffers: &[LanceBuffer],
908 chunks: &[MiniBlockChunk],
909 buffer_index: usize,
910 bits_per_value: u64,
911 ) -> Result<Option<RleChildCandidate>> {
912 let original = &buffers[buffer_index];
913 if original.is_empty() {
914 return Ok(None);
915 }
916 let packed_bits = Self::required_bits(original, bits_per_value)?;
917 if packed_bits >= bits_per_value {
918 return Ok(None);
919 }
920
921 let compressor = crate::encodings::physical::bitpacking::OutOfLineBitpacking::new(
922 packed_bits,
923 bits_per_value,
924 );
925 let mut packed = Vec::new();
926 let mut offset = 0usize;
927 let mut packed_sizes = Vec::with_capacity(chunks.len());
928 let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
929 Error::invalid_input_source(
930 format!("RLE child bit width is too large: {bits_per_value}").into(),
931 )
932 })?;
933
934 for chunk in chunks {
935 let chunk_size = chunk.buffer_sizes[buffer_index] as usize;
936 let end = offset.checked_add(chunk_size).ok_or_else(|| {
937 Error::invalid_input_source("RLE child buffer offset overflow".into())
938 })?;
939 if end > original.len() {
940 return Err(Error::invalid_input_source(
941 format!(
942 "RLE child buffer {} chunk size exceeds buffer length: end {}, len {}",
943 buffer_index,
944 end,
945 original.len()
946 )
947 .into(),
948 ));
949 }
950 if bytes_per_value == 0 || !chunk_size.is_multiple_of(bytes_per_value) {
951 return Err(Error::invalid_input_source(
952 format!(
953 "RLE child buffer {} chunk has invalid size {} for {} bits per value",
954 buffer_index, chunk_size, bits_per_value
955 )
956 .into(),
957 ));
958 }
959
960 let child_values = (chunk_size / bytes_per_value) as u64;
961 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
962 bits_per_value,
963 data: original.slice_with_length(offset, chunk_size),
964 num_values: child_values,
965 block_info: BlockInfo::default(),
966 });
967 let (chunk_packed, _) = BlockCompressor::compress(&compressor, block)?;
968 let chunk_packed = chunk_packed.ok_or_else(|| {
969 Error::internal("RLE bitpacking child returned no payload".to_string())
970 })?;
971 let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| {
972 Error::invalid_input_source(
973 format!(
974 "RLE child buffer {} bitpacked chunk is too large: {} bytes",
975 buffer_index,
976 chunk_packed.len()
977 )
978 .into(),
979 )
980 })?;
981 packed_sizes.push(packed_size);
982 packed.extend_from_slice(chunk_packed.as_ref());
983 offset = end;
984 }
985
986 if packed.len() >= original.len() {
987 return Ok(None);
988 }
989
990 Ok(Some(
991 RleChildCandidate {
992 encoding: EncodingUtils::out_of_line_bitpacking(
993 bits_per_value,
994 EncodingUtils::flat(packed_bits, None),
995 ),
996 data: LanceBuffer::from(packed),
997 chunk_sizes: packed_sizes,
998 size: 0,
999 requires_num_values: true,
1000 }
1001 .with_size_from_data(),
1002 ))
1003 }
1004
1005 #[cfg(feature = "bitpacking")]
1006 fn required_bits(buffer: &LanceBuffer, bits_per_value: u64) -> Result<u64> {
1007 let max_value = match bits_per_value {
1008 8 => buffer.as_ref().iter().map(|value| *value as u64).max(),
1009 16 => buffer
1010 .as_ref()
1011 .as_chunks::<2>()
1012 .0
1013 .iter()
1014 .map(|value| u16::from_le_bytes(*value) as u64)
1015 .max(),
1016 32 => buffer
1017 .as_ref()
1018 .as_chunks::<4>()
1019 .0
1020 .iter()
1021 .map(|value| u32::from_le_bytes(*value) as u64)
1022 .max(),
1023 64 => buffer
1024 .as_ref()
1025 .as_chunks::<8>()
1026 .0
1027 .iter()
1028 .map(|value| u64::from_le_bytes(*value))
1029 .max(),
1030 _ => {
1031 return Err(Error::invalid_input_source(
1032 format!(
1033 "RLE child bitpacking only supports 8, 16, 32, or 64-bit values, got {bits_per_value}"
1034 )
1035 .into(),
1036 ));
1037 }
1038 }
1039 .unwrap_or(0);
1040 Ok((u64::BITS - max_value.leading_zeros()).max(1) as u64)
1041 }
1042
1043 fn child_candidates(
1044 buffers: &[LanceBuffer],
1045 chunks: &[MiniBlockChunk],
1046 buffer_index: usize,
1047 bits_per_value: u64,
1048 compression: Option<CompressionConfig>,
1049 use_child_bitpacking: bool,
1050 ) -> Result<Vec<RleChildCandidate>> {
1051 #[cfg(not(feature = "bitpacking"))]
1052 let _ = use_child_bitpacking;
1053 let mut candidates = vec![Self::flat_child_candidate(
1054 buffers,
1055 chunks,
1056 buffer_index,
1057 bits_per_value,
1058 )];
1059 if let Some(compression) = compression
1060 && let Some(candidate) = Self::general_child_candidate(
1061 buffers,
1062 chunks,
1063 buffer_index,
1064 bits_per_value,
1065 compression,
1066 )?
1067 {
1068 candidates.push(candidate);
1069 }
1070 #[cfg(feature = "bitpacking")]
1071 {
1072 if use_child_bitpacking
1073 && let Some(candidate) =
1074 Self::bitpacked_child_candidate(buffers, chunks, buffer_index, bits_per_value)?
1075 {
1076 candidates.push(candidate);
1077 }
1078 }
1079 Ok(candidates)
1080 }
1081
1082 fn select_child_candidates(
1083 values: Vec<RleChildCandidate>,
1084 run_lengths: Vec<RleChildCandidate>,
1085 ) -> (RleChildCandidate, RleChildCandidate) {
1086 let mut best: Option<(usize, usize, usize)> = None;
1087 for (value_idx, value) in values.iter().enumerate() {
1088 for (length_idx, length) in run_lengths.iter().enumerate() {
1089 if value.requires_num_values && length.requires_num_values {
1090 continue;
1091 }
1092 let size = value.size + length.size;
1093 if best.is_none_or(|(_, _, best_size)| size < best_size) {
1094 best = Some((value_idx, length_idx, size));
1095 }
1096 }
1097 }
1098 let (value_idx, length_idx, _) =
1099 best.expect("flat RLE child candidates should always be selectable");
1100 (values[value_idx].clone(), run_lengths[length_idx].clone())
1101 }
1102
1103 pub(crate) fn estimated_payload_size(&self, data: &FixedWidthDataBlock) -> Result<u128> {
1106 const SAMPLE_BLOCKS: u64 = 4;
1107 let block_values = *MAX_MINIBLOCK_VALUES;
1108 if data.num_values <= block_values * SAMPLE_BLOCKS {
1109 return self.selected_payload_size(data);
1110 }
1111 let width = (data.bits_per_value / 8) as usize;
1112 let mut bytes = 0;
1113 for index in 0..SAMPLE_BLOCKS {
1114 let start = index * (data.num_values - block_values) / (SAMPLE_BLOCKS - 1);
1115 let sample = FixedWidthDataBlock {
1116 data: data
1117 .data
1118 .slice_with_length(start as usize * width, block_values as usize * width),
1119 bits_per_value: data.bits_per_value,
1120 num_values: block_values,
1121 block_info: BlockInfo::new(),
1122 };
1123 bytes += self.selected_payload_size(&sample)?;
1124 }
1125 Ok((bytes * data.num_values as u128).div_ceil((block_values * SAMPLE_BLOCKS) as u128))
1126 }
1127
1128 fn selected_payload_size(&self, data: &FixedWidthDataBlock) -> Result<u128> {
1129 let (all_buffers, chunks) =
1130 self.encode_data(&data.data, data.num_values, data.bits_per_value)?;
1131 if all_buffers.is_empty() {
1132 return Ok(0);
1133 }
1134
1135 let values_candidates = Self::child_candidates(
1136 &all_buffers,
1137 &chunks,
1138 0,
1139 data.bits_per_value,
1140 self.values_compression,
1141 self.use_child_bitpacking,
1142 )?;
1143 let run_lengths_candidates = Self::child_candidates(
1144 &all_buffers,
1145 &chunks,
1146 1,
1147 self.run_length_width.bits_per_value(),
1148 self.run_lengths_compression,
1149 self.use_child_bitpacking,
1150 )?;
1151 let (values, run_lengths) =
1152 Self::select_child_candidates(values_candidates, run_lengths_candidates);
1153 Ok((values.size as u128).saturating_add(run_lengths.size as u128))
1154 }
1155}
1156
1157impl RleChildCandidate {
1158 fn with_size_from_data(mut self) -> Self {
1159 self.size = self.data.len();
1160 self
1161 }
1162}
1163
1164impl MiniBlockCompressor for RleEncoder {
1165 fn compress(
1166 &self,
1167 _context: MiniBlockCompressionContext,
1168 data: DataBlock,
1169 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1170 match data {
1171 DataBlock::FixedWidth(fixed_width) => {
1172 let num_values = fixed_width.num_values;
1173 let bits_per_value = fixed_width.bits_per_value;
1174
1175 let (all_buffers, chunks) =
1176 self.encode_data(&fixed_width.data, num_values, bits_per_value)?;
1177 if all_buffers.is_empty() {
1178 let compressed = MiniBlockCompressed {
1179 data: all_buffers,
1180 chunks,
1181 num_values,
1182 };
1183 let encoding = EncodingUtils::rle(
1184 EncodingUtils::flat(bits_per_value, None),
1185 EncodingUtils::flat(self.run_length_width.bits_per_value(), None),
1186 );
1187 return Ok((compressed, encoding));
1188 }
1189
1190 let values_candidates = Self::child_candidates(
1191 &all_buffers,
1192 &chunks,
1193 0,
1194 bits_per_value,
1195 self.values_compression,
1196 self.use_child_bitpacking,
1197 )?;
1198 let run_lengths_candidates = Self::child_candidates(
1199 &all_buffers,
1200 &chunks,
1201 1,
1202 self.run_length_width.bits_per_value(),
1203 self.run_lengths_compression,
1204 self.use_child_bitpacking,
1205 )?;
1206 let (values, run_lengths) =
1207 Self::select_child_candidates(values_candidates, run_lengths_candidates);
1208 let chunks = chunks
1209 .into_iter()
1210 .enumerate()
1211 .map(|(idx, chunk)| MiniBlockChunk {
1212 buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]],
1213 log_num_values: chunk.log_num_values,
1214 })
1215 .collect();
1216
1217 let compressed = MiniBlockCompressed {
1218 data: vec![values.data, run_lengths.data],
1219 chunks,
1220 num_values,
1221 };
1222
1223 let encoding = EncodingUtils::rle(values.encoding, run_lengths.encoding);
1224
1225 Ok((compressed, encoding))
1226 }
1227 _ => Err(Error::invalid_input_source(
1228 "RLE encoding only supports FixedWidth data blocks".into(),
1229 )),
1230 }
1231 }
1232}
1233
1234impl BlockCompressor for RleEncoder {
1235 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
1237 match data {
1238 DataBlock::FixedWidth(fixed_width) => {
1239 let num_values = fixed_width.num_values;
1240 let bits_per_value = fixed_width.bits_per_value;
1241
1242 let all_buffers =
1243 self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?;
1244
1245 let values_size = all_buffers[0].len() as u64;
1246
1247 let mut combined = Vec::new();
1248 combined.extend_from_slice(&values_size.to_le_bytes());
1249 combined.extend_from_slice(&all_buffers[0]);
1250 combined.extend_from_slice(&all_buffers[1]);
1251 Ok((
1252 Some(LanceBuffer::from(combined)),
1253 EncodingUtils::rle(
1254 EncodingUtils::flat(bits_per_value, None),
1255 EncodingUtils::flat(self.run_length_width.bits_per_value(), None),
1256 ),
1257 ))
1258 }
1259 _ => Err(Error::invalid_input_source(
1260 "RLE encoding only supports FixedWidth data blocks".into(),
1261 )),
1262 }
1263 }
1264}
1265
1266#[derive(Debug)]
1268pub struct RleDecompressor {
1269 bits_per_value: u64,
1270 run_length_width: RunLengthWidth,
1271 values: RleChildDecompressor,
1272 run_lengths: RleChildDecompressor,
1273}
1274
1275#[derive(Debug)]
1276pub(crate) struct RleChildDecompressor {
1277 bits_per_value: u64,
1278 inner: RleChildDecompressorInner,
1279}
1280
1281#[derive(Debug)]
1282enum RleChildDecompressorInner {
1283 Flat,
1284 Block {
1285 decompressor: Box<dyn BlockDecompressor>,
1286 requires_num_values: bool,
1287 },
1288}
1289
1290impl RleChildDecompressor {
1291 pub(crate) fn flat(bits_per_value: u64) -> Self {
1292 Self {
1293 bits_per_value,
1294 inner: RleChildDecompressorInner::Flat,
1295 }
1296 }
1297
1298 pub(crate) fn block(
1299 bits_per_value: u64,
1300 decompressor: Box<dyn BlockDecompressor>,
1301 requires_num_values: bool,
1302 ) -> Self {
1303 Self {
1304 bits_per_value,
1305 inner: RleChildDecompressorInner::Block {
1306 decompressor,
1307 requires_num_values,
1308 },
1309 }
1310 }
1311
1312 pub(crate) fn bits_per_value(&self) -> u64 {
1313 self.bits_per_value
1314 }
1315
1316 pub(crate) fn requires_num_values(&self) -> bool {
1317 match &self.inner {
1318 RleChildDecompressorInner::Flat => false,
1319 RleChildDecompressorInner::Block {
1320 requires_num_values,
1321 ..
1322 } => *requires_num_values,
1323 }
1324 }
1325
1326 pub(crate) fn is_identity(&self) -> bool {
1327 matches!(self.inner, RleChildDecompressorInner::Flat)
1328 }
1329
1330 fn decode(
1331 &self,
1332 data: LanceBuffer,
1333 num_values: Option<u64>,
1334 label: &str,
1335 ) -> Result<LanceBuffer> {
1336 match &self.inner {
1337 RleChildDecompressorInner::Flat => Ok(data),
1338 RleChildDecompressorInner::Block {
1339 decompressor,
1340 requires_num_values,
1341 } => {
1342 let num_values = if *requires_num_values {
1343 num_values.ok_or_else(|| {
1344 Error::invalid_input_source(
1345 format!("RLE {label} child compression requires the run count").into(),
1346 )
1347 })?
1348 } else {
1349 num_values.unwrap_or(0)
1350 };
1351 let decoded = decompressor.decompress(Some(data), num_values)?;
1352 self.extract_fixed_width(decoded, num_values, label)
1353 }
1354 }
1355 }
1356
1357 fn extract_fixed_width(
1358 &self,
1359 data: DataBlock,
1360 expected_num_values: u64,
1361 label: &str,
1362 ) -> Result<LanceBuffer> {
1363 match data {
1364 DataBlock::FixedWidth(block) => {
1365 if block.bits_per_value != self.bits_per_value {
1366 return Err(Error::invalid_input_source(
1367 format!(
1368 "RLE {label} child decoded {}-bit values, expected {}",
1369 block.bits_per_value, self.bits_per_value
1370 )
1371 .into(),
1372 ));
1373 }
1374 if expected_num_values != 0 && block.num_values != expected_num_values {
1375 return Err(Error::invalid_input_source(
1376 format!(
1377 "RLE {label} child decoded {} values, expected {}",
1378 block.num_values, expected_num_values
1379 )
1380 .into(),
1381 ));
1382 }
1383 Ok(block.data)
1384 }
1385 _ => Err(Error::invalid_input_source(
1386 format!("RLE {label} child decoded to a non fixed-width block").into(),
1387 )),
1388 }
1389 }
1390}
1391
1392impl RleDecompressor {
1393 pub fn new(bits_per_value: u64) -> Self {
1394 Self {
1395 bits_per_value,
1396 run_length_width: RunLengthWidth::U8,
1397 values: RleChildDecompressor::flat(bits_per_value),
1398 run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()),
1399 }
1400 }
1401
1402 pub(crate) fn with_run_length_width(
1403 bits_per_value: u64,
1404 run_length_width: RunLengthWidth,
1405 ) -> Self {
1406 Self {
1407 bits_per_value,
1408 run_length_width,
1409 values: RleChildDecompressor::flat(bits_per_value),
1410 run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()),
1411 }
1412 }
1413
1414 pub(crate) fn with_child_decompressors(
1415 bits_per_value: u64,
1416 run_length_width: RunLengthWidth,
1417 values: RleChildDecompressor,
1418 run_lengths: RleChildDecompressor,
1419 ) -> Self {
1420 Self {
1421 bits_per_value,
1422 run_length_width,
1423 values,
1424 run_lengths,
1425 }
1426 }
1427
1428 fn decode_data(
1429 &self,
1430 data: Vec<LanceBuffer>,
1431 num_values: u64,
1432 clamp_overflow: bool,
1433 ) -> Result<DataBlock> {
1434 if num_values == 0 {
1435 return Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1436 bits_per_value: self.bits_per_value,
1437 data: LanceBuffer::from(vec![]),
1438 num_values: 0,
1439 block_info: BlockInfo::default(),
1440 }));
1441 }
1442
1443 if data.len() != 2 {
1444 return Err(Error::invalid_input_source(
1445 format!(
1446 "RLE decompressor expects exactly 2 buffers, got {}",
1447 data.len()
1448 )
1449 .into(),
1450 ));
1451 }
1452
1453 let mut data_iter = data.into_iter();
1454 let values_buffer = data_iter.next().unwrap();
1455 let lengths_buffer = data_iter.next().unwrap();
1456 let (values_buffer, lengths_buffer) =
1457 self.decode_child_buffers(values_buffer, lengths_buffer)?;
1458
1459 self.decode_child_data(&values_buffer, &lengths_buffer, num_values, clamp_overflow)
1460 }
1461
1462 fn decode_child_data(
1463 &self,
1464 values_buffer: &LanceBuffer,
1465 lengths_buffer: &LanceBuffer,
1466 num_values: u64,
1467 clamp_overflow: bool,
1468 ) -> Result<DataBlock> {
1469 let decoded_data = match self.bits_per_value {
1470 8 => self.decode_generic::<u8>(
1471 values_buffer,
1472 lengths_buffer,
1473 num_values,
1474 clamp_overflow,
1475 )?,
1476 16 => self.decode_generic::<u16>(
1477 values_buffer,
1478 lengths_buffer,
1479 num_values,
1480 clamp_overflow,
1481 )?,
1482 32 => self.decode_generic::<u32>(
1483 values_buffer,
1484 lengths_buffer,
1485 num_values,
1486 clamp_overflow,
1487 )?,
1488 64 => self.decode_generic::<u64>(
1489 values_buffer,
1490 lengths_buffer,
1491 num_values,
1492 clamp_overflow,
1493 )?,
1494 _ => {
1495 return Err(Error::invalid_input_source(
1496 format!(
1497 "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}",
1498 self.bits_per_value
1499 )
1500 .into(),
1501 ));
1502 }
1503 };
1504
1505 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1506 bits_per_value: self.bits_per_value,
1507 data: decoded_data,
1508 num_values,
1509 block_info: BlockInfo::default(),
1510 }))
1511 }
1512
1513 fn sum_run_lengths(
1514 &self,
1515 values_buffer: &LanceBuffer,
1516 lengths_buffer: &LanceBuffer,
1517 ) -> Result<u64> {
1518 let (value_size, value_type) = match self.bits_per_value {
1519 8 => (1, "u8"),
1520 16 => (2, "u16"),
1521 32 => (4, "u32"),
1522 64 => (8, "u64"),
1523 _ => {
1524 return Err(Error::invalid_input_source(
1525 format!(
1526 "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}",
1527 self.bits_per_value
1528 )
1529 .into(),
1530 ));
1531 }
1532 };
1533 let length_size =
1534 self.validate_buffer_sizes(values_buffer, lengths_buffer, value_size, value_type)?;
1535
1536 lengths_buffer
1537 .chunks_exact(length_size)
1538 .try_fold(0_u64, |num_values, length_bytes| {
1539 let length = self.run_length_width.read_length(length_bytes);
1540 if length == 0 {
1541 return Err(Error::invalid_input_source(
1542 "RLE decoding encountered a zero run length".into(),
1543 ));
1544 }
1545 num_values.checked_add(length).ok_or_else(|| {
1546 Error::invalid_input_source("RLE run length sum overflowed u64".into())
1547 })
1548 })
1549 }
1550
1551 fn validate_buffer_sizes(
1552 &self,
1553 values_buffer: &LanceBuffer,
1554 lengths_buffer: &LanceBuffer,
1555 value_size: usize,
1556 value_type: &str,
1557 ) -> Result<usize> {
1558 let length_size = self.run_length_width.bytes_per_value();
1559 if !values_buffer.len().is_multiple_of(value_size)
1560 || !lengths_buffer.len().is_multiple_of(length_size)
1561 {
1562 return Err(Error::invalid_input_source(format!(
1563 "Invalid buffer sizes for RLE {value_type} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})",
1564 values_buffer.len(),
1565 value_size,
1566 lengths_buffer.len(),
1567 length_size
1568 )
1569 .into()));
1570 }
1571
1572 let num_runs = values_buffer.len() / value_size;
1573 let num_length_entries = lengths_buffer.len() / length_size;
1574 if num_runs != num_length_entries {
1575 return Err(Error::invalid_input_source(
1576 format!(
1577 "Inconsistent RLE buffers: {} runs but {} length entries",
1578 num_runs, num_length_entries
1579 )
1580 .into(),
1581 ));
1582 }
1583 Ok(length_size)
1584 }
1585
1586 fn decode_child_buffers(
1587 &self,
1588 values_buffer: LanceBuffer,
1589 lengths_buffer: LanceBuffer,
1590 ) -> Result<(LanceBuffer, LanceBuffer)> {
1591 let values_requires_num_runs = self.values.requires_num_values();
1592 let lengths_requires_num_runs = self.run_lengths.requires_num_values();
1593 if values_requires_num_runs && lengths_requires_num_runs {
1594 return Err(Error::invalid_input_source(
1595 "RLE values and run lengths child compression both require the run count".into(),
1596 ));
1597 }
1598
1599 if values_requires_num_runs {
1600 let lengths_buffer = self
1601 .run_lengths
1602 .decode(lengths_buffer, None, "run lengths")?;
1603 let num_runs = Self::num_child_values(
1604 &lengths_buffer,
1605 self.run_lengths.bits_per_value(),
1606 "run lengths",
1607 )?;
1608 let values_buffer = self
1609 .values
1610 .decode(values_buffer, Some(num_runs), "values")?;
1611 Ok((values_buffer, lengths_buffer))
1612 } else if lengths_requires_num_runs {
1613 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1614 let num_runs =
1615 Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?;
1616 let lengths_buffer =
1617 self.run_lengths
1618 .decode(lengths_buffer, Some(num_runs), "run lengths")?;
1619 Ok((values_buffer, lengths_buffer))
1620 } else {
1621 let values_buffer = self.values.decode(values_buffer, None, "values")?;
1622 let lengths_buffer = self
1623 .run_lengths
1624 .decode(lengths_buffer, None, "run lengths")?;
1625 Ok((values_buffer, lengths_buffer))
1626 }
1627 }
1628
1629 fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result<u64> {
1630 let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
1631 Error::invalid_input_source(
1632 format!("RLE {label} child bit width is too large: {bits_per_value}").into(),
1633 )
1634 })?;
1635 if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) {
1636 return Err(Error::invalid_input_source(
1637 format!(
1638 "RLE {label} child decoded to {} bytes, not divisible by {}",
1639 buffer.len(),
1640 bytes_per_value
1641 )
1642 .into(),
1643 ));
1644 }
1645 Ok((buffer.len() / bytes_per_value) as u64)
1646 }
1647
1648 fn decode_generic<T>(
1649 &self,
1650 values_buffer: &LanceBuffer,
1651 lengths_buffer: &LanceBuffer,
1652 num_values: u64,
1653 clamp_overflow: bool,
1654 ) -> Result<LanceBuffer>
1655 where
1656 T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType,
1657 {
1658 let type_size = std::mem::size_of::<T>();
1659
1660 if values_buffer.is_empty() || lengths_buffer.is_empty() {
1661 if num_values == 0 {
1662 return Ok(LanceBuffer::empty());
1663 } else {
1664 return Err(Error::invalid_input_source(
1665 format!("Empty buffers but expected {} values", num_values).into(),
1666 ));
1667 }
1668 }
1669
1670 let length_size = self.validate_buffer_sizes(
1671 values_buffer,
1672 lengths_buffer,
1673 type_size,
1674 std::any::type_name::<T>(),
1675 )?;
1676
1677 let values_ref = values_buffer.borrow_to_typed_slice::<T>();
1678 let values: &[T] = values_ref.as_ref();
1679 let lengths = lengths_buffer.as_ref();
1680
1681 let expected_value_count = usize::try_from(num_values).map_err(|_| {
1682 Error::invalid_input_source(
1683 format!("RLE num_values does not fit in usize: {num_values}").into(),
1684 )
1685 })?;
1686 let mut decoded: Vec<T> = Vec::new();
1693 decoded
1694 .try_reserve_exact(expected_value_count)
1695 .map_err(|_| {
1696 Error::invalid_input_source(
1697 format!("RLE decoding cannot allocate {expected_value_count} values").into(),
1698 )
1699 })?;
1700 for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) {
1701 let length = self.run_length_width.read_length(length_bytes);
1702 if length == 0 {
1703 return Err(Error::invalid_input_source(
1704 "RLE decoding encountered a zero run length".into(),
1705 ));
1706 }
1707 let length = usize::try_from(length).map_err(|_| {
1708 Error::invalid_input_source(
1709 format!("RLE run length does not fit in usize: {length}").into(),
1710 )
1711 })?;
1712 let remaining = expected_value_count - decoded.len();
1713 if length > remaining && !clamp_overflow {
1714 return Err(Error::invalid_input_source(
1715 format!(
1716 "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1717 decoded.len() + length,
1718 expected_value_count
1719 )
1720 .into(),
1721 ));
1722 }
1723 let count = length.min(remaining);
1724 let initialized = decoded.len() + count;
1725 decoded.spare_capacity_mut()[..count].fill(std::mem::MaybeUninit::new(*value));
1728 unsafe { decoded.set_len(initialized) };
1731 if length > remaining {
1732 break;
1733 }
1734 }
1735
1736 if decoded.len() != expected_value_count {
1737 return Err(Error::invalid_input_source(
1738 format!(
1739 "RLE decoding produced {} values, expected {}",
1740 decoded.len(),
1741 expected_value_count
1742 )
1743 .into(),
1744 ));
1745 }
1746
1747 trace!(
1748 "RLE decoded {} {} values",
1749 num_values,
1750 std::any::type_name::<T>()
1751 );
1752 Ok(LanceBuffer::reinterpret_vec(decoded))
1753 }
1754}
1755
1756impl MiniBlockDecompressor for RleDecompressor {
1757 fn num_buffers(&self) -> usize {
1758 2
1759 }
1760
1761 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
1762 self.decode_data(data, num_values, true)
1763 }
1764
1765 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
1766 num_values
1767 .checked_mul(self.bits_per_value)
1768 .map(|bits| bits.div_ceil(8))
1769 }
1770}
1771
1772impl BlockDecompressor for RleDecompressor {
1773 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
1774 let data = require_block_payload(data, "RLE")?;
1775 let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
1776 self.decode_data(vec![values_buffer, lengths_buffer], num_values, false)
1777 }
1778
1779 fn infer_num_values(&self, data: &LanceBuffer) -> Result<Option<u64>> {
1780 if self.bits_per_value != 16
1784 || self.run_length_width != RunLengthWidth::U8
1785 || !self.values.is_identity()
1786 || !self.run_lengths.is_identity()
1787 {
1788 return Ok(None);
1789 }
1790 let (values_buffer, lengths_buffer) = parse_rle_block_frame(data)?;
1791 self.sum_run_lengths(&values_buffer, &lengths_buffer)
1792 .map(Some)
1793 }
1794}
1795
1796fn parse_rle_block_frame(data: &LanceBuffer) -> Result<(LanceBuffer, LanceBuffer)> {
1799 if data.len() < 8 {
1801 return Err(Error::invalid_input_source(
1802 format!("Insufficient data size: {}", data.len()).into(),
1803 ));
1804 }
1805
1806 let values_size_bytes: [u8; 8] = data[..8].try_into().expect("slice length already checked");
1807 let values_size: usize = u64::from_le_bytes(values_size_bytes)
1808 .try_into()
1809 .map_err(|_| {
1810 Error::invalid_input_source(
1811 format!(
1812 "Invalid values buffer size: {}",
1813 u64::from_le_bytes(values_size_bytes)
1814 )
1815 .into(),
1816 )
1817 })?;
1818
1819 let values_start: usize = 8;
1821 let lengths_start = values_start
1822 .checked_add(values_size)
1823 .ok_or_else(|| Error::invalid_input_source("Invalid RLE values buffer size".into()))?;
1824
1825 if data.len() < lengths_start {
1826 return Err(Error::invalid_input_source(
1827 format!("Insufficient data size: {}", data.len()).into(),
1828 ));
1829 }
1830
1831 let values_buffer = data.slice_with_length(values_start, values_size);
1832 let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start);
1833 Ok((values_buffer, lengths_buffer))
1834}
1835
1836#[derive(Clone, Debug)]
1837enum RleRunLengths {
1838 U8(ScalarBuffer<u8>),
1839 U16(ScalarBuffer<u16>),
1840 U32(ScalarBuffer<u32>),
1841}
1842
1843impl RleRunLengths {
1844 fn try_new(buffer: LanceBuffer, width: RunLengthWidth) -> Result<Self> {
1845 let width_bytes = width.bytes_per_value();
1846 if !buffer.len().is_multiple_of(width_bytes) {
1847 return Err(Error::invalid_input_source(
1848 format!(
1849 "Invalid RLE run lengths buffer: {} bytes (not divisible by {})",
1850 buffer.len(),
1851 width_bytes
1852 )
1853 .into(),
1854 ));
1855 }
1856 Ok(match width {
1857 RunLengthWidth::U8 => Self::U8(buffer.borrow_to_typed_slice()),
1858 RunLengthWidth::U16 => Self::U16(buffer.borrow_to_typed_slice()),
1859 RunLengthWidth::U32 => Self::U32(buffer.borrow_to_typed_slice()),
1860 })
1861 }
1862
1863 fn len(&self) -> usize {
1864 match self {
1865 Self::U8(lengths) => lengths.len(),
1866 Self::U16(lengths) => lengths.len(),
1867 Self::U32(lengths) => lengths.len(),
1868 }
1869 }
1870
1871 fn get(&self, index: usize) -> usize {
1872 match self {
1873 Self::U8(lengths) => lengths[index] as usize,
1874 Self::U16(lengths) => lengths[index] as usize,
1875 Self::U32(lengths) => lengths[index] as usize,
1876 }
1877 }
1878
1879 fn owned_size(&self) -> usize {
1880 match self {
1881 Self::U8(lengths) => std::mem::size_of_val(lengths.as_ref()),
1882 Self::U16(lengths) => std::mem::size_of_val(lengths.as_ref()),
1883 Self::U32(lengths) => std::mem::size_of_val(lengths.as_ref()),
1884 }
1885 }
1886
1887 fn into_owned(self) -> Self {
1888 match self {
1889 Self::U8(lengths) => Self::U8(ScalarBuffer::from(lengths.as_ref().to_vec())),
1890 Self::U16(lengths) => Self::U16(ScalarBuffer::from(lengths.as_ref().to_vec())),
1891 Self::U32(lengths) => Self::U32(ScalarBuffer::from(lengths.as_ref().to_vec())),
1892 }
1893 }
1894
1895 fn deep_size(&self) -> usize {
1896 match self {
1897 Self::U8(lengths) => lengths.inner().capacity(),
1898 Self::U16(lengths) => lengths.inner().capacity(),
1899 Self::U32(lengths) => lengths.inner().capacity(),
1900 }
1901 }
1902}
1903
1904#[derive(Clone, Debug)]
1910pub(crate) struct RleRuns {
1911 values: ScalarBuffer<u16>,
1912 lengths: RleRunLengths,
1913 num_values: usize,
1914 coalesced_runs: usize,
1915}
1916
1917impl RleRuns {
1918 fn try_new(
1919 values_buffer: LanceBuffer,
1920 lengths_buffer: LanceBuffer,
1921 run_length_width: RunLengthWidth,
1922 num_values: u64,
1923 ) -> Result<Self> {
1924 let num_values = usize::try_from(num_values).map_err(|_| {
1925 Error::invalid_input_source(
1926 format!("RLE num_values does not fit in usize: {num_values}").into(),
1927 )
1928 })?;
1929 let type_size = std::mem::size_of::<u16>();
1930 if !values_buffer.len().is_multiple_of(type_size) {
1931 return Err(Error::invalid_input_source(
1932 format!(
1933 "Invalid RLE u16 values buffer: {} bytes (not divisible by {})",
1934 values_buffer.len(),
1935 type_size
1936 )
1937 .into(),
1938 ));
1939 }
1940
1941 let values = values_buffer.borrow_to_typed_slice::<u16>();
1942 let lengths = RleRunLengths::try_new(lengths_buffer, run_length_width)?;
1943 if values.len() != lengths.len() {
1944 return Err(Error::invalid_input_source(
1945 format!(
1946 "Inconsistent RLE buffers: {} runs but {} length entries",
1947 values.len(),
1948 lengths.len()
1949 )
1950 .into(),
1951 ));
1952 }
1953 if values.is_empty() && num_values != 0 {
1954 return Err(Error::invalid_input_source(
1955 format!("Empty RLE buffers but expected {num_values} values").into(),
1956 ));
1957 }
1958
1959 let mut decoded_values = 0usize;
1960 let mut coalesced_runs = 0usize;
1961 let mut previous_value = None;
1962 for run in 0..values.len() {
1963 let length = lengths.get(run);
1964 if length == 0 {
1965 return Err(Error::invalid_input_source(
1966 "RLE decoding encountered a zero run length".into(),
1967 ));
1968 }
1969 decoded_values = decoded_values.checked_add(length).ok_or_else(|| {
1970 Error::invalid_input_source("RLE run length sum overflowed usize".into())
1971 })?;
1972 if decoded_values > num_values {
1973 return Err(Error::invalid_input_source(
1974 format!(
1975 "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1976 decoded_values, num_values
1977 )
1978 .into(),
1979 ));
1980 }
1981 if previous_value != Some(values[run]) {
1982 coalesced_runs += 1;
1983 previous_value = Some(values[run]);
1984 }
1985 }
1986 if decoded_values != num_values {
1987 return Err(Error::invalid_input_source(
1988 format!(
1989 "RLE decoding produced {} values, expected {}",
1990 decoded_values, num_values
1991 )
1992 .into(),
1993 ));
1994 }
1995
1996 Ok(Self {
1997 values,
1998 lengths,
1999 num_values,
2000 coalesced_runs,
2001 })
2002 }
2003
2004 pub(crate) fn num_values(&self) -> usize {
2005 self.num_values
2006 }
2007
2008 pub(crate) fn num_runs(&self) -> usize {
2009 self.values.len()
2010 }
2011
2012 pub(crate) fn coalesced_runs(&self) -> usize {
2013 self.coalesced_runs
2014 }
2015
2016 pub(crate) fn owned_size(&self) -> usize {
2017 std::mem::size_of_val(self.values.as_ref()) + self.lengths.owned_size()
2018 }
2019
2020 pub(crate) fn into_owned(self) -> Self {
2021 Self {
2022 values: ScalarBuffer::from(self.values.as_ref().to_vec()),
2023 lengths: self.lengths.into_owned(),
2024 num_values: self.num_values,
2025 coalesced_runs: self.coalesced_runs,
2026 }
2027 }
2028
2029 pub(crate) fn deep_size(&self) -> usize {
2030 self.values.inner().capacity() + self.lengths.deep_size()
2031 }
2032
2033 pub(crate) fn value(&self, run: usize) -> u16 {
2034 self.values[run]
2035 }
2036
2037 pub(crate) fn length(&self, run: usize) -> usize {
2038 self.lengths.get(run)
2039 }
2040
2041 pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = (u16, usize)> + '_ {
2042 (0..self.num_runs()).map(|run| (self.value(run), self.length(run)))
2043 }
2044}
2045
2046impl RleDecompressor {
2047 pub(crate) fn decode_u16_runs(&self, data: LanceBuffer, num_values: u64) -> Result<RleRuns> {
2049 if self.bits_per_value != 16 {
2050 return Err(Error::invalid_input_source(
2051 format!(
2052 "RLE level values must be 16 bits, got {}",
2053 self.bits_per_value
2054 )
2055 .into(),
2056 ));
2057 }
2058 let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
2059 let (values_buffer, lengths_buffer) =
2060 self.decode_child_buffers(values_buffer, lengths_buffer)?;
2061 RleRuns::try_new(
2062 values_buffer,
2063 lengths_buffer,
2064 self.run_length_width,
2065 num_values,
2066 )
2067 }
2068}
2069
2070pub(crate) const GENERIC_BLOCK_RLE_HEADER_BYTES: usize = 8;
2073
2074#[derive(Debug)]
2076pub struct BlockRleEncoder {
2077 bits_per_value: u64,
2078 values: Box<dyn BlockCompressor>,
2079 run_lengths: Box<dyn BlockCompressor>,
2080}
2081
2082impl BlockRleEncoder {
2083 pub fn try_new(
2084 bits_per_value: u64,
2085 values: Box<dyn BlockCompressor>,
2086 run_lengths: Box<dyn BlockCompressor>,
2087 ) -> Result<Self> {
2088 if !matches!(bits_per_value, 32 | 64) {
2089 return Err(Error::invalid_input(format!(
2090 "Generic block RLE only supports 32 or 64-bit values, got {bits_per_value}"
2091 )));
2092 }
2093 Ok(Self {
2094 bits_per_value,
2095 values,
2096 run_lengths,
2097 })
2098 }
2099}
2100
2101impl BlockCompressor for BlockRleEncoder {
2102 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
2103 let DataBlock::FixedWidth(data) = data else {
2104 return Err(Error::invalid_input(
2105 "Generic block RLE requires fixed-width data",
2106 ));
2107 };
2108 let (run_values, run_lengths) = materialize_generic_runs(&data, self.bits_per_value)?;
2109 let (values_payload, values_encoding) =
2110 self.values.compress(DataBlock::FixedWidth(run_values))?;
2111 let (lengths_payload, lengths_encoding) = self
2112 .run_lengths
2113 .compress(DataBlock::FixedWidth(run_lengths))?;
2114 let encoding = EncodingUtils::rle(values_encoding, lengths_encoding);
2115 if values_payload.is_none() && lengths_payload.is_none() {
2116 return Ok((None, encoding));
2117 }
2118
2119 let values_payload = values_payload.unwrap_or_else(LanceBuffer::empty);
2120 let lengths_payload = lengths_payload.unwrap_or_else(LanceBuffer::empty);
2121 let capacity = GENERIC_BLOCK_RLE_HEADER_BYTES
2122 .checked_add(values_payload.len())
2123 .and_then(|capacity| capacity.checked_add(lengths_payload.len()))
2124 .ok_or_else(|| {
2125 Error::invalid_input("Generic block RLE frame length overflows usize")
2126 })?;
2127 let mut output = Vec::new();
2128 output.try_reserve_exact(capacity).map_err(|error| {
2129 Error::invalid_input(format!(
2130 "Generic block RLE could not reserve {capacity} frame bytes: {error}"
2131 ))
2132 })?;
2133 output.extend_from_slice(&(values_payload.len() as u64).to_le_bytes());
2134 output.extend_from_slice(&values_payload);
2135 output.extend_from_slice(&lengths_payload);
2136 Ok((Some(LanceBuffer::from(output)), encoding))
2137 }
2138}
2139
2140fn materialize_generic_runs(
2141 data: &FixedWidthDataBlock,
2142 bits_per_value: u64,
2143) -> Result<(FixedWidthDataBlock, FixedWidthDataBlock)> {
2144 if data.bits_per_value != bits_per_value {
2145 return Err(Error::invalid_input(format!(
2146 "Generic block RLE expects {bits_per_value}-bit values, got {}",
2147 data.bits_per_value
2148 )));
2149 }
2150 if data.num_values == 0 {
2151 return Err(Error::invalid_input(
2152 "Generic block RLE cannot encode an empty sequence",
2153 ));
2154 }
2155
2156 fn collect<T: ArrowNativeType + Copy + Eq>(values: &[T]) -> (Vec<T>, Vec<u32>) {
2157 let mut run_values = Vec::new();
2158 let mut run_lengths = Vec::new();
2159 let mut current = values[0];
2160 let mut length = 1_u32;
2161 for value in values.iter().copied().skip(1) {
2162 if value == current && length < u32::MAX {
2163 length += 1;
2164 } else {
2165 run_values.push(current);
2166 run_lengths.push(length);
2167 current = value;
2168 length = 1;
2169 }
2170 }
2171 run_values.push(current);
2172 run_lengths.push(length);
2173 (run_values, run_lengths)
2174 }
2175
2176 let (values, lengths, num_runs) = match bits_per_value {
2177 32 => {
2178 let input = crate::encodings::physical::checked_fixed_values::<u32>(
2179 data,
2180 "Generic block RLE input",
2181 )?;
2182 let (values, lengths) = collect(&input);
2183 let num_runs = values.len() as u64;
2184 (LanceBuffer::reinterpret_vec(values), lengths, num_runs)
2185 }
2186 64 => {
2187 let input = crate::encodings::physical::checked_fixed_values::<u64>(
2188 data,
2189 "Generic block RLE input",
2190 )?;
2191 let (values, lengths) = collect(&input);
2192 let num_runs = values.len() as u64;
2193 (LanceBuffer::reinterpret_vec(values), lengths, num_runs)
2194 }
2195 _ => unreachable!("generic block RLE width was validated at construction"),
2196 };
2197 Ok((
2198 FixedWidthDataBlock {
2199 bits_per_value,
2200 data: values,
2201 num_values: num_runs,
2202 block_info: BlockInfo::default(),
2203 },
2204 FixedWidthDataBlock {
2205 bits_per_value: 32,
2206 data: LanceBuffer::reinterpret_vec(lengths),
2207 num_values: num_runs,
2208 block_info: BlockInfo::default(),
2209 },
2210 ))
2211}
2212
2213#[derive(Debug, Clone, Copy)]
2215pub(crate) enum BlockRleRunCount {
2216 ConstantRunLength(u64),
2217 RangeRunLengths { start: u64, step: u64 },
2218 FlatValues,
2219 FlatRunLengths,
2220}
2221
2222#[derive(Debug)]
2224pub(crate) struct BlockRleDecompressor {
2225 bits_per_value: u64,
2226 values: Box<dyn BlockDecompressor>,
2227 run_lengths: Box<dyn BlockDecompressor>,
2228 run_count: BlockRleRunCount,
2229}
2230
2231impl BlockRleDecompressor {
2232 pub(crate) fn try_new(
2233 bits_per_value: u64,
2234 values: Box<dyn BlockDecompressor>,
2235 run_lengths: Box<dyn BlockDecompressor>,
2236 run_count: BlockRleRunCount,
2237 ) -> Result<Self> {
2238 if !matches!(bits_per_value, 32 | 64) {
2239 return Err(Error::invalid_input(format!(
2240 "Generic block RLE only supports 32 or 64-bit values, got {bits_per_value}"
2241 )));
2242 }
2243 Ok(Self {
2244 bits_per_value,
2245 values,
2246 run_lengths,
2247 run_count,
2248 })
2249 }
2250}
2251
2252impl BlockDecompressor for BlockRleDecompressor {
2253 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
2254 if num_values == 0 {
2255 return Err(Error::invalid_input(
2256 "Generic block RLE cannot decode an empty sequence",
2257 ));
2258 }
2259 let values_have_payload = self.values.requires_payload();
2260 let lengths_have_payload = self.run_lengths.requires_payload();
2261 let (values_payload, lengths_payload) =
2262 split_generic_rle_payload(data, values_have_payload, lengths_have_payload)?;
2263 let run_count = infer_generic_run_count(
2264 self.run_count,
2265 values_payload.as_ref(),
2266 lengths_payload.as_ref(),
2267 self.bits_per_value,
2268 num_values,
2269 )?;
2270 let values = self.values.decompress(values_payload, run_count)?;
2271 let lengths = self.run_lengths.decompress(lengths_payload, run_count)?;
2272 expand_generic_runs(values, lengths, self.bits_per_value, num_values, run_count)
2273 }
2274
2275 fn requires_payload(&self) -> bool {
2276 self.values.requires_payload() || self.run_lengths.requires_payload()
2277 }
2278}
2279
2280fn split_generic_rle_payload(
2281 data: Option<LanceBuffer>,
2282 values_have_payload: bool,
2283 lengths_have_payload: bool,
2284) -> Result<(Option<LanceBuffer>, Option<LanceBuffer>)> {
2285 if !values_have_payload && !lengths_have_payload {
2286 if data.is_some() {
2287 return Err(Error::invalid_input(
2288 "Metadata-only generic block RLE expects no payload",
2289 ));
2290 }
2291 return Ok((None, None));
2292 }
2293 let data = require_block_payload(data, "Generic block RLE")?;
2294 if data.len() < GENERIC_BLOCK_RLE_HEADER_BYTES {
2295 return Err(Error::invalid_input(format!(
2296 "Generic block RLE payload has {} bytes, shorter than its {GENERIC_BLOCK_RLE_HEADER_BYTES}-byte header",
2297 data.len()
2298 )));
2299 }
2300 let values_size = u64::from_le_bytes(
2301 data[..GENERIC_BLOCK_RLE_HEADER_BYTES]
2302 .try_into()
2303 .expect("generic block RLE header length was checked"),
2304 );
2305 let values_size = usize::try_from(values_size).map_err(|_| {
2306 Error::invalid_input("Generic block RLE values payload length does not fit usize")
2307 })?;
2308 let lengths_start = GENERIC_BLOCK_RLE_HEADER_BYTES
2309 .checked_add(values_size)
2310 .ok_or_else(|| Error::invalid_input("Generic block RLE values payload end overflows"))?;
2311 if lengths_start > data.len() {
2312 return Err(Error::invalid_input(format!(
2313 "Generic block RLE values payload ends at {lengths_start}, beyond {} bytes",
2314 data.len()
2315 )));
2316 }
2317 if !values_have_payload && values_size != 0 {
2318 return Err(Error::invalid_input(format!(
2319 "Metadata-only RLE values child has {values_size} framed payload bytes"
2320 )));
2321 }
2322 if !lengths_have_payload && lengths_start != data.len() {
2323 return Err(Error::invalid_input(format!(
2324 "Metadata-only RLE run-length child has {} framed payload bytes",
2325 data.len() - lengths_start
2326 )));
2327 }
2328 Ok((
2329 values_have_payload
2330 .then(|| data.slice_with_length(GENERIC_BLOCK_RLE_HEADER_BYTES, values_size)),
2331 lengths_have_payload
2332 .then(|| data.slice_with_length(lengths_start, data.len() - lengths_start)),
2333 ))
2334}
2335
2336fn infer_generic_run_count(
2337 source: BlockRleRunCount,
2338 values_payload: Option<&LanceBuffer>,
2339 lengths_payload: Option<&LanceBuffer>,
2340 bits_per_value: u64,
2341 num_values: u64,
2342) -> Result<u64> {
2343 let run_count = match source {
2344 BlockRleRunCount::ConstantRunLength(length) => {
2345 if length == 0 || !num_values.is_multiple_of(length) {
2346 return Err(Error::invalid_input(format!(
2347 "RLE constant run length {length} does not divide {num_values} values"
2348 )));
2349 }
2350 num_values / length
2351 }
2352 BlockRleRunCount::RangeRunLengths { start, step } => {
2353 if start == 0 || step == 0 {
2354 return Err(Error::invalid_input(
2355 "RLE range run lengths must be positive",
2356 ));
2357 }
2358 let target = u128::from(num_values);
2359 let mut low = 1_u64;
2360 let mut high = num_values;
2361 let mut found = None;
2362 while low <= high {
2363 let count = low + (high - low) / 2;
2364 let count128 = u128::from(count);
2365 let sum = count128
2366 .checked_mul(
2367 u128::from(start)
2368 .checked_mul(2)
2369 .and_then(|start| {
2370 u128::from(count - 1)
2371 .checked_mul(u128::from(step))
2372 .and_then(|tail| start.checked_add(tail))
2373 })
2374 .unwrap_or(u128::MAX),
2375 )
2376 .map(|sum| sum / 2)
2377 .unwrap_or(u128::MAX);
2378 match sum.cmp(&target) {
2379 std::cmp::Ordering::Less => low = count + 1,
2380 std::cmp::Ordering::Greater => high = count - 1,
2381 std::cmp::Ordering::Equal => {
2382 found = Some(count);
2383 break;
2384 }
2385 }
2386 }
2387 found.ok_or_else(|| {
2388 Error::invalid_input(format!(
2389 "RLE range run lengths start={start} step={step} do not sum to {num_values}"
2390 ))
2391 })?
2392 }
2393 BlockRleRunCount::FlatValues => flat_generic_run_count(
2394 values_payload.ok_or_else(|| {
2395 Error::invalid_input("RLE values payload is required to infer the run count")
2396 })?,
2397 bits_per_value,
2398 num_values,
2399 "values",
2400 )?,
2401 BlockRleRunCount::FlatRunLengths => flat_generic_run_count(
2402 lengths_payload.ok_or_else(|| {
2403 Error::invalid_input("RLE run-length payload is required to infer the run count")
2404 })?,
2405 32,
2406 num_values,
2407 "run lengths",
2408 )?,
2409 };
2410 if run_count == 0 {
2411 return Err(Error::invalid_input(
2412 "Generic block RLE describes zero runs",
2413 ));
2414 }
2415 Ok(run_count)
2416}
2417
2418fn flat_generic_run_count(
2419 payload: &LanceBuffer,
2420 bits_per_value: u64,
2421 max_runs: u64,
2422 label: &str,
2423) -> Result<u64> {
2424 let bytes_per_value = usize::try_from(bits_per_value / 8)
2425 .map_err(|_| Error::invalid_input("RLE child width does not fit usize"))?;
2426 if bytes_per_value == 0 || !payload.len().is_multiple_of(bytes_per_value) {
2427 return Err(Error::invalid_input(format!(
2428 "RLE {label} payload has {} bytes, not divisible by {bytes_per_value}",
2429 payload.len()
2430 )));
2431 }
2432 let run_count = (payload.len() / bytes_per_value) as u64;
2433 if run_count > max_runs {
2434 return Err(Error::invalid_input(format!(
2435 "RLE {label} payload contains {run_count} runs, exceeding {max_runs}"
2436 )));
2437 }
2438 Ok(run_count)
2439}
2440
2441fn expand_generic_runs(
2442 values: DataBlock,
2443 lengths: DataBlock,
2444 bits_per_value: u64,
2445 num_values: u64,
2446 run_count: u64,
2447) -> Result<DataBlock> {
2448 let DataBlock::FixedWidth(values) = values else {
2449 return Err(Error::invalid_input(
2450 "RLE values decoded to a non fixed-width block",
2451 ));
2452 };
2453 let DataBlock::FixedWidth(lengths) = lengths else {
2454 return Err(Error::invalid_input(
2455 "RLE run lengths decoded to a non fixed-width block",
2456 ));
2457 };
2458 if values.bits_per_value != bits_per_value
2459 || values.num_values != run_count
2460 || lengths.bits_per_value != 32
2461 || lengths.num_values != run_count
2462 {
2463 return Err(Error::invalid_input(
2464 "RLE child cardinality or bit width does not match its descriptor",
2465 ));
2466 }
2467 let lengths =
2468 crate::encodings::physical::checked_fixed_values::<u32>(&lengths, "RLE run lengths")?;
2469 let total = lengths
2470 .iter()
2471 .enumerate()
2472 .try_fold(0_u64, |total, (index, length)| {
2473 if *length == 0 {
2474 return Err(Error::invalid_input(format!(
2475 "RLE run length at index {index} is zero"
2476 )));
2477 }
2478 total.checked_add(u64::from(*length)).ok_or_else(|| {
2479 Error::invalid_input(format!("RLE run length sum overflows at index {index}"))
2480 })
2481 })?;
2482 if total != num_values {
2483 return Err(Error::invalid_input(format!(
2484 "RLE run lengths sum to {total}, expected {num_values}"
2485 )));
2486 }
2487 let output = match bits_per_value {
2488 32 => {
2489 let values =
2490 crate::encodings::physical::checked_fixed_values::<u32>(&values, "RLE values")?;
2491 let mut output =
2492 crate::encodings::physical::try_vec_with_capacity::<u32>(num_values, "RLE output")?;
2493 for (value, length) in values.iter().zip(lengths.iter()) {
2494 output.extend(std::iter::repeat_n(*value, *length as usize));
2495 }
2496 LanceBuffer::reinterpret_vec(output)
2497 }
2498 64 => {
2499 let values =
2500 crate::encodings::physical::checked_fixed_values::<u64>(&values, "RLE values")?;
2501 let mut output =
2502 crate::encodings::physical::try_vec_with_capacity::<u64>(num_values, "RLE output")?;
2503 for (value, length) in values.iter().zip(lengths.iter()) {
2504 output.extend(std::iter::repeat_n(*value, *length as usize));
2505 }
2506 LanceBuffer::reinterpret_vec(output)
2507 }
2508 _ => unreachable!("generic block RLE width was validated at construction"),
2509 };
2510 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
2511 bits_per_value,
2512 data: output,
2513 num_values,
2514 block_info: BlockInfo::default(),
2515 }))
2516}
2517
2518#[cfg(test)]
2519mod tests {
2520 use std::sync::Arc;
2521
2522 use super::*;
2523 use crate::compression::{
2524 DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor,
2525 };
2526 use crate::data::DataBlock;
2527 use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES;
2528 use crate::encodings::physical::block::{CompressionConfig, CompressionScheme};
2529 use crate::{
2530 buffer::LanceBuffer,
2531 compression::{BlockCompressor, BlockDecompressor},
2532 };
2533 use arrow_array::Int32Array;
2534 use rstest::rstest;
2535
2536 fn compress_miniblock(
2537 compressor: &dyn MiniBlockCompressor,
2538 data: DataBlock,
2539 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
2540 compressor.compress(MiniBlockCompressionContext::new(0, true, true), data)
2541 }
2542
2543 fn expand_u16_runs(runs: &RleRuns) -> Vec<u16> {
2544 let mut expanded = Vec::with_capacity(runs.num_values());
2545 for (value, length) in runs.iter() {
2546 expanded.extend(std::iter::repeat_n(value, length));
2547 }
2548 expanded
2549 }
2550
2551 #[test]
2552 fn decode_u16_runs_matches_eager() {
2553 let mut levels: Vec<u16> = vec![1u16; 1000];
2555 levels.extend(std::iter::repeat_n(0u16, 500));
2556 levels.extend(std::iter::repeat_n(2u16, 300));
2557 let num_values = levels.len() as u64;
2558
2559 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2560 data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
2561 bits_per_value: 16,
2562 num_values,
2563 block_info: BlockInfo::new(),
2564 });
2565 let frame = BlockCompressor::compress(&RleEncoder::new(), block)
2566 .unwrap()
2567 .0
2568 .unwrap();
2569
2570 let eager = BlockDecompressor::decompress(
2571 &RleDecompressor::new(16),
2572 Some(frame.clone()),
2573 num_values,
2574 )
2575 .unwrap();
2576 let DataBlock::FixedWidth(eager) = eager else {
2577 panic!("expected fixed-width block");
2578 };
2579 assert_eq!(
2580 eager.data.borrow_to_typed_slice::<u16>().as_ref(),
2581 levels.as_slice()
2582 );
2583
2584 let runs = RleDecompressor::new(16)
2586 .decode_u16_runs(frame, num_values)
2587 .unwrap();
2588 assert_eq!(runs.num_values(), num_values as usize);
2589 assert_eq!(runs.num_runs(), 8);
2592 assert_eq!(runs.coalesced_runs(), 3);
2593 assert_eq!(expand_u16_runs(&runs), levels);
2594 }
2595
2596 #[test]
2597 fn decode_u16_runs_empty() {
2598 let mut empty_frame = Vec::new();
2599 empty_frame.extend_from_slice(&0u64.to_le_bytes());
2600 let runs = RleDecompressor::new(16)
2601 .decode_u16_runs(LanceBuffer::from(empty_frame), 0)
2602 .unwrap();
2603 assert_eq!(runs.num_values(), 0);
2604 assert_eq!(runs.num_runs(), 0);
2605 assert_eq!(runs.coalesced_runs(), 0);
2606
2607 let error = RleDecompressor::new(16)
2608 .decode_u16_runs(LanceBuffer::empty(), 0)
2609 .unwrap_err();
2610 assert!(error.to_string().contains("Insufficient data size: 0"));
2611 }
2612
2613 #[test]
2614 fn legacy_block_rle_infers_value_count_without_materializing() {
2615 let num_values = u64::from(u16::MAX) + 8;
2616 let full_runs = num_values / u64::from(u8::MAX);
2617 let remainder = num_values % u64::from(u8::MAX);
2618 let num_runs = full_runs + u64::from(remainder != 0);
2619 let mut frame = Vec::new();
2620 frame.extend_from_slice(&(num_runs * 2).to_le_bytes());
2621 frame.extend(std::iter::repeat_n(7_u16, num_runs as usize).flat_map(u16::to_le_bytes));
2622 frame.extend(std::iter::repeat_n(u8::MAX, full_runs as usize));
2623 if remainder != 0 {
2624 frame.push(remainder as u8);
2625 }
2626
2627 let inferred_num_values = RleDecompressor::new(16)
2628 .infer_num_values(&LanceBuffer::from(frame))
2629 .unwrap();
2630 assert_eq!(inferred_num_values, Some(num_values));
2631 }
2632
2633 #[test]
2634 fn newer_block_rle_does_not_infer_untrusted_run_sum() {
2635 let mut frame = Vec::new();
2636 frame.extend_from_slice(&2_u64.to_le_bytes());
2637 frame.extend_from_slice(&7_u16.to_le_bytes());
2638 frame.extend_from_slice(&u32::MAX.to_le_bytes());
2639 let frame = LanceBuffer::from(frame);
2640 let decompressor = RleDecompressor::with_run_length_width(16, RunLengthWidth::U32);
2641
2642 assert_eq!(decompressor.infer_num_values(&frame).unwrap(), None);
2643 let error = BlockDecompressor::decompress(&decompressor, Some(frame), u64::from(u16::MAX))
2644 .unwrap_err();
2645 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
2646 assert!(
2647 error
2648 .to_string()
2649 .contains("RLE decoding overflowed expected value count")
2650 );
2651 }
2652
2653 #[rstest]
2654 #[case::zero(0, 1, "zero run length")]
2655 #[case::underflow(1, 2, "produced 1 values, expected 2")]
2656 #[case::overflow(2, 1, "overflowed expected value count")]
2657 #[case::nonempty_for_empty(1, 0, "overflowed expected value count")]
2658 fn decode_u16_runs_rejects_invalid_coverage(
2659 #[case] run_length: u8,
2660 #[case] num_values: u64,
2661 #[case] expected_message: &str,
2662 ) {
2663 let mut frame = Vec::new();
2664 frame.extend_from_slice(&2u64.to_le_bytes());
2665 frame.extend_from_slice(&7u16.to_le_bytes());
2666 frame.push(run_length);
2667
2668 let error = RleDecompressor::new(16)
2669 .decode_u16_runs(LanceBuffer::from(frame), num_values)
2670 .unwrap_err();
2671 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
2672 assert!(error.to_string().contains(expected_message));
2673 }
2674
2675 #[test]
2676 #[cfg(any(feature = "lz4", feature = "zstd"))]
2677 fn decode_u16_runs_supports_compressed_values_child() {
2678 let levels: Vec<u16> = (0..1024)
2679 .flat_map(|run| std::iter::repeat_n((run % 8) as u16, 4))
2680 .collect();
2681 let num_values = levels.len() as u64;
2682 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2683 data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
2684 bits_per_value: 16,
2685 num_values,
2686 block_info: BlockInfo::new(),
2687 });
2688 let frame = BlockCompressor::compress(&RleEncoder::new(), block)
2689 .unwrap()
2690 .0
2691 .unwrap();
2692 let (values, lengths) = parse_rle_block_frame(&frame).unwrap();
2693
2694 let compression = test_general_compression();
2695 let compressor = GeneralBufferCompressor::get_compressor(compression).unwrap();
2696 let mut compressed_values = Vec::new();
2697 compressor
2698 .compress(values.as_ref(), &mut compressed_values)
2699 .unwrap();
2700 let mut compressed_frame = Vec::new();
2701 compressed_frame.extend_from_slice(&(compressed_values.len() as u64).to_le_bytes());
2702 compressed_frame.extend_from_slice(&compressed_values);
2703 compressed_frame.extend_from_slice(lengths.as_ref());
2704
2705 let encoding = EncodingUtils::rle(
2706 EncodingUtils::wrapped(compression, EncodingUtils::flat(16, None)).unwrap(),
2707 EncodingUtils::flat(8, None),
2708 );
2709 let decompressor = create_rle_decompressor(
2710 expect_rle(&encoding),
2711 &DefaultDecompressionStrategy::default(),
2712 )
2713 .unwrap();
2714 let runs = decompressor
2715 .decode_u16_runs(LanceBuffer::from(compressed_frame), num_values)
2716 .unwrap();
2717 assert_eq!(expand_u16_runs(&runs), levels);
2718 }
2719
2720 #[test]
2721 fn decode_u16_runs_counts_coalesced_runs() {
2722 let num_values = 5000u64;
2726 let constant: Vec<u16> = vec![7u16; num_values as usize];
2727 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2728 data: LanceBuffer::reinterpret_slice(Arc::from(constant)),
2729 bits_per_value: 16,
2730 num_values,
2731 block_info: BlockInfo::new(),
2732 });
2733 let frame = BlockCompressor::compress(&RleEncoder::new(), block)
2734 .unwrap()
2735 .0
2736 .unwrap();
2737 let runs = RleDecompressor::new(16)
2738 .decode_u16_runs(frame, num_values)
2739 .unwrap();
2740 assert_eq!(
2741 runs.num_runs(),
2742 num_values.div_ceil(u8::MAX as u64) as usize
2743 );
2744 assert_eq!(runs.coalesced_runs(), 1);
2745 assert_eq!(expand_u16_runs(&runs), vec![7u16; num_values as usize]);
2746
2747 let alternating: Vec<u16> = (0..200u16).map(|i| i % 2).collect();
2750 let n = alternating.len() as u64;
2751 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2752 data: LanceBuffer::reinterpret_slice(Arc::from(alternating.clone())),
2753 bits_per_value: 16,
2754 num_values: n,
2755 block_info: BlockInfo::new(),
2756 });
2757 let frame = BlockCompressor::compress(&RleEncoder::new(), block)
2758 .unwrap()
2759 .0
2760 .unwrap();
2761 let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap();
2762 assert_eq!(
2763 runs.coalesced_runs() as u64,
2764 n,
2765 "no two adjacent values are equal"
2766 );
2767 assert_eq!(expand_u16_runs(&runs), alternating);
2768 }
2769
2770 #[test]
2773 fn test_basic_miniblock_rle_encoding() {
2774 let encoder = RleEncoder::new();
2775
2776 let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]);
2778 let data_block = DataBlock::from_array(array);
2779
2780 let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap();
2781
2782 assert_eq!(compressed.num_values, 9);
2783 assert_eq!(compressed.chunks.len(), 1);
2784
2785 let values_buffer = &compressed.data[0];
2787 let lengths_buffer = &compressed.data[1];
2788 assert_eq!(values_buffer.len(), 12); assert_eq!(lengths_buffer.len(), 3); }
2791
2792 #[test]
2793 fn test_long_run_splitting() {
2794 let encoder = RleEncoder::new();
2795
2796 let mut data = vec![42i32; 1000]; data.extend(&[100i32; 300]); let array = Int32Array::from(data);
2801 let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2802
2803 let lengths_buffer = &compressed.data[1];
2805 assert_eq!(lengths_buffer.len(), 6);
2806 }
2807
2808 #[test]
2809 fn test_rle_v2_u16_miniblock_encoding() {
2810 let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16);
2811
2812 let data = vec![42i32; 1000];
2813 let array = Int32Array::from(data);
2814 let (compressed, encoding) =
2815 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2816
2817 assert_eq!(compressed.data[0].len(), 4);
2818 assert_eq!(compressed.data[1].len(), 2);
2819 assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes());
2820
2821 let rle = match encoding.compression.as_ref().unwrap() {
2822 crate::format::metadata::compressive_encoding::Compression::Rle(rle) => rle,
2823 other => panic!("expected RLE encoding, got {other:?}"),
2824 };
2825 let run_lengths = rle.run_lengths.as_ref().unwrap();
2826 let flat = match run_lengths.compression.as_ref().unwrap() {
2827 crate::format::metadata::compressive_encoding::Compression::Flat(flat) => flat,
2828 other => panic!("expected flat run lengths, got {other:?}"),
2829 };
2830 assert_eq!(flat.bits_per_value, 16);
2831
2832 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2833 let decompressed = MiniBlockDecompressor::decompress(
2834 &decompressor,
2835 compressed.data,
2836 compressed.num_values,
2837 )
2838 .unwrap();
2839 match decompressed {
2840 DataBlock::FixedWidth(block) => {
2841 let values = block.data.borrow_to_typed_slice::<i32>();
2842 assert_eq!(values.as_ref(), vec![42i32; 1000]);
2843 }
2844 _ => panic!("Expected FixedWidth block"),
2845 }
2846 }
2847
2848 #[test]
2849 #[cfg(any(feature = "lz4", feature = "zstd"))]
2850 fn test_rle_miniblock_compressed_values_child() {
2851 let compression = test_general_compression();
2852 let encoder =
2853 RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false);
2854 let array = Int32Array::from(repeating_runs(1024, 4));
2855 let (compressed, encoding) =
2856 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2857
2858 let rle = expect_rle(&encoding);
2859 assert!(matches!(
2860 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2861 crate::format::metadata::compressive_encoding::Compression::General(_)
2862 ));
2863 assert!(matches!(
2864 rle.run_lengths
2865 .as_ref()
2866 .unwrap()
2867 .compression
2868 .as_ref()
2869 .unwrap(),
2870 crate::format::metadata::compressive_encoding::Compression::Flat(_)
2871 ));
2872
2873 let decompressor = DefaultDecompressionStrategy::default()
2874 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2875 .unwrap();
2876 let decoded =
2877 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2878 .unwrap();
2879 assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4));
2880 }
2881
2882 #[test]
2883 #[cfg(any(feature = "lz4", feature = "zstd"))]
2884 fn test_rle_miniblock_compressed_run_lengths_child() {
2885 let compression = test_general_compression();
2886 let encoder =
2887 RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false);
2888 let expected = repeating_runs(1024, 4);
2889 let (compressed, encoding) = compress_miniblock(
2890 &encoder,
2891 DataBlock::from_array(Int32Array::from(expected.clone())),
2892 )
2893 .unwrap();
2894
2895 let rle = expect_rle(&encoding);
2896 assert!(matches!(
2897 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2898 crate::format::metadata::compressive_encoding::Compression::Flat(_)
2899 ));
2900 assert!(matches!(
2901 rle.run_lengths
2902 .as_ref()
2903 .unwrap()
2904 .compression
2905 .as_ref()
2906 .unwrap(),
2907 crate::format::metadata::compressive_encoding::Compression::General(_)
2908 ));
2909
2910 let decompressor = DefaultDecompressionStrategy::default()
2911 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2912 .unwrap();
2913 let decoded =
2914 MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2915 .unwrap();
2916 assert_decoded_i32_eq(decoded, &expected);
2917 }
2918
2919 #[test]
2920 #[cfg(feature = "bitpacking")]
2921 fn test_rle_miniblock_bitpacked_run_lengths_child() {
2922 use crate::encodings::physical::bitpacking::OutOfLineBitpacking;
2923
2924 let expected = repeating_runs(1024, 4);
2925 let (compressed, _) = compress_miniblock(
2926 &RleEncoder::new(),
2927 DataBlock::from_array(Int32Array::from(expected.clone())),
2928 )
2929 .unwrap();
2930 let run_lengths = compressed.data[1].clone();
2931 let num_runs = run_lengths.len() as u64;
2932 let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2933 bits_per_value: 8,
2934 data: run_lengths,
2935 num_values: num_runs,
2936 block_info: BlockInfo::default(),
2937 });
2938 let bitpacked_run_lengths =
2939 BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block)
2940 .unwrap()
2941 .0
2942 .unwrap();
2943 let encoding = EncodingUtils::rle(
2944 EncodingUtils::flat(32, None),
2945 EncodingUtils::out_of_line_bitpacking(8, EncodingUtils::flat(3, None)),
2946 );
2947
2948 let decompressor = DefaultDecompressionStrategy::default()
2949 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2950 .unwrap();
2951 let decoded = MiniBlockDecompressor::decompress(
2952 decompressor.as_ref(),
2953 vec![compressed.data[0].clone(), bitpacked_run_lengths],
2954 expected.len() as u64,
2955 )
2956 .unwrap();
2957 assert_decoded_i32_eq(decoded, &expected);
2958 }
2959
2960 #[test]
2961 #[cfg(feature = "bitpacking")]
2962 fn test_rle_rejects_two_count_dependent_child_encodings() {
2963 let encoding = EncodingUtils::rle(
2964 EncodingUtils::out_of_line_bitpacking(32, EncodingUtils::flat(3, None)),
2965 EncodingUtils::out_of_line_bitpacking(8, EncodingUtils::flat(3, None)),
2966 );
2967
2968 let err = DefaultDecompressionStrategy::default()
2969 .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2970 .unwrap_err();
2971 assert!(
2972 err.to_string()
2973 .contains("cannot both require the run count")
2974 );
2975 }
2976
2977 #[cfg(any(feature = "lz4", feature = "zstd"))]
2978 fn test_general_compression() -> CompressionConfig {
2979 if cfg!(feature = "zstd") {
2980 CompressionConfig::new(CompressionScheme::Zstd, Some(3))
2981 } else {
2982 CompressionConfig::new(CompressionScheme::Lz4, None)
2983 }
2984 }
2985
2986 fn repeating_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2987 let mut values = Vec::with_capacity(num_runs * run_length);
2988 for run in 0..num_runs {
2989 values.extend(std::iter::repeat_n((run % 8) as i32, run_length));
2990 }
2991 values
2992 }
2993
2994 fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::metadata::Rle {
2995 match encoding.compression.as_ref().unwrap() {
2996 crate::format::metadata::compressive_encoding::Compression::Rle(rle) => rle,
2997 other => panic!("expected RLE encoding, got {other:?}"),
2998 }
2999 }
3000
3001 fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) {
3002 match decoded {
3003 DataBlock::FixedWidth(block) => {
3004 let values = block.data.borrow_to_typed_slice::<i32>();
3005 assert_eq!(values.as_ref(), expected);
3006 }
3007 _ => panic!("Expected FixedWidth block"),
3008 }
3009 }
3010
3011 #[test]
3012 #[cfg(any(feature = "lz4", feature = "zstd"))]
3013 fn test_rle_miniblock_compressed_children_multiple_chunks() {
3014 let compression = test_general_compression();
3015 let encoder = RleEncoder::with_child_encoding(
3016 RunLengthWidth::U8,
3017 Some(compression),
3018 Some(compression),
3019 false,
3020 );
3021 let expected = repeating_runs(8192, 4);
3022 let (compressed, encoding) = compress_miniblock(
3023 &encoder,
3024 DataBlock::from_array(Int32Array::from(expected.clone())),
3025 )
3026 .unwrap();
3027
3028 assert!(compressed.chunks.len() > 1);
3029 let rle = expect_rle(&encoding);
3030 assert!(matches!(
3031 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
3032 crate::format::metadata::compressive_encoding::Compression::General(_)
3033 ));
3034 assert!(matches!(
3035 rle.run_lengths
3036 .as_ref()
3037 .unwrap()
3038 .compression
3039 .as_ref()
3040 .unwrap(),
3041 crate::format::metadata::compressive_encoding::Compression::General(_)
3042 ));
3043
3044 let decoded = decompress_i32_chunks(&compressed, &encoding);
3045 assert_eq!(decoded, expected);
3046 }
3047
3048 #[test]
3049 #[cfg(feature = "bitpacking")]
3050 fn test_rle_miniblock_bitpacks_values_child_when_smaller() {
3051 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
3052 let expected = monotonic_runs(2048, 4);
3053 let (compressed, encoding) = compress_miniblock(
3054 &encoder,
3055 DataBlock::from_array(Int32Array::from(expected.clone())),
3056 )
3057 .unwrap();
3058
3059 let rle = expect_rle(&encoding);
3060 assert!(matches!(
3061 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
3062 crate::format::metadata::compressive_encoding::Compression::OutOfLineBitpacking(_)
3063 ));
3064 assert!(matches!(
3065 rle.run_lengths
3066 .as_ref()
3067 .unwrap()
3068 .compression
3069 .as_ref()
3070 .unwrap(),
3071 crate::format::metadata::compressive_encoding::Compression::Flat(_)
3072 ));
3073
3074 let decoded = decompress_i32_chunks(&compressed, &encoding);
3075 assert_eq!(decoded, expected);
3076 }
3077
3078 #[test]
3079 #[cfg(feature = "bitpacking")]
3080 fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() {
3081 let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
3082 let expected = high_entropy_runs(2048, 4);
3083 let (compressed, encoding) = compress_miniblock(
3084 &encoder,
3085 DataBlock::from_array(Int32Array::from(expected.clone())),
3086 )
3087 .unwrap();
3088
3089 let rle = expect_rle(&encoding);
3090 assert!(matches!(
3091 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
3092 crate::format::metadata::compressive_encoding::Compression::Flat(_)
3093 ));
3094 assert!(matches!(
3095 rle.run_lengths
3096 .as_ref()
3097 .unwrap()
3098 .compression
3099 .as_ref()
3100 .unwrap(),
3101 crate::format::metadata::compressive_encoding::Compression::OutOfLineBitpacking(_)
3102 ));
3103
3104 let decoded = decompress_i32_chunks(&compressed, &encoding);
3105 assert_eq!(decoded, expected);
3106 }
3107
3108 fn decompress_i32_chunks(
3109 compressed: &MiniBlockCompressed,
3110 encoding: &CompressiveEncoding,
3111 ) -> Vec<i32> {
3112 let strategy = DefaultDecompressionStrategy::default();
3113 let decompressor = strategy
3114 .create_miniblock_decompressor(encoding, &strategy)
3115 .unwrap();
3116 let mut offsets = vec![0usize; compressed.data.len()];
3117 let mut values_processed = 0u64;
3118 let mut decoded_values = Vec::new();
3119
3120 for chunk in &compressed.chunks {
3121 let chunk_values = chunk.num_values(values_processed, compressed.num_values);
3122 let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len());
3123 for (idx, size) in chunk.buffer_sizes.iter().enumerate() {
3124 let size = *size as usize;
3125 chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size));
3126 offsets[idx] += size;
3127 }
3128
3129 let decoded = decompressor
3130 .decompress(chunk_buffers, chunk_values)
3131 .unwrap();
3132 match decoded {
3133 DataBlock::FixedWidth(block) => {
3134 let values = block.data.borrow_to_typed_slice::<i32>();
3135 decoded_values.extend_from_slice(values.as_ref());
3136 }
3137 _ => panic!("Expected FixedWidth block"),
3138 }
3139 values_processed += chunk_values;
3140 }
3141
3142 assert_eq!(values_processed, compressed.num_values);
3143 decoded_values
3144 }
3145
3146 #[cfg(feature = "bitpacking")]
3147 fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
3148 let mut values = Vec::with_capacity(num_runs * run_length);
3149 for run in 0..num_runs {
3150 values.extend(std::iter::repeat_n(run as i32, run_length));
3151 }
3152 values
3153 }
3154
3155 #[cfg(feature = "bitpacking")]
3156 fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
3157 let mut values = Vec::with_capacity(num_runs * run_length);
3158 let mut state = 7u64;
3159 for _ in 0..num_runs {
3160 state = state
3161 .wrapping_mul(6364136223846793005)
3162 .wrapping_add(1442695040888963407);
3163 values.extend(std::iter::repeat_n((state >> 32) as i32, run_length));
3164 }
3165 values
3166 }
3167
3168 #[test]
3169 fn test_select_run_length_width_prefers_u16_for_long_runs() {
3170 let mut entries = [0u64; 3];
3171 accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries);
3172 let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap();
3173 assert_eq!(width, RunLengthWidth::U16);
3174 }
3175
3176 #[test]
3179 fn test_round_trip_all_types() {
3180 test_round_trip_helper(vec![42u8, 42, 42, 100, 100, 255, 255, 255, 255], 8);
3182
3183 test_round_trip_helper(vec![1000u16, 1000, 2000, 2000, 2000, 3000], 16);
3185
3186 test_round_trip_helper(vec![100i32, 100, 100, -200, -200, 300, 300, 300, 300], 32);
3188
3189 test_round_trip_helper(vec![1_000_000_000u64; 5], 64);
3191 }
3192
3193 fn test_round_trip_helper<T>(data: Vec<T>, bits_per_value: u64)
3194 where
3195 T: bytemuck::Pod + PartialEq + std::fmt::Debug,
3196 {
3197 let encoder = RleEncoder::new();
3198 let bytes: Vec<u8> = data
3199 .iter()
3200 .flat_map(|v| bytemuck::bytes_of(v))
3201 .copied()
3202 .collect();
3203
3204 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
3205 bits_per_value,
3206 data: LanceBuffer::from(bytes),
3207 num_values: data.len() as u64,
3208 block_info: BlockInfo::default(),
3209 });
3210
3211 let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
3212 let decompressor = RleDecompressor::new(bits_per_value);
3213 let decompressed = MiniBlockDecompressor::decompress(
3214 &decompressor,
3215 compressed.data,
3216 compressed.num_values,
3217 )
3218 .unwrap();
3219
3220 match decompressed {
3221 DataBlock::FixedWidth(ref block) => {
3222 assert_eq!(block.data.len(), data.len() * std::mem::size_of::<T>());
3224 }
3225 _ => panic!("Expected FixedWidth block"),
3226 }
3227 }
3228
3229 #[test]
3232 fn test_power_of_two_chunking() {
3233 let encoder = RleEncoder::new();
3234
3235 let test_sizes = vec![1000, 2500, 5000, 10000];
3237
3238 for size in test_sizes {
3239 let data: Vec<i32> = (0..size)
3240 .map(|i| i / 50) .collect();
3242
3243 let array = Int32Array::from(data);
3244 let (compressed, _) =
3245 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
3246
3247 for (i, chunk) in compressed.chunks.iter().enumerate() {
3249 if i < compressed.chunks.len() - 1 {
3250 assert!(chunk.log_num_values > 0);
3251 let chunk_values = 1u64 << chunk.log_num_values;
3252 assert!(chunk_values.is_power_of_two());
3253 assert!(chunk_values <= *MAX_MINIBLOCK_VALUES);
3254 } else {
3255 assert_eq!(chunk.log_num_values, 0);
3256 }
3257 }
3258 }
3259 }
3260
3261 #[rstest]
3262 #[case::u8_lengths(RunLengthWidth::U8)]
3263 #[case::u16_lengths(RunLengthWidth::U16)]
3264 #[case::u32_lengths(RunLengthWidth::U32)]
3265 fn test_miniblock_chunk_counts_match_encoded_runs(#[case] run_length_width: RunLengthWidth) {
3266 let levels = (0..4098)
3268 .map(|index| if index % 3 == 0 { 1u16 } else { 0u16 })
3269 .collect::<Vec<_>>();
3270 let num_values = levels.len() as u64;
3271 let encoder = RleEncoder::with_run_length_width(run_length_width);
3272 let (buffers, chunks) = encoder
3273 .encode_data(
3274 &LanceBuffer::reinterpret_vec(levels),
3275 num_values,
3276 u16::BITS as u64,
3277 )
3278 .unwrap();
3279
3280 assert_eq!(buffers.len(), 2);
3281 let bytes_per_length = run_length_width.bytes_per_value();
3282 let mut values_offset = 0usize;
3283 let mut lengths_offset = 0usize;
3284 let mut values_processed = 0u64;
3285
3286 for chunk in &chunks {
3287 let values_size = chunk.buffer_sizes[0] as usize;
3288 let lengths_size = chunk.buffer_sizes[1] as usize;
3289 let lengths_end = lengths_offset + lengths_size;
3290 let chunk_lengths = &buffers[1].as_ref()[lengths_offset..lengths_end];
3291 let length_chunks = chunk_lengths.chunks_exact(bytes_per_length);
3292 assert!(length_chunks.remainder().is_empty());
3293 let num_runs = length_chunks.len();
3294 let encoded_values = length_chunks
3295 .map(|bytes| run_length_width.read_length(bytes))
3296 .sum::<u64>();
3297 let declared_values = chunk.num_values(values_processed, num_values);
3298
3299 assert_eq!(values_size, num_runs * size_of::<u16>());
3300 assert_eq!(encoded_values, declared_values);
3301
3302 values_offset += values_size;
3303 lengths_offset = lengths_end;
3304 values_processed += declared_values;
3305 }
3306
3307 assert_eq!(values_processed, num_values);
3308 assert_eq!(values_offset, buffers[0].len());
3309 assert_eq!(lengths_offset, buffers[1].len());
3310 }
3311
3312 #[test]
3315 fn test_encoder_rejects_zero_progress() {
3316 let error = RleEncoder::new()
3317 .encode_data(&LanceBuffer::empty(), 1, u16::BITS as u64)
3318 .unwrap_err();
3319
3320 assert!(
3321 matches!(&error, Error::Internal { .. }),
3322 "expected internal error, got: {error:?}"
3323 );
3324 assert!(error.to_string().contains("made no progress"));
3325 assert!(error.to_string().contains("values_remaining=1"));
3326 }
3327
3328 #[test]
3329 fn test_invalid_buffer_count() {
3330 let decompressor = RleDecompressor::new(32);
3331 let result = MiniBlockDecompressor::decompress(
3332 &decompressor,
3333 vec![LanceBuffer::from(vec![1, 2, 3, 4])],
3334 10,
3335 );
3336 assert!(result.is_err());
3337 assert!(
3338 result
3339 .unwrap_err()
3340 .to_string()
3341 .contains("expects exactly 2 buffers")
3342 );
3343 }
3344
3345 #[test]
3346 fn test_buffer_consistency() {
3347 let decompressor = RleDecompressor::new(32);
3348 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);
3351 assert!(result.is_err());
3352 assert!(
3353 result
3354 .unwrap_err()
3355 .to_string()
3356 .contains("Inconsistent RLE buffers")
3357 );
3358 }
3359
3360 #[test]
3361 fn test_u16_length_buffer_must_be_aligned() {
3362 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
3363 let values = LanceBuffer::from(vec![1, 0, 0, 0]);
3364 let lengths = LanceBuffer::from(vec![5]);
3365 let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5);
3366 assert!(matches!(&result, Err(Error::InvalidInput { .. })));
3367 assert!(
3368 result
3369 .unwrap_err()
3370 .to_string()
3371 .contains("not divisible by 2")
3372 );
3373 }
3374
3375 #[test]
3376 fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() {
3377 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
3378 let value = LanceBuffer::from(1i32.to_le_bytes().to_vec());
3379
3380 let underflow = MiniBlockDecompressor::decompress(
3381 &decompressor,
3382 vec![
3383 value.clone(),
3384 LanceBuffer::from(4u16.to_le_bytes().to_vec()),
3385 ],
3386 5,
3387 )
3388 .unwrap_err();
3389 assert!(underflow.to_string().contains("produced 4 values"));
3390
3391 let overflow = MiniBlockDecompressor::decompress(
3392 &decompressor,
3393 vec![
3394 value.clone(),
3395 LanceBuffer::from(6u16.to_le_bytes().to_vec()),
3396 ],
3397 5,
3398 )
3399 .unwrap();
3400 match overflow {
3401 DataBlock::FixedWidth(block) => {
3402 assert_eq!(block.num_values, 5);
3403 let decoded = block.data.borrow_to_typed_slice::<i32>();
3404 assert_eq!(decoded.as_ref(), &[1i32; 5]);
3405 }
3406 _ => panic!("Expected FixedWidth block"),
3407 }
3408
3409 let zero = MiniBlockDecompressor::decompress(
3410 &decompressor,
3411 vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())],
3412 5,
3413 )
3414 .unwrap_err();
3415 assert!(zero.to_string().contains("zero run length"));
3416 }
3417
3418 #[test]
3419 fn test_block_rle_rejects_overflow() {
3420 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
3423 let values = 1i32.to_le_bytes();
3424 let lengths = 6u16.to_le_bytes();
3425 let mut payload = Vec::new();
3426 payload.extend_from_slice(&(values.len() as u64).to_le_bytes());
3427 payload.extend_from_slice(&values);
3428 payload.extend_from_slice(&lengths);
3429
3430 let error =
3431 BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(payload)), 5)
3432 .unwrap_err();
3433 assert!(matches!(&error, Error::InvalidInput { .. }));
3434 assert!(
3435 error
3436 .to_string()
3437 .contains("overflowed expected value count")
3438 );
3439 }
3440
3441 #[test]
3442 fn test_rle_truncates_legacy_chunk_boundary_overflow() {
3443 let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
3447 let mut values = Vec::new();
3448 values.extend_from_slice(&7i32.to_le_bytes());
3449 values.extend_from_slice(&8i32.to_le_bytes());
3450 let mut lengths = Vec::new();
3451 lengths.extend_from_slice(&2000u16.to_le_bytes());
3452 lengths.extend_from_slice(&80u16.to_le_bytes());
3453
3454 let decoded = MiniBlockDecompressor::decompress(
3455 &decompressor,
3456 vec![LanceBuffer::from(values), LanceBuffer::from(lengths)],
3457 2048,
3458 )
3459 .unwrap();
3460 match decoded {
3461 DataBlock::FixedWidth(block) => {
3462 assert_eq!(block.num_values, 2048);
3463 let decoded = block.data.borrow_to_typed_slice::<i32>();
3464 let decoded = decoded.as_ref();
3465 assert_eq!(decoded.len(), 2048);
3466 assert!(decoded[..2000].iter().all(|&v| v == 7));
3467 assert!(decoded[2000..].iter().all(|&v| v == 8));
3468 }
3469 _ => panic!("Expected FixedWidth block"),
3470 }
3471 }
3472
3473 #[test]
3474 fn test_empty_data_handling() {
3475 let encoder = RleEncoder::new();
3476
3477 let empty_block = DataBlock::FixedWidth(FixedWidthDataBlock {
3479 bits_per_value: 32,
3480 data: LanceBuffer::from(vec![]),
3481 num_values: 0,
3482 block_info: BlockInfo::default(),
3483 });
3484
3485 let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap();
3486 assert_eq!(compressed.num_values, 0);
3487 assert!(compressed.data.is_empty());
3488
3489 let decompressor = RleDecompressor::new(32);
3491 let decompressed = MiniBlockDecompressor::decompress(&decompressor, vec![], 0).unwrap();
3492
3493 match decompressed {
3494 DataBlock::FixedWidth(ref block) => {
3495 assert_eq!(block.num_values, 0);
3496 assert_eq!(block.data.len(), 0);
3497 }
3498 _ => panic!("Expected FixedWidth block"),
3499 }
3500 }
3501
3502 #[test]
3505 fn test_multi_chunk_round_trip() {
3506 let encoder = RleEncoder::new();
3507
3508 let mut data = Vec::new();
3510
3511 data.extend(vec![999i32; 2000]);
3513 data.extend(0..1000);
3515 data.extend(vec![777i32; 2000]);
3517
3518 let array = Int32Array::from(data.clone());
3519 let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
3520
3521 let mut reconstructed = Vec::new();
3523 let mut values_offset = 0usize;
3524 let mut lengths_offset = 0usize;
3525 let mut values_processed = 0u64;
3526
3527 assert_eq!(compressed.data.len(), 2);
3529 let global_values = &compressed.data[0];
3530 let global_lengths = &compressed.data[1];
3531
3532 for chunk in &compressed.chunks {
3533 let chunk_values = if chunk.log_num_values > 0 {
3534 1u64 << chunk.log_num_values
3535 } else {
3536 compressed.num_values - values_processed
3537 };
3538
3539 let values_size = chunk.buffer_sizes[0] as usize;
3541 let lengths_size = chunk.buffer_sizes[1] as usize;
3542
3543 let chunk_values_buffer = global_values.slice_with_length(values_offset, values_size);
3544 let chunk_lengths_buffer =
3545 global_lengths.slice_with_length(lengths_offset, lengths_size);
3546
3547 let decompressor = RleDecompressor::new(32);
3548 let chunk_data = MiniBlockDecompressor::decompress(
3549 &decompressor,
3550 vec![chunk_values_buffer, chunk_lengths_buffer],
3551 chunk_values,
3552 )
3553 .unwrap();
3554
3555 values_offset += values_size;
3556 lengths_offset += lengths_size;
3557 values_processed += chunk_values;
3558
3559 match chunk_data {
3560 DataBlock::FixedWidth(ref block) => {
3561 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3562 reconstructed.extend_from_slice(values);
3563 }
3564 _ => panic!("Expected FixedWidth block"),
3565 }
3566 }
3567
3568 assert_eq!(reconstructed, data);
3569 }
3570
3571 #[test]
3572 fn test_1024_boundary_conditions() {
3573 let encoder = RleEncoder::new();
3576 let decompressor = RleDecompressor::new(32);
3577
3578 let test_cases = [
3579 ("runs_of_2", {
3580 let mut data = Vec::new();
3581 for i in 0..512 {
3582 data.push(i);
3583 data.push(i);
3584 }
3585 data
3586 }),
3587 ("single_run_1024", vec![42i32; 1024]),
3588 ("alternating_values", {
3589 let mut data = Vec::new();
3590 for i in 0..1024 {
3591 data.push(i % 2);
3592 }
3593 data
3594 }),
3595 ("run_boundary_255s", {
3596 let mut data = Vec::new();
3597 data.extend(vec![1i32; 255]);
3598 data.extend(vec![2i32; 255]);
3599 data.extend(vec![3i32; 255]);
3600 data.extend(vec![4i32; 255]);
3601 data.extend(vec![5i32; 4]);
3602 data
3603 }),
3604 ("unique_values_1024", (0..1024).collect::<Vec<_>>()),
3605 ("unique_plus_duplicate", {
3606 let mut data = Vec::new();
3608 for i in 0..1023 {
3609 data.push(i);
3610 }
3611 data.push(1022i32); data
3613 }),
3614 ("bug_4092_pattern", {
3615 let mut data = Vec::new();
3617 for i in 0..1022 {
3618 data.push(i);
3619 }
3620 data.push(999999i32);
3621 data.push(999999i32);
3622 data
3623 }),
3624 ];
3625
3626 for (test_name, data) in test_cases.iter() {
3627 assert_eq!(data.len(), 1024, "Test case {} has wrong length", test_name);
3628
3629 let array = Int32Array::from(data.clone());
3631 let (compressed, _) =
3632 compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
3633
3634 match MiniBlockDecompressor::decompress(
3636 &decompressor,
3637 compressed.data,
3638 compressed.num_values,
3639 ) {
3640 Ok(decompressed) => match decompressed {
3641 DataBlock::FixedWidth(ref block) => {
3642 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3643 assert_eq!(
3644 values.len(),
3645 1024,
3646 "Test case {} got {} values, expected 1024",
3647 test_name,
3648 values.len()
3649 );
3650 assert_eq!(
3651 block.data.len(),
3652 4096,
3653 "Test case {} got {} bytes, expected 4096",
3654 test_name,
3655 block.data.len()
3656 );
3657 assert_eq!(values, &data[..], "Test case {} data mismatch", test_name);
3658 }
3659 _ => panic!("Test case {} expected FixedWidth block", test_name),
3660 },
3661 Err(e) => {
3662 if e.to_string().contains("4092") {
3663 panic!("Test case {} found bug 4092: {}", test_name, e);
3664 }
3665 panic!("Test case {} failed with error: {}", test_name, e);
3666 }
3667 }
3668 }
3669 }
3670
3671 #[test]
3672 fn test_low_repetition_50pct_bug() {
3673 let encoder = RleEncoder::new();
3676
3677 let num_values = 1_048_576; let mut data = Vec::with_capacity(num_values);
3680 let mut value = 0i32;
3681 let mut rng = 12345u64; for _ in 0..num_values {
3684 data.push(value);
3685 rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
3687 if (rng >> 16) & 1 == 1 {
3689 value += 1;
3690 }
3691 }
3692
3693 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3694
3695 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
3696 bits_per_value: 32,
3697 data: LanceBuffer::from(bytes),
3698 num_values: num_values as u64,
3699 block_info: BlockInfo::default(),
3700 });
3701
3702 let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
3703
3704 for (i, chunk) in compressed.chunks.iter().take(5).enumerate() {
3706 let _chunk_values = if chunk.log_num_values > 0 {
3707 1 << chunk.log_num_values
3708 } else {
3709 let prev_total: usize = compressed.chunks[..i]
3711 .iter()
3712 .map(|c| 1usize << c.log_num_values)
3713 .sum();
3714 num_values - prev_total
3715 };
3716 }
3717
3718 let decompressor = RleDecompressor::new(32);
3720 match MiniBlockDecompressor::decompress(
3721 &decompressor,
3722 compressed.data,
3723 compressed.num_values,
3724 ) {
3725 Ok(decompressed) => match decompressed {
3726 DataBlock::FixedWidth(ref block) => {
3727 assert_eq!(
3728 block.data.len(),
3729 num_values * 4,
3730 "Expected {} bytes but got {}",
3731 num_values * 4,
3732 block.data.len()
3733 );
3734 }
3735 _ => panic!("Expected FixedWidth block"),
3736 },
3737 Err(e) => {
3738 if e.to_string().contains("4092") {
3739 panic!("Bug reproduced! {}", e);
3740 } else {
3741 panic!("Unexpected error: {}", e);
3742 }
3743 }
3744 }
3745 }
3746
3747 #[test_log::test(tokio::test)]
3750 async fn test_rle_encoding_verification() {
3751 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
3752 use arrow_array::{Array, Int32Array};
3753 use lance_datagen::{ArrayGenerator, RowCount};
3754 use std::collections::HashMap;
3755 use std::sync::Arc;
3756
3757 let test_cases = TestCases::default()
3758 .with_expected_encoding("rle")
3759 .with_structural_encodings();
3760
3761 let mut metadata_explicit = HashMap::new();
3764 metadata_explicit.insert(
3765 "lance-encoding:rle-threshold".to_string(),
3766 "0.8".to_string(),
3767 );
3768 metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
3769
3770 let mut generator = RleDataGenerator::new(vec![
3771 i32::MIN,
3772 i32::MIN,
3773 i32::MIN,
3774 i32::MIN,
3775 i32::MIN + 1,
3776 i32::MIN + 1,
3777 i32::MIN + 1,
3778 i32::MIN + 1,
3779 i32::MIN + 2,
3780 i32::MIN + 2,
3781 i32::MIN + 2,
3782 i32::MIN + 2,
3783 ]);
3784 let data_explicit = generator.generate_default(RowCount::from(10000)).unwrap();
3785 check_round_trip_encoding_of_data(vec![data_explicit], &test_cases, metadata_explicit)
3786 .await;
3787
3788 let mut metadata = HashMap::new();
3794 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3795
3796 let mut values = vec![i32::MIN; 8000]; values.extend(
3798 [
3799 i32::MIN + 1,
3800 i32::MIN + 2,
3801 i32::MIN + 3,
3802 i32::MIN + 4,
3803 i32::MIN + 5,
3804 ]
3805 .repeat(400),
3806 ); let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3808 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3809
3810 #[cfg(any(feature = "lz4", feature = "zstd"))]
3811 {
3812 let mut metadata = HashMap::new();
3813 metadata.insert(
3814 "lance-encoding:rle-threshold".to_string(),
3815 "0.8".to_string(),
3816 );
3817 metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3818 metadata.insert(
3819 "lance-encoding:compression".to_string(),
3820 if cfg!(feature = "zstd") {
3821 "zstd".to_string()
3822 } else {
3823 "lz4".to_string()
3824 },
3825 );
3826 let mut values = Vec::with_capacity(2048 * 4);
3827 for run in 0..2048 {
3828 values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4));
3829 }
3830 let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3831 check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3832 }
3833 }
3834
3835 #[derive(Debug)]
3837 struct RleDataGenerator {
3838 pattern: Vec<i32>,
3839 idx: usize,
3840 }
3841
3842 impl RleDataGenerator {
3843 fn new(pattern: Vec<i32>) -> Self {
3844 Self { pattern, idx: 0 }
3845 }
3846 }
3847
3848 impl lance_datagen::ArrayGenerator for RleDataGenerator {
3849 fn generate(
3850 &mut self,
3851 _length: lance_datagen::RowCount,
3852 _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
3853 ) -> std::result::Result<std::sync::Arc<dyn arrow_array::Array>, arrow_schema::ArrowError>
3854 {
3855 use arrow_array::Int32Array;
3856 use std::sync::Arc;
3857
3858 let mut values = Vec::new();
3860 for _ in 0..10000 {
3861 values.push(self.pattern[self.idx]);
3862 self.idx = (self.idx + 1) % self.pattern.len();
3863 }
3864 Ok(Arc::new(Int32Array::from(values)))
3865 }
3866
3867 fn data_type(&self) -> &arrow_schema::DataType {
3868 &arrow_schema::DataType::Int32
3869 }
3870
3871 fn element_size_bytes(&self) -> Option<lance_datagen::ByteCount> {
3872 Some(lance_datagen::ByteCount::from(4))
3873 }
3874 }
3875
3876 #[test]
3878 fn test_block_decompressor_rejects_overflowing_values_size() {
3879 let decompressor = RleDecompressor::new(32);
3880
3881 let mut data = Vec::new();
3882 data.extend_from_slice(&u64::MAX.to_le_bytes());
3883 let result = BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(data)), 1);
3884 assert!(result.is_err());
3885 assert!(
3886 result
3887 .unwrap_err()
3888 .to_string()
3889 .contains("Invalid RLE values buffer size")
3890 );
3891 }
3892
3893 #[test]
3894 fn test_block_decompressor_too_small() {
3895 let decompressor = RleDecompressor::new(32);
3896 let result = BlockDecompressor::decompress(
3897 &decompressor,
3898 Some(LanceBuffer::from(vec![1, 2, 3])),
3899 10,
3900 );
3901 assert!(result.is_err());
3902 assert!(
3903 result
3904 .unwrap_err()
3905 .to_string()
3906 .contains("Insufficient data size: 3")
3907 );
3908 }
3909
3910 #[test]
3911 fn test_block_compressor_header_format() {
3912 let encoder = RleEncoder::new();
3913
3914 let data = vec![1i32, 1, 1];
3915 let array = Int32Array::from(data);
3916 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array))
3917 .unwrap()
3918 .0
3919 .unwrap();
3920
3921 assert!(compressed.len() >= 8);
3923 let values_size_bytes: [u8; 8] = compressed.as_ref()[..8].try_into().unwrap();
3924 let values_size = u64::from_le_bytes(values_size_bytes);
3925
3926 assert_eq!(values_size, 4);
3928
3929 assert_eq!(compressed.len(), 13);
3931 }
3932
3933 #[test]
3934 fn test_block_compressor_round_trip() {
3935 let encoder = RleEncoder::new();
3936 let decompressor = RleDecompressor::new(32);
3937
3938 let data = vec![1i32, 1, 1, 2, 2, 3, 3, 3, 3];
3940 let array = Int32Array::from(data.clone());
3941 let data_block = DataBlock::from_array(array);
3942
3943 let compressed = BlockCompressor::compress(&encoder, data_block).unwrap().0;
3944 let decompressed =
3945 BlockDecompressor::decompress(&decompressor, compressed, data.len() as u64).unwrap();
3946
3947 match decompressed {
3948 DataBlock::FixedWidth(block) => {
3949 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3950 assert_eq!(values, &data[..]);
3951 }
3952 _ => panic!("Expected FixedWidth block"),
3953 }
3954 }
3955
3956 #[test]
3957 fn test_block_compressor_large_data() {
3958 let encoder = RleEncoder::new();
3959 let decompressor = RleDecompressor::new(32);
3960
3961 let mut data = Vec::new();
3964 data.extend(vec![999i32; 3000]); data.extend(vec![777i32; 3000]); data.extend(vec![555i32; 4000]); let total_values = data.len();
3969 assert_eq!(total_values, 10000);
3970
3971 let array = Int32Array::from(data.clone());
3972 let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array))
3973 .unwrap()
3974 .0;
3975 let decompressed =
3976 BlockDecompressor::decompress(&decompressor, compressed, total_values as u64).unwrap();
3977
3978 match decompressed {
3979 DataBlock::FixedWidth(block) => {
3980 let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3981 assert_eq!(values.len(), total_values);
3982 assert_eq!(values, &data[..]);
3983 }
3984 _ => panic!("Expected FixedWidth block"),
3985 }
3986 }
3987}