parquet/encodings/encoding/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Contains all supported encoders for Parquet.
19
20use std::{cmp, marker::PhantomData};
21
22use crate::basic::*;
23use crate::data_type::private::ParquetValueType;
24use crate::data_type::*;
25use crate::encodings::rle::RleEncoder;
26use crate::errors::{ParquetError, Result};
27use crate::file::properties::ResolvedColumnProperties;
28use crate::schema::types::ColumnDescPtr;
29use crate::util::bit_util::{BitWriter, num_required_bits};
30use crate::util::prefix::common_prefix_length;
31
32use alp_encoder::AlpEncoder;
33use byte_stream_split_encoder::{ByteStreamSplitEncoder, VariableWidthByteStreamSplitEncoder};
34use bytes::Bytes;
35pub use dict_encoder::DictEncoder;
36
37mod alp_encoder;
38mod byte_stream_split_encoder;
39mod dict_encoder;
40
41// ----------------------------------------------------------------------
42// Encoders
43
44/// An Parquet encoder for the data type `T`.
45///
46/// Currently this allocates internal buffers for the encoded values. After done putting
47/// values, caller should call `flush_buffer()` to get an immutable buffer pointer.
48pub trait Encoder<T: DataType>: Send {
49    /// Encodes data from `values`.
50    fn put(&mut self, values: &[T::T]) -> Result<()>;
51
52    /// Encodes data from `values`, which contains spaces for null values, that is
53    /// identified by `valid_bits`.
54    ///
55    /// Returns the number of non-null values encoded.
56    #[cfg(test)]
57    fn put_spaced(&mut self, values: &[T::T], valid_bits: &[u8]) -> Result<usize> {
58        let num_values = values.len();
59        let mut buffer = Vec::with_capacity(num_values);
60        // TODO: this is pretty inefficient. Revisit in future.
61        for (i, item) in values.iter().enumerate().take(num_values) {
62            if crate::util::bit_util::get_bit(valid_bits, i) {
63                buffer.push(item.clone());
64            }
65        }
66        self.put(&buffer[..])?;
67        Ok(buffer.len())
68    }
69
70    /// Returns the encoding type of this encoder.
71    fn encoding(&self) -> Encoding;
72
73    /// Returns an estimate of the encoded data, in bytes.
74    /// Method call must be O(1).
75    fn estimated_data_encoded_size(&self) -> usize;
76
77    /// Returns an estimate of the memory use of this encoder, in bytes
78    fn estimated_memory_size(&self) -> usize;
79
80    /// Flushes the underlying byte buffer that's being processed by this encoder, and
81    /// return the immutable copy of it. This will also reset the internal state.
82    fn flush_buffer(&mut self) -> Result<Bytes>;
83}
84
85/// Gets a encoder for the particular data type `T` and encoding `encoding`. Memory usage
86/// for the encoder instance is tracked by `mem_tracker`.
87#[cfg_attr(all(not(feature = "experimental"), not(test)), expect(dead_code))]
88pub fn get_encoder<T: DataType>(
89    encoding: Encoding,
90    descr: &ColumnDescPtr,
91) -> Result<Box<dyn Encoder<T>>> {
92    <T::T as private::GetEncoder>::get_encoder(descr, encoding, None)
93}
94
95pub(crate) fn get_encoder_with_properties<T: DataType>(
96    encoding: Encoding,
97    descr: &ColumnDescPtr,
98    column_props: &ResolvedColumnProperties,
99) -> Result<Box<dyn Encoder<T>>> {
100    <T::T as private::GetEncoder>::get_encoder(descr, encoding, Some(column_props))
101}
102
103pub(crate) mod private {
104    use super::*;
105
106    /// A trait that allows getting an [`Encoder`] implementation for a [`DataType`]
107    /// with the corresponding [`ParquetValueType`]. This is necessary to support
108    /// [`Encoder`] implementations that may not be applicable for all [`DataType`]
109    /// and by extension all [`ParquetValueType`], such as ALP, which encodes only
110    /// floating-point columns.
111    ///
112    /// [`ParquetValueType`]: crate::data_type::private::ParquetValueType
113    pub trait GetEncoder {
114        #[expect(
115            private_interfaces,
116            reason = "this trait is not nameable outside the crate"
117        )]
118        fn get_encoder<T: DataType<T = Self>>(
119            descr: &ColumnDescPtr,
120            encoding: Encoding,
121            options: Option<&ResolvedColumnProperties>,
122        ) -> Result<Box<dyn Encoder<T>>> {
123            get_encoder_default(descr, encoding, options)
124        }
125    }
126
127    fn get_encoder_default<T: DataType>(
128        descr: &ColumnDescPtr,
129        encoding: Encoding,
130        options: Option<&ResolvedColumnProperties>,
131    ) -> Result<Box<dyn Encoder<T>>> {
132        let encoder: Box<dyn Encoder<T>> = match encoding {
133            Encoding::PLAIN => Box::new(PlainEncoder::new()),
134            Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => {
135                return Err(general_err!(
136                    "Cannot initialize this encoding through this function"
137                ));
138            }
139            Encoding::RLE => match T::get_physical_type() {
140                Type::BOOLEAN => Box::new(RleValueEncoder::new()),
141                physical_type => return Err(unsupported_column_encoding(encoding, physical_type)),
142            },
143            Encoding::DELTA_BINARY_PACKED => match T::get_physical_type() {
144                Type::INT32 | Type::INT64 => Box::new(
145                    match options.and_then(|props| props.delta_binary_packed_encoder_options) {
146                        Some(options) => DeltaBitPackEncoder::new_with_options(options),
147                        None => DeltaBitPackEncoder::new(),
148                    },
149                ),
150                physical_type => return Err(unsupported_column_encoding(encoding, physical_type)),
151            },
152            Encoding::DELTA_LENGTH_BYTE_ARRAY => match T::get_physical_type() {
153                Type::BYTE_ARRAY => Box::new(
154                    match options.and_then(|props| props.delta_binary_packed_encoder_options) {
155                        Some(options) => DeltaLengthByteArrayEncoder::new_with_options(options),
156                        None => DeltaLengthByteArrayEncoder::new(),
157                    },
158                ),
159                physical_type => return Err(unsupported_column_encoding(encoding, physical_type)),
160            },
161            Encoding::DELTA_BYTE_ARRAY => match T::get_physical_type() {
162                Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => Box::new(
163                    match options.and_then(|props| props.delta_binary_packed_encoder_options) {
164                        Some(options) => DeltaByteArrayEncoder::new_with_options(options),
165                        None => DeltaByteArrayEncoder::new(),
166                    },
167                ),
168                physical_type => return Err(unsupported_column_encoding(encoding, physical_type)),
169            },
170            Encoding::BYTE_STREAM_SPLIT => match T::get_physical_type() {
171                Type::FIXED_LEN_BYTE_ARRAY => Box::new(VariableWidthByteStreamSplitEncoder::new(
172                    descr.type_length(),
173                )),
174                Type::INT32 | Type::INT64 | Type::FLOAT | Type::DOUBLE => {
175                    Box::new(ByteStreamSplitEncoder::new())
176                }
177                physical_type => return Err(unsupported_column_encoding(encoding, physical_type)),
178            },
179            Encoding::ALP => {
180                return Err(general_err!(
181                    "Encoding {} only supports FLOAT and DOUBLE, got {}",
182                    encoding,
183                    T::get_physical_type()
184                ));
185            }
186            #[expect(deprecated, reason = "BIT_PACKED is the encoding we reject here")]
187            e @ Encoding::BIT_PACKED => return Err(nyi_err!("Encoding {} is not supported", e)),
188        };
189        Ok(encoder)
190    }
191
192    impl GetEncoder for bool {}
193    impl GetEncoder for i32 {}
194    impl GetEncoder for i64 {}
195    impl GetEncoder for Int96 {}
196    impl GetEncoder for ByteArray {}
197    impl GetEncoder for FixedLenByteArray {}
198
199    impl GetEncoder for f32 {
200        #[expect(
201            private_interfaces,
202            reason = "this trait is not nameable outside the crate"
203        )]
204        fn get_encoder<T: DataType<T = Self>>(
205            descr: &ColumnDescPtr,
206            encoding: Encoding,
207            options: Option<&ResolvedColumnProperties>,
208        ) -> Result<Box<dyn Encoder<T>>> {
209            match encoding {
210                Encoding::ALP => Ok(Box::new(AlpEncoder::new())),
211                _ => get_encoder_default(descr, encoding, options),
212            }
213        }
214    }
215
216    impl GetEncoder for f64 {
217        #[expect(
218            private_interfaces,
219            reason = "this trait is not nameable outside the crate"
220        )]
221        fn get_encoder<T: DataType<T = Self>>(
222            descr: &ColumnDescPtr,
223            encoding: Encoding,
224            options: Option<&ResolvedColumnProperties>,
225        ) -> Result<Box<dyn Encoder<T>>> {
226            match encoding {
227                Encoding::ALP => Ok(Box::new(AlpEncoder::new())),
228                _ => get_encoder_default(descr, encoding, options),
229            }
230        }
231    }
232}
233
234fn unsupported_column_encoding(encoding: Encoding, physical_type: Type) -> ParquetError {
235    nyi_err!(
236        "Encoding {} is not supported for physical type {:?}",
237        encoding,
238        physical_type
239    )
240}
241
242// ----------------------------------------------------------------------
243// Plain encoding
244
245/// Plain encoding that supports all types.
246/// Values are encoded back to back.
247/// The plain encoding is used whenever a more efficient encoding can not be used.
248/// It stores the data in the following format:
249/// - BOOLEAN - 1 bit per value, 0 is false; 1 is true.
250/// - INT32 - 4 bytes per value, stored as little-endian.
251/// - INT64 - 8 bytes per value, stored as little-endian.
252/// - FLOAT - 4 bytes per value, stored as IEEE little-endian.
253/// - DOUBLE - 8 bytes per value, stored as IEEE little-endian.
254/// - BYTE_ARRAY - 4 byte length stored as little endian, followed by bytes.
255/// - FIXED_LEN_BYTE_ARRAY - just the bytes are stored.
256pub struct PlainEncoder<T: DataType> {
257    buffer: Vec<u8>,
258    bit_writer: BitWriter,
259    _phantom: PhantomData<T>,
260}
261
262impl<T: DataType> Default for PlainEncoder<T> {
263    fn default() -> Self {
264        Self::new()
265    }
266}
267
268impl<T: DataType> PlainEncoder<T> {
269    /// Creates new plain encoder.
270    pub fn new() -> Self {
271        Self {
272            buffer: vec![],
273            bit_writer: BitWriter::new(256),
274            _phantom: PhantomData,
275        }
276    }
277}
278
279impl<T: DataType> Encoder<T> for PlainEncoder<T> {
280    // Performance Note:
281    // As far as can be seen these functions are rarely called and as such we can hint to the
282    // compiler that they dont need to be folded into hot locations in the final output.
283    #[cold]
284    fn encoding(&self) -> Encoding {
285        Encoding::PLAIN
286    }
287
288    fn estimated_data_encoded_size(&self) -> usize {
289        self.buffer.len() + self.bit_writer.bytes_written()
290    }
291
292    #[inline]
293    fn flush_buffer(&mut self) -> Result<Bytes> {
294        self.buffer
295            .extend_from_slice(self.bit_writer.flush_buffer());
296        self.bit_writer.clear();
297        Ok(std::mem::take(&mut self.buffer).into())
298    }
299
300    #[inline]
301    fn put(&mut self, values: &[T::T]) -> Result<()> {
302        T::T::encode(values, &mut self.buffer, &mut self.bit_writer)?;
303        Ok(())
304    }
305
306    /// Return the estimated memory size of this encoder.
307    fn estimated_memory_size(&self) -> usize {
308        self.buffer.capacity() * std::mem::size_of::<u8>() + self.bit_writer.estimated_memory_size()
309    }
310}
311
312// ----------------------------------------------------------------------
313// RLE encoding
314
315const DEFAULT_RLE_BUFFER_LEN: usize = 1024;
316
317/// RLE/Bit-Packing hybrid encoding for values.
318/// Currently is used only for data pages v2 and supports boolean types.
319pub struct RleValueEncoder<T: DataType> {
320    // Buffer with raw values that we collect,
321    // when flushing buffer they are encoded using RLE encoder
322    encoder: Option<RleEncoder>,
323    _phantom: PhantomData<T>,
324}
325
326impl<T: DataType> Default for RleValueEncoder<T> {
327    fn default() -> Self {
328        Self::new()
329    }
330}
331
332impl<T: DataType> RleValueEncoder<T> {
333    /// Creates new rle value encoder.
334    pub fn new() -> Self {
335        Self {
336            encoder: None,
337            _phantom: PhantomData,
338        }
339    }
340}
341
342impl<T: DataType> Encoder<T> for RleValueEncoder<T> {
343    #[inline]
344    fn put(&mut self, values: &[T::T]) -> Result<()> {
345        ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType");
346
347        let rle_encoder = self.encoder.get_or_insert_with(|| {
348            let mut buffer = Vec::with_capacity(DEFAULT_RLE_BUFFER_LEN);
349            // Reserve space for length
350            buffer.extend_from_slice(&[0; 4]);
351            RleEncoder::new_from_buf(1, buffer)
352        });
353
354        let mut buf = [0_u64; 64];
355        for chunk in values.chunks(buf.len()) {
356            let buf = &mut buf[..chunk.len()];
357            for (b, value) in buf.iter_mut().zip(chunk) {
358                *b = value.as_u64()?;
359            }
360            rle_encoder.put_batch(buf);
361        }
362        Ok(())
363    }
364
365    // Performance Note:
366    // As far as can be seen these functions are rarely called and as such we can hint to the
367    // compiler that they dont need to be folded into hot locations in the final output.
368    #[cold]
369    fn encoding(&self) -> Encoding {
370        Encoding::RLE
371    }
372
373    #[inline]
374    fn estimated_data_encoded_size(&self) -> usize {
375        match self.encoder {
376            Some(ref enc) => enc.len(),
377            None => 0,
378        }
379    }
380
381    #[inline]
382    fn flush_buffer(&mut self) -> Result<Bytes> {
383        ensure_phys_ty!(Type::BOOLEAN, "RleValueEncoder only supports BoolType");
384        let rle_encoder = self
385            .encoder
386            .take()
387            .expect("RLE value encoder is not initialized");
388
389        // Flush all encoder buffers and raw values
390        let mut buf = rle_encoder.consume();
391        assert!(buf.len() >= 4, "should have had padding inserted");
392
393        // Note that buf does not have any offset, all data is encoded bytes
394        let len = (buf.len() - 4) as i32;
395        buf[..4].copy_from_slice(&len.to_le_bytes());
396
397        Ok(buf.into())
398    }
399
400    /// return the estimated memory size of this encoder.
401    fn estimated_memory_size(&self) -> usize {
402        self.encoder
403            .as_ref()
404            .map_or(0, |enc| enc.estimated_memory_size())
405    }
406}
407
408// ----------------------------------------------------------------------
409// DELTA_BINARY_PACKED encoding
410
411const MAX_PAGE_HEADER_WRITER_SIZE: usize = 32;
412const DEFAULT_NUM_MINI_BLOCKS: usize = 4;
413
414/// Controls the block layout emitted by [`DeltaBitPackEncoder`] and the integer sub-encoders used
415/// by [`DeltaLengthByteArrayEncoder`] and [`DeltaByteArrayEncoder`].
416///
417/// In the [Parquet DELTA_BINARY_PACKED encoding], each block records a minimum delta and each mini
418/// block records its own bit width. The defaults use 128 values per block and 4 mini blocks for
419/// `INT32`, and 256 values per block and 4 mini blocks for `INT64`.
420///
421/// Smaller blocks or mini blocks can reduce encoded size for bursty data, where the required bit
422/// width changes over short ranges, and for short pages that would otherwise need more padding.
423/// They also add more minimum-delta and bit-width metadata. Larger blocks or mini blocks can reduce
424/// that metadata and block-processing overhead for long inputs with stable delta distributions,
425/// improving compression or throughput, but an outlier then affects more values and the encoder
426/// uses a larger scratch buffer. Benchmark representative data before changing the defaults.
427///
428/// The format requires block sizes to be multiples of 128 and mini block sizes to be multiples
429/// of 32.
430///
431/// [Parquet DELTA_BINARY_PACKED encoding]: https://github.com/apache/parquet-format/blob/master/Encodings.md#delta-encoding-delta_binary_packed--5
432#[derive(Debug, Clone, Copy, PartialEq, Eq)]
433pub struct DeltaBinaryPackedEncoderOptions {
434    block_size: usize,
435    mini_blocks_per_block: usize,
436}
437
438impl DeltaBinaryPackedEncoderOptions {
439    /// Creates validated delta binary packed encoder options.
440    pub fn try_new(block_size: usize, mini_blocks_per_block: usize) -> Result<Self> {
441        if block_size == 0 || !block_size.is_multiple_of(128) {
442            return Err(general_err!(
443                "delta binary packed block size must be a positive multiple of 128, got {block_size}"
444            ));
445        }
446        if mini_blocks_per_block == 0 || !block_size.is_multiple_of(mini_blocks_per_block) {
447            return Err(general_err!(
448                "delta binary packed mini block count must be a positive divisor of the block size, got {mini_blocks_per_block} for block size {block_size}"
449            ));
450        }
451
452        let mini_block_size = block_size / mini_blocks_per_block;
453        if !mini_block_size.is_multiple_of(32) {
454            return Err(general_err!(
455                "delta binary packed mini block size must be a multiple of 32, got {mini_block_size}"
456            ));
457        }
458
459        Ok(Self {
460            block_size,
461            mini_blocks_per_block,
462        })
463    }
464
465    /// Returns the number of values in each block.
466    pub fn block_size(&self) -> usize {
467        self.block_size
468    }
469
470    /// Returns the number of mini blocks in each block.
471    pub fn mini_blocks_per_block(&self) -> usize {
472        self.mini_blocks_per_block
473    }
474}
475
476/// Delta bit packed encoder.
477/// Consists of a header followed by blocks of delta encoded values binary packed.
478///
479/// Delta-binary-packing:
480/// ```shell
481///   [page-header] [block 1], [block 2], ... [block N]
482/// ```
483///
484/// Each page header consists of:
485/// ```shell
486///   [block size] [number of miniblocks in a block] [total value count] [first value]
487/// ```
488///
489/// Each block consists of:
490/// ```shell
491///   [min delta] [list of bitwidths of miniblocks] [miniblocks]
492/// ```
493///
494/// Current implementation writes values in `put` method, multiple calls to `put` to
495/// existing block or start new block if block size is exceeded. Calling `flush_buffer`
496/// writes out all data and resets internal state, including page header.
497///
498/// Supports only INT32 and INT64.
499pub struct DeltaBitPackEncoder<T: DataType> {
500    page_header_writer: BitWriter,
501    bit_writer: BitWriter,
502    total_values: usize,
503    first_value: i64,
504    current_value: i64,
505    block_size: usize,
506    mini_block_size: usize,
507    num_mini_blocks: usize,
508    values_in_block: usize,
509    deltas: Vec<i64>,
510    _phantom: PhantomData<T>,
511}
512
513impl<T: DataType> Default for DeltaBitPackEncoder<T> {
514    fn default() -> Self {
515        Self::new()
516    }
517}
518
519impl<T: DataType> DeltaBitPackEncoder<T> {
520    /// Creates new delta bit packed encoder.
521    pub fn new() -> Self {
522        Self::assert_supported_type();
523
524        // Size miniblocks so that they can be efficiently decoded
525        let mini_block_size = match T::T::PHYSICAL_TYPE {
526            Type::INT32 => 32,
527            Type::INT64 => 64,
528            _ => unreachable!(),
529        };
530
531        let num_mini_blocks = DEFAULT_NUM_MINI_BLOCKS;
532        let block_size = mini_block_size * num_mini_blocks;
533        assert_eq!(block_size % 128, 0);
534
535        Self::new_with_layout(block_size, num_mini_blocks)
536    }
537
538    fn new_with_layout(block_size: usize, num_mini_blocks: usize) -> Self {
539        DeltaBitPackEncoder {
540            page_header_writer: BitWriter::new(MAX_PAGE_HEADER_WRITER_SIZE),
541            // Don't pre-allocate: encoders are often created eagerly (e.g. as a
542            // byte array fallback encoder) and may never be used. The buffer
543            // grows on demand and retains its capacity across pages.
544            bit_writer: BitWriter::new(0),
545            total_values: 0,
546            first_value: 0,
547            current_value: 0, // current value to keep adding deltas
548            block_size,       // can write fewer values than block size for last block
549            mini_block_size: block_size / num_mini_blocks,
550            num_mini_blocks,
551            values_in_block: 0, // will be at most block_size
552            deltas: vec![0; block_size],
553            _phantom: PhantomData,
554        }
555    }
556
557    /// Creates a delta bit packed encoder with a custom block layout.
558    pub fn new_with_options(options: DeltaBinaryPackedEncoderOptions) -> Self {
559        Self::assert_supported_type();
560
561        let block_size = options.block_size();
562        let num_mini_blocks = options.mini_blocks_per_block();
563        Self::new_with_layout(block_size, num_mini_blocks)
564    }
565
566    /// Writes page header for blocks, this method is invoked when we are done encoding
567    /// values. It is also okay to encode when no values have been provided
568    fn write_page_header(&mut self) {
569        // We ignore the result of each 'put' operation, because
570        // MAX_PAGE_HEADER_WRITER_SIZE is chosen to fit all header values and
571        // guarantees that writes will not fail.
572
573        // Write the size of each block
574        self.page_header_writer.put_vlq_int(self.block_size as u64);
575        // Write the number of mini blocks
576        self.page_header_writer
577            .put_vlq_int(self.num_mini_blocks as u64);
578        // Write the number of all values (including non-encoded first value)
579        self.page_header_writer
580            .put_vlq_int(self.total_values as u64);
581        // Write first value
582        self.page_header_writer.put_zigzag_vlq_int(self.first_value);
583    }
584
585    // Write current delta buffer (<= 'block size' values) into bit writer
586    #[inline(never)]
587    fn flush_block_values(&mut self) -> Result<()> {
588        if self.values_in_block == 0 {
589            return Ok(());
590        }
591
592        let mut min_delta = i64::MAX;
593        for i in 0..self.values_in_block {
594            min_delta = cmp::min(min_delta, self.deltas[i]);
595        }
596
597        // Write min delta
598        self.bit_writer.put_zigzag_vlq_int(min_delta);
599
600        // Slice to store bit width for each mini block
601        let offset = self.bit_writer.skip(self.num_mini_blocks);
602
603        for i in 0..self.num_mini_blocks {
604            // Find how many values we need to encode - either block size or whatever
605            // values left
606            let n = cmp::min(self.mini_block_size, self.values_in_block);
607            if n == 0 {
608                // Decoders should be agnostic to the padding value, we therefore use 0xFF
609                // when running tests. However, not all implementations may handle this correctly
610                // so pad with 0 when not running tests
611                let pad_value = cfg!(test).then(|| 0xFF).unwrap_or(0);
612                for j in i..self.num_mini_blocks {
613                    self.bit_writer.write_at(offset + j, pad_value);
614                }
615                break;
616            }
617
618            // Compute the max delta in current mini block
619            let mut max_delta = i64::MIN;
620            for j in 0..n {
621                max_delta = cmp::max(max_delta, self.deltas[i * self.mini_block_size + j]);
622            }
623
624            // Compute bit width to store (max_delta - min_delta)
625            let bit_width = num_required_bits(self.subtract_u64(max_delta, min_delta)) as usize;
626            self.bit_writer.write_at(offset + i, bit_width as u8);
627
628            // Encode values in current mini block using min_delta and bit_width. This
629            // mini block's deltas are not read again, so they can be rewritten in place
630            // with the values to pack
631            let start = i * self.mini_block_size;
632            for j in start..start + n {
633                self.deltas[j] = self.subtract_u64(self.deltas[j], min_delta) as i64;
634            }
635            self.bit_writer
636                .put_batch(&self.deltas[start..start + n], bit_width);
637
638            // Pad the last block (n < mini_block_size)
639            for _ in n..self.mini_block_size {
640                self.bit_writer.put_value(0, bit_width);
641            }
642
643            self.values_in_block -= n;
644        }
645
646        assert_eq!(
647            self.values_in_block, 0,
648            "Expected 0 values in block, found {}",
649            self.values_in_block
650        );
651        Ok(())
652    }
653}
654
655// Implementation is shared between Int32Type and Int64Type,
656// see `DeltaBitPackEncoderConversion` below for specifics.
657impl<T: DataType> Encoder<T> for DeltaBitPackEncoder<T> {
658    fn put(&mut self, values: &[T::T]) -> Result<()> {
659        if values.is_empty() {
660            return Ok(());
661        }
662
663        // Define values to encode, initialize state
664        let mut idx = if self.total_values == 0 {
665            self.first_value = self.as_i64(values, 0);
666            self.current_value = self.first_value;
667            1
668        } else {
669            0
670        };
671        // Add all values (including first value)
672        self.total_values += values.len();
673
674        // Write block
675        while idx < values.len() {
676            let value = self.as_i64(values, idx);
677            self.deltas[self.values_in_block] = self.subtract(value, self.current_value);
678            self.current_value = value;
679            idx += 1;
680            self.values_in_block += 1;
681            if self.values_in_block == self.block_size {
682                self.flush_block_values()?;
683            }
684        }
685        Ok(())
686    }
687
688    // Performance Note:
689    // As far as can be seen these functions are rarely called and as such we can hint to the
690    // compiler that they dont need to be folded into hot locations in the final output.
691    #[cold]
692    fn encoding(&self) -> Encoding {
693        Encoding::DELTA_BINARY_PACKED
694    }
695
696    fn estimated_data_encoded_size(&self) -> usize {
697        self.bit_writer.bytes_written()
698    }
699
700    fn flush_buffer(&mut self) -> Result<Bytes> {
701        // Write remaining values
702        self.flush_block_values()?;
703        // Write page header with total values
704        self.write_page_header();
705
706        let mut buffer = Vec::new();
707        buffer.extend_from_slice(self.page_header_writer.flush_buffer());
708        buffer.extend_from_slice(self.bit_writer.flush_buffer());
709
710        // Reset state
711        self.page_header_writer.clear();
712        self.bit_writer.clear();
713        self.total_values = 0;
714        self.first_value = 0;
715        self.current_value = 0;
716        self.values_in_block = 0;
717
718        Ok(buffer.into())
719    }
720
721    /// return the estimated memory size of this encoder.
722    fn estimated_memory_size(&self) -> usize {
723        self.page_header_writer.estimated_memory_size()
724            + self.bit_writer.estimated_memory_size()
725            + self.deltas.capacity() * std::mem::size_of::<i64>()
726            + std::mem::size_of::<Self>()
727    }
728}
729
730/// Helper trait to define specific conversions and subtractions when computing deltas
731trait DeltaBitPackEncoderConversion<T: DataType> {
732    // Method should panic if type is not supported, otherwise no-op
733    fn assert_supported_type();
734
735    fn as_i64(&self, values: &[T::T], index: usize) -> i64;
736
737    fn subtract(&self, left: i64, right: i64) -> i64;
738
739    fn subtract_u64(&self, left: i64, right: i64) -> u64;
740}
741
742const DELTA_BIT_PACK_TYPE_ERROR: &str =
743    "DeltaBitPackDecoder only supports Int32Type, UInt32Type, Int64Type, and UInt64Type";
744
745impl<T: DataType> DeltaBitPackEncoderConversion<T> for DeltaBitPackEncoder<T> {
746    #[inline]
747    fn assert_supported_type() {
748        ensure_phys_ty!(Type::INT32 | Type::INT64, "{}", DELTA_BIT_PACK_TYPE_ERROR);
749    }
750
751    #[inline]
752    fn as_i64(&self, values: &[T::T], index: usize) -> i64 {
753        values[index].as_i64().expect(DELTA_BIT_PACK_TYPE_ERROR)
754    }
755
756    #[inline]
757    fn subtract(&self, left: i64, right: i64) -> i64 {
758        // It is okay for values to overflow, wrapping_sub wrapping around at the boundary
759        match T::get_physical_type() {
760            Type::INT32 => (left as i32).wrapping_sub(right as i32) as i64,
761            Type::INT64 => left.wrapping_sub(right),
762            _ => panic!("{}", DELTA_BIT_PACK_TYPE_ERROR),
763        }
764    }
765
766    #[inline]
767    fn subtract_u64(&self, left: i64, right: i64) -> u64 {
768        match T::get_physical_type() {
769            // Conversion of i32 -> u32 -> u64 is to avoid non-zero left most bytes in int repr
770            Type::INT32 => (left as i32).wrapping_sub(right as i32) as u32 as u64,
771            Type::INT64 => left.wrapping_sub(right) as u64,
772            _ => panic!("{}", DELTA_BIT_PACK_TYPE_ERROR),
773        }
774    }
775}
776
777// ----------------------------------------------------------------------
778// DELTA_LENGTH_BYTE_ARRAY encoding
779
780/// Encoding for byte arrays to separate the length values and the data.
781/// The lengths are encoded using DELTA_BINARY_PACKED encoding, data is
782/// stored as raw bytes.
783pub struct DeltaLengthByteArrayEncoder<T: DataType> {
784    // length encoder
785    len_encoder: DeltaBitPackEncoder<Int32Type>,
786    // byte array data
787    data: Vec<ByteArray>,
788    // data size in bytes of encoded values
789    encoded_size: usize,
790    _phantom: PhantomData<T>,
791}
792
793impl<T: DataType> Default for DeltaLengthByteArrayEncoder<T> {
794    fn default() -> Self {
795        Self::new()
796    }
797}
798
799impl<T: DataType> DeltaLengthByteArrayEncoder<T> {
800    /// Creates new delta length byte array encoder.
801    pub fn new() -> Self {
802        Self {
803            len_encoder: DeltaBitPackEncoder::new(),
804            data: vec![],
805            encoded_size: 0,
806            _phantom: PhantomData,
807        }
808    }
809
810    /// Creates a delta length byte array encoder with a custom layout for its length encoder.
811    pub fn new_with_options(options: DeltaBinaryPackedEncoderOptions) -> Self {
812        Self {
813            len_encoder: DeltaBitPackEncoder::new_with_options(options),
814            data: vec![],
815            encoded_size: 0,
816            _phantom: PhantomData,
817        }
818    }
819}
820
821impl<T: DataType> Encoder<T> for DeltaLengthByteArrayEncoder<T> {
822    fn put(&mut self, values: &[T::T]) -> Result<()> {
823        ensure_phys_ty!(
824            Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY,
825            "DeltaLengthByteArrayEncoder only supports ByteArrayType"
826        );
827
828        let val_it = || {
829            values
830                .iter()
831                .map(|x| x.as_any().downcast_ref::<ByteArray>().unwrap())
832        };
833
834        let lengths: Vec<i32> = val_it().map(|byte_array| byte_array.len() as i32).collect();
835        self.len_encoder.put(&lengths)?;
836        for byte_array in val_it() {
837            self.encoded_size += byte_array.len();
838            self.data.push(byte_array.clone());
839        }
840
841        Ok(())
842    }
843
844    // Performance Note:
845    // As far as can be seen these functions are rarely called and as such we can hint to the
846    // compiler that they dont need to be folded into hot locations in the final output.
847    #[cold]
848    fn encoding(&self) -> Encoding {
849        Encoding::DELTA_LENGTH_BYTE_ARRAY
850    }
851
852    fn estimated_data_encoded_size(&self) -> usize {
853        self.len_encoder.estimated_data_encoded_size() + self.encoded_size
854    }
855
856    fn flush_buffer(&mut self) -> Result<Bytes> {
857        ensure_phys_ty!(
858            Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY,
859            "DeltaLengthByteArrayEncoder only supports ByteArrayType"
860        );
861
862        let mut total_bytes = vec![];
863        let lengths = self.len_encoder.flush_buffer()?;
864        total_bytes.extend_from_slice(&lengths);
865        self.data.iter().for_each(|byte_array| {
866            total_bytes.extend_from_slice(byte_array.data());
867        });
868        self.data.clear();
869        self.encoded_size = 0;
870
871        Ok(total_bytes.into())
872    }
873
874    /// return the estimated memory size of this encoder.
875    fn estimated_memory_size(&self) -> usize {
876        self.len_encoder.estimated_memory_size() + self.data.len() + std::mem::size_of::<Self>()
877    }
878}
879
880// ----------------------------------------------------------------------
881// DELTA_BYTE_ARRAY encoding
882
883/// Encoding for byte arrays, prefix lengths are encoded using DELTA_BINARY_PACKED
884/// encoding, followed by suffixes with DELTA_LENGTH_BYTE_ARRAY encoding.
885pub struct DeltaByteArrayEncoder<T: DataType> {
886    prefix_len_encoder: DeltaBitPackEncoder<Int32Type>,
887    suffix_writer: DeltaLengthByteArrayEncoder<ByteArrayType>,
888    previous: Vec<u8>,
889    _phantom: PhantomData<T>,
890}
891
892impl<T: DataType> Default for DeltaByteArrayEncoder<T> {
893    fn default() -> Self {
894        Self::new()
895    }
896}
897
898impl<T: DataType> DeltaByteArrayEncoder<T> {
899    /// Creates new delta byte array encoder.
900    pub fn new() -> Self {
901        Self {
902            prefix_len_encoder: DeltaBitPackEncoder::new(),
903            suffix_writer: DeltaLengthByteArrayEncoder::new(),
904            previous: vec![],
905            _phantom: PhantomData,
906        }
907    }
908
909    /// Creates a delta byte array encoder with a custom layout for its prefix and suffix length
910    /// encoders.
911    pub fn new_with_options(options: DeltaBinaryPackedEncoderOptions) -> Self {
912        Self {
913            prefix_len_encoder: DeltaBitPackEncoder::new_with_options(options),
914            suffix_writer: DeltaLengthByteArrayEncoder::new_with_options(options),
915            previous: vec![],
916            _phantom: PhantomData,
917        }
918    }
919}
920
921impl<T: DataType> Encoder<T> for DeltaByteArrayEncoder<T> {
922    fn put(&mut self, values: &[T::T]) -> Result<()> {
923        let mut prefix_lengths: Vec<i32> = vec![];
924        let mut suffixes: Vec<ByteArray> = vec![];
925
926        let values = values
927            .iter()
928            .map(|x| x.as_any())
929            .map(|x| match T::get_physical_type() {
930                Type::BYTE_ARRAY => x.downcast_ref::<ByteArray>().unwrap(),
931                Type::FIXED_LEN_BYTE_ARRAY => x.downcast_ref::<FixedLenByteArray>().unwrap(),
932                _ => panic!(
933                    "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType"
934                ),
935            });
936
937        for byte_array in values {
938            let current = byte_array.data();
939            // Number of leading bytes shared with the previous value
940            let match_len = common_prefix_length(&self.previous, current);
941            prefix_lengths.push(match_len as i32);
942            suffixes.push(byte_array.slice(match_len, byte_array.len() - match_len));
943            // Update previous for the next prefix
944            self.previous.clear();
945            self.previous.extend_from_slice(current);
946        }
947        self.prefix_len_encoder.put(&prefix_lengths)?;
948        self.suffix_writer.put(&suffixes)?;
949
950        Ok(())
951    }
952
953    // Performance Note:
954    // As far as can be seen these functions are rarely called and as such we can hint to the
955    // compiler that they dont need to be folded into hot locations in the final output.
956    #[cold]
957    fn encoding(&self) -> Encoding {
958        Encoding::DELTA_BYTE_ARRAY
959    }
960
961    fn estimated_data_encoded_size(&self) -> usize {
962        self.prefix_len_encoder.estimated_data_encoded_size()
963            + self.suffix_writer.estimated_data_encoded_size()
964    }
965
966    fn flush_buffer(&mut self) -> Result<Bytes> {
967        match T::get_physical_type() {
968            Type::BYTE_ARRAY | Type::FIXED_LEN_BYTE_ARRAY => {
969                // TODO: investigate if we can merge lengths and suffixes
970                // without copying data into new vector.
971                let mut total_bytes = vec![];
972                // Insert lengths ...
973                let lengths = self.prefix_len_encoder.flush_buffer()?;
974                total_bytes.extend_from_slice(&lengths);
975                // ... followed by suffixes
976                let suffixes = self.suffix_writer.flush_buffer()?;
977                total_bytes.extend_from_slice(&suffixes);
978
979                self.previous.clear();
980                Ok(total_bytes.into())
981            }
982            _ => panic!(
983                "DeltaByteArrayEncoder only supports ByteArrayType and FixedLenByteArrayType"
984            ),
985        }
986    }
987
988    /// return the estimated memory size of this encoder.
989    fn estimated_memory_size(&self) -> usize {
990        self.prefix_len_encoder.estimated_memory_size()
991            + self.suffix_writer.estimated_memory_size()
992            + (self.previous.capacity() * std::mem::size_of::<u8>())
993    }
994}
995
996#[cfg(test)]
997mod tests {
998    use super::*;
999
1000    use std::sync::Arc;
1001
1002    use crate::encodings::decoding::{Decoder, DictDecoder, PlainDecoder, get_decoder};
1003    use crate::file::properties::WriterProperties;
1004    use crate::schema::types::{ColumnDescPtr, ColumnDescriptor, ColumnPath, Type as SchemaType};
1005    use crate::util::bit_util;
1006    use crate::util::test_common::rand_gen::{RandGen, random_bytes};
1007
1008    const TEST_SET_SIZE: usize = 1024;
1009
1010    #[test]
1011    fn test_delta_binary_packed_encoder_options() {
1012        let options = DeltaBinaryPackedEncoderOptions::try_new(256, 4).unwrap();
1013        assert_eq!(options.block_size(), 256);
1014        assert_eq!(options.mini_blocks_per_block(), 4);
1015
1016        let encoder = DeltaBitPackEncoder::<Int32Type>::new_with_options(options);
1017        assert_eq!(encoder.block_size, 256);
1018        assert_eq!(encoder.mini_block_size, 64);
1019        assert_eq!(encoder.num_mini_blocks, 4);
1020
1021        let props = WriterProperties::builder()
1022            .set_delta_binary_packed_encoder_options(options)
1023            .build();
1024        let column_props = props.resolve_column_properties(&ColumnPath::new(vec![]));
1025        let desc = create_test_col_desc_ptr(-1, Type::INT32);
1026        let mut encoder = get_encoder_with_properties::<Int32Type>(
1027            Encoding::DELTA_BINARY_PACKED,
1028            &desc,
1029            &column_props,
1030        )
1031        .unwrap();
1032        let input: Vec<i32> = (0..600).map(|value| value * value).collect();
1033        encoder.put(&input).unwrap();
1034        let encoded = encoder.flush_buffer().unwrap();
1035        let mut header = bit_util::BitReader::new(encoded.clone());
1036        assert_eq!(header.get_vlq_int(), Some(256));
1037        assert_eq!(header.get_vlq_int(), Some(4));
1038
1039        let mut decoder = create_test_decoder::<Int32Type>(-1, Encoding::DELTA_BINARY_PACKED);
1040        decoder.set_data(encoded, input.len()).unwrap();
1041        let mut output = vec![0; input.len()];
1042        assert_eq!(decoder.get(&mut output).unwrap(), input.len());
1043        assert_eq!(output, input);
1044
1045        let encoder = DeltaLengthByteArrayEncoder::<ByteArrayType>::new_with_options(options);
1046        assert_eq!(encoder.len_encoder.block_size, 256);
1047        let encoder = DeltaByteArrayEncoder::<ByteArrayType>::new_with_options(options);
1048        assert_eq!(encoder.prefix_len_encoder.block_size, 256);
1049        assert_eq!(encoder.suffix_writer.len_encoder.block_size, 256);
1050
1051        let input: Vec<ByteArray> = (0..600)
1052            .map(|value| ByteArray::from(format!("prefix-{value:04}-suffix").into_bytes()))
1053            .collect();
1054        let desc = create_test_col_desc_ptr(-1, Type::BYTE_ARRAY);
1055        for encoding in [
1056            Encoding::DELTA_LENGTH_BYTE_ARRAY,
1057            Encoding::DELTA_BYTE_ARRAY,
1058        ] {
1059            let mut encoder =
1060                get_encoder_with_properties::<ByteArrayType>(encoding, &desc, &column_props)
1061                    .unwrap();
1062            encoder.put(&input).unwrap();
1063            let encoded = encoder.flush_buffer().unwrap();
1064            let mut header = bit_util::BitReader::new(encoded.clone());
1065            assert_eq!(header.get_vlq_int(), Some(256));
1066            assert_eq!(header.get_vlq_int(), Some(4));
1067
1068            let mut decoder = create_test_decoder::<ByteArrayType>(-1, encoding);
1069            decoder.set_data(encoded, input.len()).unwrap();
1070            let mut output = vec![ByteArray::default(); input.len()];
1071            assert_eq!(decoder.get(&mut output).unwrap(), input.len());
1072            assert_eq!(output, input);
1073        }
1074    }
1075
1076    #[test]
1077    fn test_delta_binary_packed_encoder_options_validation() {
1078        for (block_size, mini_blocks_per_block, message) in [
1079            (0, 4, "positive multiple of 128"),
1080            (127, 4, "positive multiple of 128"),
1081            (128, 0, "positive divisor"),
1082            (128, 3, "positive divisor"),
1083            (256, 16, "must be a multiple of 32"),
1084        ] {
1085            let error = DeltaBinaryPackedEncoderOptions::try_new(block_size, mini_blocks_per_block)
1086                .unwrap_err();
1087            assert!(error.to_string().contains(message), "{error}");
1088        }
1089    }
1090
1091    #[test]
1092    fn test_get_encoders_with_properties() {
1093        fn check<T: DataType>(
1094            encoding: Encoding,
1095            values: &[T::T],
1096            props: &ResolvedColumnProperties,
1097        ) {
1098            let desc = create_test_col_desc_ptr(-1, T::get_physical_type());
1099            let mut default = get_encoder::<T>(encoding, &desc).unwrap();
1100            let mut resolved = get_encoder_with_properties::<T>(encoding, &desc, props).unwrap();
1101            default.put(values).unwrap();
1102            resolved.put(values).unwrap();
1103            assert_eq!(resolved.encoding(), encoding);
1104            assert_eq!(
1105                resolved.flush_buffer().unwrap(),
1106                default.flush_buffer().unwrap()
1107            );
1108        }
1109
1110        let path = ColumnPath::from("col");
1111        let props = WriterProperties::builder().build();
1112        let defaults = props.resolve_column_properties(&path);
1113        let ints: Vec<i32> = (0..600).map(|value| value * value).collect();
1114        check::<Int32Type>(Encoding::DELTA_BINARY_PACKED, &ints, &defaults);
1115        let longs: Vec<i64> = ints.iter().map(|&value| i64::from(value)).collect();
1116        check::<Int64Type>(Encoding::DELTA_BINARY_PACKED, &longs, &defaults);
1117        let bytes: Vec<ByteArray> = ints
1118            .iter()
1119            .map(|value| ByteArray::from(format!("prefix-{value}").into_bytes()))
1120            .collect();
1121        for encoding in [
1122            Encoding::DELTA_LENGTH_BYTE_ARRAY,
1123            Encoding::DELTA_BYTE_ARRAY,
1124        ] {
1125            check::<ByteArrayType>(encoding, &bytes, &defaults);
1126        }
1127
1128        let props = WriterProperties::builder()
1129            .set_delta_binary_packed_encoder_options(
1130                DeltaBinaryPackedEncoderOptions::try_new(256, 4).unwrap(),
1131            )
1132            .build();
1133        let custom = props.resolve_column_properties(&path);
1134        for props in [&defaults, &custom] {
1135            check::<Int32Type>(Encoding::PLAIN, &ints, props);
1136            for encoding in [Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT, Encoding::ALP] {
1137                check::<FloatType>(encoding, &[1.25, -2.5, 3.75], props);
1138                check::<DoubleType>(encoding, &[1.25, -2.5, 3.75], props);
1139            }
1140        }
1141    }
1142
1143    #[test]
1144    fn test_get_encoders() {
1145        // supported encodings
1146        create_and_check_encoder::<Int32Type>(0, Encoding::PLAIN, None);
1147        create_and_check_encoder::<Int32Type>(0, Encoding::DELTA_BINARY_PACKED, None);
1148        create_and_check_encoder::<BoolType>(0, Encoding::RLE, None);
1149
1150        // error when initializing
1151        create_and_check_encoder::<Int32Type>(
1152            0,
1153            Encoding::RLE_DICTIONARY,
1154            Some(general_err!(
1155                "Cannot initialize this encoding through this function"
1156            )),
1157        );
1158        create_and_check_encoder::<Int32Type>(
1159            0,
1160            Encoding::PLAIN_DICTIONARY,
1161            Some(general_err!(
1162                "Cannot initialize this encoding through this function"
1163            )),
1164        );
1165        create_and_check_encoder::<Int32Type>(
1166            0,
1167            Encoding::DELTA_LENGTH_BYTE_ARRAY,
1168            Some(unsupported_column_encoding(
1169                Encoding::DELTA_LENGTH_BYTE_ARRAY,
1170                Type::INT32,
1171            )),
1172        );
1173        create_and_check_encoder::<Int32Type>(
1174            0,
1175            Encoding::DELTA_BYTE_ARRAY,
1176            Some(unsupported_column_encoding(
1177                Encoding::DELTA_BYTE_ARRAY,
1178                Type::INT32,
1179            )),
1180        );
1181        create_and_check_encoder::<Int32Type>(
1182            0,
1183            Encoding::ALP,
1184            Some(general_err!(
1185                "Encoding ALP only supports FLOAT and DOUBLE, got INT32"
1186            )),
1187        );
1188
1189        // unsupported
1190        #[expect(deprecated)]
1191        create_and_check_encoder::<Int32Type>(
1192            0,
1193            Encoding::BIT_PACKED,
1194            Some(nyi_err!("Encoding BIT_PACKED is not supported")),
1195        );
1196    }
1197
1198    #[test]
1199    fn column_construction_validates_encoding_fallback_matrix() {
1200        use crate::file::properties::{WriterProperties, WriterVersion};
1201        use crate::file::writer::SerializedFileWriter;
1202        use std::panic::{AssertUnwindSafe, catch_unwind};
1203        use std::sync::Arc;
1204
1205        fn check_encoder<T: DataType>(encoding: Encoding, supported: bool) {
1206            let error = if supported {
1207                None
1208            } else if encoding == Encoding::ALP {
1209                Some(general_err!(
1210                    "Encoding ALP only supports FLOAT and DOUBLE, got {}",
1211                    T::get_physical_type()
1212                ))
1213            } else {
1214                Some(unsupported_column_encoding(
1215                    encoding,
1216                    T::get_physical_type(),
1217                ))
1218            };
1219            create_and_check_encoder::<T>(4, encoding, error);
1220        }
1221
1222        let encodings = [
1223            Encoding::PLAIN,
1224            Encoding::RLE,
1225            Encoding::DELTA_BINARY_PACKED,
1226            Encoding::DELTA_LENGTH_BYTE_ARRAY,
1227            Encoding::DELTA_BYTE_ARRAY,
1228            Encoding::BYTE_STREAM_SPLIT,
1229            Encoding::ALP,
1230        ];
1231        for physical in [
1232            Type::BOOLEAN,
1233            Type::INT32,
1234            Type::INT64,
1235            Type::INT96,
1236            Type::FLOAT,
1237            Type::DOUBLE,
1238            Type::BYTE_ARRAY,
1239            Type::FIXED_LEN_BYTE_ARRAY,
1240        ] {
1241            let valid: &[Encoding] = match physical {
1242                Type::BOOLEAN => &[Encoding::PLAIN, Encoding::RLE],
1243                Type::INT32 | Type::INT64 => &[
1244                    Encoding::PLAIN,
1245                    Encoding::DELTA_BINARY_PACKED,
1246                    Encoding::BYTE_STREAM_SPLIT,
1247                ],
1248                Type::INT96 => &[Encoding::PLAIN],
1249                Type::FLOAT | Type::DOUBLE => {
1250                    &[Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT, Encoding::ALP]
1251                }
1252                Type::BYTE_ARRAY => &[
1253                    Encoding::PLAIN,
1254                    Encoding::DELTA_LENGTH_BYTE_ARRAY,
1255                    Encoding::DELTA_BYTE_ARRAY,
1256                ],
1257                Type::FIXED_LEN_BYTE_ARRAY => &[
1258                    Encoding::PLAIN,
1259                    Encoding::DELTA_BYTE_ARRAY,
1260                    Encoding::BYTE_STREAM_SPLIT,
1261                ],
1262            };
1263            // Check the factory's returned errors directly, independently of the
1264            // infallible column writer constructor that unwraps those errors.
1265            for encoding in encodings {
1266                let supported = valid.contains(&encoding);
1267                match physical {
1268                    Type::BOOLEAN => check_encoder::<BoolType>(encoding, supported),
1269                    Type::INT32 => check_encoder::<Int32Type>(encoding, supported),
1270                    Type::INT64 => check_encoder::<Int64Type>(encoding, supported),
1271                    Type::INT96 => check_encoder::<Int96Type>(encoding, supported),
1272                    Type::FLOAT => check_encoder::<FloatType>(encoding, supported),
1273                    Type::DOUBLE => check_encoder::<DoubleType>(encoding, supported),
1274                    Type::BYTE_ARRAY => check_encoder::<ByteArrayType>(encoding, supported),
1275                    Type::FIXED_LEN_BYTE_ARRAY => {
1276                        check_encoder::<FixedLenByteArrayType>(encoding, supported)
1277                    }
1278                }
1279            }
1280            for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
1281                for dictionary in [false, true] {
1282                    for encoding in encodings {
1283                        let field = Arc::new(
1284                            SchemaType::primitive_type_builder("a", physical)
1285                                .with_length(4)
1286                                .build()
1287                                .unwrap(),
1288                        );
1289                        let schema = Arc::new(
1290                            SchemaType::group_type_builder("schema")
1291                                .with_fields(vec![field])
1292                                .build()
1293                                .unwrap(),
1294                        );
1295                        let props = Arc::new(
1296                            WriterProperties::builder()
1297                                .set_writer_version(version)
1298                                .set_dictionary_enabled(dictionary)
1299                                .set_encoding(encoding)
1300                                .build(),
1301                        );
1302                        let created = catch_unwind(AssertUnwindSafe(|| {
1303                            let mut writer =
1304                                SerializedFileWriter::new(Vec::new(), schema, props).unwrap();
1305                            let mut group = writer.next_row_group().unwrap();
1306                            group.next_column().unwrap().unwrap();
1307                        }));
1308                        assert_eq!(
1309                            created.is_ok(),
1310                            valid.contains(&encoding),
1311                            "{physical:?} {version:?} dictionary={dictionary} {encoding:?}"
1312                        );
1313                    }
1314                }
1315            }
1316        }
1317    }
1318
1319    #[test]
1320    fn test_bool() {
1321        BoolType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1322        BoolType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1323        BoolType::test(Encoding::RLE, TEST_SET_SIZE, -1);
1324    }
1325
1326    #[test]
1327    #[cfg_attr(miri, ignore)] // Takes too long
1328    fn test_i32() {
1329        Int32Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1330        Int32Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1331        Int32Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1);
1332        Int32Type::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1);
1333    }
1334
1335    #[test]
1336    #[cfg_attr(miri, ignore)] // Takes too long
1337    fn test_i64() {
1338        Int64Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1339        Int64Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1340        Int64Type::test(Encoding::DELTA_BINARY_PACKED, TEST_SET_SIZE, -1);
1341        Int64Type::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1);
1342    }
1343
1344    #[test]
1345    #[cfg_attr(miri, ignore)] // Takes too long
1346    fn test_i96() {
1347        Int96Type::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1348        Int96Type::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1349    }
1350
1351    #[test]
1352    fn test_float() {
1353        FloatType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1354        FloatType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1355        FloatType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1);
1356    }
1357
1358    #[test]
1359    #[cfg_attr(miri, ignore)] // Takes too long
1360    fn test_double() {
1361        DoubleType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1362        DoubleType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1363        DoubleType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, -1);
1364    }
1365
1366    #[test]
1367    #[cfg_attr(miri, ignore)] // Takes too long
1368    fn test_byte_array() {
1369        ByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, -1);
1370        ByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1);
1371        ByteArrayType::test(Encoding::DELTA_LENGTH_BYTE_ARRAY, TEST_SET_SIZE, -1);
1372        ByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, -1);
1373    }
1374
1375    #[test]
1376    #[cfg_attr(miri, ignore)] // Takes too long
1377    fn test_fixed_len_byte_array() {
1378        FixedLenByteArrayType::test(Encoding::PLAIN, TEST_SET_SIZE, 100);
1379        FixedLenByteArrayType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, 100);
1380        FixedLenByteArrayType::test(Encoding::DELTA_BYTE_ARRAY, TEST_SET_SIZE, 100);
1381        FixedLenByteArrayType::test(Encoding::BYTE_STREAM_SPLIT, TEST_SET_SIZE, 100);
1382    }
1383
1384    #[test]
1385    fn test_dict_encoded_size() {
1386        fn run_test<T: DataType>(type_length: i32, values: &[T::T], expected_size: usize) {
1387            let mut encoder = create_test_dict_encoder::<T>(type_length);
1388            assert_eq!(encoder.dict_encoded_size(), 0);
1389            encoder.put(values).unwrap();
1390            assert_eq!(encoder.dict_encoded_size(), expected_size);
1391            // We do not reset encoded size of the dictionary keys after flush_buffer
1392            encoder.flush_buffer().unwrap();
1393            assert_eq!(encoder.dict_encoded_size(), expected_size);
1394        }
1395
1396        // Only 2 variations of values 1 byte each
1397        run_test::<BoolType>(-1, &[true, false, true, false, true], 2);
1398        run_test::<Int32Type>(-1, &[1i32, 2i32, 3i32, 4i32, 5i32], 20);
1399        run_test::<Int64Type>(-1, &[1i64, 2i64, 3i64, 4i64, 5i64], 40);
1400        run_test::<FloatType>(-1, &[1f32, 2f32, 3f32, 4f32, 5f32], 20);
1401        run_test::<DoubleType>(-1, &[1f64, 2f64, 3f64, 4f64, 5f64], 40);
1402        // Int96: len + reference
1403        run_test::<Int96Type>(
1404            -1,
1405            &[Int96::from(vec![1, 2, 3]), Int96::from(vec![2, 3, 4])],
1406            24,
1407        );
1408        run_test::<ByteArrayType>(-1, &[ByteArray::from("abcd"), ByteArray::from("efj")], 15);
1409        run_test::<FixedLenByteArrayType>(
1410            2,
1411            &[ByteArray::from("ab").into(), ByteArray::from("bc").into()],
1412            4,
1413        );
1414    }
1415
1416    #[test]
1417    fn test_estimated_data_encoded_size() {
1418        fn run_test<T: DataType>(
1419            encoding: Encoding,
1420            type_length: i32,
1421            values: &[T::T],
1422            initial_size: usize,
1423            max_size: usize,
1424            flush_size: usize,
1425        ) {
1426            let mut encoder = match encoding {
1427                Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => {
1428                    Box::new(create_test_dict_encoder::<T>(type_length))
1429                }
1430                _ => create_test_encoder::<T>(type_length, encoding),
1431            };
1432            assert_eq!(encoder.estimated_data_encoded_size(), initial_size);
1433
1434            encoder.put(values).unwrap();
1435            assert_eq!(encoder.estimated_data_encoded_size(), max_size);
1436
1437            encoder.flush_buffer().unwrap();
1438            assert_eq!(encoder.estimated_data_encoded_size(), flush_size);
1439        }
1440
1441        // PLAIN
1442        run_test::<Int32Type>(Encoding::PLAIN, -1, &[123; 1024], 0, 4096, 0);
1443
1444        // DICTIONARY
1445        // NOTE: The final size is almost the same because the dictionary entries are
1446        // preserved after encoded values have been written.
1447        run_test::<Int32Type>(Encoding::RLE_DICTIONARY, -1, &[123, 1024], 0, 2, 0);
1448
1449        // DELTA_BINARY_PACKED
1450        run_test::<Int32Type>(Encoding::DELTA_BINARY_PACKED, -1, &[123; 1024], 0, 35, 0);
1451
1452        // RLE
1453        let mut values = vec![];
1454        values.extend_from_slice(&[true; 16]);
1455        values.extend_from_slice(&[false; 16]);
1456        run_test::<BoolType>(Encoding::RLE, -1, &values, 0, 6, 0);
1457
1458        // DELTA_LENGTH_BYTE_ARRAY
1459        run_test::<ByteArrayType>(
1460            Encoding::DELTA_LENGTH_BYTE_ARRAY,
1461            -1,
1462            &[ByteArray::from("ab"), ByteArray::from("abc")],
1463            0,
1464            5, // only value bytes, length encoder is not flushed yet
1465            0,
1466        );
1467
1468        // DELTA_BYTE_ARRAY
1469        run_test::<ByteArrayType>(
1470            Encoding::DELTA_BYTE_ARRAY,
1471            -1,
1472            &[ByteArray::from("ab"), ByteArray::from("abc")],
1473            0,
1474            3, // only suffix bytes, length encoder is not flushed yet
1475            0,
1476        );
1477
1478        // BYTE_STREAM_SPLIT
1479        run_test::<FloatType>(Encoding::BYTE_STREAM_SPLIT, -1, &[0.1, 0.2], 0, 8, 0);
1480    }
1481
1482    #[test]
1483    fn test_byte_stream_split_example_f32() {
1484        // Test data from https://github.com/apache/parquet-format/blob/2a481fe1aad64ff770e21734533bb7ef5a057dac/Encodings.md#byte-stream-split-byte_stream_split--9
1485        let mut encoder = create_test_encoder::<FloatType>(0, Encoding::BYTE_STREAM_SPLIT);
1486        let mut decoder = create_test_decoder::<FloatType>(0, Encoding::BYTE_STREAM_SPLIT);
1487
1488        let input = vec![
1489            f32::from_le_bytes([0xAA, 0xBB, 0xCC, 0xDD]),
1490            f32::from_le_bytes([0x00, 0x11, 0x22, 0x33]),
1491            f32::from_le_bytes([0xA3, 0xB4, 0xC5, 0xD6]),
1492        ];
1493
1494        encoder.put(&input).unwrap();
1495        let encoded = encoder.flush_buffer().unwrap();
1496
1497        assert_eq!(
1498            encoded,
1499            Bytes::from(vec![
1500                0xAA_u8, 0x00, 0xA3, 0xBB, 0x11, 0xB4, 0xCC, 0x22, 0xC5, 0xDD, 0x33, 0xD6
1501            ])
1502        );
1503
1504        let mut decoded = vec![0.0; input.len()];
1505        decoder.set_data(encoded, input.len()).unwrap();
1506        decoder.get(&mut decoded).unwrap();
1507
1508        assert_eq!(decoded, input);
1509    }
1510
1511    // See: https://github.com/sunchao/parquet-rs/issues/47
1512    #[test]
1513    fn test_issue_47() {
1514        let mut encoder = create_test_encoder::<ByteArrayType>(0, Encoding::DELTA_BYTE_ARRAY);
1515        let mut decoder = create_test_decoder::<ByteArrayType>(0, Encoding::DELTA_BYTE_ARRAY);
1516
1517        let input = vec![
1518            ByteArray::from("aa"),
1519            ByteArray::from("aaa"),
1520            ByteArray::from("aa"),
1521            ByteArray::from("aaa"),
1522        ];
1523
1524        let mut output = vec![ByteArray::default(); input.len()];
1525
1526        let mut result = put_and_get(&mut encoder, &mut decoder, &input[..2], &mut output[..2]);
1527        assert!(
1528            result.is_ok(),
1529            "first put_and_get() failed with: {}",
1530            result.unwrap_err()
1531        );
1532        result = put_and_get(&mut encoder, &mut decoder, &input[2..], &mut output[2..]);
1533        assert!(
1534            result.is_ok(),
1535            "second put_and_get() failed with: {}",
1536            result.unwrap_err()
1537        );
1538        assert_eq!(output, input);
1539    }
1540
1541    trait EncodingTester<T: DataType> {
1542        fn test(enc: Encoding, total: usize, type_length: i32) {
1543            let result = match enc {
1544                Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY => {
1545                    Self::test_dict_internal(total, type_length)
1546                }
1547                enc => Self::test_internal(enc, total, type_length),
1548            };
1549
1550            assert!(
1551                result.is_ok(),
1552                "Expected result to be OK but got err:\n {}",
1553                result.unwrap_err()
1554            );
1555        }
1556
1557        fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()>;
1558
1559        fn test_dict_internal(total: usize, type_length: i32) -> Result<()>;
1560    }
1561
1562    impl<T: DataType + RandGen<T>> EncodingTester<T> for T {
1563        fn test_internal(enc: Encoding, total: usize, type_length: i32) -> Result<()> {
1564            let mut encoder = create_test_encoder::<T>(type_length, enc);
1565            let mut decoder = create_test_decoder::<T>(type_length, enc);
1566            let mut values = <T as RandGen<T>>::gen_vec(type_length, total);
1567            let mut result_data = vec![T::T::default(); total];
1568
1569            // Test put/get spaced.
1570            let num_bytes = bit_util::ceil(total as i64, 8);
1571            let valid_bits = random_bytes(num_bytes as usize);
1572            let values_written = encoder.put_spaced(&values[..], &valid_bits[..])?;
1573            let data = encoder.flush_buffer()?;
1574            decoder.set_data(data, values_written)?;
1575            let _ = decoder.get_spaced(
1576                &mut result_data[..],
1577                values.len() - values_written,
1578                &valid_bits[..],
1579            )?;
1580
1581            // Check equality
1582            for i in 0..total {
1583                if bit_util::get_bit(&valid_bits[..], i) {
1584                    assert_eq!(result_data[i], values[i]);
1585                } else {
1586                    assert_eq!(result_data[i], T::T::default());
1587                }
1588            }
1589
1590            let mut actual_total = put_and_get(
1591                &mut encoder,
1592                &mut decoder,
1593                &values[..],
1594                &mut result_data[..],
1595            )?;
1596            assert_eq!(actual_total, total);
1597            assert_eq!(result_data, values);
1598
1599            // Encode more data after flush and test with decoder
1600
1601            values = <T as RandGen<T>>::gen_vec(type_length, total);
1602            actual_total = put_and_get(
1603                &mut encoder,
1604                &mut decoder,
1605                &values[..],
1606                &mut result_data[..],
1607            )?;
1608            assert_eq!(actual_total, total);
1609            assert_eq!(result_data, values);
1610
1611            Ok(())
1612        }
1613
1614        fn test_dict_internal(total: usize, type_length: i32) -> Result<()> {
1615            let mut encoder = create_test_dict_encoder::<T>(type_length);
1616            let mut values = <T as RandGen<T>>::gen_vec(type_length, total);
1617            encoder.put(&values[..])?;
1618
1619            let mut data = encoder.flush_buffer()?;
1620            let mut decoder = create_test_dict_decoder::<T>();
1621            let mut dict_decoder = PlainDecoder::<T>::new(type_length);
1622            dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?;
1623            decoder.set_dict(Box::new(dict_decoder))?;
1624            let mut result_data = vec![T::T::default(); total];
1625            decoder.set_data(data, total)?;
1626            let mut actual_total = decoder.get(&mut result_data)?;
1627
1628            assert_eq!(actual_total, total);
1629            assert_eq!(result_data, values);
1630
1631            // Encode more data after flush and test with decoder
1632
1633            values = <T as RandGen<T>>::gen_vec(type_length, total);
1634            encoder.put(&values[..])?;
1635            data = encoder.flush_buffer()?;
1636
1637            let mut dict_decoder = PlainDecoder::<T>::new(type_length);
1638            dict_decoder.set_data(encoder.write_dict()?, encoder.num_entries())?;
1639            decoder.set_dict(Box::new(dict_decoder))?;
1640            decoder.set_data(data, total)?;
1641            actual_total = decoder.get(&mut result_data)?;
1642
1643            assert_eq!(actual_total, total);
1644            assert_eq!(result_data, values);
1645
1646            Ok(())
1647        }
1648    }
1649
1650    fn put_and_get<T: DataType>(
1651        encoder: &mut Box<dyn Encoder<T>>,
1652        decoder: &mut Box<dyn Decoder<T>>,
1653        input: &[T::T],
1654        output: &mut [T::T],
1655    ) -> Result<usize> {
1656        encoder.put(input)?;
1657        let data = encoder.flush_buffer()?;
1658        decoder.set_data(data, input.len())?;
1659        decoder.get(output)
1660    }
1661
1662    fn create_and_check_encoder<T: DataType>(
1663        type_length: i32,
1664        encoding: Encoding,
1665        err: Option<ParquetError>,
1666    ) {
1667        let desc = create_test_col_desc_ptr(type_length, T::get_physical_type());
1668        let defaults = WriterProperties::builder().build();
1669        let custom = WriterProperties::builder()
1670            .set_delta_binary_packed_encoder_options(
1671                DeltaBinaryPackedEncoderOptions::try_new(256, 4).unwrap(),
1672            )
1673            .build();
1674        let path = ColumnPath::from("t");
1675        for encoder in [
1676            get_encoder::<T>(encoding, &desc),
1677            get_encoder_with_properties::<T>(
1678                encoding,
1679                &desc,
1680                &defaults.resolve_column_properties(&path),
1681            ),
1682            get_encoder_with_properties::<T>(
1683                encoding,
1684                &desc,
1685                &custom.resolve_column_properties(&path),
1686            ),
1687        ] {
1688            match &err {
1689                Some(parquet_error) => {
1690                    assert_eq!(
1691                        encoder.err().unwrap().to_string(),
1692                        parquet_error.to_string()
1693                    )
1694                }
1695                None => assert_eq!(encoder.unwrap().encoding(), encoding),
1696            }
1697        }
1698    }
1699
1700    // Creates test column descriptor.
1701    fn create_test_col_desc_ptr(type_len: i32, t: Type) -> ColumnDescPtr {
1702        let ty = SchemaType::primitive_type_builder("t", t)
1703            .with_length(type_len)
1704            .build()
1705            .unwrap();
1706        Arc::new(ColumnDescriptor::new(
1707            Arc::new(ty),
1708            0,
1709            0,
1710            ColumnPath::new(vec![]),
1711        ))
1712    }
1713
1714    fn create_test_encoder<T: DataType>(type_len: i32, enc: Encoding) -> Box<dyn Encoder<T>> {
1715        let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
1716        get_encoder(enc, &desc).unwrap()
1717    }
1718
1719    fn create_test_decoder<T: DataType>(type_len: i32, enc: Encoding) -> Box<dyn Decoder<T>> {
1720        let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
1721        get_decoder(desc, enc).unwrap()
1722    }
1723
1724    fn create_test_dict_encoder<T: DataType>(type_len: i32) -> DictEncoder<T> {
1725        let desc = create_test_col_desc_ptr(type_len, T::get_physical_type());
1726        DictEncoder::<T>::new(desc)
1727    }
1728
1729    fn create_test_dict_decoder<T: DataType>() -> DictDecoder<T> {
1730        DictDecoder::<T>::new()
1731    }
1732}
parquet/encodings/encoding/mod.rs

parquet/encodings/encoding/
mod.rs
