1use arrow_array::OffsetSizeTrait;
13use byteorder::{ByteOrder, LittleEndian};
14use core::panic;
15
16use crate::compression::{
17 BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor,
18 require_block_payload,
19};
20
21use crate::buffer::LanceBuffer;
22use crate::data::{BlockInfo, DataBlock, VariableWidthBlock};
23use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock};
24use crate::encodings::logical::primitive::miniblock::{
25 MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext,
26 MiniBlockCompressor,
27};
28use crate::format::pb21::CompressiveEncoding;
29use crate::format::pb21::compressive_encoding::Compression;
30use crate::format::{ProtobufUtils21, pb21};
31
32use lance_core::utils::bit::pad_bytes_to;
33use lance_core::{Error, Result};
34
35#[derive(Debug)]
36pub struct BinaryMiniBlockEncoder {
37 minichunk_size: i64,
38}
39
40impl Default for BinaryMiniBlockEncoder {
41 fn default() -> Self {
42 Self {
43 minichunk_size: *AIM_MINICHUNK_SIZE,
44 }
45 }
46}
47
48const DEFAULT_AIM_MINICHUNK_SIZE: i64 = 4 * 1024;
49
50pub static AIM_MINICHUNK_SIZE: std::sync::LazyLock<i64> = std::sync::LazyLock::new(|| {
51 std::env::var("LANCE_BINARY_MINIBLOCK_CHUNK_SIZE")
52 .unwrap_or_else(|_| DEFAULT_AIM_MINICHUNK_SIZE.to_string())
53 .parse::<i64>()
54 .unwrap_or(DEFAULT_AIM_MINICHUNK_SIZE)
55});
56
57fn chunk_offsets<N: OffsetSizeTrait>(
59 offsets: &[N],
60 data: &[u8],
61 alignment: usize,
62 minichunk_size: i64,
63) -> (Vec<LanceBuffer>, Vec<MiniBlockChunk>) {
64 #[derive(Debug)]
65 struct ChunkInfo {
66 chunk_start_offset_in_orig_idx: usize,
67 chunk_last_offset_in_orig_idx: usize,
68 bytes_start_offset: usize,
70 padded_chunk_size: usize,
75 }
76
77 let byte_width: usize = N::get_byte_width();
78 let mut chunks_info = vec![];
79 let mut chunks = vec![];
80 let mut last_offset_in_orig_idx = 0;
81 loop {
82 let this_last_offset_in_orig_idx =
83 search_next_offset_idx(offsets, last_offset_in_orig_idx, minichunk_size);
84
85 let num_values_in_this_chunk = this_last_offset_in_orig_idx - last_offset_in_orig_idx;
86 let chunk_bytes = offsets[this_last_offset_in_orig_idx] - offsets[last_offset_in_orig_idx];
87 let this_chunk_size =
88 (num_values_in_this_chunk + 1) * byte_width + chunk_bytes.to_usize().unwrap();
89
90 let padded_chunk_size = this_chunk_size.next_multiple_of(alignment);
91 debug_assert!(padded_chunk_size > 0);
92
93 let this_chunk_bytes_start_offset = (num_values_in_this_chunk + 1) * byte_width;
94 chunks_info.push(ChunkInfo {
95 chunk_start_offset_in_orig_idx: last_offset_in_orig_idx,
96 chunk_last_offset_in_orig_idx: this_last_offset_in_orig_idx,
97 bytes_start_offset: this_chunk_bytes_start_offset,
98 padded_chunk_size,
99 });
100 chunks.push(MiniBlockChunk {
101 log_num_values: if this_last_offset_in_orig_idx == offsets.len() - 1 {
102 0
103 } else {
104 num_values_in_this_chunk.trailing_zeros() as u8
105 },
106 buffer_sizes: vec![padded_chunk_size as u32],
107 });
108 if this_last_offset_in_orig_idx == offsets.len() - 1 {
109 break;
110 }
111 last_offset_in_orig_idx = this_last_offset_in_orig_idx;
112 }
113
114 let output_total_bytes = chunks_info
115 .iter()
116 .map(|chunk_info| chunk_info.padded_chunk_size)
117 .sum::<usize>();
118
119 let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
120
121 for chunk in chunks_info {
122 let this_chunk_offsets: Vec<N> = offsets
123 [chunk.chunk_start_offset_in_orig_idx..=chunk.chunk_last_offset_in_orig_idx]
124 .iter()
125 .map(|offset| {
126 *offset - offsets[chunk.chunk_start_offset_in_orig_idx]
127 + N::from_usize(chunk.bytes_start_offset).unwrap()
128 })
129 .collect();
130
131 let this_chunk_offsets = LanceBuffer::reinterpret_vec(this_chunk_offsets);
132 output.extend_from_slice(&this_chunk_offsets);
133
134 let start_in_orig = offsets[chunk.chunk_start_offset_in_orig_idx]
135 .to_usize()
136 .unwrap();
137 let end_in_orig = offsets[chunk.chunk_last_offset_in_orig_idx]
138 .to_usize()
139 .unwrap();
140 output.extend_from_slice(&data[start_in_orig..end_in_orig]);
141
142 const PAD_BYTE: u8 = 72;
144 let pad_len = pad_bytes_to(output.len(), alignment);
145
146 if pad_len > 0_usize {
148 output.extend(std::iter::repeat_n(PAD_BYTE, pad_len));
149 }
150 }
151 (vec![LanceBuffer::reinterpret_vec(output)], chunks)
152}
153
154fn search_next_offset_idx<N: OffsetSizeTrait>(
159 offsets: &[N],
160 last_offset_idx: usize,
161 minichunk_size: i64,
162) -> usize {
163 let remaining_values = offsets.len().saturating_sub(last_offset_idx + 1);
167 if remaining_values <= 1 {
168 return offsets.len() - 1;
169 }
170
171 let mut num_values = 2;
172 let mut new_num_values = num_values * 2;
173 loop {
174 if last_offset_idx + new_num_values >= offsets.len() {
175 let existing_bytes = offsets[offsets.len() - 1] - offsets[last_offset_idx];
176 let new_size = existing_bytes
178 + N::from_usize((offsets.len() - last_offset_idx) * N::get_byte_width()).unwrap();
179 if new_size.to_i64().unwrap() <= minichunk_size {
180 return offsets.len() - 1;
182 } else {
183 return last_offset_idx + num_values;
185 }
186 }
187 let existing_bytes = offsets[last_offset_idx + new_num_values] - offsets[last_offset_idx];
188 let new_size =
189 existing_bytes + N::from_usize((new_num_values + 1) * N::get_byte_width()).unwrap();
190 if new_size.to_i64().unwrap() <= minichunk_size {
191 if new_num_values * 2 > *MAX_MINIBLOCK_VALUES as usize {
192 break;
194 }
195 num_values = new_num_values;
196 new_num_values *= 2;
197 } else {
198 break;
199 }
200 }
201 last_offset_idx + num_values
202}
203
204impl BinaryMiniBlockEncoder {
205 pub fn new(minichunk_size: Option<i64>) -> Self {
206 Self {
207 minichunk_size: minichunk_size.unwrap_or(*AIM_MINICHUNK_SIZE),
208 }
209 }
210
211 fn chunk_data(&self, data: VariableWidthBlock) -> (MiniBlockCompressed, CompressiveEncoding) {
215 match data.bits_per_offset {
218 32 => {
219 let offsets = data.offsets.borrow_to_typed_slice::<i32>();
220 let (buffers, chunks) =
221 chunk_offsets(offsets.as_ref(), &data.data, 4, self.minichunk_size);
222 (
223 MiniBlockCompressed {
224 data: buffers,
225 chunks,
226 num_values: data.num_values,
227 },
228 ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None),
229 )
230 }
231 64 => {
232 let offsets = data.offsets.borrow_to_typed_slice::<i64>();
233 let (buffers, chunks) =
234 chunk_offsets(offsets.as_ref(), &data.data, 8, self.minichunk_size);
235 (
236 MiniBlockCompressed {
237 data: buffers,
238 chunks,
239 num_values: data.num_values,
240 },
241 ProtobufUtils21::variable(ProtobufUtils21::flat(64, None), None),
242 )
243 }
244 _ => panic!("Unsupported bits_per_offset={}", data.bits_per_offset),
245 }
246 }
247}
248
249impl MiniBlockCompressor for BinaryMiniBlockEncoder {
250 fn compress(
251 &self,
252 _context: MiniBlockCompressionContext,
253 data: DataBlock,
254 ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
255 match data {
256 DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)),
257 _ => Err(Error::invalid_input_source(
258 format!(
259 "Cannot compress a data block of type {} with BinaryMiniBlockEncoder",
260 data.name()
261 )
262 .into(),
263 )),
264 }
265 }
266}
267
268#[derive(Debug)]
269pub struct BinaryMiniBlockDecompressor {
270 bits_per_offset: u8,
271}
272
273impl BinaryMiniBlockDecompressor {
274 pub fn new(bits_per_offset: u8) -> Self {
275 assert!(bits_per_offset == 32 || bits_per_offset == 64);
276 Self { bits_per_offset }
277 }
278
279 pub fn from_variable(variable: &pb21::Variable) -> Self {
280 if let Compression::Flat(flat) = variable
281 .offsets
282 .as_ref()
283 .unwrap()
284 .compression
285 .as_ref()
286 .unwrap()
287 {
288 Self {
289 bits_per_offset: flat.bits_per_value as u8,
290 }
291 } else {
292 panic!("Unsupported offsets compression: {:?}", variable.offsets);
293 }
294 }
295}
296
297fn chunk_offset_violation_error<T: Copy + Into<u64>>(offsets: &[T], chunk_len: usize) -> Error {
300 let mut previous: u64 = offsets[0].into();
301 for (position, &offset) in offsets.iter().enumerate().skip(1) {
302 let offset: u64 = offset.into();
303 if offset < previous {
304 return Error::corrupt_file_named(
305 "binary mini-block",
306 format!(
307 "value offset at position {position} decreases: {offset} < {previous} \
308 (chunk is {chunk_len} bytes)"
309 ),
310 );
311 }
312 previous = offset;
313 }
314 Error::corrupt_file_named(
315 "binary mini-block",
316 format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"),
317 )
318}
319
320impl MiniBlockDecompressor for BinaryMiniBlockDecompressor {
321 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
332 assert_eq!(data.len(), 1);
333 let data = data.into_iter().next().unwrap();
334
335 let bytes_per_offset = self.bits_per_offset as usize / 8;
336 if !data.len().is_multiple_of(bytes_per_offset) {
337 return Err(Error::corrupt_file_named(
338 "binary mini-block",
339 format!(
340 "chunk size {} is not a multiple of the {}-byte offset width",
341 data.len(),
342 bytes_per_offset
343 ),
344 ));
345 }
346 let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| {
347 Error::corrupt_file_named(
348 "binary mini-block",
349 format!("cannot decode {num_values} values from a single chunk"),
350 )
351 })?;
352 if data.len() / bytes_per_offset < num_offsets {
353 return Err(Error::corrupt_file_named(
354 "binary mini-block",
355 format!(
356 "chunk of {} bytes holds {} offsets but decoding {} values requires {}",
357 data.len(),
358 data.len() / bytes_per_offset,
359 num_values,
360 num_offsets
361 ),
362 ));
363 }
364
365 let min_value_region_start = num_offsets * bytes_per_offset;
370 let value_region_overlap_error = |first: u64| {
371 Error::corrupt_file_named(
372 "binary mini-block",
373 format!(
374 "value region starts at offset {first} which overlaps the {num_offsets} \
375 requested offsets ({min_value_region_start} bytes)"
376 ),
377 )
378 };
379
380 if self.bits_per_offset == 64 {
381 let offsets_buffer = data.borrow_to_typed_slice::<u64>();
382 let offsets = &offsets_buffer.as_ref()[..num_offsets];
383
384 let first = offsets[0];
385 if first < min_value_region_start as u64 {
386 return Err(value_region_overlap_error(first));
387 }
388 let mut previous = first;
389 let mut is_monotonic = true;
390 let result_offsets = offsets
391 .iter()
392 .map(|&offset| {
393 is_monotonic &= previous <= offset;
394 previous = offset;
395 offset.wrapping_sub(first)
396 })
397 .collect::<Vec<u64>>();
398 let last = offsets[num_offsets - 1];
399 if !is_monotonic || last as usize > data.len() {
400 return Err(chunk_offset_violation_error(offsets, data.len()));
401 }
402
403 Ok(DataBlock::VariableWidth(VariableWidthBlock {
404 data: LanceBuffer::from(data[first as usize..last as usize].to_vec()),
405 offsets: LanceBuffer::reinterpret_vec(result_offsets),
406 bits_per_offset: 64,
407 num_values,
408 block_info: BlockInfo::new(),
409 }))
410 } else {
411 let offsets_buffer = data.borrow_to_typed_slice::<u32>();
412 let offsets = &offsets_buffer.as_ref()[..num_offsets];
413
414 let first = offsets[0];
415 if (first as u64) < min_value_region_start as u64 {
416 return Err(value_region_overlap_error(first as u64));
417 }
418 let mut previous = first;
419 let mut is_monotonic = true;
420 let result_offsets = offsets
421 .iter()
422 .map(|&offset| {
423 is_monotonic &= previous <= offset;
424 previous = offset;
425 offset.wrapping_sub(first)
426 })
427 .collect::<Vec<u32>>();
428 let last = offsets[num_offsets - 1];
429 if !is_monotonic || last as usize > data.len() {
430 return Err(chunk_offset_violation_error(offsets, data.len()));
431 }
432
433 Ok(DataBlock::VariableWidth(VariableWidthBlock {
434 data: LanceBuffer::from(data[first as usize..last as usize].to_vec()),
435 offsets: LanceBuffer::reinterpret_vec(result_offsets),
436 bits_per_offset: 32,
437 num_values,
438 block_info: BlockInfo::new(),
439 }))
440 }
441 }
442}
443
444#[derive(Debug, Default)]
454pub struct VariableEncoder {}
455
456impl BlockCompressor for VariableEncoder {
457 fn compress(&self, mut data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
458 let bits_per_offset = match &data {
459 DataBlock::VariableWidth(data) => data.bits_per_offset,
460 _ => {
461 return Err(Error::invalid_input(
462 "BinaryBlockEncoder requires a variable-width block",
463 ));
464 }
465 };
466 match data {
467 DataBlock::VariableWidth(ref mut variable_width_data) => {
468 match variable_width_data.bits_per_offset {
469 32 => {
470 let offsets = variable_width_data.offsets.borrow_to_typed_slice::<u32>();
471 let offsets = offsets.as_ref();
472 let bytes_start_offset = 4 + 4 + std::mem::size_of_val(offsets) as u32;
475
476 let output_total_bytes =
477 bytes_start_offset as usize + variable_width_data.data.len();
478 let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
479
480 output.extend_from_slice(&(32_u32).to_le_bytes());
482
483 output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
485
486 output.extend_from_slice(&variable_width_data.offsets);
488
489 output.extend_from_slice(&variable_width_data.data);
491 Ok(LanceBuffer::from(output))
492 }
493 64 => {
494 let offsets = variable_width_data.offsets.borrow_to_typed_slice::<u64>();
495 let offsets = offsets.as_ref();
496 let bytes_start_offset = 8 + 8 + std::mem::size_of_val(offsets) as u64;
499
500 let output_total_bytes =
501 bytes_start_offset as usize + variable_width_data.data.len();
502 let mut output: Vec<u8> = Vec::with_capacity(output_total_bytes);
503
504 output.extend_from_slice(&(64_u64).to_le_bytes());
506
507 output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
509
510 output.extend_from_slice(&variable_width_data.offsets);
512
513 output.extend_from_slice(&variable_width_data.data);
515 Ok(LanceBuffer::from(output))
516 }
517 _ => Err(Error::invalid_input(format!(
518 "BinaryBlockEncoder does not support {}-bit offsets",
519 variable_width_data.bits_per_offset
520 ))),
521 }
522 }
523 _ => unreachable!("variable-width input was validated above"),
524 }
525 .map(|payload| {
526 (
527 Some(payload),
528 ProtobufUtils21::variable(
529 ProtobufUtils21::flat(bits_per_offset as u64, None),
530 None,
531 ),
532 )
533 })
534 }
535}
536
537impl PerValueCompressor for VariableEncoder {
538 fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
539 let DataBlock::VariableWidth(variable) = data else {
540 panic!("BinaryPerValueCompressor can only work with Variable Width DataBlock.");
541 };
542
543 let encoding = ProtobufUtils21::variable(
544 ProtobufUtils21::flat(variable.bits_per_offset as u64, None),
545 None,
546 );
547 Ok((PerValueDataBlock::Variable(variable), encoding))
548 }
549}
550
551#[derive(Debug, Default)]
552pub struct VariableDecoder {}
553
554impl VariablePerValueDecompressor for VariableDecoder {
555 fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
556 Ok(DataBlock::VariableWidth(data))
557 }
558}
559
560#[derive(Debug, Default)]
561pub struct BinaryBlockDecompressor {}
562
563impl BlockDecompressor for BinaryBlockDecompressor {
564 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
565 let data = require_block_payload(data, "Binary block")?;
566 if data.len() < 4 {
584 return Err(Error::corrupt_file_named(
585 "variable-width block",
586 format!(
587 "block of {} bytes is too small to hold a header",
588 data.len()
589 ),
590 ));
591 }
592 let is_old_scheme = data[1] != 0 || data[2] != 0 || data[3] != 0;
593
594 let ensure_header = |header_len: usize| {
595 if data.len() < header_len {
596 return Err(Error::corrupt_file_named(
597 "variable-width block",
598 format!(
599 "block of {} bytes is too small for a {} byte header",
600 data.len(),
601 header_len
602 ),
603 ));
604 }
605 Ok(())
606 };
607 let (bits_per_offset, bytes_start_offset, offset_start) = if is_old_scheme {
608 let bits_per_offset = data[0];
610 match bits_per_offset {
611 32 => {
612 ensure_header(9)?;
613 debug_assert_eq!(LittleEndian::read_u32(&data[1..5]), num_values as u32);
614 let bytes_start_offset = LittleEndian::read_u32(&data[5..9]);
615 (bits_per_offset, bytes_start_offset as u64, 9_u64)
616 }
617 64 => {
618 ensure_header(17)?;
619 debug_assert_eq!(LittleEndian::read_u64(&data[1..9]), num_values);
620 let bytes_start_offset = LittleEndian::read_u64(&data[9..17]);
621 (bits_per_offset, bytes_start_offset, 17)
622 }
623 _ => {
624 return Err(Error::invalid_input_source(
625 format!("Unsupported bits_per_offset={}", bits_per_offset).into(),
626 ));
627 }
628 }
629 } else {
630 let bits_per_offset = LittleEndian::read_u32(&data[0..4]) as u8;
632 match bits_per_offset {
633 32 => {
634 ensure_header(8)?;
635 let bytes_start_offset = LittleEndian::read_u32(&data[4..8]);
636 (bits_per_offset, bytes_start_offset as u64, 8)
637 }
638 64 => {
639 ensure_header(16)?;
640 let bytes_start_offset = LittleEndian::read_u64(&data[8..16]);
641 (bits_per_offset, bytes_start_offset, 16)
642 }
643 _ => {
644 return Err(Error::invalid_input_source(
645 format!("Unsupported bits_per_offset={}", bits_per_offset).into(),
646 ));
647 }
648 }
649 };
650
651 let expected_offsets_bytes = num_values
654 .checked_add(1)
655 .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8))
656 .ok_or_else(|| {
657 Error::corrupt_file_named(
658 "variable-width block",
659 format!("offsets region size overflows for {num_values} values"),
660 )
661 })?;
662 if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 {
663 return Err(Error::corrupt_file_named(
664 "variable-width block",
665 format!(
666 "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)",
667 bytes_start_offset,
668 offset_start,
669 data.len()
670 ),
671 ));
672 }
673 if bytes_start_offset - offset_start != expected_offsets_bytes {
674 return Err(Error::corrupt_file_named(
675 "variable-width block",
676 format!(
677 "expected {} offset bytes for {} values but found {}",
678 expected_offsets_bytes,
679 num_values,
680 bytes_start_offset - offset_start
681 ),
682 ));
683 }
684
685 let offsets = data.slice_with_length(
687 offset_start as usize,
688 (bytes_start_offset - offset_start) as usize,
689 );
690 let first_offset = match bits_per_offset {
691 32 => LittleEndian::read_u32(&offsets[0..4]) as u64,
692 _ => LittleEndian::read_u64(&offsets[0..8]),
693 };
694 if first_offset != 0 {
695 return Err(Error::corrupt_file_named(
696 "variable-width block",
697 format!("first offset must be 0 but found {first_offset}"),
698 ));
699 }
700
701 let data = data.slice_with_length(
703 bytes_start_offset as usize,
704 data.len() - bytes_start_offset as usize,
705 );
706
707 Ok(DataBlock::VariableWidth(VariableWidthBlock {
708 data,
709 offsets,
710 bits_per_offset,
711 num_values,
712 block_info: BlockInfo::new(),
713 }))
714 }
715}
716
717#[cfg(test)]
718mod tests {
719 use arrow_array::{
720 ArrayRef, StringArray,
721 builder::{LargeStringBuilder, StringBuilder},
722 };
723 use arrow_schema::{DataType, Field};
724
725 use crate::{
726 buffer::LanceBuffer,
727 constants::{
728 COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY,
729 STRUCTURAL_ENCODING_MINIBLOCK,
730 },
731 data::{BlockInfo, DataBlock, VariableWidthBlock},
732 testing::check_specific_random,
733 };
734 use rstest::rstest;
735 use std::{collections::HashMap, sync::Arc, vec};
736
737 use crate::testing::{
738 FnArrayGeneratorProvider, TestCases, TestEncoding, check_basic_random_case,
739 check_round_trip_encoding_generated, check_round_trip_encoding_of_data,
740 };
741
742 #[rstest]
743 #[test_log::test(tokio::test)]
744 async fn test_utf8_binary(
745 #[values(
746 TestEncoding::StructuralU16,
747 TestEncoding::StructuralU32,
748 TestEncoding::StructuralSparse
749 )]
750 encoding: TestEncoding,
751 #[values(4096, 1024 * 1024)] page_size: u64,
752 #[values(false, true)] use_slicing: bool,
753 ) {
754 let field = Field::new("", DataType::Utf8, false);
755 check_specific_random(
756 field,
757 TestCases::basic()
758 .with_encoding(encoding)
759 .with_page_sizes(vec![page_size])
760 .with_slicing_modes([use_slicing]),
761 )
762 .await;
763 }
764
765 #[rstest]
766 #[test_log::test(tokio::test)]
767 async fn test_binary(
768 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
769 structural_encoding: &str,
770 #[values(DataType::Utf8, DataType::Binary)] data_type: DataType,
771 #[values(
772 TestEncoding::Array,
773 TestEncoding::StructuralU16,
774 TestEncoding::StructuralU32,
775 TestEncoding::StructuralSparse
776 )]
777 encoding: TestEncoding,
778 #[values(4096, 1024 * 1024)] page_size: u64,
779 #[values(false, true)] use_slicing: bool,
780 ) {
781 let mut field_metadata = HashMap::new();
782 field_metadata.insert(
783 STRUCTURAL_ENCODING_META_KEY.to_string(),
784 structural_encoding.into(),
785 );
786
787 let field = Field::new("", data_type, false).with_metadata(field_metadata);
788 check_basic_random_case(field, encoding, page_size, use_slicing).await;
789 }
790
791 #[rstest]
792 #[test_log::test(tokio::test)]
793 async fn test_binary_fsst(
794 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
795 structural_encoding: &str,
796 #[values(DataType::Binary, DataType::Utf8)] data_type: DataType,
797 #[values(
798 TestEncoding::StructuralU16,
799 TestEncoding::StructuralU32,
800 TestEncoding::StructuralSparse
801 )]
802 encoding: TestEncoding,
803 #[values(4096, 1024 * 1024)] page_size: u64,
804 #[values(false, true)] use_slicing: bool,
805 ) {
806 let mut field_metadata = HashMap::new();
807 field_metadata.insert(
808 STRUCTURAL_ENCODING_META_KEY.to_string(),
809 structural_encoding.into(),
810 );
811 field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into());
812 let field = Field::new("", data_type, true).with_metadata(field_metadata);
813 let test_cases = TestCases::default()
815 .with_encoding(encoding)
816 .with_page_sizes(vec![page_size])
817 .with_slicing_modes([use_slicing]);
818 check_specific_random(field, test_cases).await;
819 }
820
821 #[rstest]
822 #[test_log::test(tokio::test)]
823 async fn test_fsst_large_binary(
824 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
825 structural_encoding: &str,
826 #[values(DataType::LargeBinary, DataType::LargeUtf8)] data_type: DataType,
827 #[values(
828 TestEncoding::StructuralU16,
829 TestEncoding::StructuralU32,
830 TestEncoding::StructuralSparse
831 )]
832 encoding: TestEncoding,
833 #[values(4096, 1024 * 1024)] page_size: u64,
834 #[values(false, true)] use_slicing: bool,
835 ) {
836 let mut field_metadata = HashMap::new();
837 field_metadata.insert(
838 STRUCTURAL_ENCODING_META_KEY.to_string(),
839 structural_encoding.into(),
840 );
841 field_metadata.insert(COMPRESSION_META_KEY.to_string(), "fsst".into());
842 let field = Field::new("", data_type, true).with_metadata(field_metadata);
843 check_specific_random(
844 field,
845 TestCases::basic()
846 .with_encoding(encoding)
847 .with_page_sizes(vec![page_size])
848 .with_slicing_modes([use_slicing]),
849 )
850 .await;
851 }
852
853 #[rstest]
854 #[test_log::test(tokio::test)]
855 async fn test_large_binary_types(
856 #[values(DataType::LargeBinary, DataType::LargeUtf8)] data_type: DataType,
857 #[values(
858 TestEncoding::Array,
859 TestEncoding::StructuralU16,
860 TestEncoding::StructuralU32,
861 TestEncoding::StructuralSparse
862 )]
863 encoding: TestEncoding,
864 #[values(4096, 1024 * 1024)] page_size: u64,
865 #[values(false, true)] use_slicing: bool,
866 ) {
867 let field = Field::new("", data_type, true);
868 check_basic_random_case(field, encoding, page_size, use_slicing).await;
869 }
870
871 #[rstest]
872 #[test_log::test(tokio::test)]
873 async fn test_small_strings(
874 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
875 structural_encoding: &str,
876 #[values(
877 TestEncoding::Array,
878 TestEncoding::StructuralU16,
879 TestEncoding::StructuralU32,
880 TestEncoding::StructuralSparse
881 )]
882 encoding: TestEncoding,
883 #[values(4096, 1024 * 1024)] page_size: u64,
884 #[values(false, true)] use_slicing: bool,
885 ) {
886 let mut field_metadata = HashMap::new();
887 field_metadata.insert(
888 STRUCTURAL_ENCODING_META_KEY.to_string(),
889 structural_encoding.into(),
890 );
891 let field = Field::new("", DataType::Utf8, true).with_metadata(field_metadata);
892 check_round_trip_encoding_generated(
893 field,
894 Box::new(FnArrayGeneratorProvider::new(move || {
895 lance_datagen::array::utf8_prefix_plus_counter("user_", false)
896 })),
897 TestCases::basic()
898 .with_encoding(encoding)
899 .with_page_sizes(vec![page_size])
900 .with_slicing_modes([use_slicing]),
901 )
902 .await;
903 }
904
905 #[rstest]
906 #[test_log::test(tokio::test)]
907 async fn test_simple_binary(
908 #[values(STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_FULLZIP)]
909 structural_encoding: &str,
910 #[values(DataType::Utf8, DataType::Binary)] data_type: DataType,
911 ) {
912 let string_array = StringArray::from(vec![Some("abc"), None, Some("pqr"), None, Some("m")]);
913 let string_array = arrow_cast::cast(&string_array, &data_type).unwrap();
914
915 let mut field_metadata = HashMap::new();
916 field_metadata.insert(
917 STRUCTURAL_ENCODING_META_KEY.to_string(),
918 structural_encoding.into(),
919 );
920
921 let test_cases = TestCases::default()
922 .with_range(0..2)
923 .with_range(0..3)
924 .with_range(1..3)
925 .with_indices(vec![0, 1, 3, 4]);
926 check_round_trip_encoding_of_data(
927 vec![Arc::new(string_array)],
928 &test_cases,
929 field_metadata,
930 )
931 .await;
932 }
933
934 #[test_log::test(tokio::test)]
935 async fn test_sliced_utf8() {
936 let string_array = StringArray::from(vec![Some("abc"), Some("de"), None, Some("fgh")]);
937 let string_array = string_array.slice(1, 3);
938
939 let test_cases = TestCases::default()
940 .with_range(0..1)
941 .with_range(0..2)
942 .with_range(1..2);
943 check_round_trip_encoding_of_data(
944 vec![Arc::new(string_array)],
945 &test_cases,
946 HashMap::new(),
947 )
948 .await;
949 }
950
951 #[rstest]
952 #[test_log::test(tokio::test)]
953 async fn test_value_bigger_than_max_page_size(
954 #[values(
955 TestEncoding::Array,
956 TestEncoding::StructuralU16,
957 TestEncoding::StructuralU32,
958 TestEncoding::StructuralSparse
959 )]
960 encoding: TestEncoding,
961 ) {
962 let big_string = String::from_iter((0..(2 * 1024 * 1024)).map(|_| '0'));
964 let string_array = StringArray::from(vec![
965 Some(big_string),
966 Some("abc".to_string()),
967 None,
968 None,
969 Some("xyz".to_string()),
970 ]);
971
972 let test_cases = TestCases::default()
974 .with_max_page_size(1024 * 1024)
975 .with_encoding(encoding);
976
977 check_round_trip_encoding_of_data(
978 vec![Arc::new(string_array)],
979 &test_cases,
980 HashMap::new(),
981 )
982 .await;
983 }
984
985 #[rstest]
986 #[test_log::test(tokio::test)]
987 async fn test_page_split_parts_do_not_evenly_divide_rows(
988 #[values(
989 TestEncoding::Array,
990 TestEncoding::StructuralU16,
991 TestEncoding::StructuralU32,
992 TestEncoding::StructuralSparse
993 )]
994 encoding: TestEncoding,
995 ) {
996 let big_string = String::from_iter((0..45_000).map(|_| '0'));
999 let string_array = StringArray::from_iter_values((0..90).map(|_| big_string.clone()));
1000
1001 check_round_trip_encoding_of_data(
1002 vec![Arc::new(string_array)],
1003 &TestCases::default()
1004 .with_max_page_size(1024 * 1024)
1005 .with_encoding(encoding),
1006 HashMap::new(),
1007 )
1008 .await;
1009 }
1010
1011 #[test_log::test(tokio::test)]
1012 async fn test_empty_strings() {
1013 let values = [Some("abc"), Some(""), None];
1016 for order in [[0, 1, 2], [1, 0, 2], [2, 0, 1]] {
1018 let mut string_builder = StringBuilder::new();
1019 for idx in order {
1020 string_builder.append_option(values[idx]);
1021 }
1022 let string_array = Arc::new(string_builder.finish());
1023 let test_cases = TestCases::default()
1024 .with_indices(vec![1])
1025 .with_indices(vec![0])
1026 .with_indices(vec![2])
1027 .with_indices(vec![0, 1]);
1028 check_round_trip_encoding_of_data(
1029 vec![string_array.clone()],
1030 &test_cases,
1031 HashMap::new(),
1032 )
1033 .await;
1034 let test_cases = test_cases.with_batch_size(1);
1035 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new())
1036 .await;
1037 }
1038
1039 let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")]));
1044
1045 let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
1046 check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new())
1047 .await;
1048 let test_cases = test_cases.with_batch_size(1);
1049 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
1050 }
1051
1052 #[test_log::test(tokio::test)]
1053 #[ignore] async fn test_jumbo_string() {
1055 let mut string_builder = LargeStringBuilder::new();
1059 let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0'));
1061 for _ in 0..5000 {
1062 string_builder.append_option(Some(&giant_string));
1063 }
1064 let giant_array = Arc::new(string_builder.finish()) as ArrayRef;
1065 let arrs = vec![giant_array];
1066
1067 let test_cases = TestCases::default().without_validation();
1069 check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await;
1070 }
1071
1072 #[rstest]
1073 #[test_log::test(tokio::test)]
1074 async fn test_binary_dictionary_encoding(
1075 #[values(true, false)] with_nulls: bool,
1076 #[values(100, 500, 35000)] dict_size: u32,
1077 ) {
1078 let test_cases = TestCases::default().with_structural_encodings();
1079 let strings = (0..dict_size)
1080 .map(|i| i.to_string())
1081 .collect::<Vec<String>>();
1082
1083 let repeated_strings: Vec<_> = strings
1084 .iter()
1085 .cycle()
1086 .take(70000)
1087 .enumerate()
1088 .map(|(i, s)| {
1089 if with_nulls && i % 7 == 0 {
1090 None
1091 } else {
1092 Some(s.clone())
1093 }
1094 })
1095 .collect();
1096 let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
1097 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
1098 }
1099
1100 #[test_log::test(tokio::test)]
1101 async fn test_binary_encoding_verification() {
1102 use lance_datagen::{ByteCount, RowCount};
1103
1104 let test_cases = TestCases::default()
1105 .with_expected_encoding("variable")
1106 .with_structural_encodings();
1107
1108 let arr_small = lance_datagen::gen_batch()
1111 .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(10), false))
1112 .into_batch_rows(RowCount::from(1000))
1113 .unwrap()
1114 .column(0)
1115 .clone();
1116 check_round_trip_encoding_of_data(vec![arr_small], &test_cases, HashMap::new()).await;
1117
1118 let metadata_explicit =
1120 HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]);
1121 let arr_large = lance_datagen::gen_batch()
1122 .anon_col(lance_datagen::array::rand_utf8(ByteCount::from(50), false))
1123 .into_batch_rows(RowCount::from(2000))
1124 .unwrap()
1125 .column(0)
1126 .clone();
1127 check_round_trip_encoding_of_data(vec![arr_large], &test_cases, metadata_explicit).await;
1128 }
1129
1130 #[test]
1131 fn test_binary_miniblock_with_misaligned_buffer() {
1132 use super::BinaryMiniBlockDecompressor;
1133 use crate::buffer::LanceBuffer;
1134 use crate::compression::MiniBlockDecompressor;
1135 use crate::data::DataBlock;
1136
1137 {
1139 let decompressor = BinaryMiniBlockDecompressor {
1140 bits_per_offset: 32,
1141 };
1142
1143 let mut test_data = Vec::new();
1147
1148 test_data.extend_from_slice(&12u32.to_le_bytes()); test_data.extend_from_slice(&15u32.to_le_bytes()); test_data.extend_from_slice(&20u32.to_le_bytes()); test_data.extend_from_slice(b"ABCXYZ"); test_data.extend_from_slice(&[0, 0]); let mut padded = Vec::with_capacity(test_data.len() + 1);
1159 padded.push(0xFF); padded.extend_from_slice(&test_data);
1161
1162 let bytes = bytes::Bytes::from(padded);
1163 let misaligned = bytes.slice(1..); let buffer = LanceBuffer::from_bytes(misaligned, 1);
1167
1168 let ptr = buffer.as_ref().as_ptr();
1170 assert_ne!(
1171 ptr.align_offset(4),
1172 0,
1173 "Test setup: buffer should be misaligned for u32"
1174 );
1175
1176 let result = decompressor.decompress(vec![buffer], 2);
1178 assert!(
1179 result.is_ok(),
1180 "Decompression should succeed with misaligned buffer"
1181 );
1182
1183 if let Ok(DataBlock::VariableWidth(block)) = result {
1185 assert_eq!(block.num_values, 2);
1186 assert_eq!(&block.data.as_ref()[..6], b"ABCXYZ");
1188 } else {
1189 panic!("Expected VariableWidth block");
1190 }
1191 }
1192
1193 {
1195 let decompressor = BinaryMiniBlockDecompressor {
1196 bits_per_offset: 64,
1197 };
1198
1199 let mut test_data = Vec::new();
1201
1202 test_data.extend_from_slice(&24u64.to_le_bytes()); test_data.extend_from_slice(&29u64.to_le_bytes()); test_data.extend_from_slice(&40u64.to_le_bytes()); test_data.extend_from_slice(b"HelloWorld"); test_data.extend_from_slice(&[0, 0, 0, 0, 0, 0]); let mut padded = Vec::with_capacity(test_data.len() + 3);
1213 padded.extend_from_slice(&[0xFF, 0xFF, 0xFF]); padded.extend_from_slice(&test_data);
1215
1216 let bytes = bytes::Bytes::from(padded);
1217 let misaligned = bytes.slice(3..); let buffer = LanceBuffer::from_bytes(misaligned, 1);
1220
1221 let ptr = buffer.as_ref().as_ptr();
1223 assert_ne!(
1224 ptr.align_offset(8),
1225 0,
1226 "Test setup: buffer should be misaligned for u64"
1227 );
1228
1229 let result = decompressor.decompress(vec![buffer], 2);
1231 assert!(
1232 result.is_ok(),
1233 "Decompression should succeed with misaligned u64 buffer"
1234 );
1235
1236 if let Ok(DataBlock::VariableWidth(block)) = result {
1237 assert_eq!(block.num_values, 2);
1238 assert_eq!(&block.data.as_ref()[..10], b"HelloWorld");
1240 } else {
1241 panic!("Expected VariableWidth block");
1242 }
1243 }
1244 }
1245
1246 #[test]
1247 fn test_binary_miniblock_rejects_corrupt_offsets() {
1248 use super::BinaryMiniBlockDecompressor;
1249 use crate::compression::MiniBlockDecompressor;
1250 use lance_core::Error;
1251
1252 fn chunk_u32(offsets: &[u32], values: &[u8]) -> LanceBuffer {
1256 let mut chunk = offsets
1257 .iter()
1258 .flat_map(|offset| offset.to_le_bytes())
1259 .collect::<Vec<u8>>();
1260 chunk.extend_from_slice(values);
1261 chunk.resize(chunk.len().next_multiple_of(8), 0);
1262 LanceBuffer::from(chunk)
1263 }
1264
1265 let decompressor = BinaryMiniBlockDecompressor::new(32);
1266
1267 let err = decompressor
1269 .decompress(
1270 vec![chunk_u32(&[16, 21, 25, 100_000], b"alphabetagamma")],
1271 3,
1272 )
1273 .unwrap_err();
1274 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1275 assert!(err.to_string().contains("out of bounds"), "{err}");
1276
1277 let err = decompressor
1279 .decompress(vec![chunk_u32(&[16, 25, 21, 30], b"alphabetagamma")], 3)
1280 .unwrap_err();
1281 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1282 assert!(err.to_string().contains("decreases"), "{err}");
1283
1284 let err = decompressor
1287 .decompress(vec![chunk_u32(&[0, 21, 25, 30], b"alphabetagamma")], 3)
1288 .unwrap_err();
1289 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1290 assert!(err.to_string().contains("overlaps"), "{err}");
1291
1292 let err = decompressor
1294 .decompress(vec![chunk_u32(&[8, 8], &[])], 3)
1295 .unwrap_err();
1296 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1297 assert!(err.to_string().contains("requires 4"), "{err}");
1298
1299 let err = decompressor
1301 .decompress(vec![LanceBuffer::from(vec![0u8; 10])], 1)
1302 .unwrap_err();
1303 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1304 assert!(err.to_string().contains("multiple"), "{err}");
1305
1306 fn chunk_u64(offsets: &[u64], values: &[u8]) -> LanceBuffer {
1308 let mut chunk = offsets
1309 .iter()
1310 .flat_map(|offset| offset.to_le_bytes())
1311 .collect::<Vec<u8>>();
1312 chunk.extend_from_slice(values);
1313 chunk.resize(chunk.len().next_multiple_of(8), 0);
1314 LanceBuffer::from(chunk)
1315 }
1316 let decompressor = BinaryMiniBlockDecompressor::new(64);
1317 let err = decompressor
1318 .decompress(
1319 vec![chunk_u64(&[32, 37, 41, 100_000], b"alphabetagamma")],
1320 3,
1321 )
1322 .unwrap_err();
1323 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1324 assert!(err.to_string().contains("out of bounds"), "{err}");
1325 let err = decompressor
1326 .decompress(vec![chunk_u64(&[0, 37, 41, 46], b"alphabetagamma")], 3)
1327 .unwrap_err();
1328 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1329 assert!(err.to_string().contains("overlaps"), "{err}");
1330
1331 let decompressor = BinaryMiniBlockDecompressor::new(32);
1333 let block = decompressor
1334 .decompress(vec![chunk_u32(&[16, 21, 25, 30], b"alphabetagamma")], 3)
1335 .unwrap();
1336 let DataBlock::VariableWidth(block) = block else {
1337 panic!("expected a variable-width block");
1338 };
1339 assert_eq!(block.data.as_ref(), b"alphabetagamma");
1340 assert_eq!(
1341 block.offsets,
1342 LanceBuffer::reinterpret_vec(vec![0_u32, 5, 9, 14])
1343 );
1344 }
1345
1346 fn encoded_binary_block(bits_per_offset: u8) -> Vec<u8> {
1347 use crate::compression::BlockCompressor;
1348
1349 let offsets = match bits_per_offset {
1350 32 => LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 14]),
1351 64 => LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 14]),
1352 _ => unreachable!(),
1353 };
1354 let block = DataBlock::VariableWidth(VariableWidthBlock {
1355 data: LanceBuffer::copy_slice(b"alphabetagamma"),
1356 offsets,
1357 bits_per_offset,
1358 num_values: 3,
1359 block_info: BlockInfo::new(),
1360 });
1361 BlockCompressor::compress(&super::VariableEncoder::default(), block)
1362 .unwrap()
1363 .0
1364 .as_ref()
1365 .unwrap()
1366 .to_vec()
1367 }
1368
1369 #[rstest]
1373 #[case::i32_tail_out_of_bounds(32, 3, 100_000, "out of bounds")]
1374 #[case::i64_tail_out_of_bounds(64, 3, 15, "out of bounds")]
1375 #[case::i32_non_monotonic(32, 2, 4, "non-monotonic")]
1376 #[case::i64_non_monotonic(64, 2, 4, "non-monotonic")]
1377 fn test_binary_block_bad_offsets_rejected_at_arrow_conversion(
1378 #[case] bits_per_offset: u8,
1379 #[case] mutated_offset_index: usize,
1380 #[case] mutated_offset_value: u64,
1381 #[case] expected_message: &str,
1382 ) {
1383 use crate::compression::BlockDecompressor;
1384 use lance_core::Error;
1385
1386 let mut encoded = encoded_binary_block(bits_per_offset);
1387 let bytes_per_offset = (bits_per_offset / 8) as usize;
1388 let mutated_offset_start = bytes_per_offset * (2 + mutated_offset_index);
1390 encoded[mutated_offset_start..mutated_offset_start + bytes_per_offset]
1391 .copy_from_slice(&mutated_offset_value.to_le_bytes()[..bytes_per_offset]);
1392
1393 let block = super::BinaryBlockDecompressor::default()
1394 .decompress(Some(LanceBuffer::from(encoded)), 3)
1395 .unwrap();
1396 let data_type = match bits_per_offset {
1397 32 => DataType::Binary,
1398 _ => DataType::LargeBinary,
1399 };
1400 let err = block.into_arrow(data_type, false).unwrap_err();
1401 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1402 assert!(err.to_string().contains(expected_message), "{err}");
1403 }
1404
1405 #[test]
1406 fn test_binary_block_rejects_corrupt_structure() {
1407 use crate::compression::BlockDecompressor;
1408 use lance_core::Error;
1409
1410 let decompressor = super::BinaryBlockDecompressor::default();
1411
1412 let mut encoded = encoded_binary_block(32);
1414 encoded[8..12].copy_from_slice(&5_u32.to_le_bytes());
1415 let err = decompressor
1416 .decompress(Some(LanceBuffer::from(encoded)), 3)
1417 .unwrap_err();
1418 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1419 assert!(err.to_string().contains("first offset"), "{err}");
1420
1421 let encoded = encoded_binary_block(32);
1423 let err = decompressor
1424 .decompress(Some(LanceBuffer::from(encoded)), 4)
1425 .unwrap_err();
1426 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1427 assert!(err.to_string().contains("offset bytes"), "{err}");
1428
1429 let err = decompressor
1431 .decompress(Some(LanceBuffer::from(vec![0_u8; 2])), 1)
1432 .unwrap_err();
1433 assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}");
1434 assert!(err.to_string().contains("too small"), "{err}");
1435 }
1436}