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