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