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::ConstantDecompressor,
46 fsst::{
47 FsstMiniBlockDecompressor, FsstMiniBlockEncoder, FsstPerValueDecompressor,
48 FsstPerValueEncoder,
49 },
50 general::{GeneralMiniBlockCompressor, GeneralMiniBlockDecompressor},
51 packed::{
52 PackedStructFixedPerValueDecompressor, PackedStructFixedPerValueEncoder,
53 PackedStructFixedWidthMiniBlockDecompressor,
54 PackedStructFixedWidthMiniBlockEncoder, PackedStructVariablePerValueDecompressor,
55 PackedStructVariablePerValueEncoder, VariablePackedStructFieldDecoder,
56 VariablePackedStructFieldKind,
57 },
58 rle::{
59 RleChildDecompressor, RleDecompressor, RleEncoder, RunLengthWidth,
60 rle_encoded_size, select_run_length_width,
61 },
62 value::{ValueDecompressor, ValueEncoder},
63 },
64 },
65 format::pb21::{CompressiveEncoding, compressive_encoding::Compression},
66 statistics::{GetStat, Stat},
67};
68
69use arrow_array::{cast::AsArray, types::UInt64Type};
70use arrow_schema::DataType;
71use fsst::fsst::{FSST_LEAST_INPUT_MAX_LENGTH, FSST_LEAST_INPUT_SIZE};
72use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt};
73use std::{str::FromStr, sync::Arc};
74
75const DEFAULT_RLE_COMPRESSION_THRESHOLD: f64 = 0.5;
81
82const MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION: u64 = 32 * 1024;
84const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::<u64>() as u128;
85
86pub trait BlockCompressor: std::fmt::Debug + Send + Sync {
99 fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)>;
104}
105
106pub trait CompressionStrategy: Send + Sync + std::fmt::Debug {
119 fn create_block_compressor(
121 &self,
122 field: &Field,
123 data: &DataBlock,
124 ) -> Result<Box<dyn BlockCompressor>>;
125
126 fn create_per_value(
128 &self,
129 field: &Field,
130 data: &DataBlock,
131 ) -> Result<Box<dyn PerValueCompressor>>;
132
133 fn create_miniblock_compressor(
135 &self,
136 field: &Field,
137 data: &DataBlock,
138 ) -> Result<Box<dyn MiniBlockCompressor>>;
139}
140
141pub(crate) fn compress_required_block(
142 strategy: &dyn CompressionStrategy,
143 field: &Field,
144 data: DataBlock,
145) -> Result<(LanceBuffer, CompressiveEncoding)> {
146 let compressor = strategy.create_block_compressor(field, &data)?;
147 let (payload, encoding) = compressor.compress(data)?;
148 let payload = payload.ok_or_else(|| {
149 Error::internal("Required block compressor selected a metadata-only codec".to_string())
150 })?;
151 Ok((payload, encoding))
152}
153
154fn try_bss_for_mini_block(
155 data: &FixedWidthDataBlock,
156 params: &CompressionFieldParams,
157) -> Option<Box<dyn MiniBlockCompressor>> {
158 if params.compression.is_none() || params.compression.as_deref() == Some("none") {
161 return None;
162 }
163
164 let mode = params.bss.unwrap_or(BssMode::Auto);
165 if should_use_bss(data, mode) {
167 return Some(Box::new(ByteStreamSplitEncoder::new(
168 data.bits_per_value as usize,
169 )));
170 }
171 None
172}
173
174fn rle_is_applicable(data: &FixedWidthDataBlock, params: &CompressionFieldParams) -> Option<u128> {
175 let bits = data.bits_per_value;
176 if !matches!(bits, 8 | 16 | 32 | 64) {
177 return None;
178 }
179
180 let type_size = bits / 8;
181 let run_count = data.expect_single_stat::<UInt64Type>(Stat::RunCount);
182 let threshold = params
183 .rle_threshold
184 .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD);
185
186 let passes_threshold = match params.rle_threshold {
189 Some(_) => (run_count as f64) < (data.num_values as f64) * threshold,
190 None => true,
191 };
192
193 if !passes_threshold {
194 return None;
195 }
196
197 Some((data.num_values as u128) * (type_size as u128))
198}
199
200fn rle_beats_raw_and_bitpacking(
201 data: &FixedWidthDataBlock,
202 encoded_bytes: u128,
203 raw_bytes: u128,
204) -> bool {
205 if encoded_bytes >= raw_bytes {
206 return false;
207 }
208
209 #[cfg(feature = "bitpacking")]
210 {
211 if let Some(bitpack_bytes) = estimate_inline_bitpacking_bytes(data).map(u128::from)
212 && bitpack_bytes < encoded_bytes
213 {
214 return false;
215 }
216 }
217 true
218}
219
220fn try_fixed_u8_rle_for_mini_block(
221 data: &FixedWidthDataBlock,
222 params: &CompressionFieldParams,
223) -> Option<Box<dyn MiniBlockCompressor>> {
224 let raw_bytes = rle_is_applicable(data, params)?;
225 let rle_bytes = estimate_rle_size_for_width_from_data(
226 data,
227 Some(*MAX_MINIBLOCK_VALUES),
228 RunLengthWidth::U8,
229 )
230 .ok()?;
231 rle_beats_raw_and_bitpacking(data, rle_bytes, raw_bytes)
232 .then(|| Box::new(RleEncoder::with_run_length_width(RunLengthWidth::U8)) as _)
233}
234
235fn try_child_rle_for_mini_block(
236 data: &FixedWidthDataBlock,
237 params: &CompressionFieldParams,
238) -> Option<Box<dyn MiniBlockCompressor>> {
239 let raw_bytes = rle_is_applicable(data, params)?;
240 let (run_length_width, estimated_bytes) =
241 estimate_rle_width_and_size_from_data(data, Some(*MAX_MINIBLOCK_VALUES)).ok()?;
242 let child_compression = rle_child_compression_config(params);
243 let encoder = || {
244 RleEncoder::with_child_encoding(
245 run_length_width,
246 child_compression,
247 child_compression,
248 true,
249 )
250 };
251
252 #[cfg(feature = "bitpacking")]
253 let bitpack_bytes = estimate_inline_bitpacking_bytes(data).map(u128::from);
254 #[cfg(not(feature = "bitpacking"))]
255 let bitpack_bytes = None::<u128>;
256
257 let should_measure_children = (child_compression.is_some() || cfg!(feature = "bitpacking"))
258 && (estimated_bytes >= raw_bytes
259 || bitpack_bytes.is_some_and(|bytes| bytes < estimated_bytes));
260 let selected_bytes = if should_measure_children {
261 encoder().selected_payload_size(data).ok()?
262 } else {
263 estimated_bytes
264 };
265
266 rle_beats_raw_and_bitpacking(data, selected_bytes, raw_bytes).then(|| Box::new(encoder()) as _)
267}
268
269fn rle_child_compression_config(params: &CompressionFieldParams) -> Option<CompressionConfig> {
270 let raw = params.compression.as_deref()?;
271 if matches!(raw, "none" | "fsst") {
272 return None;
273 }
274 let scheme = CompressionScheme::from_str(raw).ok()?;
275 Some(CompressionConfig::new(scheme, params.compression_level))
276}
277
278fn try_rle_for_block_with_width(
279 data: &FixedWidthDataBlock,
280 params: &CompressionFieldParams,
281 run_length_width: RunLengthWidth,
282 rle_payload_bytes: u128,
283) -> Result<Option<Box<dyn BlockCompressor>>> {
284 let bits = data.bits_per_value;
285 if !matches!(bits, 8 | 16 | 32 | 64) {
286 return Ok(None);
287 }
288
289 let run_count = data.expect_single_stat::<UInt64Type>(Stat::RunCount);
290 let threshold = params
291 .rle_threshold
292 .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD);
293
294 let passes_threshold = match params.rle_threshold {
295 Some(_) => (run_count as f64) < (data.num_values as f64) * threshold,
296 None => true,
297 };
298
299 if !passes_threshold {
300 return Ok(None);
301 }
302
303 let raw_bytes = (data.num_values as u128) * ((bits / 8) as u128);
304 let rle_bytes = rle_payload_bytes.saturating_add(RLE_BLOCK_HEADER_BYTES);
305
306 if rle_bytes >= raw_bytes {
307 return Ok(None);
308 }
309
310 #[cfg(feature = "bitpacking")]
311 {
312 if let Some(bitpack_bytes) = estimate_block_bitpacking_bytes(data)
313 && bitpack_bytes < rle_bytes
314 {
315 return Ok(None);
316 }
317 }
318
319 Ok(Some(Box::new(RleEncoder::with_run_length_width(
320 run_length_width,
321 ))))
322}
323
324fn try_fixed_u8_rle_for_block(
325 data: &FixedWidthDataBlock,
326 params: &CompressionFieldParams,
327) -> Result<Option<Box<dyn BlockCompressor>>> {
328 if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) {
329 return Ok(None);
330 }
331 let encoded_bytes = estimate_rle_size_for_width_from_data(data, None, RunLengthWidth::U8)?;
332 try_rle_for_block_with_width(data, params, RunLengthWidth::U8, encoded_bytes)
333}
334
335fn try_variable_rle_for_block(
336 data: &FixedWidthDataBlock,
337 params: &CompressionFieldParams,
338) -> Result<Option<Box<dyn BlockCompressor>>> {
339 if !matches!(data.bits_per_value, 8 | 16 | 32 | 64) {
340 return Ok(None);
341 }
342 let (width, encoded_bytes) = estimate_rle_width_and_size_from_data(data, None)?;
343 try_rle_for_block_with_width(data, params, width, encoded_bytes)
344}
345
346fn estimate_rle_width_and_size_from_data(
347 data: &FixedWidthDataBlock,
348 max_segment_values: Option<u64>,
349) -> Result<(RunLengthWidth, u128)> {
350 select_run_length_width(
351 &data.data,
352 data.num_values,
353 data.bits_per_value,
354 max_segment_values,
355 )
356}
357
358fn estimate_rle_size_for_width_from_data(
359 data: &FixedWidthDataBlock,
360 max_segment_values: Option<u64>,
361 run_length_width: RunLengthWidth,
362) -> Result<u128> {
363 rle_encoded_size(
364 &data.data,
365 data.num_values,
366 data.bits_per_value,
367 max_segment_values,
368 run_length_width,
369 )
370}
371
372fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
373 #[cfg(feature = "bitpacking")]
374 {
375 let bits = _data.bits_per_value;
376 if estimate_inline_bitpacking_bytes(_data).is_some() {
377 return Some(Box::new(InlineBitpacking::new(bits)));
378 }
379 None
380 }
381 #[cfg(not(feature = "bitpacking"))]
382 {
383 None
384 }
385}
386
387#[cfg(feature = "bitpacking")]
388fn estimate_inline_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option<u64> {
389 use arrow_array::cast::AsArray;
390
391 let bits = data.bits_per_value;
392 if !matches!(bits, 8 | 16 | 32 | 64) {
393 return None;
394 }
395 if data.num_values == 0 {
396 return None;
397 }
398
399 let bit_widths = data.expect_stat(Stat::BitWidth);
400 let widths = bit_widths.as_primitive::<UInt64Type>();
401
402 let words_per_chunk: u128 = 1;
403 let word_bytes: u128 = (bits / 8) as u128;
404 let mut total_words: u128 = 0;
405 for i in 0..widths.len() {
406 let bit_width = widths.value(i) as u128;
407 let packed_words = (1024u128 * bit_width) / (bits as u128);
408 total_words = total_words.saturating_add(words_per_chunk.saturating_add(packed_words));
409 }
410
411 let estimated_bytes = total_words.saturating_mul(word_bytes);
412 let raw_bytes = data.data_size() as u128;
413
414 if estimated_bytes >= raw_bytes {
415 return None;
416 }
417
418 u64::try_from(estimated_bytes).ok()
419}
420
421fn try_bitpack_for_block(data: &FixedWidthDataBlock) -> Option<Box<dyn BlockCompressor>> {
422 let bits = data.bits_per_value;
423 if !matches!(bits, 8 | 16 | 32 | 64) {
424 return None;
425 }
426
427 let bit_widths = data.expect_stat(Stat::BitWidth);
428 let widths = bit_widths.as_primitive::<UInt64Type>();
429 let max_bit_width = *widths.values().iter().max().unwrap();
430
431 if max_bit_width >= bits {
433 return None;
434 }
435
436 let too_small =
437 widths.len() == 1 && InlineBitpacking::min_size_bytes(widths.value(0)) >= data.data_size();
438
439 if too_small {
440 return None;
441 }
442
443 if data.num_values <= 1024 {
444 Some(Box::new(InlineBitpacking::new(bits)))
445 } else {
446 Some(Box::new(OutOfLineBitpacking::new(max_bit_width, bits)))
447 }
448}
449
450#[cfg(feature = "bitpacking")]
451fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option<u128> {
452 let bits = data.bits_per_value;
453 if !matches!(bits, 8 | 16 | 32 | 64) || data.num_values == 0 {
454 return None;
455 }
456
457 let bit_widths = data.expect_stat(Stat::BitWidth);
458 let widths = bit_widths.as_primitive::<UInt64Type>();
459 let max_bit_width = *widths.values().iter().max()?;
460 let word_bytes = (bits / 8) as u128;
461
462 let bitpacked_words = if data.num_values <= 1024 {
463 1 + (1024u128 * (max_bit_width as u128)) / (bits as u128)
464 } else {
465 estimate_out_of_line_bitpacking_words(data.num_values, max_bit_width, bits)?
466 };
467 let bitpacked_bytes = bitpacked_words.saturating_mul(word_bytes);
468 if bitpacked_bytes >= data.data_size() as u128 {
469 return None;
470 }
471
472 Some(bitpacked_bytes)
473}
474
475#[cfg(feature = "bitpacking")]
476fn estimate_out_of_line_bitpacking_words(
477 num_values: u64,
478 compressed_bits_per_value: u64,
479 bits_per_value: u64,
480) -> Option<u128> {
481 let num_values = usize::try_from(num_values).ok()?;
482 let compressed_bits_per_value = usize::try_from(compressed_bits_per_value).ok()?;
483 let bits_per_value = usize::try_from(bits_per_value).ok()?;
484 if compressed_bits_per_value >= bits_per_value {
485 return None;
486 }
487
488 let elems_per_chunk = 1024usize;
489 let num_chunks = num_values.div_ceil(elems_per_chunk);
490 let words_per_chunk = (elems_per_chunk * compressed_bits_per_value).div_ceil(bits_per_value);
491 let last_chunk_is_runt = !num_values.is_multiple_of(elems_per_chunk);
492
493 if !last_chunk_is_runt {
494 return Some((num_chunks * words_per_chunk) as u128);
495 }
496
497 let num_whole_chunks = num_chunks - 1;
498 let remaining_items = num_values - num_whole_chunks * elems_per_chunk;
499 let tail_bit_savings = bits_per_value - compressed_bits_per_value;
500 let padding_cost = compressed_bits_per_value * (elems_per_chunk - remaining_items);
501 let tail_pack_savings = tail_bit_savings * remaining_items;
502 let tail_words = if padding_cost < tail_pack_savings {
503 words_per_chunk
504 } else {
505 remaining_items
506 };
507
508 Some((num_whole_chunks * words_per_chunk + tail_words) as u128)
509}
510
511fn maybe_wrap_general_for_mini_block(
512 inner: Box<dyn MiniBlockCompressor>,
513 params: &CompressionFieldParams,
514) -> Result<Box<dyn MiniBlockCompressor>> {
515 match params.compression.as_deref() {
516 None | Some("none") | Some("fsst") => Ok(inner),
517 Some(raw) => {
518 let scheme = CompressionScheme::from_str(raw)
519 .map_err(|_| Error::invalid_input(format!("Unknown compression scheme: {raw}")))?;
520 let cfg = CompressionConfig::new(scheme, params.compression_level);
521 Ok(Box::new(GeneralMiniBlockCompressor::new(inner, cfg)))
522 }
523 }
524}
525
526fn try_general_compression(
527 field_params: &CompressionFieldParams,
528 data: &DataBlock,
529) -> Result<Option<(Box<dyn BlockCompressor>, CompressionConfig)>> {
530 if field_params.compression.as_deref() == Some("none") {
532 return Ok(None);
533 }
534
535 if let Some(compression_scheme) = &field_params.compression {
538 let scheme: CompressionScheme = compression_scheme.parse()?;
539 let config = CompressionConfig::new(scheme, field_params.compression_level);
540 let compressor = Box::new(CompressedBufferEncoder::try_new(config)?);
541 return Ok(Some((compressor, config)));
542 }
543
544 if data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION {
546 let compressor = Box::new(CompressedBufferEncoder::default());
547 let config = compressor.compressor.config();
548 return Ok(Some((compressor, config)));
549 }
550
551 Ok(None)
552}
553
554pub fn field_metadata_params(field: &Field) -> CompressionFieldParams {
556 let mut params = CompressionFieldParams::default();
557
558 if let Some(compression) = field.metadata.get(COMPRESSION_META_KEY) {
559 params.compression = Some(compression.clone());
560 }
561 if let Some(level) = field.metadata.get(COMPRESSION_LEVEL_META_KEY) {
562 params.compression_level = level.parse().ok();
563 }
564 if let Some(threshold) = field.metadata.get(RLE_THRESHOLD_META_KEY) {
565 params.rle_threshold = threshold.parse().ok();
566 }
567 if let Some(bss_str) = field.metadata.get(BSS_META_KEY) {
568 match BssMode::parse(bss_str) {
569 Some(mode) => params.bss = Some(mode),
570 None => log::warn!("Invalid BSS mode '{}', using default", bss_str),
571 }
572 }
573 if let Some(minichunk_size_str) = field
574 .metadata
575 .get(super::constants::MINICHUNK_SIZE_META_KEY)
576 {
577 if let Ok(minichunk_size) = minichunk_size_str.parse::<i64>() {
578 params.minichunk_size = Some(minichunk_size);
579 } else {
580 log::warn!("Invalid minichunk_size '{}', skipping", minichunk_size_str);
581 }
582 }
583
584 params
585}
586
587pub fn finalize_miniblock_compressor(
589 data: &DataBlock,
590 compressor: Box<dyn MiniBlockCompressor>,
591 params: &CompressionFieldParams,
592) -> Result<Box<dyn MiniBlockCompressor>> {
593 if matches!(data, DataBlock::FixedWidth(_)) {
594 maybe_wrap_general_for_mini_block(compressor, params)
595 } else {
596 Ok(compressor)
597 }
598}
599
600pub fn try_uncompressed_fixed_width_miniblock(
602 data: &DataBlock,
603 params: &CompressionFieldParams,
604) -> Option<Box<dyn MiniBlockCompressor>> {
605 (matches!(data, DataBlock::FixedWidth(_)) && params.compression.as_deref() == Some("none"))
606 .then(|| Box::new(ValueEncoder::default()) as _)
607}
608
609pub fn try_byte_stream_split_miniblock(
611 data: &DataBlock,
612 params: &CompressionFieldParams,
613) -> Option<Box<dyn MiniBlockCompressor>> {
614 let DataBlock::FixedWidth(data) = data else {
615 return None;
616 };
617 try_bss_for_mini_block(data, params)
618}
619
620pub fn try_fixed_u8_rle_miniblock(
622 data: &DataBlock,
623 params: &CompressionFieldParams,
624) -> Option<Box<dyn MiniBlockCompressor>> {
625 let DataBlock::FixedWidth(data) = data else {
626 return None;
627 };
628 try_fixed_u8_rle_for_mini_block(data, params)
629}
630
631pub fn try_child_rle_miniblock(
633 data: &DataBlock,
634 params: &CompressionFieldParams,
635) -> Option<Box<dyn MiniBlockCompressor>> {
636 let DataBlock::FixedWidth(data) = data else {
637 return None;
638 };
639 try_child_rle_for_mini_block(data, params)
640}
641
642pub fn try_bitpacking_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
644 let DataBlock::FixedWidth(data) = data else {
645 return None;
646 };
647 try_bitpack_for_mini_block(data)
648}
649
650pub fn try_raw_fixed_width_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
652 matches!(data, DataBlock::FixedWidth(_)).then(|| Box::new(ValueEncoder::default()) as _)
653}
654
655pub fn try_variable_width_miniblock(
657 field: &Field,
658 data: &DataBlock,
659 params: &CompressionFieldParams,
660) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
661 let DataBlock::VariableWidth(data) = data else {
662 return Ok(None);
663 };
664 if data.bits_per_offset != 32 && data.bits_per_offset != 64 {
665 return Err(Error::invalid_input(format!(
666 "Variable width compression not supported for {} bit offsets",
667 data.bits_per_offset
668 )));
669 }
670
671 let compression = params.compression.as_deref();
672 let data_size = data.expect_single_stat::<UInt64Type>(Stat::DataSize);
673 let max_len = data.expect_single_stat::<UInt64Type>(Stat::MaxLength);
674 if compression == Some("none") {
675 return Ok(Some(Box::new(BinaryMiniBlockEncoder::new(
676 params.minichunk_size,
677 ))));
678 }
679
680 let use_fsst = compression == Some("fsst")
681 || (compression.is_none()
682 && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary)
683 && max_len >= FSST_LEAST_INPUT_MAX_LENGTH
684 && data_size >= FSST_LEAST_INPUT_SIZE as u64);
685 let mut encoder: Box<dyn MiniBlockCompressor> = if use_fsst {
686 Box::new(FsstMiniBlockEncoder::new(params.minichunk_size))
687 } else {
688 Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size))
689 };
690 if let Some(compression_scheme) = compression.filter(|scheme| *scheme != "fsst") {
691 let scheme: CompressionScheme = compression_scheme.parse()?;
692 let config = CompressionConfig::new(scheme, params.compression_level);
693 encoder = Box::new(GeneralMiniBlockCompressor::new(encoder, config));
694 }
695 Ok(Some(encoder))
696}
697
698pub fn try_fixed_packed_struct_miniblock(
700 data: &DataBlock,
701) -> Result<Option<Box<dyn MiniBlockCompressor>>> {
702 let DataBlock::Struct(data) = data else {
703 return Ok(None);
704 };
705 if data.has_variable_width_child() {
706 return Err(Error::invalid_input(
707 "Packed struct mini-block encoding supports only fixed-width children",
708 ));
709 }
710 Ok(Some(Box::new(
711 PackedStructFixedWidthMiniBlockEncoder::default(),
712 )))
713}
714
715pub fn try_raw_fixed_size_list_miniblock(data: &DataBlock) -> Option<Box<dyn MiniBlockCompressor>> {
717 matches!(data, DataBlock::FixedSizeList(_)).then(|| Box::new(ValueEncoder::default()) as _)
718}
719
720pub fn try_raw_per_value(data: &DataBlock) -> Option<Box<dyn PerValueCompressor>> {
722 matches!(data, DataBlock::FixedWidth(_) | DataBlock::FixedSizeList(_))
723 .then(|| Box::new(ValueEncoder::default()) as _)
724}
725
726fn validate_packed_struct(field: &Field, data: &DataBlock) -> Result<Option<bool>> {
727 let DataBlock::Struct(data) = data else {
728 return Ok(None);
729 };
730 if field.children.len() != data.children.len() {
731 return Err(Error::invalid_input(
732 "Struct field metadata does not match data block children",
733 ));
734 }
735 Ok(Some(data.has_variable_width_child()))
736}
737
738pub fn reject_packed_struct_per_value(
740 field: &Field,
741 data: &DataBlock,
742) -> Result<Option<Box<dyn PerValueCompressor>>> {
743 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
744 return Ok(None);
745 };
746 if has_variable_child {
747 return Err(Error::not_supported_source(
748 "Variable packed struct encoding is not enabled by the selected file format".into(),
749 ));
750 }
751 Err(Error::invalid_input(
752 "Packed struct per-value compression should not be used for fixed-width-only structs",
753 ))
754}
755
756pub fn try_variable_packed_struct_per_value(
758 strategy: Arc<dyn CompressionStrategy>,
759 field: &Field,
760 data: &DataBlock,
761) -> Result<Option<Box<dyn PerValueCompressor>>> {
762 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
763 return Ok(None);
764 };
765 if !has_variable_child {
766 return Err(Error::invalid_input(
767 "Packed struct per-value compression should not be used for fixed-width-only structs",
768 ));
769 }
770 Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new(
771 strategy,
772 field.children.clone(),
773 ))))
774}
775
776pub fn try_packed_struct_per_value(
778 strategy: Arc<dyn CompressionStrategy>,
779 field: &Field,
780 data: &DataBlock,
781) -> Result<Option<Box<dyn PerValueCompressor>>> {
782 let Some(has_variable_child) = validate_packed_struct(field, data)? else {
783 return Ok(None);
784 };
785 if has_variable_child {
786 return Ok(Some(Box::new(PackedStructVariablePerValueEncoder::new(
787 strategy,
788 field.children.clone(),
789 ))));
790 }
791
792 Ok(Some(Box::new(PackedStructFixedPerValueEncoder::new(
793 field.children.clone(),
794 ))))
795}
796
797pub fn try_variable_width_per_value(
800 field: &Field,
801 data: &DataBlock,
802 params: &CompressionFieldParams,
803) -> Result<Option<Box<dyn PerValueCompressor>>> {
804 let DataBlock::VariableWidth(data) = data else {
805 return Ok(None);
806 };
807 let compression = params.compression.as_deref();
808 if compression == Some("none") {
809 return Ok(Some(Box::new(VariableEncoder::default())));
810 }
811
812 let max_len = data.expect_single_stat::<UInt64Type>(Stat::MaxLength);
813 let data_size = data.expect_single_stat::<UInt64Type>(Stat::DataSize);
814 let per_value_requested = compression.is_some_and(|compression| compression != "fsst");
815 if (max_len > 32 * 1024 || per_value_requested) && data_size >= FSST_LEAST_INPUT_SIZE as u64 {
816 if compression == Some("zstd") {
817 let config = CompressionConfig::new(CompressionScheme::Zstd, params.compression_level);
818 return Ok(Some(Box::new(CompressedBufferEncoder::try_new(config)?)));
819 }
820 return Ok(Some(Box::new(CompressedBufferEncoder::default())));
821 }
822
823 if data.bits_per_offset != 32 && data.bits_per_offset != 64 {
824 return Err(Error::invalid_input(format!(
825 "Per-value compression does not support variable-width data with {}-bit offsets",
826 data.bits_per_offset
827 )));
828 }
829 let encoder = Box::new(VariableEncoder::default());
830 let use_fsst = compression == Some("fsst")
831 || (compression.is_none()
832 && !matches!(field.data_type(), DataType::Binary | DataType::LargeBinary)
833 && max_len >= FSST_LEAST_INPUT_MAX_LENGTH
834 && data_size >= FSST_LEAST_INPUT_SIZE as u64);
835 Ok(Some(if use_fsst {
836 Box::new(FsstPerValueEncoder::new(encoder))
837 } else {
838 encoder
839 }))
840}
841
842pub fn try_fixed_u8_rle_block(
844 data: &DataBlock,
845 params: &CompressionFieldParams,
846) -> Result<Option<Box<dyn BlockCompressor>>> {
847 let DataBlock::FixedWidth(data) = data else {
848 return Ok(None);
849 };
850 try_fixed_u8_rle_for_block(data, params)
851}
852
853pub fn try_variable_rle_block(
855 data: &DataBlock,
856 params: &CompressionFieldParams,
857) -> Result<Option<Box<dyn BlockCompressor>>> {
858 let DataBlock::FixedWidth(data) = data else {
859 return Ok(None);
860 };
861 try_variable_rle_for_block(data, params)
862}
863
864pub fn try_bitpacking_block(data: &DataBlock) -> Option<Box<dyn BlockCompressor>> {
866 let DataBlock::FixedWidth(data) = data else {
867 return None;
868 };
869 try_bitpack_for_block(data)
870}
871
872pub fn try_general_block(
874 data: &DataBlock,
875 params: &CompressionFieldParams,
876) -> Result<Option<Box<dyn BlockCompressor>>> {
877 let Some((compressor, _config)) = try_general_compression(params, data)? else {
878 return Ok(None);
879 };
880 Ok(Some(compressor))
881}
882
883pub fn try_raw_block(data: &DataBlock) -> Option<Box<dyn BlockCompressor>> {
885 match data {
886 DataBlock::FixedWidth(_) => {
887 Some(Box::new(ValueEncoder::default()) as Box<dyn BlockCompressor>)
888 }
889 DataBlock::VariableWidth(_) => {
890 Some(Box::new(VariableEncoder::default()) as Box<dyn BlockCompressor>)
891 }
892 _ => None,
893 }
894}
895
896pub trait MiniBlockDecompressor: std::fmt::Debug + Send + Sync {
897 fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock>;
898
899 fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
906 None
907 }
908}
909
910pub trait FixedPerValueDecompressor: std::fmt::Debug + Send + Sync {
911 fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result<DataBlock>;
913 fn bits_per_value(&self) -> u64;
917
918 fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
925 None
926 }
927}
928
929pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync {
930 fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock>;
932}
933
934pub trait BlockDecompressor: std::fmt::Debug + Send + Sync {
935 fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock>;
936
937 fn requires_payload(&self) -> bool {
939 true
940 }
941
942 fn infer_num_values(&self, _data: &LanceBuffer) -> Result<Option<u64>> {
947 Ok(None)
948 }
949}
950
951pub(crate) fn require_block_payload(data: Option<LanceBuffer>, codec: &str) -> Result<LanceBuffer> {
952 data.ok_or_else(|| Error::invalid_input(format!("{codec} requires one payload")))
953}
954
955pub(crate) fn require_no_block_payload(data: Option<LanceBuffer>, codec: &str) -> Result<()> {
956 if data.is_some() {
957 return Err(Error::invalid_input(format!("{codec} expects no payload")));
958 }
959 Ok(())
960}
961
962pub trait DecompressionStrategy: std::fmt::Debug + Send + Sync {
963 fn create_miniblock_decompressor(
964 &self,
965 description: &CompressiveEncoding,
966 decompression_strategy: &dyn DecompressionStrategy,
967 ) -> Result<Box<dyn MiniBlockDecompressor>>;
968
969 fn create_fixed_per_value_decompressor(
970 &self,
971 description: &CompressiveEncoding,
972 ) -> Result<Box<dyn FixedPerValueDecompressor>>;
973
974 fn create_variable_per_value_decompressor(
975 &self,
976 description: &CompressiveEncoding,
977 ) -> Result<Box<dyn VariablePerValueDecompressor>>;
978
979 fn create_block_decompressor(
980 &self,
981 description: &CompressiveEncoding,
982 ) -> Result<Box<dyn BlockDecompressor>>;
983}
984
985#[derive(Debug, Default)]
986pub struct DefaultDecompressionStrategy {}
987
988impl DecompressionStrategy for DefaultDecompressionStrategy {
989 fn create_miniblock_decompressor(
990 &self,
991 description: &CompressiveEncoding,
992 decompression_strategy: &dyn DecompressionStrategy,
993 ) -> Result<Box<dyn MiniBlockDecompressor>> {
994 match description.compression.as_ref().unwrap() {
995 Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))),
996 #[cfg(feature = "bitpacking")]
997 Compression::InlineBitpacking(description) => {
998 Ok(Box::new(InlineBitpacking::from_description(description)))
999 }
1000 #[cfg(not(feature = "bitpacking"))]
1001 Compression::InlineBitpacking(_) => Err(Error::not_supported_source(
1002 "this runtime was not built with bitpacking support".into(),
1003 )),
1004 Compression::Variable(variable) => {
1005 let Compression::Flat(offsets) = variable
1006 .offsets
1007 .as_ref()
1008 .unwrap()
1009 .compression
1010 .as_ref()
1011 .unwrap()
1012 else {
1013 panic!("Variable compression only supports flat offsets")
1014 };
1015 Ok(Box::new(BinaryMiniBlockDecompressor::new(
1016 offsets.bits_per_value as u8,
1017 )))
1018 }
1019 Compression::Fsst(description) => {
1020 let inner_decompressor = decompression_strategy.create_miniblock_decompressor(
1021 description.values.as_ref().unwrap(),
1022 decompression_strategy,
1023 )?;
1024 Ok(Box::new(FsstMiniBlockDecompressor::new(
1025 description,
1026 inner_decompressor,
1027 )))
1028 }
1029 Compression::PackedStruct(description) => Ok(Box::new(
1030 PackedStructFixedWidthMiniBlockDecompressor::new(description),
1031 )),
1032 Compression::VariablePackedStruct(_) => Err(Error::not_supported_source(
1033 "variable packed struct decoding is not yet implemented".into(),
1034 )),
1035 Compression::FixedSizeList(fsl) => {
1036 Ok(Box::new(ValueDecompressor::from_fsl(fsl)?))
1039 }
1040 Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(
1041 rle,
1042 decompression_strategy,
1043 )?)),
1044 Compression::ByteStreamSplit(bss) => {
1045 let Compression::Flat(values) =
1046 bss.values.as_ref().unwrap().compression.as_ref().unwrap()
1047 else {
1048 panic!("ByteStreamSplit compression only supports flat values")
1049 };
1050 Ok(Box::new(ByteStreamSplitDecompressor::new(
1051 values.bits_per_value as usize,
1052 )))
1053 }
1054 Compression::General(general) => {
1055 let inner_decompressor = self.create_miniblock_decompressor(
1057 general.values.as_ref().ok_or_else(|| {
1058 Error::invalid_input("GeneralMiniBlock missing inner encoding")
1059 })?,
1060 decompression_strategy,
1061 )?;
1062
1063 let compression = general.compression.as_ref().ok_or_else(|| {
1065 Error::invalid_input("GeneralMiniBlock missing compression config")
1066 })?;
1067
1068 let scheme = compression.scheme().try_into()?;
1069
1070 let compression_config = CompressionConfig::new(scheme, compression.level);
1071
1072 Ok(Box::new(GeneralMiniBlockDecompressor::new(
1073 inner_decompressor,
1074 compression_config,
1075 )))
1076 }
1077 _ => todo!(),
1078 }
1079 }
1080
1081 fn create_fixed_per_value_decompressor(
1082 &self,
1083 description: &CompressiveEncoding,
1084 ) -> Result<Box<dyn FixedPerValueDecompressor>> {
1085 match description.compression.as_ref().unwrap() {
1086 Compression::Constant(constant) => Ok(Box::new(ConstantDecompressor::new(
1087 constant
1088 .value
1089 .as_ref()
1090 .map(|v| LanceBuffer::from_bytes(v.clone(), 1)),
1091 ))),
1092 Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))),
1093 Compression::FixedSizeList(fsl) => Ok(Box::new(ValueDecompressor::from_fsl(fsl)?)),
1094 Compression::PackedStruct(description) => Ok(Box::new(
1095 PackedStructFixedPerValueDecompressor::new(description)?,
1096 )),
1097 _ => todo!("fixed-per-value decompressor for {:?}", description),
1098 }
1099 }
1100
1101 fn create_variable_per_value_decompressor(
1102 &self,
1103 description: &CompressiveEncoding,
1104 ) -> Result<Box<dyn VariablePerValueDecompressor>> {
1105 match description.compression.as_ref().unwrap() {
1106 Compression::Variable(variable) => {
1107 let Compression::Flat(offsets) = variable
1108 .offsets
1109 .as_ref()
1110 .unwrap()
1111 .compression
1112 .as_ref()
1113 .unwrap()
1114 else {
1115 panic!("Variable compression only supports flat offsets")
1116 };
1117 assert!(offsets.bits_per_value < u8::MAX as u64);
1118 Ok(Box::new(VariableDecoder::default()))
1119 }
1120 Compression::Fsst(fsst) => Ok(Box::new(FsstPerValueDecompressor::new(
1121 LanceBuffer::from_bytes(fsst.symbol_table.clone(), 1),
1122 Box::new(VariableDecoder::default()),
1123 ))),
1124 Compression::General(general) => Ok(Box::new(CompressedBufferEncoder::from_scheme(
1125 general.compression.as_ref().expect_ok()?.scheme(),
1126 )?)),
1127 Compression::VariablePackedStruct(description) => {
1128 let mut fields = Vec::with_capacity(description.fields.len());
1129 for field in &description.fields {
1130 let value_encoding = field.value.as_ref().ok_or_else(|| {
1131 Error::invalid_input("VariablePackedStruct field is missing value encoding")
1132 })?;
1133 let decoder = match field.layout.as_ref().ok_or_else(|| {
1134 Error::invalid_input("VariablePackedStruct field is missing layout details")
1135 })? {
1136 crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerValue(
1137 bits_per_value,
1138 ) => {
1139 let decompressor =
1140 self.create_fixed_per_value_decompressor(value_encoding)?;
1141 VariablePackedStructFieldDecoder {
1142 kind: VariablePackedStructFieldKind::Fixed {
1143 bits_per_value: *bits_per_value,
1144 decompressor: Arc::from(decompressor),
1145 },
1146 }
1147 }
1148 crate::format::pb21::variable_packed_struct::field_encoding::Layout::BitsPerLength(
1149 bits_per_length,
1150 ) => {
1151 let decompressor =
1152 self.create_variable_per_value_decompressor(value_encoding)?;
1153 VariablePackedStructFieldDecoder {
1154 kind: VariablePackedStructFieldKind::Variable {
1155 bits_per_length: *bits_per_length,
1156 decompressor: Arc::from(decompressor),
1157 },
1158 }
1159 }
1160 };
1161 fields.push(decoder);
1162 }
1163 Ok(Box::new(PackedStructVariablePerValueDecompressor::new(
1164 fields,
1165 )))
1166 }
1167 _ => todo!("variable-per-value decompressor for {:?}", description),
1168 }
1169 }
1170
1171 fn create_block_decompressor(
1172 &self,
1173 description: &CompressiveEncoding,
1174 ) -> Result<Box<dyn BlockDecompressor>> {
1175 match description.compression.as_ref().unwrap() {
1176 Compression::InlineBitpacking(inline_bitpacking) => Ok(Box::new(
1177 InlineBitpacking::from_description(inline_bitpacking),
1178 )),
1179 Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))),
1180 Compression::Constant(constant) => {
1181 let scalar = constant
1182 .value
1183 .as_ref()
1184 .map(|v| LanceBuffer::from_bytes(v.clone(), 1));
1185 Ok(Box::new(ConstantDecompressor::new(scalar)))
1186 }
1187 Compression::Variable(_) => Ok(Box::new(BinaryBlockDecompressor::default())),
1188 Compression::FixedSizeList(fsl) => {
1189 Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref())?))
1190 }
1191 Compression::OutOfLineBitpacking(out_of_line) => {
1192 let compressed_bit_width = match out_of_line
1194 .values
1195 .as_ref()
1196 .unwrap()
1197 .compression
1198 .as_ref()
1199 .unwrap()
1200 {
1201 Compression::Flat(flat) => flat.bits_per_value,
1202 _ => {
1203 return Err(Error::invalid_input_source(
1204 "OutOfLineBitpacking values must use Flat encoding".into(),
1205 ));
1206 }
1207 };
1208 Ok(Box::new(OutOfLineBitpacking::new(
1209 compressed_bit_width,
1210 out_of_line.uncompressed_bits_per_value,
1211 )))
1212 }
1213 Compression::General(general) => {
1214 let inner_desc = general
1215 .values
1216 .as_ref()
1217 .ok_or_else(|| {
1218 Error::invalid_input("General compression missing inner encoding")
1219 })?
1220 .as_ref();
1221 let inner_decompressor = self.create_block_decompressor(inner_desc)?;
1222
1223 let compression = general.compression.as_ref().ok_or_else(|| {
1224 Error::invalid_input("General compression missing compression config")
1225 })?;
1226 let scheme = compression.scheme().try_into()?;
1227 let config = CompressionConfig::new(scheme, compression.level);
1228 let general_decompressor =
1229 GeneralBlockDecompressor::try_new(inner_decompressor, config)?;
1230
1231 Ok(Box::new(general_decompressor))
1232 }
1233 Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)),
1234 _ => todo!(),
1235 }
1236 }
1237}
1238pub(crate) fn create_rle_decompressor(
1239 rle: &crate::format::pb21::Rle,
1240 decompression_strategy: &dyn DecompressionStrategy,
1241) -> Result<RleDecompressor> {
1242 let values = rle
1243 .values
1244 .as_ref()
1245 .ok_or_else(|| Error::invalid_input("RLE compression missing values encoding"))?;
1246 let run_lengths = rle
1247 .run_lengths
1248 .as_ref()
1249 .ok_or_else(|| Error::invalid_input("RLE compression missing run lengths encoding"))?;
1250
1251 let values = create_rle_child_decompressor(values, "values", decompression_strategy)?;
1252 let run_lengths =
1253 create_rle_child_decompressor(run_lengths, "run lengths", decompression_strategy)?;
1254
1255 if !matches!(values.bits_per_value(), 8 | 16 | 32 | 64) {
1256 return Err(Error::invalid_input(format!(
1257 "RLE compression only supports 8, 16, 32, or 64-bit values, got {}",
1258 values.bits_per_value()
1259 )));
1260 }
1261
1262 let run_length_width =
1263 RunLengthWidth::from_bits(run_lengths.bits_per_value()).ok_or_else(|| {
1264 Error::invalid_input(format!(
1265 "RLE compression only supports 8, 16, or 32-bit run lengths, got {}",
1266 run_lengths.bits_per_value()
1267 ))
1268 })?;
1269
1270 if values.requires_num_values() && run_lengths.requires_num_values() {
1271 return Err(Error::invalid_input(
1272 "RLE values and run lengths child encodings cannot both require the run count",
1273 ));
1274 }
1275
1276 if values.is_identity() && run_lengths.is_identity() {
1277 return Ok(RleDecompressor::with_run_length_width(
1278 values.bits_per_value(),
1279 run_length_width,
1280 ));
1281 }
1282
1283 Ok(RleDecompressor::with_child_decompressors(
1284 values.bits_per_value(),
1285 run_length_width,
1286 values,
1287 run_lengths,
1288 ))
1289}
1290
1291fn create_rle_child_decompressor(
1292 encoding: &CompressiveEncoding,
1293 role: &str,
1294 decompression_strategy: &dyn DecompressionStrategy,
1295) -> Result<RleChildDecompressor> {
1296 let compression = encoding
1297 .compression
1298 .as_ref()
1299 .ok_or_else(|| Error::invalid_input(format!("RLE {role} missing child compression")))?;
1300 let (bits_per_value, requires_num_values, needs_decompressor) =
1301 validate_rle_child_compression(compression, role)?;
1302
1303 if needs_decompressor {
1304 Ok(RleChildDecompressor::block(
1305 bits_per_value,
1306 decompression_strategy.create_block_decompressor(encoding)?,
1307 requires_num_values,
1308 ))
1309 } else {
1310 Ok(RleChildDecompressor::flat(bits_per_value))
1311 }
1312}
1313
1314fn validate_rle_child_compression(
1315 compression: &Compression,
1316 role: &str,
1317) -> Result<(u64, bool, bool)> {
1318 match compression {
1319 Compression::Flat(flat) => Ok((flat.bits_per_value, false, false)),
1320 Compression::General(general) => {
1321 general.compression.as_ref().ok_or_else(|| {
1322 Error::invalid_input(format!(
1323 "RLE {role} general child missing compression config"
1324 ))
1325 })?;
1326 let values = general.values.as_ref().ok_or_else(|| {
1327 Error::invalid_input(format!("RLE {role} general child missing inner encoding"))
1328 })?;
1329 let inner = values.compression.as_ref().ok_or_else(|| {
1330 Error::invalid_input(format!(
1331 "RLE {role} general child missing inner compression"
1332 ))
1333 })?;
1334 let (bits_per_value, requires_num_values) =
1335 validate_rle_block_child_inner(inner, role)?;
1336 Ok((bits_per_value, requires_num_values, true))
1337 }
1338 Compression::OutOfLineBitpacking(out_of_line) => {
1339 let values = out_of_line.values.as_ref().ok_or_else(|| {
1340 Error::invalid_input(format!(
1341 "RLE {role} bitpacking child missing values encoding"
1342 ))
1343 })?;
1344 let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| {
1345 Error::invalid_input(format!(
1346 "RLE {role} bitpacking child missing values compression"
1347 ))
1348 })?
1349 else {
1350 return Err(Error::invalid_input(format!(
1351 "RLE {role} bitpacking child only supports flat values"
1352 )));
1353 };
1354 Ok((out_of_line.uncompressed_bits_per_value, true, true))
1355 }
1356 other => Err(Error::invalid_input(format!(
1357 "RLE {role} only supports flat, general, or out-of-line bitpacking child encodings, got {}",
1358 compression_name(other)
1359 ))),
1360 }
1361}
1362
1363fn validate_rle_block_child_inner(compression: &Compression, role: &str) -> Result<(u64, bool)> {
1364 match compression {
1365 Compression::Flat(flat) => Ok((flat.bits_per_value, false)),
1366 Compression::OutOfLineBitpacking(out_of_line) => {
1367 let values = out_of_line.values.as_ref().ok_or_else(|| {
1368 Error::invalid_input(format!(
1369 "RLE {role} bitpacking child missing values encoding"
1370 ))
1371 })?;
1372 let Compression::Flat(_) = values.compression.as_ref().ok_or_else(|| {
1373 Error::invalid_input(format!(
1374 "RLE {role} bitpacking child missing values compression"
1375 ))
1376 })?
1377 else {
1378 return Err(Error::invalid_input(format!(
1379 "RLE {role} bitpacking child only supports flat values"
1380 )));
1381 };
1382 Ok((out_of_line.uncompressed_bits_per_value, true))
1383 }
1384 other => Err(Error::invalid_input(format!(
1385 "RLE {role} general child only supports flat or out-of-line bitpacking inner encodings, got {}",
1386 compression_name(other)
1387 ))),
1388 }
1389}
1390
1391fn compression_name(compression: &Compression) -> &'static str {
1392 match compression {
1393 Compression::Flat(_) => "flat",
1394 Compression::Variable(_) => "variable",
1395 Compression::Fsst(_) => "fsst",
1396 Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking",
1397 Compression::InlineBitpacking(_) => "inline bitpacking",
1398 Compression::General(_) => "general",
1399 Compression::Constant(_) => "constant",
1400 Compression::Dictionary(_) => "dictionary",
1401 Compression::ByteStreamSplit(_) => "byte stream split",
1402 Compression::PackedStruct(_) => "packed struct",
1403 Compression::FixedSizeList(_) => "fixed-size list",
1404 Compression::VariablePackedStruct(_) => "variable packed struct",
1405 Compression::Rle(_) => "rle",
1406 }
1407}
1408
1409#[cfg(test)]
1410mod tests {
1411 use super::*;
1412 use crate::buffer::LanceBuffer;
1413 use crate::compression_config::CompressionParams;
1414 use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock};
1415 use crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext;
1416 use crate::statistics::ComputeStat;
1417 use crate::testing::{TestEncoding, extract_array_encoding_chain, test_compression_strategy};
1418 use arrow_schema::{DataType, Field as ArrowField};
1419 use std::collections::HashMap;
1420
1421 fn strategy(encoding: TestEncoding, params: CompressionParams) -> Arc<dyn CompressionStrategy> {
1422 test_compression_strategy(encoding, params)
1423 }
1424
1425 fn baseline_strategy(params: CompressionParams) -> Arc<dyn CompressionStrategy> {
1426 strategy(TestEncoding::StructuralU16, params)
1427 }
1428
1429 fn selected_block_codec(
1430 strategy: &Arc<dyn CompressionStrategy>,
1431 field: &Field,
1432 data: &DataBlock,
1433 ) -> (Box<dyn BlockCompressor>, CompressiveEncoding) {
1434 let compressor = strategy.create_block_compressor(field, data).unwrap();
1435 let (_, encoding) = compressor.compress(data.clone()).unwrap();
1436 (compressor, encoding)
1437 }
1438
1439 fn miniblock_context() -> MiniBlockCompressionContext {
1440 MiniBlockCompressionContext::new(0, true, true)
1441 }
1442
1443 fn create_test_field(name: &str, data_type: DataType) -> Field {
1444 let arrow_field = ArrowField::new(name, data_type, true);
1445 let mut field = Field::try_from(&arrow_field).unwrap();
1446 field.id = -1;
1447 field
1448 }
1449
1450 fn create_fixed_width_block_with_stats(
1451 bits_per_value: u64,
1452 num_values: u64,
1453 run_count: u64,
1454 ) -> DataBlock {
1455 let bytes_per_value = (bits_per_value / 8) as usize;
1457 let total_bytes = bytes_per_value * num_values as usize;
1458 let mut data = vec![0u8; total_bytes];
1459
1460 let values_per_run = (num_values / run_count).max(1);
1462 let mut run_value = 0u8;
1463
1464 for i in 0..num_values as usize {
1465 if i % values_per_run as usize == 0 {
1466 run_value = run_value.wrapping_add(17); }
1468 for j in 0..bytes_per_value {
1470 let byte_offset = i * bytes_per_value + j;
1471 if byte_offset < data.len() {
1472 data[byte_offset] = run_value.wrapping_add(j as u8);
1473 }
1474 }
1475 }
1476
1477 let mut block = FixedWidthDataBlock {
1478 bits_per_value,
1479 data: LanceBuffer::reinterpret_vec(data),
1480 num_values,
1481 block_info: BlockInfo::default(),
1482 };
1483
1484 use crate::statistics::ComputeStat;
1486 block.compute_stat();
1487
1488 DataBlock::FixedWidth(block)
1489 }
1490
1491 fn create_fixed_width_block(bits_per_value: u64, num_values: u64) -> DataBlock {
1492 let bytes_per_value = (bits_per_value / 8) as usize;
1494 let total_bytes = bytes_per_value * num_values as usize;
1495 let mut data = vec![0u8; total_bytes];
1496
1497 for i in 0..num_values as usize {
1499 let byte_offset = i * bytes_per_value;
1500 if byte_offset < data.len() {
1501 data[byte_offset] = (i % 256) as u8;
1502 }
1503 }
1504
1505 let mut block = FixedWidthDataBlock {
1506 bits_per_value,
1507 data: LanceBuffer::reinterpret_vec(data),
1508 num_values,
1509 block_info: BlockInfo::default(),
1510 };
1511
1512 use crate::statistics::ComputeStat;
1514 block.compute_stat();
1515
1516 DataBlock::FixedWidth(block)
1517 }
1518
1519 fn rle_run_length_bits(encoding: &CompressiveEncoding) -> u64 {
1520 let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else {
1521 panic!("expected RLE encoding");
1522 };
1523 let Compression::Flat(run_lengths) = rle
1524 .run_lengths
1525 .as_ref()
1526 .unwrap()
1527 .compression
1528 .as_ref()
1529 .unwrap()
1530 else {
1531 panic!("expected flat run lengths");
1532 };
1533 run_lengths.bits_per_value
1534 }
1535
1536 fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle {
1537 match encoding.compression.as_ref().unwrap() {
1538 Compression::Rle(rle) => rle,
1539 Compression::General(general) => {
1540 let inner = general.values.as_ref().unwrap();
1541 let Compression::Rle(rle) = inner.compression.as_ref().unwrap() else {
1542 panic!("expected wrapped RLE encoding");
1543 };
1544 rle
1545 }
1546 other => panic!("expected RLE encoding, got {}", compression_name(other)),
1547 }
1548 }
1549
1550 fn create_variable_width_block(
1551 bits_per_offset: u8,
1552 num_values: u64,
1553 avg_value_size: usize,
1554 ) -> DataBlock {
1555 use crate::statistics::ComputeStat;
1556
1557 let mut offsets = Vec::with_capacity((num_values + 1) as usize);
1559 let mut current_offset = 0i64;
1560 offsets.push(current_offset);
1561
1562 for i in 0..num_values {
1564 let value_size = if avg_value_size == 0 {
1565 1
1566 } else {
1567 ((avg_value_size as i64 + (i as i64 % 8) - 4).max(1) as usize)
1568 .min(avg_value_size * 2)
1569 };
1570 current_offset += value_size as i64;
1571 offsets.push(current_offset);
1572 }
1573
1574 let total_data_size = current_offset as usize;
1576 let mut data = vec![0u8; total_data_size];
1577
1578 for i in 0..num_values {
1580 let start_offset = offsets[i as usize] as usize;
1581 let end_offset = offsets[(i + 1) as usize] as usize;
1582
1583 let content = (i % 256) as u8;
1584 for j in 0..end_offset - start_offset {
1585 data[start_offset + j] = content.wrapping_add(j as u8);
1586 }
1587 }
1588
1589 let offsets_buffer = match bits_per_offset {
1591 32 => {
1592 let offsets_32: Vec<i32> = offsets.iter().map(|&o| o as i32).collect();
1593 LanceBuffer::reinterpret_vec(offsets_32)
1594 }
1595 64 => LanceBuffer::reinterpret_vec(offsets),
1596 _ => panic!("Unsupported bits_per_offset: {}", bits_per_offset),
1597 };
1598
1599 let mut block = VariableWidthBlock {
1600 data: LanceBuffer::from(data),
1601 offsets: offsets_buffer,
1602 bits_per_offset,
1603 num_values,
1604 block_info: BlockInfo::default(),
1605 };
1606
1607 block.compute_stat();
1608 DataBlock::VariableWidth(block)
1609 }
1610
1611 fn create_fsst_candidate_variable_width_block() -> DataBlock {
1612 create_variable_width_block(32, 4096, FSST_LEAST_INPUT_MAX_LENGTH as usize + 16)
1613 }
1614
1615 #[test]
1616 fn test_parameter_based_compression() {
1617 let mut params = CompressionParams::new();
1618
1619 params.columns.insert(
1621 "*_id".to_string(),
1622 CompressionFieldParams {
1623 rle_threshold: Some(0.3),
1624 compression: Some("lz4".to_string()),
1625 compression_level: None,
1626 bss: Some(BssMode::Off), minichunk_size: None,
1628 },
1629 );
1630
1631 let strategy = baseline_strategy(params);
1632 let field = create_test_field("user_id", DataType::Int32);
1633
1634 let data = create_fixed_width_block_with_stats(32, 1000, 100); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
1639 let debug_str = format!("{:?}", compressor);
1641
1642 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
1644 assert!(debug_str.contains("RleEncoder"));
1645 }
1646
1647 #[test]
1648 fn test_type_level_parameters() {
1649 let mut params = CompressionParams::new();
1650
1651 params.types.insert(
1653 "Int32".to_string(),
1654 CompressionFieldParams {
1655 rle_threshold: Some(0.1), compression: Some("zstd".to_string()),
1657 compression_level: Some(3),
1658 bss: Some(BssMode::Off), minichunk_size: None,
1660 },
1661 );
1662
1663 let strategy = baseline_strategy(params);
1664 let field = create_test_field("some_column", DataType::Int32);
1665 let data = create_fixed_width_block_with_stats(32, 1000, 50);
1667
1668 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
1669 assert!(format!("{:?}", compressor).contains("RleEncoder"));
1671 }
1672
1673 #[test]
1676 #[cfg(feature = "bitpacking")]
1677 fn test_block_bitpacks_with_zero_segment() {
1678 let strategy = baseline_strategy(CompressionParams::default());
1679 let field = create_test_field("levels", DataType::UInt16);
1680
1681 let mut values: Vec<u16> = vec![0; 1024];
1683 values.extend(std::iter::repeat_n(1u16, 1024));
1684 let mut block = FixedWidthDataBlock {
1685 bits_per_value: 16,
1686 data: LanceBuffer::reinterpret_vec(values),
1687 num_values: 2048,
1688 block_info: BlockInfo::default(),
1689 };
1690 block.compute_stat();
1691 let data = DataBlock::FixedWidth(block);
1692
1693 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
1694 let debug_str = format!("{:?}", compressor);
1695 assert!(
1696 debug_str.contains("OutOfLineBitpacking"),
1697 "expected OutOfLineBitpacking, got: {debug_str}"
1698 );
1699 }
1700
1701 #[rstest::rstest]
1702 #[case::runt_tail(1600)]
1703 #[case::whole_chunks(2048)]
1704 #[cfg(feature = "bitpacking")]
1705 fn test_block_skips_full_width_bitpacking(#[case] num_values: usize) {
1706 let mut block = FixedWidthDataBlock {
1707 bits_per_value: 64,
1708 data: LanceBuffer::reinterpret_vec(vec![-1_i64; num_values]),
1709 num_values: num_values as u64,
1710 block_info: BlockInfo::default(),
1711 };
1712 block.compute_stat();
1713
1714 assert!(try_bitpacking_block(&DataBlock::FixedWidth(block)).is_none());
1715 }
1716
1717 #[test]
1718 fn test_rle_block_accounts_for_header_before_selecting() {
1719 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
1720 let field = create_test_field("small_constant", DataType::Int32);
1721 let values = vec![42i32; 2];
1722 let mut block = FixedWidthDataBlock {
1723 bits_per_value: 32,
1724 data: LanceBuffer::reinterpret_vec(values),
1725 num_values: 2,
1726 block_info: BlockInfo::default(),
1727 };
1728 block.compute_stat();
1729 let data = DataBlock::FixedWidth(block);
1730
1731 let (compressor, encoding) = selected_block_codec(&strategy, &field, &data);
1732
1733 assert!(format!("{compressor:?}").contains("ValueEncoder"));
1734 assert!(matches!(
1735 encoding.compression.as_ref(),
1736 Some(Compression::Flat(_))
1737 ));
1738 }
1739
1740 #[test]
1741 #[cfg(feature = "bitpacking")]
1742 fn test_rle_block_prefers_bitpacking_when_smaller() {
1743 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
1744 let field = create_test_field("levels", DataType::UInt16);
1745
1746 let mut values = Vec::with_capacity(2048);
1747 for run_idx in 0..1024 {
1748 values.extend(std::iter::repeat_n((run_idx % 2) as u16, 2));
1749 }
1750 let mut block = FixedWidthDataBlock {
1751 bits_per_value: 16,
1752 data: LanceBuffer::reinterpret_vec(values),
1753 num_values: 2048,
1754 block_info: BlockInfo::default(),
1755 };
1756 block.compute_stat();
1757 let data = DataBlock::FixedWidth(block);
1758
1759 let (compressor, encoding) = selected_block_codec(&strategy, &field, &data);
1760 let debug_str = format!("{compressor:?}");
1761 assert!(
1762 debug_str.contains("OutOfLineBitpacking"),
1763 "expected OutOfLineBitpacking, got: {debug_str}"
1764 );
1765 assert!(matches!(
1766 encoding.compression.as_ref(),
1767 Some(Compression::OutOfLineBitpacking(_))
1768 ));
1769 }
1770
1771 #[test]
1772 #[cfg(feature = "bitpacking")]
1773 fn test_low_cardinality_prefers_bitpacking_over_rle() {
1774 let strategy = baseline_strategy(CompressionParams::default());
1775 let field = create_test_field("int_score", DataType::Int64);
1776
1777 let mut values: Vec<u64> = Vec::with_capacity(256);
1780 for run_idx in 0..64 {
1781 let value = match run_idx % 3 {
1782 0 => 3u64,
1783 1 => 4u64,
1784 _ => 5u64,
1785 };
1786 values.extend(std::iter::repeat_n(value, 4));
1787 }
1788
1789 let mut block = FixedWidthDataBlock {
1790 bits_per_value: 64,
1791 data: LanceBuffer::reinterpret_vec(values),
1792 num_values: 256,
1793 block_info: BlockInfo::default(),
1794 };
1795
1796 use crate::statistics::ComputeStat;
1797 block.compute_stat();
1798
1799 let data = DataBlock::FixedWidth(block);
1800 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
1801 let debug_str = format!("{:?}", compressor);
1802 assert!(
1803 debug_str.contains("InlineBitpacking"),
1804 "expected InlineBitpacking, got: {debug_str}"
1805 );
1806 assert!(
1807 !debug_str.contains("RleEncoder"),
1808 "expected RLE to be skipped when bitpacking is smaller, got: {debug_str}"
1809 );
1810 }
1811
1812 fn check_uncompressed_encoding(encoding: &CompressiveEncoding, variable: bool) {
1813 let chain = extract_array_encoding_chain(encoding);
1814 if variable {
1815 assert_eq!(chain.len(), 2);
1816 assert_eq!(chain.first().unwrap().as_str(), "variable");
1817 assert_eq!(chain.get(1).unwrap().as_str(), "flat");
1818 } else {
1819 assert_eq!(chain.len(), 1);
1820 assert_eq!(chain.first().unwrap().as_str(), "flat");
1821 }
1822 }
1823
1824 #[test]
1825 fn test_none_compression() {
1826 let mut params = CompressionParams::new();
1827
1828 params.columns.insert(
1830 "embeddings".to_string(),
1831 CompressionFieldParams {
1832 compression: Some("none".to_string()),
1833 ..Default::default()
1834 },
1835 );
1836
1837 let strategy = baseline_strategy(params);
1838 let field = create_test_field("embeddings", DataType::Float32);
1839 let fixed_data = create_fixed_width_block(32, 1000);
1840 let variable_data = create_variable_width_block(32, 10, 32 * 1024);
1841
1842 let compressor = strategy
1844 .create_miniblock_compressor(&field, &fixed_data)
1845 .unwrap();
1846 let (_block, encoding) = compressor
1847 .compress(miniblock_context(), fixed_data.clone())
1848 .unwrap();
1849 check_uncompressed_encoding(&encoding, false);
1850 let compressor = strategy
1851 .create_miniblock_compressor(&field, &variable_data)
1852 .unwrap();
1853 let (_block, encoding) = compressor
1854 .compress(miniblock_context(), variable_data.clone())
1855 .unwrap();
1856 check_uncompressed_encoding(&encoding, true);
1857
1858 let compressor = strategy.create_per_value(&field, &fixed_data).unwrap();
1860 let (_block, encoding) = compressor.compress(fixed_data).unwrap();
1861 check_uncompressed_encoding(&encoding, false);
1862 let compressor = strategy.create_per_value(&field, &variable_data).unwrap();
1863 let (_block, encoding) = compressor.compress(variable_data).unwrap();
1864 check_uncompressed_encoding(&encoding, true);
1865 }
1866
1867 #[test]
1868 fn test_field_metadata_none_compression() {
1869 let mut arrow_field = ArrowField::new("simple_col", DataType::Binary, true);
1871 let mut metadata = HashMap::new();
1872 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
1873 arrow_field = arrow_field.with_metadata(metadata);
1874 let field = Field::try_from(&arrow_field).unwrap();
1875
1876 let strategy = baseline_strategy(CompressionParams::new());
1877
1878 let fixed_data = create_fixed_width_block(32, 1000);
1880 let variable_data = create_variable_width_block(32, 10, 32 * 1024);
1881
1882 let compressor = strategy
1883 .create_miniblock_compressor(&field, &fixed_data)
1884 .unwrap();
1885 let (_block, encoding) = compressor
1886 .compress(miniblock_context(), fixed_data.clone())
1887 .unwrap();
1888 check_uncompressed_encoding(&encoding, false);
1889
1890 let compressor = strategy
1891 .create_miniblock_compressor(&field, &variable_data)
1892 .unwrap();
1893 let (_block, encoding) = compressor
1894 .compress(miniblock_context(), variable_data.clone())
1895 .unwrap();
1896 check_uncompressed_encoding(&encoding, true);
1897
1898 let compressor = strategy.create_per_value(&field, &fixed_data).unwrap();
1900 let (_block, encoding) = compressor.compress(fixed_data).unwrap();
1901 check_uncompressed_encoding(&encoding, false);
1902
1903 let compressor = strategy.create_per_value(&field, &variable_data).unwrap();
1904 let (_block, encoding) = compressor.compress(variable_data).unwrap();
1905 check_uncompressed_encoding(&encoding, true);
1906 }
1907
1908 #[test]
1909 fn test_auto_fsst_disabled_for_binary_fields() {
1910 let strategy = baseline_strategy(CompressionParams::default());
1911 let field = create_test_field("bytes", DataType::Binary);
1912 let variable_data = create_fsst_candidate_variable_width_block();
1913
1914 let miniblock = strategy
1915 .create_miniblock_compressor(&field, &variable_data)
1916 .unwrap();
1917 let miniblock_debug = format!("{:?}", miniblock);
1918 assert!(
1919 miniblock_debug.contains("BinaryMiniBlockEncoder"),
1920 "expected BinaryMiniBlockEncoder, got: {miniblock_debug}"
1921 );
1922 assert!(
1923 !miniblock_debug.contains("FsstMiniBlockEncoder"),
1924 "did not expect FsstMiniBlockEncoder, got: {miniblock_debug}"
1925 );
1926
1927 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
1928 let per_value_debug = format!("{:?}", per_value);
1929 assert!(
1930 per_value_debug.contains("VariableEncoder"),
1931 "expected VariableEncoder, got: {per_value_debug}"
1932 );
1933 assert!(
1934 !per_value_debug.contains("FsstPerValueEncoder"),
1935 "did not expect FsstPerValueEncoder, got: {per_value_debug}"
1936 );
1937 }
1938
1939 #[test]
1940 fn test_auto_fsst_still_enabled_for_utf8_fields() {
1941 let strategy = baseline_strategy(CompressionParams::default());
1942 let field = create_test_field("text", DataType::Utf8);
1943 let variable_data = create_fsst_candidate_variable_width_block();
1944
1945 let miniblock = strategy
1946 .create_miniblock_compressor(&field, &variable_data)
1947 .unwrap();
1948 let miniblock_debug = format!("{:?}", miniblock);
1949 assert!(
1950 miniblock_debug.contains("FsstMiniBlockEncoder"),
1951 "expected FsstMiniBlockEncoder, got: {miniblock_debug}"
1952 );
1953
1954 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
1955 let per_value_debug = format!("{:?}", per_value);
1956 assert!(
1957 per_value_debug.contains("FsstPerValueEncoder"),
1958 "expected FsstPerValueEncoder, got: {per_value_debug}"
1959 );
1960 }
1961
1962 #[test]
1963 fn test_explicit_fsst_still_supported_for_binary_fields() {
1964 let mut params = CompressionParams::new();
1965 params.columns.insert(
1966 "bytes".to_string(),
1967 CompressionFieldParams {
1968 compression: Some("fsst".to_string()),
1969 ..Default::default()
1970 },
1971 );
1972
1973 let strategy = baseline_strategy(params);
1974 let field = create_test_field("bytes", DataType::Binary);
1975 let variable_data = create_fsst_candidate_variable_width_block();
1976
1977 let miniblock = strategy
1978 .create_miniblock_compressor(&field, &variable_data)
1979 .unwrap();
1980 let miniblock_debug = format!("{:?}", miniblock);
1981 assert!(
1982 miniblock_debug.contains("FsstMiniBlockEncoder"),
1983 "expected FsstMiniBlockEncoder, got: {miniblock_debug}"
1984 );
1985
1986 let per_value = strategy.create_per_value(&field, &variable_data).unwrap();
1987 let per_value_debug = format!("{:?}", per_value);
1988 assert!(
1989 per_value_debug.contains("FsstPerValueEncoder"),
1990 "expected FsstPerValueEncoder, got: {per_value_debug}"
1991 );
1992 }
1993
1994 #[test]
1995 #[cfg(feature = "zstd")]
1996 fn test_compression_level_honored_for_large_per_value() {
1997 let mut params = CompressionParams::new();
1998 params.columns.insert(
1999 "html".to_string(),
2000 CompressionFieldParams {
2001 compression: Some("zstd".to_string()),
2002 compression_level: Some(19),
2003 ..Default::default()
2004 },
2005 );
2006 let strategy = baseline_strategy(params);
2007 let field = create_test_field("html", DataType::Utf8);
2008 let large = create_variable_width_block(32, 64, 40 * 1024);
2009
2010 let per_value = strategy.create_per_value(&field, &large).unwrap();
2011 let debug = format!("{per_value:?}");
2012 assert!(
2013 debug.contains("ZstdBufferCompressor") && debug.contains("compression_level: 19"),
2014 "expected zstd level 19 to reach the per-value compressor, got: {debug}"
2015 );
2016 }
2017
2018 #[test]
2019 fn test_parameter_merge_priority() {
2020 let mut params = CompressionParams::new();
2021
2022 params.types.insert(
2024 "Int32".to_string(),
2025 CompressionFieldParams {
2026 rle_threshold: Some(0.5),
2027 compression: Some("lz4".to_string()),
2028 ..Default::default()
2029 },
2030 );
2031
2032 params.columns.insert(
2034 "user_id".to_string(),
2035 CompressionFieldParams {
2036 rle_threshold: Some(0.2),
2037 compression: Some("zstd".to_string()),
2038 compression_level: Some(6),
2039 bss: None,
2040 minichunk_size: None,
2041 },
2042 );
2043
2044 let merged = params.get_field_params("user_id", &DataType::Int32);
2046
2047 assert_eq!(merged.rle_threshold, Some(0.2));
2049 assert_eq!(merged.compression, Some("zstd".to_string()));
2050 assert_eq!(merged.compression_level, Some(6));
2051
2052 let merged = params.get_field_params("other_field", &DataType::Int32);
2054 assert_eq!(merged.rle_threshold, Some(0.5));
2055 assert_eq!(merged.compression, Some("lz4".to_string()));
2056 assert_eq!(merged.compression_level, None);
2057 }
2058
2059 #[test]
2060 fn test_pattern_matching() {
2061 let mut params = CompressionParams::new();
2062
2063 params.columns.insert(
2065 "log_*".to_string(),
2066 CompressionFieldParams {
2067 compression: Some("zstd".to_string()),
2068 compression_level: Some(6),
2069 ..Default::default()
2070 },
2071 );
2072
2073 let merged = params.get_field_params("log_messages", &DataType::Utf8);
2075 assert_eq!(merged.compression, Some("zstd".to_string()));
2076 assert_eq!(merged.compression_level, Some(6));
2077
2078 let merged = params.get_field_params("messages_log", &DataType::Utf8);
2080 assert_eq!(merged.compression, None);
2081 }
2082
2083 #[test]
2084 fn test_legacy_metadata_support() {
2085 let params = CompressionParams::new();
2086 let strategy = baseline_strategy(params);
2087
2088 let mut metadata = HashMap::new();
2090 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
2091 let mut field = create_test_field("some_column", DataType::Int32);
2092 field.metadata = metadata;
2093
2094 let data = create_fixed_width_block(32, 1000);
2095 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2096
2097 assert!(format!("{:?}", compressor).contains("ValueEncoder"));
2099 }
2100
2101 #[test]
2102 fn test_default_behavior() {
2103 let params = CompressionParams::new();
2105 let strategy = baseline_strategy(params);
2106
2107 let field = create_test_field("random_column", DataType::Int32);
2108 let data = create_fixed_width_block_with_stats(32, 1000, 600);
2110
2111 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2112 let debug_str = format!("{:?}", compressor);
2114 assert!(debug_str.contains("ValueEncoder") || debug_str.contains("InlineBitpacking"));
2115 }
2116
2117 #[test]
2118 fn test_field_metadata_compression() {
2119 let params = CompressionParams::new();
2120 let strategy = baseline_strategy(params);
2121
2122 let mut metadata = HashMap::new();
2124 metadata.insert(COMPRESSION_META_KEY.to_string(), "zstd".to_string());
2125 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "6".to_string());
2126 let mut field = create_test_field("test_column", DataType::Int32);
2127 field.metadata = metadata;
2128
2129 let data = create_fixed_width_block(32, 1000);
2130 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2131
2132 let debug_str = format!("{:?}", compressor);
2134 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
2135 }
2136
2137 #[test]
2138 fn test_field_metadata_rle_threshold() {
2139 let params = CompressionParams::new();
2140 let strategy = baseline_strategy(params);
2141
2142 let mut metadata = HashMap::new();
2144 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "0.8".to_string());
2145 metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); let mut field = create_test_field("test_column", DataType::Int32);
2147 field.metadata = metadata;
2148
2149 let data = create_fixed_width_block_with_stats(32, 1000, 100);
2152
2153 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2154
2155 let debug_str = format!("{:?}", compressor);
2157 assert!(debug_str.contains("RleEncoder"));
2158 }
2159
2160 #[test]
2161 fn test_rle_v2_miniblock_selects_u16_run_lengths() {
2162 let mut metadata = HashMap::new();
2163 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2164 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2165 let mut field = create_test_field("test_column", DataType::Int32);
2166 field.metadata = metadata;
2167
2168 let values = vec![7i32; 1000];
2169 let mut data = FixedWidthDataBlock {
2170 bits_per_value: 32,
2171 data: LanceBuffer::reinterpret_vec(values),
2172 num_values: 1000,
2173 block_info: BlockInfo::default(),
2174 };
2175 data.compute_stat();
2176 let data = DataBlock::FixedWidth(data);
2177
2178 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2179 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2180 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2181 assert_eq!(rle_run_length_bits(&encoding), 16);
2182 }
2183
2184 #[test]
2185 fn test_rle_v2_miniblock_keeps_u8_run_lengths_before_v2_3() {
2186 for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] {
2187 let mut metadata = HashMap::new();
2188 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2189 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2190 let mut field = create_test_field("test_column", DataType::Int32);
2191 field.metadata = metadata;
2192
2193 let values = vec![7i32; 1000];
2194 let mut data = FixedWidthDataBlock {
2195 bits_per_value: 32,
2196 data: LanceBuffer::reinterpret_vec(values),
2197 num_values: 1000,
2198 block_info: BlockInfo::default(),
2199 };
2200 data.compute_stat();
2201 let data = DataBlock::FixedWidth(data);
2202
2203 let strategy = strategy(version, CompressionParams::default());
2204 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2205 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2206 assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}");
2207 }
2208 }
2209
2210 #[test]
2211 fn test_rle_v2_uses_selected_width_cost_before_bitpacking() {
2212 let mut metadata = HashMap::new();
2213 metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string());
2214 metadata.insert(BSS_META_KEY.to_string(), "off".to_string());
2215 let mut field = create_test_field("test_column", DataType::Int32);
2216 field.metadata = metadata;
2217
2218 let values = vec![0i32; 4096];
2219 let mut data = FixedWidthDataBlock {
2220 bits_per_value: 32,
2221 data: LanceBuffer::reinterpret_vec(values),
2222 num_values: 4096,
2223 block_info: BlockInfo::default(),
2224 };
2225 data.compute_stat();
2226 let data = DataBlock::FixedWidth(data);
2227
2228 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2229 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2230 let debug_str = format!("{compressor:?}");
2231 assert!(debug_str.contains("RleEncoder"));
2232
2233 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2234 assert_eq!(rle_run_length_bits(&encoding), 16);
2235 }
2236
2237 #[test]
2238 fn test_rle_v2_sorted_dictionary_indices_select_u16_run_lengths() {
2239 let field = create_test_field("dict_indices", DataType::Int32);
2240
2241 let mut values = Vec::with_capacity(1_200);
2242 for value in 0..4 {
2243 values.extend(std::iter::repeat_n(value, 300));
2244 }
2245 let mut data = FixedWidthDataBlock {
2246 bits_per_value: 32,
2247 data: LanceBuffer::reinterpret_vec(values),
2248 num_values: 1_200,
2249 block_info: BlockInfo::default(),
2250 };
2251 data.compute_stat();
2252 let data = DataBlock::FixedWidth(data);
2253
2254 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2255 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2256 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2257 assert_eq!(rle_run_length_bits(&encoding), 16);
2258 }
2259
2260 #[test]
2261 fn test_rle_v2_short_runs_keep_u8_run_lengths() {
2262 let field = create_test_field("dict_indices", DataType::Int32);
2263
2264 let mut values = Vec::with_capacity(1_280);
2265 for value in 0..10 {
2266 values.extend(std::iter::repeat_n(value, 128));
2267 }
2268 let mut data = FixedWidthDataBlock {
2269 bits_per_value: 32,
2270 data: LanceBuffer::reinterpret_vec(values),
2271 num_values: 1_280,
2272 block_info: BlockInfo::default(),
2273 };
2274 data.compute_stat();
2275 let data = DataBlock::FixedWidth(data);
2276
2277 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2278 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2279 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2280 assert_eq!(rle_run_length_bits(&encoding), 8);
2281 }
2282
2283 #[test]
2284 #[cfg(any(feature = "lz4", feature = "zstd"))]
2285 fn test_rle_miniblock_released_versions_keep_flat_children_when_compression_requested() {
2286 for version in [TestEncoding::StructuralU16, TestEncoding::StructuralU32] {
2287 let mut params = CompressionParams::new();
2288 params.columns.insert(
2289 "dict_indices".to_string(),
2290 CompressionFieldParams {
2291 compression: Some(
2292 if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string(),
2293 ),
2294 rle_threshold: Some(1.0),
2295 bss: Some(BssMode::Off),
2296 ..Default::default()
2297 },
2298 );
2299 let strategy = strategy(version, params);
2300 let field = create_test_field("dict_indices", DataType::UInt32);
2301
2302 let mut values = Vec::with_capacity(8192 * 4);
2303 for value in 0..8192u32 {
2304 values.extend(std::iter::repeat_n(value, 4));
2305 }
2306 let mut data = FixedWidthDataBlock {
2307 bits_per_value: 32,
2308 data: LanceBuffer::reinterpret_vec(values),
2309 num_values: 8192 * 4,
2310 block_info: BlockInfo::default(),
2311 };
2312 data.compute_stat();
2313 let data = DataBlock::FixedWidth(data);
2314
2315 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2316 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2317 let rle = expect_rle_encoding(&encoding);
2318
2319 assert!(
2320 matches!(
2321 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2322 Compression::Flat(_)
2323 ),
2324 "version={version}"
2325 );
2326 assert!(
2327 matches!(
2328 rle.run_lengths
2329 .as_ref()
2330 .unwrap()
2331 .compression
2332 .as_ref()
2333 .unwrap(),
2334 Compression::Flat(_)
2335 ),
2336 "version={version}"
2337 );
2338 }
2339 }
2340
2341 #[test]
2342 #[cfg(feature = "bitpacking")]
2343 fn test_rle_miniblock_strategy_bitpacks_child_values_when_smaller() {
2344 let field = create_test_field("dict_indices", DataType::Int32);
2345
2346 let mut values = Vec::with_capacity(8192 * 4);
2347 for value in 0..8192 {
2348 values.extend(std::iter::repeat_n(value, 4));
2349 }
2350 let mut data = FixedWidthDataBlock {
2351 bits_per_value: 32,
2352 data: LanceBuffer::reinterpret_vec(values),
2353 num_values: 8192 * 4,
2354 block_info: BlockInfo::default(),
2355 };
2356 data.compute_stat();
2357 let data = DataBlock::FixedWidth(data);
2358
2359 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2360 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2361 let debug_str = format!("{compressor:?}");
2362 assert!(debug_str.contains("RleEncoder"));
2363
2364 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2365 let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else {
2366 panic!("expected RLE encoding");
2367 };
2368 assert!(matches!(
2369 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2370 Compression::OutOfLineBitpacking(_)
2371 ));
2372 assert!(matches!(
2373 rle.run_lengths
2374 .as_ref()
2375 .unwrap()
2376 .compression
2377 .as_ref()
2378 .unwrap(),
2379 Compression::Flat(_)
2380 ));
2381 }
2382
2383 #[test]
2384 #[cfg(feature = "bitpacking")]
2385 fn test_rle_miniblock_keeps_child_bitpacked_rle_when_smaller_than_inline_bitpacking() {
2386 let field = create_test_field("int_score", DataType::UInt64);
2387
2388 let mut values = Vec::with_capacity(8192 * 8);
2389 for run_idx in 0..8192 {
2390 let value = match run_idx % 3 {
2391 0 => 3u64,
2392 1 => 4u64,
2393 _ => 5u64,
2394 };
2395 values.extend(std::iter::repeat_n(value, 8));
2396 }
2397 let mut data = FixedWidthDataBlock {
2398 bits_per_value: 64,
2399 data: LanceBuffer::reinterpret_vec(values),
2400 num_values: 8192 * 8,
2401 block_info: BlockInfo::default(),
2402 };
2403 data.compute_stat();
2404 let data = DataBlock::FixedWidth(data);
2405
2406 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::default());
2407 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2408 let debug_str = format!("{compressor:?}");
2409 assert!(
2410 debug_str.contains("RleEncoder"),
2411 "expected RLE to beat inline bitpacking after child selection, got: {debug_str}"
2412 );
2413
2414 let (_compressed, encoding) = compressor.compress(miniblock_context(), data).unwrap();
2415 let rle = expect_rle_encoding(&encoding);
2416 assert!(matches!(
2417 rle.values.as_ref().unwrap().compression.as_ref().unwrap(),
2418 Compression::OutOfLineBitpacking(_)
2419 ));
2420 assert!(matches!(
2421 rle.run_lengths
2422 .as_ref()
2423 .unwrap()
2424 .compression
2425 .as_ref()
2426 .unwrap(),
2427 Compression::Flat(_)
2428 ));
2429 }
2430
2431 #[test]
2432 fn test_field_metadata_override_params() {
2433 let mut params = CompressionParams::new();
2435 params.columns.insert(
2436 "test_column".to_string(),
2437 CompressionFieldParams {
2438 rle_threshold: Some(0.3),
2439 compression: Some("lz4".to_string()),
2440 compression_level: None,
2441 bss: None,
2442 minichunk_size: None,
2443 },
2444 );
2445
2446 let strategy = baseline_strategy(params);
2447
2448 let mut metadata = HashMap::new();
2450 metadata.insert(COMPRESSION_META_KEY.to_string(), "none".to_string());
2451 let mut field = create_test_field("test_column", DataType::Int32);
2452 field.metadata = metadata;
2453
2454 let data = create_fixed_width_block(32, 1000);
2455 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2456
2457 assert!(format!("{:?}", compressor).contains("ValueEncoder"));
2459 }
2460
2461 #[test]
2462 fn test_field_metadata_mixed_configuration() {
2463 let mut params = CompressionParams::new();
2465 params.types.insert(
2466 "Int32".to_string(),
2467 CompressionFieldParams {
2468 rle_threshold: Some(0.5),
2469 compression: Some("lz4".to_string()),
2470 ..Default::default()
2471 },
2472 );
2473
2474 let strategy = baseline_strategy(params);
2475
2476 let mut metadata = HashMap::new();
2478 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), "3".to_string());
2479 let mut field = create_test_field("test_column", DataType::Int32);
2480 field.metadata = metadata;
2481
2482 let data = create_fixed_width_block(32, 1000);
2483 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2484
2485 let debug_str = format!("{:?}", compressor);
2487 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
2488 }
2489
2490 #[test]
2491 fn test_bss_field_metadata() {
2492 let params = CompressionParams::new();
2493 let strategy = baseline_strategy(params);
2494
2495 let mut metadata = HashMap::new();
2497 metadata.insert(BSS_META_KEY.to_string(), "on".to_string());
2498 metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string());
2499 let arrow_field =
2500 ArrowField::new("temperature", DataType::Float32, false).with_metadata(metadata);
2501 let field = Field::try_from(&arrow_field).unwrap();
2502
2503 let data = create_fixed_width_block(32, 100);
2505
2506 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2507 let debug_str = format!("{:?}", compressor);
2508 assert!(debug_str.contains("ByteStreamSplitEncoder"));
2509 }
2510
2511 #[test]
2512 fn test_bss_with_compression() {
2513 let params = CompressionParams::new();
2514 let strategy = baseline_strategy(params);
2515
2516 let mut metadata = HashMap::new();
2518 metadata.insert(BSS_META_KEY.to_string(), "on".to_string());
2519 metadata.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string());
2520 let arrow_field =
2521 ArrowField::new("sensor_data", DataType::Float64, false).with_metadata(metadata);
2522 let field = Field::try_from(&arrow_field).unwrap();
2523
2524 let data = create_fixed_width_block(64, 100);
2526
2527 let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2528 let debug_str = format!("{:?}", compressor);
2529 assert!(debug_str.contains("GeneralMiniBlockCompressor"));
2531 assert!(debug_str.contains("ByteStreamSplitEncoder"));
2532 }
2533
2534 #[test]
2535 #[cfg(any(feature = "lz4", feature = "zstd"))]
2536 fn test_general_block_decompression_fixed_width_v2_2() {
2537 let mut params = CompressionParams::new();
2539 params.columns.insert(
2540 "dict_values".to_string(),
2541 CompressionFieldParams {
2542 compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()),
2543 ..Default::default()
2544 },
2545 );
2546
2547 let strategy = strategy(TestEncoding::StructuralU32, params);
2548
2549 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
2550 let data = create_fixed_width_block(24, 1024);
2551 let DataBlock::FixedWidth(expected_block) = &data else {
2552 panic!("expected fixed width block");
2553 };
2554 let expected_bits = expected_block.bits_per_value;
2555 let expected_num_values = expected_block.num_values;
2556 let num_values = expected_num_values;
2557
2558 let compressor = strategy
2559 .create_block_compressor(&field, &data)
2560 .expect("general compression should be selected");
2561 let (compressed_buffer, encoding) = compressor
2562 .compress(data.clone())
2563 .expect("write path general compression should succeed");
2564 match encoding.compression.as_ref() {
2565 Some(Compression::General(_)) => {}
2566 other => panic!("expected general compression, got {:?}", other),
2567 }
2568
2569 let decompressor = DefaultDecompressionStrategy::default()
2570 .create_block_decompressor(&encoding)
2571 .expect("general block decompressor should be created");
2572
2573 let decoded = decompressor
2574 .decompress(compressed_buffer, num_values)
2575 .expect("decompression should succeed");
2576
2577 match decoded {
2578 DataBlock::FixedWidth(block) => {
2579 assert_eq!(block.bits_per_value, expected_bits);
2580 assert_eq!(block.num_values, expected_num_values);
2581 assert_eq!(block.data.as_ref(), expected_block.data.as_ref());
2582 }
2583 _ => panic!("expected fixed width block"),
2584 }
2585 }
2586
2587 #[cfg(any(feature = "lz4", feature = "zstd"))]
2588 fn assert_general_block_preserves_compression_level(
2589 compression: &str,
2590 expected_scheme: crate::format::pb21::CompressionScheme,
2591 compression_level: Option<i32>,
2592 ) {
2593 let mut params = CompressionParams::new();
2594 params.columns.insert(
2595 "dict_values".to_string(),
2596 CompressionFieldParams {
2597 compression: Some(compression.to_string()),
2598 compression_level,
2599 ..Default::default()
2600 },
2601 );
2602 let strategy = strategy(TestEncoding::StructuralU32, params);
2603 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
2604 let data = create_fixed_width_block(24, 1024);
2605
2606 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
2607 let (_, encoding) = compressor.compress(data).unwrap();
2608 let Some(Compression::General(general)) = encoding.compression.as_ref() else {
2609 panic!("expected general compression");
2610 };
2611
2612 assert_eq!(
2613 general.compression.as_ref(),
2614 Some(&crate::format::pb21::BufferCompression {
2615 scheme: expected_scheme as i32,
2616 level: compression_level,
2617 })
2618 );
2619 }
2620
2621 #[test]
2622 #[cfg(feature = "zstd")]
2623 fn test_general_block_preserves_absent_zstd_level() {
2624 assert_general_block_preserves_compression_level(
2625 "zstd",
2626 crate::format::pb21::CompressionScheme::CompressionAlgorithmZstd,
2627 None,
2628 );
2629 }
2630
2631 #[test]
2632 #[cfg(feature = "lz4")]
2633 fn test_general_block_preserves_explicit_lz4_level() {
2634 assert_general_block_preserves_compression_level(
2635 "lz4",
2636 crate::format::pb21::CompressionScheme::CompressionAlgorithmLz4,
2637 Some(7),
2638 );
2639 }
2640
2641 #[test]
2642 #[cfg(any(feature = "lz4", feature = "zstd"))]
2643 fn test_general_compression_not_selected_for_v2_1_even_if_requested() {
2644 let mut params = CompressionParams::new();
2645 params.columns.insert(
2646 "dict_values".to_string(),
2647 CompressionFieldParams {
2648 compression: Some(if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string()),
2649 ..Default::default()
2650 },
2651 );
2652
2653 let strategy = strategy(TestEncoding::StructuralU16, params);
2654 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
2655 let data = create_fixed_width_block(24, 1024);
2656
2657 let compressor = strategy
2658 .create_block_compressor(&field, &data)
2659 .expect("block compressor selection should succeed");
2660 let (_, encoding) = compressor.compress(data).unwrap();
2661
2662 assert!(
2663 !matches!(encoding.compression.as_ref(), Some(Compression::General(_))),
2664 "general compression should not be selected for V2.1"
2665 );
2666 }
2667
2668 #[test]
2669 fn test_none_compression_disables_auto_general_block_compression() {
2670 let mut params = CompressionParams::new();
2671 params.columns.insert(
2672 "dict_values".to_string(),
2673 CompressionFieldParams {
2674 compression: Some("none".to_string()),
2675 ..Default::default()
2676 },
2677 );
2678
2679 let strategy = strategy(TestEncoding::StructuralU32, params);
2680 let field = create_test_field("dict_values", DataType::FixedSizeBinary(3));
2681 let data = create_fixed_width_block(24, 20_000);
2682
2683 assert!(
2684 data.data_size() > MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION,
2685 "test requires block size above automatic general compression threshold"
2686 );
2687
2688 let compressor = strategy
2689 .create_block_compressor(&field, &data)
2690 .expect("block compressor selection should succeed");
2691 let (_, encoding) = compressor.compress(data).unwrap();
2692
2693 assert!(
2694 !matches!(encoding.compression.as_ref(), Some(Compression::General(_))),
2695 "compression=none should disable automatic block general compression"
2696 );
2697 }
2698
2699 #[test]
2700 fn test_rle_v2_block_selects_u32_run_lengths() {
2701 let field = create_test_field("dict_indices", DataType::Int32);
2702 let expected_values = vec![42i32; 70_000];
2703 let mut block = FixedWidthDataBlock {
2704 bits_per_value: 32,
2705 data: LanceBuffer::reinterpret_vec(expected_values.clone()),
2706 num_values: expected_values.len() as u64,
2707 block_info: BlockInfo::default(),
2708 };
2709 block.compute_stat();
2710 let data = DataBlock::FixedWidth(block);
2711
2712 let strategy = strategy(TestEncoding::StructuralSparse, CompressionParams::new());
2713 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
2714 let (compressed, encoding) = compressor.compress(data).unwrap();
2715 assert_eq!(rle_run_length_bits(&encoding), 32);
2716 let decompressor = DefaultDecompressionStrategy::default()
2717 .create_block_decompressor(&encoding)
2718 .unwrap();
2719 let decoded = decompressor
2720 .decompress(compressed, expected_values.len() as u64)
2721 .unwrap();
2722
2723 match decoded {
2724 DataBlock::FixedWidth(block) => {
2725 let values = block.data.borrow_to_typed_slice::<i32>();
2726 assert_eq!(values.as_ref(), expected_values);
2727 }
2728 _ => panic!("expected fixed-width block"),
2729 }
2730 }
2731
2732 #[test]
2733 fn test_rle_v2_block_keeps_u8_run_lengths_for_v2_2() {
2734 let field = create_test_field("dict_indices", DataType::Int32);
2735 let values = vec![42i32; 70_000];
2736 let mut block = FixedWidthDataBlock {
2737 bits_per_value: 32,
2738 data: LanceBuffer::reinterpret_vec(values),
2739 num_values: 70_000,
2740 block_info: BlockInfo::default(),
2741 };
2742 block.compute_stat();
2743 let data = DataBlock::FixedWidth(block);
2744
2745 let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new());
2746 let compressor = strategy.create_block_compressor(&field, &data).unwrap();
2747 let (_, encoding) = compressor.compress(data).unwrap();
2748 assert_eq!(rle_run_length_bits(&encoding), 8);
2749 }
2750
2751 #[test]
2752 fn test_rle_block_used_for_version_v2_2() {
2753 let field = create_test_field("test_repdef", DataType::UInt16);
2754
2755 let num_values = 1000u64;
2757 let mut data = Vec::with_capacity(num_values as usize);
2758 for i in 0..10 {
2759 for _ in 0..100 {
2760 data.push(i as u16);
2761 }
2762 }
2763
2764 let mut block = FixedWidthDataBlock {
2765 bits_per_value: 16,
2766 data: LanceBuffer::reinterpret_vec(data),
2767 num_values,
2768 block_info: BlockInfo::default(),
2769 };
2770
2771 block.compute_stat();
2772
2773 let data_block = DataBlock::FixedWidth(block);
2774
2775 let strategy = strategy(TestEncoding::StructuralU32, CompressionParams::new());
2776
2777 let compressor = strategy
2778 .create_block_compressor(&field, &data_block)
2779 .unwrap();
2780
2781 let debug_str = format!("{:?}", compressor);
2782 assert!(debug_str.contains("RleEncoder"));
2783 }
2784
2785 #[test]
2786 fn test_rle_block_not_used_for_version_v2_1() {
2787 let field = create_test_field("test_repdef", DataType::UInt16);
2788
2789 let num_values = 1000u64;
2791 let mut data = Vec::with_capacity(num_values as usize);
2792 for i in 0..10 {
2793 for _ in 0..100 {
2794 data.push(i as u16);
2795 }
2796 }
2797
2798 let mut block = FixedWidthDataBlock {
2799 bits_per_value: 16,
2800 data: LanceBuffer::reinterpret_vec(data),
2801 num_values,
2802 block_info: BlockInfo::default(),
2803 };
2804
2805 block.compute_stat();
2806
2807 let data_block = DataBlock::FixedWidth(block);
2808
2809 let strategy = strategy(TestEncoding::StructuralU16, CompressionParams::new());
2810
2811 let compressor = strategy
2812 .create_block_compressor(&field, &data_block)
2813 .unwrap();
2814
2815 let debug_str = format!("{:?}", compressor);
2816 assert!(
2817 !debug_str.contains("RleEncoder"),
2818 "RLE should not be used for V2.1"
2819 );
2820 }
2821}