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