lance_encoding/encodings/physical/
packed.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Packed encoding
5//!
6//! These encodings take struct data and compress it in a way that all fields are collected
7//! together.
8//!
9//! This encoding can be transparent or opaque.  In order to be transparent we must use transparent
10//! compression on all children.  Then we can zip together the compressed children.
11
12use std::{convert::TryInto, sync::Arc};
13
14use arrow_array::types::UInt64Type;
15use lance_core::{Error, Result, datatypes::Field};
16
17use crate::{
18    buffer::LanceBuffer,
19    compression::{
20        CompressionStrategy, FixedPerValueDecompressor, MiniBlockDecompressor,
21        VariablePerValueDecompressor,
22    },
23    data::{
24        BlockInfo, DataBlock, DataBlockBuilder, FixedWidthDataBlock, StructDataBlock,
25        VariableWidthBlock,
26    },
27    encodings::logical::primitive::{
28        fullzip::{PerValueCompressor, PerValueDataBlock},
29        miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor},
30    },
31    format::{
32        EncodingUtils,
33        metadata::{CompressiveEncoding, PackedStruct, compressive_encoding::Compression},
34    },
35    statistics::{GetStat, Stat},
36};
37
38use super::value::{ValueDecompressor, ValueEncoder};
39
40// Transforms a `StructDataBlock` into a row major `FixedWidthDataBlock`.
41// Only fields with fixed-width fields are supported for now, and the
42// assumption that all fields has `bits_per_value % 8 == 0` is made.
43fn struct_data_block_to_fixed_width_data_block(
44    struct_data_block: StructDataBlock,
45    bits_per_values: &[u64],
46) -> DataBlock {
47    let data_size = struct_data_block.expect_single_stat::<UInt64Type>(Stat::DataSize);
48    let mut output = Vec::with_capacity(data_size as usize);
49    let num_values = struct_data_block.children[0].num_values();
50
51    for i in 0..num_values as usize {
52        for (j, child) in struct_data_block.children.iter().enumerate() {
53            let bytes_per_value = (bits_per_values[j] / 8) as usize;
54            let this_data = child
55                .as_fixed_width_ref()
56                .unwrap()
57                .data
58                .slice_with_length(bytes_per_value * i, bytes_per_value);
59            output.extend_from_slice(&this_data);
60        }
61    }
62
63    DataBlock::FixedWidth(FixedWidthDataBlock {
64        bits_per_value: bits_per_values.iter().copied().sum(),
65        data: LanceBuffer::from(output),
66        num_values,
67        block_info: BlockInfo::default(),
68    })
69}
70
71#[derive(Debug, Default)]
72pub struct PackedStructFixedWidthMiniBlockEncoder {}
73
74impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder {
75    fn compress(
76        &self,
77        context: MiniBlockCompressionContext,
78        data: DataBlock,
79    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
80        match data {
81            DataBlock::Struct(struct_data_block) => {
82                let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::<Vec<_>>();
83
84                // transform struct datablock to fixed-width data block.
85                let data_block = struct_data_block_to_fixed_width_data_block(struct_data_block, &bits_per_values);
86
87                // store and transformed fixed-width data block.
88                let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box<dyn MiniBlockCompressor>;
89                let (value_miniblock_compressed, value_array_encoding) =
90                value_miniblock_compressor.compress(context, data_block)?;
91
92                Ok((
93                    value_miniblock_compressed,
94                    EncodingUtils::packed_struct(value_array_encoding, bits_per_values),
95                ))
96            }
97            _ => Err(Error::invalid_input_source(format!(
98                "Cannot compress a data block of type {} with PackedStructFixedWidthBlockEncoder",
99                data.name()
100            )
101            .into())),
102        }
103    }
104}
105
106#[derive(Debug)]
107pub struct PackedStructFixedWidthMiniBlockDecompressor {
108    bits_per_values: Vec<u64>,
109    array_encoding: Box<dyn MiniBlockDecompressor>,
110}
111
112impl PackedStructFixedWidthMiniBlockDecompressor {
113    /// Validates a fixed-width packed descriptor before constructing its decoder.
114    pub(crate) fn try_new(description: &PackedStruct) -> Result<Self> {
115        let Some(Compression::Flat(flat)) = description
116            .values
117            .as_deref()
118            .and_then(|v| v.compression.as_ref())
119        else {
120            return Err(Error::invalid_input("PackedStruct requires Flat values"));
121        };
122        let bits = description
123            .bits_per_value
124            .iter()
125            .try_fold(0_u64, |sum, &width| {
126                if width == 0 || !width.is_multiple_of(8) {
127                    return Err(Error::invalid_input(
128                        "PackedStruct field widths must be positive byte widths",
129                    ));
130                }
131                sum.checked_add(width)
132                    .ok_or_else(|| Error::invalid_input("PackedStruct bit width sum overflows"))
133            })?;
134        if bits == 0 || bits != flat.bits_per_value || flat.data.is_some() {
135            return Err(Error::invalid_input(
136                "PackedStruct field widths must match its plain Flat values",
137            ));
138        }
139        Ok(Self {
140            bits_per_values: description.bits_per_value.clone(),
141            array_encoding: Box::new(ValueDecompressor::from_flat(flat)),
142        })
143    }
144
145    pub fn new(description: &PackedStruct) -> Self {
146        let array_encoding: Box<dyn MiniBlockDecompressor> = match description
147            .values
148            .as_ref()
149            .unwrap()
150            .compression
151            .as_ref()
152            .unwrap()
153        {
154            Compression::Flat(flat) => Box::new(ValueDecompressor::from_flat(flat)),
155            _ => panic!(
156                "Currently only `ArrayEncoding::Flat` is supported in packed struct encoding in Lance 2.1."
157            ),
158        };
159        Self {
160            bits_per_values: description.bits_per_value.clone(),
161            array_encoding,
162        }
163    }
164}
165
166impl MiniBlockDecompressor for PackedStructFixedWidthMiniBlockDecompressor {
167    fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
168        let encoded_data_block = self.array_encoding.decompress(data, num_values)?;
169        let DataBlock::FixedWidth(encoded_data_block) = encoded_data_block else {
170            panic!("ValueDecompressor should output FixedWidth DataBlock")
171        };
172
173        let bytes_per_values = self
174            .bits_per_values
175            .iter()
176            .map(|bits_per_value| *bits_per_value as usize / 8)
177            .collect::<Vec<_>>();
178
179        assert!(encoded_data_block.bits_per_value % 8 == 0);
180        let encoded_bytes_per_row = (encoded_data_block.bits_per_value / 8) as usize;
181
182        // use a prefix_sum vector as a helper to reconstruct to `StructDataBlock`.
183        let mut prefix_sum = vec![0; self.bits_per_values.len()];
184        for i in 0..(self.bits_per_values.len() - 1) {
185            prefix_sum[i + 1] = prefix_sum[i] + bytes_per_values[i];
186        }
187
188        let mut children_data_block = vec![];
189        for i in 0..self.bits_per_values.len() {
190            let child_buf_size = bytes_per_values[i] * num_values as usize;
191            let mut child_buf: Vec<u8> = Vec::with_capacity(child_buf_size);
192
193            for j in 0..num_values as usize {
194                // the start of the data at this row is `j * encoded_bytes_per_row`, and the offset for this field is `prefix_sum[i]`, this field has length `bytes_per_values[i]`.
195                let this_value = encoded_data_block.data.slice_with_length(
196                    prefix_sum[i] + (j * encoded_bytes_per_row),
197                    bytes_per_values[i],
198                );
199
200                child_buf.extend_from_slice(&this_value);
201            }
202
203            let child = DataBlock::FixedWidth(FixedWidthDataBlock {
204                data: LanceBuffer::from(child_buf),
205                bits_per_value: self.bits_per_values[i],
206                num_values,
207                block_info: BlockInfo::default(),
208            });
209            children_data_block.push(child);
210        }
211        Ok(DataBlock::Struct(StructDataBlock {
212            children: children_data_block,
213            block_info: BlockInfo::default(),
214            validity: None,
215        }))
216    }
217}
218
219#[derive(Debug)]
220struct FixedPackedFieldData {
221    block: FixedWidthDataBlock,
222}
223
224impl FixedPackedFieldData {
225    fn append_row_bytes(&self, row_idx: usize, output: &mut Vec<u8>) -> Result<()> {
226        let bits_per_value = self.block.bits_per_value;
227        if !bits_per_value.is_multiple_of(8) {
228            return Err(Error::invalid_input(
229                "Packed struct encoding requires byte-aligned fixed-width children",
230            ));
231        }
232        let bytes_per_value = (bits_per_value / 8) as usize;
233        let start = row_idx
234            .checked_mul(bytes_per_value)
235            .ok_or_else(|| Error::invalid_input("Packed struct row size overflow"))?;
236        let end = start.checked_add(bytes_per_value).ok_or_else(|| {
237            Error::invalid_input(format!(
238                "Packed struct fixed child range overflow: row_idx={row_idx}, \
239                 bytes_per_value={bytes_per_value}"
240            ))
241        })?;
242        let data = self.block.data.as_ref();
243        if end > data.len() {
244            return Err(Error::invalid_input(
245                "Packed struct fixed child out of bounds",
246            ));
247        }
248        output.extend_from_slice(&data[start..end]);
249        Ok(())
250    }
251}
252
253#[derive(Debug)]
254enum VariablePackedFieldData {
255    Fixed(FixedPackedFieldData),
256    Variable {
257        block: VariableWidthBlock,
258        bits_per_length: u64,
259    },
260}
261
262impl VariablePackedFieldData {
263    fn append_row_bytes(&self, row_idx: usize, output: &mut Vec<u8>) -> Result<()> {
264        match self {
265            Self::Fixed(fixed_data) => fixed_data.append_row_bytes(row_idx, output),
266            Self::Variable {
267                block,
268                bits_per_length,
269            } => {
270                if !bits_per_length.is_multiple_of(8) {
271                    return Err(Error::invalid_input(
272                        "Packed struct variable children must have byte-aligned length prefixes",
273                    ));
274                }
275                let prefix_bytes = (*bits_per_length / 8) as usize;
276                if !(prefix_bytes == 4 || prefix_bytes == 8) {
277                    return Err(Error::invalid_input(
278                        "Packed struct variable children must use 32 or 64-bit length prefixes",
279                    ));
280                }
281                match block.bits_per_offset {
282                    32 => {
283                        let offsets = block.offsets.borrow_to_typed_slice::<u32>();
284                        let start = offsets[row_idx] as usize;
285                        let end = offsets[row_idx + 1] as usize;
286                        if end > block.data.len() {
287                            return Err(Error::invalid_input(
288                                "Packed struct variable child offsets out of bounds",
289                            ));
290                        }
291                        let len = (end - start) as u32;
292                        if prefix_bytes != std::mem::size_of::<u32>() {
293                            return Err(Error::invalid_input(
294                                "Packed struct variable child length prefix mismatch",
295                            ));
296                        }
297                        output.extend_from_slice(&len.to_le_bytes());
298                        output.extend_from_slice(&block.data[start..end]);
299                        Ok(())
300                    }
301                    64 => {
302                        let offsets = block.offsets.borrow_to_typed_slice::<u64>();
303                        let start = offsets[row_idx] as usize;
304                        let end = offsets[row_idx + 1] as usize;
305                        if end > block.data.len() {
306                            return Err(Error::invalid_input(
307                                "Packed struct variable child offsets out of bounds",
308                            ));
309                        }
310                        let len = (end - start) as u64;
311                        if prefix_bytes != std::mem::size_of::<u64>() {
312                            return Err(Error::invalid_input(
313                                "Packed struct variable child length prefix mismatch",
314                            ));
315                        }
316                        output.extend_from_slice(&len.to_le_bytes());
317                        output.extend_from_slice(&block.data[start..end]);
318                        Ok(())
319                    }
320                    _ => Err(Error::invalid_input(
321                        "Packed struct variable child must use 32 or 64-bit offsets",
322                    )),
323                }
324            }
325        }
326    }
327}
328
329fn check_struct_validity(data: DataBlock, field_length: usize) -> Result<StructDataBlock> {
330    let DataBlock::Struct(struct_block) = data else {
331        return Err(Error::invalid_input(
332            "Packed struct encoder requires Struct data block",
333        ));
334    };
335
336    if struct_block.children.is_empty() {
337        return Err(Error::invalid_input(
338            "Packed struct encoder requires at least one child field",
339        ));
340    }
341    if struct_block.children.len() != field_length {
342        return Err(Error::invalid_input(
343            "Struct field metadata does not match number of children",
344        ));
345    }
346
347    let num_values = struct_block.children[0].num_values();
348    for child in struct_block.children.iter() {
349        if child.num_values() != num_values {
350            return Err(Error::invalid_input(
351                "Packed struct children must have matching value counts",
352            ));
353        }
354    }
355    Ok(struct_block)
356}
357
358#[derive(Debug)]
359pub struct PackedStructVariablePerValueEncoder {
360    strategy: Arc<dyn CompressionStrategy>,
361    fields: Vec<Field>,
362}
363
364impl PackedStructVariablePerValueEncoder {
365    pub fn new(strategy: Arc<dyn CompressionStrategy>, fields: Vec<Field>) -> Self {
366        Self { strategy, fields }
367    }
368}
369
370impl PerValueCompressor for PackedStructVariablePerValueEncoder {
371    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
372        let struct_block = check_struct_validity(data, self.fields.len())?;
373        let num_values = struct_block.children[0].num_values();
374        let mut field_data = Vec::with_capacity(self.fields.len());
375        let mut field_metadata = Vec::with_capacity(self.fields.len());
376
377        for (field, child_block) in self.fields.iter().zip(struct_block.children) {
378            let compressor = self.strategy.create_per_value(field, &child_block)?;
379            let (compressed, encoding) = compressor.compress(child_block)?;
380            match compressed {
381                PerValueDataBlock::Fixed(block) => {
382                    field_metadata.push(EncodingUtils::packed_struct_field_fixed(
383                        encoding,
384                        block.bits_per_value,
385                    ));
386                    let block = FixedPackedFieldData { block };
387                    field_data.push(VariablePackedFieldData::Fixed(block));
388                }
389                PerValueDataBlock::Variable(block) => {
390                    let bits_per_length = block.bits_per_offset as u64;
391                    field_metadata.push(EncodingUtils::packed_struct_field_variable(
392                        encoding,
393                        bits_per_length,
394                    ));
395                    field_data.push(VariablePackedFieldData::Variable {
396                        block,
397                        bits_per_length,
398                    });
399                }
400            }
401        }
402
403        let mut row_data: Vec<u8> = Vec::new();
404        let mut row_offsets: Vec<u64> = Vec::with_capacity(num_values as usize + 1);
405        row_offsets.push(0);
406        let mut total_bytes: usize = 0;
407        let mut max_row_len: usize = 0;
408        for row in 0..num_values as usize {
409            let start = row_data.len();
410            for field in &field_data {
411                field.append_row_bytes(row, &mut row_data)?;
412            }
413            let end = row_data.len();
414            let row_len = end - start;
415            max_row_len = max_row_len.max(row_len);
416            total_bytes = total_bytes
417                .checked_add(row_len)
418                .ok_or_else(|| Error::invalid_input("Packed struct row data size overflow"))?;
419            row_offsets.push(end as u64);
420        }
421        debug_assert_eq!(total_bytes, row_data.len());
422
423        let use_u32_offsets = total_bytes <= u32::MAX as usize && max_row_len <= u32::MAX as usize;
424        let bits_per_offset = if use_u32_offsets { 32 } else { 64 };
425        let offsets_buffer = if use_u32_offsets {
426            let offsets_u32 = row_offsets
427                .iter()
428                .map(|&offset| offset as u32)
429                .collect::<Vec<_>>();
430            LanceBuffer::reinterpret_vec(offsets_u32)
431        } else {
432            LanceBuffer::reinterpret_vec(row_offsets)
433        };
434
435        let data_block = VariableWidthBlock {
436            data: LanceBuffer::from(row_data),
437            bits_per_offset,
438            offsets: offsets_buffer,
439            num_values,
440            block_info: BlockInfo::new(),
441        };
442
443        Ok((
444            PerValueDataBlock::Variable(data_block),
445            EncodingUtils::packed_struct_variable(field_metadata),
446        ))
447    }
448}
449
450#[derive(Debug)]
451pub(crate) struct PackedStructFixedPerValueEncoder {
452    field_len: usize,
453}
454
455impl PackedStructFixedPerValueEncoder {
456    pub(crate) fn new(fields: Vec<Field>) -> Self {
457        Self {
458            field_len: fields.len(),
459        }
460    }
461}
462
463impl PerValueCompressor for PackedStructFixedPerValueEncoder {
464    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
465        let struct_block = check_struct_validity(data, self.field_len)?;
466
467        if struct_block.has_variable_width_child() {
468            return Err(Error::invalid_input(
469                "Packed struct fixed encoding requires all children to be fixed-width",
470            ));
471        }
472
473        let num_values = struct_block.children[0].num_values();
474        // Fixed length - supporting only flat encoding
475        let compressor = Box::new(ValueEncoder::default()) as Box<dyn PerValueCompressor>;
476        let mut field_data = Vec::with_capacity(self.field_len);
477        let mut field_bits_per_value = Vec::with_capacity(self.field_len);
478        let mut bits_per_row: u64 = 0;
479
480        for child_block in struct_block.children.into_iter() {
481            let (compressed, ..) = compressor.compress(child_block)?;
482            match compressed {
483                PerValueDataBlock::Fixed(block) => {
484                    bits_per_row = bits_per_row
485                        .checked_add(block.bits_per_value)
486                        .ok_or_else(|| Error::invalid_input("Packed struct row width overflow"))?;
487                    field_bits_per_value.push(block.bits_per_value);
488                    field_data.push(FixedPackedFieldData { block });
489                }
490                _ => {
491                    return Err(Error::invalid_input(
492                        "Packed struct fixed encoding requires all children to be fixed-width",
493                    ));
494                }
495            }
496        }
497
498        // Children are validated byte-aligned in `append_row_bytes`, so the row width
499        // is an exact number of bytes.
500        let bytes_per_row = (bits_per_row / 8) as usize;
501        let mut row_data: Vec<u8> =
502            Vec::with_capacity(bytes_per_row.saturating_mul(num_values as usize));
503        for row in 0..num_values as usize {
504            for field in &field_data {
505                field.append_row_bytes(row, &mut row_data)?;
506            }
507            debug_assert_eq!(row_data.len(), bytes_per_row * (row + 1));
508        }
509
510        let data_block = FixedWidthDataBlock {
511            data: LanceBuffer::from(row_data),
512            bits_per_value: bits_per_row,
513            num_values,
514            block_info: BlockInfo::new(),
515        };
516
517        Ok((
518            PerValueDataBlock::Fixed(data_block),
519            EncodingUtils::packed_struct(
520                EncodingUtils::flat(bits_per_row, None),
521                field_bits_per_value,
522            ),
523        ))
524    }
525}
526
527#[derive(Debug)]
528pub(crate) enum VariablePackedStructFieldKind {
529    Fixed {
530        bits_per_value: u64,
531        decompressor: Arc<dyn FixedPerValueDecompressor>,
532    },
533    Variable {
534        bits_per_length: u64,
535        decompressor: Arc<dyn VariablePerValueDecompressor>,
536    },
537}
538
539#[derive(Debug)]
540pub(crate) struct VariablePackedStructFieldDecoder {
541    pub(crate) kind: VariablePackedStructFieldKind,
542}
543
544#[derive(Debug)]
545pub struct PackedStructVariablePerValueDecompressor {
546    fields: Vec<VariablePackedStructFieldDecoder>,
547}
548
549impl PackedStructVariablePerValueDecompressor {
550    pub(crate) fn new(fields: Vec<VariablePackedStructFieldDecoder>) -> Self {
551        Self { fields }
552    }
553}
554
555#[derive(Debug)]
556struct FixedFieldAccumulator {
557    builder: DataBlockBuilder,
558    bits_per_value: u64,
559    empty_value: DataBlock,
560}
561
562impl FixedFieldAccumulator {
563    fn append_empty(&mut self) -> Result<()> {
564        self.builder.append(&self.empty_value, 0..1)
565    }
566
567    fn new(bits_per_value: u64, num_values: u64) -> Result<Self> {
568        if !bits_per_value.is_multiple_of(8) {
569            return Err(Error::invalid_input(
570                "Packed struct fixed child must be byte-aligned",
571            ));
572        }
573
574        let bytes_per_value = bits_per_value.checked_div(8).ok_or_else(|| {
575            Error::invalid_input("Invalid bits per value for packed struct field")
576        })?;
577
578        let estimate = bytes_per_value
579            .checked_mul(num_values)
580            .ok_or_else(|| Error::invalid_input("Packed struct fixed child allocation overflow"))?;
581
582        let empty_value = DataBlock::FixedWidth(FixedWidthDataBlock {
583            data: LanceBuffer::from(vec![0_u8; bytes_per_value as usize]),
584            bits_per_value,
585            num_values: 1,
586            block_info: BlockInfo::new(),
587        });
588
589        Ok(Self {
590            builder: DataBlockBuilder::with_capacity_estimate(estimate),
591            bits_per_value,
592            empty_value,
593        })
594    }
595
596    fn finish(self) -> Result<FixedWidthDataBlock> {
597        let DataBlock::FixedWidth(block) = self.builder.finish() else {
598            return Err(Error::invalid_input(
599                "Expected fixed-width datablock from builder",
600            ));
601        };
602        Ok(block)
603    }
604}
605
606enum FieldAccumulator {
607    Fixed(FixedFieldAccumulator),
608    Variable32 {
609        builder: DataBlockBuilder,
610        empty_value: DataBlock,
611    },
612    Variable64 {
613        builder: DataBlockBuilder,
614        empty_value: DataBlock,
615    },
616}
617
618impl FieldAccumulator {
619    // In full-zip variable packed decoding, rep/def may produce a visible row
620    // with an empty payload (e.g. null/invalid item). We still need to append
621    // one placeholder per child so child row counts remain aligned.
622    fn append_empty(&mut self) -> Result<()> {
623        match self {
624            Self::Fixed(fixed_field_accumulator) => fixed_field_accumulator.append_empty(),
625            Self::Variable32 {
626                builder,
627                empty_value,
628            } => builder.append(empty_value, 0..1),
629            Self::Variable64 {
630                builder,
631                empty_value,
632            } => builder.append(empty_value, 0..1),
633        }
634    }
635}
636
637impl VariablePerValueDecompressor for PackedStructVariablePerValueDecompressor {
638    fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
639        let num_values = data.num_values;
640        let offsets_u64 = match data.bits_per_offset {
641            32 => data
642                .offsets
643                .borrow_to_typed_slice::<u32>()
644                .iter()
645                .map(|v| *v as u64)
646                .collect::<Vec<_>>(),
647            64 => data
648                .offsets
649                .borrow_to_typed_slice::<u64>()
650                .as_ref()
651                .to_vec(),
652            _ => {
653                return Err(Error::invalid_input(
654                    "Packed struct row offsets must be 32 or 64 bits",
655                ));
656            }
657        };
658
659        if offsets_u64.len() != num_values as usize + 1 {
660            return Err(Error::invalid_input(
661                "Packed struct row offsets length mismatch",
662            ));
663        }
664
665        let mut accumulators = Vec::with_capacity(self.fields.len());
666        for field in &self.fields {
667            match &field.kind {
668                VariablePackedStructFieldKind::Fixed { bits_per_value, .. } => {
669                    let accumulator = FixedFieldAccumulator::new(*bits_per_value, num_values)?;
670                    accumulators.push(FieldAccumulator::Fixed(accumulator));
671                }
672                VariablePackedStructFieldKind::Variable {
673                    bits_per_length, ..
674                } => match bits_per_length {
675                    32 => accumulators.push(FieldAccumulator::Variable32 {
676                        builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64),
677                        empty_value: DataBlock::VariableWidth(VariableWidthBlock {
678                            data: LanceBuffer::empty(),
679                            bits_per_offset: 32,
680                            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 0_u32]),
681                            num_values: 1,
682                            block_info: BlockInfo::new(),
683                        }),
684                    }),
685                    64 => accumulators.push(FieldAccumulator::Variable64 {
686                        builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64),
687                        empty_value: DataBlock::VariableWidth(VariableWidthBlock {
688                            data: LanceBuffer::empty(),
689                            bits_per_offset: 64,
690                            offsets: LanceBuffer::reinterpret_vec(vec![0_u64, 0_u64]),
691                            num_values: 1,
692                            block_info: BlockInfo::new(),
693                        }),
694                    }),
695                    _ => {
696                        return Err(Error::invalid_input(
697                            "Packed struct variable child must use 32 or 64-bit length prefixes",
698                        ));
699                    }
700                },
701            }
702        }
703
704        for row_idx in 0..num_values as usize {
705            let row_start = offsets_u64[row_idx] as usize;
706            let row_end = offsets_u64[row_idx + 1] as usize;
707            if row_end > data.data.len() || row_start > row_end {
708                return Err(Error::invalid_input(
709                    "Packed struct row bounds exceed buffer",
710                ));
711            }
712            if row_start == row_end {
713                for accumulator in accumulators.iter_mut() {
714                    accumulator.append_empty()?;
715                }
716                continue;
717            }
718            let mut cursor = row_start;
719            for (field, accumulator) in self.fields.iter().zip(accumulators.iter_mut()) {
720                match (&field.kind, accumulator) {
721                    (
722                        VariablePackedStructFieldKind::Fixed { bits_per_value, .. },
723                        FieldAccumulator::Fixed(fixed_accumulator),
724                    ) => {
725                        debug_assert_eq!(*bits_per_value, fixed_accumulator.bits_per_value);
726                        let bytes_per_value = (*bits_per_value / 8) as usize;
727                        let end = cursor + bytes_per_value;
728                        if end > row_end {
729                            return Err(Error::invalid_input(
730                                "Packed struct fixed child exceeds row bounds",
731                            ));
732                        }
733                        let value_block = DataBlock::FixedWidth(FixedWidthDataBlock {
734                            data: LanceBuffer::from(data.data[cursor..end].to_vec()),
735                            bits_per_value: *bits_per_value,
736                            num_values: 1,
737                            block_info: BlockInfo::new(),
738                        });
739                        fixed_accumulator.builder.append(&value_block, 0..1)?;
740                        cursor = end;
741                    }
742                    (
743                        VariablePackedStructFieldKind::Variable {
744                            bits_per_length, ..
745                        },
746                        FieldAccumulator::Variable32 { builder, .. },
747                    ) => {
748                        if *bits_per_length != 32 {
749                            return Err(Error::invalid_input(
750                                "Packed struct length prefix size mismatch",
751                            ));
752                        }
753                        let end = cursor + std::mem::size_of::<u32>();
754                        if end > row_end {
755                            return Err(Error::invalid_input(
756                                "Packed struct variable child length prefix out of bounds",
757                            ));
758                        }
759                        let len = u32::from_le_bytes(
760                            data.data[cursor..end]
761                                .try_into()
762                                .expect("slice has exact length"),
763                        ) as usize;
764                        cursor = end;
765                        let value_end = cursor + len;
766                        if value_end > row_end {
767                            return Err(Error::invalid_input(
768                                "Packed struct variable child exceeds row bounds",
769                            ));
770                        }
771                        let value_block = DataBlock::VariableWidth(VariableWidthBlock {
772                            data: LanceBuffer::from(data.data[cursor..value_end].to_vec()),
773                            bits_per_offset: 32,
774                            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, len as u32]),
775                            num_values: 1,
776                            block_info: BlockInfo::new(),
777                        });
778                        builder.append(&value_block, 0..1)?;
779                        cursor = value_end;
780                    }
781                    (
782                        VariablePackedStructFieldKind::Variable {
783                            bits_per_length, ..
784                        },
785                        FieldAccumulator::Variable64 { builder, .. },
786                    ) => {
787                        if *bits_per_length != 64 {
788                            return Err(Error::invalid_input(
789                                "Packed struct length prefix size mismatch",
790                            ));
791                        }
792                        let end = cursor + std::mem::size_of::<u64>();
793                        if end > row_end {
794                            return Err(Error::invalid_input(
795                                "Packed struct variable child length prefix out of bounds",
796                            ));
797                        }
798                        let len = u64::from_le_bytes(
799                            data.data[cursor..end]
800                                .try_into()
801                                .expect("slice has exact length"),
802                        ) as usize;
803                        cursor = end;
804                        let value_end = cursor + len;
805                        if value_end > row_end {
806                            return Err(Error::invalid_input(
807                                "Packed struct variable child exceeds row bounds",
808                            ));
809                        }
810                        let value_block = DataBlock::VariableWidth(VariableWidthBlock {
811                            data: LanceBuffer::from(data.data[cursor..value_end].to_vec()),
812                            bits_per_offset: 64,
813                            offsets: LanceBuffer::reinterpret_vec(vec![0_u64, len as u64]),
814                            num_values: 1,
815                            block_info: BlockInfo::new(),
816                        });
817                        builder.append(&value_block, 0..1)?;
818                        cursor = value_end;
819                    }
820                    _ => {
821                        return Err(Error::invalid_input(
822                            "Packed struct accumulator kind mismatch",
823                        ));
824                    }
825                }
826            }
827            if cursor != row_end {
828                return Err(Error::invalid_input(
829                    "Packed struct row parsing did not consume full row",
830                ));
831            }
832        }
833
834        let mut children = Vec::with_capacity(self.fields.len());
835        for (field, accumulator) in self.fields.iter().zip(accumulators) {
836            match (field, accumulator) {
837                (
838                    VariablePackedStructFieldDecoder {
839                        kind: VariablePackedStructFieldKind::Fixed { decompressor, .. },
840                    },
841                    FieldAccumulator::Fixed(fixed_accumulator),
842                ) => {
843                    let finished_accumulator = fixed_accumulator.finish()?;
844                    let decoded = decompressor.decompress(finished_accumulator, num_values)?;
845                    children.push(decoded);
846                }
847                (
848                    VariablePackedStructFieldDecoder {
849                        kind:
850                            VariablePackedStructFieldKind::Variable {
851                                bits_per_length,
852                                decompressor,
853                            },
854                    },
855                    FieldAccumulator::Variable32 { builder, .. },
856                ) => {
857                    let DataBlock::VariableWidth(mut block) = builder.finish() else {
858                        panic!("Expected variable-width datablock from builder");
859                    };
860                    debug_assert_eq!(block.bits_per_offset, 32);
861                    block.bits_per_offset = (*bits_per_length) as u8;
862                    let decoded = decompressor.decompress(block)?;
863                    children.push(decoded);
864                }
865                (
866                    VariablePackedStructFieldDecoder {
867                        kind:
868                            VariablePackedStructFieldKind::Variable {
869                                bits_per_length,
870                                decompressor,
871                            },
872                    },
873                    FieldAccumulator::Variable64 { builder, .. },
874                ) => {
875                    let DataBlock::VariableWidth(mut block) = builder.finish() else {
876                        panic!("Expected variable-width datablock from builder");
877                    };
878                    debug_assert_eq!(block.bits_per_offset, 64);
879                    block.bits_per_offset = (*bits_per_length) as u8;
880                    let decoded = decompressor.decompress(block)?;
881                    children.push(decoded);
882                }
883                _ => {
884                    return Err(Error::invalid_input(
885                        "Packed struct accumulator mismatch during finalize",
886                    ));
887                }
888            }
889        }
890
891        Ok(DataBlock::Struct(StructDataBlock {
892            children,
893            block_info: BlockInfo::new(),
894            validity: None,
895        }))
896    }
897}
898
899#[derive(Debug)]
900struct PackedStructFixedFieldDecoder {
901    bits_per_value: u64,
902    decompressor: Box<dyn FixedPerValueDecompressor>,
903}
904
905#[derive(Debug)]
906pub(crate) struct PackedStructFixedPerValueDecompressor {
907    decoders: Vec<PackedStructFixedFieldDecoder>,
908}
909
910impl PackedStructFixedPerValueDecompressor {
911    pub(crate) fn new(description: &PackedStruct) -> Result<Self> {
912        let compression = description
913            .values
914            .as_ref()
915            .ok_or_else(|| Error::invalid_input("PackedStruct missing values encoding"))?
916            .compression
917            .as_ref()
918            .ok_or_else(|| {
919                Error::invalid_input("PackedStruct values missing compression encoding")
920            })?;
921
922        // The encoder always flat-encodes each child, so that is the only layout we can decode.
923        if !matches!(compression, Compression::Flat(..)) {
924            return Err(Error::invalid_input(
925                "PackedStruct fixed encoding currently requires flat compression",
926            ));
927        }
928
929        let decoders = description
930            .bits_per_value
931            .iter()
932            .map(|&bits_per_value| {
933                let flat = crate::format::metadata::Flat {
934                    bits_per_value,
935                    data: None,
936                };
937                PackedStructFixedFieldDecoder {
938                    bits_per_value,
939                    decompressor: Box::new(ValueDecompressor::from_flat(&flat)),
940                }
941            })
942            .collect();
943        Ok(Self { decoders })
944    }
945}
946
947impl FixedPerValueDecompressor for PackedStructFixedPerValueDecompressor {
948    fn decompress(&self, data: FixedWidthDataBlock, num_values: u64) -> Result<DataBlock> {
949        if !data.bits_per_value.is_multiple_of(8) {
950            return Err(Error::invalid_input(
951                "Packed struct fixed encoding requires byte-aligned children",
952            ));
953        }
954        let bytes_per_row = (data.bits_per_value / 8) as usize;
955
956        // Byte offset of each child within a packed row (a running prefix sum of the
957        // child widths). The final offset must equal the packed row width.
958        let mut child_bytes = Vec::with_capacity(self.decoders.len());
959        for decoder in &self.decoders {
960            if !decoder.bits_per_value.is_multiple_of(8) {
961                return Err(Error::invalid_input(
962                    "Packed struct fixed child must be byte-aligned",
963                ));
964            }
965            child_bytes.push((decoder.bits_per_value / 8) as usize);
966        }
967        if child_bytes.iter().sum::<usize>() != bytes_per_row {
968            return Err(Error::invalid_input(
969                "Packed struct child widths do not sum to the packed row width",
970            ));
971        }
972        if bytes_per_row.saturating_mul(num_values as usize) > data.data.len() {
973            return Err(Error::invalid_input(
974                "Packed struct row bounds exceed buffer",
975            ));
976        }
977
978        // Un-zip the row-major buffer one child at a time by gathering that child's
979        // slice out of every row, then hand the column to the child decompressor.
980        let bytes = data.data.as_ref();
981        let mut children = Vec::with_capacity(self.decoders.len());
982        let mut field_offset = 0;
983        for (decoder, &field_bytes) in self.decoders.iter().zip(child_bytes.iter()) {
984            let mut child_buf = Vec::with_capacity(field_bytes * num_values as usize);
985            for row_idx in 0..num_values as usize {
986                let start = row_idx * bytes_per_row + field_offset;
987                child_buf.extend_from_slice(&bytes[start..start + field_bytes]);
988            }
989            let child_block = FixedWidthDataBlock {
990                data: LanceBuffer::from(child_buf),
991                bits_per_value: decoder.bits_per_value,
992                num_values,
993                block_info: BlockInfo::new(),
994            };
995            children.push(decoder.decompressor.decompress(child_block, num_values)?);
996            field_offset += field_bytes;
997        }
998
999        Ok(DataBlock::Struct(StructDataBlock {
1000            children,
1001            block_info: BlockInfo::new(),
1002            validity: None,
1003        }))
1004    }
1005
1006    fn bits_per_value(&self) -> u64 {
1007        self.decoders
1008            .iter()
1009            .map(|decoder| decoder.bits_per_value)
1010            .sum()
1011    }
1012}
1013
1014#[cfg(test)]
1015mod tests {
1016    use super::*;
1017    use crate::{
1018        compression::CompressionStrategy,
1019        compression::DefaultDecompressionStrategy,
1020        compression_config::CompressionParams,
1021        constants::{
1022            PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY,
1023        },
1024        statistics::ComputeStat,
1025        testing::{
1026            TestCases, TestEncoding, check_round_trip_encoding_of_data, test_compression_strategy,
1027        },
1028    };
1029    use arrow_array::{
1030        Array, ArrayRef, BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray,
1031        StructArray, UInt32Array,
1032    };
1033    use arrow_schema::{DataType, Field as ArrowField, Fields};
1034    use std::collections::HashMap;
1035    use std::sync::Arc;
1036
1037    fn fixed_block_from_array(array: Int64Array) -> FixedWidthDataBlock {
1038        let num_values = array.len() as u64;
1039        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
1040        match block {
1041            DataBlock::FixedWidth(block) => block,
1042            _ => panic!("Expected fixed-width data block"),
1043        }
1044    }
1045
1046    fn fixed_i32_block_from_array(array: Int32Array) -> FixedWidthDataBlock {
1047        let num_values = array.len() as u64;
1048        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
1049        match block {
1050            DataBlock::FixedWidth(block) => block,
1051            _ => panic!("Expected fixed-width data block"),
1052        }
1053    }
1054
1055    fn variable_block_from_string_array(array: StringArray) -> VariableWidthBlock {
1056        let num_values = array.len() as u64;
1057        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
1058        match block {
1059            DataBlock::VariableWidth(block) => block,
1060            _ => panic!("Expected variable-width block"),
1061        }
1062    }
1063
1064    fn variable_block_from_large_string_array(array: LargeStringArray) -> VariableWidthBlock {
1065        let num_values = array.len() as u64;
1066        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
1067        match block {
1068            DataBlock::VariableWidth(block) => block,
1069            _ => panic!("Expected variable-width block"),
1070        }
1071    }
1072
1073    fn variable_block_from_binary_array(array: BinaryArray) -> VariableWidthBlock {
1074        let num_values = array.len() as u64;
1075        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
1076        match block {
1077            DataBlock::VariableWidth(block) => block,
1078            _ => panic!("Expected variable-width block"),
1079        }
1080    }
1081
1082    #[test]
1083    fn variable_packed_struct_round_trip() -> Result<()> {
1084        let arrow_fields: Fields = vec![
1085            ArrowField::new("id", DataType::UInt32, false),
1086            ArrowField::new("name", DataType::Utf8, true),
1087        ]
1088        .into();
1089        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1090        let struct_field = Field::try_from(&arrow_struct)?;
1091
1092        let ids = vec![1_u32, 2, 42];
1093        let id_bytes = ids
1094            .iter()
1095            .flat_map(|value| value.to_le_bytes())
1096            .collect::<Vec<_>>();
1097        let mut id_block = FixedWidthDataBlock {
1098            data: LanceBuffer::reinterpret_vec(ids),
1099            bits_per_value: 32,
1100            num_values: 3,
1101            block_info: BlockInfo::new(),
1102        };
1103        id_block.compute_stat();
1104        let id_block = DataBlock::FixedWidth(id_block);
1105
1106        let name_offsets = vec![0_i32, 1, 4, 4];
1107        let name_bytes = b"abcz".to_vec();
1108        let mut name_block = VariableWidthBlock {
1109            data: LanceBuffer::from(name_bytes.clone()),
1110            bits_per_offset: 32,
1111            offsets: LanceBuffer::reinterpret_vec(name_offsets.clone()),
1112            num_values: 3,
1113            block_info: BlockInfo::new(),
1114        };
1115        name_block.compute_stat();
1116        let name_block = DataBlock::VariableWidth(name_block);
1117
1118        let struct_block = StructDataBlock {
1119            children: vec![id_block, name_block],
1120            block_info: BlockInfo::new(),
1121            validity: None,
1122        };
1123
1124        let data_block = DataBlock::Struct(struct_block);
1125
1126        let compression_strategy =
1127            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
1128        let compressor = crate::compression::CompressionStrategy::create_per_value(
1129            compression_strategy.as_ref(),
1130            &struct_field,
1131            &data_block,
1132        )?;
1133        let (compressed, encoding) = compressor.compress(data_block)?;
1134
1135        let PerValueDataBlock::Variable(zipped) = compressed else {
1136            panic!("expected variable-width packed struct output");
1137        };
1138
1139        let decompression_strategy = DefaultDecompressionStrategy::default();
1140        let decompressor =
1141            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1142                &decompression_strategy,
1143                &encoding,
1144            )?;
1145        let decoded = decompressor.decompress(zipped)?;
1146
1147        let DataBlock::Struct(decoded_struct) = decoded else {
1148            panic!("expected struct datablock after decode");
1149        };
1150
1151        let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap();
1152        assert_eq!(decoded_id.bits_per_value, 32);
1153        assert_eq!(decoded_id.data.as_ref(), id_bytes.as_slice());
1154
1155        let decoded_name = decoded_struct.children[1].as_variable_width_ref().unwrap();
1156        assert_eq!(decoded_name.bits_per_offset, 32);
1157        let decoded_offsets = decoded_name.offsets.borrow_to_typed_slice::<i32>();
1158        assert_eq!(decoded_offsets.as_ref(), name_offsets.as_slice());
1159        assert_eq!(decoded_name.data.as_ref(), name_bytes.as_slice());
1160
1161        Ok(())
1162    }
1163
1164    #[test]
1165    fn variable_packed_struct_large_utf8_round_trip() -> Result<()> {
1166        let arrow_fields: Fields = vec![
1167            ArrowField::new("value", DataType::Int64, false),
1168            ArrowField::new("text", DataType::LargeUtf8, false),
1169        ]
1170        .into();
1171        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1172        let struct_field = Field::try_from(&arrow_struct)?;
1173
1174        let id_block = fixed_block_from_array(Int64Array::from(vec![10, 20, 30, 40]));
1175        let payload_array = LargeStringArray::from(vec![
1176            "alpha",
1177            "a considerably longer payload for testing",
1178            "mid",
1179            "z",
1180        ]);
1181        let payload_block = variable_block_from_large_string_array(payload_array);
1182
1183        let struct_block = StructDataBlock {
1184            children: vec![
1185                DataBlock::FixedWidth(id_block.clone()),
1186                DataBlock::VariableWidth(payload_block.clone()),
1187            ],
1188            block_info: BlockInfo::new(),
1189            validity: None,
1190        };
1191
1192        let data_block = DataBlock::Struct(struct_block);
1193
1194        let compression_strategy =
1195            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
1196        let compressor = crate::compression::CompressionStrategy::create_per_value(
1197            compression_strategy.as_ref(),
1198            &struct_field,
1199            &data_block,
1200        )?;
1201        let (compressed, encoding) = compressor.compress(data_block)?;
1202
1203        let PerValueDataBlock::Variable(zipped) = compressed else {
1204            panic!("expected variable-width packed struct output");
1205        };
1206
1207        let decompression_strategy = DefaultDecompressionStrategy::default();
1208        let decompressor =
1209            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1210                &decompression_strategy,
1211                &encoding,
1212            )?;
1213        let decoded = decompressor.decompress(zipped)?;
1214
1215        let DataBlock::Struct(decoded_struct) = decoded else {
1216            panic!("expected struct datablock after decode");
1217        };
1218
1219        let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap();
1220        assert_eq!(decoded_id.bits_per_value, 64);
1221        assert_eq!(decoded_id.data.as_ref(), id_block.data.as_ref());
1222
1223        let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap();
1224        assert_eq!(decoded_payload.bits_per_offset, 64);
1225        assert_eq!(
1226            decoded_payload
1227                .offsets
1228                .borrow_to_typed_slice::<i64>()
1229                .as_ref(),
1230            payload_block
1231                .offsets
1232                .borrow_to_typed_slice::<i64>()
1233                .as_ref()
1234        );
1235        assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref());
1236
1237        Ok(())
1238    }
1239
1240    #[tokio::test]
1241    async fn variable_packed_struct_utf8_round_trip() {
1242        // schema: Struct<id: UInt32, uri: Utf8, long_text: LargeUtf8>
1243        let fields = Fields::from(vec![
1244            Arc::new(ArrowField::new("id", DataType::UInt32, false)),
1245            Arc::new(ArrowField::new("uri", DataType::Utf8, false)),
1246            Arc::new(ArrowField::new("long_text", DataType::LargeUtf8, false)),
1247        ]);
1248
1249        // mark struct as packed
1250        let mut meta = HashMap::new();
1251        meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
1252
1253        let array = Arc::new(StructArray::from(vec![
1254            (
1255                fields[0].clone(),
1256                Arc::new(UInt32Array::from(vec![1, 2, 3])) as ArrayRef,
1257            ),
1258            (
1259                fields[1].clone(),
1260                Arc::new(StringArray::from(vec![
1261                    Some("a"),
1262                    Some("b"),
1263                    Some("/tmp/x"),
1264                ])) as ArrayRef,
1265            ),
1266            (
1267                fields[2].clone(),
1268                Arc::new(LargeStringArray::from(vec![
1269                    Some("alpha"),
1270                    Some("a considerably longer payload for testing"),
1271                    Some("mid"),
1272                ])) as ArrayRef,
1273            ),
1274        ]));
1275
1276        let test_cases = TestCases::default()
1277            .with_u32_structural_encodings()
1278            .with_expected_encoding("variable_packed_struct");
1279
1280        check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await;
1281    }
1282
1283    #[test]
1284    fn variable_packed_struct_multi_variable_round_trip() -> Result<()> {
1285        let arrow_fields: Fields = vec![
1286            ArrowField::new("category", DataType::Utf8, false),
1287            ArrowField::new("payload", DataType::Binary, false),
1288            ArrowField::new("count", DataType::Int32, false),
1289        ]
1290        .into();
1291        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1292        let struct_field = Field::try_from(&arrow_struct)?;
1293
1294        let category_array = StringArray::from(vec!["red", "blue", "green", "red"]);
1295        let category_block = variable_block_from_string_array(category_array);
1296        let payload_values: Vec<Vec<u8>> =
1297            vec![vec![0x01, 0x02], vec![], vec![0x05, 0x06, 0x07], vec![0xff]];
1298        let payload_array =
1299            BinaryArray::from_iter_values(payload_values.iter().map(|v| v.as_slice()));
1300        let payload_block = variable_block_from_binary_array(payload_array);
1301        let count_block = fixed_i32_block_from_array(Int32Array::from(vec![1, 2, 3, 4]));
1302
1303        let struct_block = StructDataBlock {
1304            children: vec![
1305                DataBlock::VariableWidth(category_block.clone()),
1306                DataBlock::VariableWidth(payload_block.clone()),
1307                DataBlock::FixedWidth(count_block.clone()),
1308            ],
1309            block_info: BlockInfo::new(),
1310            validity: None,
1311        };
1312
1313        let data_block = DataBlock::Struct(struct_block);
1314
1315        let compression_strategy =
1316            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
1317        let compressor = crate::compression::CompressionStrategy::create_per_value(
1318            compression_strategy.as_ref(),
1319            &struct_field,
1320            &data_block,
1321        )?;
1322        let (compressed, encoding) = compressor.compress(data_block)?;
1323
1324        let PerValueDataBlock::Variable(zipped) = compressed else {
1325            panic!("expected variable-width packed struct output");
1326        };
1327
1328        let decompression_strategy = DefaultDecompressionStrategy::default();
1329        let decompressor =
1330            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1331                &decompression_strategy,
1332                &encoding,
1333            )?;
1334        let decoded = decompressor.decompress(zipped)?;
1335
1336        let DataBlock::Struct(decoded_struct) = decoded else {
1337            panic!("expected struct datablock after decode");
1338        };
1339
1340        let decoded_category = decoded_struct.children[0].as_variable_width_ref().unwrap();
1341        assert_eq!(decoded_category.bits_per_offset, 32);
1342        assert_eq!(
1343            decoded_category
1344                .offsets
1345                .borrow_to_typed_slice::<i32>()
1346                .as_ref(),
1347            category_block
1348                .offsets
1349                .borrow_to_typed_slice::<i32>()
1350                .as_ref()
1351        );
1352        assert_eq!(decoded_category.data.as_ref(), category_block.data.as_ref());
1353
1354        let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap();
1355        assert_eq!(decoded_payload.bits_per_offset, 32);
1356        assert_eq!(
1357            decoded_payload
1358                .offsets
1359                .borrow_to_typed_slice::<i32>()
1360                .as_ref(),
1361            payload_block
1362                .offsets
1363                .borrow_to_typed_slice::<i32>()
1364                .as_ref()
1365        );
1366        assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref());
1367
1368        let decoded_count = decoded_struct.children[2].as_fixed_width_ref().unwrap();
1369        assert_eq!(decoded_count.bits_per_value, 32);
1370        assert_eq!(decoded_count.data.as_ref(), count_block.data.as_ref());
1371
1372        Ok(())
1373    }
1374
1375    #[test]
1376    fn variable_packed_struct_requires_v22() {
1377        let arrow_fields: Fields = vec![
1378            ArrowField::new("value", DataType::Int64, false),
1379            ArrowField::new("text", DataType::Utf8, false),
1380        ]
1381        .into();
1382        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1383        let struct_field = Field::try_from(&arrow_struct).unwrap();
1384
1385        let value_block = fixed_block_from_array(Int64Array::from(vec![1, 2, 3]));
1386        let text_block =
1387            variable_block_from_string_array(StringArray::from(vec!["a", "bb", "ccc"]));
1388
1389        let struct_block = StructDataBlock {
1390            children: vec![
1391                DataBlock::FixedWidth(value_block),
1392                DataBlock::VariableWidth(text_block),
1393            ],
1394            block_info: BlockInfo::new(),
1395            validity: None,
1396        };
1397
1398        let compression_strategy =
1399            test_compression_strategy(TestEncoding::StructuralU16, CompressionParams::default());
1400        let result =
1401            compression_strategy.create_per_value(&struct_field, &DataBlock::Struct(struct_block));
1402
1403        assert!(matches!(result, Err(Error::NotSupported { .. })));
1404    }
1405
1406    #[test]
1407    fn variable_packed_struct_decompress_empty_row() -> Result<()> {
1408        let strategy = DefaultDecompressionStrategy::default();
1409        let fixed_decompressor = Arc::from(
1410            crate::compression::DecompressionStrategy::create_fixed_per_value_decompressor(
1411                &strategy,
1412                &EncodingUtils::flat(32, None),
1413            )?,
1414        );
1415        let variable_decompressor = Arc::from(
1416            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1417                &strategy,
1418                &EncodingUtils::variable(EncodingUtils::flat(32, None), None),
1419            )?,
1420        );
1421
1422        let decompressor = PackedStructVariablePerValueDecompressor::new(vec![
1423            VariablePackedStructFieldDecoder {
1424                kind: VariablePackedStructFieldKind::Fixed {
1425                    bits_per_value: 32,
1426                    decompressor: fixed_decompressor,
1427                },
1428            },
1429            VariablePackedStructFieldDecoder {
1430                kind: VariablePackedStructFieldKind::Variable {
1431                    bits_per_length: 32,
1432                    decompressor: variable_decompressor,
1433                },
1434            },
1435        ]);
1436
1437        let mut row_data = Vec::new();
1438        row_data.extend_from_slice(&1_u32.to_le_bytes());
1439        row_data.extend_from_slice(&1_u32.to_le_bytes());
1440        row_data.extend_from_slice(b"a");
1441        row_data.extend_from_slice(&2_u32.to_le_bytes());
1442        row_data.extend_from_slice(&0_u32.to_le_bytes());
1443
1444        let input = VariableWidthBlock {
1445            data: LanceBuffer::from(row_data),
1446            bits_per_offset: 32,
1447            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 9_u32, 9_u32, 17_u32]),
1448            num_values: 3,
1449            block_info: BlockInfo::new(),
1450        };
1451
1452        let decoded = decompressor.decompress(input)?;
1453        let DataBlock::Struct(decoded_struct) = decoded else {
1454            panic!("expected struct output");
1455        };
1456
1457        let fixed = decoded_struct.children[0].as_fixed_width_ref().unwrap();
1458        assert_eq!(fixed.bits_per_value, 32);
1459        assert_eq!(
1460            fixed.data.borrow_to_typed_slice::<u32>().as_ref(),
1461            &[1, 0, 2]
1462        );
1463
1464        let variable = decoded_struct.children[1].as_variable_width_ref().unwrap();
1465        assert_eq!(variable.bits_per_offset, 32);
1466        assert_eq!(
1467            variable.offsets.borrow_to_typed_slice::<u32>().as_ref(),
1468            &[0_u32, 1_u32, 1_u32, 1_u32]
1469        );
1470        assert_eq!(variable.data.as_ref(), b"a");
1471
1472        Ok(())
1473    }
1474
1475    #[test]
1476    fn fixed_packed_struct_round_trip() -> Result<()> {
1477        let arrow_fields: Fields = vec![
1478            ArrowField::new("id", DataType::Int32, false),
1479            ArrowField::new("value", DataType::Int64, false),
1480        ]
1481        .into();
1482        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1483        let struct_field = Field::try_from(&arrow_struct)?;
1484
1485        let id_block = fixed_i32_block_from_array(Int32Array::from(vec![1, 2, 3, 4]));
1486        let value_block = fixed_block_from_array(Int64Array::from(vec![10, 20, 30, 40]));
1487
1488        let struct_block = StructDataBlock {
1489            children: vec![
1490                DataBlock::FixedWidth(id_block.clone()),
1491                DataBlock::FixedWidth(value_block.clone()),
1492            ],
1493            block_info: BlockInfo::new(),
1494            validity: None,
1495        };
1496
1497        let data_block = DataBlock::Struct(struct_block);
1498
1499        let compression_strategy =
1500            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
1501        let compressor = CompressionStrategy::create_per_value(
1502            compression_strategy.as_ref(),
1503            &struct_field,
1504            &data_block,
1505        )?;
1506        let (compressed, encoding) = compressor.compress(data_block)?;
1507
1508        let PerValueDataBlock::Fixed(zipped) = compressed else {
1509            panic!("expected fixed-width packed struct output");
1510        };
1511
1512        let decompression_strategy = DefaultDecompressionStrategy::default();
1513        let decompressor =
1514            crate::compression::DecompressionStrategy::create_fixed_per_value_decompressor(
1515                &decompression_strategy,
1516                &encoding,
1517            )?;
1518        let decoded = decompressor.decompress(zipped, 4)?;
1519
1520        let DataBlock::Struct(decoded_struct) = decoded else {
1521            panic!("expected struct datablock after decode");
1522        };
1523
1524        let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap();
1525        assert_eq!(decoded_id.bits_per_value, 32);
1526        assert_eq!(decoded_id.data.as_ref(), id_block.data.as_ref());
1527
1528        let decoded_value = decoded_struct.children[1].as_fixed_width_ref().unwrap();
1529        assert_eq!(decoded_value.bits_per_value, 64);
1530        assert_eq!(decoded_value.data.as_ref(), value_block.data.as_ref());
1531
1532        Ok(())
1533    }
1534
1535    // End-to-end round trip through the file writer. Requesting full-zip on an
1536    // all-fixed-width struct routes to `PackedStructFixedPerValueEncoder` (mini-block
1537    // is only chosen for narrow structs), exercising the writer/reader wiring rather
1538    // than just the block-level compress/decompress above.
1539    #[tokio::test]
1540    async fn fixed_packed_struct_full_zip_round_trip() {
1541        let fields = Fields::from(vec![
1542            Arc::new(ArrowField::new("id", DataType::Int32, false)),
1543            Arc::new(ArrowField::new("value", DataType::Int64, false)),
1544        ]);
1545
1546        let mut meta = HashMap::new();
1547        meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
1548        meta.insert(
1549            STRUCTURAL_ENCODING_META_KEY.to_string(),
1550            STRUCTURAL_ENCODING_FULLZIP.to_string(),
1551        );
1552
1553        let array = Arc::new(StructArray::from(vec![
1554            (
1555                fields[0].clone(),
1556                Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef,
1557            ),
1558            (
1559                fields[1].clone(),
1560                Arc::new(Int64Array::from(vec![10, 20, 30, 40])) as ArrayRef,
1561            ),
1562        ]));
1563
1564        let test_cases = TestCases::default()
1565            .with_u32_structural_encodings()
1566            .with_expected_encoding("packed_struct");
1567
1568        check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await;
1569    }
1570
1571    #[test]
1572    fn fixed_packed_struct_rejects_variable_child() -> Result<()> {
1573        let arrow_fields: Fields = vec![
1574            ArrowField::new("id", DataType::Int32, false),
1575            ArrowField::new("name", DataType::Utf8, false),
1576        ]
1577        .into();
1578        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1579        let struct_field = Field::try_from(&arrow_struct)?;
1580
1581        let struct_block = DataBlock::Struct(StructDataBlock {
1582            children: vec![
1583                DataBlock::FixedWidth(fixed_i32_block_from_array(Int32Array::from(vec![1, 2]))),
1584                DataBlock::VariableWidth(variable_block_from_string_array(StringArray::from(
1585                    vec!["a", "bb"],
1586                ))),
1587            ],
1588            block_info: BlockInfo::new(),
1589            validity: None,
1590        });
1591
1592        let encoder = PackedStructFixedPerValueEncoder::new(struct_field.children);
1593        let err = encoder.compress(struct_block).unwrap_err();
1594        assert!(matches!(&err, Error::InvalidInput { .. }));
1595        assert!(
1596            err.to_string().contains("fixed-width"),
1597            "unexpected error: {err}"
1598        );
1599
1600        Ok(())
1601    }
1602}
lance_encoding/encodings/physical/packed.rs

lance_encoding/encodings/physical/
packed.rs
