Skip to main content

lance_encoding/encodings/physical/
binary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Basic encodings for variable width data
5//!
6//! These are not compression but represent the "leaf" encodings for variable length data
7//! where we simply match the data with the rules of the structural encoding.
8//!
9//! These encodings are transparent since we aren't actually doing any compression.  No information
10//! is needed in the encoding description.
11
12use arrow_array::OffsetSizeTrait;
13use byteorder::{ByteOrder, LittleEndian};
14use core::panic;
15
16use crate::compression::{
17    BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor,
18};
19
20use crate::buffer::LanceBuffer;
21use crate::data::{BlockInfo, DataBlock, VariableWidthBlock};
22use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock};
23use crate::encodings::logical::primitive::miniblock::{
24    MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext,
25    MiniBlockCompressor,
26};
27use crate::format::pb21::CompressiveEncoding;
28use crate::format::pb21::compressive_encoding::Compression;
29use crate::format::{ProtobufUtils21, pb21};
30
31use lance_core::utils::bit::pad_bytes_to;
32use lance_core::{Error, Result};
33
34#[derive(Debug)]
35pub struct BinaryMiniBlockEncoder {
36    minichunk_size: i64,
37}
38
39impl Default for BinaryMiniBlockEncoder {
40    fn default() -> Self {
41        Self {
42            minichunk_size: *AIM_MINICHUNK_SIZE,
43        }
44    }
45}
46
47const DEFAULT_AIM_MINICHUNK_SIZE: i64 = 4 * 1024;
48
49pub static AIM_MINICHUNK_SIZE: std::sync::LazyLock<i64> = std::sync::LazyLock::new(|| {
50    std::env::var("LANCE_BINARY_MINIBLOCK_CHUNK_SIZE")
51        .unwrap_or_else(|_| DEFAULT_AIM_MINICHUNK_SIZE.to_string())
52        .parse::<i64>()
53        .unwrap_or(DEFAULT_AIM_MINICHUNK_SIZE)
54});
55
56// Make it to support both u32 and u64
57fn chunk_offsets<N: OffsetSizeTrait>(
58    offsets: &[N],
59    data: &[u8],
60    alignment: usize,
61    minichunk_size: i64,
62) -> (Vec<LanceBuffer>, Vec<MiniBlockChunk>) {
63    #[derive(Debug)]
64    struct ChunkInfo {
65        chunk_start_offset_in_orig_idx: usize,
66        chunk_last_offset_in_orig_idx: usize,
67        // the bytes in every chunk starts at `chunk.bytes_start_offset`
68        bytes_start_offset: usize,
69        // every chunk is padded to 8 bytes.
70        // we need to interpret every chunk as &[u32] so we need it to padded at least to 4 bytes,
71        // this field can actually be eliminated and I can use `num_bytes` in `MiniBlockChunk` to compute
72        // the `output_total_bytes`.
73        padded_chunk_size: usize,
74    }
75
76    let byte_width: usize = N::get_byte_width();
77    let mut chunks_info = vec![];
78    let mut chunks = vec![];
79    let mut last_offset_in_orig_idx = 0;
80    loop {
81        let this_last_offset_in_orig_idx =
82            search_next_offset_idx(offsets, last_offset_in_orig_idx, minichunk_size);
83
84        let num_values_in_this_chunk = this_last_offset_in_orig_idx - last_offset_in_orig_idx;
85        let chunk_bytes = offsets[this_last_offset_in_orig_idx] - offsets[last_offset_in_orig_idx];
86        let this_chunk_size =
87            (num_values_in_this_chunk + 1) * byte_width + chunk_bytes.to_usize().unwrap();
88
89        let padded_chunk_size = this_chunk_size.next_multiple_of(alignment);
90        debug_assert!(padded_chunk_size > 0);
91
92        let this_chunk_bytes_start_offset = (num_values_in_this_chunk + 1) * byte_width;
93        chunks_info.push(ChunkInfo {
94            chunk_start_offset_in_orig_idx: last_offset_in_orig_idx,
95            chunk_last_offset_in_orig_idx: this_last_offset_in_orig_idx,
96            bytes_start_offset: this_chunk_bytes_start_offset,
97            padded_chunk_size,
98        });
99        chunks.push(MiniBlockChunk {
100            log_num_values: if this_last_offset_in_orig_idx == offsets.len() - 1 {
101                0
102            } else {
103                num_values_in_this_chunk.trailing_zeros() as u8
104            },
105            buffer_sizes: vec![padded_chunk_size as u32],
106        });
107        if this_last_offset_in_orig_idx == offsets.len() - 1 {
108            break;
109        }
110        last_offset_in_orig_idx = this_last_offset_in_orig_idx;
111    }
112
113    let output_total_bytes = chunks_info
114        .iter()
115        .map(|chunk_info| chunk_info.padded_chunk_size)
116        .sum::<usize>();
117
118    let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
119
120    for chunk in chunks_info {
121        let this_chunk_offsets: Vec<N> = offsets
122            [chunk.chunk_start_offset_in_orig_idx..=chunk.chunk_last_offset_in_orig_idx]
123            .iter()
124            .map(|offset| {
125                *offset - offsets[chunk.chunk_start_offset_in_orig_idx]
126                    + N::from_usize(chunk.bytes_start_offset).unwrap()
127            })
128            .collect();
129
130        let this_chunk_offsets = LanceBuffer::reinterpret_vec(this_chunk_offsets);
131        output.extend_from_slice(&this_chunk_offsets);
132
133        let start_in_orig = offsets[chunk.chunk_start_offset_in_orig_idx]
134            .to_usize()
135            .unwrap();
136        let end_in_orig = offsets[chunk.chunk_last_offset_in_orig_idx]
137            .to_usize()
138            .unwrap();
139        output.extend_from_slice(&data[start_in_orig..end_in_orig]);
140
141        // pad this chunk to make it align to desired bytes.
142        const PAD_BYTE: u8 = 72;
143        let pad_len = pad_bytes_to(output.len(), alignment);
144
145        // Compare with usize literal to avoid type mismatch with N
146        if pad_len > 0_usize {
147            output.extend(std::iter::repeat_n(PAD_BYTE, pad_len));
148        }
149    }
150    (vec![LanceBuffer::reinterpret_vec(output)], chunks)
151}
152
153// search for the next offset index to cut the values into a chunk.
154// this function incrementally peek the number of values in a chunk,
155// each time multiplies the number of values by 2.
156// It returns the offset_idx in `offsets` that belongs to this chunk.
157fn search_next_offset_idx<N: OffsetSizeTrait>(
158    offsets: &[N],
159    last_offset_idx: usize,
160    minichunk_size: i64,
161) -> usize {
162    // MiniBlockChunk uses `log_num_values == 0` as a sentinel for the final chunk. This means we
163    // must avoid creating 1-value chunks except for the final chunk, even if the configured
164    // `minichunk_size` is too small to fit more than one value.
165    let remaining_values = offsets.len().saturating_sub(last_offset_idx + 1);
166    if remaining_values <= 1 {
167        return offsets.len() - 1;
168    }
169
170    let mut num_values = 2;
171    let mut new_num_values = num_values * 2;
172    loop {
173        if last_offset_idx + new_num_values >= offsets.len() {
174            let existing_bytes = offsets[offsets.len() - 1] - offsets[last_offset_idx];
175            // existing bytes plus the new offset size
176            let new_size = existing_bytes
177                + N::from_usize((offsets.len() - last_offset_idx) * N::get_byte_width()).unwrap();
178            if new_size.to_i64().unwrap() <= minichunk_size {
179                // case 1: can fit the rest of all data into a miniblock
180                return offsets.len() - 1;
181            } else {
182                // case 2: can only fit the last tried `num_values` into a miniblock
183                return last_offset_idx + num_values;
184            }
185        }
186        let existing_bytes = offsets[last_offset_idx + new_num_values] - offsets[last_offset_idx];
187        let new_size =
188            existing_bytes + N::from_usize((new_num_values + 1) * N::get_byte_width()).unwrap();
189        if new_size.to_i64().unwrap() <= minichunk_size {
190            if new_num_values * 2 > *MAX_MINIBLOCK_VALUES as usize {
191                // hit the max number of values limit
192                break;
193            }
194            num_values = new_num_values;
195            new_num_values *= 2;
196        } else {
197            break;
198        }
199    }
200    last_offset_idx + num_values
201}
202
203impl BinaryMiniBlockEncoder {
204    pub fn new(minichunk_size: Option<i64>) -> Self {
205        Self {
206            minichunk_size: minichunk_size.unwrap_or(*AIM_MINICHUNK_SIZE),
207        }
208    }
209
210    // put binary data into chunks, every chunk is less than or equal to `minichunk_size`.
211    // In each chunk, offsets are put first then followed by binary bytes data, each chunk is padded to 8 bytes.
212    // the offsets in the chunk points to the bytes offset in this chunk.
213    fn chunk_data(&self, data: VariableWidthBlock) -> (MiniBlockCompressed, CompressiveEncoding) {
214        // TODO: Support compression of offsets
215        // TODO: Support general compression of data
216        match data.bits_per_offset {
217            32 => {
218                let offsets = data.offsets.borrow_to_typed_slice::<i32>();
219                let (buffers, chunks) =
220                    chunk_offsets(offsets.as_ref(), &data.data, 4, self.minichunk_size);
221                (
222                    MiniBlockCompressed {
223                        data: buffers,
224                        chunks,
225                        num_values: data.num_values,
226                    },
227                    ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None),
228                )
229            }
230            64 => {
231                let offsets = data.offsets.borrow_to_typed_slice::<i64>();
232                let (buffers, chunks) =
233                    chunk_offsets(offsets.as_ref(), &data.data, 8, self.minichunk_size);
234                (
235                    MiniBlockCompressed {
236                        data: buffers,
237                        chunks,
238                        num_values: data.num_values,
239                    },
240                    ProtobufUtils21::variable(ProtobufUtils21::flat(64, None), None),
241                )
242            }
243            _ => panic!("Unsupported bits_per_offset={}", data.bits_per_offset),
244        }
245    }
246}
247
248impl MiniBlockCompressor for BinaryMiniBlockEncoder {
249    fn compress(
250        &self,
251        _context: MiniBlockCompressionContext,
252        data: DataBlock,
253    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
254        match data {
255            DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)),
256            _ => Err(Error::invalid_input_source(
257                format!(
258                    "Cannot compress a data block of type {} with BinaryMiniBlockEncoder",
259                    data.name()
260                )
261                .into(),
262            )),
263        }
264    }
265}
266
267#[derive(Debug)]
268pub struct BinaryMiniBlockDecompressor {
269    bits_per_offset: u8,
270}
271
272impl BinaryMiniBlockDecompressor {
273    pub fn new(bits_per_offset: u8) -> Self {
274        assert!(bits_per_offset == 32 || bits_per_offset == 64);
275        Self { bits_per_offset }
276    }
277
278    pub fn from_variable(variable: &pb21::Variable) -> Self {
279        if let Compression::Flat(flat) = variable
280            .offsets
281            .as_ref()
282            .unwrap()
283            .compression
284            .as_ref()
285            .unwrap()
286        {
287            Self {
288                bits_per_offset: flat.bits_per_value as u8,
289            }
290        } else {
291            panic!("Unsupported offsets compression: {:?}", variable.offsets);
292        }
293    }
294}
295
296/// Cold path: pinpoint why the chunk-relative offsets of a binary mini-block
297/// chunk failed validation.
298fn chunk_offset_violation_error<T: Copy + Into<u64>>(offsets: &[T], chunk_len: usize) -> Error {
299    let mut previous: u64 = offsets[0].into();
300    for (position, &offset) in offsets.iter().enumerate().skip(1) {
301        let offset: u64 = offset.into();
302        if offset < previous {
303            return Error::corrupt_file_named(
304                "binary mini-block",
305                format!(
306                    "value offset at position {position} decreases: {offset} < {previous} \
307                     (chunk is {chunk_len} bytes)"
308                ),
309            );
310        }
311        previous = offset;
312    }
313    Error::corrupt_file_named(
314        "binary mini-block",
315        format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"),
316    )
317}
318
319impl MiniBlockDecompressor for BinaryMiniBlockDecompressor {
320    // decompress a MiniBlock of binary data, the num_values must be less than or equal
321    // to the number of values this MiniBlock has, BinaryMiniBlock doesn't store `the number of values`
322    // it has so assertion can not be done here and the caller of `decompress` must ensure
323    // `num_values` <= number of values in the chunk.
324    //
325    // The chunk-relative value offsets at the front of the chunk come straight
326    // from the file and are used to slice the chunk buffer, so corrupt values
327    // must surface as a typed error instead of a panic or an out-of-bounds
328    // read.  The monotonicity check rides along the existing rebase loop (the
329    // `&=` accumulation keeps it branchless) so validation adds no extra pass.
330    fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
331        assert_eq!(data.len(), 1);
332        let data = data.into_iter().next().unwrap();
333
334        let bytes_per_offset = self.bits_per_offset as usize / 8;
335        if !data.len().is_multiple_of(bytes_per_offset) {
336            return Err(Error::corrupt_file_named(
337                "binary mini-block",
338                format!(
339                    "chunk size {} is not a multiple of the {}-byte offset width",
340                    data.len(),
341                    bytes_per_offset
342                ),
343            ));
344        }
345        let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| {
346            Error::corrupt_file_named(
347                "binary mini-block",
348                format!("cannot decode {num_values} values from a single chunk"),
349            )
350        })?;
351        if data.len() / bytes_per_offset < num_offsets {
352            return Err(Error::corrupt_file_named(
353                "binary mini-block",
354                format!(
355                    "chunk of {} bytes holds {} offsets but decoding {} values requires {}",
356                    data.len(),
357                    data.len() / bytes_per_offset,
358                    num_values,
359                    num_offsets
360                ),
361            ));
362        }
363
364        // The value region must start past the offsets being decoded, otherwise
365        // the offset table itself aliases into the value bytes.  A lower bound
366        // (not equality) because a prefix read of the chunk legitimately leaves
367        // unrequested offsets between the requested prefix and the values.
368        let min_value_region_start = num_offsets * bytes_per_offset;
369        let value_region_overlap_error = |first: u64| {
370            Error::corrupt_file_named(
371                "binary mini-block",
372                format!(
373                    "value region starts at offset {first} which overlaps the {num_offsets} \
374                     requested offsets ({min_value_region_start} bytes)"
375                ),
376            )
377        };
378
379        if self.bits_per_offset == 64 {
380            let offsets_buffer = data.borrow_to_typed_slice::<u64>();
381            let offsets = &offsets_buffer.as_ref()[..num_offsets];
382
383            let first = offsets[0];
384            if first < min_value_region_start as u64 {
385                return Err(value_region_overlap_error(first));
386            }
387            let mut previous = first;
388            let mut is_monotonic = true;
389            let result_offsets = offsets
390                .iter()
391                .map(|&offset| {
392                    is_monotonic &= previous <= offset;
393                    previous = offset;
394                    offset.wrapping_sub(first)
395                })
396                .collect::<Vec<u64>>();
397            let last = offsets[num_offsets - 1];
398            if !is_monotonic || last as usize > data.len() {
399                return Err(chunk_offset_violation_error(offsets, data.len()));
400            }
401
402            Ok(DataBlock::VariableWidth(VariableWidthBlock {
403                data: LanceBuffer::from(data[first as usize..last as usize].to_vec()),
404                offsets: LanceBuffer::reinterpret_vec(result_offsets),
405                bits_per_offset: 64,
406                num_values,
407                block_info: BlockInfo::new(),
408            }))
409        } else {
410            let offsets_buffer = data.borrow_to_typed_slice::<u32>();
411            let offsets = &offsets_buffer.as_ref()[..num_offsets];
412
413            let first = offsets[0];
414            if (first as u64) < min_value_region_start as u64 {
415                return Err(value_region_overlap_error(first as u64));
416            }
417            let mut previous = first;
418            let mut is_monotonic = true;
419            let result_offsets = offsets
420                .iter()
421                .map(|&offset| {
422                    is_monotonic &= previous <= offset;
423                    previous = offset;
424                    offset.wrapping_sub(first)
425                })
426                .collect::<Vec<u32>>();
427            let last = offsets[num_offsets - 1];
428            if !is_monotonic || last as usize > data.len() {
429                return Err(chunk_offset_violation_error(offsets, data.len()));
430            }
431
432            Ok(DataBlock::VariableWidth(VariableWidthBlock {
433                data: LanceBuffer::from(data[first as usize..last as usize].to_vec()),
434                offsets: LanceBuffer::reinterpret_vec(result_offsets),
435                bits_per_offset: 32,
436                num_values,
437                block_info: BlockInfo::new(),
438            }))
439        }
440    }
441}
442
443/// Most basic encoding for variable-width data which does no compression at all
444/// The DataBlock memory layout looks like below:
445///
446/// | bits_per_offset           | bytes_start_offset        | offsets data | bytes data |
447/// | ------------------------- | ------------------------- | ------------ | ---------- |
448/// | <bits_per_offset>/8 bytes | <bits_per_offset>/8 bytes | offsets_len  | data_len   |
449///
450/// It's used in VariableEncoder and BinaryBlockDecompressor
451///
452#[derive(Debug, Default)]
453pub struct VariableEncoder {}
454
455impl BlockCompressor for VariableEncoder {
456    fn compress(&self, mut data: DataBlock) -> Result<LanceBuffer> {
457        match data {
458            DataBlock::VariableWidth(ref mut variable_width_data) => {
459                match variable_width_data.bits_per_offset {
460                    32 => {
461                        let offsets = variable_width_data.offsets.borrow_to_typed_slice::<u32>();
462                        let offsets = offsets.as_ref();
463                        // The first 4 bytes store the bits per offset, the next 4 bytes store the start
464                        // offset of the bytes data, then offsets data, then bytes data.
465                        let bytes_start_offset = 4 + 4 + std::mem::size_of_val(offsets) as u32;
466
467                        let output_total_bytes =
468                            bytes_start_offset as usize + variable_width_data.data.len();
469                        let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
470
471                        // Store bit_per_offset info
472                        output.extend_from_slice(&(32_u32).to_le_bytes());
473
474                        // store `bytes_start_offset` in the next 4 bytes of output buffer
475                        output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
476
477                        // store offsets
478                        output.extend_from_slice(&variable_width_data.offsets);
479
480                        // store bytes
481                        output.extend_from_slice(&variable_width_data.data);
482                        Ok(LanceBuffer::from(output))
483                    }
484                    64 => {
485                        let offsets = variable_width_data.offsets.borrow_to_typed_slice::<u64>();
486                        let offsets = offsets.as_ref();
487                        // The first 8 bytes store the bits per offset, the next 8 bytes store the start
488                        // offset of the bytes data, then offsets data, then bytes data.
489                        let bytes_start_offset = 8 + 8 + std::mem::size_of_val(offsets) as u64;
490
491                        let output_total_bytes =
492                            bytes_start_offset as usize + variable_width_data.data.len();
493                        let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
494
495                        // Store bit_per_offset info
496                        output.extend_from_slice(&(64_u64).to_le_bytes());
497
498                        // store `bytes_start_offset` in the next 8 bytes of output buffer
499                        output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
500
501                        // store offsets
502                        output.extend_from_slice(&variable_width_data.offsets);
503
504                        // store bytes
505                        output.extend_from_slice(&variable_width_data.data);
506                        Ok(LanceBuffer::from(output))
507                    }
508                    _ => {
509                        panic!(
510                            "BinaryBlockEncoder does not work with {} bits per offset VariableWidth DataBlock.",
511                            variable_width_data.bits_per_offset
512                        );
513                    }
514                }
515            }
516            _ => {
517                panic!("BinaryBlockEncoder can only work with Variable Width DataBlock.");
518            }
519        }
520    }
521}
522
523impl PerValueCompressor for VariableEncoder {
524    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
525        let DataBlock::VariableWidth(variable) = data else {
526            panic!("BinaryPerValueCompressor can only work with Variable Width DataBlock.");
527        };
528
529        let encoding = ProtobufUtils21::variable(
530            ProtobufUtils21::flat(variable.bits_per_offset as u64, None),
531            None,
532        );
533        Ok((PerValueDataBlock::Variable(variable), encoding))
534    }
535}
536
537#[derive(Debug, Default)]
538pub struct VariableDecoder {}
539
540impl VariablePerValueDecompressor for VariableDecoder {
541    fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
542        Ok(DataBlock::VariableWidth(data))
543    }
544}
545
546#[derive(Debug, Default)]
547pub struct BinaryBlockDecompressor {}
548
549impl BlockDecompressor for BinaryBlockDecompressor {
550    fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result<DataBlock> {
551        // In older (not quite stable) versions we stored the bits per offset as a single byte and then the num_values
552        // as four bytes.  However, this led to alignment problems and was wasteful since we already store the num_values
553        // in higher layers.
554        //
555        // In the standard scheme we use 4 bytes for the bits per offset and 4 bytes for the bytes_start_offset and we
556        // rely on the passed in num_values to be correct.
557
558        // This isn't perfect but it's probably good enough and the best I think we can do.  The bits per offset will
559        // never be more than 255 and it's little endian so the last 3 bytes will always be 0.  These will be the least
560        // significant 3 bytes of the number of values in the old scheme.  It's pretty unlikely these are all 0 (that would
561        // mean there are at least 16M values in a single page) so we'll use this to determine if the old scheme is used.
562        //
563        // The header fields and the offsets themselves come straight from the file.
564        // The structural checks below (all O(1)) reject blocks whose regions do not
565        // line up; the offset *values* are validated later, by the mandatory layout
566        // validation in `VariableWidthBlock::into_arrow`, so they are not rescanned
567        // here.
568        if data.len() < 4 {
569            return Err(Error::corrupt_file_named(
570                "variable-width block",
571                format!(
572                    "block of {} bytes is too small to hold a header",
573                    data.len()
574                ),
575            ));
576        }
577        let is_old_scheme = data[1] != 0 || data[2] != 0 || data[3] != 0;
578
579        let ensure_header = |header_len: usize| {
580            if data.len() < header_len {
581                return Err(Error::corrupt_file_named(
582                    "variable-width block",
583                    format!(
584                        "block of {} bytes is too small for a {} byte header",
585                        data.len(),
586                        header_len
587                    ),
588                ));
589            }
590            Ok(())
591        };
592        let (bits_per_offset, bytes_start_offset, offset_start) = if is_old_scheme {
593            // Old scheme
594            let bits_per_offset = data[0];
595            match bits_per_offset {
596                32 => {
597                    ensure_header(9)?;
598                    debug_assert_eq!(LittleEndian::read_u32(&data[1..5]), num_values as u32);
599                    let bytes_start_offset = LittleEndian::read_u32(&data[5..9]);
600                    (bits_per_offset, bytes_start_offset as u64, 9_u64)
601                }
602                64 => {
603                    ensure_header(17)?;
604                    debug_assert_eq!(LittleEndian::read_u64(&data[1..9]), num_values);
605                    let bytes_start_offset = LittleEndian::read_u64(&data[9..17]);
606                    (bits_per_offset, bytes_start_offset, 17)
607                }
608                _ => {
609                    return Err(Error::invalid_input_source(
610                        format!("Unsupported bits_per_offset={}", bits_per_offset).into(),
611                    ));
612                }
613            }
614        } else {
615            // Standard scheme
616            let bits_per_offset = LittleEndian::read_u32(&data[0..4]) as u8;
617            match bits_per_offset {
618                32 => {
619                    ensure_header(8)?;
620                    let bytes_start_offset = LittleEndian::read_u32(&data[4..8]);
621                    (bits_per_offset, bytes_start_offset as u64, 8)
622                }
623                64 => {
624                    ensure_header(16)?;
625                    let bytes_start_offset = LittleEndian::read_u64(&data[8..16]);
626                    (bits_per_offset, bytes_start_offset, 16)
627                }
628                _ => {
629                    return Err(Error::invalid_input_source(
630                        format!("Unsupported bits_per_offset={}", bits_per_offset).into(),
631                    ));
632                }
633            }
634        };
635
636        // The offsets region sits between the header and `bytes_start_offset`
637        // and must hold exactly `num_values + 1` offsets starting at zero.
638        let expected_offsets_bytes = num_values
639            .checked_add(1)
640            .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8))
641            .ok_or_else(|| {
642                Error::corrupt_file_named(
643                    "variable-width block",
644                    format!("offsets region size overflows for {num_values} values"),
645                )
646            })?;
647        if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 {
648            return Err(Error::corrupt_file_named(
649                "variable-width block",
650                format!(
651                    "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)",
652                    bytes_start_offset,
653                    offset_start,
654                    data.len()
655                ),
656            ));
657        }
658        if bytes_start_offset - offset_start != expected_offsets_bytes {
659            return Err(Error::corrupt_file_named(
660                "variable-width block",
661                format!(
662                    "expected {} offset bytes for {} values but found {}",
663                    expected_offsets_bytes,
664                    num_values,
665                    bytes_start_offset - offset_start
666                ),
667            ));
668        }
669
670        // the next `bytes_start_offset - offset_start` stores the offsets.
671        let offsets = data.slice_with_length(
672            offset_start as usize,
673            (bytes_start_offset - offset_start) as usize,
674        );
675        let first_offset = match bits_per_offset {
676            32 => LittleEndian::read_u32(&offsets[0..4]) as u64,
677            _ => LittleEndian::read_u64(&offsets[0..8]),
678        };
679        if first_offset != 0 {
680            return Err(Error::corrupt_file_named(
681                "variable-width block",
682                format!("first offset must be 0 but found {first_offset}"),
683            ));
684        }
685
686        // the rest are the binary bytes.
687        let data = data.slice_with_length(
688            bytes_start_offset as usize,
689            data.len() - bytes_start_offset as usize,
690        );
691
692        Ok(DataBlock::VariableWidth(VariableWidthBlock {
693            data,
694            offsets,
695            bits_per_offset,
696            num_values,
697            block_info: BlockInfo::new(),
698        }))
699    }
700}
701
702#[cfg(test)]
703mod tests {
704    use arrow_array::{
705        ArrayRef, StringArray,
706        builder::{LargeStringBuilder, StringBuilder},
707    };
708    use arrow_schema::{DataType, Field};
709
710    use crate::{
711        buffer::LanceBuffer,
712        constants::{
713            COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY,
714            STRUCTURAL_ENCODING_MINIBLOCK,
715        },
716        data::{BlockInfo, DataBlock, VariableWidthBlock},
717        testing::check_specific_random,
718    };
719    use rstest::rstest;
720    use std::{collections::HashMap, sync::Arc, vec};
721
722    use crate::testing::{
723        FnArrayGeneratorProvider, TestCases, check_basic_random, check_round_trip_encoding_of_data,
724    };
725
726    #[test_log::test(tokio::test)]
727    async fn test_utf8_binary() {
728        let field = Field::new("", DataType::Utf8, false);
729        check_specific_random(field, TestCases::basic().with_structural_encodings()).await;
730    }
731
732    #[rstest]
733    #[test_log::test(tokio::test)]
734    async fn test_binary(
735        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
736        structural_encoding: &str,
737        #[values(DataType::Utf8, DataType::Binary)] data_type: DataType,
738    ) {
739        let mut field_metadata = HashMap::new();
740        field_metadata.insert(
741            STRUCTURAL_ENCODING_META_KEY.to_string(),
742            structural_encoding.into(),
743        );
744
745        let field = Field::new("", data_type, false).with_metadata(field_metadata);
746        check_basic_random(field).await;
747    }
748
749    #[rstest]
750    #[test_log::test(tokio::test)]
751    async fn test_binary_fsst(
752        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
753        structural_encoding: &str,
754        #[values(DataType::Binary, DataType::Utf8)] data_type: DataType,
755    ) {
756        let mut field_metadata = HashMap::new();
757        field_metadata.insert(
758            STRUCTURAL_ENCODING_META_KEY.to_string(),
759            structural_encoding.into(),
760        );
761        field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into());
762        let field = Field::new("", data_type, true).with_metadata(field_metadata);
763        // TODO (https://github.com/lance-format/lance/issues/4783)
764        let test_cases = TestCases::default().with_structural_encodings();
765        check_specific_random(field, test_cases).await;
766    }
767
768    #[rstest]
769    #[test_log::test(tokio::test)]
770    async fn test_fsst_large_binary(
771        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
772        structural_encoding: &str,
773        #[values(DataType::LargeBinary, DataType::LargeUtf8)] data_type: DataType,
774    ) {
775        let mut field_metadata = HashMap::new();
776        field_metadata.insert(
777            STRUCTURAL_ENCODING_META_KEY.to_string(),
778            structural_encoding.into(),
779        );
780        field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into());
781        let field = Field::new("", data_type, true).with_metadata(field_metadata);
782        check_specific_random(field, TestCases::basic().with_structural_encodings()).await;
783    }
784
785    #[test_log::test(tokio::test)]
786    async fn test_large_binary() {
787        let field = Field::new("", DataType::LargeBinary, true);
788        check_basic_random(field).await;
789    }
790
791    #[test_log::test(tokio::test)]
792    async fn test_large_utf8() {
793        let field = Field::new("", DataType::LargeUtf8, true);
794        check_basic_random(field).await;
795    }
796
797    #[rstest]
798    #[test_log::test(tokio::test)]
799    async fn test_small_strings(
800        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
801        structural_encoding: &str,
802    ) {
803        use crate::testing::check_basic_generated;
804
805        let mut field_metadata = HashMap::new();
806        field_metadata.insert(
807            STRUCTURAL_ENCODING_META_KEY.to_string(),
808            structural_encoding.into(),
809        );
810        let field = Field::new("", DataType::Utf8, true).with_metadata(field_metadata);
811        check_basic_generated(
812            field,
813            Box::new(FnArrayGeneratorProvider::new(move || {
814                lance_datagen::array::utf8_prefix_plus_counter("user_", /*is_large=*/ false)
815            })),
816        )
817        .await;
818    }
819
820    #[rstest]
821    #[test_log::test(tokio::test)]
822    async fn test_simple_binary(
823        #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
824        structural_encoding: &str,
825        #[values(DataType::Utf8, DataType::Binary)] data_type: DataType,
826    ) {
827        let string_array = StringArray::from(vec![Some("abc"), None, Some("pqr"), None, Some("m")]);
828        let string_array = arrow_cast::cast(&string_array, &data_type).unwrap();
829
830        let mut field_metadata = HashMap::new();
831        field_metadata.insert(
832            STRUCTURAL_ENCODING_META_KEY.to_string(),
833            structural_encoding.into(),
834        );
835
836        let test_cases = TestCases::default()
837            .with_range(0..2)
838            .with_range(0..3)
839            .with_range(1..3)
840            .with_indices(vec![0, 1, 3, 4]);
841        check_round_trip_encoding_of_data(
842            vec![Arc::new(string_array)],
843            &test_cases,
844            field_metadata,
845        )
846        .await;
847    }
848
849    #[test_log::test(tokio::test)]
850    async fn test_sliced_utf8() {
851        let string_array = StringArray::from(vec![Some("abc"), Some("de"), None, Some("fgh")]);
852        let string_array = string_array.slice(1, 3);
853
854        let test_cases = TestCases::default()
855            .with_range(0..1)
856            .with_range(0..2)
857            .with_range(1..2);
858        check_round_trip_encoding_of_data(
859            vec![Arc::new(string_array)],
860            &test_cases,
861            HashMap::new(),
862        )
863        .await;
864    }
865
866    #[test_log::test(tokio::test)]
867    async fn test_bigger_than_max_page_size() {
868        // Create an array with one single 32MiB string
869        let big_string = String::from_iter((0..(32 * 1024 * 1024)).map(|_| '0'));
870        let string_array = StringArray::from(vec![
871            Some(big_string),
872            Some("abc".to_string()),
873            None,
874            None,
875            Some("xyz".to_string()),
876        ]);
877
878        // Drop the max page size to 1MiB
879        let test_cases = TestCases::default().with_max_page_size(1024 * 1024);
880
881        check_round_trip_encoding_of_data(
882            vec![Arc::new(string_array)],
883            &test_cases,
884            HashMap::new(),
885        )
886        .await;
887
888        // This is a regression testing the case where a page with X rows is split into Y parts
889        // where the number of parts is not evenly divisible by the number of rows.  In this
890        // case we are splitting 90 rows into 4 parts.
891        let big_string = String::from_iter((0..(1000 * 1000)).map(|_| '0'));
892        let string_array = StringArray::from_iter_values((0..90).map(|_| big_string.clone()));
893
894        check_round_trip_encoding_of_data(
895            vec![Arc::new(string_array)],
896            &TestCases::default(),
897            HashMap::new(),
898        )
899        .await;
900    }
901
902    #[test_log::test(tokio::test)]
903    async fn test_empty_strings() {
904        // Scenario 1: Some strings are empty
905
906        let values = [Some("abc"), Some(""), None];
907        // Test empty list at beginning, middle, and end
908        for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] {
909            let mut string_builder = StringBuilder::new();
910            for idx in order {
911                string_builder.append_option(values[idx]);
912            }
913            let string_array = Arc::new(string_builder.finish());
914            let test_cases = TestCases::default()
915                .with_indices(vec![1])
916                .with_indices(vec![0])
917                .with_indices(vec![2])
918                .with_indices(vec![0, 1]);
919            check_round_trip_encoding_of_data(
920                vec![string_array.clone()],
921                &test_cases,
922                HashMap::new(),
923            )
924            .await;
925            let test_cases = test_cases.with_batch_size(1);
926            check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new())
927                .await;
928        }
929
930        // Scenario 2: All strings are empty
931
932        // When encoding an array of empty strings there are no bytes to encode
933        // which is strange and we want to ensure we handle it
934        let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")]));
935
936        let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
937        check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new())
938            .await;
939        let test_cases = test_cases.with_batch_size(1);
940        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
941    }
942
943    #[test_log::test(tokio::test)]
944    #[ignore] // This test is quite slow in debug mode
945    async fn test_jumbo_string() {
946        // This is an overflow test.  We have a list of lists where each list
947        // has 1Mi items.  We encode 5000 of these lists and so we have over 4Gi in the
948        // offsets range
949        let mut string_builder = LargeStringBuilder::new();
950        // a 1 MiB string
951        let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0'));
952        for _ in 0..5000 {
953            string_builder.append_option(Some(&giant_string));
954        }
955        let giant_array = Arc::new(string_builder.finish()) as ArrayRef;
956        let arrs = vec![giant_array];
957
958        // // We can't validate because our validation relies on concatenating all input arrays
959        let test_cases = TestCases::default().without_validation();
960        check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await;
961    }
962
963    #[rstest]
964    #[test_log::test(tokio::test)]
965    async fn test_binary_dictionary_encoding(
966        #[values(true, false)] with_nulls: bool,
967        #[values(100, 500, 35000)] dict_size: u32,
968    ) {
969        let test_cases = TestCases::default().with_structural_encodings();
970        let strings = (0..dict_size)
971            .map(|i| i.to_string())
972            .collect::<Vec<String>>();
973
974        let repeated_strings: Vec<_> = strings
975            .iter()
976            .cycle()
977            .take(70000)
978            .enumerate()
979            .map(|(i, s)| {
980                if with_nulls && i % 7 == 0 {
981                    None
982                } else {
983                    Some(s.clone())
984                }
985            })
986            .collect();
987        let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
988        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
989    }
990
991    #[test_log::test(tokio::test)]
992    async fn test_binary_encoding_verification() {
993        use lance_datagen::{ByteCount, RowCount};
994
995        let test_cases = TestCases::default()
996            .with_expected_encoding("variable")
997            .with_structural_encodings();
998
999        // Test both automatic selection and explicit configuration
1000        // 1. Test automatic binary encoding selection (small strings that won't trigger FSST)
1001        let arr_small = lance_datagen::gen_batch()
1002            .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(10), false))
1003            .into_batch_rows(RowCount::from(1000))
1004            .unwrap()
1005            .column(0)
1006            .clone();
1007        check_round_trip_encoding_of_data(vec![arr_small], &test_cases, HashMap::new()).await;
1008
1009        // 2. Test explicit "none" compression to force binary encoding
1010        let metadata_explicit =
1011            HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]);
1012        let arr_large = lance_datagen::gen_batch()
1013            .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(50), false))
1014            .into_batch_rows(RowCount::from(2000))
1015            .unwrap()
1016            .column(0)
1017            .clone();
1018        check_round_trip_encoding_of_data(vec![arr_large], &test_cases, metadata_explicit).await;
1019    }
1020
1021    #[test]
1022    fn test_binary_miniblock_with_misaligned_buffer() {
1023        use super::BinaryMiniBlockDecompressor;
1024        use crate::buffer::LanceBuffer;
1025        use crate::compression::MiniBlockDecompressor;
1026        use crate::data::DataBlock;
1027
1028        // Test case 1: u32 offsets
1029        {
1030            let decompressor = BinaryMiniBlockDecompressor {
1031                bits_per_offset: 32,
1032            };
1033
1034            // Create test data with u32 offsets
1035            // BinaryMiniBlock format: all offsets followed by all string data
1036            // Need to ensure total size is divisible by 4 for u32
1037            let mut test_data = Vec::new();
1038
1039            // Offsets section (3 offsets for 2 values + 1 end offset)
1040            test_data.extend_from_slice(&12u32.to_le_bytes()); // offset to start of strings (after offsets)
1041            test_data.extend_from_slice(&15u32.to_le_bytes()); // offset to second string
1042            test_data.extend_from_slice(&20u32.to_le_bytes()); // offset to end
1043
1044            // String data section
1045            test_data.extend_from_slice(b"ABCXYZ"); // 6 bytes of string data
1046            test_data.extend_from_slice(&[0, 0]); // 2 bytes padding to make total 20 bytes (divisible by 4)
1047
1048            // Create a misaligned buffer by adding padding and slicing
1049            let mut padded = Vec::with_capacity(test_data.len() + 1);
1050            padded.push(0xFF); // Padding byte to misalign
1051            padded.extend_from_slice(&test_data);
1052
1053            let bytes = bytes::Bytes::from(padded);
1054            let misaligned = bytes.slice(1..); // Skip first byte to create misalignment
1055
1056            // Create LanceBuffer with bytes_per_value=1 to bypass alignment check
1057            let buffer = LanceBuffer::from_bytes(misaligned, 1);
1058
1059            // Verify the buffer is actually misaligned
1060            let ptr = buffer.as_ref().as_ptr();
1061            assert_ne!(
1062                ptr.align_offset(4),
1063                0,
1064                "Test setup: buffer should be misaligned for u32"
1065            );
1066
1067            // Decompress with misaligned buffer - should work with borrow_to_typed_slice
1068            let result = decompressor.decompress(vec![buffer], 2);
1069            assert!(
1070                result.is_ok(),
1071                "Decompression should succeed with misaligned buffer"
1072            );
1073
1074            // Verify the data is correct
1075            if let Ok(DataBlock::VariableWidth(block)) = result {
1076                assert_eq!(block.num_values, 2);
1077                // Data should be the strings (including padding from the original buffer)
1078                assert_eq!(&block.data.as_ref()[..6], b"ABCXYZ");
1079            } else {
1080                panic!("Expected VariableWidth block");
1081            }
1082        }
1083
1084        // Test case 2: u64 offsets
1085        {
1086            let decompressor = BinaryMiniBlockDecompressor {
1087                bits_per_offset: 64,
1088            };
1089
1090            // Create test data with u64 offsets
1091            let mut test_data = Vec::new();
1092
1093            // Offsets section (3 offsets for 2 values + 1 end offset)
1094            test_data.extend_from_slice(&24u64.to_le_bytes()); // offset to start of strings (after offsets)
1095            test_data.extend_from_slice(&29u64.to_le_bytes()); // offset to second string
1096            test_data.extend_from_slice(&40u64.to_le_bytes()); // offset to end (divisible by 8)
1097
1098            // String data section
1099            test_data.extend_from_slice(b"HelloWorld"); // 10 bytes of string data
1100            test_data.extend_from_slice(&[0, 0, 0, 0, 0, 0]); // 6 bytes padding to make total 40 bytes (divisible by 8)
1101
1102            // Create misaligned buffer
1103            let mut padded = Vec::with_capacity(test_data.len() + 3);
1104            padded.extend_from_slice(&[0xFF, 0xFF, 0xFF]); // 3 bytes padding for misalignment
1105            padded.extend_from_slice(&test_data);
1106
1107            let bytes = bytes::Bytes::from(padded);
1108            let misaligned = bytes.slice(3..); // Skip 3 bytes
1109
1110            let buffer = LanceBuffer::from_bytes(misaligned, 1);
1111
1112            // Verify misalignment for u64
1113            let ptr = buffer.as_ref().as_ptr();
1114            assert_ne!(
1115                ptr.align_offset(8),
1116                0,
1117                "Test setup: buffer should be misaligned for u64"
1118            );
1119
1120            // Decompress should succeed
1121            let result = decompressor.decompress(vec![buffer], 2);
1122            assert!(
1123                result.is_ok(),
1124                "Decompression should succeed with misaligned u64 buffer"
1125            );
1126
1127            if let Ok(DataBlock::VariableWidth(block)) = result {
1128                assert_eq!(block.num_values, 2);
1129                // Data should be the strings (including padding from the original buffer)
1130                assert_eq!(&block.data.as_ref()[..10], b"HelloWorld");
1131            } else {
1132                panic!("Expected VariableWidth block");
1133            }
1134        }
1135    }
1136
1137    #[test]
1138    fn test_binary_miniblock_rejects_corrupt_offsets() {
1139        use super::BinaryMiniBlockDecompressor;
1140        use crate::compression::MiniBlockDecompressor;
1141        use lance_core::Error;
1142
1143        // Chunk layout mirrors the on-disk format for ["alpha", "beta", "gamma"]:
1144        // LE u32 offsets [16, 21, 25, 30] followed by the value bytes, padded to
1145        // a multiple of 8 bytes.
1146        fn chunk_u32(offsets: &[u32], values: &[u8]) -> LanceBuffer {
1147            let mut chunk = offsets
1148                .iter()
1149                .flat_map(|offset| offset.to_le_bytes())
1150                .collect::<Vec<u8>>();
1151            chunk.extend_from_slice(values);
1152            chunk.resize(chunk.len().next_multiple_of(8), 0);
1153            LanceBuffer::from(chunk)
1154        }
1155
1156        let decompressor = BinaryMiniBlockDecompressor::new(32);
1157
1158        // The tail offset points past the end of the 32-byte chunk.
1159        let err = decompressor
1160            .decompress(
1161                vec![chunk_u32(&[16, 21, 25, 100_000], b"alphabetagamma")],
1162                3,
1163            )
1164            .unwrap_err();
1165        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1166        assert!(err.to_string().contains("out of bounds"), "{err}");
1167
1168        // Offsets go backwards, which would underflow the rebase subtraction.
1169        let err = decompressor
1170            .decompress(vec![chunk_u32(&[16, 25, 21, 30], b"alphabetagamma")], 3)
1171            .unwrap_err();
1172        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1173        assert!(err.to_string().contains("decreases"), "{err}");
1174
1175        // The first offset points inside the offset table, which would alias
1176        // the serialized offsets into the value bytes.
1177        let err = decompressor
1178            .decompress(vec![chunk_u32(&[0, 21, 25, 30], b"alphabetagamma")], 3)
1179            .unwrap_err();
1180        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1181        assert!(err.to_string().contains("overlaps"), "{err}");
1182
1183        // The chunk stores fewer offsets than the requested value count needs.
1184        let err = decompressor
1185            .decompress(vec![chunk_u32(&[8, 8], &[])], 3)
1186            .unwrap_err();
1187        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1188        assert!(err.to_string().contains("requires 4"), "{err}");
1189
1190        // The chunk size is not a multiple of the offset width.
1191        let err = decompressor
1192            .decompress(vec![LanceBuffer::from(vec![0u8; 10])], 1)
1193            .unwrap_err();
1194        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1195        assert!(err.to_string().contains("multiple"), "{err}");
1196
1197        // 64-bit offsets take the same validation path.
1198        fn chunk_u64(offsets: &[u64], values: &[u8]) -> LanceBuffer {
1199            let mut chunk = offsets
1200                .iter()
1201                .flat_map(|offset| offset.to_le_bytes())
1202                .collect::<Vec<u8>>();
1203            chunk.extend_from_slice(values);
1204            chunk.resize(chunk.len().next_multiple_of(8), 0);
1205            LanceBuffer::from(chunk)
1206        }
1207        let decompressor = BinaryMiniBlockDecompressor::new(64);
1208        let err = decompressor
1209            .decompress(
1210                vec![chunk_u64(&[32, 37, 41, 100_000], b"alphabetagamma")],
1211                3,
1212            )
1213            .unwrap_err();
1214        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1215        assert!(err.to_string().contains("out of bounds"), "{err}");
1216        let err = decompressor
1217            .decompress(vec![chunk_u64(&[0, 37, 41, 46], b"alphabetagamma")], 3)
1218            .unwrap_err();
1219        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1220        assert!(err.to_string().contains("overlaps"), "{err}");
1221
1222        // A valid chunk still decodes: offsets rebase to [0, 5, 9, 14].
1223        let decompressor = BinaryMiniBlockDecompressor::new(32);
1224        let block = decompressor
1225            .decompress(vec![chunk_u32(&[16, 21, 25, 30], b"alphabetagamma")], 3)
1226            .unwrap();
1227        let DataBlock::VariableWidth(block) = block else {
1228            panic!("expected a variable-width block");
1229        };
1230        assert_eq!(block.data.as_ref(), b"alphabetagamma");
1231        assert_eq!(
1232            block.offsets,
1233            LanceBuffer::reinterpret_vec(vec![0_u32, 5, 9, 14])
1234        );
1235    }
1236
1237    fn encoded_binary_block(bits_per_offset: u8) -> Vec<u8> {
1238        use crate::compression::BlockCompressor;
1239
1240        let offsets = match bits_per_offset {
1241            32 => LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 14]),
1242            64 => LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 14]),
1243            _ => unreachable!(),
1244        };
1245        let block = DataBlock::VariableWidth(VariableWidthBlock {
1246            data: LanceBuffer::copy_slice(b"alphabetagamma"),
1247            offsets,
1248            bits_per_offset,
1249            num_values: 3,
1250            block_info: BlockInfo::new(),
1251        });
1252        BlockCompressor::compress(&super::VariableEncoder::default(), block)
1253            .unwrap()
1254            .as_ref()
1255            .to_vec()
1256    }
1257
1258    /// The block decompressor only checks the block structure (all O(1)); bad
1259    /// offset values inside a structurally-sound block are rejected by the
1260    /// mandatory layout validation when the block is converted to Arrow.
1261    #[rstest]
1262    #[case::i32_tail_out_of_bounds(32, 3, 100_000, "out of bounds")]
1263    #[case::i64_tail_out_of_bounds(64, 3, 15, "out of bounds")]
1264    #[case::i32_non_monotonic(32, 2, 4, "non-monotonic")]
1265    #[case::i64_non_monotonic(64, 2, 4, "non-monotonic")]
1266    fn test_binary_block_bad_offsets_rejected_at_arrow_conversion(
1267        #[case] bits_per_offset: u8,
1268        #[case] mutated_offset_index: usize,
1269        #[case] mutated_offset_value: u64,
1270        #[case] expected_message: &str,
1271    ) {
1272        use crate::compression::BlockDecompressor;
1273        use lance_core::Error;
1274
1275        let mut encoded = encoded_binary_block(bits_per_offset);
1276        let bytes_per_offset = (bits_per_offset / 8) as usize;
1277        // The standard scheme header is two offset-width fields.
1278        let mutated_offset_start = bytes_per_offset * (2 + mutated_offset_index);
1279        encoded[mutated_offset_start..mutated_offset_start + bytes_per_offset]
1280            .copy_from_slice(&mutated_offset_value.to_le_bytes()[..bytes_per_offset]);
1281
1282        let block = super::BinaryBlockDecompressor::default()
1283            .decompress(LanceBuffer::from(encoded), 3)
1284            .unwrap();
1285        let data_type = match bits_per_offset {
1286            32 => DataType::Binary,
1287            _ => DataType::LargeBinary,
1288        };
1289        let err = block.into_arrow(data_type, false).unwrap_err();
1290        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1291        assert!(err.to_string().contains(expected_message), "{err}");
1292    }
1293
1294    #[test]
1295    fn test_binary_block_rejects_corrupt_structure() {
1296        use crate::compression::BlockDecompressor;
1297        use lance_core::Error;
1298
1299        let decompressor = super::BinaryBlockDecompressor::default();
1300
1301        // The first offset must be zero.
1302        let mut encoded = encoded_binary_block(32);
1303        encoded[8..12].copy_from_slice(&5_u32.to_le_bytes());
1304        let err = decompressor
1305            .decompress(LanceBuffer::from(encoded), 3)
1306            .unwrap_err();
1307        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1308        assert!(err.to_string().contains("first offset"), "{err}");
1309
1310        // The offsets region must hold exactly num_values + 1 offsets.
1311        let encoded = encoded_binary_block(32);
1312        let err = decompressor
1313            .decompress(LanceBuffer::from(encoded), 4)
1314            .unwrap_err();
1315        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1316        assert!(err.to_string().contains("offset bytes"), "{err}");
1317
1318        // A block too small to hold its header is rejected, not a panic.
1319        let err = decompressor
1320            .decompress(LanceBuffer::from(vec![0_u8; 2]), 1)
1321            .unwrap_err();
1322        assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1323        assert!(err.to_string().contains("too small"), "{err}");
1324    }
1325}