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