Skip to main content

lance_encoding/encodings/physical/
rle.rs

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