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, ScalarBuffer};
58use log::trace;
59
60use crate::buffer::LanceBuffer;
61use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor};
62use crate::data::DataBlock;
63use crate::data::{BlockInfo, FixedWidthDataBlock};
64use crate::encodings::logical::primitive::miniblock::{
65    MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed,
66    MiniBlockCompressionContext, MiniBlockCompressor,
67};
68use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor};
69use crate::format::ProtobufUtils21;
70use crate::format::pb21::CompressiveEncoding;
71
72use lance_core::{Error, Result};
73
74/// 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(
1040        &self,
1041        _context: MiniBlockCompressionContext,
1042        data: DataBlock,
1043    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1044        match data {
1045            DataBlock::FixedWidth(fixed_width) => {
1046                let num_values = fixed_width.num_values;
1047                let bits_per_value = fixed_width.bits_per_value;
1048
1049                let (all_buffers, chunks) =
1050                    self.encode_data(&fixed_width.data, num_values, bits_per_value)?;
1051                if all_buffers.is_empty() {
1052                    let compressed = MiniBlockCompressed {
1053                        data: all_buffers,
1054                        chunks,
1055                        num_values,
1056                    };
1057                    let encoding = ProtobufUtils21::rle(
1058                        ProtobufUtils21::flat(bits_per_value, None),
1059                        ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None),
1060                    );
1061                    return Ok((compressed, encoding));
1062                }
1063
1064                let values_candidates = Self::child_candidates(
1065                    &all_buffers,
1066                    &chunks,
1067                    0,
1068                    bits_per_value,
1069                    self.values_compression,
1070                    self.use_child_bitpacking,
1071                )?;
1072                let run_lengths_candidates = Self::child_candidates(
1073                    &all_buffers,
1074                    &chunks,
1075                    1,
1076                    self.run_length_width.bits_per_value(),
1077                    self.run_lengths_compression,
1078                    self.use_child_bitpacking,
1079                )?;
1080                let (values, run_lengths) =
1081                    Self::select_child_candidates(values_candidates, run_lengths_candidates);
1082                let chunks = chunks
1083                    .into_iter()
1084                    .enumerate()
1085                    .map(|(idx, chunk)| MiniBlockChunk {
1086                        buffer_sizes: vec![values.chunk_sizes[idx], run_lengths.chunk_sizes[idx]],
1087                        log_num_values: chunk.log_num_values,
1088                    })
1089                    .collect();
1090
1091                let compressed = MiniBlockCompressed {
1092                    data: vec![values.data, run_lengths.data],
1093                    chunks,
1094                    num_values,
1095                };
1096
1097                let encoding = ProtobufUtils21::rle(values.encoding, run_lengths.encoding);
1098
1099                Ok((compressed, encoding))
1100            }
1101            _ => Err(Error::invalid_input_source(
1102                "RLE encoding only supports FixedWidth data blocks".into(),
1103            )),
1104        }
1105    }
1106}
1107
1108impl BlockCompressor for RleEncoder {
1109    // Block format: [8-byte header: values buffer size][values buffer][run_lengths buffer]
1110    fn compress(&self, data: DataBlock) -> Result<LanceBuffer> {
1111        match data {
1112            DataBlock::FixedWidth(fixed_width) => {
1113                let num_values = fixed_width.num_values;
1114                let bits_per_value = fixed_width.bits_per_value;
1115
1116                let all_buffers =
1117                    self.encode_block_data(&fixed_width.data, num_values, bits_per_value)?;
1118
1119                let values_size = all_buffers[0].len() as u64;
1120
1121                let mut combined = Vec::new();
1122                combined.extend_from_slice(&values_size.to_le_bytes());
1123                combined.extend_from_slice(&all_buffers[0]);
1124                combined.extend_from_slice(&all_buffers[1]);
1125                Ok(LanceBuffer::from(combined))
1126            }
1127            _ => Err(Error::invalid_input_source(
1128                "RLE encoding only supports FixedWidth data blocks".into(),
1129            )),
1130        }
1131    }
1132}
1133
1134/// RLE decompressor for miniblock format
1135#[derive(Debug)]
1136pub struct RleDecompressor {
1137    bits_per_value: u64,
1138    run_length_width: RunLengthWidth,
1139    values: RleChildDecompressor,
1140    run_lengths: RleChildDecompressor,
1141}
1142
1143#[derive(Debug)]
1144pub(crate) struct RleChildDecompressor {
1145    bits_per_value: u64,
1146    inner: RleChildDecompressorInner,
1147}
1148
1149#[derive(Debug)]
1150enum RleChildDecompressorInner {
1151    Flat,
1152    Block {
1153        decompressor: Box<dyn BlockDecompressor>,
1154        requires_num_values: bool,
1155    },
1156}
1157
1158impl RleChildDecompressor {
1159    pub(crate) fn flat(bits_per_value: u64) -> Self {
1160        Self {
1161            bits_per_value,
1162            inner: RleChildDecompressorInner::Flat,
1163        }
1164    }
1165
1166    pub(crate) fn block(
1167        bits_per_value: u64,
1168        decompressor: Box<dyn BlockDecompressor>,
1169        requires_num_values: bool,
1170    ) -> Self {
1171        Self {
1172            bits_per_value,
1173            inner: RleChildDecompressorInner::Block {
1174                decompressor,
1175                requires_num_values,
1176            },
1177        }
1178    }
1179
1180    pub(crate) fn bits_per_value(&self) -> u64 {
1181        self.bits_per_value
1182    }
1183
1184    pub(crate) fn requires_num_values(&self) -> bool {
1185        match &self.inner {
1186            RleChildDecompressorInner::Flat => false,
1187            RleChildDecompressorInner::Block {
1188                requires_num_values,
1189                ..
1190            } => *requires_num_values,
1191        }
1192    }
1193
1194    pub(crate) fn is_identity(&self) -> bool {
1195        matches!(self.inner, RleChildDecompressorInner::Flat)
1196    }
1197
1198    fn decode(
1199        &self,
1200        data: LanceBuffer,
1201        num_values: Option<u64>,
1202        label: &str,
1203    ) -> Result<LanceBuffer> {
1204        match &self.inner {
1205            RleChildDecompressorInner::Flat => Ok(data),
1206            RleChildDecompressorInner::Block {
1207                decompressor,
1208                requires_num_values,
1209            } => {
1210                let num_values = if *requires_num_values {
1211                    num_values.ok_or_else(|| {
1212                        Error::invalid_input_source(
1213                            format!("RLE {label} child compression requires the run count").into(),
1214                        )
1215                    })?
1216                } else {
1217                    num_values.unwrap_or(0)
1218                };
1219                let decoded = decompressor.decompress(data, num_values)?;
1220                self.extract_fixed_width(decoded, num_values, label)
1221            }
1222        }
1223    }
1224
1225    fn extract_fixed_width(
1226        &self,
1227        data: DataBlock,
1228        expected_num_values: u64,
1229        label: &str,
1230    ) -> Result<LanceBuffer> {
1231        match data {
1232            DataBlock::FixedWidth(block) => {
1233                if block.bits_per_value != self.bits_per_value {
1234                    return Err(Error::invalid_input_source(
1235                        format!(
1236                            "RLE {label} child decoded {}-bit values, expected {}",
1237                            block.bits_per_value, self.bits_per_value
1238                        )
1239                        .into(),
1240                    ));
1241                }
1242                if expected_num_values != 0 && block.num_values != expected_num_values {
1243                    return Err(Error::invalid_input_source(
1244                        format!(
1245                            "RLE {label} child decoded {} values, expected {}",
1246                            block.num_values, expected_num_values
1247                        )
1248                        .into(),
1249                    ));
1250                }
1251                Ok(block.data)
1252            }
1253            _ => Err(Error::invalid_input_source(
1254                format!("RLE {label} child decoded to a non fixed-width block").into(),
1255            )),
1256        }
1257    }
1258}
1259
1260impl RleDecompressor {
1261    pub fn new(bits_per_value: u64) -> Self {
1262        Self {
1263            bits_per_value,
1264            run_length_width: RunLengthWidth::U8,
1265            values: RleChildDecompressor::flat(bits_per_value),
1266            run_lengths: RleChildDecompressor::flat(RunLengthWidth::U8.bits_per_value()),
1267        }
1268    }
1269
1270    pub(crate) fn with_run_length_width(
1271        bits_per_value: u64,
1272        run_length_width: RunLengthWidth,
1273    ) -> Self {
1274        Self {
1275            bits_per_value,
1276            run_length_width,
1277            values: RleChildDecompressor::flat(bits_per_value),
1278            run_lengths: RleChildDecompressor::flat(run_length_width.bits_per_value()),
1279        }
1280    }
1281
1282    pub(crate) fn with_child_decompressors(
1283        bits_per_value: u64,
1284        run_length_width: RunLengthWidth,
1285        values: RleChildDecompressor,
1286        run_lengths: RleChildDecompressor,
1287    ) -> Self {
1288        Self {
1289            bits_per_value,
1290            run_length_width,
1291            values,
1292            run_lengths,
1293        }
1294    }
1295
1296    fn decode_data(
1297        &self,
1298        data: Vec<LanceBuffer>,
1299        num_values: u64,
1300        clamp_overflow: bool,
1301    ) -> Result<DataBlock> {
1302        if num_values == 0 {
1303            return Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1304                bits_per_value: self.bits_per_value,
1305                data: LanceBuffer::from(vec![]),
1306                num_values: 0,
1307                block_info: BlockInfo::default(),
1308            }));
1309        }
1310
1311        if data.len() != 2 {
1312            return Err(Error::invalid_input_source(
1313                format!(
1314                    "RLE decompressor expects exactly 2 buffers, got {}",
1315                    data.len()
1316                )
1317                .into(),
1318            ));
1319        }
1320
1321        let mut data_iter = data.into_iter();
1322        let values_buffer = data_iter.next().unwrap();
1323        let lengths_buffer = data_iter.next().unwrap();
1324        let (values_buffer, lengths_buffer) =
1325            self.decode_child_buffers(values_buffer, lengths_buffer)?;
1326
1327        let decoded_data = match self.bits_per_value {
1328            8 => self.decode_generic::<u8>(
1329                &values_buffer,
1330                &lengths_buffer,
1331                num_values,
1332                clamp_overflow,
1333            )?,
1334            16 => self.decode_generic::<u16>(
1335                &values_buffer,
1336                &lengths_buffer,
1337                num_values,
1338                clamp_overflow,
1339            )?,
1340            32 => self.decode_generic::<u32>(
1341                &values_buffer,
1342                &lengths_buffer,
1343                num_values,
1344                clamp_overflow,
1345            )?,
1346            64 => self.decode_generic::<u64>(
1347                &values_buffer,
1348                &lengths_buffer,
1349                num_values,
1350                clamp_overflow,
1351            )?,
1352            _ => {
1353                return Err(Error::invalid_input_source(
1354                    format!(
1355                        "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}",
1356                        self.bits_per_value
1357                    )
1358                    .into(),
1359                ));
1360            }
1361        };
1362
1363        Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
1364            bits_per_value: self.bits_per_value,
1365            data: decoded_data,
1366            num_values,
1367            block_info: BlockInfo::default(),
1368        }))
1369    }
1370
1371    fn decode_child_buffers(
1372        &self,
1373        values_buffer: LanceBuffer,
1374        lengths_buffer: LanceBuffer,
1375    ) -> Result<(LanceBuffer, LanceBuffer)> {
1376        let values_requires_num_runs = self.values.requires_num_values();
1377        let lengths_requires_num_runs = self.run_lengths.requires_num_values();
1378        if values_requires_num_runs && lengths_requires_num_runs {
1379            return Err(Error::invalid_input_source(
1380                "RLE values and run lengths child compression both require the run count".into(),
1381            ));
1382        }
1383
1384        if values_requires_num_runs {
1385            let lengths_buffer = self
1386                .run_lengths
1387                .decode(lengths_buffer, None, "run lengths")?;
1388            let num_runs = Self::num_child_values(
1389                &lengths_buffer,
1390                self.run_lengths.bits_per_value(),
1391                "run lengths",
1392            )?;
1393            let values_buffer = self
1394                .values
1395                .decode(values_buffer, Some(num_runs), "values")?;
1396            Ok((values_buffer, lengths_buffer))
1397        } else if lengths_requires_num_runs {
1398            let values_buffer = self.values.decode(values_buffer, None, "values")?;
1399            let num_runs =
1400                Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?;
1401            let lengths_buffer =
1402                self.run_lengths
1403                    .decode(lengths_buffer, Some(num_runs), "run lengths")?;
1404            Ok((values_buffer, lengths_buffer))
1405        } else {
1406            let values_buffer = self.values.decode(values_buffer, None, "values")?;
1407            let lengths_buffer = self
1408                .run_lengths
1409                .decode(lengths_buffer, None, "run lengths")?;
1410            Ok((values_buffer, lengths_buffer))
1411        }
1412    }
1413
1414    fn num_child_values(buffer: &LanceBuffer, bits_per_value: u64, label: &str) -> Result<u64> {
1415        let bytes_per_value = usize::try_from(bits_per_value / 8).map_err(|_| {
1416            Error::invalid_input_source(
1417                format!("RLE {label} child bit width is too large: {bits_per_value}").into(),
1418            )
1419        })?;
1420        if bytes_per_value == 0 || !buffer.len().is_multiple_of(bytes_per_value) {
1421            return Err(Error::invalid_input_source(
1422                format!(
1423                    "RLE {label} child decoded to {} bytes, not divisible by {}",
1424                    buffer.len(),
1425                    bytes_per_value
1426                )
1427                .into(),
1428            ));
1429        }
1430        Ok((buffer.len() / bytes_per_value) as u64)
1431    }
1432
1433    fn decode_generic<T>(
1434        &self,
1435        values_buffer: &LanceBuffer,
1436        lengths_buffer: &LanceBuffer,
1437        num_values: u64,
1438        clamp_overflow: bool,
1439    ) -> Result<LanceBuffer>
1440    where
1441        T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType,
1442    {
1443        let type_size = std::mem::size_of::<T>();
1444        let length_size = self.run_length_width.bytes_per_value();
1445
1446        if values_buffer.is_empty() || lengths_buffer.is_empty() {
1447            if num_values == 0 {
1448                return Ok(LanceBuffer::empty());
1449            } else {
1450                return Err(Error::invalid_input_source(
1451                    format!("Empty buffers but expected {} values", num_values).into(),
1452                ));
1453            }
1454        }
1455
1456        if !values_buffer.len().is_multiple_of(type_size)
1457            || !lengths_buffer.len().is_multiple_of(length_size)
1458        {
1459            return Err(Error::invalid_input_source(format!(
1460                "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})",
1461                std::any::type_name::<T>(),
1462                values_buffer.len(),
1463                type_size,
1464                lengths_buffer.len(),
1465                length_size
1466            )
1467            .into()));
1468        }
1469
1470        let num_runs = values_buffer.len() / type_size;
1471        let num_length_entries = lengths_buffer.len() / length_size;
1472        if num_runs != num_length_entries {
1473            return Err(Error::invalid_input_source(
1474                format!(
1475                    "Inconsistent RLE buffers: {} runs but {} length entries",
1476                    num_runs, num_length_entries
1477                )
1478                .into(),
1479            ));
1480        }
1481
1482        let values_ref = values_buffer.borrow_to_typed_slice::<T>();
1483        let values: &[T] = values_ref.as_ref();
1484        let lengths = lengths_buffer.as_ref();
1485
1486        let expected_value_count = usize::try_from(num_values).map_err(|_| {
1487            Error::invalid_input_source(
1488                format!("RLE num_values does not fit in usize: {num_values}").into(),
1489            )
1490        })?;
1491        // Legacy miniblock encoders rolled back to a power-of-2 checkpoint after a run
1492        // had already crossed it, so a chunk's run lengths can sum past its declared
1493        // value count (the excess values are re-encoded at the start of the next chunk).
1494        // The pre-run-length-width decoder truncated the excess, so miniblock decoding
1495        // clamps rather than rejects to keep those files readable. Block payloads never
1496        // legitimately overflow, so they decode strictly.
1497        let mut decoded: Vec<T> = Vec::new();
1498        decoded
1499            .try_reserve_exact(expected_value_count)
1500            .map_err(|_| {
1501                Error::invalid_input_source(
1502                    format!("RLE decoding cannot allocate {expected_value_count} values").into(),
1503                )
1504            })?;
1505        for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) {
1506            let length = self.run_length_width.read_length(length_bytes);
1507            if length == 0 {
1508                return Err(Error::invalid_input_source(
1509                    "RLE decoding encountered a zero run length".into(),
1510                ));
1511            }
1512            let length = usize::try_from(length).map_err(|_| {
1513                Error::invalid_input_source(
1514                    format!("RLE run length does not fit in usize: {length}").into(),
1515                )
1516            })?;
1517            let remaining = expected_value_count - decoded.len();
1518            if length > remaining {
1519                if !clamp_overflow {
1520                    return Err(Error::invalid_input_source(
1521                        format!(
1522                            "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1523                            decoded.len() + length,
1524                            expected_value_count
1525                        )
1526                        .into(),
1527                    ));
1528                }
1529                decoded.resize(expected_value_count, *value);
1530                break;
1531            }
1532            decoded.resize(decoded.len() + length, *value);
1533        }
1534
1535        if decoded.len() != expected_value_count {
1536            return Err(Error::invalid_input_source(
1537                format!(
1538                    "RLE decoding produced {} values, expected {}",
1539                    decoded.len(),
1540                    expected_value_count
1541                )
1542                .into(),
1543            ));
1544        }
1545
1546        trace!(
1547            "RLE decoded {} {} values",
1548            num_values,
1549            std::any::type_name::<T>()
1550        );
1551        Ok(LanceBuffer::reinterpret_vec(decoded))
1552    }
1553}
1554
1555impl MiniBlockDecompressor for RleDecompressor {
1556    fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
1557        self.decode_data(data, num_values, true)
1558    }
1559
1560    fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
1561        num_values
1562            .checked_mul(self.bits_per_value)
1563            .map(|bits| bits.div_ceil(8))
1564    }
1565}
1566
1567impl BlockDecompressor for RleDecompressor {
1568    fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result<DataBlock> {
1569        let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
1570        self.decode_data(vec![values_buffer, lengths_buffer], num_values, false)
1571    }
1572}
1573
1574/// Split an RLE block-format buffer into its `(values, lengths)` sub-buffers.
1575/// Frame: `[values_size: u64-le][values bytes][run-length bytes]`.
1576fn parse_rle_block_frame(data: &LanceBuffer) -> Result<(LanceBuffer, LanceBuffer)> {
1577    // fetch the values_size
1578    if data.len() < 8 {
1579        return Err(Error::invalid_input_source(
1580            format!("Insufficient data size: {}", data.len()).into(),
1581        ));
1582    }
1583
1584    let values_size_bytes: [u8; 8] = data[..8].try_into().expect("slice length already checked");
1585    let values_size: usize = u64::from_le_bytes(values_size_bytes)
1586        .try_into()
1587        .map_err(|_| {
1588            Error::invalid_input_source(
1589                format!(
1590                    "Invalid values buffer size: {}",
1591                    u64::from_le_bytes(values_size_bytes)
1592                )
1593                .into(),
1594            )
1595        })?;
1596
1597    // parse values
1598    let values_start: usize = 8;
1599    let lengths_start = values_start
1600        .checked_add(values_size)
1601        .ok_or_else(|| Error::invalid_input_source("Invalid RLE values buffer size".into()))?;
1602
1603    if data.len() < lengths_start {
1604        return Err(Error::invalid_input_source(
1605            format!("Insufficient data size: {}", data.len()).into(),
1606        ));
1607    }
1608
1609    let values_buffer = data.slice_with_length(values_start, values_size);
1610    let lengths_buffer = data.slice_with_length(lengths_start, data.len() - lengths_start);
1611    Ok((values_buffer, lengths_buffer))
1612}
1613
1614#[derive(Clone, Debug)]
1615enum RleRunLengths {
1616    U8(ScalarBuffer<u8>),
1617    U16(ScalarBuffer<u16>),
1618    U32(ScalarBuffer<u32>),
1619}
1620
1621impl RleRunLengths {
1622    fn try_new(buffer: LanceBuffer, width: RunLengthWidth) -> Result<Self> {
1623        let width_bytes = width.bytes_per_value();
1624        if !buffer.len().is_multiple_of(width_bytes) {
1625            return Err(Error::invalid_input_source(
1626                format!(
1627                    "Invalid RLE run lengths buffer: {} bytes (not divisible by {})",
1628                    buffer.len(),
1629                    width_bytes
1630                )
1631                .into(),
1632            ));
1633        }
1634        Ok(match width {
1635            RunLengthWidth::U8 => Self::U8(buffer.borrow_to_typed_slice()),
1636            RunLengthWidth::U16 => Self::U16(buffer.borrow_to_typed_slice()),
1637            RunLengthWidth::U32 => Self::U32(buffer.borrow_to_typed_slice()),
1638        })
1639    }
1640
1641    fn len(&self) -> usize {
1642        match self {
1643            Self::U8(lengths) => lengths.len(),
1644            Self::U16(lengths) => lengths.len(),
1645            Self::U32(lengths) => lengths.len(),
1646        }
1647    }
1648
1649    fn get(&self, index: usize) -> usize {
1650        match self {
1651            Self::U8(lengths) => lengths[index] as usize,
1652            Self::U16(lengths) => lengths[index] as usize,
1653            Self::U32(lengths) => lengths[index] as usize,
1654        }
1655    }
1656
1657    fn owned_size(&self) -> usize {
1658        match self {
1659            Self::U8(lengths) => std::mem::size_of_val(lengths.as_ref()),
1660            Self::U16(lengths) => std::mem::size_of_val(lengths.as_ref()),
1661            Self::U32(lengths) => std::mem::size_of_val(lengths.as_ref()),
1662        }
1663    }
1664
1665    fn into_owned(self) -> Self {
1666        match self {
1667            Self::U8(lengths) => Self::U8(ScalarBuffer::from(lengths.as_ref().to_vec())),
1668            Self::U16(lengths) => Self::U16(ScalarBuffer::from(lengths.as_ref().to_vec())),
1669            Self::U32(lengths) => Self::U32(ScalarBuffer::from(lengths.as_ref().to_vec())),
1670        }
1671    }
1672
1673    fn deep_size(&self) -> usize {
1674        match self {
1675            Self::U8(lengths) => lengths.inner().capacity(),
1676            Self::U16(lengths) => lengths.inner().capacity(),
1677            Self::U32(lengths) => lengths.inner().capacity(),
1678        }
1679    }
1680}
1681
1682/// Validated physical RLE runs for `u16` values.
1683///
1684/// The values and original-width lengths remain unexpanded. The constructor
1685/// verifies that every length is non-zero and that the runs cover exactly
1686/// `num_values` logical values.
1687#[derive(Clone, Debug)]
1688pub(crate) struct RleRuns {
1689    values: ScalarBuffer<u16>,
1690    lengths: RleRunLengths,
1691    num_values: usize,
1692    coalesced_runs: usize,
1693}
1694
1695impl RleRuns {
1696    fn try_new(
1697        values_buffer: LanceBuffer,
1698        lengths_buffer: LanceBuffer,
1699        run_length_width: RunLengthWidth,
1700        num_values: u64,
1701    ) -> Result<Self> {
1702        let num_values = usize::try_from(num_values).map_err(|_| {
1703            Error::invalid_input_source(
1704                format!("RLE num_values does not fit in usize: {num_values}").into(),
1705            )
1706        })?;
1707        let type_size = std::mem::size_of::<u16>();
1708        if !values_buffer.len().is_multiple_of(type_size) {
1709            return Err(Error::invalid_input_source(
1710                format!(
1711                    "Invalid RLE u16 values buffer: {} bytes (not divisible by {})",
1712                    values_buffer.len(),
1713                    type_size
1714                )
1715                .into(),
1716            ));
1717        }
1718
1719        let values = values_buffer.borrow_to_typed_slice::<u16>();
1720        let lengths = RleRunLengths::try_new(lengths_buffer, run_length_width)?;
1721        if values.len() != lengths.len() {
1722            return Err(Error::invalid_input_source(
1723                format!(
1724                    "Inconsistent RLE buffers: {} runs but {} length entries",
1725                    values.len(),
1726                    lengths.len()
1727                )
1728                .into(),
1729            ));
1730        }
1731        if values.is_empty() && num_values != 0 {
1732            return Err(Error::invalid_input_source(
1733                format!("Empty RLE buffers but expected {num_values} values").into(),
1734            ));
1735        }
1736
1737        let mut decoded_values = 0usize;
1738        let mut coalesced_runs = 0usize;
1739        let mut previous_value = None;
1740        for run in 0..values.len() {
1741            let length = lengths.get(run);
1742            if length == 0 {
1743                return Err(Error::invalid_input_source(
1744                    "RLE decoding encountered a zero run length".into(),
1745                ));
1746            }
1747            decoded_values = decoded_values.checked_add(length).ok_or_else(|| {
1748                Error::invalid_input_source("RLE run length sum overflowed usize".into())
1749            })?;
1750            if decoded_values > num_values {
1751                return Err(Error::invalid_input_source(
1752                    format!(
1753                        "RLE decoding overflowed expected value count: produced at least {}, expected {}",
1754                        decoded_values, num_values
1755                    )
1756                    .into(),
1757                ));
1758            }
1759            if previous_value != Some(values[run]) {
1760                coalesced_runs += 1;
1761                previous_value = Some(values[run]);
1762            }
1763        }
1764        if decoded_values != num_values {
1765            return Err(Error::invalid_input_source(
1766                format!(
1767                    "RLE decoding produced {} values, expected {}",
1768                    decoded_values, num_values
1769                )
1770                .into(),
1771            ));
1772        }
1773
1774        Ok(Self {
1775            values,
1776            lengths,
1777            num_values,
1778            coalesced_runs,
1779        })
1780    }
1781
1782    pub(crate) fn num_values(&self) -> usize {
1783        self.num_values
1784    }
1785
1786    pub(crate) fn num_runs(&self) -> usize {
1787        self.values.len()
1788    }
1789
1790    pub(crate) fn coalesced_runs(&self) -> usize {
1791        self.coalesced_runs
1792    }
1793
1794    pub(crate) fn owned_size(&self) -> usize {
1795        std::mem::size_of_val(self.values.as_ref()) + self.lengths.owned_size()
1796    }
1797
1798    pub(crate) fn into_owned(self) -> Self {
1799        Self {
1800            values: ScalarBuffer::from(self.values.as_ref().to_vec()),
1801            lengths: self.lengths.into_owned(),
1802            num_values: self.num_values,
1803            coalesced_runs: self.coalesced_runs,
1804        }
1805    }
1806
1807    pub(crate) fn deep_size(&self) -> usize {
1808        self.values.inner().capacity() + self.lengths.deep_size()
1809    }
1810
1811    pub(crate) fn value(&self, run: usize) -> u16 {
1812        self.values[run]
1813    }
1814
1815    pub(crate) fn length(&self, run: usize) -> usize {
1816        self.lengths.get(run)
1817    }
1818
1819    pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = (u16, usize)> + '_ {
1820        (0..self.num_runs()).map(|run| (self.value(run), self.length(run)))
1821    }
1822}
1823
1824impl RleDecompressor {
1825    /// Decode and validate a block frame while preserving its physical runs.
1826    pub(crate) fn decode_u16_runs(&self, data: LanceBuffer, num_values: u64) -> Result<RleRuns> {
1827        if self.bits_per_value != 16 {
1828            return Err(Error::invalid_input_source(
1829                format!(
1830                    "RLE level values must be 16 bits, got {}",
1831                    self.bits_per_value
1832                )
1833                .into(),
1834            ));
1835        }
1836        let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?;
1837        let (values_buffer, lengths_buffer) =
1838            self.decode_child_buffers(values_buffer, lengths_buffer)?;
1839        RleRuns::try_new(
1840            values_buffer,
1841            lengths_buffer,
1842            self.run_length_width,
1843            num_values,
1844        )
1845    }
1846}
1847
1848#[cfg(test)]
1849mod tests {
1850    use std::sync::Arc;
1851
1852    use super::*;
1853    use crate::compression::{
1854        DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor,
1855    };
1856    use crate::data::DataBlock;
1857    use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES;
1858    use crate::encodings::physical::block::{CompressionConfig, CompressionScheme};
1859    use crate::{
1860        buffer::LanceBuffer,
1861        compression::{BlockCompressor, BlockDecompressor},
1862    };
1863    use arrow_array::Int32Array;
1864    use rstest::rstest;
1865
1866    fn compress_miniblock(
1867        compressor: &dyn MiniBlockCompressor,
1868        data: DataBlock,
1869    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
1870        compressor.compress(MiniBlockCompressionContext::new(0, true, true), data)
1871    }
1872
1873    fn expand_u16_runs(runs: &RleRuns) -> Vec<u16> {
1874        let mut expanded = Vec::with_capacity(runs.num_values());
1875        for (value, length) in runs.iter() {
1876            expanded.extend(std::iter::repeat_n(value, length));
1877        }
1878        expanded
1879    }
1880
1881    #[test]
1882    fn decode_u16_runs_matches_eager() {
1883        // Near-constant u16 levels (the all-null shape): a few long runs.
1884        let mut levels: Vec<u16> = vec![1u16; 1000];
1885        levels.extend(std::iter::repeat_n(0u16, 500));
1886        levels.extend(std::iter::repeat_n(2u16, 300));
1887        let num_values = levels.len() as u64;
1888
1889        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
1890            data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
1891            bits_per_value: 16,
1892            num_values,
1893            block_info: BlockInfo::new(),
1894        });
1895        let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
1896
1897        let eager =
1898            BlockDecompressor::decompress(&RleDecompressor::new(16), frame.clone(), num_values)
1899                .unwrap();
1900        let DataBlock::FixedWidth(eager) = eager else {
1901            panic!("expected fixed-width block");
1902        };
1903        assert_eq!(
1904            eager.data.borrow_to_typed_slice::<u16>().as_ref(),
1905            levels.as_slice()
1906        );
1907
1908        // Lazy run form preserves boundaries and expands identically.
1909        let runs = RleDecompressor::new(16)
1910            .decode_u16_runs(frame, num_values)
1911            .unwrap();
1912        assert_eq!(runs.num_values(), num_values as usize);
1913        // The encoder splits each value into <=255-length runs (4 + 2 + 2 = 8
1914        // on-disk runs here); the scan identifies 3 coalesced logical runs.
1915        assert_eq!(runs.num_runs(), 8);
1916        assert_eq!(runs.coalesced_runs(), 3);
1917        assert_eq!(expand_u16_runs(&runs), levels);
1918    }
1919
1920    #[test]
1921    fn decode_u16_runs_empty() {
1922        let mut empty_frame = Vec::new();
1923        empty_frame.extend_from_slice(&0u64.to_le_bytes());
1924        let runs = RleDecompressor::new(16)
1925            .decode_u16_runs(LanceBuffer::from(empty_frame), 0)
1926            .unwrap();
1927        assert_eq!(runs.num_values(), 0);
1928        assert_eq!(runs.num_runs(), 0);
1929        assert_eq!(runs.coalesced_runs(), 0);
1930
1931        let error = RleDecompressor::new(16)
1932            .decode_u16_runs(LanceBuffer::empty(), 0)
1933            .unwrap_err();
1934        assert!(error.to_string().contains("Insufficient data size: 0"));
1935    }
1936
1937    #[rstest]
1938    #[case::zero(0, 1, "zero run length")]
1939    #[case::underflow(1, 2, "produced 1 values, expected 2")]
1940    #[case::overflow(2, 1, "overflowed expected value count")]
1941    #[case::nonempty_for_empty(1, 0, "overflowed expected value count")]
1942    fn decode_u16_runs_rejects_invalid_coverage(
1943        #[case] run_length: u8,
1944        #[case] num_values: u64,
1945        #[case] expected_message: &str,
1946    ) {
1947        let mut frame = Vec::new();
1948        frame.extend_from_slice(&2u64.to_le_bytes());
1949        frame.extend_from_slice(&7u16.to_le_bytes());
1950        frame.push(run_length);
1951
1952        let error = RleDecompressor::new(16)
1953            .decode_u16_runs(LanceBuffer::from(frame), num_values)
1954            .unwrap_err();
1955        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
1956        assert!(error.to_string().contains(expected_message));
1957    }
1958
1959    #[test]
1960    #[cfg(any(feature = "lz4", feature = "zstd"))]
1961    fn decode_u16_runs_supports_compressed_values_child() {
1962        let levels: Vec<u16> = (0..1024)
1963            .flat_map(|run| std::iter::repeat_n((run % 8) as u16, 4))
1964            .collect();
1965        let num_values = levels.len() as u64;
1966        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
1967            data: LanceBuffer::reinterpret_slice(Arc::from(levels.clone())),
1968            bits_per_value: 16,
1969            num_values,
1970            block_info: BlockInfo::new(),
1971        });
1972        let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
1973        let (values, lengths) = parse_rle_block_frame(&frame).unwrap();
1974
1975        let compression = test_general_compression();
1976        let compressor = GeneralBufferCompressor::get_compressor(compression).unwrap();
1977        let mut compressed_values = Vec::new();
1978        compressor
1979            .compress(values.as_ref(), &mut compressed_values)
1980            .unwrap();
1981        let mut compressed_frame = Vec::new();
1982        compressed_frame.extend_from_slice(&(compressed_values.len() as u64).to_le_bytes());
1983        compressed_frame.extend_from_slice(&compressed_values);
1984        compressed_frame.extend_from_slice(lengths.as_ref());
1985
1986        let encoding = ProtobufUtils21::rle(
1987            ProtobufUtils21::wrapped(compression, ProtobufUtils21::flat(16, None)).unwrap(),
1988            ProtobufUtils21::flat(8, None),
1989        );
1990        let decompressor = create_rle_decompressor(
1991            expect_rle(&encoding),
1992            &DefaultDecompressionStrategy::default(),
1993        )
1994        .unwrap();
1995        let runs = decompressor
1996            .decode_u16_runs(LanceBuffer::from(compressed_frame), num_values)
1997            .unwrap();
1998        assert_eq!(expand_u16_runs(&runs), levels);
1999    }
2000
2001    #[test]
2002    fn decode_u16_runs_counts_coalesced_runs() {
2003        // A logically constant page is emitted as ceil(N / 255) equal-valued
2004        // runs (the encoder caps run lengths at 255); the validated view records
2005        // that they can collapse to a single logical run.
2006        let num_values = 5000u64;
2007        let constant: Vec<u16> = vec![7u16; num_values as usize];
2008        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2009            data: LanceBuffer::reinterpret_slice(Arc::from(constant)),
2010            bits_per_value: 16,
2011            num_values,
2012            block_info: BlockInfo::new(),
2013        });
2014        let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
2015        let runs = RleDecompressor::new(16)
2016            .decode_u16_runs(frame, num_values)
2017            .unwrap();
2018        assert_eq!(
2019            runs.num_runs(),
2020            num_values.div_ceil(u8::MAX as u64) as usize
2021        );
2022        assert_eq!(runs.coalesced_runs(), 1);
2023        assert_eq!(expand_u16_runs(&runs), vec![7u16; num_values as usize]);
2024
2025        // Distinct adjacent values must not be merged: alternating single-value
2026        // runs stay separate and still expand to the original.
2027        let alternating: Vec<u16> = (0..200u16).map(|i| i % 2).collect();
2028        let n = alternating.len() as u64;
2029        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2030            data: LanceBuffer::reinterpret_slice(Arc::from(alternating.clone())),
2031            bits_per_value: 16,
2032            num_values: n,
2033            block_info: BlockInfo::new(),
2034        });
2035        let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap();
2036        let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap();
2037        assert_eq!(
2038            runs.coalesced_runs() as u64,
2039            n,
2040            "no two adjacent values are equal"
2041        );
2042        assert_eq!(expand_u16_runs(&runs), alternating);
2043    }
2044
2045    // ========== Core Functionality Tests ==========
2046
2047    #[test]
2048    fn test_basic_miniblock_rle_encoding() {
2049        let encoder = RleEncoder::new();
2050
2051        // Test basic RLE pattern: [1, 1, 1, 2, 2, 3, 3, 3, 3]
2052        let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]);
2053        let data_block = DataBlock::from_array(array);
2054
2055        let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap();
2056
2057        assert_eq!(compressed.num_values, 9);
2058        assert_eq!(compressed.chunks.len(), 1);
2059
2060        // Verify compression happened (3 runs instead of 9 values)
2061        let values_buffer = &compressed.data[0];
2062        let lengths_buffer = &compressed.data[1];
2063        assert_eq!(values_buffer.len(), 12); // 3 i32 values
2064        assert_eq!(lengths_buffer.len(), 3); // 3 u8 lengths
2065    }
2066
2067    #[test]
2068    fn test_long_run_splitting() {
2069        let encoder = RleEncoder::new();
2070
2071        // Create a run longer than 255 to test splitting
2072        let mut data = vec![42i32; 1000]; // Will be split into 255+255+255+235
2073        data.extend(&[100i32; 300]); // Will be split into 255+45
2074
2075        let array = Int32Array::from(data);
2076        let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2077
2078        // Should have 6 runs total (4 for first value, 2 for second)
2079        let lengths_buffer = &compressed.data[1];
2080        assert_eq!(lengths_buffer.len(), 6);
2081    }
2082
2083    #[test]
2084    fn test_rle_v2_u16_miniblock_encoding() {
2085        let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16);
2086
2087        let data = vec![42i32; 1000];
2088        let array = Int32Array::from(data);
2089        let (compressed, encoding) =
2090            compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2091
2092        assert_eq!(compressed.data[0].len(), 4);
2093        assert_eq!(compressed.data[1].len(), 2);
2094        assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes());
2095
2096        let rle = match encoding.compression.as_ref().unwrap() {
2097            crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
2098            other => panic!("expected RLE encoding, got {other:?}"),
2099        };
2100        let run_lengths = rle.run_lengths.as_ref().unwrap();
2101        let flat = match run_lengths.compression.as_ref().unwrap() {
2102            crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat,
2103            other => panic!("expected flat run lengths, got {other:?}"),
2104        };
2105        assert_eq!(flat.bits_per_value, 16);
2106
2107        let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2108        let decompressed = MiniBlockDecompressor::decompress(
2109            &decompressor,
2110            compressed.data,
2111            compressed.num_values,
2112        )
2113        .unwrap();
2114        match decompressed {
2115            DataBlock::FixedWidth(block) => {
2116                let values = block.data.borrow_to_typed_slice::<i32>();
2117                assert_eq!(values.as_ref(), vec![42i32; 1000]);
2118            }
2119            _ => panic!("Expected FixedWidth block"),
2120        }
2121    }
2122
2123    #[test]
2124    #[cfg(any(feature = "lz4", feature = "zstd"))]
2125    fn test_rle_miniblock_compressed_values_child() {
2126        let compression = test_general_compression();
2127        let encoder =
2128            RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false);
2129        let array = Int32Array::from(repeating_runs(1024, 4));
2130        let (compressed, encoding) =
2131            compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2132
2133        let rle = expect_rle(&encoding);
2134        assert!(matches!(
2135            rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2136            crate::format::pb21::compressive_encoding::Compression::General(_)
2137        ));
2138        assert!(matches!(
2139            rle.run_lengths
2140                .as_ref()
2141                .unwrap()
2142                .compression
2143                .as_ref()
2144                .unwrap(),
2145            crate::format::pb21::compressive_encoding::Compression::Flat(_)
2146        ));
2147
2148        let decompressor = DefaultDecompressionStrategy::default()
2149            .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2150            .unwrap();
2151        let decoded =
2152            MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2153                .unwrap();
2154        assert_decoded_i32_eq(decoded, &repeating_runs(1024, 4));
2155    }
2156
2157    #[test]
2158    #[cfg(any(feature = "lz4", feature = "zstd"))]
2159    fn test_rle_miniblock_compressed_run_lengths_child() {
2160        let compression = test_general_compression();
2161        let encoder =
2162            RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false);
2163        let expected = repeating_runs(1024, 4);
2164        let (compressed, encoding) = compress_miniblock(
2165            &encoder,
2166            DataBlock::from_array(Int32Array::from(expected.clone())),
2167        )
2168        .unwrap();
2169
2170        let rle = expect_rle(&encoding);
2171        assert!(matches!(
2172            rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2173            crate::format::pb21::compressive_encoding::Compression::Flat(_)
2174        ));
2175        assert!(matches!(
2176            rle.run_lengths
2177                .as_ref()
2178                .unwrap()
2179                .compression
2180                .as_ref()
2181                .unwrap(),
2182            crate::format::pb21::compressive_encoding::Compression::General(_)
2183        ));
2184
2185        let decompressor = DefaultDecompressionStrategy::default()
2186            .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2187            .unwrap();
2188        let decoded =
2189            MiniBlockDecompressor::decompress(decompressor.as_ref(), compressed.data, 1024 * 4)
2190                .unwrap();
2191        assert_decoded_i32_eq(decoded, &expected);
2192    }
2193
2194    #[test]
2195    #[cfg(feature = "bitpacking")]
2196    fn test_rle_miniblock_bitpacked_run_lengths_child() {
2197        use crate::encodings::physical::bitpacking::OutOfLineBitpacking;
2198
2199        let expected = repeating_runs(1024, 4);
2200        let (compressed, _) = compress_miniblock(
2201            &RleEncoder::new(),
2202            DataBlock::from_array(Int32Array::from(expected.clone())),
2203        )
2204        .unwrap();
2205        let run_lengths = compressed.data[1].clone();
2206        let num_runs = run_lengths.len() as u64;
2207        let run_lengths_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2208            bits_per_value: 8,
2209            data: run_lengths,
2210            num_values: num_runs,
2211            block_info: BlockInfo::default(),
2212        });
2213        let bitpacked_run_lengths =
2214            BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap();
2215        let encoding = ProtobufUtils21::rle(
2216            ProtobufUtils21::flat(32, None),
2217            ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
2218        );
2219
2220        let decompressor = DefaultDecompressionStrategy::default()
2221            .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2222            .unwrap();
2223        let decoded = MiniBlockDecompressor::decompress(
2224            decompressor.as_ref(),
2225            vec![compressed.data[0].clone(), bitpacked_run_lengths],
2226            expected.len() as u64,
2227        )
2228        .unwrap();
2229        assert_decoded_i32_eq(decoded, &expected);
2230    }
2231
2232    #[test]
2233    #[cfg(feature = "bitpacking")]
2234    fn test_rle_rejects_two_count_dependent_child_encodings() {
2235        let encoding = ProtobufUtils21::rle(
2236            ProtobufUtils21::out_of_line_bitpacking(32, ProtobufUtils21::flat(3, None)),
2237            ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)),
2238        );
2239
2240        let err = DefaultDecompressionStrategy::default()
2241            .create_miniblock_decompressor(&encoding, &DefaultDecompressionStrategy::default())
2242            .unwrap_err();
2243        assert!(
2244            err.to_string()
2245                .contains("cannot both require the run count")
2246        );
2247    }
2248
2249    #[cfg(any(feature = "lz4", feature = "zstd"))]
2250    fn test_general_compression() -> CompressionConfig {
2251        if cfg!(feature = "zstd") {
2252            CompressionConfig::new(CompressionScheme::Zstd, Some(3))
2253        } else {
2254            CompressionConfig::new(CompressionScheme::Lz4, None)
2255        }
2256    }
2257
2258    fn repeating_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2259        let mut values = Vec::with_capacity(num_runs * run_length);
2260        for run in 0..num_runs {
2261            values.extend(std::iter::repeat_n((run % 8) as i32, run_length));
2262        }
2263        values
2264    }
2265
2266    fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle {
2267        match encoding.compression.as_ref().unwrap() {
2268            crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle,
2269            other => panic!("expected RLE encoding, got {other:?}"),
2270        }
2271    }
2272
2273    fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) {
2274        match decoded {
2275            DataBlock::FixedWidth(block) => {
2276                let values = block.data.borrow_to_typed_slice::<i32>();
2277                assert_eq!(values.as_ref(), expected);
2278            }
2279            _ => panic!("Expected FixedWidth block"),
2280        }
2281    }
2282
2283    #[test]
2284    #[cfg(any(feature = "lz4", feature = "zstd"))]
2285    fn test_rle_miniblock_compressed_children_multiple_chunks() {
2286        let compression = test_general_compression();
2287        let encoder = RleEncoder::with_child_encoding(
2288            RunLengthWidth::U8,
2289            Some(compression),
2290            Some(compression),
2291            false,
2292        );
2293        let expected = repeating_runs(8192, 4);
2294        let (compressed, encoding) = compress_miniblock(
2295            &encoder,
2296            DataBlock::from_array(Int32Array::from(expected.clone())),
2297        )
2298        .unwrap();
2299
2300        assert!(compressed.chunks.len() > 1);
2301        let rle = expect_rle(&encoding);
2302        assert!(matches!(
2303            rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2304            crate::format::pb21::compressive_encoding::Compression::General(_)
2305        ));
2306        assert!(matches!(
2307            rle.run_lengths
2308                .as_ref()
2309                .unwrap()
2310                .compression
2311                .as_ref()
2312                .unwrap(),
2313            crate::format::pb21::compressive_encoding::Compression::General(_)
2314        ));
2315
2316        let decoded = decompress_i32_chunks(&compressed, &encoding);
2317        assert_eq!(decoded, expected);
2318    }
2319
2320    #[test]
2321    #[cfg(feature = "bitpacking")]
2322    fn test_rle_miniblock_bitpacks_values_child_when_smaller() {
2323        let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
2324        let expected = monotonic_runs(2048, 4);
2325        let (compressed, encoding) = compress_miniblock(
2326            &encoder,
2327            DataBlock::from_array(Int32Array::from(expected.clone())),
2328        )
2329        .unwrap();
2330
2331        let rle = expect_rle(&encoding);
2332        assert!(matches!(
2333            rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2334            crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
2335        ));
2336        assert!(matches!(
2337            rle.run_lengths
2338                .as_ref()
2339                .unwrap()
2340                .compression
2341                .as_ref()
2342                .unwrap(),
2343            crate::format::pb21::compressive_encoding::Compression::Flat(_)
2344        ));
2345
2346        let decoded = decompress_i32_chunks(&compressed, &encoding);
2347        assert_eq!(decoded, expected);
2348    }
2349
2350    #[test]
2351    #[cfg(feature = "bitpacking")]
2352    fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() {
2353        let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true);
2354        let expected = high_entropy_runs(2048, 4);
2355        let (compressed, encoding) = compress_miniblock(
2356            &encoder,
2357            DataBlock::from_array(Int32Array::from(expected.clone())),
2358        )
2359        .unwrap();
2360
2361        let rle = expect_rle(&encoding);
2362        assert!(matches!(
2363            rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2364            crate::format::pb21::compressive_encoding::Compression::Flat(_)
2365        ));
2366        assert!(matches!(
2367            rle.run_lengths
2368                .as_ref()
2369                .unwrap()
2370                .compression
2371                .as_ref()
2372                .unwrap(),
2373            crate::format::pb21::compressive_encoding::Compression::OutOfLineBitpacking(_)
2374        ));
2375
2376        let decoded = decompress_i32_chunks(&compressed, &encoding);
2377        assert_eq!(decoded, expected);
2378    }
2379
2380    fn decompress_i32_chunks(
2381        compressed: &MiniBlockCompressed,
2382        encoding: &CompressiveEncoding,
2383    ) -> Vec<i32> {
2384        let strategy = DefaultDecompressionStrategy::default();
2385        let decompressor = strategy
2386            .create_miniblock_decompressor(encoding, &strategy)
2387            .unwrap();
2388        let mut offsets = vec![0usize; compressed.data.len()];
2389        let mut values_processed = 0u64;
2390        let mut decoded_values = Vec::new();
2391
2392        for chunk in &compressed.chunks {
2393            let chunk_values = chunk.num_values(values_processed, compressed.num_values);
2394            let mut chunk_buffers = Vec::with_capacity(chunk.buffer_sizes.len());
2395            for (idx, size) in chunk.buffer_sizes.iter().enumerate() {
2396                let size = *size as usize;
2397                chunk_buffers.push(compressed.data[idx].slice_with_length(offsets[idx], size));
2398                offsets[idx] += size;
2399            }
2400
2401            let decoded = decompressor
2402                .decompress(chunk_buffers, chunk_values)
2403                .unwrap();
2404            match decoded {
2405                DataBlock::FixedWidth(block) => {
2406                    let values = block.data.borrow_to_typed_slice::<i32>();
2407                    decoded_values.extend_from_slice(values.as_ref());
2408                }
2409                _ => panic!("Expected FixedWidth block"),
2410            }
2411            values_processed += chunk_values;
2412        }
2413
2414        assert_eq!(values_processed, compressed.num_values);
2415        decoded_values
2416    }
2417
2418    #[cfg(feature = "bitpacking")]
2419    fn monotonic_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2420        let mut values = Vec::with_capacity(num_runs * run_length);
2421        for run in 0..num_runs {
2422            values.extend(std::iter::repeat_n(run as i32, run_length));
2423        }
2424        values
2425    }
2426
2427    #[cfg(feature = "bitpacking")]
2428    fn high_entropy_runs(num_runs: usize, run_length: usize) -> Vec<i32> {
2429        let mut values = Vec::with_capacity(num_runs * run_length);
2430        let mut state = 7u64;
2431        for _ in 0..num_runs {
2432            state = state
2433                .wrapping_mul(6364136223846793005)
2434                .wrapping_add(1442695040888963407);
2435            values.extend(std::iter::repeat_n((state >> 32) as i32, run_length));
2436        }
2437        values
2438    }
2439
2440    #[test]
2441    fn test_select_run_length_width_prefers_u16_for_long_runs() {
2442        let mut entries = [0u64; 3];
2443        accumulate_run_length_entries(300, Some(*MAX_MINIBLOCK_VALUES), &mut entries);
2444        let (width, _) = select_run_length_width_from_entries(&entries, 32).unwrap();
2445        assert_eq!(width, RunLengthWidth::U16);
2446    }
2447
2448    // ========== Round-trip Tests for Different Types ==========
2449
2450    #[test]
2451    fn test_round_trip_all_types() {
2452        // Test u8
2453        test_round_trip_helper(vec![42u8, 42, 42, 100, 100, 255, 255, 255, 255], 8);
2454
2455        // Test u16
2456        test_round_trip_helper(vec![1000u16, 1000, 2000, 2000, 2000, 3000], 16);
2457
2458        // Test i32
2459        test_round_trip_helper(vec![100i32, 100, 100, -200, -200, 300, 300, 300, 300], 32);
2460
2461        // Test u64
2462        test_round_trip_helper(vec![1_000_000_000u64; 5], 64);
2463    }
2464
2465    fn test_round_trip_helper<T>(data: Vec<T>, bits_per_value: u64)
2466    where
2467        T: bytemuck::Pod + PartialEq + std::fmt::Debug,
2468    {
2469        let encoder = RleEncoder::new();
2470        let bytes: Vec<u8> = data
2471            .iter()
2472            .flat_map(|v| bytemuck::bytes_of(v))
2473            .copied()
2474            .collect();
2475
2476        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2477            bits_per_value,
2478            data: LanceBuffer::from(bytes),
2479            num_values: data.len() as u64,
2480            block_info: BlockInfo::default(),
2481        });
2482
2483        let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
2484        let decompressor = RleDecompressor::new(bits_per_value);
2485        let decompressed = MiniBlockDecompressor::decompress(
2486            &decompressor,
2487            compressed.data,
2488            compressed.num_values,
2489        )
2490        .unwrap();
2491
2492        match decompressed {
2493            DataBlock::FixedWidth(ref block) => {
2494                // Verify the decompressed data length matches expected
2495                assert_eq!(block.data.len(), data.len() * std::mem::size_of::<T>());
2496            }
2497            _ => panic!("Expected FixedWidth block"),
2498        }
2499    }
2500
2501    // ========== Chunk Boundary Tests ==========
2502
2503    #[test]
2504    fn test_power_of_two_chunking() {
2505        let encoder = RleEncoder::new();
2506
2507        // Create data that will require multiple chunks
2508        let test_sizes = vec![1000, 2500, 5000, 10000];
2509
2510        for size in test_sizes {
2511            let data: Vec<i32> = (0..size)
2512                .map(|i| i / 50) // Create runs of 50
2513                .collect();
2514
2515            let array = Int32Array::from(data);
2516            let (compressed, _) =
2517                compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2518
2519            // Verify all non-last chunks have power-of-2 values
2520            for (i, chunk) in compressed.chunks.iter().enumerate() {
2521                if i < compressed.chunks.len() - 1 {
2522                    assert!(chunk.log_num_values > 0);
2523                    let chunk_values = 1u64 << chunk.log_num_values;
2524                    assert!(chunk_values.is_power_of_two());
2525                    assert!(chunk_values <= *MAX_MINIBLOCK_VALUES);
2526                } else {
2527                    assert_eq!(chunk.log_num_values, 0);
2528                }
2529            }
2530        }
2531    }
2532
2533    #[rstest]
2534    #[case::u8_lengths(RunLengthWidth::U8)]
2535    #[case::u16_lengths(RunLengthWidth::U16)]
2536    #[case::u32_lengths(RunLengthWidth::U32)]
2537    fn test_miniblock_chunk_counts_match_encoded_runs(#[case] run_length_width: RunLengthWidth) {
2538        // This pattern crosses the 2,048-value boundary in the middle of a two-value run.
2539        let levels = (0..4098)
2540            .map(|index| if index % 3 == 0 { 1u16 } else { 0u16 })
2541            .collect::<Vec<_>>();
2542        let num_values = levels.len() as u64;
2543        let encoder = RleEncoder::with_run_length_width(run_length_width);
2544        let (buffers, chunks) = encoder
2545            .encode_data(
2546                &LanceBuffer::reinterpret_vec(levels),
2547                num_values,
2548                u16::BITS as u64,
2549            )
2550            .unwrap();
2551
2552        assert_eq!(buffers.len(), 2);
2553        let bytes_per_length = run_length_width.bytes_per_value();
2554        let mut values_offset = 0usize;
2555        let mut lengths_offset = 0usize;
2556        let mut values_processed = 0u64;
2557
2558        for chunk in &chunks {
2559            let values_size = chunk.buffer_sizes[0] as usize;
2560            let lengths_size = chunk.buffer_sizes[1] as usize;
2561            let lengths_end = lengths_offset + lengths_size;
2562            let chunk_lengths = &buffers[1].as_ref()[lengths_offset..lengths_end];
2563            let length_chunks = chunk_lengths.chunks_exact(bytes_per_length);
2564            assert!(length_chunks.remainder().is_empty());
2565            let num_runs = length_chunks.len();
2566            let encoded_values = length_chunks
2567                .map(|bytes| run_length_width.read_length(bytes))
2568                .sum::<u64>();
2569            let declared_values = chunk.num_values(values_processed, num_values);
2570
2571            assert_eq!(values_size, num_runs * size_of::<u16>());
2572            assert_eq!(encoded_values, declared_values);
2573
2574            values_offset += values_size;
2575            lengths_offset = lengths_end;
2576            values_processed += declared_values;
2577        }
2578
2579        assert_eq!(values_processed, num_values);
2580        assert_eq!(values_offset, buffers[0].len());
2581        assert_eq!(lengths_offset, buffers[1].len());
2582    }
2583
2584    // ========== Error Handling Tests ==========
2585
2586    #[test]
2587    fn test_encoder_rejects_zero_progress() {
2588        let error = RleEncoder::new()
2589            .encode_data(&LanceBuffer::empty(), 1, u16::BITS as u64)
2590            .unwrap_err();
2591
2592        assert!(
2593            matches!(&error, Error::Internal { .. }),
2594            "expected internal error, got: {error:?}"
2595        );
2596        assert!(error.to_string().contains("made no progress"));
2597        assert!(error.to_string().contains("values_remaining=1"));
2598    }
2599
2600    #[test]
2601    fn test_invalid_buffer_count() {
2602        let decompressor = RleDecompressor::new(32);
2603        let result = MiniBlockDecompressor::decompress(
2604            &decompressor,
2605            vec![LanceBuffer::from(vec![1, 2, 3, 4])],
2606            10,
2607        );
2608        assert!(result.is_err());
2609        assert!(
2610            result
2611                .unwrap_err()
2612                .to_string()
2613                .contains("expects exactly 2 buffers")
2614        );
2615    }
2616
2617    #[test]
2618    fn test_buffer_consistency() {
2619        let decompressor = RleDecompressor::new(32);
2620        let values = LanceBuffer::from(vec![1, 0, 0, 0]); // 1 i32 value
2621        let lengths = LanceBuffer::from(vec![5, 10]); // 2 lengths - mismatch!
2622        let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 15);
2623        assert!(result.is_err());
2624        assert!(
2625            result
2626                .unwrap_err()
2627                .to_string()
2628                .contains("Inconsistent RLE buffers")
2629        );
2630    }
2631
2632    #[test]
2633    fn test_u16_length_buffer_must_be_aligned() {
2634        let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2635        let values = LanceBuffer::from(vec![1, 0, 0, 0]);
2636        let lengths = LanceBuffer::from(vec![5]);
2637        let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5);
2638        assert!(matches!(&result, Err(Error::InvalidInput { .. })));
2639        assert!(
2640            result
2641                .unwrap_err()
2642                .to_string()
2643                .contains("not divisible by 2")
2644        );
2645    }
2646
2647    #[test]
2648    fn test_rle_rejects_underflow_and_zero_lengths_and_clamps_overflow() {
2649        let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2650        let value = LanceBuffer::from(1i32.to_le_bytes().to_vec());
2651
2652        let underflow = MiniBlockDecompressor::decompress(
2653            &decompressor,
2654            vec![
2655                value.clone(),
2656                LanceBuffer::from(4u16.to_le_bytes().to_vec()),
2657            ],
2658            5,
2659        )
2660        .unwrap_err();
2661        assert!(underflow.to_string().contains("produced 4 values"));
2662
2663        let overflow = MiniBlockDecompressor::decompress(
2664            &decompressor,
2665            vec![
2666                value.clone(),
2667                LanceBuffer::from(6u16.to_le_bytes().to_vec()),
2668            ],
2669            5,
2670        )
2671        .unwrap();
2672        match overflow {
2673            DataBlock::FixedWidth(block) => {
2674                assert_eq!(block.num_values, 5);
2675                let decoded = block.data.borrow_to_typed_slice::<i32>();
2676                assert_eq!(decoded.as_ref(), &[1i32; 5]);
2677            }
2678            _ => panic!("Expected FixedWidth block"),
2679        }
2680
2681        let zero = MiniBlockDecompressor::decompress(
2682            &decompressor,
2683            vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())],
2684            5,
2685        )
2686        .unwrap_err();
2687        assert!(zero.to_string().contains("zero run length"));
2688    }
2689
2690    #[test]
2691    fn test_block_rle_rejects_overflow() {
2692        // Block payloads have no chunk boundaries, so run lengths summing past
2693        // num_values can only be corruption and must stay a hard error.
2694        let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2695        let values = 1i32.to_le_bytes();
2696        let lengths = 6u16.to_le_bytes();
2697        let mut payload = Vec::new();
2698        payload.extend_from_slice(&(values.len() as u64).to_le_bytes());
2699        payload.extend_from_slice(&values);
2700        payload.extend_from_slice(&lengths);
2701
2702        let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5)
2703            .unwrap_err();
2704        assert!(matches!(&error, Error::InvalidInput { .. }));
2705        assert!(
2706            error
2707                .to_string()
2708                .contains("overflowed expected value count")
2709        );
2710    }
2711
2712    #[test]
2713    fn test_rle_truncates_legacy_chunk_boundary_overflow() {
2714        // Legacy encoders emitted chunks declaring 2048 values whose final run crossed
2715        // the checkpoint boundary (e.g. run lengths summing to 2080); the excess values
2716        // are duplicated at the start of the next chunk and must be ignored here.
2717        let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16);
2718        let mut values = Vec::new();
2719        values.extend_from_slice(&7i32.to_le_bytes());
2720        values.extend_from_slice(&8i32.to_le_bytes());
2721        let mut lengths = Vec::new();
2722        lengths.extend_from_slice(&2000u16.to_le_bytes());
2723        lengths.extend_from_slice(&80u16.to_le_bytes());
2724
2725        let decoded = MiniBlockDecompressor::decompress(
2726            &decompressor,
2727            vec![LanceBuffer::from(values), LanceBuffer::from(lengths)],
2728            2048,
2729        )
2730        .unwrap();
2731        match decoded {
2732            DataBlock::FixedWidth(block) => {
2733                assert_eq!(block.num_values, 2048);
2734                let decoded = block.data.borrow_to_typed_slice::<i32>();
2735                let decoded = decoded.as_ref();
2736                assert_eq!(decoded.len(), 2048);
2737                assert!(decoded[..2000].iter().all(|&v| v == 7));
2738                assert!(decoded[2000..].iter().all(|&v| v == 8));
2739            }
2740            _ => panic!("Expected FixedWidth block"),
2741        }
2742    }
2743
2744    #[test]
2745    fn test_empty_data_handling() {
2746        let encoder = RleEncoder::new();
2747
2748        // Test empty block
2749        let empty_block = DataBlock::FixedWidth(FixedWidthDataBlock {
2750            bits_per_value: 32,
2751            data: LanceBuffer::from(vec![]),
2752            num_values: 0,
2753            block_info: BlockInfo::default(),
2754        });
2755
2756        let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap();
2757        assert_eq!(compressed.num_values, 0);
2758        assert!(compressed.data.is_empty());
2759
2760        // Test decompression of empty data
2761        let decompressor = RleDecompressor::new(32);
2762        let decompressed = MiniBlockDecompressor::decompress(&decompressor, vec![], 0).unwrap();
2763
2764        match decompressed {
2765            DataBlock::FixedWidth(ref block) => {
2766                assert_eq!(block.num_values, 0);
2767                assert_eq!(block.data.len(), 0);
2768            }
2769            _ => panic!("Expected FixedWidth block"),
2770        }
2771    }
2772
2773    // ========== Integration Test ==========
2774
2775    #[test]
2776    fn test_multi_chunk_round_trip() {
2777        let encoder = RleEncoder::new();
2778
2779        // Create data that spans multiple chunks with mixed patterns
2780        let mut data = Vec::new();
2781
2782        // High compression section
2783        data.extend(vec![999i32; 2000]);
2784        // Low compression section
2785        data.extend(0..1000);
2786        // Another high compression section
2787        data.extend(vec![777i32; 2000]);
2788
2789        let array = Int32Array::from(data.clone());
2790        let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2791
2792        // Manually decompress all chunks
2793        let mut reconstructed = Vec::new();
2794        let mut values_offset = 0usize;
2795        let mut lengths_offset = 0usize;
2796        let mut values_processed = 0u64;
2797
2798        // We now have exactly 2 global buffers
2799        assert_eq!(compressed.data.len(), 2);
2800        let global_values = &compressed.data[0];
2801        let global_lengths = &compressed.data[1];
2802
2803        for chunk in &compressed.chunks {
2804            let chunk_values = if chunk.log_num_values > 0 {
2805                1u64 << chunk.log_num_values
2806            } else {
2807                compressed.num_values - values_processed
2808            };
2809
2810            // Extract chunk buffers from global buffers using buffer_sizes
2811            let values_size = chunk.buffer_sizes[0] as usize;
2812            let lengths_size = chunk.buffer_sizes[1] as usize;
2813
2814            let chunk_values_buffer = global_values.slice_with_length(values_offset, values_size);
2815            let chunk_lengths_buffer =
2816                global_lengths.slice_with_length(lengths_offset, lengths_size);
2817
2818            let decompressor = RleDecompressor::new(32);
2819            let chunk_data = MiniBlockDecompressor::decompress(
2820                &decompressor,
2821                vec![chunk_values_buffer, chunk_lengths_buffer],
2822                chunk_values,
2823            )
2824            .unwrap();
2825
2826            values_offset += values_size;
2827            lengths_offset += lengths_size;
2828            values_processed += chunk_values;
2829
2830            match chunk_data {
2831                DataBlock::FixedWidth(ref block) => {
2832                    let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2833                    reconstructed.extend_from_slice(values);
2834                }
2835                _ => panic!("Expected FixedWidth block"),
2836            }
2837        }
2838
2839        assert_eq!(reconstructed, data);
2840    }
2841
2842    #[test]
2843    fn test_1024_boundary_conditions() {
2844        // Comprehensive test for various boundary conditions at 1024 values
2845        // This consolidates multiple bug tests that were previously separate
2846        let encoder = RleEncoder::new();
2847        let decompressor = RleDecompressor::new(32);
2848
2849        let test_cases = [
2850            ("runs_of_2", {
2851                let mut data = Vec::new();
2852                for i in 0..512 {
2853                    data.push(i);
2854                    data.push(i);
2855                }
2856                data
2857            }),
2858            ("single_run_1024", vec![42i32; 1024]),
2859            ("alternating_values", {
2860                let mut data = Vec::new();
2861                for i in 0..1024 {
2862                    data.push(i % 2);
2863                }
2864                data
2865            }),
2866            ("run_boundary_255s", {
2867                let mut data = Vec::new();
2868                data.extend(vec![1i32; 255]);
2869                data.extend(vec![2i32; 255]);
2870                data.extend(vec![3i32; 255]);
2871                data.extend(vec![4i32; 255]);
2872                data.extend(vec![5i32; 4]);
2873                data
2874            }),
2875            ("unique_values_1024", (0..1024).collect::<Vec<_>>()),
2876            ("unique_plus_duplicate", {
2877                // 1023 unique values followed by one duplicate (regression test)
2878                let mut data = Vec::new();
2879                for i in 0..1023 {
2880                    data.push(i);
2881                }
2882                data.push(1022i32); // Last value same as second-to-last
2883                data
2884            }),
2885            ("bug_4092_pattern", {
2886                // Test exact scenario that produces 4092 bytes instead of 4096
2887                let mut data = Vec::new();
2888                for i in 0..1022 {
2889                    data.push(i);
2890                }
2891                data.push(999999i32);
2892                data.push(999999i32);
2893                data
2894            }),
2895        ];
2896
2897        for (test_name, data) in test_cases.iter() {
2898            assert_eq!(data.len(), 1024, "Test case {} has wrong length", test_name);
2899
2900            // Compress the data
2901            let array = Int32Array::from(data.clone());
2902            let (compressed, _) =
2903                compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap();
2904
2905            // Decompress and verify
2906            match MiniBlockDecompressor::decompress(
2907                &decompressor,
2908                compressed.data,
2909                compressed.num_values,
2910            ) {
2911                Ok(decompressed) => match decompressed {
2912                    DataBlock::FixedWidth(ref block) => {
2913                        let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
2914                        assert_eq!(
2915                            values.len(),
2916                            1024,
2917                            "Test case {} got {} values, expected 1024",
2918                            test_name,
2919                            values.len()
2920                        );
2921                        assert_eq!(
2922                            block.data.len(),
2923                            4096,
2924                            "Test case {} got {} bytes, expected 4096",
2925                            test_name,
2926                            block.data.len()
2927                        );
2928                        assert_eq!(values, &data[..], "Test case {} data mismatch", test_name);
2929                    }
2930                    _ => panic!("Test case {} expected FixedWidth block", test_name),
2931                },
2932                Err(e) => {
2933                    if e.to_string().contains("4092") {
2934                        panic!("Test case {} found bug 4092: {}", test_name, e);
2935                    }
2936                    panic!("Test case {} failed with error: {}", test_name, e);
2937                }
2938            }
2939        }
2940    }
2941
2942    #[test]
2943    fn test_low_repetition_50pct_bug() {
2944        // Test case that reproduces the 4092 bytes bug with low repetition (50%)
2945        // This simulates the 1M benchmark case
2946        let encoder = RleEncoder::new();
2947
2948        // Create 1M values with low repetition (50% chance of change)
2949        let num_values = 1_048_576; // 1M values
2950        let mut data = Vec::with_capacity(num_values);
2951        let mut value = 0i32;
2952        let mut rng = 12345u64; // Simple deterministic RNG
2953
2954        for _ in 0..num_values {
2955            data.push(value);
2956            // Simple LCG for deterministic randomness
2957            rng = rng.wrapping_mul(1664525).wrapping_add(1013904223);
2958            // 50% chance to increment value
2959            if (rng >> 16) & 1 == 1 {
2960                value += 1;
2961            }
2962        }
2963
2964        let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
2965
2966        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
2967            bits_per_value: 32,
2968            data: LanceBuffer::from(bytes),
2969            num_values: num_values as u64,
2970            block_info: BlockInfo::default(),
2971        });
2972
2973        let (compressed, _) = compress_miniblock(&encoder, block).unwrap();
2974
2975        // Debug first few chunks
2976        for (i, chunk) in compressed.chunks.iter().take(5).enumerate() {
2977            let _chunk_values = if chunk.log_num_values > 0 {
2978                1 << chunk.log_num_values
2979            } else {
2980                // Last chunk - calculate remaining
2981                let prev_total: usize = compressed.chunks[..i]
2982                    .iter()
2983                    .map(|c| 1usize << c.log_num_values)
2984                    .sum();
2985                num_values - prev_total
2986            };
2987        }
2988
2989        // Try to decompress
2990        let decompressor = RleDecompressor::new(32);
2991        match MiniBlockDecompressor::decompress(
2992            &decompressor,
2993            compressed.data,
2994            compressed.num_values,
2995        ) {
2996            Ok(decompressed) => match decompressed {
2997                DataBlock::FixedWidth(ref block) => {
2998                    assert_eq!(
2999                        block.data.len(),
3000                        num_values * 4,
3001                        "Expected {} bytes but got {}",
3002                        num_values * 4,
3003                        block.data.len()
3004                    );
3005                }
3006                _ => panic!("Expected FixedWidth block"),
3007            },
3008            Err(e) => {
3009                if e.to_string().contains("4092") {
3010                    panic!("Bug reproduced! {}", e);
3011                } else {
3012                    panic!("Unexpected error: {}", e);
3013                }
3014            }
3015        }
3016    }
3017
3018    // ========== Encoding Verification Tests ==========
3019
3020    #[test_log::test(tokio::test)]
3021    async fn test_rle_encoding_verification() {
3022        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
3023        use arrow_array::{Array, Int32Array};
3024        use lance_datagen::{ArrayGenerator, RowCount};
3025        use std::collections::HashMap;
3026        use std::sync::Arc;
3027
3028        let test_cases = TestCases::default()
3029            .with_expected_encoding("rle")
3030            .with_structural_encodings();
3031
3032        // Test both explicit metadata and automatic selection
3033        // 1. Test with explicit RLE threshold metadata (also disable BSS)
3034        let mut metadata_explicit = HashMap::new();
3035        metadata_explicit.insert(
3036            "lance-encoding:rle-threshold".to_string(),
3037            "0.8".to_string(),
3038        );
3039        metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
3040
3041        let mut generator = RleDataGenerator::new(vec![
3042            i32::MIN,
3043            i32::MIN,
3044            i32::MIN,
3045            i32::MIN,
3046            i32::MIN + 1,
3047            i32::MIN + 1,
3048            i32::MIN + 1,
3049            i32::MIN + 1,
3050            i32::MIN + 2,
3051            i32::MIN + 2,
3052            i32::MIN + 2,
3053            i32::MIN + 2,
3054        ]);
3055        let data_explicit = generator.generate_default(RowCount::from(10000)).unwrap();
3056        check_round_trip_encoding_of_data(vec![data_explicit], &test_cases, metadata_explicit)
3057            .await;
3058
3059        // 2. Test automatic RLE selection based on data characteristics
3060        // 80% repetition should trigger RLE (> default 50% threshold).
3061        //
3062        // Use values with the high bit set so bitpacking can't shrink the values.
3063        // Explicitly disable BSS to ensure RLE is tested
3064        let mut metadata = HashMap::new();
3065        metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3066
3067        let mut values = vec![i32::MIN; 8000]; // 80% repetition
3068        values.extend(
3069            [
3070                i32::MIN + 1,
3071                i32::MIN + 2,
3072                i32::MIN + 3,
3073                i32::MIN + 4,
3074                i32::MIN + 5,
3075            ]
3076            .repeat(400),
3077        ); // 20% variety
3078        let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3079        check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3080
3081        #[cfg(any(feature = "lz4", feature = "zstd"))]
3082        {
3083            let mut metadata = HashMap::new();
3084            metadata.insert(
3085                "lance-encoding:rle-threshold".to_string(),
3086                "0.8".to_string(),
3087            );
3088            metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
3089            metadata.insert(
3090                "lance-encoding:compression".to_string(),
3091                if cfg!(feature = "zstd") {
3092                    "zstd".to_string()
3093                } else {
3094                    "lz4".to_string()
3095                },
3096            );
3097            let mut values = Vec::with_capacity(2048 * 4);
3098            for run in 0..2048 {
3099                values.extend(std::iter::repeat_n(i32::MIN + (run % 8), 4));
3100            }
3101            let arr = Arc::new(Int32Array::from(values)) as Arc<dyn Array>;
3102            check_round_trip_encoding_of_data(vec![arr], &test_cases, metadata).await;
3103        }
3104    }
3105
3106    /// Generator that produces repetitive patterns suitable for RLE
3107    #[derive(Debug)]
3108    struct RleDataGenerator {
3109        pattern: Vec<i32>,
3110        idx: usize,
3111    }
3112
3113    impl RleDataGenerator {
3114        fn new(pattern: Vec<i32>) -> Self {
3115            Self { pattern, idx: 0 }
3116        }
3117    }
3118
3119    impl lance_datagen::ArrayGenerator for RleDataGenerator {
3120        fn generate(
3121            &mut self,
3122            _length: lance_datagen::RowCount,
3123            _rng: &mut rand_xoshiro::Xoshiro256PlusPlus,
3124        ) -> std::result::Result<std::sync::Arc<dyn arrow_array::Array>, arrow_schema::ArrowError>
3125        {
3126            use arrow_array::Int32Array;
3127            use std::sync::Arc;
3128
3129            // Generate enough repetitive data to trigger RLE
3130            let mut values = Vec::new();
3131            for _ in 0..10000 {
3132                values.push(self.pattern[self.idx]);
3133                self.idx = (self.idx + 1) % self.pattern.len();
3134            }
3135            Ok(Arc::new(Int32Array::from(values)))
3136        }
3137
3138        fn data_type(&self) -> &arrow_schema::DataType {
3139            &arrow_schema::DataType::Int32
3140        }
3141
3142        fn element_size_bytes(&self) -> Option<lance_datagen::ByteCount> {
3143            Some(lance_datagen::ByteCount::from(4))
3144        }
3145    }
3146
3147    // ========== Block Related tests ==========
3148    #[test]
3149    fn test_block_decompressor_rejects_overflowing_values_size() {
3150        let decompressor = RleDecompressor::new(32);
3151
3152        let mut data = Vec::new();
3153        data.extend_from_slice(&u64::MAX.to_le_bytes());
3154        let result = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(data), 1);
3155        assert!(result.is_err());
3156        assert!(
3157            result
3158                .unwrap_err()
3159                .to_string()
3160                .contains("Invalid RLE values buffer size")
3161        );
3162    }
3163
3164    #[test]
3165    fn test_block_decompressor_too_small() {
3166        let decompressor = RleDecompressor::new(32);
3167        let result =
3168            BlockDecompressor::decompress(&decompressor, LanceBuffer::from(vec![1, 2, 3]), 10);
3169        assert!(result.is_err());
3170        assert!(
3171            result
3172                .unwrap_err()
3173                .to_string()
3174                .contains("Insufficient data size: 3")
3175        );
3176    }
3177
3178    #[test]
3179    fn test_block_compressor_header_format() {
3180        let encoder = RleEncoder::new();
3181
3182        let data = vec![1i32, 1, 1];
3183        let array = Int32Array::from(data);
3184        let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
3185
3186        // Verify header format: first 8 bytes should be values_size as u64
3187        assert!(compressed.len() >= 8);
3188        let values_size_bytes: [u8; 8] = compressed.as_ref()[..8].try_into().unwrap();
3189        let values_size = u64::from_le_bytes(values_size_bytes);
3190
3191        // Values buffer should contain 1 i32 value (4 bytes)
3192        assert_eq!(values_size, 4);
3193
3194        // Total size should be: 8 (header) + 4 (values) + 1 (lengths)
3195        assert_eq!(compressed.len(), 13);
3196    }
3197
3198    #[test]
3199    fn test_block_compressor_round_trip() {
3200        let encoder = RleEncoder::new();
3201        let decompressor = RleDecompressor::new(32);
3202
3203        // Test basic pattern
3204        let data = vec![1i32, 1, 1, 2, 2, 3, 3, 3, 3];
3205        let array = Int32Array::from(data.clone());
3206        let data_block = DataBlock::from_array(array);
3207
3208        let compressed = BlockCompressor::compress(&encoder, data_block).unwrap();
3209        let decompressed =
3210            BlockDecompressor::decompress(&decompressor, compressed, data.len() as u64).unwrap();
3211
3212        match decompressed {
3213            DataBlock::FixedWidth(block) => {
3214                let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3215                assert_eq!(values, &data[..]);
3216            }
3217            _ => panic!("Expected FixedWidth block"),
3218        }
3219    }
3220
3221    #[test]
3222    fn test_block_compressor_large_data() {
3223        let encoder = RleEncoder::new();
3224        let decompressor = RleDecompressor::new(32);
3225
3226        // Create data that will span multiple chunks
3227        // Each chunks can handle ~2048 values, so use 10K values
3228        let mut data = Vec::new();
3229        data.extend(vec![999i32; 3000]); // First ~2 chunks
3230        data.extend(vec![777i32; 3000]); // Next ~2 chunks
3231        data.extend(vec![555i32; 4000]); // Final ~2 chunks
3232
3233        let total_values = data.len();
3234        assert_eq!(total_values, 10000);
3235
3236        let array = Int32Array::from(data.clone());
3237        let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap();
3238        let decompressed =
3239            BlockDecompressor::decompress(&decompressor, compressed, total_values as u64).unwrap();
3240
3241        match decompressed {
3242            DataBlock::FixedWidth(block) => {
3243                let values: &[i32] = bytemuck::cast_slice(block.data.as_ref());
3244                assert_eq!(values.len(), total_values);
3245                assert_eq!(values, &data[..]);
3246            }
3247            _ => panic!("Expected FixedWidth block"),
3248        }
3249    }
3250}