1use 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
59fn 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 bytes_start_offset: usize,
72 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 const PAD_BYTE: u8 = 72;
146 let pad_len = pad_bytes_to(output.len(), alignment);
147
148 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
156fn search_next_offset_idx<N: OffsetSizeTrait>(
161 offsets: &[N],
162 last_offset_idx: usize,
163 minichunk_size: i64,
164) -> usize {
165 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 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 return offsets.len() - 1;
184 } else {
185 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 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 fn chunk_data(&self, data: VariableWidthBlock) -> (MiniBlockCompressed, CompressiveEncoding) {
217 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
295fn 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 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 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#[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 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 output.extend_from_slice(&(32_u32).to_le_bytes());
472
473 output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
475
476 output.extend_from_slice(&variable_width_data.offsets);
478
479 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 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 output.extend_from_slice(&(64_u64).to_le_bytes());
496
497 output.extend_from_slice(&(bytes_start_offset).to_le_bytes());
499
500 output.extend_from_slice(&variable_width_data.offsets);
502
503 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 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 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 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 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 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 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 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_", 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 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 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 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 let values = [Some("abc"), Some(""), None];
918 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 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] async fn test_jumbo_string() {
957 let mut string_builder = LargeStringBuilder::new();
961 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 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 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 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 {
1041 let decompressor = BinaryMiniBlockDecompressor {
1042 bits_per_offset: 32,
1043 };
1044
1045 let mut test_data = Vec::new();
1049
1050 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);
1061 padded.push(0xFF); padded.extend_from_slice(&test_data);
1063
1064 let bytes = bytes::Bytes::from(padded);
1065 let misaligned = bytes.slice(1..); let buffer = LanceBuffer::from_bytes(misaligned, 1);
1069
1070 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 let result = decompressor.decompress(vec![buffer], 2);
1080 assert!(
1081 result.is_ok(),
1082 "Decompression should succeed with misaligned buffer"
1083 );
1084
1085 if let Ok(DataBlock::VariableWidth(block)) = result {
1087 assert_eq!(block.num_values, 2);
1088 assert_eq!(&block.data.as_ref()[..6], b"ABCXYZ");
1090 } else {
1091 panic!("Expected VariableWidth block");
1092 }
1093 }
1094
1095 {
1097 let decompressor = BinaryMiniBlockDecompressor {
1098 bits_per_offset: 64,
1099 };
1100
1101 let mut test_data = Vec::new();
1103
1104 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);
1115 padded.extend_from_slice(&[0xFF, 0xFF, 0xFF]); padded.extend_from_slice(&test_data);
1117
1118 let bytes = bytes::Bytes::from(padded);
1119 let misaligned = bytes.slice(3..); let buffer = LanceBuffer::from_bytes(misaligned, 1);
1122
1123 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 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 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 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 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 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 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 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 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 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 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 #[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 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 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 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 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}