1#[cfg(feature = "bitpacking")]
20use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking};
21use crate::{
22 buffer::LanceBuffer,
23 compression_config::{BssMode, CompressionFieldParams},
24 constants::{
25 BSS_META_KEY, COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, RLE_THRESHOLD_META_KEY,
26 },
27 data::{DataBlock, FixedWidthDataBlock, VariableWidthBlock},
28 encodings::{
29 logical::primitive::{
30 fullzip::PerValueCompressor,
31 miniblock::{MAX_MINIBLOCK_VALUES, MiniBlockCompressor},
32 },
33 physical::{
34 binary::{
35 BinaryBlockDecompressor, BinaryMiniBlockDecompressor, BinaryMiniBlockEncoder,
36 VariableDecoder, VariableEncoder,
37 },
38 block::{
39 CompressedBufferEncoder, CompressionConfig, CompressionScheme,
40 GeneralBlockDecompressor,
41 },
42 byte_stream_split::{
43 ByteStreamSplitDecompressor, ByteStreamSplitEncoder, should_use_bss,
44 },
45 constant::{ConstantBlockDecompressor, ConstantDecompressor},
46 delta::DeltaDecompressor,
47 dictionary::BlockDictionaryDecompressor,
48 fsst::{
49 FsstMiniBlockDecompressor, FsstMiniBlockEncoder, FsstPerValueDecompressor,
50 FsstPerValueEncoder,
51 },
52 general::{
53 FixedWidthGeneralBlockDecompressor, GeneralBlockCompressor,
54 GeneralMiniBlockCompressor, GeneralMiniBlockDecompressor,
55 },
56 packed::{
57 PackedStructFixedPerValueDecompressor, PackedStructFixedPerValueEncoder,
58 PackedStructFixedWidthMiniBlockDecompressor,
59 PackedStructFixedWidthMiniBlockEncoder, PackedStructVariablePerValueDecompressor,
60 PackedStructVariablePerValueEncoder, VariablePackedStructFieldDecoder,
61 VariablePackedStructFieldKind,
62 },
63 range::RangeDecompressor,
64 rle::{
65 BlockRleDecompressor, BlockRleRunCount, RleChildDecompressor, RleDecompressor,
66 RleEncoder, RunLengthWidth, rle_encoded_size, select_miniblock_run_length_width,
67 select_run_length_width,
68 },
69 value::{ValueDecompressor, ValueEncoder},
70 },
71 },
72 format::metadata::{CompressiveEncoding, compressive_encoding::Compression},
73 statistics::{GetStat, Stat},
74};
75
76use arrow_array::{cast::AsArray, types::UInt64Type};
77use arrow_schema::DataType;
78use fsst::fsst::{FSST_LEAST_INPUT_MAX_LENGTH, FSST_LEAST_INPUT_SIZE};
79use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt};
80use std::{str::FromStr, sync::Arc};
81
82const DEFAULT_RLE_COMPRESSION_THRESHOLD: f64 = 0.5;
88
89const MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION: u64 = 32 * 1024;
91const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::<u64>() as u128;
92
93pub trait BlockCompressor: std::fmt::Debug + Send + Sync {
106 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)>;
111}
112
113pub trait CompressionStrategy: Send + Sync + std::fmt::Debug {
126 fn create_block_compressor(
128 &self,
129 field: &Field,
130 data: &DataBlock,
131 ) -> Result<Box<dyn BlockCompressor>>;
132
133 fn create_per_value(
135 &self,
136 field: &Field,
137 data: &DataBlock,
138 ) -> Result<Box<dyn PerValueCompressor>>;
139
140 fn create_miniblock_compressor(
142 &self,
143 field: &Field,
144 data: &DataBlock,
145 ) -> Result<Box<dyn MiniBlockCompressor>>;
146
147 fn create_dictionary_index_compressor(
151 &self,
152 field: &Field,
153 data: &DataBlock,
154 ) -> Result<Box<dyn MiniBlockCompressor>> {
155 self.create_miniblock_compressor(field, data)
156 }
157}
158
159pub(crate) fn compress_required_block(
160 strategy: &dyn CompressionStrategy,
161 field: &Field,
162 data: DataBlock,
163) -> Result<(LanceBuffer, CompressiveEncoding)> {
164 let compressor = strategy.create_block_compressor(field, &data)?;
165 let (payload, encoding) = compressor.compress(data)?;
166 let payload = payload.ok_or_else(|| {
167 Error::internal("Required block compressor selected a metadata-only codec".to_string())
168 })?;
169 Ok((payload, encoding))
170}
171
172fn try_bss_for_mini_block(
173 data: &FixedWidthDataBlock,
174 params: &CompressionFieldParams,
175) -> Option<Box<dyn MiniBlockCompressor>> {
176 if params.compression.is_none() || params.compression.as_deref() == Some("none") {
179 return None;
180 }
181
182 let mode = params.bss.unwrap_or(BssMode::Auto);
183 if should_use_bss(data, mode) {
185 return Some(Box::new(ByteStreamSplitEncoder::new(
186 data.bits_per_value as usize,
187 )));
188 }
189 None
190}
191
192fn rle_is_applicable(data: &FixedWidthDataBlock, params: &CompressionFieldParams) -> Option<u128> {
193 let bits = data.bits_per_value;
194 if !matches!(bits, 8 | 16 | 32 | 64) {
195 return None;
196 }
197
198 let type_size = bits / 8;
199 let run_count = data.expect_single_stat::<UInt64Type>(Stat::RunCount);
200 let threshold = params
201 .rle_threshold
202 .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD);
203
204 let passes_threshold = match params.rle_threshold {
207 Some(_) => (run_count as f64) < (data.num_values as f64) * threshold,
208 None => true,
209 };
210
211 if !passes_threshold {
212 return None;
213 }
214
215 Some((data.num_values as u128) * (type_size as u128))
216}
217
218fn rle_beats_raw_and_bitpacking(
219 data: &FixedWidthDataBlock,
220 encoded_bytes: u128,
221 raw_bytes: u128,
222) -> bool {
223 if encoded_bytes >= raw_bytes {
224 return false;
225 }
226
227 #[cfg(feature = "bitpacking")]
228 {
229 if let Some(bitpack_bytes) = estimate_inline_bitpacking_bytes(data).map(u128::from)
230 && bitpack_bytes < encoded_bytes
231 {
232 return false;
233 }
234 }
235 true
236}
237
238fn try_fixed_u8_rle_for_mini_block(
239 data: &FixedWidthDataBlock,
240 params: &CompressionFieldParams,
241) -> Option<Box<dyn MiniBlockCompressor>> {
242 let raw_bytes = rle_is_applicable(data, params)?;
243 let rle_bytes = estimate_rle_size_for_width_from_data(
244 data,
245 Some(*MAX_MINIBLOCK_VALUES),
246 RunLengthWidth::U8,
247 )
248 .ok()?;
249 rle_beats_raw_and_bitpacking(data, rle_bytes, raw_bytes)
250 .then(|| Box::new(RleEncoder::with_run_length_width(RunLengthWidth::U8)) as _)
251}
252
253fn try_child_rle_for_mini_block(
254 data: &FixedWidthDataBlock,
255 params: &CompressionFieldParams,
256) -> Option<Box<dyn MiniBlockCompressor>> {
257 let raw_bytes = rle_is_applicable(data, params)?;
258 let (run_length_width, estimated_bytes, max_values) = select_miniblock_run_length_width(
261 data,
262 *MAX_MINIBLOCK_VALUES,
263 (*MAX_MINIBLOCK_VALUES).min(2048),
264 )
265 .ok()?;
266 let child_compression = rle_child_compression_config(params);
267 let encoder = || {
268 RleEncoder::with_child_encoding(
269 run_length_width,
270 child_compression,
271 child_compression,
272 true,
273 )
274 .with_max_miniblock_values(max_values as usize)
275 };
276
277 #[cfg(feature = "bitpacking")]
278 let bitpack_bytes = estimate_inline_bitpacking_bytes(data).map(u128::from);
279 #[cfg(not(feature = "bitpacking"))]
280 let bitpack_bytes = None::<u128>;
281
282 let should_measure_children = (child_compression.is_some() || cfg!(feature = "bitpacking"))
283 && (estimated_bytes >= raw_bytes
284 || bitpack_bytes.is_some_and(|bytes| bytes < estimated_bytes));
285 let selected_bytes = if should_measure_children {
286 encoder().estimated_payload_size(data).ok()?
287 } else {
288 estimated_bytes
289 };
290
291 rle_beats_raw_and_bitpacking(data, selected_bytes, raw_bytes).then(|| Box::new(encoder()) as _)
292}
293
294fn rle_child_compression_config(params: &CompressionFieldParams) -> Option<CompressionConfig> {
295 let raw = params.compression.as_deref()?;
296 if matches!(raw, "none" | "fsst") {
297 return None;
298 }
299 let scheme = CompressionScheme::from_str(raw).ok()?;
300 Some(CompressionConfig::new(scheme, params.compression_level))
301}
302
303fn try_rle_for_block_with_width(
304 data: &FixedWidthDataBlock,
305 params: &CompressionFieldParams,
306 run_length_width: RunLengthWidth,
307 rle_payload_bytes: u128,
308) -> Result<Option<Box<dyn BlockCompressor>>> {
309 let bits = data.bits_per_value;
310 if !matches!(bits, 8 | 16 | 32 | 64) {
311 return Ok(None);
312 }
313
314 let run_count = data.expect_single_stat::<UInt64Type>(Stat::RunCount);
315 let threshold = params
316 .rle_threshold
317 .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD);
318
319 let passes_threshold = match params.rle_threshold {
320 Some(_) => (run_count as f64) < (data.num_values as f64) * threshold,
321 None => true,
322 };
323
324 if !passes_threshold {
325 return Ok(None);
326 }
327
328 let raw_bytes = (data.num_values as u128) * ((bits / 8) as u128);
329 let rle_bytes = rle_payload_bytes.saturating_add(RLE_BLOCK_HEADER_BYTES);
330
331 if rle_bytes >= raw_bytes {
332 return Ok(None);
333 }
334
335 #[cfg(feature = "bitpacking")]
336 {
337 if let Some(bitpack_bytes) = estimate_block_bitpacking_bytes(data)
338 && bitpack_bytes < rle_bytes
339 {
340 return Ok(None);
341 }
342 }
343
344 Ok(Some(Box::new(RleEncoder::with_run_length_width(
345 run_length_width,
346 ))))
347}
348
349fn try_fixed_u8_rle_for_block(
350 data: &FixedWidthDataBlock,
351 params: &CompressionFieldParams,
352) -> Result<Option<Box<dyn BlockCompressor>>> {
353 if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) {
354 return Ok(None);
355 }
356 let encoded_bytes = estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?;
357 try_rle_for_block_with_width(data, params, RunLengthWidth::U8, encoded_bytes)
358}
359
360fn try_variable_rle_for_block(
361 data: &FixedWidthDataBlock,
362 params: &CompressionFieldParams,
363) -> Result<Option<Box<dyn BlockCompressor>>> {
364 if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) {
365 return Ok(None);
366 }
367 let (width, encoded_bytes) = estimate_rle_width_and_size_from_data(data, None)?;
368 try_rle_for_block_with_width(data, params, width, encoded_bytes)
369}
370
371fn estimate_rle_width_and_size_from_data(
372 data: &FixedWidthDataBlock,
373 max_segment_values: Option<u64>,
374) -> Result<(RunLengthWidth, u128)> {
375 select_run_length_width(
376 &data.data,
377 data.num_values,
378 data.bits_per_value,
379 max_segment_values,
380 )
381}
382
383fn estimate_rle_size_for_width_from_data(
384 data: &FixedWidthDataBlock,
385 max_segment_values: Option<u64>,
386 run_length_width: RunLengthWidth,
387) -> Result<u128> {
388 rle_encoded_size(
389 &data.data,
390 data.num_values,
391 data.bits_per_value,
392 max_segment_values,
393 run_length_width,
394 )
395}
396
397fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
398 #[cfg(feature = "bitpacking")]
399 {
400 let bits = _data.bits_per_value;
401 if estimate_inline_bitpacking_bytes(_data).is_some() {
402 return Some(Box::new(InlineBitpacking::new(bits)));
403 }
404 None
405 }
406 #[cfg(not(feature = "bitpacking"))]
407 {
408 None
409 }
410}
411
412#[cfg(feature = "bitpacking")]
413fn estimate_inline_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option<u64> {
414 use arrow_array::cast::AsArray;
415
416 let bits = data.bits_per_value;
417 if !matches!(bits, 8 | 16 | 32 | 64) {
418 return None;
419 }
420 if data.num_values == 0 {
421 return None;
422 }
423
424 let bit_widths = data.expect_stat(Stat::BitWidth);
425 let widths = bit_widths.as_primitive::<UInt64Type>();
426
427 let words_per_chunk: u128 = 1;
428 let word_bytes: u128 = (bits / 8) as u128;
429 let mut total_words: u128 = 0;
430 for i in 0..widths.len() {
431 let bit_width = widths.value(i) as u128;
432 let packed_words = (1024u128 * bit_width) / (bits as u128);
433 total_words = total_words.saturating_add(words_per_chunk.saturating_add(packed_words));
434 }
435
436 let estimated_bytes = total_words.saturating_mul(word_bytes);
437 let raw_bytes = data.data_size() as u128;
438
439 if estimated_bytes >= raw_bytes {
440 return None;
441 }
442
443 u64::try_from(estimated_bytes).ok()
444}
445
446fn try_bitpack_for_block(data: &FixedWidthDataBlock) -> Option<Box<dyn BlockCompressor>> {
447 let bits = data.bits_per_value;
448 if !matches!(bits, 8 | 16 | 32 | 64) {
449 return None;
450 }
451
452 let bit_widths = data.expect_stat(Stat::BitWidth);
453 let widths = bit_widths.as_primitive::<UInt64Type>();
454 let max_bit_width = *widths.values().iter().max().unwrap();
455
456 if max_bit_width >= bits {
458 return None;
459 }
460
461 let too_small =
462 widths.len() == 1 && InlineBitpacking::min_size_bytes(widths.value(0)) >= data.data_size();
463
464 if too_small {
465 return None;
466 }
467
468 if data.num_values <= 1024 {
469 Some(Box::new(InlineBitpacking::new(bits)))
470 } else {
471 Some(Box::new(OutOfLineBitpacking::new(max_bit_width, bits)))
472 }
473}
474
475#[cfg(feature = "bitpacking")]
476fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option<u128> {
477 let bits = data.bits_per_value;
478 if !matches!(bits, 8 | 16 | 32 | 64) || data.num_values == 0 {
479 return None;
480 }
481
482 let bit_widths = data.expect_stat(Stat::BitWidth);
483 let widths = bit_widths.as_primitive::<UInt64Type>();
484 let max_bit_width = *widths.values().iter().max()?;
485 let word_bytes = (bits / 8) as u128;
486
487 let bitpacked_words = if data.num_values <= 1024 {
488 1 + (1024u128 * (max_bit_width as u128)) / (bits as u128)
489 } else {
490 estimate_out_of_line_bitpacking_words(data.num_values, max_bit_width, bits)?
491 };
492 let bitpacked_bytes = bitpacked_words.saturating_mul(word_bytes);
493 if bitpacked_bytes >= data.data_size() as u128 {
494 return None;
495 }
496
497 Some(bitpacked_bytes)
498}
499
500#[cfg(feature = "bitpacking")]
501fn estimate_out_of_line_bitpacking_words(
502 num_values: u64,
503 compressed_bits_per_value: u64,
504 bits_per_value: u64,
505) -> Option<u128> {
506 let num_values = usize::try_from(num_values).ok()?;
507 let compressed_bits_per_value = usize::try_from(compressed_bits_per_value).ok()?;
508 let bits_per_value = usize::try_from(bits_per_value).ok()?;
509 if compressed_bits_per_value >= bits_per_value {
510 return None;
511 }
512
513 let elems_per_chunk = 1024usize;
514 let num_chunks = num_values.div_ceil(elems_per_chunk);
515 let words_per_chunk = (elems_per_chunk * compressed_bits_per_value).div_ceil(bits_per_value);
516 let last_chunk_is_runt = !num_values.is_multiple_of(elems_per_chunk);
517
518 if !last_chunk_is_runt {
519 return Some((num_chunks * words_per_chunk) as u128);
520 }
521
522 let num_whole_chunks = num_chunks - 1;
523 let remaining_items = num_values - num_whole_chunks * elems_per_chunk;
524 let tail_bit_savings = bits_per_value - compressed_bits_per_value;
525 let padding_cost = compressed_bits_per_value * (elems_per_chunk - remaining_items);
526 let tail_pack_savings = tail_bit_savings * remaining_items;
527 let tail_words = if padding_cost < tail_pack_savings {
528 words_per_chunk
529 } else {
530 remaining_items
531 };
532
533 Some((num_whole_chunks * words_per_chunk + tail_words) as u128)
534}
535
536fn maybe_wrap_general_for_mini_block(
537 inner: Box<dyn MiniBlockCompressor>,
538 params: &CompressionFieldParams,
539) -> Result<Box<dyn MiniBlockCompressor>> {
540 match params.compression.as_deref() {
541 None | Some("none") | Some("fsst") => Ok(inner),
542 Some(raw) => {
543 let scheme = CompressionScheme::from_str(raw)
544 .map_err(|_| Error::invalid_input(format!("Unknown compression scheme: {raw}")))?;
545 let cfg = CompressionConfig::new(scheme, params.compression_level);
546 Ok(Box::new(GeneralMiniBlockCompressor::new(inner, cfg)))
547 }
548 }
549}
550
551fn try_general_compression(
552 field_params: &CompressionFieldParams,
553 data: &DataBlock,
554) -> Result<Option<(Box<dyn BlockCompressor>, CompressionConfig)>> {
555 if field_params.compression.as_deref() == Some("none") {
557 return Ok(None);
558 }
559
560 if let Some(compression_scheme) = &field_params.compression {
563 let scheme: CompressionScheme = compression_scheme.parse()?;
564 let config = CompressionConfig::new(scheme, field_params.compression_level);
565 let compressor = Box::new(CompressedBufferEncoder::try_new(config)?);
566 return Ok(Some((compressor, config)));
567 }
568
569 if data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION {
571 let compressor = Box::new(CompressedBufferEncoder::default());
572 let config = compressor.compressor.config();
573 return Ok(Some((compressor, config)));
574 }
575
576 Ok(None)
577}
578
579pub fn field_metadata_params(field: &Field) -> CompressionFieldParams {
581 let mut params = CompressionFieldParams::default();
582
583 if let Some(compression) = field.metadata.get(COMPRESSION_META_KEY) {
584 params.compression = Some(compression.clone());
585 }
586 if let Some(level) = field.metadata.get(COMPRESSION_LEVEL_META_KEY) {
587 params.compression_level = level.parse().ok();
588 }
589 if let Some(threshold) = field.metadata.get(RLE_THRESHOLD_META_KEY) {
590 params.rle_threshold = threshold.parse().ok();
591 }
592 if let Some(bss_str) = field.metadata.get(BSS_META_KEY) {
593 match BssMode::parse(bss_str) {
594 Some(mode) => params.bss = Some(mode),
595 None => log::warn!("Invalid BSS mode '{}', using default", bss_str),
596 }
597 }
598 if let Some(minichunk_size_str) = field
599 .metadata
600 .get(super::constants::MINICHUNK_SIZE_META_KEY)
601 {
602 if let Ok(minichunk_size) = minichunk_size_str.parse::<i64>() {
603 params.minichunk_size = Some(minichunk_size);
604 } else {
605 log::warn!("Invalid minichunk_size '{}', skipping", minichunk_size_str);
606 }
607 }
608
609 params
610}
611
612pub fn finalize_miniblock_compressor(
614 data: &DataBlock,
615 compressor: Box<dyn MiniBlockCompressor>,
616 params: &CompressionFieldParams,
617) -> Result<Box<dyn MiniBlockCompressor>> {
618 if matches!(data, DataBlock::FixedWidth(_)) {
619 maybe_wrap_general_for_mini_block(compressor, params)
620 } else {
621 Ok(compressor)
622 }
623}
624
625pub fn try_uncompressed_fixed_width_miniblock(
627 data: &DataBlock,
628 params: &CompressionFieldParams,
629) -> Option<Box<dyn MiniBlockCompressor>> {
630 (matches!(data, DataBlock::FixedWidth(_)) && params.compression.as_deref() == Some("none"))
631 .then(|| Box::new(ValueEncoder::default()) as _)
632}
633
634pub fn try_byte_stream_split_miniblock(
636 data: &DataBlock,
637 params: &CompressionFieldParams,
638) -> Option<Box<dyn MiniBlockCompressor>> {
639 let DataBlock::FixedWidth(data) = data else {
640 return None;
641 };
642 try_bss_for_mini_block(data, params)
643}
644
645pub fn try_fixed_u8_rle_miniblock(
647 data: &DataBlock,
648 params: &CompressionFieldParams,
649) -> Option<Box<dyn MiniBlockCompressor>> {
650 let DataBlock::FixedWidth(data) = data else {
651 return None;
652 };
653 try_fixed_u8_rle_for_mini_block(data, params)
654}
655
656pub fn try_child_rle_miniblock(
658 data: &DataBlock,
659 params: &CompressionFieldParams,
660) -> Option<Box<dyn MiniBlockCompressor>> {
661 let DataBlock::FixedWidth(data) = data else {
662 return None;
663 };
664 try_child_rle_for_mini_block(data, params)
665}
666
667pub fn try_bitpacking_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
669 let DataBlock::FixedWidth(data) = data else {
670 return None;
671 };
672 try_bitpack_for_mini_block(data)
673}
674
675pub fn try_raw_fixed_width_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
677 matches!(data, DataBlock::FixedWidth(_)).then(|| Box::new(ValueEncoder::default()) as _)
678}
679
680pub fn try_variable_width_miniblock(
682 field: &Field,
683 data: &DataBlock,
684 params: &CompressionFieldParams,
685) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
686 try_variable_width_miniblock_impl(field, data, params, false)
687}
688
689pub fn try_variable_width_miniblock_with_generic_offsets(
694 field: &Field,
695 data: &DataBlock,
696 params: &CompressionFieldParams,
697) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
698 try_variable_width_miniblock_impl(field, data, params, true)
699}
700
701fn try_variable_width_miniblock_impl(
702 field: &Field,
703 data: &DataBlock,
704 params: &CompressionFieldParams,
705 allow_generic_offsets: bool,
706) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
707 let DataBlock::VariableWidth(data) = data else {
708 return Ok(None);
709 };
710 if data.bits_per_offset != 32 && data.bits_per_offset != 64 {
711 return Err(Error::invalid_input(format!(
712 "Variable width compression not supported for {} bit offsets",
713 data.bits_per_offset
714 )));
715 }
716
717 let compression = params.compression.as_deref();
718 let data_size = data.expect_single_stat::<UInt64Type>(Stat::DataSize);
719 let max_len = data.expect_single_stat::<UInt64Type>(Stat::MaxLength);
720 if compression == Some("none") {
721 return Ok(Some(if allow_generic_offsets {
722 Box::new(BinaryMiniBlockEncoder::with_generic_offsets(
723 params.minichunk_size,
724 params.clone(),
725 ))
726 } else {
727 Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size))
728 }));
729 }
730
731 let use_fsst = compression == Some("fsst")
732 || (compression.is_none()
733 && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary)
734 && max_len >= FSST_LEAST_INPUT_MAX_LENGTH
735 && data_size >= FSST_LEAST_INPUT_SIZE as u64);
736 let mut encoder: Box<dyn MiniBlockCompressor> = if use_fsst {
737 Box::new(FsstMiniBlockEncoder::new(params.minichunk_size))
738 } else if allow_generic_offsets && compression.is_none() {
739 Box::new(BinaryMiniBlockEncoder::with_generic_offsets(
743 params.minichunk_size,
744 params.clone(),
745 ))
746 } else {
747 Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size))
748 };
749 if let Some(compression_scheme) = compression.filter(|scheme| *scheme != "fsst") {
750 let scheme: CompressionScheme = compression_scheme.parse()?;
751 let config = CompressionConfig::new(scheme, params.compression_level);
752 encoder = Box::new(GeneralMiniBlockCompressor::new(encoder, config));
753 }
754 Ok(Some(encoder))
755}
756
757pub fn try_fixed_packed_struct_miniblock(
759 data: &DataBlock,
760) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
761 let DataBlock::Struct(data) = data else {
762 return Ok(None);
763 };
764 if data.has_variable_width_child() {
765 return Err(Error::invalid_input(
766 "Packed struct mini-block encoding supports only fixed-width children",
767 ));
768 }
769 Ok(Some(Box::new(
770 PackedStructFixedWidthMiniBlockEncoder::default(),
771 )))
772}
773
774pub fn try_raw_fixed_size_list_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
776 matches!(data, DataBlock::FixedSizeList(_)).then(|| Box::new(ValueEncoder::default()) as _)
777}
778
779pub fn try_raw_per_value(data: &DataBlock) -> Option<Box<dyn PerValueCompressor>> {
781 matches!(data, DataBlock::FixedWidth(_) | DataBlock::FixedSizeList(_))
782 .then(|| Box::new(ValueEncoder::default()) as _)
783}
784
785fn validate_packed_struct(field: &Field, data: &DataBlock) -> Result<Option<bool>> {
786 let DataBlock::Struct(data) = data else {
787 return Ok(None);
788 };
789 if field.children.len() != data.children.len() {
790 return Err(Error::invalid_input(
791 "Struct field metadata does not match data block children",
792 ));
793 }
794 Ok(Some(data.has_variable_width_child()))
795}
796
797pub fn reject_packed_struct_per_value(
799 field: &Field,
800 data: &DataBlock,
801) -> Result<Option<Box<dyn PerValueCompressor>>> {
802 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
803 return Ok(None);
804 };
805 if has_variable_child {
806 return Err(Error::not_supported_source(
807 "Variable packed struct encoding is not enabled by the selected file format".into(),
808 ));
809 }
810 Err(Error::invalid_input(
811 "Packed struct per-value compression should not be used for fixed-width-only structs",
812 ))
813}
814
815pub fn try_variable_packed_struct_per_value(
817 strategy: Arc<dyn CompressionStrategy>,
818 field: &Field,
819 data: &DataBlock,
820) -> Result<Option<Box<dyn PerValueCompressor>>> {
821 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
822 return Ok(None);
823 };
824 if !has_variable_child {
825 return Err(Error::invalid_input(
826 "Packed struct per-value compression should not be used for fixed-width-only structs",
827 ));
828 }
829 Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new(
830 strategy,
831 field.children.clone(),
832 ))))
833}
834
835pub fn try_packed_struct_per_value(
837 strategy: Arc<dyn CompressionStrategy>,
838 field: &Field,
839 data: &DataBlock,
840) -> Result<Option<Box<dyn PerValueCompressor>>> {
841 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
842 return Ok(None);
843 };
844 if has_variable_child {
845 return Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new(
846 strategy,
847 field.children.clone(),
848 ))));
849 }
850
851 Ok(Some(Box::new(PackedStructFixedPerValueEncoder::new(
852 field.children.clone(),
853 ))))
854}
855
856pub fn try_variable_width_per_value(
859 field: &Field,
860 data: &DataBlock,
861 params: &CompressionFieldParams,
862) -> Result<Option<Box<dyn PerValueCompressor>>> {
863 let DataBlock::VariableWidth(data) = data else {
864 return Ok(None);
865 };
866 let compression = params.compression.as_deref();
867 if compression == Some("none") {
868 return Ok(Some(Box::new(VariableEncoder::default())));
869 }
870
871 let max_len = data.expect_single_stat::<UInt64Type>(Stat::MaxLength);
872 let data_size = data.expect_single_stat::<UInt64Type>(Stat::DataSize);
873 let per_value_requested = compression.is_some_and(|compression| compression != "fsst");
874 if (max_len > 32 * 1024 || per_value_requested) && data_size >= FSST_LEAST_INPUT_SIZE as u64 {
875 if compression == Some("zstd") {
876 let config = CompressionConfig::new(CompressionScheme::Zstd, params.compression_level);
877 return Ok(Some(Box::new(CompressedBufferEncoder::try_new(config)?)));
878 }
879 return Ok(Some(Box::new(CompressedBufferEncoder::default())));
880 }
881
882 if data.bits_per_offset != 32 && data.bits_per_offset != 64 {
883 return Err(Error::invalid_input(format!(
884 "Per-value compression does not support variable-width data with {}-bit offsets",
885 data.bits_per_offset
886 )));
887 }
888 let encoder = Box::new(VariableEncoder::default());
889 let use_fsst = compression == Some("fsst")
890 || (compression.is_none()
891 && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary)
892 && max_len >= FSST_LEAST_INPUT_MAX_LENGTH
893 && data_size >= FSST_LEAST_INPUT_SIZE as u64);
894 Ok(Some(if use_fsst {
895 Box::new(FsstPerValueEncoder::new(encoder))
896 } else {
897 encoder
898 }))
899}
900
901pub fn try_fixed_u8_rle_block(
903 data: &DataBlock,
904 params: &CompressionFieldParams,
905) -> Result<Option<Box<dyn BlockCompressor>>> {
906 let DataBlock::FixedWidth(data) = data else {
907 return Ok(None);
908 };
909 try_fixed_u8_rle_for_block(data, params)
910}
911
912pub fn try_variable_rle_block(
914 data: &DataBlock,
915 params: &CompressionFieldParams,
916) -> Result<Option<Box<dyn BlockCompressor>>> {
917 let DataBlock::FixedWidth(data) = data else {
918 return Ok(None);
919 };
920 try_variable_rle_for_block(data, params)
921}
922
923pub fn try_bitpacking_block(data: &DataBlock) -> Option<Box<dyn BlockCompressor>> {
925 let DataBlock::FixedWidth(data) = data else {
926 return None;
927 };
928 try_bitpack_for_block(data)
929}
930
931pub fn try_general_block(
933 data: &DataBlock,
934 params: &CompressionFieldParams,
935) -> Result<Option<Box<dyn BlockCompressor>>> {
936 let Some((compressor, _config)) = try_general_compression(params, data)? else {
937 return Ok(None);
938 };
939 Ok(Some(compressor))
940}
941
942pub fn try_raw_block(data: &DataBlock) -> Option<Box<dyn BlockCompressor>> {
944 match data {
945 DataBlock::FixedWidth(_) => {
946 Some(Box::new(ValueEncoder::default()) as Box<dyn BlockCompressor>)
947 }
948 DataBlock::VariableWidth(_) => {
949 Some(Box::new(VariableEncoder::default()) as Box<dyn BlockCompressor>)
950 }
951 _ => None,
952 }
953}
954
955pub fn try_fixed_size_list_block(
960 data: &DataBlock,
961 params: &CompressionFieldParams,
962) -> Result<Option<Box<dyn BlockCompressor>>> {
963 if !matches!(data, DataBlock::FixedSizeList(_)) {
964 return Ok(None);
965 }
966 let encoder: Box<dyn BlockCompressor> = Box::new(ValueEncoder::default());
967 match params.compression.as_deref() {
968 None | Some("none") => Ok(Some(encoder)),
969 Some(scheme) => Ok(Some(Box::new(GeneralBlockCompressor::new(
970 encoder,
971 CompressionConfig::new(scheme.parse()?, params.compression_level),
972 )))),
973 }
974}
975
976pub trait MiniBlockDecompressor: std::fmt::Debug + Send + Sync {
977 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock>;
978
979 fn num_buffers(&self) -> usize {
982 1
983 }
984
985 fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
992 None
993 }
994}
995
996pub trait FixedPerValueDecompressor: std::fmt::Debug + Send + Sync {
997 fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result<DataBlock>;
999 fn bits_per_value(&self) -> u64;
1003
1004 fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
1011 None
1012 }
1013}
1014
1015pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync {
1016 fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock>;
1018}
1019
1020pub trait BlockDecompressor: std::fmt::Debug + Send + Sync {
1021 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock>;
1022
1023 fn requires_payload(&self) -> bool {
1025 true
1026 }
1027
1028 fn infer_num_values(&self, _data: &LanceBuffer) -> Result<Option<u64>> {
1033 Ok(None)
1034 }
1035}
1036
1037pub(crate) fn require_block_payload(data: Option<LanceBuffer>, codec: &str) -> Result<LanceBuffer> {
1038 data.ok_or_else(|| Error::invalid_input(format!("{codec} requires one payload")))
1039}
1040
1041pub(crate) fn require_no_block_payload(data: Option<LanceBuffer>, codec: &str) -> Result<()> {
1042 if data.is_some() {
1043 return Err(Error::invalid_input(format!("{codec} expects no payload")));
1044 }
1045 Ok(())
1046}
1047
1048pub trait DecompressionStrategy: std::fmt::Debug + Send + Sync {
1049 fn create_miniblock_decompressor(
1050 &self,
1051 description: &CompressiveEncoding,
1052 decompression_strategy: &dyn DecompressionStrategy,
1053 ) -> Result<Box<dyn MiniBlockDecompressor>>;
1054
1055 fn create_fixed_per_value_decompressor(
1056 &self,
1057 description: &CompressiveEncoding,
1058 ) -> Result<Box<dyn FixedPerValueDecompressor>>;
1059
1060 fn create_variable_per_value_decompressor(
1061 &self,
1062 description: &CompressiveEncoding,
1063 ) -> Result<Box<dyn VariablePerValueDecompressor>>;
1064
1065 fn create_block_decompressor(
1066 &self,
1067 description: &CompressiveEncoding,
1068 ) -> Result<Box<dyn BlockDecompressor>>;
1069}
1070
1071#[derive(Debug, Default)]
1072pub struct DefaultDecompressionStrategy {}
1073
1074impl DecompressionStrategy for DefaultDecompressionStrategy {
1075 fn create_miniblock_decompressor(
1076 &self,
1077 description: &CompressiveEncoding,
1078 decompression_strategy: &dyn DecompressionStrategy,
1079 ) -> Result<Box<dyn MiniBlockDecompressor>> {
1080 let compression = description.compression.as_ref().ok_or_else(|| {
1081 Error::invalid_input("Mini-block encoding is missing its compression variant")
1082 })?;
1083 match compression {
1084 Compression::Flat(flat) => {
1085 if flat.bits_per_value == 0 || flat.data.is_some() {
1086 return Err(Error::invalid_input(
1087 "Flat mini-block requires a positive bit width and no leaf compression",
1088 ));
1089 }
1090 Ok(Box::new(ValueDecompressor::from_flat(flat)))
1091 }
1092 #[cfg(feature = "bitpacking")]
1093 Compression::InlineBitpacking(description) => {
1094 if !matches!(description.uncompressed_bits_per_value, 8 | 16 | 32 | 64)
1095 || description.values.is_some()
1096 {
1097 return Err(Error::invalid_input(
1098 "Inline bitpacking requires an 8, 16, 32, or 64-bit width and no leaf compression",
1099 ));
1100 }
1101 Ok(Box::new(InlineBitpacking::from_description(description)))
1102 }
1103 #[cfg(not(feature = "bitpacking"))]
1104 Compression::InlineBitpacking(_) => Err(Error::not_supported_source(
1105 "this runtime was not built with bitpacking support".into(),
1106 )),
1107 Compression::Variable(variable) => Ok(Box::new(
1108 BinaryMiniBlockDecompressor::from_variable(variable)?,
1109 )),
1110 Compression::Fsst(description) => {
1111 if description.symbol_table.is_empty()
1112 || !matches!(
1113 description
1114 .values
1115 .as_deref()
1116 .and_then(|v| v.compression.as_ref()),
1117 Some(Compression::Variable(_))
1118 )
1119 {
1120 return Err(Error::invalid_input(
1121 "FSST requires a symbol table and Variable values",
1122 ));
1123 }
1124 let inner_decompressor = decompression_strategy.create_miniblock_decompressor(
1125 description.values.as_ref().ok_or_else(|| {
1126 Error::invalid_input("FSST mini-block encoding is missing its values codec")
1127 })?,
1128 decompression_strategy,
1129 )?;
1130 Ok(Box::new(FsstMiniBlockDecompressor::new(
1131 description,
1132 inner_decompressor,
1133 )))
1134 }
1135 Compression::PackedStruct(description) => Ok(Box::new(
1136 PackedStructFixedWidthMiniBlockDecompressor::try_new(description)?,
1137 )),
1138 Compression::VariablePackedStruct(_) => Err(Error::not_supported_source(
1139 "variable packed struct decoding is not yet implemented".into(),
1140 )),
1141 Compression::FixedSizeList(fsl) => {
1142 Ok(Box::new(ValueDecompressor::from_fsl(fsl)?))
1145 }
1146 Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(
1147 rle,
1148 decompression_strategy,
1149 )?)),
1150 Compression::ByteStreamSplit(bss) => {
1151 let values = bss.values.as_ref().ok_or_else(|| {
1152 Error::invalid_input("ByteStreamSplit encoding is missing its values codec")
1153 })?;
1154 let Compression::Flat(values) = values.compression.as_ref().ok_or_else(|| {
1155 Error::invalid_input(
1156 "ByteStreamSplit values codec is missing its compression variant",
1157 )
1158 })?
1159 else {
1160 return Err(Error::invalid_input(
1161 "ByteStreamSplit compression only supports Flat values",
1162 ));
1163 };
1164 if !matches!(values.bits_per_value, 32 | 64) || values.data.is_some() {
1165 return Err(Error::invalid_input(
1166 "ByteStreamSplit requires plain 32 or 64-bit Flat values",
1167 ));
1168 }
1169 Ok(Box::new(ByteStreamSplitDecompressor::new(
1170 values.bits_per_value as usize,
1171 )))
1172 }
1173 Compression::General(general) => {
1174 let inner_decompressor = self.create_miniblock_decompressor(
1176 general.values.as_ref().ok_or_else(|| {
1177 Error::invalid_input("GeneralMiniBlock missing inner encoding")
1178 })?,
1179 decompression_strategy,
1180 )?;
1181
1182 let compression = general.compression.as_ref().ok_or_else(|| {
1184 Error::invalid_input("GeneralMiniBlock missing compression config")
1185 })?;
1186
1187 let scheme = compression.scheme().try_into()?;
1188
1189 let compression_config = CompressionConfig::new(scheme, compression.level);
1190
1191 Ok(Box::new(GeneralMiniBlockDecompressor::new(
1192 inner_decompressor,
1193 compression_config,
1194 )))
1195 }
1196 compression @ (Compression::Range(_) | Compression::Delta(_)) => {
1197 Err(Error::not_supported_source(
1198 format!(
1199 "{} compression is only supported in block positions",
1200 compression_name(compression)
1201 )
1202 .into(),
1203 ))
1204 }
1205 other => Err(Error::not_supported_source(
1206 format!(
1207 "{} is not supported for mini-block decompression",
1208 compression_name(other)
1209 )
1210 .into(),
1211 )),
1212 }
1213 }
1214
1215 fn create_fixed_per_value_decompressor(
1216 &self,
1217 description: &CompressiveEncoding,
1218 ) -> Result<Box<dyn FixedPerValueDecompressor>> {
1219 let compression = description.compression.as_ref().ok_or_else(|| {
1220 Error::invalid_input("Fixed per-value encoding is missing its compression variant")
1221 })?;
1222 match compression {
1223 Compression::Constant(constant) => Ok(Box::new(ConstantDecompressor::new(
1224 constant
1225 .value
1226 .as_ref()
1227 .map(|v| LanceBuffer::from_bytes(v.clone(), 1)),
1228 ))),
1229 Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))),
1230 Compression::FixedSizeList(fsl) => Ok(Box::new(ValueDecompressor::from_fsl(fsl)?)),
1231 Compression::PackedStruct(description) => Ok(Box::new(
1232 PackedStructFixedPerValueDecompressor::new(description)?,
1233 )),
1234 other => Err(Error::not_supported_source(
1235 format!(
1236 "{} is not supported for fixed per-value decompression",
1237 compression_name(other)
1238 )
1239 .into(),
1240 )),
1241 }
1242 }
1243
1244 fn create_variable_per_value_decompressor(
1245 &self,
1246 description: &CompressiveEncoding,
1247 ) -> Result<Box<dyn VariablePerValueDecompressor>> {
1248 let compression = description.compression.as_ref().ok_or_else(|| {
1249 Error::invalid_input("Variable per-value encoding is missing its compression variant")
1250 })?;
1251 match compression {
1252 Compression::Variable(variable) => {
1253 let offsets = variable.offsets.as_ref().ok_or_else(|| {
1254 Error::invalid_input("Variable per-value encoding is missing offsets")
1255 })?;
1256 let Compression::Flat(offsets) = offsets.compression.as_ref().ok_or_else(|| {
1257 Error::invalid_input(
1258 "Variable per-value offsets are missing a compression variant",
1259 )
1260 })?
1261 else {
1262 return Err(Error::invalid_input(
1263 "Variable per-value compression only supports Flat offsets",
1264 ));
1265 };
1266 if offsets.bits_per_value >= u8::MAX as u64 {
1267 return Err(Error::invalid_input(format!(
1268 "Variable per-value offset width {} does not fit u8",
1269 offsets.bits_per_value
1270 )));
1271 }
1272 Ok(Box::new(VariableDecoder::default()))
1273 }
1274 Compression::Fsst(fsst) => Ok(Box::new(FsstPerValueDecompressor::new(
1275 LanceBuffer::from_bytes(fsst.symbol_table.clone(), 1),
1276 Box::new(VariableDecoder::default()),
1277 ))),
1278 Compression::General(general) => Ok(Box::new(CompressedBufferEncoder::from_scheme(
1279 general.compression.as_ref().expect_ok()?.scheme(),
1280 )?)),
1281 Compression::VariablePackedStruct(description) => {
1282 let mut fields = Vec::with_capacity(description.fields.len());
1283 for field in &description.fields {
1284 let value_encoding = field.value.as_ref().ok_or_else(|| {
1285 Error::invalid_input("VariablePackedStruct field is missing value encoding")
1286 })?;
1287 let decoder = match field.layout.as_ref().ok_or_else(|| {
1288 Error::invalid_input("VariablePackedStruct field is missing layout details")
1289 })? {
1290 crate::format::metadata::variable_packed_struct::field_encoding::Layout::BitsPerValue(
1291 bits_per_value,
1292 ) => {
1293 let decompressor =
1294 self.create_fixed_per_value_decompressor(value_encoding)?;
1295 VariablePackedStructFieldDecoder {
1296 kind: VariablePackedStructFieldKind::Fixed {
1297 bits_per_value: *bits_per_value,
1298 decompressor: Arc::from(decompressor),
1299 },
1300 }
1301 }
1302 crate::format::metadata::variable_packed_struct::field_encoding::Layout::BitsPerLength(
1303 bits_per_length,
1304 ) => {
1305 let decompressor =
1306 self.create_variable_per_value_decompressor(value_encoding)?;
1307 VariablePackedStructFieldDecoder {
1308 kind: VariablePackedStructFieldKind::Variable {
1309 bits_per_length: *bits_per_length,
1310 decompressor: Arc::from(decompressor),
1311 },
1312 }
1313 }
1314 };
1315 fields.push(decoder);
1316 }
1317 Ok(Box::new(PackedStructVariablePerValueDecompressor::new(
1318 fields,
1319 )))
1320 }
1321 other => Err(Error::not_supported_source(
1322 format!(
1323 "{} is not supported for variable per-value decompression",
1324 compression_name(other)
1325 )
1326 .into(),
1327 )),
1328 }
1329 }
1330
1331 fn create_block_decompressor(
1332 &self,
1333 description: &CompressiveEncoding,
1334 ) -> Result<Box<dyn BlockDecompressor>> {
1335 let compression = description.compression.as_ref().ok_or_else(|| {
1336 Error::invalid_input("Block encoding is missing its compression variant")
1337 })?;
1338 match compression {
1339 Compression::InlineBitpacking(inline_bitpacking) => Ok(Box::new(
1340 InlineBitpacking::from_description(inline_bitpacking),
1341 )),
1342 Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))),
1343 Compression::Constant(constant) => {
1344 let scalar = constant
1345 .value
1346 .as_ref()
1347 .map(|v| LanceBuffer::from_bytes(v.clone(), 1));
1348 Ok(Box::new(ConstantDecompressor::new(scalar)))
1349 }
1350 Compression::Variable(_) => Ok(Box::new(BinaryBlockDecompressor::default())),
1351 Compression::FixedSizeList(fsl) => {
1352 Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref())?))
1353 }
1354 Compression::OutOfLineBitpacking(out_of_line) => {
1355 let values = out_of_line.values.as_ref().ok_or_else(|| {
1356 Error::invalid_input("OutOfLineBitpacking is missing its values encoding")
1357 })?;
1358 let compressed_bit_width = match values.compression.as_ref().ok_or_else(|| {
1359 Error::invalid_input(
1360 "OutOfLineBitpacking values are missing a compression variant",
1361 )
1362 })? {
1363 Compression::Flat(flat) => flat.bits_per_value,
1364 _ => {
1365 return Err(Error::invalid_input_source(
1366 "OutOfLineBitpacking values must use Flat encoding".into(),
1367 ));
1368 }
1369 };
1370 Ok(Box::new(OutOfLineBitpacking::new(
1371 compressed_bit_width,
1372 out_of_line.uncompressed_bits_per_value,
1373 )))
1374 }
1375 Compression::General(general) => {
1376 let inner_desc = general
1377 .values
1378 .as_ref()
1379 .ok_or_else(|| {
1380 Error::invalid_input("General compression missing inner encoding")
1381 })?
1382 .as_ref();
1383 let inner_decompressor = self.create_block_decompressor(inner_desc)?;
1384
1385 let compression = general.compression.as_ref().ok_or_else(|| {
1386 Error::invalid_input("General compression missing compression config")
1387 })?;
1388 let scheme = compression.scheme().try_into()?;
1389 let config = CompressionConfig::new(scheme, compression.level);
1390 if let Some(Compression::FixedSizeList(fsl)) = &inner_desc.compression {
1391 let decoder = ValueDecompressor::from_fsl(fsl)?;
1392 return Ok(Box::new(FixedWidthGeneralBlockDecompressor::new(
1393 inner_decompressor,
1394 config,
1395 FixedPerValueDecompressor::bits_per_value(&decoder),
1396 )));
1397 }
1398 let general_decompressor =
1399 GeneralBlockDecompressor::try_new(inner_decompressor, config)?;
1400
1401 Ok(Box::new(general_decompressor))
1402 }
1403 Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)),
1404 Compression::Range(range) => create_fixed_width_block_decompressor(
1405 description,
1406 range.uncompressed_bits_per_value,
1407 ),
1408 Compression::Delta(delta) => create_fixed_width_block_decompressor(
1409 description,
1410 delta.uncompressed_bits_per_value,
1411 ),
1412 other => Err(Error::not_supported_source(
1413 format!(
1414 "{} is not supported for block decompression",
1415 compression_name(other)
1416 )
1417 .into(),
1418 )),
1419 }
1420 }
1421}
1422
1423pub(crate) fn infer_fixed_width_block_bits(description: &CompressiveEncoding) -> Result<u64> {
1424 fn infer(description: &CompressiveEncoding, depth: u8) -> Result<u64> {
1425 if depth > 2 {
1426 return Err(Error::invalid_input(
1427 "Generic block encoding exceeds the supported nesting depth",
1428 ));
1429 }
1430 let compression = description.compression.as_ref().ok_or_else(|| {
1431 Error::invalid_input("Block encoding is missing its compression variant")
1432 })?;
1433 match compression {
1434 Compression::Flat(flat) => Ok(flat.bits_per_value),
1435 Compression::Constant(constant) => constant
1436 .value
1437 .as_ref()
1438 .map(|value| value.len() as u64 * 8)
1439 .ok_or_else(|| Error::invalid_input("Typed Constant is missing its scalar value")),
1440 Compression::Range(range) => Ok(range.uncompressed_bits_per_value),
1441 Compression::Delta(delta) => Ok(delta.uncompressed_bits_per_value),
1442 Compression::InlineBitpacking(bitpacking) => Ok(bitpacking.uncompressed_bits_per_value),
1443 Compression::OutOfLineBitpacking(bitpacking) => {
1444 Ok(bitpacking.uncompressed_bits_per_value)
1445 }
1446 Compression::General(general) => infer(
1447 general.values.as_deref().ok_or_else(|| {
1448 Error::invalid_input("General block compression is missing its values encoding")
1449 })?,
1450 depth + 1,
1451 ),
1452 Compression::Rle(rle) => infer(
1453 rle.values.as_deref().ok_or_else(|| {
1454 Error::invalid_input("RLE compression is missing its values encoding")
1455 })?,
1456 depth + 1,
1457 ),
1458 Compression::Dictionary(dictionary) => infer(
1459 dictionary.items.as_deref().ok_or_else(|| {
1460 Error::invalid_input("Dictionary is missing its items encoding")
1461 })?,
1462 depth + 1,
1463 ),
1464 other => Err(Error::invalid_input(format!(
1465 "Cannot infer fixed-width output from {}",
1466 compression_name(other)
1467 ))),
1468 }
1469 }
1470
1471 let bits = infer(description, 0)?;
1472 if !matches!(bits, 32 | 64) {
1473 return Err(Error::invalid_input(format!(
1474 "Generic block offsets require 32 or 64-bit values, got {bits}"
1475 )));
1476 }
1477 Ok(bits)
1478}
1479
1480pub(crate) fn create_fixed_width_block_decompressor(
1486 description: &CompressiveEncoding,
1487 expected_bits_per_value: u64,
1488) -> Result<Box<dyn BlockDecompressor>> {
1489 if !matches!(expected_bits_per_value, 32 | 64) {
1490 return Err(Error::invalid_input(format!(
1491 "Generic block decoding only supports 32 or 64-bit values, got {expected_bits_per_value}"
1492 )));
1493 }
1494 create_fixed_width_block_decompressor_inner(description, expected_bits_per_value, true)
1495}
1496
1497pub fn validate_fixed_width_block_sequence(
1502 description: &CompressiveEncoding,
1503 bits_per_value: u64,
1504 num_values: u64,
1505) -> Result<()> {
1506 crate::encodings::physical::sequence::SequenceDecoder::try_new(
1507 description,
1508 bits_per_value,
1509 num_values,
1510 )
1511 .map(|_| ())
1512}
1513
1514fn create_fixed_width_block_decompressor_inner(
1515 description: &CompressiveEncoding,
1516 expected_bits_per_value: u64,
1517 allow_composite: bool,
1518) -> Result<Box<dyn BlockDecompressor>> {
1519 let compression = description
1520 .compression
1521 .as_ref()
1522 .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?;
1523 match compression {
1524 Compression::Flat(flat) => {
1525 if flat.bits_per_value != expected_bits_per_value {
1526 return Err(Error::invalid_input(format!(
1527 "Flat declares {}-bit values, expected {expected_bits_per_value}",
1528 flat.bits_per_value
1529 )));
1530 }
1531 if flat.data.is_some() {
1532 return Err(Error::invalid_input(
1533 "Generic block Flat cannot contain buffer compression",
1534 ));
1535 }
1536 Ok(Box::new(ValueDecompressor::from_flat(flat)))
1537 }
1538 Compression::Constant(constant) => {
1539 let value = constant
1540 .value
1541 .as_ref()
1542 .map(|_| decode_fixed_width_constant(constant, expected_bits_per_value))
1543 .transpose()?;
1544 Ok(Box::new(ConstantBlockDecompressor::new(
1545 expected_bits_per_value,
1546 value,
1547 )))
1548 }
1549 Compression::Range(range) => {
1550 if range.uncompressed_bits_per_value != expected_bits_per_value {
1551 return Err(Error::invalid_input(format!(
1552 "Range declares {}-bit values, expected {expected_bits_per_value}",
1553 range.uncompressed_bits_per_value
1554 )));
1555 }
1556 Ok(Box::new(RangeDecompressor::new(
1557 expected_bits_per_value,
1558 range.start,
1559 range.step,
1560 )))
1561 }
1562 Compression::InlineBitpacking(bitpacking) => {
1563 if bitpacking.uncompressed_bits_per_value != expected_bits_per_value {
1564 return Err(Error::invalid_input(format!(
1565 "Inline bitpacking declares {}-bit values, expected {expected_bits_per_value}",
1566 bitpacking.uncompressed_bits_per_value
1567 )));
1568 }
1569 if bitpacking.values.is_some() {
1570 return Err(Error::invalid_input(
1571 "Generic Inline bitpacking cannot contain buffer compression",
1572 ));
1573 }
1574 #[cfg(feature = "bitpacking")]
1575 {
1576 Ok(Box::new(InlineBitpacking::from_description(bitpacking)))
1577 }
1578 #[cfg(not(feature = "bitpacking"))]
1579 {
1580 Err(Error::not_supported_source(
1581 "this runtime was not built with bitpacking support".into(),
1582 ))
1583 }
1584 }
1585 Compression::OutOfLineBitpacking(bitpacking) => {
1586 if bitpacking.uncompressed_bits_per_value != expected_bits_per_value {
1587 return Err(Error::invalid_input(format!(
1588 "Out-of-line bitpacking declares {}-bit values, expected {expected_bits_per_value}",
1589 bitpacking.uncompressed_bits_per_value
1590 )));
1591 }
1592 let values = bitpacking.values.as_deref().ok_or_else(|| {
1593 Error::invalid_input("Out-of-line bitpacking is missing its values encoding")
1594 })?;
1595 let compressed_bits = match values.compression.as_ref() {
1596 Some(Compression::Flat(flat)) if flat.data.is_none() => flat.bits_per_value,
1597 _ => {
1598 return Err(Error::invalid_input(
1599 "Out-of-line bitpacking values must use plain Flat encoding",
1600 ));
1601 }
1602 };
1603 if compressed_bits == 0 || compressed_bits >= expected_bits_per_value {
1604 return Err(Error::invalid_input(format!(
1605 "Out-of-line bitpacking width {compressed_bits} is invalid for {expected_bits_per_value}-bit values"
1606 )));
1607 }
1608 #[cfg(feature = "bitpacking")]
1609 {
1610 Ok(Box::new(OutOfLineBitpacking::new(
1611 compressed_bits,
1612 expected_bits_per_value,
1613 )))
1614 }
1615 #[cfg(not(feature = "bitpacking"))]
1616 {
1617 Err(Error::not_supported_source(
1618 "this runtime was not built with bitpacking support".into(),
1619 ))
1620 }
1621 }
1622 Compression::Delta(delta) if allow_composite => {
1623 if delta.uncompressed_bits_per_value != expected_bits_per_value {
1624 return Err(Error::invalid_input(format!(
1625 "Delta declares {}-bit values, expected {expected_bits_per_value}",
1626 delta.uncompressed_bits_per_value
1627 )));
1628 }
1629 let child = delta.deltas.as_deref().ok_or_else(|| {
1630 Error::invalid_input("Delta is missing its deltas child encoding")
1631 })?;
1632 let child =
1633 create_fixed_width_block_decompressor_inner(child, expected_bits_per_value, false)?;
1634 Ok(Box::new(DeltaDecompressor::new(
1635 expected_bits_per_value,
1636 delta.base,
1637 child,
1638 )))
1639 }
1640 Compression::General(general) if allow_composite => {
1641 let child = general.values.as_deref().ok_or_else(|| {
1642 Error::invalid_input("General block compression is missing its values encoding")
1643 })?;
1644 if !matches!(
1645 child.compression.as_ref(),
1646 Some(Compression::Flat(flat))
1647 if flat.bits_per_value == expected_bits_per_value && flat.data.is_none()
1648 ) {
1649 return Err(Error::invalid_input(
1650 "Generic General compression must wrap plain Flat values",
1651 ));
1652 }
1653 let child =
1654 create_fixed_width_block_decompressor_inner(child, expected_bits_per_value, false)?;
1655 let compression = general.compression.as_ref().ok_or_else(|| {
1656 Error::invalid_input("General block compression is missing its compression config")
1657 })?;
1658 let scheme = compression.scheme().try_into()?;
1659 let config = CompressionConfig::new(scheme, compression.level);
1660 Ok(Box::new(FixedWidthGeneralBlockDecompressor::new(
1661 child,
1662 config,
1663 expected_bits_per_value,
1664 )))
1665 }
1666 Compression::Rle(rle) if allow_composite => {
1667 let values = rle.values.as_deref().ok_or_else(|| {
1668 Error::invalid_input("RLE compression is missing its values encoding")
1669 })?;
1670 let run_lengths = rle.run_lengths.as_deref().ok_or_else(|| {
1671 Error::invalid_input("RLE compression is missing its run-length encoding")
1672 })?;
1673 let run_count = match run_lengths.compression.as_ref() {
1674 Some(Compression::Constant(constant)) => {
1675 BlockRleRunCount::ConstantRunLength(decode_fixed_width_constant(constant, 32)?)
1676 }
1677 Some(Compression::Range(range)) if range.uncompressed_bits_per_value == 32 => {
1678 BlockRleRunCount::RangeRunLengths {
1679 start: range.start,
1680 step: range.step,
1681 }
1682 }
1683 _ if matches!(values.compression.as_ref(), Some(Compression::Flat(flat)) if flat.bits_per_value == expected_bits_per_value && flat.data.is_none()) => {
1684 BlockRleRunCount::FlatValues
1685 }
1686 _ if matches!(run_lengths.compression.as_ref(), Some(Compression::Flat(flat)) if flat.bits_per_value == 32 && flat.data.is_none()) => {
1687 BlockRleRunCount::FlatRunLengths
1688 }
1689 _ => {
1690 return Err(Error::invalid_input(
1691 "RLE requires metadata run lengths or one plain Flat child to recover the run count",
1692 ));
1693 }
1694 };
1695 let values = create_fixed_width_block_decompressor_inner(
1696 values,
1697 expected_bits_per_value,
1698 false,
1699 )?;
1700 let run_lengths = create_fixed_width_block_decompressor_inner(run_lengths, 32, false)?;
1701 Ok(Box::new(BlockRleDecompressor::try_new(
1702 expected_bits_per_value,
1703 values,
1704 run_lengths,
1705 run_count,
1706 )?))
1707 }
1708 Compression::Dictionary(dictionary) if allow_composite => {
1709 let indices = dictionary.indices.as_deref().ok_or_else(|| {
1710 Error::invalid_input("Dictionary is missing its indices encoding")
1711 })?;
1712 let items = dictionary
1713 .items
1714 .as_deref()
1715 .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?;
1716 let indices = create_fixed_width_block_decompressor_inner(indices, 32, false)?;
1717 let items =
1718 create_fixed_width_block_decompressor_inner(items, expected_bits_per_value, false)?;
1719 Ok(Box::new(BlockDictionaryDecompressor::try_new(
1720 expected_bits_per_value,
1721 dictionary.num_dictionary_items,
1722 indices,
1723 items,
1724 )?))
1725 }
1726 other => Err(Error::invalid_input(format!(
1727 "{} is not allowed at this position in a generic block encoding",
1728 compression_name(other)
1729 ))),
1730 }
1731}
1732
1733fn decode_fixed_width_constant(
1734 constant: &crate::format::metadata::Constant,
1735 bits_per_value: u64,
1736) -> Result<u64> {
1737 let value = constant
1738 .value
1739 .as_ref()
1740 .ok_or_else(|| Error::invalid_input("Typed Constant is missing its scalar value"))?;
1741 let expected_bytes = usize::try_from(bits_per_value / 8)
1742 .map_err(|_| Error::invalid_input("Constant scalar width does not fit usize"))?;
1743 if value.len() != expected_bytes {
1744 return Err(Error::invalid_input(format!(
1745 "Constant scalar has {} bytes, expected {expected_bytes}",
1746 value.len()
1747 )));
1748 }
1749 Ok(match bits_per_value {
1750 32 => u64::from(u32::from_le_bytes(
1751 value
1752 .as_ref()
1753 .try_into()
1754 .expect("constant width was checked"),
1755 )),
1756 64 => u64::from_le_bytes(
1757 value
1758 .as_ref()
1759 .try_into()
1760 .expect("constant width was checked"),
1761 ),
1762 _ => {
1763 return Err(Error::invalid_input(format!(
1764 "Constant only supports 32 or 64-bit values, got {bits_per_value}"
1765 )));
1766 }
1767 })
1768}
1769
1770pub(crate) fn validate_delta_child_encoding(
1771 child: &CompressiveEncoding,
1772 expected_bits_per_value: u64,
1773) -> Result<()> {
1774 let child_bits_per_value = match child.compression.as_ref() {
1775 Some(Compression::Flat(flat)) => flat.bits_per_value,
1776 Some(Compression::Constant(constant)) => {
1777 decode_fixed_width_constant(constant, expected_bits_per_value)?;
1778 expected_bits_per_value
1779 }
1780 Some(Compression::Range(range)) => range.uncompressed_bits_per_value,
1781 Some(Compression::InlineBitpacking(bitpacking)) => bitpacking.uncompressed_bits_per_value,
1782 Some(Compression::OutOfLineBitpacking(bitpacking)) => {
1783 bitpacking.uncompressed_bits_per_value
1784 }
1785 Some(other) => {
1786 return Err(Error::invalid_input(format!(
1787 "Delta does not support a {} child",
1788 compression_name(other)
1789 )));
1790 }
1791 None => {
1792 return Err(Error::invalid_input(
1793 "Delta child is missing its compression variant",
1794 ));
1795 }
1796 };
1797 if child_bits_per_value != expected_bits_per_value {
1798 return Err(Error::invalid_input(format!(
1799 "Delta child declares {child_bits_per_value}-bit values, expected {expected_bits_per_value}"
1800 )));
1801 }
1802 Ok(())
1803}
1804
1805pub(crate) fn create_rle_decompressor(
1806 rle: &crate::format::metadata::Rle,
1807 decompression_strategy: &dyn DecompressionStrategy,
1808) -> Result<RleDecompressor> {
1809 let values = rle
1810 .values
1811 .as_ref()
1812 .ok_or_else(|| Error::invalid_input("RLE compression missing values encoding"))?;
1813 let run_lengths = rle
1814 .run_lengths
1815 .as_ref()
1816 .ok_or_else(|| Error::invalid_input("RLE compression missing run lengths encoding"))?;
1817
1818 let values = create_rle_child_decompressor(values, "values", decompression_strategy)?;
1819 let run_lengths =
1820 create_rle_child_decompressor(run_lengths, "run lengths", decompression_strategy)?;
1821
1822 if !matches!(values.bits_per_value(), 8 | 16 | 32 | 64) {
1823 return Err(Error::invalid_input(format!(
1824 "RLE compression only supports 8, 16, 32, or 64-bit values, got {}",
1825 values.bits_per_value()
1826 )));
1827 }
1828
1829 let run_length_width =
1830 RunLengthWidth::from_bits(run_lengths.bits_per_value()).ok_or_else(|| {
1831 Error::invalid_input(format!(
1832 "RLE compression only supports 8, 16, or 32-bit run lengths, got {}",
1833 run_lengths.bits_per_value()
1834 ))
1835 })?;
1836
1837 if values.requires_num_values() && run_lengths.requires_num_values() {
1838 return Err(Error::invalid_input(
1839 "RLE values and run lengths child encodings cannot both require the run count",
1840 ));
1841 }
1842
1843 if values.is_identity() && run_lengths.is_identity() {
1844 return Ok(RleDecompressor::with_run_length_width(
1845 values.bits_per_value(),
1846 run_length_width,
1847 ));
1848 }
1849
1850 Ok(RleDecompressor::with_child_decompressors(
1851 values.bits_per_value(),
1852 run_length_width,
1853 values,
1854 run_lengths,
1855 ))
1856}
1857
1858fn create_rle_child_decompressor(
1859 encoding: &CompressiveEncoding,
1860 role: &str,
1861 decompression_strategy: &dyn DecompressionStrategy,
1862) -> Result<RleChildDecompressor> {
1863 let compression = encoding
1864 .compression
1865 .as_ref()
1866 .ok_or_else(|| Error::invalid_input(format!("RLE {role} missing child compression")))?;
1867 let (bits_per_value, requires_num_values, needs_decompressor) =
1868 validate_rle_child_compression(compression, role)?;
1869
1870 if needs_decompressor {
1871 Ok(RleChildDecompressor::block(
1872 bits_per_value,
1873 decompression_strategy.create_block_decompressor(encoding)?,
1874 requires_num_values,
1875 ))
1876 } else {
1877 Ok(RleChildDecompressor::flat(bits_per_value))
1878 }
1879}
1880
1881fn validate_rle_child_compression(
1882 compression: &Compression,
1883 role: &str,
1884) -> Result<(u64, bool, bool)> {
1885 match compression {
1886 Compression::Flat(flat) => Ok((flat.bits_per_value, false, false)),
1887 Compression::General(general) => {
1888 general.compression.as_ref().ok_or_else(|| {
1889 Error::invalid_input(format!(
1890 "RLE {role} general child missing compression config"
1891 ))
1892 })?;
1893 let values = general.values.as_ref().ok_or_else(|| {
1894 Error::invalid_input(format!("RLE {role} general child missing inner encoding"))
1895 })?;
1896 let inner = values.compression.as_ref().ok_or_else(|| {
1897 Error::invalid_input(format!(
1898 "RLE {role} general child missing inner compression"
1899 ))
1900 })?;
1901 let (bits_per_value, requires_num_values) =
1902 validate_rle_block_child_inner(inner, role)?;
1903 Ok((bits_per_value, requires_num_values, true))
1904 }
1905 Compression::OutOfLineBitpacking(out_of_line) => {
1906 let values = out_of_line.values.as_ref().ok_or_else(|| {
1907 Error::invalid_input(format!(
1908 "RLE {role} bitpacking child missing values encoding"
1909 ))
1910 })?;
1911 let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| {
1912 Error::invalid_input(format!(
1913 "RLE {role} bitpacking child missing values compression"
1914 ))
1915 })?
1916 else {
1917 return Err(Error::invalid_input(format!(
1918 "RLE {role} bitpacking child only supports flat values"
1919 )));
1920 };
1921 Ok((out_of_line.uncompressed_bits_per_value, true, true))
1922 }
1923 other => Err(Error::invalid_input(format!(
1924 "RLE {role} only supports flat, general, or out-of-line bitpacking child encodings, got {}",
1925 compression_name(other)
1926 ))),
1927 }
1928}
1929
1930fn validate_rle_block_child_inner(compression: &Compression, role: &str) -> Result<(u64, bool)> {
1931 match compression {
1932 Compression::Flat(flat) => Ok((flat.bits_per_value, false)),
1933 Compression::OutOfLineBitpacking(out_of_line) => {
1934 let values = out_of_line.values.as_ref().ok_or_else(|| {
1935 Error::invalid_input(format!(
1936 "RLE {role} bitpacking child missing values encoding"
1937 ))
1938 })?;
1939 let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| {
1940 Error::invalid_input(format!(
1941 "RLE {role} bitpacking child missing values compression"
1942 ))
1943 })?
1944 else {
1945 return Err(Error::invalid_input(format!(
1946 "RLE {role} bitpacking child only supports flat values"
1947 )));
1948 };
1949 Ok((out_of_line.uncompressed_bits_per_value, true))
1950 }
1951 other => Err(Error::invalid_input(format!(
1952 "RLE {role} general child only supports flat or out-of-line bitpacking inner encodings, got {}",
1953 compression_name(other)
1954 ))),
1955 }
1956}
1957
1958fn compression_name(compression: &Compression) -> &'static str {
1959 match compression {
1960 Compression::Flat(_) => "flat",
1961 Compression::Variable(_) => "variable",
1962 Compression::Fsst(_) => "fsst",
1963 Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking",
1964 Compression::InlineBitpacking(_) => "inline bitpacking",
1965 Compression::General(_) => "general",
1966 Compression::Constant(_) => "constant",
1967 Compression::Dictionary(_) => "dictionary",
1968 Compression::ByteStreamSplit(_) => "byte stream split",
1969 Compression::PackedStruct(_) => "packed struct",
1970 Compression::FixedSizeList(_) => "fixed-size list",
1971 Compression::VariablePackedStruct(_) => "variable packed struct",
1972 Compression::Rle(_) => "rle",
1973 Compression::Range(_) => "range",
1974 Compression::Delta(_) => "delta",
1975 }
1976}
1977
1978#[cfg(test)]
1979mod tests {
1980 use super::*;
1981 use crate::buffer::LanceBuffer;
1982 use crate::compression_config::CompressionParams;
1983 use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock};
1984 use crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext;
1985 use crate::format::EncodingUtils;
1986 use crate::statistics::ComputeStat;
1987 use crate::testing::{TestEncoding, extract_array_encoding_chain, test_compression_strategy};
1988 use arrow_schema::{DataType, Field as ArrowField};
1989 use std::collections::HashMap;
1990
1991 #[rstest::rstest]
1992 #[case(EncodingUtils::flat(0, None), "positive bit width")]
1993 #[case(
1994 EncodingUtils::fsl(0, false, EncodingUtils::flat(32, None)),
1995 "dimension must be positive"
1996 )]
1997 #[case(
1998 EncodingUtils::fsl(u64::MAX, false, EncodingUtils::flat(32, None)),
1999 "overflows"
2000 )]
2001 #[case(EncodingUtils::packed_struct(EncodingUtils::flat(32, None), vec![16]), "must match")]
2002 #[case(EncodingUtils::packed_struct(EncodingUtils::flat(1, None), vec![1]), "positive byte widths")]
2003 fn mini_block_descriptor_errors_are_owned_by_codecs(
2004 #[case] encoding: CompressiveEncoding,
2005 #[case] message: &str,
2006 ) {
2007 let strategy = DefaultDecompressionStrategy::default();
2008 let error = strategy
2009 .create_miniblock_decompressor(&encoding, &strategy)
2010 .unwrap_err();
2011 assert!(matches!(error, Error::InvalidInput { .. }));
2012 assert!(error.to_string().contains(message), "{error}");
2013 }
2014
2015 #[rstest::rstest]
2016 #[case(EncodingUtils::flat(32, None), vec![], 1, "buffers")]
2017 #[case(EncodingUtils::flat(32, None), vec![vec![0; 3]], 1, "expected 4")]
2018 #[case(EncodingUtils::fsl(2, true, EncodingUtils::flat(32, None)), vec![vec![], vec![0; 8]], 1, "expected 1")]
2019 #[case(EncodingUtils::fsl(2, false, EncodingUtils::flat(32, None)), vec![vec![]], u64::MAX, "overflows")]
2020 fn mini_block_payload_errors_are_owned_by_codecs(
2021 #[case] encoding: CompressiveEncoding,
2022 #[case] buffers: Vec<Vec<u8>>,
2023 #[case] count: u64,
2024 #[case] message: &str,
2025 ) {
2026 let strategy = DefaultDecompressionStrategy::default();
2027 let decoder = strategy
2028 .create_miniblock_decompressor(&encoding, &strategy)
2029 .unwrap();
2030 let error = decoder
2031 .decompress(buffers.into_iter().map(LanceBuffer::from).collect(), count)
2032 .unwrap_err();
2033 assert!(matches!(error, Error::InvalidInput { .. }));
2034 assert!(error.to_string().contains(message), "{error}");
2035 }
2036 fn strategy(encoding: TestEncoding, params: CompressionParams) -> Arc<dyn CompressionStrategy> {
2037 test_compression_strategy(encoding, params)
2038 }
2039
2040 fn baseline_strategy(params: CompressionParams) -> Arc<dyn CompressionStrategy> {
2041 strategy(TestEncoding::StructuralU16, params)
2042 }
2043
2044 fn selected_block_codec(
2045 strategy: &Arc<dyn CompressionStrategy>,
2046 field: &Field,
2047 data: &DataBlock,
2048 ) -> (Box<dyn BlockCompressor>, CompressiveEncoding) {
2049 let compressor = strategy.create_block_compressor(field, data).unwrap();
2050 let (_, encoding) = compressor.compress(data.clone()).unwrap();
2051 (compressor, encoding)
2052 }
2053
2054 fn miniblock_context() -> MiniBlockCompressionContext {
2055 MiniBlockCompressionContext::new(0, true, true)
2056 }
2057
2058 fn create_test_field(name: &str, data_type: DataType) -> Field {
2059 let arrow_field = ArrowField::new(name, data_type, true);
2060 let mut field = Field::try_from(&arrow_field).unwrap();
2061 field.id = -1;
2062 field
2063 }
2064
2065 fn create_fixed_width_block_with_stats(
2066 bits_per_value: u64,
2067 num_values: u64,
2068 run_count: u64,
2069 ) -> DataBlock {
2070 let bytes_per_value = (bits_per_value / 8) as usize;
2072 let total_bytes = bytes_per_value * num_values as usize;
2073 let mut data = vec![0u8; total_bytes];
2074
2075 let values_per_run = (num_values / run_count).max(1);
2077 let mut run_value = 0u8;
2078
2079 for i in 0..num_values as usize {
2080 if i % values_per_run as usize == 0 {
2081 run_value = run_value.wrapping_add(17); }
2083 for j in 0..bytes_per_value {
2085 let byte_offset = i * bytes_per_value + j;
2086 if byte_offset < data.len() {
2087 data[byte_offset] = run_value.wrapping_add(j as u8);
2088 }
2089 }
2090 }
2091
2092 let mut block = FixedWidthDataBlock {
2093 bits_per_value,
2094 data: LanceBuffer::reinterpret_vec(data),
2095 num_values,
2096 block_info: BlockInfo::default(),
2097 };
2098
2099 use crate::statistics::ComputeStat;
2101 block.compute_stat();
2102
2103 DataBlock::FixedWidth(block)
2104 }
2105
2106 fn create_fixed_width_block(bits_per_value: u64, num_values: u64) -> DataBlock {
2107 let bytes_per_value = (bits_per_value / 8) as usize;
2109 let total_bytes = bytes_per_value * num_values as usize;
2110 let mut data = vec![0u8; total_bytes];
2111
2112 for i in 0..num_values as usize {
2114 let byte_offset = i * bytes_per_value;
2115 if byte_offset < data.len() {
2116 data[byte_offset] = (i % 256) as u8;
2117 }
2118 }
2119
2120 let mut block = FixedWidthDataBlock {
2121 bits_per_value,
2122 data: LanceBuffer::reinterpret_vec(data),
2123 num_values,
2124 block_info: BlockInfo::default(),
2125 };
2126
2127 use crate::statistics::ComputeStat;
2129 block.compute_stat();
2130
2131 DataBlock::FixedWidth(block)
2132 }
2133
2134 fn rle_run_length_bits(encoding: &CompressiveEncoding) -> u64 {
2135 let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else {
2136 panic!("expected RLE encoding");
2137 };
2138 let Compression::Flat(run_lengths) = rle
2139 .run_lengths
2140 .as_ref()
2141 .unwrap()
2142 .compression
2143 .as_ref()
2144 .unwrap()
2145 else {
2146 panic!("expected flat run lengths");
2147 };
2148 run_lengths.bits_per_value
2149 }
2150
2151 fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::metadata::Rle {
2152 match encoding.compression.as_ref().unwrap() {
2153 Compression::Rle(rle) => rle,
2154 Compression::General(general) => {
2155 let inner = general.values.as_ref().unwrap();
2156 let Compression::Rle(rle) = inner.compression.as_ref().unwrap() else {
2157 panic!("expected wrapped RLE encoding");
2158 };
2159 rle
2160 }
2161 other => panic!("expected RLE encoding, got {}", compression_name(other)),
2162 }
2163 }
2164
2165 fn create_variable_width_block(
2166 bits_per_offset: u8,
2167 num_values: u64,
2168 avg_value_size: usize,
2169 ) -> DataBlock {
2170 use crate::statistics::ComputeStat;
2171
2172 let mut offsets = Vec::with_capacity((num_values + 1) as usize);
2174 let mut current_offset = 0i64;
2175 offsets.push(current_offset);
2176
2177 for i in 0..num_values {
2179 let value_size = if avg_value_size == 0 {
2180 1
2181 } else {
2182 ((avg_value_size as i64 + (i as i64 % 8) - 4).max(1) as usize)
2183 .min(avg_value_size * 2)
2184 };
2185 current_offset += value_size as i64;
2186 offsets.push(current_offset);
2187 }
2188
2189 let total_data_size = current_offset as usize;
2191 let mut data = vec![0u8; total_data_size];
2192
2193 for i in 0..num_values {
2195 let start_offset = offsets[i as usize] as usize;
2196 let end_offset = offsets[(i + 1) as usize] as usize;
2197
2198 let content = (i % 256) as u8;
2199 for j in 0..end_offset - start_offset {
2200 data[start_offset + j] = content.wrapping_add(j as u8);
2201 }
2202 }
2203
2204 let offsets_buffer = match bits_per_offset {
2206 32 => {
2207 let offsets_32: Vec<i32> = offsets.iter().map(|&o| o as i32).collect();
2208 LanceBuffer::reinterpret_vec(offsets_32)
2209 }
2210 64 => LanceBuffer::reinterpret_vec(offsets),
2211 _ => panic!("Unsupported bits_per_offset: {}", bits_per_offset),
2212 };
2213
2214 let mut block = VariableWidthBlock {
2215 data: LanceBuffer::from(data),
2216 offsets: offsets_buffer,
2217 bits_per_offset,
2218 num_values,
2219 block_info: BlockInfo::default(),
2220 };
2221
2222 block.compute_stat();
2223 DataBlock::VariableWidth(block)
2224 }
2225
2226 fn create_fsst_candidate_variable_width_block() -> DataBlock {
2227 create_variable_width_block(32, 4096, FSST_LEAST_INPUT_MAX_LENGTH as usize + 16)
2228 }
2229
2230 #[test]
2231 fn test_parameter_based_compression() {
2232 let mut params = CompressionParams::new();
2233
2234 params.columns.insert(
2236 "*_id".to_string(),
2237 CompressionFieldParams {
2238 rle_threshold: Some(0.3),
2239 compression: Some("lz4".to_string()),
2240 compression_level: None,
2241 bss: Some(BssMode::Off), minichunk_size: None,
2243 },
2244 );
2245
2246 let strategy = baseline_strategy(params);
2247 let field = create_test_field("user_id", DataType::Int32);
2248
2249 let data = create_fixed_width_block_with_stats(32, 1000, 100); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2254 let debug_str = format!("{:?}", compressor);
2256
2257 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
2259 assert!(debug_str.contains("RleEncoder"));
2260 }
2261
2262 #[test]
2263 fn test_type_level_parameters() {
2264 let mut params = CompressionParams::new();
2265
2266 params.types.insert(
2268 "Int32".to_string(),
2269 CompressionFieldParams {
2270 rle_threshold: Some(0.1), compression: Some("zstd".to_string()),
2272 compression_level: Some(3),
2273 bss: Some(BssMode::Off), minichunk_size: None,
2275 },
2276 );
2277
2278 let strategy = baseline_strategy(params);
2279 let field = create_test_field("some_column", DataType::Int32);
2280 let data = create_fixed_width_block_with_stats(32, 1000, 50);
2282
2283 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2284 assert!(format!("{:?}", compressor).contains("RleEncoder"));
2286 }
2287
2288 #[test]
2291 #[cfg(feature = "bitpacking")]
2292 fn test_block_bitpacks_with_zero_segment() {
2293 let strategy = baseline_strategy(CompressionParams::default());
2294 let field = create_test_field("levels", DataType::UInt16);
2295
2296 let mut values: Vec<u16> = vec![0; 1024];
2298 values.extend(std::iter::repeat_n(1u16, 1024));
2299 let mut block = FixedWidthDataBlock {
2300 bits_per_value: 16,
2301 data: LanceBuffer::reinterpret_vec(values),
2302 num_values: 2048,
2303 block_info: BlockInfo::default(),
2304 };
2305 block.compute_stat();
2306 let data = DataBlock::FixedWidth(block);
2307
2308 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
2309 let debug_str = format!("{:?}", compressor);
2310 assert!(
2311 debug_str.contains("OutOfLineBitpacking"),
2312 "expected OutOfLineBitpacking, got: {debug_str}"
2313 );
2314 }
2315
2316 #[rstest::rstest]
2317 #[case::runt_tail(1600)]
2318 #[case::whole_chunks(2048)]
2319 #[cfg(feature = "bitpacking")]
2320 fn test_block_skips_full_width_bitpacking(#[case] num_values: usize) {
2321 let mut block = FixedWidthDataBlock {
2322 bits_per_value: 64,
2323 data: LanceBuffer::reinterpret_vec(vec![-1_i64; num_values]),
2324 num_values: num_values as u64,
2325 block_info: BlockInfo::default(),
2326 };
2327 block.compute_stat();
2328
2329 assert!(try_bitpacking_block(&DataBlock::FixedWidth(block)).is_none());
2330 }
2331
2332 #[test]
2333 fn test_rle_block_accounts_for_header_before_selecting() {
2334 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2335 let field = create_test_field("small_constant", DataType::Int32);
2336 let values = vec![42i32; 2];
2337 let mut block = FixedWidthDataBlock {
2338 bits_per_value: 32,
2339 data: LanceBuffer::reinterpret_vec(values),
2340 num_values: 2,
2341 block_info: BlockInfo::default(),
2342 };
2343 block.compute_stat();
2344 let data = DataBlock::FixedWidth(block);
2345
2346 let (compressor, encoding) = selected_block_codec(&strategy, &field, &data);
2347
2348 assert!(format!("{compressor:?}").contains("ValueEncoder"));
2349 assert!(matches!(
2350 encoding.compression.as_ref(),
2351 Some(Compression::Flat(_))
2352 ));
2353 }
2354
2355 #[test]
2356 #[cfg(feature = "bitpacking")]
2357 fn test_rle_block_prefers_bitpacking_when_smaller() {
2358 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2359 let field = create_test_field("levels", DataType::UInt16);
2360
2361 let mut values = Vec::with_capacity(2048);
2362 for run_idx in 0..1024 {
2363 values.extend(std::iter::repeat_n((run_idx % 2) as u16, 2));
2364 }
2365 let mut block = FixedWidthDataBlock {
2366 bits_per_value: 16,
2367 data: LanceBuffer::reinterpret_vec(values),
2368 num_values: 2048,
2369 block_info: BlockInfo::default(),
2370 };
2371 block.compute_stat();
2372 let data = DataBlock::FixedWidth(block);
2373
2374 let (compressor, encoding) = selected_block_codec(&strategy, &field, &data);
2375 let debug_str = format!("{compressor:?}");
2376 assert!(
2377 debug_str.contains("OutOfLineBitpacking"),
2378 "expected OutOfLineBitpacking, got: {debug_str}"
2379 );
2380 assert!(matches!(
2381 encoding.compression.as_ref(),
2382 Some(Compression::OutOfLineBitpacking(_))
2383 ));
2384 }
2385
2386 #[test]
2387 #[cfg(feature = "bitpacking")]
2388 fn test_low_cardinality_prefers_bitpacking_over_rle() {
2389 let strategy = baseline_strategy(CompressionParams::default());
2390 let field = create_test_field("int_score", DataType::Int64);
2391
2392 let mut values: Vec<u64> = Vec::with_capacity(256);
2395 for run_idx in 0..64 {
2396 let value = match run_idx % 3 {
2397 0 => 3u64,
2398 1 => 4u64,
2399 _ => 5u64,
2400 };
2401 values.extend(std::iter::repeat_n(value, 4));
2402 }
2403
2404 let mut block = FixedWidthDataBlock {
2405 bits_per_value: 64,
2406 data: LanceBuffer::reinterpret_vec(values),
2407 num_values: 256,
2408 block_info: BlockInfo::default(),
2409 };
2410
2411 use crate::statistics::ComputeStat;
2412 block.compute_stat();
2413
2414 let data = DataBlock::FixedWidth(block);
2415 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2416 let debug_str = format!("{:?}", compressor);
2417 assert!(
2418 debug_str.contains("InlineBitpacking"),
2419 "expected InlineBitpacking, got: {debug_str}"
2420 );
2421 assert!(
2422 !debug_str.contains("RleEncoder"),
2423 "expected RLE to be skipped when bitpacking is smaller, got: {debug_str}"
2424 );
2425 }
2426
2427 fn check_uncompressed_encoding(encoding: &CompressiveEncoding, variable: bool) {
2428 let chain = extract_array_encoding_chain(encoding);
2429 if variable {
2430 assert_eq!(chain.len(), 2);
2431 assert_eq!(chain.first().unwrap().as_str(), "variable");
2432 assert_eq!(chain.get(1).unwrap().as_str(), "flat");
2433 } else {
2434 assert_eq!(chain.len(), 1);
2435 assert_eq!(chain.first().unwrap().as_str(), "flat");
2436 }
2437 }
2438
2439 #[test]
2440 fn test_none_compression() {
2441 let mut params = CompressionParams::new();
2442
2443 params.columns.insert(
2445 "embeddings".to_string(),
2446 CompressionFieldParams {
2447 compression: Some("none".to_string()),
2448 ..Default::default()
2449 },
2450 );
2451
2452 let strategy = baseline_strategy(params);
2453 let field = create_test_field("embeddings", DataType::Float32);
2454 let fixed_data = create_fixed_width_block(32, 1000);
2455 let variable_data = create_variable_width_block(32, 10, 32 * 1024);
2456
2457 let compressor = strategy
2459 .create_miniblock_compressor(&field, &fixed_data)
2460 .unwrap();
2461 let (_block, encoding) = compressor
2462 .compress(miniblock_context(), fixed_data.clone())
2463 .unwrap();
2464 check_uncompressed_encoding(&encoding, false);
2465 let compressor = strategy
2466 .create_miniblock_compressor(&field, &variable_data)
2467 .unwrap();
2468 let (_block, encoding) = compressor
2469 .compress(miniblock_context(), variable_data.clone())
2470 .unwrap();
2471 check_uncompressed_encoding(&encoding, true);
2472
2473 let compressor = strategy.create_per_value(&field, &fixed_data).unwrap();
2475 let (_block, encoding) = compressor.compress(fixed_data).unwrap();
2476 check_uncompressed_encoding(&encoding, false);
2477 let compressor = strategy.create_per_value(&field, &variable_data).unwrap();
2478 let (_block, encoding) = compressor.compress(variable_data).unwrap();
2479 check_uncompressed_encoding(&encoding, true);
2480 }
2481
2482 #[test]
2483 fn test_field_metadata_none_compression() {
2484 let mut arrow_field = ArrowField::new("simple_col", DataType::Binary, true);
2486 let mut metadata = HashMap::new();
2487 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
2488 arrow_field = arrow_field.with_metadata(metadata);
2489 let field = Field::try_from(&arrow_field).unwrap();
2490
2491 let strategy = baseline_strategy(CompressionParams::new());
2492
2493 let fixed_data = create_fixed_width_block(32, 1000);
2495 let variable_data = create_variable_width_block(32, 10, 32 * 1024);
2496
2497 let compressor = strategy
2498 .create_miniblock_compressor(&field, &fixed_data)
2499 .unwrap();
2500 let (_block, encoding) = compressor
2501 .compress(miniblock_context(), fixed_data.clone())
2502 .unwrap();
2503 check_uncompressed_encoding(&encoding, false);
2504
2505 let compressor = strategy
2506 .create_miniblock_compressor(&field, &variable_data)
2507 .unwrap();
2508 let (_block, encoding) = compressor
2509 .compress(miniblock_context(), variable_data.clone())
2510 .unwrap();
2511 check_uncompressed_encoding(&encoding, true);
2512
2513 let compressor = strategy.create_per_value(&field, &fixed_data).unwrap();
2515 let (_block, encoding) = compressor.compress(fixed_data).unwrap();
2516 check_uncompressed_encoding(&encoding, false);
2517
2518 let compressor = strategy.create_per_value(&field, &variable_data).unwrap();
2519 let (_block, encoding) = compressor.compress(variable_data).unwrap();
2520 check_uncompressed_encoding(&encoding, true);
2521 }
2522
2523 #[test]
2524 fn test_auto_fsst_disabled_for_binary_fields() {
2525 let strategy = baseline_strategy(CompressionParams::default());
2526 let field = create_test_field("bytes", DataType::Binary);
2527 let variable_data = create_fsst_candidate_variable_width_block();
2528
2529 let miniblock = strategy
2530 .create_miniblock_compressor(&field, &variable_data)
2531 .unwrap();
2532 let miniblock_debug = format!("{:?}", miniblock);
2533 assert!(
2534 miniblock_debug.contains("BinaryMiniBlockEncoder"),
2535 "expected BinaryMiniBlockEncoder, got: {miniblock_debug}"
2536 );
2537 assert!(
2538 !miniblock_debug.contains("FsstMiniBlockEncoder"),
2539 "did not expect FsstMiniBlockEncoder, got: {miniblock_debug}"
2540 );
2541
2542 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
2543 let per_value_debug = format!("{:?}", per_value);
2544 assert!(
2545 per_value_debug.contains("VariableEncoder"),
2546 "expected VariableEncoder, got: {per_value_debug}"
2547 );
2548 assert!(
2549 !per_value_debug.contains("FsstPerValueEncoder"),
2550 "did not expect FsstPerValueEncoder, got: {per_value_debug}"
2551 );
2552 }
2553
2554 #[test]
2555 fn test_auto_fsst_still_enabled_for_utf8_fields() {
2556 let strategy = baseline_strategy(CompressionParams::default());
2557 let field = create_test_field("text", DataType::Utf8);
2558 let variable_data = create_fsst_candidate_variable_width_block();
2559
2560 let miniblock = strategy
2561 .create_miniblock_compressor(&field, &variable_data)
2562 .unwrap();
2563 let miniblock_debug = format!("{:?}", miniblock);
2564 assert!(
2565 miniblock_debug.contains("FsstMiniBlockEncoder"),
2566 "expected FsstMiniBlockEncoder, got: {miniblock_debug}"
2567 );
2568
2569 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
2570 let per_value_debug = format!("{:?}", per_value);
2571 assert!(
2572 per_value_debug.contains("FsstPerValueEncoder"),
2573 "expected FsstPerValueEncoder, got: {per_value_debug}"
2574 );
2575 }
2576
2577 #[test]
2578 fn test_explicit_fsst_still_supported_for_binary_fields() {
2579 let mut params = CompressionParams::new();
2580 params.columns.insert(
2581 "bytes".to_string(),
2582 CompressionFieldParams {
2583 compression: Some("fsst".to_string()),
2584 ..Default::default()
2585 },
2586 );
2587
2588 let strategy = baseline_strategy(params);
2589 let field = create_test_field("bytes", DataType::Binary);
2590 let variable_data = create_fsst_candidate_variable_width_block();
2591
2592 let miniblock = strategy
2593 .create_miniblock_compressor(&field, &variable_data)
2594 .unwrap();
2595 let miniblock_debug = format!("{:?}", miniblock);
2596 assert!(
2597 miniblock_debug.contains("FsstMiniBlockEncoder"),
2598 "expected FsstMiniBlockEncoder, got: {miniblock_debug}"
2599 );
2600
2601 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
2602 let per_value_debug = format!("{:?}", per_value);
2603 assert!(
2604 per_value_debug.contains("FsstPerValueEncoder"),
2605 "expected FsstPerValueEncoder, got: {per_value_debug}"
2606 );
2607 }
2608
2609 #[test]
2610 #[cfg(feature = "zstd")]
2611 fn test_compression_level_honored_for_large_per_value() {
2612 let mut params = CompressionParams::new();
2613 params.columns.insert(
2614 "html".to_string(),
2615 CompressionFieldParams {
2616 compression: Some("zstd".to_string()),
2617 compression_level: Some(19),
2618 ..Default::default()
2619 },
2620 );
2621 let strategy = baseline_strategy(params);
2622 let field = create_test_field("html", DataType::Utf8);
2623 let large = create_variable_width_block(32, 64, 40 * 1024);
2624
2625 let per_value = strategy.create_per_value(&field, &large).unwrap();
2626 let debug = format!("{per_value:?}");
2627 assert!(
2628 debug.contains("ZstdBufferCompressor") && debug.contains("compression_level: 19"),
2629 "expected zstd level 19 to reach the per-value compressor, got: {debug}"
2630 );
2631 }
2632
2633 #[test]
2634 fn test_parameter_merge_priority() {
2635 let mut params = CompressionParams::new();
2636
2637 params.types.insert(
2639 "Int32".to_string(),
2640 CompressionFieldParams {
2641 rle_threshold: Some(0.5),
2642 compression: Some("lz4".to_string()),
2643 ..Default::default()
2644 },
2645 );
2646
2647 params.columns.insert(
2649 "user_id".to_string(),
2650 CompressionFieldParams {
2651 rle_threshold: Some(0.2),
2652 compression: Some("zstd".to_string()),
2653 compression_level: Some(6),
2654 bss: None,
2655 minichunk_size: None,
2656 },
2657 );
2658
2659 let merged = params.get_field_params("user_id", &DataType::Int32);
2661
2662 assert_eq!(merged.rle_threshold, Some(0.2));
2664 assert_eq!(merged.compression, Some("zstd".to_string()));
2665 assert_eq!(merged.compression_level, Some(6));
2666
2667 let merged = params.get_field_params("other_field", &DataType::Int32);
2669 assert_eq!(merged.rle_threshold, Some(0.5));
2670 assert_eq!(merged.compression, Some("lz4".to_string()));
2671 assert_eq!(merged.compression_level, None);
2672 }
2673
2674 #[test]
2675 fn test_pattern_matching() {
2676 let mut params = CompressionParams::new();
2677
2678 params.columns.insert(
2680 "log_*".to_string(),
2681 CompressionFieldParams {
2682 compression: Some("zstd".to_string()),
2683 compression_level: Some(6),
2684 ..Default::default()
2685 },
2686 );
2687
2688 let merged = params.get_field_params("log_messages", &DataType::Utf8);
2690 assert_eq!(merged.compression, Some("zstd".to_string()));
2691 assert_eq!(merged.compression_level, Some(6));
2692
2693 let merged = params.get_field_params("messages_log", &DataType::Utf8);
2695 assert_eq!(merged.compression, None);
2696 }
2697
2698 #[test]
2699 fn test_legacy_metadata_support() {
2700 let params = CompressionParams::new();
2701 let strategy = baseline_strategy(params);
2702
2703 let mut metadata = HashMap::new();
2705 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
2706 let mut field = create_test_field("some_column", DataType::Int32);
2707 field.metadata = metadata;
2708
2709 let data = create_fixed_width_block(32, 1000);
2710 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2711
2712 assert!(format!("{:?}", compressor).contains("ValueEncoder"));
2714 }
2715
2716 #[test]
2717 fn test_default_behavior() {
2718 let params = CompressionParams::new();
2720 let strategy = baseline_strategy(params);
2721
2722 let field = create_test_field("random_column", DataType::Int32);
2723 let data = create_fixed_width_block_with_stats(32, 1000, 600);
2725
2726 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2727 let debug_str = format!("{:?}", compressor);
2729 assert!(debug_str.contains("ValueEncoder") || debug_str.contains("InlineBitpacking"));
2730 }
2731
2732 #[test]
2733 fn test_field_metadata_compression() {
2734 let params = CompressionParams::new();
2735 let strategy = baseline_strategy(params);
2736
2737 let mut metadata = HashMap::new();
2739 metadata.insert(COMPRESSION_META_KEY.to_string(), "zstd".to_string());
2740 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "6".to_string());
2741 let mut field = create_test_field("test_column", DataType::Int32);
2742 field.metadata = metadata;
2743
2744 let data = create_fixed_width_block(32, 1000);
2745 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2746
2747 let debug_str = format!("{:?}", compressor);
2749 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
2750 }
2751
2752 #[test]
2753 fn test_field_metadata_rle_threshold() {
2754 let params = CompressionParams::new();
2755 let strategy = baseline_strategy(params);
2756
2757 let mut metadata = HashMap::new();
2759 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "0.8".to_string());
2760 metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); let mut field = create_test_field("test_column", DataType::Int32);
2762 field.metadata = metadata;
2763
2764 let data = create_fixed_width_block_with_stats(32, 1000, 100);
2767
2768 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2769
2770 let debug_str = format!("{:?}", compressor);
2772 assert!(debug_str.contains("RleEncoder"));
2773 }
2774
2775 #[test]
2776 fn test_rle_v2_miniblock_selects_u16_run_lengths() {
2777 let mut metadata = HashMap::new();
2778 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2779 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2780 let mut field = create_test_field("test_column", DataType::Int32);
2781 field.metadata = metadata;
2782
2783 let values = vec![7i32; 1000];
2784 let mut data = FixedWidthDataBlock {
2785 bits_per_value: 32,
2786 data: LanceBuffer::reinterpret_vec(values),
2787 num_values: 1000,
2788 block_info: BlockInfo::default(),
2789 };
2790 data.compute_stat();
2791 let data = DataBlock::FixedWidth(data);
2792
2793 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2794 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2795 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2796 assert_eq!(rle_run_length_bits(&encoding), 16);
2797 }
2798
2799 #[test]
2800 fn test_rle_v2_miniblock_keeps_u8_run_lengths_before_v2_3() {
2801 for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] {
2802 let mut metadata = HashMap::new();
2803 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2804 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2805 let mut field = create_test_field("test_column", DataType::Int32);
2806 field.metadata = metadata;
2807
2808 let values = vec![7i32; 1000];
2809 let mut data = FixedWidthDataBlock {
2810 bits_per_value: 32,
2811 data: LanceBuffer::reinterpret_vec(values),
2812 num_values: 1000,
2813 block_info: BlockInfo::default(),
2814 };
2815 data.compute_stat();
2816 let data = DataBlock::FixedWidth(data);
2817
2818 let strategy = strategy(version, CompressionParams::default());
2819 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2820 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2821 assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}");
2822 }
2823 }
2824
2825 #[test]
2826 fn test_rle_v2_uses_selected_width_cost_before_bitpacking() {
2827 let mut metadata = HashMap::new();
2828 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2829 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2830 let mut field = create_test_field("test_column", DataType::Int32);
2831 field.metadata = metadata;
2832
2833 let values = vec![0i32; 4096];
2834 let mut data = FixedWidthDataBlock {
2835 bits_per_value: 32,
2836 data: LanceBuffer::reinterpret_vec(values),
2837 num_values: 4096,
2838 block_info: BlockInfo::default(),
2839 };
2840 data.compute_stat();
2841 let data = DataBlock::FixedWidth(data);
2842
2843 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2844 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2845 let debug_str = format!("{compressor:?}");
2846 assert!(debug_str.contains("RleEncoder"));
2847
2848 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2849 assert_eq!(rle_run_length_bits(&encoding), 16);
2850 }
2851
2852 #[test]
2853 fn test_rle_v2_sorted_dictionary_indices_select_u16_run_lengths() {
2854 let field = create_test_field("dict_indices", DataType::Int32);
2855
2856 let mut values = Vec::with_capacity(1_200);
2857 for value in 0..4 {
2858 values.extend(std::iter::repeat_n(value, 300));
2859 }
2860 let mut data = FixedWidthDataBlock {
2861 bits_per_value: 32,
2862 data: LanceBuffer::reinterpret_vec(values),
2863 num_values: 1_200,
2864 block_info: BlockInfo::default(),
2865 };
2866 data.compute_stat();
2867 let data = DataBlock::FixedWidth(data);
2868
2869 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2870 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2871 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2872 assert_eq!(rle_run_length_bits(&encoding), 16);
2873 }
2874
2875 #[test]
2876 fn test_rle_v2_short_runs_keep_u8_run_lengths() {
2877 let field = create_test_field("dict_indices", DataType::Int32);
2878
2879 let mut values = Vec::with_capacity(1_280);
2880 for value in 0..10 {
2881 values.extend(std::iter::repeat_n(value, 128));
2882 }
2883 let mut data = FixedWidthDataBlock {
2884 bits_per_value: 32,
2885 data: LanceBuffer::reinterpret_vec(values),
2886 num_values: 1_280,
2887 block_info: BlockInfo::default(),
2888 };
2889 data.compute_stat();
2890 let data = DataBlock::FixedWidth(data);
2891
2892 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2893 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2894 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2895 assert_eq!(rle_run_length_bits(&encoding), 8);
2896 }
2897
2898 #[test]
2899 #[cfg(any(feature = "lz4", feature = "zstd"))]
2900 fn test_rle_miniblock_released_versions_keep_flat_children_when_compression_requested() {
2901 for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] {
2902 let mut params = CompressionParams::new();
2903 params.columns.insert(
2904 "dict_indices".to_string(),
2905 CompressionFieldParams {
2906 compression: Some(
2907 if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string(),
2908 ),
2909 rle_threshold: Some(1.0),
2910 bss: Some(BssMode::Off),
2911 ..Default::default()
2912 },
2913 );
2914 let strategy = strategy(version, params);
2915 let field = create_test_field("dict_indices", DataType::UInt32);
2916
2917 let mut values = Vec::with_capacity(8192 * 4);
2918 for value in 0..8192u32 {
2919 values.extend(std::iter::repeat_n(value, 4));
2920 }
2921 let mut data = FixedWidthDataBlock {
2922 bits_per_value: 32,
2923 data: LanceBuffer::reinterpret_vec(values),
2924 num_values: 8192 * 4,
2925 block_info: BlockInfo::default(),
2926 };
2927 data.compute_stat();
2928 let data = DataBlock::FixedWidth(data);
2929
2930 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2931 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2932 let rle = expect_rle_encoding(&encoding);
2933
2934 assert!(
2935 matches!(
2936 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2937 Compression::Flat(_)
2938 ),
2939 "version={version}"
2940 );
2941 assert!(
2942 matches!(
2943 rle.run_lengths
2944 .as_ref()
2945 .unwrap()
2946 .compression
2947 .as_ref()
2948 .unwrap(),
2949 Compression::Flat(_)
2950 ),
2951 "version={version}"
2952 );
2953 }
2954 }
2955
2956 #[test]
2957 #[cfg(feature = "bitpacking")]
2958 fn test_rle_miniblock_strategy_bitpacks_child_values_when_smaller() {
2959 let field = create_test_field("dict_indices", DataType::Int32);
2960
2961 let mut values = Vec::with_capacity(8192 * 4);
2962 for value in 0..8192 {
2963 values.extend(std::iter::repeat_n(value, 4));
2964 }
2965 let mut data = FixedWidthDataBlock {
2966 bits_per_value: 32,
2967 data: LanceBuffer::reinterpret_vec(values),
2968 num_values: 8192 * 4,
2969 block_info: BlockInfo::default(),
2970 };
2971 data.compute_stat();
2972 let data = DataBlock::FixedWidth(data);
2973
2974 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2975 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2976 let debug_str = format!("{compressor:?}");
2977 assert!(debug_str.contains("RleEncoder"));
2978
2979 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2980 let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else {
2981 panic!("expected RLE encoding");
2982 };
2983 assert!(matches!(
2984 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2985 Compression::OutOfLineBitpacking(_)
2986 ));
2987 assert!(matches!(
2988 rle.run_lengths
2989 .as_ref()
2990 .unwrap()
2991 .compression
2992 .as_ref()
2993 .unwrap(),
2994 Compression::Flat(_)
2995 ));
2996 }
2997
2998 #[test]
2999 #[cfg(feature = "bitpacking")]
3000 fn test_rle_miniblock_keeps_child_bitpacked_rle_when_smaller_than_inline_bitpacking() {
3001 let field = create_test_field("int_score", DataType::UInt64);
3002
3003 let mut values = Vec::with_capacity(8192 * 8);
3004 for run_idx in 0..8192 {
3005 let value = match run_idx % 3 {
3006 0 => 3u64,
3007 1 => 4u64,
3008 _ => 5u64,
3009 };
3010 values.extend(std::iter::repeat_n(value, 8));
3011 }
3012 let mut data = FixedWidthDataBlock {
3013 bits_per_value: 64,
3014 data: LanceBuffer::reinterpret_vec(values),
3015 num_values: 8192 * 8,
3016 block_info: BlockInfo::default(),
3017 };
3018 data.compute_stat();
3019 let data = DataBlock::FixedWidth(data);
3020
3021 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
3022 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
3023 let debug_str = format!("{compressor:?}");
3024 assert!(
3025 debug_str.contains("RleEncoder"),
3026 "expected RLE to beat inline bitpacking after child selection, got: {debug_str}"
3027 );
3028
3029 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
3030 let rle = expect_rle_encoding(&encoding);
3031 assert!(matches!(
3032 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
3033 Compression::OutOfLineBitpacking(_)
3034 ));
3035 assert!(matches!(
3036 rle.run_lengths
3037 .as_ref()
3038 .unwrap()
3039 .compression
3040 .as_ref()
3041 .unwrap(),
3042 Compression::Flat(_)
3043 ));
3044 }
3045
3046 #[test]
3047 fn test_field_metadata_override_params() {
3048 let mut params = CompressionParams::new();
3050 params.columns.insert(
3051 "test_column".to_string(),
3052 CompressionFieldParams {
3053 rle_threshold: Some(0.3),
3054 compression: Some("lz4".to_string()),
3055 compression_level: None,
3056 bss: None,
3057 minichunk_size: None,
3058 },
3059 );
3060
3061 let strategy = baseline_strategy(params);
3062
3063 let mut metadata = HashMap::new();
3065 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
3066 let mut field = create_test_field("test_column", DataType::Int32);
3067 field.metadata = metadata;
3068
3069 let data = create_fixed_width_block(32, 1000);
3070 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
3071
3072 assert!(format!("{:?}", compressor).contains("ValueEncoder"));
3074 }
3075
3076 #[test]
3077 fn test_field_metadata_mixed_configuration() {
3078 let mut params = CompressionParams::new();
3080 params.types.insert(
3081 "Int32".to_string(),
3082 CompressionFieldParams {
3083 rle_threshold: Some(0.5),
3084 compression: Some("lz4".to_string()),
3085 ..Default::default()
3086 },
3087 );
3088
3089 let strategy = baseline_strategy(params);
3090
3091 let mut metadata = HashMap::new();
3093 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "3".to_string());
3094 let mut field = create_test_field("test_column", DataType::Int32);
3095 field.metadata = metadata;
3096
3097 let data = create_fixed_width_block(32, 1000);
3098 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
3099
3100 let debug_str = format!("{:?}", compressor);
3102 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
3103 }
3104
3105 #[test]
3106 fn test_bss_field_metadata() {
3107 let params = CompressionParams::new();
3108 let strategy = baseline_strategy(params);
3109
3110 let mut metadata = HashMap::new();
3112 metadata.insert(BSS_META_KEY.to_string(), "on".to_string());
3113 metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string());
3114 let arrow_field =
3115 ArrowField::new("temperature", DataType::Float32, false).with_metadata(metadata);
3116 let field = Field::try_from(&arrow_field).unwrap();
3117
3118 let data = create_fixed_width_block(32, 100);
3120
3121 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
3122 let debug_str = format!("{:?}", compressor);
3123 assert!(debug_str.contains("ByteStreamSplitEncoder"));
3124 }
3125
3126 #[test]
3127 fn test_bss_with_compression() {
3128 let params = CompressionParams::new();
3129 let strategy = baseline_strategy(params);
3130
3131 let mut metadata = HashMap::new();
3133 metadata.insert(BSS_META_KEY.to_string(), "on".to_string());
3134 metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string());
3135 let arrow_field =
3136 ArrowField::new("sensor_data", DataType::Float64, false).with_metadata(metadata);
3137 let field = Field::try_from(&arrow_field).unwrap();
3138
3139 let data = create_fixed_width_block(64, 100);
3141
3142 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
3143 let debug_str = format!("{:?}", compressor);
3144 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
3146 assert!(debug_str.contains("ByteStreamSplitEncoder"));
3147 }
3148
3149 #[test]
3150 #[cfg(any(feature = "lz4", feature = "zstd"))]
3151 fn test_general_block_decompression_fixed_width_v2_2() {
3152 let mut params = CompressionParams::new();
3154 params.columns.insert(
3155 "dict_values".to_string(),
3156 CompressionFieldParams {
3157 compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()),
3158 ..Default::default()
3159 },
3160 );
3161
3162 let strategy = strategy(TestEncoding::StructuralU32, params);
3163
3164 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
3165 let data = create_fixed_width_block(24, 1024);
3166 let DataBlock::FixedWidth(expected_block) = &data else {
3167 panic!("expected fixed width block");
3168 };
3169 let expected_bits = expected_block.bits_per_value;
3170 let expected_num_values = expected_block.num_values;
3171 let num_values = expected_num_values;
3172
3173 let compressor = strategy
3174 .create_block_compressor(&field, &data)
3175 .expect("general compression should be selected");
3176 let (compressed_buffer, encoding) = compressor
3177 .compress(data.clone())
3178 .expect("write path general compression should succeed");
3179 match encoding.compression.as_ref() {
3180 Some(Compression::General(_)) => {}
3181 other => panic!("expected general compression, got {:?}", other),
3182 }
3183
3184 let decompressor = DefaultDecompressionStrategy::default()
3185 .create_block_decompressor(&encoding)
3186 .expect("general block decompressor should be created");
3187
3188 let decoded = decompressor
3189 .decompress(compressed_buffer, num_values)
3190 .expect("decompression should succeed");
3191
3192 match decoded {
3193 DataBlock::FixedWidth(block) => {
3194 assert_eq!(block.bits_per_value, expected_bits);
3195 assert_eq!(block.num_values, expected_num_values);
3196 assert_eq!(block.data.as_ref(), expected_block.data.as_ref());
3197 }
3198 _ => panic!("expected fixed width block"),
3199 }
3200 }
3201
3202 #[cfg(any(feature = "lz4", feature = "zstd"))]
3203 fn assert_general_block_preserves_compression_level(
3204 compression: &str,
3205 expected_scheme: crate::format::metadata::CompressionScheme,
3206 compression_level: Option<i32>,
3207 ) {
3208 let mut params = CompressionParams::new();
3209 params.columns.insert(
3210 "dict_values".to_string(),
3211 CompressionFieldParams {
3212 compression: Some(compression.to_string()),
3213 compression_level,
3214 ..Default::default()
3215 },
3216 );
3217 let strategy = strategy(TestEncoding::StructuralU32, params);
3218 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
3219 let data = create_fixed_width_block(24, 1024);
3220
3221 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
3222 let (_, encoding) = compressor.compress(data).unwrap();
3223 let Some(Compression::General(general)) = encoding.compression.as_ref() else {
3224 panic!("expected general compression");
3225 };
3226
3227 assert_eq!(
3228 general.compression.as_ref(),
3229 Some(&crate::format::metadata::BufferCompression {
3230 scheme: expected_scheme as i32,
3231 level: compression_level,
3232 })
3233 );
3234 }
3235
3236 #[test]
3237 #[cfg(feature = "zstd")]
3238 fn test_general_block_preserves_absent_zstd_level() {
3239 assert_general_block_preserves_compression_level(
3240 "zstd",
3241 crate::format::metadata::CompressionScheme::CompressionAlgorithmZstd,
3242 None,
3243 );
3244 }
3245
3246 #[test]
3247 #[cfg(feature = "lz4")]
3248 fn test_general_block_preserves_explicit_lz4_level() {
3249 assert_general_block_preserves_compression_level(
3250 "lz4",
3251 crate::format::metadata::CompressionScheme::CompressionAlgorithmLz4,
3252 Some(7),
3253 );
3254 }
3255
3256 #[test]
3257 #[cfg(any(feature = "lz4", feature = "zstd"))]
3258 fn test_general_compression_not_selected_for_v2_1_even_if_requested() {
3259 let mut params = CompressionParams::new();
3260 params.columns.insert(
3261 "dict_values".to_string(),
3262 CompressionFieldParams {
3263 compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()),
3264 ..Default::default()
3265 },
3266 );
3267
3268 let strategy = strategy(TestEncoding::StructuralU16, params);
3269 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
3270 let data = create_fixed_width_block(24, 1024);
3271
3272 let compressor = strategy
3273 .create_block_compressor(&field, &data)
3274 .expect("block compressor selection should succeed");
3275 let (_, encoding) = compressor.compress(data).unwrap();
3276
3277 assert!(
3278 !matches!(encoding.compression.as_ref(), Some(Compression::General(_))),
3279 "general compression should not be selected for V2.1"
3280 );
3281 }
3282
3283 #[test]
3284 fn test_none_compression_disables_auto_general_block_compression() {
3285 let mut params = CompressionParams::new();
3286 params.columns.insert(
3287 "dict_values".to_string(),
3288 CompressionFieldParams {
3289 compression: Some("none".to_string()),
3290 ..Default::default()
3291 },
3292 );
3293
3294 let strategy = strategy(TestEncoding::StructuralU32, params);
3295 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
3296 let data = create_fixed_width_block(24, 20_000);
3297
3298 assert!(
3299 data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION,
3300 "test requires block size above automatic general compression threshold"
3301 );
3302
3303 let compressor = strategy
3304 .create_block_compressor(&field, &data)
3305 .expect("block compressor selection should succeed");
3306 let (_, encoding) = compressor.compress(data).unwrap();
3307
3308 assert!(
3309 !matches!(encoding.compression.as_ref(), Some(Compression::General(_))),
3310 "compression=none should disable automatic block general compression"
3311 );
3312 }
3313
3314 #[test]
3315 fn test_rle_v2_block_selects_u32_run_lengths() {
3316 let field = create_test_field("dict_indices", DataType::Int32);
3317 let expected_values = vec![42i32; 70_000];
3318 let mut block = FixedWidthDataBlock {
3319 bits_per_value: 32,
3320 data: LanceBuffer::reinterpret_vec(expected_values.clone()),
3321 num_values: expected_values.len() as u64,
3322 block_info: BlockInfo::default(),
3323 };
3324 block.compute_stat();
3325 let data = DataBlock::FixedWidth(block);
3326
3327 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::new());
3328 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
3329 let (compressed, encoding) = compressor.compress(data).unwrap();
3330 assert_eq!(rle_run_length_bits(&encoding), 32);
3331 let decompressor = DefaultDecompressionStrategy::default()
3332 .create_block_decompressor(&encoding)
3333 .unwrap();
3334 let decoded = decompressor
3335 .decompress(compressed, expected_values.len() as u64)
3336 .unwrap();
3337
3338 match decoded {
3339 DataBlock::FixedWidth(block) => {
3340 let values = block.data.borrow_to_typed_slice::<i32>();
3341 assert_eq!(values.as_ref(), expected_values);
3342 }
3343 _ => panic!("expected fixed-width block"),
3344 }
3345 }
3346
3347 #[test]
3348 fn test_rle_v2_block_keeps_u8_run_lengths_for_v2_2() {
3349 let field = create_test_field("dict_indices", DataType::Int32);
3350 let values = vec![42i32; 70_000];
3351 let mut block = FixedWidthDataBlock {
3352 bits_per_value: 32,
3353 data: LanceBuffer::reinterpret_vec(values),
3354 num_values: 70_000,
3355 block_info: BlockInfo::default(),
3356 };
3357 block.compute_stat();
3358 let data = DataBlock::FixedWidth(block);
3359
3360 let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new());
3361 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
3362 let (_, encoding) = compressor.compress(data).unwrap();
3363 assert_eq!(rle_run_length_bits(&encoding), 8);
3364 }
3365
3366 #[test]
3367 fn test_rle_block_used_for_version_v2_2() {
3368 let field = create_test_field("test_repdef", DataType::UInt16);
3369
3370 let num_values = 1000u64;
3372 let mut data = Vec::with_capacity(num_values as usize);
3373 for i in 0..10 {
3374 for _ in 0..100 {
3375 data.push(i as u16);
3376 }
3377 }
3378
3379 let mut block = FixedWidthDataBlock {
3380 bits_per_value: 16,
3381 data: LanceBuffer::reinterpret_vec(data),
3382 num_values,
3383 block_info: BlockInfo::default(),
3384 };
3385
3386 block.compute_stat();
3387
3388 let data_block = DataBlock::FixedWidth(block);
3389
3390 let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new());
3391
3392 let compressor = strategy
3393 .create_block_compressor(&field, &data_block)
3394 .unwrap();
3395
3396 let debug_str = format!("{:?}", compressor);
3397 assert!(debug_str.contains("RleEncoder"));
3398 }
3399
3400 #[test]
3401 fn test_rle_block_not_used_for_version_v2_1() {
3402 let field = create_test_field("test_repdef", DataType::UInt16);
3403
3404 let num_values = 1000u64;
3406 let mut data = Vec::with_capacity(num_values as usize);
3407 for i in 0..10 {
3408 for _ in 0..100 {
3409 data.push(i as u16);
3410 }
3411 }
3412
3413 let mut block = FixedWidthDataBlock {
3414 bits_per_value: 16,
3415 data: LanceBuffer::reinterpret_vec(data),
3416 num_values,
3417 block_info: BlockInfo::default(),
3418 };
3419
3420 block.compute_stat();
3421
3422 let data_block = DataBlock::FixedWidth(block);
3423
3424 let strategy = strategy(TestEncoding::StructuralU16, CompressionParams::new());
3425
3426 let compressor = strategy
3427 .create_block_compressor(&field, &data_block)
3428 .unwrap();
3429
3430 let debug_str = format!("{:?}", compressor);
3431 assert!(
3432 !debug_str.contains("RleEncoder"),
3433 "RLE should not be used for V2.1"
3434 );
3435 }
3436
3437 #[test]
3438 fn generic_block_rle_round_trips_metadata_children() {
3439 use crate::encodings::physical::{
3440 constant::ConstantEncoder, range::RangeEncoder, rle::BlockRleEncoder,
3441 };
3442
3443 let values = vec![10_u64, 10, 20, 20, 30, 30];
3444 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
3445 bits_per_value: 64,
3446 data: LanceBuffer::reinterpret_vec(values.clone()),
3447 num_values: values.len() as u64,
3448 block_info: BlockInfo::default(),
3449 });
3450 let encoder = BlockRleEncoder::try_new(
3451 64,
3452 Box::new(RangeEncoder::new(64, 10, 10)),
3453 Box::new(ConstantEncoder::new(32, 2)),
3454 )
3455 .unwrap();
3456 let (payload, encoding) = encoder.compress(block).unwrap();
3457 assert!(payload.is_none());
3458
3459 let decoder = create_fixed_width_block_decompressor(&encoding, 64).unwrap();
3460 let decoded = decoder.decompress(None, values.len() as u64).unwrap();
3461 assert_eq!(
3462 decoded
3463 .as_fixed_width()
3464 .unwrap()
3465 .data
3466 .borrow_to_typed_slice::<u64>()
3467 .as_ref(),
3468 values
3469 );
3470 }
3471
3472 #[test]
3473 fn generic_block_dictionary_uses_the_typed_decoder_tree() {
3474 use crate::encodings::physical::{dictionary::BlockDictionaryEncoder, value::ValueEncoder};
3475
3476 let values = vec![10_u64, 20, 10, 20];
3477 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
3478 bits_per_value: 64,
3479 data: LanceBuffer::reinterpret_vec(values.clone()),
3480 num_values: values.len() as u64,
3481 block_info: BlockInfo::default(),
3482 });
3483 let encoder = BlockDictionaryEncoder::try_new(
3484 64,
3485 Arc::from([10, 20]),
3486 Box::new(ValueEncoder::default()),
3487 Box::new(ValueEncoder::default()),
3488 )
3489 .unwrap();
3490 let (payload, encoding) = encoder.compress(block).unwrap();
3491 let decoder = create_fixed_width_block_decompressor(&encoding, 64).unwrap();
3492 let decoded = decoder.decompress(payload, values.len() as u64).unwrap();
3493 assert_eq!(
3494 decoded
3495 .as_fixed_width()
3496 .unwrap()
3497 .data
3498 .borrow_to_typed_slice::<u64>()
3499 .as_ref(),
3500 values
3501 );
3502 }
3503
3504 #[test]
3505 fn decompressor_dispatch_rejects_missing_variants() {
3506 let strategy = DefaultDecompressionStrategy::default();
3507 let missing = CompressiveEncoding::default();
3508 assert!(
3509 strategy
3510 .create_miniblock_decompressor(&missing, &strategy)
3511 .is_err()
3512 );
3513 assert!(
3514 strategy
3515 .create_fixed_per_value_decompressor(&missing)
3516 .is_err()
3517 );
3518 assert!(
3519 strategy
3520 .create_variable_per_value_decompressor(&missing)
3521 .is_err()
3522 );
3523 assert!(strategy.create_block_decompressor(&missing).is_err());
3524 }
3525}