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;
15
16use lance_core::{Error, Result, datatypes::Field};
17
18use crate::{
19    buffer::LanceBuffer,
20    compression::{
21        CompressionStrategy, FixedPerValueDecompressor, MiniBlockDecompressor,
22        VariablePerValueDecompressor,
23    },
24    data::{
25        BlockInfo, DataBlock, DataBlockBuilder, FixedWidthDataBlock, StructDataBlock,
26        VariableWidthBlock,
27    },
28    encodings::logical::primitive::{
29        fullzip::{PerValueCompressor, PerValueDataBlock},
30        miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor},
31    },
32    format::{
33        ProtobufUtils21,
34        pb21::{CompressiveEncoding, PackedStruct, compressive_encoding::Compression},
35    },
36    statistics::{GetStat, Stat},
37};
38
39use super::value::{ValueDecompressor, ValueEncoder};
40
41// Transforms a `StructDataBlock` into a row major `FixedWidthDataBlock`.
42// Only fields with fixed-width fields are supported for now, and the
43// assumption that all fields has `bits_per_value % 8 == 0` is made.
44fn struct_data_block_to_fixed_width_data_block(
45    struct_data_block: StructDataBlock,
46    bits_per_values: &[u64],
47) -> DataBlock {
48    let data_size = struct_data_block.expect_single_stat::<UInt64Type>(Stat::DataSize);
49    let mut output = Vec::with_capacity(data_size as usize);
50    let num_values = struct_data_block.children[0].num_values();
51
52    for i in 0..num_values as usize {
53        for (j, child) in struct_data_block.children.iter().enumerate() {
54            let bytes_per_value = (bits_per_values[j] / 8) as usize;
55            let this_data = child
56                .as_fixed_width_ref()
57                .unwrap()
58                .data
59                .slice_with_length(bytes_per_value * i, bytes_per_value);
60            output.extend_from_slice(&this_data);
61        }
62    }
63
64    DataBlock::FixedWidth(FixedWidthDataBlock {
65        bits_per_value: bits_per_values.iter().copied().sum(),
66        data: LanceBuffer::from(output),
67        num_values,
68        block_info: BlockInfo::default(),
69    })
70}
71
72#[derive(Debug, Default)]
73pub struct PackedStructFixedWidthMiniBlockEncoder {}
74
75impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder {
76    fn compress(
77        &self,
78        context: MiniBlockCompressionContext,
79        data: DataBlock,
80    ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> {
81        match data {
82            DataBlock::Struct(struct_data_block) => {
83                let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::<Vec<_>>();
84
85                // transform struct datablock to fixed-width data block.
86                let data_block = struct_data_block_to_fixed_width_data_block(struct_data_block, &bits_per_values);
87
88                // store and transformed fixed-width data block.
89                let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box<dyn MiniBlockCompressor>;
90                let (value_miniblock_compressed, value_array_encoding) =
91                value_miniblock_compressor.compress(context, data_block)?;
92
93                Ok((
94                    value_miniblock_compressed,
95                    ProtobufUtils21::packed_struct(value_array_encoding, bits_per_values),
96                ))
97            }
98            _ => Err(Error::invalid_input_source(format!(
99                "Cannot compress a data block of type {} with PackedStructFixedWidthBlockEncoder",
100                data.name()
101            )
102            .into())),
103        }
104    }
105}
106
107#[derive(Debug)]
108pub struct PackedStructFixedWidthMiniBlockDecompressor {
109    bits_per_values: Vec<u64>,
110    array_encoding: Box<dyn MiniBlockDecompressor>,
111}
112
113impl PackedStructFixedWidthMiniBlockDecompressor {
114    pub fn new(description: &PackedStruct) -> Self {
115        let array_encoding: Box<dyn MiniBlockDecompressor> = match description
116            .values
117            .as_ref()
118            .unwrap()
119            .compression
120            .as_ref()
121            .unwrap()
122        {
123            Compression::Flat(flat) => Box::new(ValueDecompressor::from_flat(flat)),
124            _ => panic!(
125                "Currently only `ArrayEncoding::Flat` is supported in packed struct encoding in Lance 2.1."
126            ),
127        };
128        Self {
129            bits_per_values: description.bits_per_value.clone(),
130            array_encoding,
131        }
132    }
133}
134
135impl MiniBlockDecompressor for PackedStructFixedWidthMiniBlockDecompressor {
136    fn decompress(&self, data: Vec<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
137        assert_eq!(data.len(), 1);
138        let encoded_data_block = self.array_encoding.decompress(data, num_values)?;
139        let DataBlock::FixedWidth(encoded_data_block) = encoded_data_block else {
140            panic!("ValueDecompressor should output FixedWidth DataBlock")
141        };
142
143        let bytes_per_values = self
144            .bits_per_values
145            .iter()
146            .map(|bits_per_value| *bits_per_value as usize / 8)
147            .collect::<Vec<_>>();
148
149        assert!(encoded_data_block.bits_per_value % 8 == 0);
150        let encoded_bytes_per_row = (encoded_data_block.bits_per_value / 8) as usize;
151
152        // use a prefix_sum vector as a helper to reconstruct to `StructDataBlock`.
153        let mut prefix_sum = vec![0; self.bits_per_values.len()];
154        for i in 0..(self.bits_per_values.len() - 1) {
155            prefix_sum[i + 1] = prefix_sum[i] + bytes_per_values[i];
156        }
157
158        let mut children_data_block = vec![];
159        for i in 0..self.bits_per_values.len() {
160            let child_buf_size = bytes_per_values[i] * num_values as usize;
161            let mut child_buf: Vec<u8> = Vec::with_capacity(child_buf_size);
162
163            for j in 0..num_values as usize {
164                // 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]`.
165                let this_value = encoded_data_block.data.slice_with_length(
166                    prefix_sum[i] + (j * encoded_bytes_per_row),
167                    bytes_per_values[i],
168                );
169
170                child_buf.extend_from_slice(&this_value);
171            }
172
173            let child = DataBlock::FixedWidth(FixedWidthDataBlock {
174                data: LanceBuffer::from(child_buf),
175                bits_per_value: self.bits_per_values[i],
176                num_values,
177                block_info: BlockInfo::default(),
178            });
179            children_data_block.push(child);
180        }
181        Ok(DataBlock::Struct(StructDataBlock {
182            children: children_data_block,
183            block_info: BlockInfo::default(),
184            validity: None,
185        }))
186    }
187}
188
189#[derive(Debug)]
190enum VariablePackedFieldData {
191    Fixed {
192        block: FixedWidthDataBlock,
193    },
194    Variable {
195        block: VariableWidthBlock,
196        bits_per_length: u64,
197    },
198}
199
200impl VariablePackedFieldData {
201    fn append_row_bytes(&self, row_idx: usize, output: &mut Vec<u8>) -> Result<()> {
202        match self {
203            Self::Fixed { block } => {
204                let bits_per_value = block.bits_per_value;
205                if bits_per_value % 8 != 0 {
206                    return Err(Error::invalid_input(
207                        "Packed struct variable encoding requires byte-aligned fixed-width children",
208                    ));
209                }
210                let bytes_per_value = (bits_per_value / 8) as usize;
211                let start = row_idx
212                    .checked_mul(bytes_per_value)
213                    .ok_or_else(|| Error::invalid_input("Packed struct row size overflow"))?;
214                let end = start + bytes_per_value;
215                let data = block.data.as_ref();
216                if end > data.len() {
217                    return Err(Error::invalid_input(
218                        "Packed struct fixed child out of bounds",
219                    ));
220                }
221                output.extend_from_slice(&data[start..end]);
222                Ok(())
223            }
224            Self::Variable {
225                block,
226                bits_per_length,
227            } => {
228                if bits_per_length % 8 != 0 {
229                    return Err(Error::invalid_input(
230                        "Packed struct variable children must have byte-aligned length prefixes",
231                    ));
232                }
233                let prefix_bytes = (*bits_per_length / 8) as usize;
234                if !(prefix_bytes == 4 || prefix_bytes == 8) {
235                    return Err(Error::invalid_input(
236                        "Packed struct variable children must use 32 or 64-bit length prefixes",
237                    ));
238                }
239                match block.bits_per_offset {
240                    32 => {
241                        let offsets = block.offsets.borrow_to_typed_slice::<u32>();
242                        let start = offsets[row_idx] as usize;
243                        let end = offsets[row_idx + 1] as usize;
244                        if end > block.data.len() {
245                            return Err(Error::invalid_input(
246                                "Packed struct variable child offsets out of bounds",
247                            ));
248                        }
249                        let len = (end - start) as u32;
250                        if prefix_bytes != std::mem::size_of::<u32>() {
251                            return Err(Error::invalid_input(
252                                "Packed struct variable child length prefix mismatch",
253                            ));
254                        }
255                        output.extend_from_slice(&len.to_le_bytes());
256                        output.extend_from_slice(&block.data[start..end]);
257                        Ok(())
258                    }
259                    64 => {
260                        let offsets = block.offsets.borrow_to_typed_slice::<u64>();
261                        let start = offsets[row_idx] as usize;
262                        let end = offsets[row_idx + 1] as usize;
263                        if end > block.data.len() {
264                            return Err(Error::invalid_input(
265                                "Packed struct variable child offsets out of bounds",
266                            ));
267                        }
268                        let len = (end - start) as u64;
269                        if prefix_bytes != std::mem::size_of::<u64>() {
270                            return Err(Error::invalid_input(
271                                "Packed struct variable child length prefix mismatch",
272                            ));
273                        }
274                        output.extend_from_slice(&len.to_le_bytes());
275                        output.extend_from_slice(&block.data[start..end]);
276                        Ok(())
277                    }
278                    _ => Err(Error::invalid_input(
279                        "Packed struct variable child must use 32 or 64-bit offsets",
280                    )),
281                }
282            }
283        }
284    }
285}
286
287#[derive(Debug)]
288pub struct PackedStructVariablePerValueEncoder {
289    strategy: Arc<dyn CompressionStrategy>,
290    fields: Vec<Field>,
291}
292
293impl PackedStructVariablePerValueEncoder {
294    pub fn new(strategy: Arc<dyn CompressionStrategy>, fields: Vec<Field>) -> Self {
295        Self { strategy, fields }
296    }
297}
298
299impl PerValueCompressor for PackedStructVariablePerValueEncoder {
300    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
301        let DataBlock::Struct(struct_block) = data else {
302            return Err(Error::invalid_input(
303                "Packed struct encoder requires Struct data block",
304            ));
305        };
306
307        if struct_block.children.is_empty() {
308            return Err(Error::invalid_input(
309                "Packed struct encoder requires at least one child field",
310            ));
311        }
312        if struct_block.children.len() != self.fields.len() {
313            return Err(Error::invalid_input(
314                "Struct field metadata does not match number of children",
315            ));
316        }
317
318        let num_values = struct_block.children[0].num_values();
319        for child in struct_block.children.iter() {
320            if child.num_values() != num_values {
321                return Err(Error::invalid_input(
322                    "Packed struct children must have matching value counts",
323                ));
324            }
325        }
326
327        let mut field_data = Vec::with_capacity(self.fields.len());
328        let mut field_metadata = Vec::with_capacity(self.fields.len());
329
330        for (field, child_block) in self.fields.iter().zip(struct_block.children) {
331            let compressor = self.strategy.create_per_value(field, &child_block)?;
332            let (compressed, encoding) = compressor.compress(child_block)?;
333            match compressed {
334                PerValueDataBlock::Fixed(block) => {
335                    field_metadata.push(ProtobufUtils21::packed_struct_field_fixed(
336                        encoding,
337                        block.bits_per_value,
338                    ));
339                    field_data.push(VariablePackedFieldData::Fixed { block });
340                }
341                PerValueDataBlock::Variable(block) => {
342                    let bits_per_length = block.bits_per_offset as u64;
343                    field_metadata.push(ProtobufUtils21::packed_struct_field_variable(
344                        encoding,
345                        bits_per_length,
346                    ));
347                    field_data.push(VariablePackedFieldData::Variable {
348                        block,
349                        bits_per_length,
350                    });
351                }
352            }
353        }
354
355        let mut row_data: Vec<u8> = Vec::new();
356        let mut row_offsets: Vec<u64> = Vec::with_capacity(num_values as usize + 1);
357        row_offsets.push(0);
358        let mut total_bytes: usize = 0;
359        let mut max_row_len: usize = 0;
360        for row in 0..num_values as usize {
361            let start = row_data.len();
362            for field in &field_data {
363                field.append_row_bytes(row, &mut row_data)?;
364            }
365            let end = row_data.len();
366            let row_len = end - start;
367            max_row_len = max_row_len.max(row_len);
368            total_bytes = total_bytes
369                .checked_add(row_len)
370                .ok_or_else(|| Error::invalid_input("Packed struct row data size overflow"))?;
371            row_offsets.push(end as u64);
372        }
373        debug_assert_eq!(total_bytes, row_data.len());
374
375        let use_u32_offsets = total_bytes <= u32::MAX as usize && max_row_len <= u32::MAX as usize;
376        let bits_per_offset = if use_u32_offsets { 32 } else { 64 };
377        let offsets_buffer = if use_u32_offsets {
378            let offsets_u32 = row_offsets
379                .iter()
380                .map(|&offset| offset as u32)
381                .collect::<Vec<_>>();
382            LanceBuffer::reinterpret_vec(offsets_u32)
383        } else {
384            LanceBuffer::reinterpret_vec(row_offsets)
385        };
386
387        let data_block = VariableWidthBlock {
388            data: LanceBuffer::from(row_data),
389            bits_per_offset,
390            offsets: offsets_buffer,
391            num_values,
392            block_info: BlockInfo::new(),
393        };
394
395        Ok((
396            PerValueDataBlock::Variable(data_block),
397            ProtobufUtils21::packed_struct_variable(field_metadata),
398        ))
399    }
400}
401
402#[derive(Debug)]
403pub(crate) enum VariablePackedStructFieldKind {
404    Fixed {
405        bits_per_value: u64,
406        decompressor: Arc<dyn FixedPerValueDecompressor>,
407    },
408    Variable {
409        bits_per_length: u64,
410        decompressor: Arc<dyn VariablePerValueDecompressor>,
411    },
412}
413
414#[derive(Debug)]
415pub(crate) struct VariablePackedStructFieldDecoder {
416    pub(crate) kind: VariablePackedStructFieldKind,
417}
418
419#[derive(Debug)]
420pub struct PackedStructVariablePerValueDecompressor {
421    fields: Vec<VariablePackedStructFieldDecoder>,
422}
423
424impl PackedStructVariablePerValueDecompressor {
425    pub(crate) fn new(fields: Vec<VariablePackedStructFieldDecoder>) -> Self {
426        Self { fields }
427    }
428}
429
430enum FieldAccumulator {
431    Fixed {
432        builder: DataBlockBuilder,
433        bits_per_value: u64,
434        empty_value: DataBlock,
435    },
436    Variable32 {
437        builder: DataBlockBuilder,
438        empty_value: DataBlock,
439    },
440    Variable64 {
441        builder: DataBlockBuilder,
442        empty_value: DataBlock,
443    },
444}
445
446impl FieldAccumulator {
447    // In full-zip variable packed decoding, rep/def may produce a visible row
448    // with an empty payload (e.g. null/invalid item). We still need to append
449    // one placeholder per child so child row counts remain aligned.
450    fn append_empty(&mut self) {
451        match self {
452            Self::Fixed {
453                builder,
454                empty_value,
455                ..
456            } => builder.append(empty_value, 0..1),
457            Self::Variable32 {
458                builder,
459                empty_value,
460            } => builder.append(empty_value, 0..1),
461            Self::Variable64 {
462                builder,
463                empty_value,
464            } => builder.append(empty_value, 0..1),
465        }
466    }
467}
468
469impl VariablePerValueDecompressor for PackedStructVariablePerValueDecompressor {
470    fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
471        let num_values = data.num_values;
472        let offsets_u64 = match data.bits_per_offset {
473            32 => data
474                .offsets
475                .borrow_to_typed_slice::<u32>()
476                .iter()
477                .map(|v| *v as u64)
478                .collect::<Vec<_>>(),
479            64 => data
480                .offsets
481                .borrow_to_typed_slice::<u64>()
482                .as_ref()
483                .to_vec(),
484            _ => {
485                return Err(Error::invalid_input(
486                    "Packed struct row offsets must be 32 or 64 bits",
487                ));
488            }
489        };
490
491        if offsets_u64.len() != num_values as usize + 1 {
492            return Err(Error::invalid_input(
493                "Packed struct row offsets length mismatch",
494            ));
495        }
496
497        let mut accumulators = Vec::with_capacity(self.fields.len());
498        for field in &self.fields {
499            match &field.kind {
500                VariablePackedStructFieldKind::Fixed { bits_per_value, .. } => {
501                    if bits_per_value % 8 != 0 {
502                        return Err(Error::invalid_input(
503                            "Packed struct fixed child must be byte-aligned",
504                        ));
505                    }
506                    let bytes_per_value = bits_per_value.checked_div(8).ok_or_else(|| {
507                        Error::invalid_input("Invalid bits per value for packed struct field")
508                    })?;
509                    let estimate = bytes_per_value.checked_mul(num_values).ok_or_else(|| {
510                        Error::invalid_input("Packed struct fixed child allocation overflow")
511                    })?;
512                    let empty_value = DataBlock::FixedWidth(FixedWidthDataBlock {
513                        data: LanceBuffer::from(vec![0_u8; bytes_per_value as usize]),
514                        bits_per_value: *bits_per_value,
515                        num_values: 1,
516                        block_info: BlockInfo::new(),
517                    });
518                    accumulators.push(FieldAccumulator::Fixed {
519                        builder: DataBlockBuilder::with_capacity_estimate(estimate),
520                        bits_per_value: *bits_per_value,
521                        empty_value,
522                    });
523                }
524                VariablePackedStructFieldKind::Variable {
525                    bits_per_length, ..
526                } => match bits_per_length {
527                    32 => accumulators.push(FieldAccumulator::Variable32 {
528                        builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64),
529                        empty_value: DataBlock::VariableWidth(VariableWidthBlock {
530                            data: LanceBuffer::empty(),
531                            bits_per_offset: 32,
532                            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 0_u32]),
533                            num_values: 1,
534                            block_info: BlockInfo::new(),
535                        }),
536                    }),
537                    64 => accumulators.push(FieldAccumulator::Variable64 {
538                        builder: DataBlockBuilder::with_capacity_estimate(data.data.len() as u64),
539                        empty_value: DataBlock::VariableWidth(VariableWidthBlock {
540                            data: LanceBuffer::empty(),
541                            bits_per_offset: 64,
542                            offsets: LanceBuffer::reinterpret_vec(vec![0_u64, 0_u64]),
543                            num_values: 1,
544                            block_info: BlockInfo::new(),
545                        }),
546                    }),
547                    _ => {
548                        return Err(Error::invalid_input(
549                            "Packed struct variable child must use 32 or 64-bit length prefixes",
550                        ));
551                    }
552                },
553            }
554        }
555
556        for row_idx in 0..num_values as usize {
557            let row_start = offsets_u64[row_idx] as usize;
558            let row_end = offsets_u64[row_idx + 1] as usize;
559            if row_end > data.data.len() || row_start > row_end {
560                return Err(Error::invalid_input(
561                    "Packed struct row bounds exceed buffer",
562                ));
563            }
564            if row_start == row_end {
565                for accumulator in accumulators.iter_mut() {
566                    accumulator.append_empty();
567                }
568                continue;
569            }
570            let mut cursor = row_start;
571            for (field, accumulator) in self.fields.iter().zip(accumulators.iter_mut()) {
572                match (&field.kind, accumulator) {
573                    (
574                        VariablePackedStructFieldKind::Fixed { bits_per_value, .. },
575                        FieldAccumulator::Fixed {
576                            builder,
577                            bits_per_value: acc_bits,
578                            ..
579                        },
580                    ) => {
581                        debug_assert_eq!(bits_per_value, acc_bits);
582                        let bytes_per_value = (bits_per_value / 8) as usize;
583                        let end = cursor + bytes_per_value;
584                        if end > row_end {
585                            return Err(Error::invalid_input(
586                                "Packed struct fixed child exceeds row bounds",
587                            ));
588                        }
589                        let value_block = DataBlock::FixedWidth(FixedWidthDataBlock {
590                            data: LanceBuffer::from(data.data[cursor..end].to_vec()),
591                            bits_per_value: *bits_per_value,
592                            num_values: 1,
593                            block_info: BlockInfo::new(),
594                        });
595                        builder.append(&value_block, 0..1);
596                        cursor = end;
597                    }
598                    (
599                        VariablePackedStructFieldKind::Variable {
600                            bits_per_length, ..
601                        },
602                        FieldAccumulator::Variable32 { builder, .. },
603                    ) => {
604                        if *bits_per_length != 32 {
605                            return Err(Error::invalid_input(
606                                "Packed struct length prefix size mismatch",
607                            ));
608                        }
609                        let end = cursor + std::mem::size_of::<u32>();
610                        if end > row_end {
611                            return Err(Error::invalid_input(
612                                "Packed struct variable child length prefix out of bounds",
613                            ));
614                        }
615                        let len = u32::from_le_bytes(
616                            data.data[cursor..end]
617                                .try_into()
618                                .expect("slice has exact length"),
619                        ) as usize;
620                        cursor = end;
621                        let value_end = cursor + len;
622                        if value_end > row_end {
623                            return Err(Error::invalid_input(
624                                "Packed struct variable child exceeds row bounds",
625                            ));
626                        }
627                        let value_block = DataBlock::VariableWidth(VariableWidthBlock {
628                            data: LanceBuffer::from(data.data[cursor..value_end].to_vec()),
629                            bits_per_offset: 32,
630                            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, len as u32]),
631                            num_values: 1,
632                            block_info: BlockInfo::new(),
633                        });
634                        builder.append(&value_block, 0..1);
635                        cursor = value_end;
636                    }
637                    (
638                        VariablePackedStructFieldKind::Variable {
639                            bits_per_length, ..
640                        },
641                        FieldAccumulator::Variable64 { builder, .. },
642                    ) => {
643                        if *bits_per_length != 64 {
644                            return Err(Error::invalid_input(
645                                "Packed struct length prefix size mismatch",
646                            ));
647                        }
648                        let end = cursor + std::mem::size_of::<u64>();
649                        if end > row_end {
650                            return Err(Error::invalid_input(
651                                "Packed struct variable child length prefix out of bounds",
652                            ));
653                        }
654                        let len = u64::from_le_bytes(
655                            data.data[cursor..end]
656                                .try_into()
657                                .expect("slice has exact length"),
658                        ) as usize;
659                        cursor = end;
660                        let value_end = cursor + len;
661                        if value_end > row_end {
662                            return Err(Error::invalid_input(
663                                "Packed struct variable child exceeds row bounds",
664                            ));
665                        }
666                        let value_block = DataBlock::VariableWidth(VariableWidthBlock {
667                            data: LanceBuffer::from(data.data[cursor..value_end].to_vec()),
668                            bits_per_offset: 64,
669                            offsets: LanceBuffer::reinterpret_vec(vec![0_u64, len as u64]),
670                            num_values: 1,
671                            block_info: BlockInfo::new(),
672                        });
673                        builder.append(&value_block, 0..1);
674                        cursor = value_end;
675                    }
676                    _ => {
677                        return Err(Error::invalid_input(
678                            "Packed struct accumulator kind mismatch",
679                        ));
680                    }
681                }
682            }
683            if cursor != row_end {
684                return Err(Error::invalid_input(
685                    "Packed struct row parsing did not consume full row",
686                ));
687            }
688        }
689
690        let mut children = Vec::with_capacity(self.fields.len());
691        for (field, accumulator) in self.fields.iter().zip(accumulators) {
692            match (field, accumulator) {
693                (
694                    VariablePackedStructFieldDecoder {
695                        kind: VariablePackedStructFieldKind::Fixed { decompressor, .. },
696                    },
697                    FieldAccumulator::Fixed { builder, .. },
698                ) => {
699                    let DataBlock::FixedWidth(block) = builder.finish() else {
700                        panic!("Expected fixed-width datablock from builder");
701                    };
702                    let decoded = decompressor.decompress(block, num_values)?;
703                    children.push(decoded);
704                }
705                (
706                    VariablePackedStructFieldDecoder {
707                        kind:
708                            VariablePackedStructFieldKind::Variable {
709                                bits_per_length,
710                                decompressor,
711                            },
712                    },
713                    FieldAccumulator::Variable32 { builder, .. },
714                ) => {
715                    let DataBlock::VariableWidth(mut block) = builder.finish() else {
716                        panic!("Expected variable-width datablock from builder");
717                    };
718                    debug_assert_eq!(block.bits_per_offset, 32);
719                    block.bits_per_offset = (*bits_per_length) as u8;
720                    let decoded = decompressor.decompress(block)?;
721                    children.push(decoded);
722                }
723                (
724                    VariablePackedStructFieldDecoder {
725                        kind:
726                            VariablePackedStructFieldKind::Variable {
727                                bits_per_length,
728                                decompressor,
729                            },
730                    },
731                    FieldAccumulator::Variable64 { builder, .. },
732                ) => {
733                    let DataBlock::VariableWidth(mut block) = builder.finish() else {
734                        panic!("Expected variable-width datablock from builder");
735                    };
736                    debug_assert_eq!(block.bits_per_offset, 64);
737                    block.bits_per_offset = (*bits_per_length) as u8;
738                    let decoded = decompressor.decompress(block)?;
739                    children.push(decoded);
740                }
741                _ => {
742                    return Err(Error::invalid_input(
743                        "Packed struct accumulator mismatch during finalize",
744                    ));
745                }
746            }
747        }
748
749        Ok(DataBlock::Struct(StructDataBlock {
750            children,
751            block_info: BlockInfo::new(),
752            validity: None,
753        }))
754    }
755}
756
757#[cfg(test)]
758mod tests {
759    use super::*;
760    use crate::{
761        compression::DefaultDecompressionStrategy,
762        compression_config::CompressionParams,
763        constants::PACKED_STRUCT_META_KEY,
764        statistics::ComputeStat,
765        testing::{
766            TestCases, TestEncoding, check_round_trip_encoding_of_data, test_compression_strategy,
767        },
768    };
769    use arrow_array::{
770        Array, ArrayRef, BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray,
771        StructArray, UInt32Array,
772    };
773    use arrow_schema::{DataType, Field as ArrowField, Fields};
774    use std::collections::HashMap;
775    use std::sync::Arc;
776
777    fn fixed_block_from_array(array: Int64Array) -> FixedWidthDataBlock {
778        let num_values = array.len() as u64;
779        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
780        match block {
781            DataBlock::FixedWidth(block) => block,
782            _ => panic!("Expected fixed-width data block"),
783        }
784    }
785
786    fn fixed_i32_block_from_array(array: Int32Array) -> FixedWidthDataBlock {
787        let num_values = array.len() as u64;
788        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
789        match block {
790            DataBlock::FixedWidth(block) => block,
791            _ => panic!("Expected fixed-width data block"),
792        }
793    }
794
795    fn variable_block_from_string_array(array: StringArray) -> VariableWidthBlock {
796        let num_values = array.len() as u64;
797        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
798        match block {
799            DataBlock::VariableWidth(block) => block,
800            _ => panic!("Expected variable-width block"),
801        }
802    }
803
804    fn variable_block_from_large_string_array(array: LargeStringArray) -> VariableWidthBlock {
805        let num_values = array.len() as u64;
806        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
807        match block {
808            DataBlock::VariableWidth(block) => block,
809            _ => panic!("Expected variable-width block"),
810        }
811    }
812
813    fn variable_block_from_binary_array(array: BinaryArray) -> VariableWidthBlock {
814        let num_values = array.len() as u64;
815        let block = DataBlock::from_arrays(&[Arc::new(array) as ArrayRef], num_values);
816        match block {
817            DataBlock::VariableWidth(block) => block,
818            _ => panic!("Expected variable-width block"),
819        }
820    }
821
822    #[test]
823    fn variable_packed_struct_round_trip() -> Result<()> {
824        let arrow_fields: Fields = vec![
825            ArrowField::new("id", DataType::UInt32, false),
826            ArrowField::new("name", DataType::Utf8, true),
827        ]
828        .into();
829        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
830        let struct_field = Field::try_from(&arrow_struct)?;
831
832        let ids = vec![1_u32, 2, 42];
833        let id_bytes = ids
834            .iter()
835            .flat_map(|value| value.to_le_bytes())
836            .collect::<Vec<_>>();
837        let mut id_block = FixedWidthDataBlock {
838            data: LanceBuffer::reinterpret_vec(ids),
839            bits_per_value: 32,
840            num_values: 3,
841            block_info: BlockInfo::new(),
842        };
843        id_block.compute_stat();
844        let id_block = DataBlock::FixedWidth(id_block);
845
846        let name_offsets = vec![0_i32, 1, 4, 4];
847        let name_bytes = b"abcz".to_vec();
848        let mut name_block = VariableWidthBlock {
849            data: LanceBuffer::from(name_bytes.clone()),
850            bits_per_offset: 32,
851            offsets: LanceBuffer::reinterpret_vec(name_offsets.clone()),
852            num_values: 3,
853            block_info: BlockInfo::new(),
854        };
855        name_block.compute_stat();
856        let name_block = DataBlock::VariableWidth(name_block);
857
858        let struct_block = StructDataBlock {
859            children: vec![id_block, name_block],
860            block_info: BlockInfo::new(),
861            validity: None,
862        };
863
864        let data_block = DataBlock::Struct(struct_block);
865
866        let compression_strategy =
867            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
868        let compressor = crate::compression::CompressionStrategy::create_per_value(
869            compression_strategy.as_ref(),
870            &struct_field,
871            &data_block,
872        )?;
873        let (compressed, encoding) = compressor.compress(data_block)?;
874
875        let PerValueDataBlock::Variable(zipped) = compressed else {
876            panic!("expected variable-width packed struct output");
877        };
878
879        let decompression_strategy = DefaultDecompressionStrategy::default();
880        let decompressor =
881            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
882                &decompression_strategy,
883                &encoding,
884            )?;
885        let decoded = decompressor.decompress(zipped)?;
886
887        let DataBlock::Struct(decoded_struct) = decoded else {
888            panic!("expected struct datablock after decode");
889        };
890
891        let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap();
892        assert_eq!(decoded_id.bits_per_value, 32);
893        assert_eq!(decoded_id.data.as_ref(), id_bytes.as_slice());
894
895        let decoded_name = decoded_struct.children[1].as_variable_width_ref().unwrap();
896        assert_eq!(decoded_name.bits_per_offset, 32);
897        let decoded_offsets = decoded_name.offsets.borrow_to_typed_slice::<i32>();
898        assert_eq!(decoded_offsets.as_ref(), name_offsets.as_slice());
899        assert_eq!(decoded_name.data.as_ref(), name_bytes.as_slice());
900
901        Ok(())
902    }
903
904    #[test]
905    fn variable_packed_struct_large_utf8_round_trip() -> Result<()> {
906        let arrow_fields: Fields = vec![
907            ArrowField::new("value", DataType::Int64, false),
908            ArrowField::new("text", DataType::LargeUtf8, false),
909        ]
910        .into();
911        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
912        let struct_field = Field::try_from(&arrow_struct)?;
913
914        let id_block = fixed_block_from_array(Int64Array::from(vec![10, 20, 30, 40]));
915        let payload_array = LargeStringArray::from(vec![
916            "alpha",
917            "a considerably longer payload for testing",
918            "mid",
919            "z",
920        ]);
921        let payload_block = variable_block_from_large_string_array(payload_array);
922
923        let struct_block = StructDataBlock {
924            children: vec![
925                DataBlock::FixedWidth(id_block.clone()),
926                DataBlock::VariableWidth(payload_block.clone()),
927            ],
928            block_info: BlockInfo::new(),
929            validity: None,
930        };
931
932        let data_block = DataBlock::Struct(struct_block);
933
934        let compression_strategy =
935            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
936        let compressor = crate::compression::CompressionStrategy::create_per_value(
937            compression_strategy.as_ref(),
938            &struct_field,
939            &data_block,
940        )?;
941        let (compressed, encoding) = compressor.compress(data_block)?;
942
943        let PerValueDataBlock::Variable(zipped) = compressed else {
944            panic!("expected variable-width packed struct output");
945        };
946
947        let decompression_strategy = DefaultDecompressionStrategy::default();
948        let decompressor =
949            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
950                &decompression_strategy,
951                &encoding,
952            )?;
953        let decoded = decompressor.decompress(zipped)?;
954
955        let DataBlock::Struct(decoded_struct) = decoded else {
956            panic!("expected struct datablock after decode");
957        };
958
959        let decoded_id = decoded_struct.children[0].as_fixed_width_ref().unwrap();
960        assert_eq!(decoded_id.bits_per_value, 64);
961        assert_eq!(decoded_id.data.as_ref(), id_block.data.as_ref());
962
963        let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap();
964        assert_eq!(decoded_payload.bits_per_offset, 64);
965        assert_eq!(
966            decoded_payload
967                .offsets
968                .borrow_to_typed_slice::<i64>()
969                .as_ref(),
970            payload_block
971                .offsets
972                .borrow_to_typed_slice::<i64>()
973                .as_ref()
974        );
975        assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref());
976
977        Ok(())
978    }
979
980    #[tokio::test]
981    async fn variable_packed_struct_utf8_round_trip() {
982        // schema: Struct<id: UInt32, uri: Utf8, long_text: LargeUtf8>
983        let fields = Fields::from(vec![
984            Arc::new(ArrowField::new("id", DataType::UInt32, false)),
985            Arc::new(ArrowField::new("uri", DataType::Utf8, false)),
986            Arc::new(ArrowField::new("long_text", DataType::LargeUtf8, false)),
987        ]);
988
989        // mark struct as packed
990        let mut meta = HashMap::new();
991        meta.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
992
993        let array = Arc::new(StructArray::from(vec![
994            (
995                fields[0].clone(),
996                Arc::new(UInt32Array::from(vec![1, 2, 3])) as ArrayRef,
997            ),
998            (
999                fields[1].clone(),
1000                Arc::new(StringArray::from(vec![
1001                    Some("a"),
1002                    Some("b"),
1003                    Some("/tmp/x"),
1004                ])) as ArrayRef,
1005            ),
1006            (
1007                fields[2].clone(),
1008                Arc::new(LargeStringArray::from(vec![
1009                    Some("alpha"),
1010                    Some("a considerably longer payload for testing"),
1011                    Some("mid"),
1012                ])) as ArrayRef,
1013            ),
1014        ]));
1015
1016        let test_cases = TestCases::default()
1017            .with_u32_structural_encodings()
1018            .with_expected_encoding("variable_packed_struct");
1019
1020        check_round_trip_encoding_of_data(vec![array], &test_cases, meta).await;
1021    }
1022
1023    #[test]
1024    fn variable_packed_struct_multi_variable_round_trip() -> Result<()> {
1025        let arrow_fields: Fields = vec![
1026            ArrowField::new("category", DataType::Utf8, false),
1027            ArrowField::new("payload", DataType::Binary, false),
1028            ArrowField::new("count", DataType::Int32, false),
1029        ]
1030        .into();
1031        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1032        let struct_field = Field::try_from(&arrow_struct)?;
1033
1034        let category_array = StringArray::from(vec!["red", "blue", "green", "red"]);
1035        let category_block = variable_block_from_string_array(category_array);
1036        let payload_values: Vec<Vec<u8>> =
1037            vec![vec![0x01, 0x02], vec![], vec![0x05, 0x06, 0x07], vec![0xff]];
1038        let payload_array =
1039            BinaryArray::from_iter_values(payload_values.iter().map(|v| v.as_slice()));
1040        let payload_block = variable_block_from_binary_array(payload_array);
1041        let count_block = fixed_i32_block_from_array(Int32Array::from(vec![1, 2, 3, 4]));
1042
1043        let struct_block = StructDataBlock {
1044            children: vec![
1045                DataBlock::VariableWidth(category_block.clone()),
1046                DataBlock::VariableWidth(payload_block.clone()),
1047                DataBlock::FixedWidth(count_block.clone()),
1048            ],
1049            block_info: BlockInfo::new(),
1050            validity: None,
1051        };
1052
1053        let data_block = DataBlock::Struct(struct_block);
1054
1055        let compression_strategy =
1056            test_compression_strategy(TestEncoding::StructuralU32, CompressionParams::default());
1057        let compressor = crate::compression::CompressionStrategy::create_per_value(
1058            compression_strategy.as_ref(),
1059            &struct_field,
1060            &data_block,
1061        )?;
1062        let (compressed, encoding) = compressor.compress(data_block)?;
1063
1064        let PerValueDataBlock::Variable(zipped) = compressed else {
1065            panic!("expected variable-width packed struct output");
1066        };
1067
1068        let decompression_strategy = DefaultDecompressionStrategy::default();
1069        let decompressor =
1070            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1071                &decompression_strategy,
1072                &encoding,
1073            )?;
1074        let decoded = decompressor.decompress(zipped)?;
1075
1076        let DataBlock::Struct(decoded_struct) = decoded else {
1077            panic!("expected struct datablock after decode");
1078        };
1079
1080        let decoded_category = decoded_struct.children[0].as_variable_width_ref().unwrap();
1081        assert_eq!(decoded_category.bits_per_offset, 32);
1082        assert_eq!(
1083            decoded_category
1084                .offsets
1085                .borrow_to_typed_slice::<i32>()
1086                .as_ref(),
1087            category_block
1088                .offsets
1089                .borrow_to_typed_slice::<i32>()
1090                .as_ref()
1091        );
1092        assert_eq!(decoded_category.data.as_ref(), category_block.data.as_ref());
1093
1094        let decoded_payload = decoded_struct.children[1].as_variable_width_ref().unwrap();
1095        assert_eq!(decoded_payload.bits_per_offset, 32);
1096        assert_eq!(
1097            decoded_payload
1098                .offsets
1099                .borrow_to_typed_slice::<i32>()
1100                .as_ref(),
1101            payload_block
1102                .offsets
1103                .borrow_to_typed_slice::<i32>()
1104                .as_ref()
1105        );
1106        assert_eq!(decoded_payload.data.as_ref(), payload_block.data.as_ref());
1107
1108        let decoded_count = decoded_struct.children[2].as_fixed_width_ref().unwrap();
1109        assert_eq!(decoded_count.bits_per_value, 32);
1110        assert_eq!(decoded_count.data.as_ref(), count_block.data.as_ref());
1111
1112        Ok(())
1113    }
1114
1115    #[test]
1116    fn variable_packed_struct_requires_v22() {
1117        let arrow_fields: Fields = vec![
1118            ArrowField::new("value", DataType::Int64, false),
1119            ArrowField::new("text", DataType::Utf8, false),
1120        ]
1121        .into();
1122        let arrow_struct = ArrowField::new("item", DataType::Struct(arrow_fields), false);
1123        let struct_field = Field::try_from(&arrow_struct).unwrap();
1124
1125        let value_block = fixed_block_from_array(Int64Array::from(vec![1, 2, 3]));
1126        let text_block =
1127            variable_block_from_string_array(StringArray::from(vec!["a", "bb", "ccc"]));
1128
1129        let struct_block = StructDataBlock {
1130            children: vec![
1131                DataBlock::FixedWidth(value_block),
1132                DataBlock::VariableWidth(text_block),
1133            ],
1134            block_info: BlockInfo::new(),
1135            validity: None,
1136        };
1137
1138        let compression_strategy =
1139            test_compression_strategy(TestEncoding::StructuralU16, CompressionParams::default());
1140        let result =
1141            compression_strategy.create_per_value(&struct_field, &DataBlock::Struct(struct_block));
1142
1143        assert!(matches!(result, Err(Error::NotSupported { .. })));
1144    }
1145
1146    #[test]
1147    fn variable_packed_struct_decompress_empty_row() -> Result<()> {
1148        let strategy = DefaultDecompressionStrategy::default();
1149        let fixed_decompressor = Arc::from(
1150            crate::compression::DecompressionStrategy::create_fixed_per_value_decompressor(
1151                &strategy,
1152                &ProtobufUtils21::flat(32, None),
1153            )?,
1154        );
1155        let variable_decompressor = Arc::from(
1156            crate::compression::DecompressionStrategy::create_variable_per_value_decompressor(
1157                &strategy,
1158                &ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None),
1159            )?,
1160        );
1161
1162        let decompressor = PackedStructVariablePerValueDecompressor::new(vec![
1163            VariablePackedStructFieldDecoder {
1164                kind: VariablePackedStructFieldKind::Fixed {
1165                    bits_per_value: 32,
1166                    decompressor: fixed_decompressor,
1167                },
1168            },
1169            VariablePackedStructFieldDecoder {
1170                kind: VariablePackedStructFieldKind::Variable {
1171                    bits_per_length: 32,
1172                    decompressor: variable_decompressor,
1173                },
1174            },
1175        ]);
1176
1177        let mut row_data = Vec::new();
1178        row_data.extend_from_slice(&1_u32.to_le_bytes());
1179        row_data.extend_from_slice(&1_u32.to_le_bytes());
1180        row_data.extend_from_slice(b"a");
1181        row_data.extend_from_slice(&2_u32.to_le_bytes());
1182        row_data.extend_from_slice(&0_u32.to_le_bytes());
1183
1184        let input = VariableWidthBlock {
1185            data: LanceBuffer::from(row_data),
1186            bits_per_offset: 32,
1187            offsets: LanceBuffer::reinterpret_vec(vec![0_u32, 9_u32, 9_u32, 17_u32]),
1188            num_values: 3,
1189            block_info: BlockInfo::new(),
1190        };
1191
1192        let decoded = decompressor.decompress(input)?;
1193        let DataBlock::Struct(decoded_struct) = decoded else {
1194            panic!("expected struct output");
1195        };
1196
1197        let fixed = decoded_struct.children[0].as_fixed_width_ref().unwrap();
1198        assert_eq!(fixed.bits_per_value, 32);
1199        assert_eq!(
1200            fixed.data.borrow_to_typed_slice::<u32>().as_ref(),
1201            &[1, 0, 2]
1202        );
1203
1204        let variable = decoded_struct.children[1].as_variable_width_ref().unwrap();
1205        assert_eq!(variable.bits_per_offset, 32);
1206        assert_eq!(
1207            variable.offsets.borrow_to_typed_slice::<u32>().as_ref(),
1208            &[0_u32, 1_u32, 1_u32, 1_u32]
1209        );
1210        assert_eq!(variable.data.as_ref(), b"a");
1211
1212        Ok(())
1213    }
1214}
lance_encoding/encodings/physical/packed.rs

lance_encoding/encodings/physical/
packed.rs
