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