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