Skip to main content

lance_encoding/
compression.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Compression traits and definitions for Lance 2.1
5//!
6//! In 2.1 the first step of encoding is structural encoding, where we shred inputs into
7//! leaf arrays and take care of the validity / offsets structure.  Then we pick a structural
8//! encoding (mini-block or full-zip) and then we compress the data.
9//!
10//! This module defines the traits for the compression step.  Each structural encoding has its
11//! own compression strategy.
12//!
13//! Miniblock compression is a block based approach for small data.  Since we introduce some read
14//! amplification and decompress entire blocks we are able to use opaque compression.
15//!
16//! Fullzip compression is a per-value approach where we require that values are transparently
17//! compressed so that we can locate them later.
18
19#[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
75/// Default threshold for RLE compression selection when the user explicitly provides a threshold.
76///
77/// If no threshold is provided, we use a size model instead of a fixed run ratio.
78/// This preserves existing behavior for users relying on the default, while making
79/// the default selection more type-aware.
80const DEFAULT_RLE_COMPRESSION_THRESHOLD: f64 = 0.5;
81
82// Minimum block size (32kb) to trigger general block compression
83const MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION: u64 = 32 * 1024;
84const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::<u64>() as u128;
85
86/// Trait for compression algorithms that compress an entire block of data into one opaque
87/// and self-described chunk.
88///
89/// This is actually a _third_ compression strategy used in a few corner cases today (TODO: remove?)
90///
91/// This is the most general type of compression.  There are no constraints on the method
92/// of compression it is assumed that the entire block of data will be present at decompression.
93///
94/// This is the least appropriate strategy for random access because we must load the entire
95/// block to access any single value.  This should only be used for cases where random access is never
96/// required (e.g. when encoding metadata buffers like a dictionary or for encoding rep/def
97/// mini-block chunks)
98pub trait BlockCompressor: std::fmt::Debug + Send + Sync {
99    /// Compress the data into zero or one buffers and describe the codec used.
100    ///
101    /// `None` represents a metadata-only codec. `Some` represents a physical
102    /// payload, including a zero-byte payload.
103    fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)>;
104}
105
106/// A trait to pick which compression to use for given data
107///
108/// There are several different kinds of compression.
109///
110/// - Block compression is the most generic, but most difficult to use efficiently
111/// - Per-value compression results in either a fixed width data block or a variable
112///   width data block.  In other words, there is some number of bits per value.
113///   In addition, each value should be independently decompressible.
114/// - Mini-block compression results in a small block of opaque data for chunks
115///   of rows.  Each block is somewhere between 0 and 16KiB in size.  This is
116///   used for narrow data types (both fixed and variable length) where we can
117///   fit many values into an 16KiB block.
118pub trait CompressionStrategy: Send + Sync + std::fmt::Debug {
119    /// Create a block compressor for the given data.
120    fn create_block_compressor(
121        &self,
122        field: &Field,
123        data: &DataBlock,
124    ) -> Result<Box<dyn BlockCompressor>>;
125
126    /// Create a per-value compressor for the given data
127    fn create_per_value(
128        &self,
129        field: &Field,
130        data: &DataBlock,
131    ) -> Result<Box<dyn PerValueCompressor>>;
132
133    /// Create a mini-block compressor for the given data
134    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    // BSS requires general compression to be effective
159    // If compression is not set or explicitly disabled, skip BSS
160    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    // should_use_bss already checks for supported bit widths (32/64)
166    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    // If the user explicitly provided a threshold then honor it as an additional guard.
187    // A lower threshold makes RLE harder to trigger and can be used to avoid CPU overhead.
188    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    // Explicitly disable general compression.
526    if field_params.compression.as_deref() == Some("none") {
527        return Ok(None);
528    }
529
530    // User-requested compression (unused today but perhaps still used
531    // in the future someday)
532    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    // Automatic compression for large blocks
540    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
549/// Parse field-level compression metadata without applying format-specific constraints.
550pub 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
582/// Apply general-purpose compression requested for a fixed-width miniblock.
583pub 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
595/// Honor an explicit `compression = none` request for fixed-width miniblocks.
596pub 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
604/// Select byte-stream-split compression for an applicable fixed-width miniblock.
605pub 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
615/// Select the original fixed-u8 RLE miniblock grammar.
616pub 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
626/// Select variable-width RLE with independently encoded children.
627pub 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
637/// Select inline bitpacking for applicable fixed-width miniblocks.
638pub 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
645/// Store fixed-width miniblock values without a value codec.
646pub 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
650/// Encode variable-width miniblocks with binary or FSST encoding.
651pub 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
693/// Encode fixed-width packed structs as miniblocks.
694pub 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
710/// Store fixed-size-list miniblocks without a value codec.
711pub 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
715/// Store fixed-width and fixed-size-list values directly in full-zip pages.
716pub 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
733/// Reject variable-width packed structs while preserving the fixed-width error.
734pub 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
751/// Encode variable-width packed structs with the exact strategy recursively.
752pub 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
771/// Encode all packed structs with the exact strategy recursively.
772pub 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
792/// Encode variable-width values directly, with FSST or per-value compression
793/// when applicable.
794pub 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
837/// Select fixed-u8 RLE for block compression.
838pub 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
848/// Select variable-width RLE for block compression.
849pub 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
859/// Select block bitpacking for applicable fixed-width values.
860pub 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
867/// Select explicitly requested or automatic general-purpose block compression.
868pub 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
878/// Store fixed- and variable-width block values without block compression.
879pub 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    /// Returns the exact aggregate decoded size when it is determined solely by the value count.
895    ///
896    /// Implementations should only return `Some` when this aggregate estimate can be used by
897    /// [`DataBlockBuilder`](crate::data::DataBlockBuilder) to preallocate the decoded output
898    /// exactly. Outputs with multiple buffers or whose layout-dependent allocation cannot be
899    /// represented by one aggregate estimate should return `None`.
900    fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
901        None
902    }
903}
904
905pub trait FixedPerValueDecompressor: std::fmt::Debug + Send + Sync {
906    /// Decompress one or more values
907    fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result<DataBlock>;
908    /// The number of bits in each value
909    ///
910    /// Currently (and probably long term) this must be a multiple of 8
911    fn bits_per_value(&self) -> u64;
912
913    /// Returns the exact aggregate decoded size when it is determined solely by the value count.
914    ///
915    /// Implementations should only return `Some` when this aggregate estimate can be used by
916    /// [`DataBlockBuilder`](crate::data::DataBlockBuilder) to preallocate the decoded output
917    /// exactly. Outputs with multiple buffers or whose layout-dependent allocation cannot be
918    /// represented by one aggregate estimate should return `None`.
919    fn decoded_size_bytes(&self, _num_values: u64) -> Option<u64> {
920        None
921    }
922}
923
924pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync {
925    /// Decompress one or more values
926    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    /// Whether this codec consumes one payload buffer.
933    fn requires_payload(&self) -> bool {
934        true
935    }
936
937    /// Inspect a block for an exact payload-derived value count when supported.
938    ///
939    /// This must not materialize the decoded values. `None` means the encoding requires an
940    /// external count or cannot safely prove one from this payload.
941    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                // In the future, we might need to do something more complex here if FSL supports
1032                // compression.
1033                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                // Create inner decompressor
1051                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                // Parse compression config
1059                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                // Extract the compressed bit width from the values encoding
1188                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        // Create varied data to avoid low entropy
1451        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        // Create data with specified run count
1456        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); // Use prime to get varied values
1462            }
1463            // Fill all bytes of the value to create high entropy
1464            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        // Compute all statistics including BytePositionEntropy
1480        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        // Create data with some variety to avoid always triggering BSS
1488        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        // Add some variation to the data to make it more realistic
1493        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        // Compute all statistics including BytePositionEntropy
1508        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        // Create offsets buffer (num_values + 1 offsets)
1553        let mut offsets = Vec::with_capacity((num_values + 1) as usize);
1554        let mut current_offset = 0i64;
1555        offsets.push(current_offset);
1556
1557        // Generate offsets with varying value sizes
1558        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        // Create data buffer with realistic content
1570        let total_data_size = current_offset as usize;
1571        let mut data = vec![0u8; total_data_size];
1572
1573        // Fill data with varied content
1574        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        // Convert offsets to appropriate lance buffer
1585        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        // Configure RLE for ID columns with BSS explicitly disabled
1615        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), // Explicitly disable BSS to test RLE
1622                minichunk_size: None,
1623            },
1624        );
1625
1626        let strategy = baseline_strategy(params);
1627        let field = create_test_field("user_id", DataType::Int32);
1628
1629        // Create data with low run count for RLE
1630        // Use create_fixed_width_block_with_stats which properly sets run count
1631        let data = create_fixed_width_block_with_stats(32, 1000, 100); // 100 runs out of 1000 values
1632
1633        let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
1634        // Should use RLE due to low threshold (0.3) and low run count (100/1000 = 0.1)
1635        let debug_str = format!("{:?}", compressor);
1636
1637        // The compressor should be RLE wrapped in general compression
1638        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        // Configure all Int32 to use specific settings
1647        params.types.insert(
1648            "Int32".to_string(),
1649            CompressionFieldParams {
1650                rle_threshold: Some(0.1), // Very low threshold
1651                compression: Some("zstd".to_string()),
1652                compression_level: Some(3),
1653                bss: Some(BssMode::Off), // Disable BSS to test RLE
1654                minichunk_size: None,
1655            },
1656        );
1657
1658        let strategy = baseline_strategy(params);
1659        let field = create_test_field("some_column", DataType::Int32);
1660        // Create data with very low run count (50 runs for 1000 values = 0.05 ratio)
1661        let data = create_fixed_width_block_with_stats(32, 1000, 50);
1662
1663        let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
1664        // Should use RLE due to very low threshold
1665        assert!(format!("{:?}", compressor).contains("RleEncoder"));
1666    }
1667
1668    // Regression for #6626: an all-zero stat segment (e.g. rep/def for a long
1669    // run of empty lists) used to disable block bitpacking entirely.
1670    #[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        // First 1024 zeros, then 1024 ones; max bit width is 1.
1677        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        // Low cardinality values (3/4/5) but with moderate run count:
1757        // RLE compresses vs raw, yet bitpacking should be smaller.
1758        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        // Disable compression for embeddings
1808        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        // Test miniblock
1822        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        // Test pervalue
1838        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        // Prepare field with metadata for none compression
1849        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        // Test miniblock
1858        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        // Test pervalue
1878        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        // Set type-level
2002        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        // Set column-level (highest priority)
2012        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        // Get merged params
2024        let merged = params.get_field_params("user_id", &DataType::Int32);
2025
2026        // Column params should override type params
2027        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        // Test field with only type params
2032        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        // Configure pattern for log files
2043        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        // Should match pattern
2053        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        // Should not match
2058        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        // Test field with "none" compression metadata
2068        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        // Should respect metadata and use ValueEncoder
2077        assert!(format!("{:?}", compressor).contains("ValueEncoder"));
2078    }
2079
2080    #[test]
2081    fn test_default_behavior() {
2082        // Empty params should fall back to default behavior
2083        let params = CompressionParams::new();
2084        let strategy = baseline_strategy(params);
2085
2086        let field = create_test_field("random_column", DataType::Int32);
2087        // Create data with high run count that won't trigger RLE (600 runs for 1000 values = 0.6 ratio)
2088        let data = create_fixed_width_block_with_stats(32, 1000, 600);
2089
2090        let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2091        // Should use default strategy's decision
2092        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        // Test field with compression metadata
2102        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        // Should use zstd with level 6
2112        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        // Test field with RLE threshold metadata
2122        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()); // Disable BSS to test RLE
2125        let mut field = create_test_field("test_column", DataType::Int32);
2126        field.metadata = metadata;
2127
2128        // Create data with low run count (e.g., 100 runs for 1000 values = 0.1 ratio)
2129        // This ensures run_count (100) < num_values * threshold (1000 * 0.8 = 800)
2130        let data = create_fixed_width_block_with_stats(32, 1000, 100);
2131
2132        let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap();
2133
2134        // Should use RLE because run_count (100) < num_values * threshold (800)
2135        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        // Set up params with one configuration
2413        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        // Field metadata should override params
2428        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        // Should use none compression (from metadata) instead of lz4 (from params)
2437        assert!(format!("{:?}", compressor).contains("ValueEncoder"));
2438    }
2439
2440    #[test]
2441    fn test_field_metadata_mixed_configuration() {
2442        // Configure type-level params
2443        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        // Field metadata provides partial override
2456        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        // Should use lz4 (from type params) with level 3 (from metadata)
2465        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        // Test BSS "on" mode with compression enabled (BSS requires compression to be effective)
2475        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        // Create float data
2483        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        // Test BSS with LZ4 compression
2496        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        // Create double data
2504        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        // Should have BSS wrapped in general compression
2509        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        // Request general compression via the write path (2.2 requirement) and ensure the read path mirrors it.
2517        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        // Create highly repetitive data
2735        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        // Create highly repetitive data
2769        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}