Skip to main content

lance_encoding/encodings/physical/
value.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use arrow_buffer::{BooleanBufferBuilder, bit_util};
5
6use crate::buffer::LanceBuffer;
7use crate::compression::{
8    BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, MiniBlockDecompressor,
9    require_block_payload,
10};
11use crate::data::{
12    AllNullDataBlock, BlockInfo, DataBlock, FixedSizeListBlock, FixedWidthDataBlock,
13    NullableDataBlock,
14};
15use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock};
16use crate::encodings::logical::primitive::miniblock::{
17    MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed,
18    MiniBlockCompressionContext, MiniBlockCompressor,
19};
20use crate::format::EncodingUtils;
21use crate::format::metadata::compressive_encoding::Compression;
22use crate::format::metadata::{self, CompressiveEncoding};
23
24use lance_core::{Error, Result};
25
26/// A compression strategy that writes fixed-width data as-is (no compression)
27#[derive(Debug, Default)]
28pub struct ValueEncoder {}
29
30impl ValueEncoder {
31    /// Use the largest chunk we can smaller than 4KiB
32    fn find_log_vals_per_chunk(bytes_per_word: u64, values_per_word: u64) -> Result<(u64, u64)> {
33        let mut size_bytes = 2 * bytes_per_word;
34        let (mut log_num_vals, mut num_vals) = match values_per_word {
35            1 => (1, 2),
36            8 => (3, 8),
37            _ => unreachable!(),
38        };
39
40        if size_bytes >= MAX_MINIBLOCK_BYTES {
41            let num_values = 2 * values_per_word;
42            return Err(Error::invalid_input(format!(
43                "Value is too wide for miniblock encoding: {} values require {} bytes but a \
44                 miniblock chunk is limited to {} bytes.",
45                num_values, size_bytes, MAX_MINIBLOCK_BYTES
46            )));
47        }
48
49        while 2 * size_bytes < MAX_MINIBLOCK_BYTES && 2 * num_vals <= *MAX_MINIBLOCK_VALUES {
50            log_num_vals += 1;
51            size_bytes *= 2;
52            num_vals *= 2;
53        }
54
55        Ok((log_num_vals, num_vals))
56    }
57
58    fn chunk_data(data: FixedWidthDataBlock) -> Result<MiniBlockCompressed> {
59        // Usually there are X bytes per value.  However, when working with boolean
60        // or FSL<boolean> we might have some number of bits per value that isn't
61        // divisible by 8.  In this case, to avoid chunking in the middle of a byte
62        // we calculate how many 8-value words we can fit in a chunk.
63        let (bytes_per_word, values_per_word) = if data.bits_per_value.is_multiple_of(8) {
64            (data.bits_per_value / 8, 1)
65        } else {
66            (data.bits_per_value, 8)
67        };
68
69        // Aim for 4KiB chunks
70        let (log_vals_per_chunk, vals_per_chunk) =
71            Self::find_log_vals_per_chunk(bytes_per_word, values_per_word)?;
72        let num_chunks = bit_util::ceil(data.num_values as usize, vals_per_chunk as usize);
73        debug_assert_eq!(vals_per_chunk % values_per_word, 0);
74        let bytes_per_chunk = bytes_per_word * (vals_per_chunk / values_per_word);
75        let bytes_per_chunk = u32::try_from(bytes_per_chunk).unwrap();
76        debug_assert!(bytes_per_chunk > 0);
77
78        let data_buffer = data.data;
79
80        let mut row_offset = 0;
81        let mut chunks = Vec::with_capacity(num_chunks);
82
83        let mut bytes_counter = 0;
84        loop {
85            if row_offset + vals_per_chunk <= data.num_values {
86                // We can make a full chunk
87                chunks.push(MiniBlockChunk {
88                    log_num_values: log_vals_per_chunk as u8,
89                    buffer_sizes: vec![bytes_per_chunk],
90                });
91                row_offset += vals_per_chunk;
92                bytes_counter += bytes_per_chunk as u64;
93            } else if row_offset < data.num_values {
94                // Final chunk, special values
95                let num_bytes = data_buffer.len() as u64 - bytes_counter;
96                let num_bytes = u32::try_from(num_bytes).unwrap();
97                chunks.push(MiniBlockChunk {
98                    log_num_values: 0,
99                    buffer_sizes: vec![num_bytes],
100                });
101                break;
102            } else {
103                // If we get here then all chunks were full chunks and we have no remainder chunk
104                break;
105            }
106        }
107
108        debug_assert_eq!(chunks.len(), num_chunks);
109
110        Ok(MiniBlockCompressed {
111            chunks,
112            data: vec![data_buffer],
113            num_values: data.num_values,
114        })
115    }
116}
117
118#[derive(Debug)]
119struct MiniblockFslLayer {
120    validity: Option<LanceBuffer>,
121    dimension: u64,
122}
123
124/// This impl deals with encoding FSL<FSL<...<FSL<FixedWidth>>>> data as a mini-block compressor.
125/// The tricky part of FSL data is that we want to include inner validity buffers (we don't want these
126/// to be part of the rep-def because that usually ends up being more expensive).
127///
128/// The resulting mini-block will, instead of having a single buffer, have X + 1 buffers where X is
129/// the number of FSL layers that contain validity.
130///
131/// In the simple case where there is no validity inside the FSL layers, all we are doing here is flattening
132/// the FSL layers into a single buffer.
133///
134/// Also: We don't allow a row to be broken across chunks.  This typically isn't too big of a deal since we
135/// are usually dealing with relatively small vectors if we are using mini-block.
136///
137/// Note: when we do have validity we have to make copies of the validity buffers because they are bit buffers
138/// and we need to bit slice them which requires copies or offsets.  Paying the price at write time to make
139/// the copies is better than paying the price at read time to do the bit slicing.
140impl ValueEncoder {
141    fn make_fsl_encoding(layers: &[MiniblockFslLayer], bits_per_value: u64) -> CompressiveEncoding {
142        let mut encoding = EncodingUtils::flat(bits_per_value, None);
143        for layer in layers.iter().rev() {
144            let has_validity = layer.validity.is_some();
145            let dimension = layer.dimension;
146            encoding = EncodingUtils::fsl(dimension, has_validity, encoding);
147        }
148        encoding
149    }
150
151    fn extract_fsl_chunk(
152        data: &FixedWidthDataBlock,
153        layers: &[MiniblockFslLayer],
154        row_offset: usize,
155        num_rows: usize,
156        validity_buffers: &mut [Vec<u8>],
157    ) -> Vec<u32> {
158        let mut row_offset = row_offset;
159        let mut num_values = num_rows;
160        let mut buffer_counter = 0;
161        let mut buffer_sizes = Vec::with_capacity(validity_buffers.len() + 1);
162        for layer in layers {
163            row_offset *= layer.dimension as usize;
164            num_values *= layer.dimension as usize;
165            if let Some(validity) = &layer.validity {
166                let validity_slice = validity
167                    .clone()
168                    .bit_slice_le_with_length(row_offset, num_values);
169                validity_buffers[buffer_counter].extend_from_slice(&validity_slice);
170                buffer_sizes.push(validity_slice.len() as u32);
171                buffer_counter += 1;
172            }
173        }
174
175        let bits_in_chunk = data.bits_per_value * num_values as u64;
176        let bytes_in_chunk = bits_in_chunk.div_ceil(8);
177        let bytes_in_chunk = u32::try_from(bytes_in_chunk).unwrap();
178        debug_assert!(bytes_in_chunk > 0);
179        buffer_sizes.push(bytes_in_chunk);
180
181        buffer_sizes
182    }
183
184    fn chunk_fsl(
185        data: FixedWidthDataBlock,
186        layers: Vec<MiniblockFslLayer>,
187        num_rows: u64,
188    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
189        // Count size to calculate rows per chunk
190        let mut ceil_bytes_validity = 0;
191        let mut cum_dim = 1;
192        let mut num_validity_buffers = 0;
193        for layer in &layers {
194            cum_dim *= layer.dimension;
195            if layer.validity.is_some() {
196                ceil_bytes_validity += cum_dim.div_ceil(8);
197                num_validity_buffers += 1;
198            }
199        }
200        // It's an estimate because validity buffers may have some padding bits
201        let cum_bits_per_value = data.bits_per_value * cum_dim;
202        let (cum_bytes_per_word, vals_per_word) = if cum_bits_per_value.is_multiple_of(8) {
203            (cum_bits_per_value / 8, 1)
204        } else {
205            (cum_bits_per_value, 8)
206        };
207        let est_bytes_per_word = (ceil_bytes_validity * vals_per_word) + cum_bytes_per_word;
208        let (log_rows_per_chunk, rows_per_chunk) =
209            Self::find_log_vals_per_chunk(est_bytes_per_word, vals_per_word)?;
210
211        let num_chunks = num_rows.div_ceil(rows_per_chunk) as usize;
212
213        // Allocate buffers for validity, these will be slightly bigger than the input validity buffers
214        let mut chunks = Vec::with_capacity(num_chunks);
215        let mut validity_buffers: Vec<Vec<u8>> = Vec::with_capacity(num_validity_buffers);
216        cum_dim = 1;
217        for layer in &layers {
218            cum_dim *= layer.dimension;
219            if let Some(validity) = &layer.validity {
220                let layer_bytes_validity = cum_dim.div_ceil(8);
221                let validity_with_padding =
222                    layer_bytes_validity as usize * num_chunks * rows_per_chunk as usize;
223                debug_assert!(validity_with_padding >= validity.len());
224                validity_buffers.push(Vec::with_capacity(
225                    layer_bytes_validity as usize * num_chunks,
226                ));
227            }
228        }
229
230        // Now go through and extract validity buffers
231        let mut row_offset = 0;
232        while row_offset + rows_per_chunk <= num_rows {
233            let buffer_sizes = Self::extract_fsl_chunk(
234                &data,
235                &layers,
236                row_offset as usize,
237                rows_per_chunk as usize,
238                &mut validity_buffers,
239            );
240            row_offset += rows_per_chunk;
241            chunks.push(MiniBlockChunk {
242                log_num_values: log_rows_per_chunk as u8,
243                buffer_sizes,
244            })
245        }
246        let rows_in_chunk = num_rows - row_offset;
247        if rows_in_chunk > 0 {
248            let buffer_sizes = Self::extract_fsl_chunk(
249                &data,
250                &layers,
251                row_offset as usize,
252                rows_in_chunk as usize,
253                &mut validity_buffers,
254            );
255            chunks.push(MiniBlockChunk {
256                log_num_values: 0,
257                buffer_sizes,
258            });
259        }
260
261        let encoding = Self::make_fsl_encoding(&layers, data.bits_per_value);
262        // Finally, add the data buffer
263        let buffers = validity_buffers
264            .into_iter()
265            .map(LanceBuffer::from)
266            .chain(std::iter::once(data.data))
267            .collect::<Vec<_>>();
268
269        Ok((
270            MiniBlockCompressed {
271                chunks,
272                data: buffers,
273                num_values: num_rows,
274            },
275            encoding,
276        ))
277    }
278
279    fn miniblock_fsl(data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
280        let num_rows = data.num_values();
281        let fsl = data.as_fixed_size_list().unwrap();
282        let mut layers = Vec::new();
283        let mut child = *fsl.child;
284        let mut cur_layer = MiniblockFslLayer {
285            validity: None,
286            dimension: fsl.dimension,
287        };
288        loop {
289            if let DataBlock::Nullable(nullable) = child {
290                cur_layer.validity = Some(nullable.nulls);
291                child = *nullable.data;
292            }
293            match child {
294                DataBlock::FixedSizeList(inner) => {
295                    layers.push(cur_layer);
296                    cur_layer = MiniblockFslLayer {
297                        validity: None,
298                        dimension: inner.dimension,
299                    };
300                    child = *inner.child;
301                }
302                DataBlock::FixedWidth(inner) => {
303                    layers.push(cur_layer);
304                    return Self::chunk_fsl(inner, layers, num_rows);
305                }
306                _ => unreachable!("Unexpected data block type in value encoder's miniblock_fsl"),
307            }
308        }
309    }
310}
311
312struct PerValueFslValidityIter {
313    buffer: LanceBuffer,
314    bits_per_row: usize,
315    offset: usize,
316}
317
318/// In this section we deal with per-value encoding of FSL<FSL<...<FSL<FixedWidth>>>> data.
319///
320/// It's easier than mini-block.  All we need to do is flatten the FSL layers into a single buffer.
321/// This includes any validity buffers we encounter on the way.
322impl ValueEncoder {
323    fn fsl_to_encoding(fsl: &FixedSizeListBlock) -> CompressiveEncoding {
324        let mut inner = fsl.child.as_ref();
325        let mut has_validity = false;
326        if let DataBlock::Nullable(nullable) = inner {
327            has_validity = true;
328            inner = nullable.data.as_ref();
329        }
330        let inner_encoding = match inner {
331            // All inner values are null.  Reserve validity bits (has_validity=true) so that
332            // the decoder knows one validity byte per cum_dim items is stored per row.
333            // constant(None) signals that every stored item decodes to null.
334            DataBlock::AllNull(_) => {
335                has_validity = true;
336                EncodingUtils::constant(None)
337            }
338            DataBlock::FixedWidth(fixed_width) => {
339                EncodingUtils::flat(fixed_width.bits_per_value, None)
340            }
341            DataBlock::FixedSizeList(inner) => Self::fsl_to_encoding(inner),
342            _ => unreachable!(
343                "Unexpected data block type in value encoder's fsl_to_encoding: {}",
344                inner.name()
345            ),
346        };
347        EncodingUtils::fsl(fsl.dimension, has_validity, inner_encoding)
348    }
349
350    fn simple_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
351        // The simple case is zero-copy, we just return the flattened inner buffer
352        let encoding = Self::fsl_to_encoding(&fsl);
353        let Some(flat) = fsl.try_into_flat() else {
354            unreachable!("per_value_fsl only sends FSL blocks without nullable children here")
355        };
356        let data = FixedWidthDataBlock {
357            block_info: BlockInfo::new(),
358            ..flat
359        };
360        (PerValueDataBlock::Fixed(data), encoding)
361    }
362
363    fn nullable_per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
364        // If there are nullable inner values then we need to zip the validity with the values
365        let encoding = Self::fsl_to_encoding(&fsl);
366        let num_values = fsl.num_values();
367        let mut bytes_per_row = 0;
368        let mut cum_dim = 1;
369        let mut current = fsl;
370        let mut validity_iters: Vec<PerValueFslValidityIter> = Vec::new();
371        let data_bytes_per_row: usize;
372        let data_buffer: LanceBuffer;
373        loop {
374            cum_dim *= current.dimension;
375            let mut child = *current.child;
376            if let DataBlock::Nullable(nullable) = child {
377                // Each item will need this many bytes of validity prepended to it
378                bytes_per_row += cum_dim.div_ceil(8) as usize;
379                validity_iters.push(PerValueFslValidityIter {
380                    buffer: nullable.nulls,
381                    bits_per_row: cum_dim as usize,
382                    offset: 0,
383                });
384                child = *nullable.data;
385            };
386            match child {
387                DataBlock::FixedSizeList(inner) => {
388                    current = inner;
389                }
390                DataBlock::FixedWidth(fixed_width) => {
391                    data_bytes_per_row =
392                        (fixed_width.bits_per_value.div_ceil(8) * cum_dim) as usize;
393                    bytes_per_row += data_bytes_per_row;
394                    data_buffer = fixed_width.data;
395                    break;
396                }
397                DataBlock::AllNull(_) => {
398                    // All inner values are null.  Add all-zero validity bits so that
399                    // bytes_per_row > 0 and the FullZip layout doesn't write bits_per_value=0,
400                    // which would crash the reader when there are also no ctrl-word bytes.
401                    bytes_per_row += cum_dim.div_ceil(8) as usize;
402                    validity_iters.push(PerValueFslValidityIter {
403                        buffer: LanceBuffer::from(vec![
404                            0u8;
405                            cum_dim.div_ceil(8) as usize
406                                * num_values as usize
407                        ]),
408                        bits_per_row: cum_dim as usize,
409                        offset: 0,
410                    });
411                    data_bytes_per_row = 0;
412                    data_buffer = LanceBuffer::empty();
413                    break;
414                }
415                _ => unreachable!(
416                    "Unexpected data block type in value encoder's nullable_per_value_fsl: {:?}",
417                    child
418                ),
419            }
420        }
421
422        let bytes_needed = bytes_per_row * num_values as usize;
423        let mut zipped = Vec::with_capacity(bytes_needed);
424        let data_slice = &data_buffer;
425        // Hopefully values are pretty large so we don't iterate this loop _too_ many times
426        for i in 0..num_values as usize {
427            for validity in validity_iters.iter_mut() {
428                let validity_slice = validity
429                    .buffer
430                    .bit_slice_le_with_length(validity.offset, validity.bits_per_row);
431                zipped.extend_from_slice(&validity_slice);
432                validity.offset += validity.bits_per_row;
433            }
434            let start = i * data_bytes_per_row;
435            let end = start + data_bytes_per_row;
436            zipped.extend_from_slice(&data_slice[start..end]);
437        }
438
439        let zipped = LanceBuffer::from(zipped);
440        let data = PerValueDataBlock::Fixed(FixedWidthDataBlock {
441            bits_per_value: bytes_per_row as u64 * 8,
442            num_values,
443            data: zipped,
444            block_info: BlockInfo::new(),
445        });
446        (data, encoding)
447    }
448
449    fn per_value_fsl(fsl: FixedSizeListBlock) -> (PerValueDataBlock, CompressiveEncoding) {
450        if !fsl.child.is_nullable() {
451            Self::simple_per_value_fsl(fsl)
452        } else {
453            Self::nullable_per_value_fsl(fsl)
454        }
455    }
456}
457
458impl BlockCompressor for ValueEncoder {
459    fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
460        if matches!(data, DataBlock::FixedSizeList(_)) {
461            let (PerValueDataBlock::Fixed(values), encoding) =
462                PerValueCompressor::compress(self, data)?
463            else {
464                return Err(Error::internal(
465                    "Fixed-size-list block has variable row width",
466                ));
467            };
468            return Ok((Some(values.data), encoding));
469        }
470        let DataBlock::FixedWidth(fixed_width) = data else {
471            return Err(Error::invalid_input(format!(
472                "ValueEncoder cannot compress a {} block",
473                data.name()
474            )));
475        };
476        Ok((
477            Some(fixed_width.data),
478            EncodingUtils::flat(fixed_width.bits_per_value, None),
479        ))
480    }
481}
482
483impl MiniBlockCompressor for ValueEncoder {
484    fn compress(
485        &self,
486        _context: MiniBlockCompressionContext,
487        chunk: DataBlock,
488    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
489        match chunk {
490            DataBlock::FixedWidth(fixed_width) => {
491                let encoding = EncodingUtils::flat(fixed_width.bits_per_value, None);
492                Ok((Self::chunk_data(fixed_width)?, encoding))
493            }
494            DataBlock::FixedSizeList(_) => Self::miniblock_fsl(chunk),
495            _ => Err(Error::invalid_input_source(
496                format!(
497                    "Cannot compress a data block of type {} with ValueEncoder",
498                    chunk.name()
499                )
500                .into(),
501            )),
502        }
503    }
504}
505
506#[derive(Debug)]
507struct ValueFslDesc {
508    dimension: u64,
509    has_validity: bool,
510}
511
512/// A decompressor for fixed-width data that has
513/// been written, as-is, to disk in single contiguous array
514#[derive(Debug)]
515pub struct ValueDecompressor {
516    /// How many bits are in each inner-most item (e.g. FSL<Int32, 100> would be 32)
517    bits_per_item: u64,
518    /// How many bits are in each value (e.g. FSL<Int32, 100> would be 3200)
519    ///
520    /// This number is a little trickier to compute because we also have to include bytes
521    /// of any inner validity
522    bits_per_value: u64,
523    layers: Vec<ValueFslDesc>,
524}
525
526impl ValueDecompressor {
527    pub fn from_flat(description: &metadata::Flat) -> Self {
528        Self {
529            bits_per_item: description.bits_per_value,
530            bits_per_value: description.bits_per_value,
531            layers: Vec::default(),
532        }
533    }
534
535    pub fn from_fsl(mut description: &metadata::FixedSizeList) -> Result<Self> {
536        let mut layers = Vec::new();
537        let mut cum_dim = 1_u64;
538        let mut bytes_per_value = 0_u64;
539        loop {
540            if description.items_per_value == 0 {
541                return Err(Error::invalid_input("FSL dimension must be positive"));
542            }
543            layers.push(ValueFslDesc {
544                has_validity: description.has_validity,
545                dimension: description.items_per_value,
546            });
547            cum_dim = cum_dim
548                .checked_mul(description.items_per_value)
549                .ok_or_else(|| Error::invalid_input("FSL dimension product overflows"))?;
550            if description.has_validity {
551                bytes_per_value = bytes_per_value
552                    .checked_add(cum_dim.div_ceil(8))
553                    .ok_or_else(|| Error::invalid_input("FSL validity size overflows"))?;
554            }
555            let encoding = description
556                .values
557                .as_ref()
558                .ok_or_else(|| Error::invalid_input("FSL encoding missing inner values field"))?
559                .compression
560                .as_ref()
561                .ok_or_else(|| {
562                    Error::invalid_input("FSL encoding missing inner compression field")
563                })?;
564            match encoding {
565                Compression::FixedSizeList(inner) => {
566                    description = inner;
567                }
568                Compression::Flat(flat) => {
569                    if flat.bits_per_value == 0 || flat.data.is_some() {
570                        return Err(Error::invalid_input(
571                            "FSL requires positive plain Flat leaf width",
572                        ));
573                    }
574                    let bits_per_value = flat
575                        .bits_per_value
576                        .checked_mul(cum_dim)
577                        .and_then(|bits| bytes_per_value.checked_mul(8)?.checked_add(bits))
578                        .ok_or_else(|| Error::invalid_input("FSL value bit width overflows"))?;
579                    return Ok(Self {
580                        bits_per_item: flat.bits_per_value,
581                        bits_per_value,
582                        layers,
583                    });
584                }
585                // All inner values are null: only validity bytes are stored per row.
586                // bits_per_item=0 signals unzip_decompress to emit AllNull for the values.
587                Compression::Constant(_) => {
588                    return Ok(Self {
589                        bits_per_item: 0,
590                        bits_per_value: bytes_per_value.checked_mul(8).ok_or_else(|| {
591                            Error::invalid_input("FSL validity bit width overflows")
592                        })?,
593                        layers,
594                    });
595                }
596                _ => {
597                    return Err(Error::invalid_input(format!(
598                        "Unexpected inner encoding type in FSL descriptor: {:?}",
599                        encoding
600                    )));
601                }
602            }
603        }
604    }
605
606    fn buffer_to_block(&self, data: LanceBuffer, num_values: u64) -> DataBlock {
607        DataBlock::FixedWidth(FixedWidthDataBlock {
608            bits_per_value: self.bits_per_item,
609            num_values,
610            data,
611            block_info: BlockInfo::new(),
612        })
613    }
614}
615
616impl BlockDecompressor for ValueDecompressor {
617    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
618        let data = require_block_payload(data, "Flat block")?;
619        if !self.layers.is_empty() {
620            let expected = num_values
621                .checked_mul(self.bits_per_value)
622                .and_then(|bits| usize::try_from(bits.div_ceil(8)).ok())
623                .ok_or_else(|| Error::invalid_input("Fixed-size-list block length overflows"))?;
624            if data.len() != expected {
625                return Err(Error::invalid_input(format!(
626                    "Fixed-size-list block has {} bytes, expected {expected}",
627                    data.len()
628                )));
629            }
630            return FixedPerValueDecompressor::decompress(
631                self,
632                FixedWidthDataBlock {
633                    data,
634                    bits_per_value: self.bits_per_value,
635                    num_values,
636                    block_info: BlockInfo::new(),
637                },
638                num_values,
639            );
640        }
641        let block = self.buffer_to_block(data, num_values);
642        assert_eq!(block.num_values(), num_values);
643        Ok(block)
644    }
645}
646
647impl MiniBlockDecompressor for ValueDecompressor {
648    fn num_buffers(&self) -> usize {
649        1 + self
650            .layers
651            .iter()
652            .filter(|layer| layer.has_validity)
653            .count()
654    }
655
656    fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
657        if data.len() != self.num_buffers() {
658            return Err(Error::invalid_input(format!(
659                "Flat/FSL mini-block has {} buffers, expected {}",
660                data.len(),
661                self.num_buffers()
662            )));
663        }
664        let mut child_values = num_values;
665        let mut buffers = data.iter();
666        for layer in &self.layers {
667            child_values = child_values
668                .checked_mul(layer.dimension)
669                .ok_or_else(|| Error::invalid_input("FSL child value count overflows"))?;
670            if layer.has_validity {
671                validate_fixed_buffer(
672                    buffers
673                        .next()
674                        .ok_or_else(|| Error::invalid_input("FSL validity buffer is missing"))?,
675                    child_values,
676                    1,
677                )?;
678            }
679        }
680        validate_fixed_buffer(
681            buffers
682                .next()
683                .ok_or_else(|| Error::invalid_input("Flat value buffer is missing"))?,
684            child_values,
685            self.bits_per_item,
686        )?;
687        let num_items = child_values;
688        let mut buffer_iter = data.into_iter().rev();
689
690        // Always at least 1 buffer
691        let data_buf = buffer_iter.next().unwrap();
692        let items = self.buffer_to_block(data_buf, num_items);
693        let mut lists = items;
694
695        for layer in self.layers.iter().rev() {
696            if layer.has_validity {
697                let validity_buf = buffer_iter.next().unwrap();
698                lists = DataBlock::Nullable(NullableDataBlock {
699                    data: Box::new(lists),
700                    nulls: validity_buf,
701                    block_info: BlockInfo::default(),
702                });
703            }
704            lists = DataBlock::FixedSizeList(FixedSizeListBlock {
705                child: Box::new(lists),
706                dimension: layer.dimension,
707            })
708        }
709
710        assert_eq!(lists.num_values(), num_values);
711        Ok(lists)
712    }
713
714    fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
715        if self.has_validity() {
716            return None;
717        }
718        num_values
719            .checked_mul(self.bits_per_value)
720            .map(|bits| bits.div_ceil(8))
721    }
722}
723
724fn validate_fixed_buffer(data: &LanceBuffer, num_values: u64, bits: u64) -> Result<()> {
725    let expected = num_values
726        .checked_mul(bits)
727        .map(|bits| bits.div_ceil(8))
728        .ok_or_else(|| Error::invalid_input("Flat decoded bit length overflows"))?;
729    if data.len() as u64 != expected {
730        return Err(Error::invalid_input(format!(
731            "Flat buffer has {} bytes, expected {expected}",
732            data.len()
733        )));
734    }
735    Ok(())
736}
737
738struct FslDecompressorValidityBuilder {
739    buffer: BooleanBufferBuilder,
740    bits_per_row: usize,
741    bytes_per_row: usize,
742}
743
744// Helper methods for per-value decompression
745impl ValueDecompressor {
746    fn has_validity(&self) -> bool {
747        self.layers.iter().any(|layer| layer.has_validity)
748    }
749
750    // If there is no validity then decompression is zero-copy, we just need to restore any FSL layers
751    fn simple_decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> DataBlock {
752        let mut cum_dim = 1;
753        for layer in &self.layers {
754            cum_dim *= layer.dimension;
755        }
756        debug_assert_eq!(self.bits_per_item, data.bits_per_value / cum_dim);
757        let mut block = DataBlock::FixedWidth(FixedWidthDataBlock {
758            bits_per_value: self.bits_per_item,
759            num_values: num_rows * cum_dim,
760            data: data.data,
761            block_info: BlockInfo::new(),
762        });
763        for layer in self.layers.iter().rev() {
764            block = DataBlock::FixedSizeList(FixedSizeListBlock {
765                child: Box::new(block),
766                dimension: layer.dimension,
767            });
768        }
769        debug_assert_eq!(num_rows, block.num_values());
770        block
771    }
772
773    // If there is validity then it has been zipped in with the values and we must unzip it
774    fn unzip_decompress(&self, data: FixedWidthDataBlock, num_rows: usize) -> DataBlock {
775        // No support for full-zip on per-value encodings
776        assert_eq!(self.bits_per_item % 8, 0);
777        let bytes_per_item = self.bits_per_item / 8;
778        let mut buffer_builders = Vec::with_capacity(self.layers.len());
779        let mut cum_dim = 1;
780        let mut total_size_bytes = 0;
781        // First, go through the layers, setup our builders, allocate space
782        for layer in &self.layers {
783            cum_dim *= layer.dimension as usize;
784            if layer.has_validity {
785                let validity_size_bits = cum_dim;
786                let validity_size_bytes = validity_size_bits.div_ceil(8);
787                total_size_bytes += num_rows * validity_size_bytes;
788                buffer_builders.push(FslDecompressorValidityBuilder {
789                    buffer: BooleanBufferBuilder::new(validity_size_bits * num_rows),
790                    bits_per_row: cum_dim,
791                    bytes_per_row: validity_size_bytes,
792                })
793            }
794        }
795        let num_items = num_rows * cum_dim;
796        let data_size = num_items * bytes_per_item as usize;
797        total_size_bytes += data_size;
798        let mut data_buffer = Vec::with_capacity(data_size);
799
800        assert_eq!(data.data.len(), total_size_bytes);
801
802        let bytes_per_value = bytes_per_item as usize;
803        let data_bytes_per_row = bytes_per_value * cum_dim;
804
805        // Next, unzip
806        let mut data_offset = 0;
807        while data_offset < total_size_bytes {
808            for builder in buffer_builders.iter_mut() {
809                let start = data_offset * 8;
810                let end = start + builder.bits_per_row;
811                builder.buffer.append_packed_range(start..end, &data.data);
812                data_offset += builder.bytes_per_row;
813            }
814            let end = data_offset + data_bytes_per_row;
815            data_buffer.extend_from_slice(&data.data[data_offset..end]);
816            data_offset += data_bytes_per_row;
817        }
818
819        // Finally, restore the structure
820        //
821        // bits_per_item=0 means the terminal encoding is constant-null (all inner values
822        // are null and only validity bits were stored).  Use AllNull so that into_arrow
823        // produces a properly typed all-null array without a zero-width data buffer.
824        let mut block = if self.bits_per_item == 0 {
825            debug_assert!(data_buffer.is_empty());
826            DataBlock::AllNull(AllNullDataBlock {
827                num_values: num_items as u64,
828            })
829        } else {
830            DataBlock::FixedWidth(FixedWidthDataBlock {
831                bits_per_value: self.bits_per_item,
832                num_values: num_items as u64,
833                data: LanceBuffer::from(data_buffer),
834                block_info: BlockInfo::new(),
835            })
836        };
837
838        let mut validity_bufs = buffer_builders
839            .into_iter()
840            .rev()
841            .map(|mut b| LanceBuffer::from(b.buffer.finish().into_inner()));
842        for layer in self.layers.iter().rev() {
843            if layer.has_validity {
844                let nullable = NullableDataBlock {
845                    data: Box::new(block),
846                    nulls: validity_bufs.next().unwrap(),
847                    block_info: BlockInfo::new(),
848                };
849                block = DataBlock::Nullable(nullable);
850            }
851            block = DataBlock::FixedSizeList(FixedSizeListBlock {
852                child: Box::new(block),
853                dimension: layer.dimension,
854            });
855        }
856
857        assert_eq!(num_rows, block.num_values() as usize);
858
859        block
860    }
861}
862
863impl FixedPerValueDecompressor for ValueDecompressor {
864    fn decompress(&self, data: FixedWidthDataBlock, num_rows: u64) -> Result<DataBlock> {
865        if self.has_validity() {
866            Ok(self.unzip_decompress(data, num_rows as usize))
867        } else {
868            Ok(self.simple_decompress(data, num_rows))
869        }
870    }
871
872    fn bits_per_value(&self) -> u64 {
873        self.bits_per_value
874    }
875
876    fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
877        if self.has_validity() {
878            return None;
879        }
880        num_values
881            .checked_mul(self.bits_per_value)
882            .map(|bits| bits.div_ceil(8))
883    }
884}
885
886impl PerValueCompressor for ValueEncoder {
887    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
888        let (data, encoding) = match data {
889            DataBlock::FixedWidth(fixed_width) => {
890                let encoding = EncodingUtils::flat(fixed_width.bits_per_value, None);
891                (PerValueDataBlock::Fixed(fixed_width), encoding)
892            }
893            DataBlock::FixedSizeList(fixed_size_list) => Self::per_value_fsl(fixed_size_list),
894            _ => unimplemented!(
895                "Cannot compress block of type {} with ValueEncoder",
896                data.name()
897            ),
898        };
899        Ok((data, encoding))
900    }
901}
902
903// public tests module because we share the PRIMITIVE_TYPES constant with fixed_size_list
904#[cfg(test)]
905mod tests {
906    use std::{collections::HashMap, sync::Arc};
907
908    use arrow_array::{
909        Array, ArrayRef, Decimal128Array, FixedSizeListArray, Int32Array, ListArray, NullArray,
910        UInt8Array, make_array, new_null_array, types::UInt32Type,
911    };
912    use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
913    use arrow_schema::{DataType, Field, TimeUnit};
914    use lance_datagen::{ArrayGeneratorExt, Dimension, RowCount, Seed, array, gen_batch};
915
916    use crate::{
917        compression::{
918            DecompressionStrategy, DefaultDecompressionStrategy, FixedPerValueDecompressor,
919            MiniBlockDecompressor, try_fixed_size_list_block,
920        },
921        compression_config::CompressionFieldParams,
922        data::DataBlock,
923        encodings::{
924            logical::primitive::{
925                fullzip::{PerValueCompressor, PerValueDataBlock},
926                miniblock::{MiniBlockCompressionContext, MiniBlockCompressor},
927            },
928            physical::value::ValueDecompressor,
929        },
930        format::metadata::compressive_encoding::Compression,
931        testing::{
932            TestCases, TestEncoding, check_basic_random_case, check_round_trip_encoding_of_data,
933        },
934    };
935
936    use super::ValueEncoder;
937
938    #[rstest::rstest]
939    #[case::raw("none", false)]
940    #[case::lz4("lz4", false)]
941    #[case::zstd("zstd", false)]
942    #[case::all_null_children("lz4", true)]
943    #[test]
944    fn test_vector_dictionary_block_lengths(#[case] compression: &str, #[case] all_null: bool) {
945        let array = FixedSizeListArray::new(
946            Arc::new(Field::new("item", DataType::Int32, true)),
947            3,
948            Arc::new(if all_null {
949                Int32Array::new_null(6)
950            } else {
951                Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), None])
952            }),
953            None,
954        );
955        let data_type = array.data_type().clone();
956        let data = DataBlock::from_array(array.clone());
957        let params = CompressionFieldParams {
958            compression: Some(compression.to_string()),
959            ..Default::default()
960        };
961        let encoder = try_fixed_size_list_block(&data, &params).unwrap().unwrap();
962        let (payload, encoding) = encoder.compress(data).unwrap();
963        let decoder = DefaultDecompressionStrategy::default()
964            .create_block_decompressor(&encoding)
965            .unwrap();
966        let decoded = decoder.decompress(payload.clone(), 2).unwrap();
967        assert_eq!(
968            decoded.into_arrow(data_type, true).unwrap(),
969            array.to_data()
970        );
971        assert!(decoder.decompress(payload.clone(), 1).is_err());
972        assert!(decoder.decompress(payload.clone(), 3).is_err());
973        let payload = payload.unwrap();
974        assert!(
975            decoder
976                .decompress(Some(payload.slice_with_length(0, payload.len() - 1)), 2)
977                .is_err()
978        );
979    }
980
981    fn miniblock_context() -> MiniBlockCompressionContext {
982        MiniBlockCompressionContext::new(0, true, true)
983    }
984
985    const PRIMITIVE_TYPES: &[DataType] = &[
986        DataType::Null,
987        DataType::FixedSizeBinary(2),
988        DataType::Date32,
989        DataType::Date64,
990        DataType::Int8,
991        DataType::Int16,
992        DataType::Int32,
993        DataType::Int64,
994        DataType::UInt8,
995        DataType::UInt16,
996        DataType::UInt32,
997        DataType::UInt64,
998        DataType::Float16,
999        DataType::Float32,
1000        DataType::Float64,
1001        DataType::Decimal128(10, 10),
1002        DataType::Decimal256(10, 10),
1003        DataType::Timestamp(TimeUnit::Nanosecond, None),
1004        DataType::Time32(TimeUnit::Second),
1005        DataType::Time64(TimeUnit::Nanosecond),
1006        DataType::Duration(TimeUnit::Second),
1007        // The Interval type is supported by the reader but the writer works with Lance schema
1008        // at the moment and Lance schema can't parse interval
1009        // DataType::Interval(IntervalUnit::DayTime),
1010    ];
1011
1012    #[test_log::test(tokio::test)]
1013    async fn test_simple_value() {
1014        let items = Arc::new(Int32Array::from(vec![
1015            Some(0),
1016            None,
1017            Some(2),
1018            Some(3),
1019            Some(4),
1020            Some(5),
1021        ]));
1022
1023        let test_cases = TestCases::default()
1024            .with_range(0..3)
1025            .with_range(0..2)
1026            .with_range(1..3)
1027            .with_indices(vec![0, 1, 2])
1028            .with_indices(vec![1])
1029            .with_indices(vec![2])
1030            .with_structural_encodings();
1031
1032        check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await;
1033    }
1034
1035    #[test_log::test(tokio::test)]
1036    async fn test_simple_range() {
1037        let items = Arc::new(Int32Array::from_iter(
1038            (0..5000).map(|i| if i % 2 == 0 { Some(i) } else { None }),
1039        ));
1040
1041        let test_cases = TestCases::default().with_structural_encodings();
1042
1043        check_round_trip_encoding_of_data(vec![items], &test_cases, HashMap::default()).await;
1044    }
1045
1046    #[test_log::test(tokio::test)]
1047    async fn test_value_primitive() {
1048        const NUM_ROWS: u32 = 1025;
1049
1050        let test_cases = TestCases::default()
1051            .with_batch_size(NUM_ROWS)
1052            .with_page_sizes(vec![4096])
1053            .with_expected_encoding("flat");
1054        let value_metadata =
1055            HashMap::from([("lance-encoding:compression".to_string(), "none".to_string())]);
1056
1057        for (seed, data_type) in PRIMITIVE_TYPES.iter().enumerate() {
1058            log::info!("Testing encoding for {:?}", data_type);
1059            let data = gen_batch()
1060                .with_seed(Seed::from(seed as u64))
1061                .anon_col(array::rand_type(data_type))
1062                .into_batch_rows(RowCount::from(NUM_ROWS as u64))
1063                .unwrap()
1064                .column(0)
1065                .clone();
1066
1067            check_round_trip_encoding_of_data(vec![data], &test_cases, value_metadata.clone())
1068                .await;
1069        }
1070    }
1071
1072    #[rstest::rstest]
1073    #[test_log::test(tokio::test)]
1074    async fn test_large_primitive(
1075        #[values(
1076            TestEncoding::Array,
1077            TestEncoding::StructuralU16,
1078            TestEncoding::StructuralU32,
1079            TestEncoding::StructuralSparse
1080        )]
1081        encoding: TestEncoding,
1082        #[values(4096, 1024 * 1024)] page_size: u64,
1083        #[values(false, true)] use_slicing: bool,
1084    ) {
1085        let data_type =
1086            DataType::FixedSizeList(Arc::new(Field::new("", DataType::Int32, false)), 128);
1087        let field = Field::new("", data_type, false);
1088        check_basic_random_case(field, encoding, page_size, use_slicing).await;
1089    }
1090
1091    #[rstest::rstest]
1092    #[test_log::test(tokio::test)]
1093    async fn test_decimal128_dictionary_encoding(
1094        #[values(
1095            TestEncoding::StructuralU16,
1096            TestEncoding::StructuralU32,
1097            TestEncoding::StructuralSparse
1098        )]
1099        encoding: TestEncoding,
1100    ) {
1101        let test_cases = TestCases::default()
1102            .with_encoding(encoding)
1103            .with_expected_encoding("dictionary");
1104        let decimals: Vec<i32> = (0..100).collect();
1105        let repeated_strings: Vec<_> = decimals
1106            .iter()
1107            .cycle()
1108            .take(decimals.len() * 1000)
1109            .map(|&v| Some(v as i128))
1110            .collect();
1111        let decimal_array = Arc::new(Decimal128Array::from(repeated_strings)) as ArrayRef;
1112        check_round_trip_encoding_of_data(vec![decimal_array], &test_cases, HashMap::new()).await;
1113    }
1114
1115    #[rstest::rstest]
1116    #[test_log::test(tokio::test)]
1117    async fn test_miniblock_stress(
1118        #[values(false, true)] mixed_validity: bool,
1119        #[values(10, 100, 1500, 15000)] batch_size: u32,
1120        #[values(1000, 2000, 3000, 60000)] page_size: u64,
1121        #[values(
1122            TestEncoding::StructuralU16,
1123            TestEncoding::StructuralU32,
1124            TestEncoding::StructuralSparse
1125        )]
1126        encoding: TestEncoding,
1127    ) {
1128        // Tests for strange page sizes and batch sizes and validity scenarios for miniblock
1129
1130        // 10K integers, 100 per array, either all valid or mixed validity.
1131        let data = (0..100)
1132            .map(|_| {
1133                if mixed_validity {
1134                    Arc::new(Int32Array::from_iter(
1135                        (0..100).map(|i| if i % 2 == 0 { Some(i) } else { None }),
1136                    )) as Arc<dyn Array>
1137                } else {
1138                    Arc::new(Int32Array::from_iter_values(0..100)) as Arc<dyn Array>
1139                }
1140            })
1141            .collect::<Vec<_>>();
1142
1143        let test_cases = TestCases::default()
1144            .with_page_sizes(vec![page_size])
1145            .with_batch_size(batch_size)
1146            .with_encoding(encoding);
1147
1148        check_round_trip_encoding_of_data(data, &test_cases, HashMap::new()).await;
1149    }
1150
1151    fn create_simple_fsl() -> FixedSizeListArray {
1152        // [[0, 1], NULL], [NULL, NULL], [[8, 9], [NULL, 11]]
1153        let items = Arc::new(Int32Array::from(vec![
1154            Some(0),
1155            Some(1),
1156            Some(2),
1157            Some(3),
1158            None,
1159            None,
1160            None,
1161            None,
1162            Some(8),
1163            Some(9),
1164            None,
1165            Some(11),
1166        ]));
1167        let items_field = Arc::new(Field::new("item", DataType::Int32, true));
1168        let inner_list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]);
1169        let inner_list = Arc::new(FixedSizeListArray::new(
1170            items_field.clone(),
1171            2,
1172            items,
1173            Some(NullBuffer::new(inner_list_nulls)),
1174        ));
1175        let inner_list_field = Arc::new(Field::new(
1176            "item",
1177            DataType::FixedSizeList(items_field, 2),
1178            true,
1179        ));
1180        FixedSizeListArray::new(inner_list_field, 2, inner_list, None)
1181    }
1182
1183    #[test]
1184    fn test_fsl_value_compression_miniblock() {
1185        let sample_list = create_simple_fsl();
1186
1187        let starting_data = DataBlock::from_array(sample_list.clone());
1188
1189        let encoder = ValueEncoder::default();
1190        let (data, compression) =
1191            MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap();
1192
1193        assert_eq!(data.num_values, 3);
1194        assert_eq!(data.data.len(), 3);
1195        assert_eq!(data.chunks.len(), 1);
1196        assert_eq!(data.chunks[0].buffer_sizes, vec![1, 2, 48]);
1197        assert_eq!(data.chunks[0].log_num_values, 0);
1198
1199        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1200            panic!()
1201        };
1202
1203        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1204
1205        let decompressed =
1206            MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap();
1207
1208        let decompressed = make_array(
1209            decompressed
1210                .into_arrow(sample_list.data_type().clone(), true)
1211                .unwrap(),
1212        );
1213
1214        assert_eq!(decompressed.as_ref(), &sample_list);
1215    }
1216
1217    fn wide_fixed_size_binary() -> ArrayRef {
1218        let wide_value = vec![0xABu8; 5000];
1219        Arc::new(
1220            arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1221                std::iter::repeat_n(Some(wide_value.as_slice()), 4),
1222                5000,
1223            )
1224            .unwrap(),
1225        )
1226    }
1227
1228    fn wide_fixed_size_list_bool() -> ArrayRef {
1229        // A wide FSL<Boolean> is sub-byte, so it chunks eight values per word and the
1230        // smallest unit is 16 values rather than 2.
1231        let dimension = 4095;
1232        let values = arrow_array::BooleanArray::from(vec![false; dimension * 2]);
1233        let field = Arc::new(Field::new("item", DataType::Boolean, true));
1234        Arc::new(FixedSizeListArray::new(
1235            field,
1236            dimension as i32,
1237            Arc::new(values),
1238            None,
1239        ))
1240    }
1241
1242    #[rstest::rstest]
1243    #[case::fixed_size_binary(wide_fixed_size_binary(), 2)]
1244    #[case::fixed_size_list_bool(wide_fixed_size_list_bool(), 16)]
1245    fn test_wide_value_miniblock_returns_error(
1246        #[case] array: ArrayRef,
1247        #[case] expected_min_values: u64,
1248    ) {
1249        let starting_data = DataBlock::from_array(array);
1250
1251        let encoder = ValueEncoder::default();
1252        let result = MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data);
1253
1254        let err = result.expect_err("wide values should not be encodable as miniblock");
1255        assert!(
1256            matches!(err, lance_core::Error::InvalidInput { .. }),
1257            "expected InvalidInput, got {err:?}"
1258        );
1259        let msg = err.to_string();
1260        assert!(
1261            msg.contains("too wide for miniblock encoding"),
1262            "unexpected error message: {msg}"
1263        );
1264        assert!(
1265            msg.contains(&format!("{expected_min_values} values require")),
1266            "unexpected error message: {msg}"
1267        );
1268    }
1269
1270    #[test]
1271    fn test_fsl_value_compression_per_value() {
1272        let sample_list = create_simple_fsl();
1273
1274        let starting_data = DataBlock::from_array(sample_list.clone());
1275
1276        let encoder = ValueEncoder::default();
1277        let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap();
1278
1279        let PerValueDataBlock::Fixed(data) = data else {
1280            panic!()
1281        };
1282
1283        assert_eq!(data.bits_per_value, 144);
1284        assert_eq!(data.num_values, 3);
1285        assert_eq!(data.data.len(), 18 * 3);
1286
1287        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1288            panic!()
1289        };
1290
1291        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1292
1293        let num_values = data.num_values;
1294        assert_eq!(
1295            FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_values),
1296            None,
1297            "nullable FSL output uses multiple buffers and requires the fallback estimate"
1298        );
1299        let decompressed =
1300            FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1301
1302        let decompressed = make_array(
1303            decompressed
1304                .into_arrow(sample_list.data_type().clone(), true)
1305                .unwrap(),
1306        );
1307
1308        assert_eq!(decompressed.as_ref(), &sample_list);
1309    }
1310
1311    /// A null-free nested FSL is flattened by `simple_per_value_fsl`; each value
1312    /// spans `outer * inner` leaves.
1313    #[rstest::rstest]
1314    #[case::outer_wider(22, 3)]
1315    #[case::inner_wider(3, 22)]
1316    fn test_nested_fsl_simple_per_value(#[case] outer_dim: i32, #[case] inner_dim: i32) {
1317        let num_rows = 3;
1318        let inner_field = Arc::new(Field::new("item", DataType::Int32, false));
1319        let leaf = Arc::new(Int32Array::from_iter_values(
1320            0..num_rows * outer_dim * inner_dim,
1321        )) as ArrayRef;
1322        let inner = Arc::new(FixedSizeListArray::new(
1323            inner_field.clone(),
1324            inner_dim,
1325            leaf,
1326            None,
1327        )) as ArrayRef;
1328        let outer = Arc::new(FixedSizeListArray::new(
1329            Arc::new(Field::new(
1330                "item",
1331                DataType::FixedSizeList(inner_field, inner_dim),
1332                false,
1333            )),
1334            outer_dim,
1335            inner,
1336            None,
1337        )) as ArrayRef;
1338
1339        let encoder = ValueEncoder::default();
1340        let (data, compression) =
1341            PerValueCompressor::compress(&encoder, DataBlock::from_array(outer.clone())).unwrap();
1342        let PerValueDataBlock::Fixed(data) = data else {
1343            panic!()
1344        };
1345        assert_eq!(data.num_values, num_rows as u64);
1346        assert_eq!(data.bits_per_value, 32 * (outer_dim * inner_dim) as u64);
1347
1348        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1349            panic!()
1350        };
1351        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1352        let num_values = data.num_values;
1353        let decompressed =
1354            FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1355        let decompressed = make_array(
1356            decompressed
1357                .into_arrow(outer.data_type().clone(), true)
1358                .unwrap(),
1359        );
1360        assert_eq!(decompressed.as_ref(), outer.as_ref());
1361    }
1362
1363    #[test_log::test(tokio::test)]
1364    async fn test_fsl_all_null() {
1365        let items = new_null_array(&DataType::Int32, 12);
1366        let items_field = Arc::new(Field::new("item", DataType::Int32, true));
1367        let list_nulls = BooleanBuffer::from(vec![true, false, false, false, true, true]);
1368        let list_array =
1369            FixedSizeListArray::new(items_field, 2, items, Some(NullBuffer::new(list_nulls)));
1370
1371        let test_cases = TestCases::default().with_structural_encodings();
1372
1373        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
1374            .await;
1375    }
1376
1377    // A FixedSizeList of the Null type wrote fine but failed to read back, because the
1378    // decoder gave the Null child a validity buffer, which arrow-rs rejects.
1379    #[rstest::rstest]
1380    #[test_log::test(tokio::test)]
1381    async fn test_fsl_of_null_type(#[values(false, true)] has_outer_nulls: bool) {
1382        let items_field = Arc::new(Field::new("item", DataType::Null, true));
1383        let outer_nulls = has_outer_nulls.then(|| NullBuffer::from(vec![true, false, true]));
1384        let list_array =
1385            FixedSizeListArray::new(items_field, 2, Arc::new(NullArray::new(6)), outer_nulls);
1386
1387        let test_cases = TestCases::default().with_structural_encodings();
1388
1389        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
1390            .await;
1391    }
1392
1393    // Regression: nullable_per_value_fsl wrote bits_per_value=0 when the child block was
1394    // AllNull but the FSL had no outer null buffer (no ctrl-word bytes), making
1395    // total_bytes_per_value=0 and crashing the reader with "per-row byte width must be > 0".
1396    #[test_log::test(tokio::test)]
1397    async fn test_fsl_nullable_child_all_null_no_outer_nulls() {
1398        // FSL<nullable Float32, dim=4>, 2 outer rows, no outer null buffer, all 8 child
1399        // Float32 values are null.  DataBlock::from_arrays returns AllNull for the child,
1400        // which routes into nullable_per_value_fsl.  Without this fix the encoder writes
1401        // bits_per_value=0 and the decoder errors on read.
1402        let items = new_null_array(&DataType::Float32, 8);
1403        let items_field = Arc::new(Field::new("item", DataType::Float32, true));
1404        let fsl = FixedSizeListArray::new(items_field, 4, items, None);
1405
1406        let test_cases = TestCases::default().with_structural_encodings();
1407        check_round_trip_encoding_of_data(vec![Arc::new(fsl)], &test_cases, HashMap::new()).await;
1408    }
1409
1410    #[test_log::test(tokio::test)]
1411    async fn regress_list_fsl() {
1412        // This regresses a case where rows are large lists that span multiple
1413        // mini-block chunks which gives us some all-premable mini-block chunks.
1414        let offsets = ScalarBuffer::<i32>::from(vec![0, 393, 755, 1156, 1536]);
1415        let data = UInt8Array::from(vec![0; 1536 * 16]);
1416        let fsl_field = Arc::new(Field::new("item", DataType::UInt8, true));
1417        let fsl = FixedSizeListArray::new(fsl_field, 16, Arc::new(data), None);
1418        let list_field = Arc::new(Field::new("item", fsl.data_type().clone(), false));
1419        let list_arr = ListArray::new(list_field, OffsetBuffer::new(offsets), Arc::new(fsl), None);
1420
1421        let test_cases = TestCases::default()
1422            .with_structural_encodings()
1423            .with_batch_size(1);
1424
1425        check_round_trip_encoding_of_data(vec![Arc::new(list_arr)], &test_cases, HashMap::new())
1426            .await;
1427    }
1428
1429    fn create_random_fsl() -> Arc<dyn Array> {
1430        // Several levels of def and multiple pages
1431        let inner = array::rand_type(&DataType::Int32).with_random_nulls(0.1);
1432        let list_one = array::cycle_vec(inner, Dimension::from(4)).with_random_nulls(0.1);
1433        let list_two = array::cycle_vec(list_one, Dimension::from(4)).with_random_nulls(0.1);
1434        let list_three = array::cycle_vec(list_two, Dimension::from(2));
1435
1436        // Should be 256Ki rows ~ 1MiB of data
1437        let batch = gen_batch()
1438            .anon_col(list_three)
1439            .into_batch_rows(RowCount::from(8 * 1024))
1440            .unwrap();
1441        batch.column(0).clone()
1442    }
1443
1444    #[test]
1445    fn fsl_value_miniblock_stress() {
1446        let sample_array = create_random_fsl();
1447
1448        let starting_data = DataBlock::from_arrays(
1449            std::slice::from_ref(&sample_array),
1450            sample_array.len() as u64,
1451        );
1452
1453        let encoder = ValueEncoder::default();
1454        let (data, compression) =
1455            MiniBlockCompressor::compress(&encoder, miniblock_context(), starting_data).unwrap();
1456
1457        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1458            panic!()
1459        };
1460
1461        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1462
1463        let decompressed =
1464            MiniBlockDecompressor::decompress(&decompressor, data.data, data.num_values).unwrap();
1465
1466        let decompressed = make_array(
1467            decompressed
1468                .into_arrow(sample_array.data_type().clone(), true)
1469                .unwrap(),
1470        );
1471
1472        assert_eq!(decompressed.as_ref(), sample_array.as_ref());
1473    }
1474
1475    #[test]
1476    fn fsl_value_per_value_stress() {
1477        let sample_array = create_random_fsl();
1478
1479        let starting_data = DataBlock::from_arrays(
1480            std::slice::from_ref(&sample_array),
1481            sample_array.len() as u64,
1482        );
1483
1484        let encoder = ValueEncoder::default();
1485        let (data, compression) = PerValueCompressor::compress(&encoder, starting_data).unwrap();
1486
1487        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
1488            panic!()
1489        };
1490
1491        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
1492
1493        let PerValueDataBlock::Fixed(data) = data else {
1494            panic!()
1495        };
1496
1497        let num_values = data.num_values;
1498        let decompressed =
1499            FixedPerValueDecompressor::decompress(&decompressor, data, num_values).unwrap();
1500
1501        let decompressed = make_array(
1502            decompressed
1503                .into_arrow(sample_array.data_type().clone(), true)
1504                .unwrap(),
1505        );
1506
1507        assert_eq!(decompressed.as_ref(), sample_array.as_ref());
1508    }
1509
1510    #[rstest::rstest]
1511    #[test_log::test(tokio::test)]
1512    async fn test_fsl_nullable_items(
1513        #[values(
1514            TestEncoding::Array,
1515            TestEncoding::StructuralU16,
1516            TestEncoding::StructuralU32,
1517            TestEncoding::StructuralSparse
1518        )]
1519        encoding: TestEncoding,
1520    ) {
1521        let mut generator =
1522            gen_batch()
1523                .with_seed(Seed::from(0))
1524                .anon_col(array::rand_vec_nullable::<UInt32Type>(
1525                    Dimension::from(128),
1526                    0.5,
1527                ));
1528        generator.with_random_nulls(0.2);
1529        let source = generator
1530            .into_batch_rows(RowCount::from(1026))
1531            .unwrap()
1532            .column(0)
1533            .clone();
1534        let test_cases = TestCases::default()
1535            .with_page_sizes(vec![4096])
1536            .with_encoding(encoding)
1537            .with_batch_size(257)
1538            .with_range(510..515)
1539            .with_indices(vec![0, 511, 512, 1024]);
1540
1541        check_round_trip_encoding_of_data(
1542            vec![source.slice(1, 512), source.slice(513, 513)],
1543            &test_cases,
1544            HashMap::new(),
1545        )
1546        .await;
1547    }
1548
1549    #[test_log::test(tokio::test)]
1550    async fn test_value_encoding_verification() {
1551        use std::collections::HashMap;
1552
1553        let test_cases = TestCases::default()
1554            .with_expected_encoding("flat")
1555            .with_structural_encodings();
1556
1557        // Test both explicit configuration and automatic fallback scenarios
1558        // 1. Test explicit "none" compression to force flat encoding
1559        // Also explicitly disable BSS to ensure value encoding is tested
1560        let mut metadata_explicit = HashMap::new();
1561        metadata_explicit.insert("lance-encoding:compression".to_string(), "none".to_string());
1562        metadata_explicit.insert("lance-encoding:bss".to_string(), "off".to_string());
1563
1564        let arr_explicit =
1565            Arc::new(Int32Array::from((0..1000).collect::<Vec<i32>>())) as Arc<dyn Array>;
1566        check_round_trip_encoding_of_data(vec![arr_explicit], &test_cases, metadata_explicit).await;
1567
1568        // 2. Test automatic fallback to flat encoding when bitpacking conditions aren't met
1569        // Use unique values to avoid RLE encoding
1570        // Explicitly disable BSS to ensure value encoding is tested
1571        let mut metadata = HashMap::new();
1572        metadata.insert("lance-encoding:bss".to_string(), "off".to_string());
1573
1574        let arr_fallback = Arc::new(Int32Array::from(
1575            (0..100).map(|i| i * 73 + 19).collect::<Vec<i32>>(),
1576        )) as Arc<dyn Array>;
1577        check_round_trip_encoding_of_data(vec![arr_fallback], &test_cases, metadata).await;
1578    }
1579
1580    #[test_log::test(tokio::test)]
1581    async fn test_mixed_page_validity() {
1582        let no_nulls = Arc::new(Int32Array::from_iter_values([1, 2]));
1583        let has_nulls = Arc::new(Int32Array::from_iter([Some(3), None, Some(5)]));
1584
1585        let test_cases = TestCases::default().with_page_sizes(vec![1]);
1586        check_round_trip_encoding_of_data(vec![no_nulls, has_nulls], &test_cases, HashMap::new())
1587            .await;
1588    }
1589}