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