Skip to main content

parquet/
compression.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 codec interface and supported codec implementations.
19//!
20//! See [`Compression`](crate::basic::Compression) enum for all available compression
21//! algorithms.
22//!
23// NOTE: this must be an inner attribute so that the example is attached to the module (and
24// therefore actually run as a doc test) rather than to the `use` statement below.
25#![cfg_attr(
26    feature = "experimental",
27    doc = r"
28# Example
29
30```no_run
31use parquet::{basic::Compression, compression::{create_codec, CodecOptionsBuilder}};
32
33let codec_options = CodecOptionsBuilder::default()
34    .set_backward_compatible_lz4(false)
35    .build();
36let mut codec = match create_codec(Compression::SNAPPY, &codec_options) {
37 Ok(Some(codec)) => codec,
38 _ => panic!(),
39};
40
41let data = vec![b'p', b'a', b'r', b'q', b'u', b'e', b't'];
42let mut compressed = vec![];
43codec.compress(&data[..], &mut compressed).unwrap();
44
45let mut output = vec![];
46codec.decompress(&compressed[..], &mut output, None).unwrap();
47
48assert_eq!(output, data);
49```
50"
51)]
52use crate::basic::Compression as CodecType;
53use crate::errors::{ParquetError, Result};
54
55/// Parquet compression codec interface.
56pub trait Codec: Send {
57    /// Compresses data stored in slice `input_buf` and appends the compressed result
58    /// to `output_buf`.
59    ///
60    /// Note that you'll need to call `clear()` before reusing the same `output_buf`
61    /// across different `compress` calls.
62    fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()>;
63
64    /// Decompresses data stored in slice `input_buf` and appends output to `output_buf`.
65    ///
66    /// If the uncompress_size is provided it will allocate the exact amount of memory.
67    /// Otherwise, it will estimate the uncompressed size, allocating an amount of memory
68    /// greater or equal to the real uncompress_size.
69    ///
70    /// Returns the total number of bytes written.
71    fn decompress(
72        &mut self,
73        input_buf: &[u8],
74        output_buf: &mut Vec<u8>,
75        uncompress_size: Option<usize>,
76    ) -> Result<usize>;
77}
78
79/// Struct to hold `Codec` creation options.
80#[derive(Debug, PartialEq, Eq)]
81pub struct CodecOptions {
82    /// Whether or not to fallback to other LZ4 older implementations on error in LZ4_HADOOP.
83    backward_compatible_lz4: bool,
84}
85
86impl Default for CodecOptions {
87    fn default() -> Self {
88        CodecOptionsBuilder::default().build()
89    }
90}
91
92pub struct CodecOptionsBuilder {
93    /// Whether or not to fallback to other LZ4 older implementations on error in LZ4_HADOOP.
94    backward_compatible_lz4: bool,
95}
96
97impl Default for CodecOptionsBuilder {
98    fn default() -> Self {
99        Self {
100            backward_compatible_lz4: true,
101        }
102    }
103}
104
105impl CodecOptionsBuilder {
106    /// Enable/disable backward compatible LZ4.
107    ///
108    /// If backward compatible LZ4 is enable, on LZ4_HADOOP error it will fallback
109    /// to the older versions LZ4 algorithms. That is LZ4_FRAME, for backward compatibility
110    /// with files generated by older versions of this library, and LZ4_RAW, for backward
111    /// compatibility with files generated by older versions of parquet-cpp.
112    ///
113    /// If backward compatible LZ4 is disabled, on LZ4_HADOOP error it will return the error.
114    pub fn set_backward_compatible_lz4(mut self, value: bool) -> CodecOptionsBuilder {
115        self.backward_compatible_lz4 = value;
116        self
117    }
118
119    pub fn build(self) -> CodecOptions {
120        CodecOptions {
121            backward_compatible_lz4: self.backward_compatible_lz4,
122        }
123    }
124}
125
126/// Defines valid compression levels.
127pub(crate) trait CompressionLevel<T: std::fmt::Display + std::cmp::PartialOrd> {
128    const MINIMUM_LEVEL: T;
129    const MAXIMUM_LEVEL: T;
130
131    /// Tests if the provided compression level is valid.
132    fn is_valid_level(level: T) -> Result<()> {
133        let compression_range = Self::MINIMUM_LEVEL..=Self::MAXIMUM_LEVEL;
134        if compression_range.contains(&level) {
135            Ok(())
136        } else {
137            Err(ParquetError::General(format!(
138                "valid compression range {}..={} exceeded.",
139                compression_range.start(),
140                compression_range.end()
141            )))
142        }
143    }
144}
145
146/// Given the compression type `codec`, returns a codec used to compress and decompress
147/// bytes for the compression type.
148/// This returns `None` if the codec type is `UNCOMPRESSED`.
149#[cfg_attr(
150    any(feature = "lz4", test),
151    expect(
152        clippy::used_underscore_binding,
153        reason = "`_options` is only read when the `lz4` feature is on"
154    )
155)]
156pub fn create_codec(codec: CodecType, _options: &CodecOptions) -> Result<Option<Box<dyn Codec>>> {
157    #[cfg_attr(
158        any(
159            test,
160            feature = "brotli",
161            feature = "flate2",
162            feature = "lz4",
163            feature = "snap",
164            feature = "zstd"
165        ),
166        expect(unreachable_code)
167    )]
168    #[cfg_attr(
169        all(
170            not(test),
171            not(all(feature = "brotli", feature = "flate2", feature = "zstd"))
172        ),
173        expect(unused_variables)
174    )]
175    match codec {
176        CodecType::BROTLI(level) => {
177            #[cfg(any(feature = "brotli", test))]
178            return Ok(Some(Box::new(BrotliCodec::new(level))));
179            Err(ParquetError::General(
180                "Disabled feature at compile time: brotli".into(),
181            ))
182        }
183        CodecType::GZIP(level) => {
184            #[cfg(any(feature = "flate2", test))]
185            return Ok(Some(Box::new(GZipCodec::new(level))));
186            Err(ParquetError::General(
187                "Disabled feature at compile time: flate2".into(),
188            ))
189        }
190        CodecType::SNAPPY => {
191            #[cfg(any(feature = "snap", test))]
192            return Ok(Some(Box::new(SnappyCodec::new())));
193            Err(ParquetError::General(
194                "Disabled feature at compile time: snap".into(),
195            ))
196        }
197        CodecType::LZ4 => {
198            #[cfg(any(feature = "lz4", test))]
199            return Ok(Some(Box::new(LZ4HadoopCodec::new(
200                _options.backward_compatible_lz4,
201            ))));
202            Err(ParquetError::General(
203                "Disabled feature at compile time: lz4".into(),
204            ))
205        }
206        CodecType::ZSTD(level) => {
207            #[cfg(any(feature = "zstd", test))]
208            return Ok(Some(Box::new(ZSTDCodec::new(level))));
209            Err(ParquetError::General(
210                "Disabled feature at compile time: zstd".into(),
211            ))
212        }
213        CodecType::LZ4_RAW => {
214            #[cfg(any(feature = "lz4", test))]
215            return Ok(Some(Box::new(LZ4RawCodec::new())));
216            Err(ParquetError::General(
217                "Disabled feature at compile time: lz4".into(),
218            ))
219        }
220        CodecType::UNCOMPRESSED => Ok(None),
221        CodecType::LZO => Err(nyi_err!("The codec type {} is not supported yet", codec)),
222    }
223}
224
225#[cfg(any(feature = "snap", test))]
226mod snappy_codec {
227    use snap::raw::{Decoder, Encoder, decompress_len, max_compress_len};
228
229    use crate::compression::Codec;
230    use crate::errors::Result;
231
232    /// Codec for Snappy compression format.
233    pub struct SnappyCodec {
234        decoder: Decoder,
235        encoder: Encoder,
236    }
237
238    impl SnappyCodec {
239        /// Creates new Snappy compression codec.
240        pub(crate) fn new() -> Self {
241            Self {
242                decoder: Decoder::new(),
243                encoder: Encoder::new(),
244            }
245        }
246    }
247
248    impl Codec for SnappyCodec {
249        fn decompress(
250            &mut self,
251            input_buf: &[u8],
252            output_buf: &mut Vec<u8>,
253            uncompress_size: Option<usize>,
254        ) -> Result<usize> {
255            let len = match uncompress_size {
256                Some(size) => size,
257                None => decompress_len(input_buf)?,
258            };
259            let offset = output_buf.len();
260            output_buf.resize(offset + len, 0);
261            self.decoder
262                .decompress(input_buf, &mut output_buf[offset..])
263                .map_err(|e| e.into())
264        }
265
266        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
267            let output_buf_len = output_buf.len();
268            let required_len = max_compress_len(input_buf.len());
269            output_buf.resize(output_buf_len + required_len, 0);
270            let n = self
271                .encoder
272                .compress(input_buf, &mut output_buf[output_buf_len..])?;
273            output_buf.truncate(output_buf_len + n);
274            Ok(())
275        }
276    }
277}
278#[cfg(any(feature = "snap", test))]
279pub use snappy_codec::*;
280
281#[cfg(any(feature = "flate2", test))]
282mod gzip_codec {
283
284    use std::io::{Read, Write};
285
286    use flate2::{Compression, read, write};
287
288    use crate::compression::Codec;
289    use crate::errors::Result;
290
291    use super::GzipLevel;
292
293    /// Codec for GZIP compression algorithm.
294    pub struct GZipCodec {
295        level: GzipLevel,
296    }
297
298    impl GZipCodec {
299        /// Creates new GZIP compression codec.
300        pub(crate) fn new(level: GzipLevel) -> Self {
301            Self { level }
302        }
303    }
304
305    impl Codec for GZipCodec {
306        fn decompress(
307            &mut self,
308            input_buf: &[u8],
309            output_buf: &mut Vec<u8>,
310            _uncompress_size: Option<usize>,
311        ) -> Result<usize> {
312            let mut decoder = read::MultiGzDecoder::new(input_buf);
313            decoder.read_to_end(output_buf).map_err(|e| e.into())
314        }
315
316        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
317            let mut encoder = write::GzEncoder::new(output_buf, Compression::new(self.level.0));
318            encoder.write_all(input_buf)?;
319            encoder.try_finish().map_err(|e| e.into())
320        }
321    }
322}
323#[cfg(any(feature = "flate2", test))]
324pub use gzip_codec::*;
325
326/// Represents a valid gzip compression level.
327///
328/// Defaults to 6.
329///
330/// * 0: least compression
331/// * 9: most compression (that other software can read)
332/// * 10: most compression (incompatible with other software, see below)
333/// #### WARNING:
334/// Level 10 compression can offer smallest file size,
335/// but Parquet files created with it will not be readable
336/// by other "standard" paquet readers.
337///
338/// Do **NOT** use level 10 if you need other software to
339/// be able to read the files. Read below for details.
340///
341/// ### IMPORTANT:
342/// There's often confusion about the compression levels in `flate2` vs `arrow`
343/// as highlighted in issue [#1011](https://github.com/apache/arrow-rs/issues/6282).
344///
345/// `flate2` supports two compression backends: `miniz_oxide` and `zlib`.
346///
347/// - `zlib` supports levels from 0 to 9.
348/// - `miniz_oxide` supports levels from 0 to 10.
349///
350/// `arrow` uses `flate2` with `rust_backend` feature,
351/// which provides `miniz_oxide` as the backend.
352/// Therefore 0-10 levels are supported.
353///
354/// `flate2` documents this behavior properly with
355/// [this commit](https://github.com/rust-lang/flate2-rs/pull/430).
356#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
357pub struct GzipLevel(u32);
358
359impl Default for GzipLevel {
360    fn default() -> Self {
361        // The default as of miniz_oxide 0.5.1 is 6 for compression level
362        // (miniz_oxide::deflate::CompressionLevel::DefaultLevel)
363        Self(6)
364    }
365}
366
367impl CompressionLevel<u32> for GzipLevel {
368    const MINIMUM_LEVEL: u32 = 0;
369    const MAXIMUM_LEVEL: u32 = 9;
370}
371
372impl GzipLevel {
373    /// Attempts to create a gzip compression level.
374    ///
375    /// Compression levels must be valid (i.e. be acceptable for [`flate2::Compression`]).
376    pub fn try_new(level: u32) -> Result<Self> {
377        Self::is_valid_level(level).map(|()| Self(level))
378    }
379
380    /// Returns the compression level.
381    pub fn compression_level(&self) -> u32 {
382        self.0
383    }
384}
385
386#[cfg(any(feature = "brotli", test))]
387mod brotli_codec {
388
389    use std::io::{Read, Write};
390
391    use crate::compression::Codec;
392    use crate::errors::Result;
393
394    use super::BrotliLevel;
395
396    const BROTLI_DEFAULT_BUFFER_SIZE: usize = 4096;
397    const BROTLI_DEFAULT_LG_WINDOW_SIZE: u32 = 22; // recommended between 20-22
398
399    /// Codec for Brotli compression algorithm.
400    pub struct BrotliCodec {
401        level: BrotliLevel,
402    }
403
404    impl BrotliCodec {
405        /// Creates new Brotli compression codec.
406        pub(crate) fn new(level: BrotliLevel) -> Self {
407            Self { level }
408        }
409    }
410
411    impl Codec for BrotliCodec {
412        fn decompress(
413            &mut self,
414            input_buf: &[u8],
415            output_buf: &mut Vec<u8>,
416            uncompress_size: Option<usize>,
417        ) -> Result<usize> {
418            let buffer_size = uncompress_size.unwrap_or(BROTLI_DEFAULT_BUFFER_SIZE);
419            brotli::Decompressor::new(input_buf, buffer_size)
420                .read_to_end(output_buf)
421                .map_err(|e| e.into())
422        }
423
424        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
425            let mut encoder = brotli::CompressorWriter::new(
426                output_buf,
427                BROTLI_DEFAULT_BUFFER_SIZE,
428                self.level.0,
429                BROTLI_DEFAULT_LG_WINDOW_SIZE,
430            );
431            encoder.write_all(input_buf)?;
432            encoder.flush().map_err(|e| e.into())
433        }
434    }
435}
436#[cfg(any(feature = "brotli", test))]
437pub use brotli_codec::*;
438
439/// Represents a valid brotli compression level.
440#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
441pub struct BrotliLevel(u32);
442
443impl Default for BrotliLevel {
444    fn default() -> Self {
445        Self(1)
446    }
447}
448
449impl CompressionLevel<u32> for BrotliLevel {
450    const MINIMUM_LEVEL: u32 = 0;
451    const MAXIMUM_LEVEL: u32 = 11;
452}
453
454impl BrotliLevel {
455    /// Attempts to create a brotli compression level.
456    ///
457    /// Compression levels must be valid.
458    pub fn try_new(level: u32) -> Result<Self> {
459        Self::is_valid_level(level).map(|()| Self(level))
460    }
461
462    /// Returns the compression level.
463    pub fn compression_level(&self) -> u32 {
464        self.0
465    }
466}
467
468#[cfg(any(feature = "lz4", test))]
469mod lz4_codec {
470    use std::io::{Read, Write};
471
472    use crate::compression::Codec;
473    use crate::errors::{ParquetError, Result};
474
475    const LZ4_BUFFER_SIZE: usize = 4096;
476
477    /// Codec for LZ4 compression algorithm.
478    pub struct LZ4Codec {}
479
480    impl LZ4Codec {
481        /// Creates new LZ4 compression codec.
482        pub(crate) fn new() -> Self {
483            Self {}
484        }
485    }
486
487    impl Codec for LZ4Codec {
488        fn decompress(
489            &mut self,
490            input_buf: &[u8],
491            output_buf: &mut Vec<u8>,
492            _uncompress_size: Option<usize>,
493        ) -> Result<usize> {
494            let mut decoder = lz4_flex::frame::FrameDecoder::new(input_buf);
495            let mut buffer: [u8; LZ4_BUFFER_SIZE] = [0; LZ4_BUFFER_SIZE];
496            let mut total_len = 0;
497            loop {
498                let len = decoder.read(&mut buffer)?;
499                if len == 0 {
500                    break;
501                }
502                total_len += len;
503                output_buf.write_all(&buffer[0..len])?;
504            }
505            Ok(total_len)
506        }
507
508        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
509            let mut encoder = lz4_flex::frame::FrameEncoder::new(output_buf);
510            let mut from = 0;
511            loop {
512                let to = std::cmp::min(from + LZ4_BUFFER_SIZE, input_buf.len());
513                encoder.write_all(&input_buf[from..to])?;
514                from += LZ4_BUFFER_SIZE;
515                if from >= input_buf.len() {
516                    break;
517                }
518            }
519            match encoder.finish() {
520                Ok(_) => Ok(()),
521                Err(e) => Err(ParquetError::External(Box::new(e))),
522            }
523        }
524    }
525}
526
527#[cfg(all(feature = "experimental", any(feature = "lz4", test)))]
528pub use lz4_codec::*;
529
530#[cfg(any(feature = "zstd", test))]
531mod zstd_codec {
532    use crate::compression::{Codec, ZstdLevel};
533    use crate::errors::Result;
534    use std::io::Cursor;
535
536    /// Codec for Zstandard compression algorithm.
537    ///
538    /// Uses `zstd::bulk` API with reusable compressor/decompressor contexts
539    /// to avoid the overhead of reinitializing contexts for each operation.
540    pub struct ZSTDCodec {
541        compressor: zstd::bulk::Compressor<'static>,
542        decompressor: zstd::bulk::Decompressor<'static>,
543    }
544
545    impl ZSTDCodec {
546        /// Creates new Zstandard compression codec.
547        pub(crate) fn new(level: ZstdLevel) -> Self {
548            Self {
549                compressor: zstd::bulk::Compressor::new(level.compression_level())
550                    .expect("valid zstd compression level"),
551                decompressor: zstd::bulk::Decompressor::new()
552                    .expect("can create zstd decompressor"),
553            }
554        }
555    }
556
557    impl Codec for ZSTDCodec {
558        fn decompress(
559            &mut self,
560            input_buf: &[u8],
561            output_buf: &mut Vec<u8>,
562            uncompress_size: Option<usize>,
563        ) -> Result<usize> {
564            let offset = output_buf.len();
565            let len = uncompress_size
566                .or_else(|| {
567                    // Get the decompressed size from the zstd frame header
568                    zstd::zstd_safe::get_frame_content_size(input_buf)
569                        .ok()
570                        .flatten()
571                        .map(|size| size as usize)
572                })
573                .unwrap_or_else(|| input_buf.len().saturating_mul(4));
574            output_buf.reserve(len);
575
576            let mut cursor = Cursor::new(output_buf);
577            cursor.set_position(offset as u64);
578            let len = self
579                .decompressor
580                .decompress_to_buffer(input_buf, &mut cursor)?;
581            Ok(len)
582        }
583
584        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
585            let offset = output_buf.len();
586            let len = zstd::zstd_safe::compress_bound(input_buf.len());
587            output_buf.reserve(len);
588
589            let mut cursor = Cursor::new(output_buf);
590            cursor.set_position(offset as u64);
591            let _written = self.compressor.compress_to_buffer(input_buf, &mut cursor)?;
592            Ok(())
593        }
594    }
595}
596#[cfg(any(feature = "zstd", test))]
597pub use zstd_codec::*;
598
599/// Represents a valid zstd compression level.
600#[derive(Debug, Eq, PartialEq, Hash, Clone, Copy)]
601pub struct ZstdLevel(i32);
602
603impl CompressionLevel<i32> for ZstdLevel {
604    // zstd binds to C, and hence zstd::compression_level_range() is not const as this calls the
605    // underlying C library.
606    const MINIMUM_LEVEL: i32 = -131072;
607    const MAXIMUM_LEVEL: i32 = 22;
608}
609
610impl ZstdLevel {
611    /// Attempts to create a zstd compression level from a given compression level.
612    ///
613    /// Compression levels must be valid (i.e. be acceptable for [`zstd::compression_level_range`]).
614    pub fn try_new(level: i32) -> Result<Self> {
615        Self::is_valid_level(level).map(|()| Self(level))
616    }
617
618    /// Returns the compression level.
619    pub fn compression_level(&self) -> i32 {
620        self.0
621    }
622}
623
624impl Default for ZstdLevel {
625    fn default() -> Self {
626        Self(1)
627    }
628}
629
630#[cfg(any(feature = "lz4", test))]
631mod lz4_raw_codec {
632    use crate::compression::Codec;
633    use crate::errors::ParquetError;
634    use crate::errors::Result;
635
636    /// Codec for LZ4 Raw compression algorithm.
637    pub struct LZ4RawCodec {}
638
639    impl LZ4RawCodec {
640        /// Creates new LZ4 Raw compression codec.
641        pub(crate) fn new() -> Self {
642            Self {}
643        }
644    }
645
646    impl Codec for LZ4RawCodec {
647        fn decompress(
648            &mut self,
649            input_buf: &[u8],
650            output_buf: &mut Vec<u8>,
651            uncompress_size: Option<usize>,
652        ) -> Result<usize> {
653            let offset = output_buf.len();
654            let Some(required_len) = uncompress_size else {
655                return Err(ParquetError::General(
656                    "LZ4RawCodec unsupported without uncompress_size".into(),
657                ));
658            };
659            output_buf.resize(offset + required_len, 0);
660            match lz4_flex::block::decompress_into(input_buf, &mut output_buf[offset..]) {
661                Ok(n) => {
662                    if n != required_len {
663                        return Err(ParquetError::General(
664                            "LZ4RawCodec uncompress_size is not the expected one".into(),
665                        ));
666                    }
667                    Ok(n)
668                }
669                Err(e) => Err(ParquetError::External(Box::new(e))),
670            }
671        }
672
673        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
674            let offset = output_buf.len();
675            let required_len = lz4_flex::block::get_maximum_output_size(input_buf.len());
676            output_buf.resize(offset + required_len, 0);
677            match lz4_flex::block::compress_into(input_buf, &mut output_buf[offset..]) {
678                Ok(n) => {
679                    output_buf.truncate(offset + n);
680                    Ok(())
681                }
682                Err(e) => Err(ParquetError::External(Box::new(e))),
683            }
684        }
685    }
686}
687#[cfg(any(feature = "lz4", test))]
688pub use lz4_raw_codec::*;
689
690#[cfg(any(feature = "lz4", test))]
691mod lz4_hadoop_codec {
692    use crate::compression::Codec;
693    use crate::compression::lz4_codec::LZ4Codec;
694    use crate::compression::lz4_raw_codec::LZ4RawCodec;
695    use crate::errors::{ParquetError, Result};
696    use std::io;
697
698    /// Size of u32 type.
699    const SIZE_U32: usize = std::mem::size_of::<u32>();
700
701    /// Length of the LZ4_HADOOP prefix.
702    const PREFIX_LEN: usize = SIZE_U32 * 2;
703
704    /// Codec for LZ4 Hadoop compression algorithm.
705    pub struct LZ4HadoopCodec {
706        /// Whether or not to fallback to other LZ4 implementations on error.
707        /// Fallback is done to be backward compatible with older versions of this
708        /// library and older versions parquet-cpp.
709        backward_compatible_lz4: bool,
710    }
711
712    impl LZ4HadoopCodec {
713        /// Creates new LZ4 Hadoop compression codec.
714        pub(crate) fn new(backward_compatible_lz4: bool) -> Self {
715            Self {
716                backward_compatible_lz4,
717            }
718        }
719    }
720
721    /// Try to decompress the buffer as if it was compressed with the Hadoop Lz4Codec.
722    /// Adapted from pola-rs [compression.rs:try_decompress_hadoop](https://pola-rs.github.io/polars/src/parquet2/compression.rs.html#225)
723    /// Translated from the apache arrow c++ function [TryDecompressHadoop](https://github.com/apache/arrow/blob/bf18e6e4b5bb6180706b1ba0d597a65a4ce5ca48/cpp/src/arrow/util/compression_lz4.cc#L474).
724    /// Returns error if decompression failed.
725    fn try_decompress_hadoop(input_buf: &[u8], output_buf: &mut [u8]) -> io::Result<usize> {
726        // Parquet files written with the Hadoop Lz4Codec use their own framing.
727        // The input buffer can contain an arbitrary number of "frames", each
728        // with the following structure:
729        // - bytes 0..3: big-endian uint32_t representing the frame decompressed size
730        // - bytes 4..7: big-endian uint32_t representing the frame compressed size
731        // - bytes 8...: frame compressed data
732        //
733        // The Hadoop Lz4Codec source code can be found here:
734        // https://github.com/apache/hadoop/blob/trunk/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-nativetask/src/main/native/src/codec/Lz4Codec.cc
735        let mut input_len = input_buf.len();
736        let mut input = input_buf;
737        let mut read_bytes = 0;
738        let mut output_len = output_buf.len();
739        let mut output: &mut [u8] = output_buf;
740        while input_len >= PREFIX_LEN {
741            let mut bytes = [0; SIZE_U32];
742            bytes.copy_from_slice(&input[0..4]);
743            let expected_decompressed_size = u32::from_be_bytes(bytes);
744            let mut bytes = [0; SIZE_U32];
745            bytes.copy_from_slice(&input[4..8]);
746            let expected_compressed_size = u32::from_be_bytes(bytes);
747            input = &input[PREFIX_LEN..];
748            input_len -= PREFIX_LEN;
749
750            if input_len < expected_compressed_size as usize {
751                return Err(io::Error::other("Not enough bytes for Hadoop frame"));
752            }
753
754            if output_len < expected_decompressed_size as usize {
755                return Err(io::Error::other(
756                    "Not enough bytes to hold advertised output",
757                ));
758            }
759            let decompressed_size =
760                lz4_flex::decompress_into(&input[..expected_compressed_size as usize], output)
761                    .map_err(|e| ParquetError::External(Box::new(e)))?;
762            if decompressed_size != expected_decompressed_size as usize {
763                return Err(io::Error::other("Unexpected decompressed size"));
764            }
765            input_len -= expected_compressed_size as usize;
766            output_len -= expected_decompressed_size as usize;
767            read_bytes += expected_decompressed_size as usize;
768            if input_len > expected_compressed_size as usize {
769                input = &input[expected_compressed_size as usize..];
770                output = &mut output[expected_decompressed_size as usize..];
771            } else {
772                break;
773            }
774        }
775        if input_len == 0 {
776            Ok(read_bytes)
777        } else {
778            Err(io::Error::other("Not all input are consumed"))
779        }
780    }
781
782    impl Codec for LZ4HadoopCodec {
783        fn decompress(
784            &mut self,
785            input_buf: &[u8],
786            output_buf: &mut Vec<u8>,
787            uncompress_size: Option<usize>,
788        ) -> Result<usize> {
789            let output_len = output_buf.len();
790            let Some(required_len) = uncompress_size else {
791                return Err(ParquetError::General(
792                    "LZ4HadoopCodec unsupported without uncompress_size".into(),
793                ));
794            };
795            output_buf.resize(output_len + required_len, 0);
796            match try_decompress_hadoop(input_buf, &mut output_buf[output_len..]) {
797                Ok(n) => {
798                    if n != required_len {
799                        return Err(ParquetError::General(
800                            "LZ4HadoopCodec uncompress_size is not the expected one".into(),
801                        ));
802                    }
803                    Ok(n)
804                }
805                Err(e) if !self.backward_compatible_lz4 => Err(e.into()),
806                // Fallback done to be backward compatible with older versions of this
807                // library and older versions of parquet-cpp.
808                Err(_) => {
809                    // Truncate any inserted element before tryingg next algorithm.
810                    output_buf.truncate(output_len);
811                    match LZ4Codec::new().decompress(input_buf, output_buf, uncompress_size) {
812                        Ok(n) => Ok(n),
813                        Err(_) => {
814                            // Truncate any inserted element before tryingg next algorithm.
815                            output_buf.truncate(output_len);
816                            LZ4RawCodec::new().decompress(input_buf, output_buf, uncompress_size)
817                        }
818                    }
819                }
820            }
821        }
822
823        fn compress(&mut self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
824            // Allocate memory to store the LZ4_HADOOP prefix.
825            let offset = output_buf.len();
826            output_buf.resize(offset + PREFIX_LEN, 0);
827
828            // Append LZ4_RAW compressed bytes after prefix.
829            LZ4RawCodec::new().compress(input_buf, output_buf)?;
830
831            // Prepend decompressed size and compressed size in big endian to be compatible
832            // with LZ4_HADOOP.
833            let output_buf = &mut output_buf[offset..];
834            let compressed_size = output_buf.len() - PREFIX_LEN;
835            let compressed_size = compressed_size as u32;
836            let uncompressed_size = input_buf.len() as u32;
837            output_buf[..SIZE_U32].copy_from_slice(&uncompressed_size.to_be_bytes());
838            output_buf[SIZE_U32..PREFIX_LEN].copy_from_slice(&compressed_size.to_be_bytes());
839
840            Ok(())
841        }
842    }
843}
844#[cfg(any(feature = "lz4", test))]
845pub use lz4_hadoop_codec::*;
846
847#[cfg(test)]
848mod tests {
849    use super::*;
850
851    use crate::util::test_common::rand_gen::random_bytes;
852
853    fn test_roundtrip(c: CodecType, data: &[u8], uncompress_size: Option<usize>) {
854        let codec_options = CodecOptionsBuilder::default()
855            .set_backward_compatible_lz4(false)
856            .build();
857        let mut c1 = create_codec(c, &codec_options).unwrap().unwrap();
858        let mut c2 = create_codec(c, &codec_options).unwrap().unwrap();
859
860        // Compress with c1
861        let mut compressed = Vec::new();
862        let mut decompressed = Vec::new();
863        c1.compress(data, &mut compressed)
864            .expect("Error when compressing");
865
866        // Decompress with c2
867        let decompressed_size = c2
868            .decompress(compressed.as_slice(), &mut decompressed, uncompress_size)
869            .expect("Error when decompressing");
870        assert_eq!(data.len(), decompressed_size);
871        assert_eq!(data, decompressed.as_slice());
872
873        decompressed.clear();
874        compressed.clear();
875
876        // Compress with c2
877        c2.compress(data, &mut compressed)
878            .expect("Error when compressing");
879
880        // Decompress with c1
881        let decompressed_size = c1
882            .decompress(compressed.as_slice(), &mut decompressed, uncompress_size)
883            .expect("Error when decompressing");
884        assert_eq!(data.len(), decompressed_size);
885        assert_eq!(data, decompressed.as_slice());
886
887        decompressed.clear();
888        compressed.clear();
889
890        // Test does not trample existing data in output buffers
891        let prefix = &[0xDE, 0xAD, 0xBE, 0xEF];
892        decompressed.extend_from_slice(prefix);
893        compressed.extend_from_slice(prefix);
894
895        c2.compress(data, &mut compressed)
896            .expect("Error when compressing");
897
898        assert_eq!(&compressed[..4], prefix);
899
900        let decompressed_size = c2
901            .decompress(&compressed[4..], &mut decompressed, uncompress_size)
902            .expect("Error when decompressing");
903
904        assert_eq!(data.len(), decompressed_size);
905        assert_eq!(data, &decompressed[4..]);
906        assert_eq!(&decompressed[..4], prefix);
907    }
908
909    fn test_codec_with_size(c: CodecType) {
910        let sizes = vec![100, 10000, 100000];
911        for size in sizes {
912            let data = random_bytes(size);
913            test_roundtrip(c, &data, Some(data.len()));
914        }
915    }
916
917    fn test_codec_without_size(c: CodecType) {
918        let sizes = vec![100, 10000, 100000];
919        for size in sizes {
920            let data = random_bytes(size);
921            test_roundtrip(c, &data, None);
922        }
923    }
924
925    #[test]
926    #[cfg_attr(miri, ignore)] // Takes too long
927    fn test_codec_snappy() {
928        test_codec_with_size(CodecType::SNAPPY);
929        test_codec_without_size(CodecType::SNAPPY);
930    }
931
932    #[test]
933    #[cfg_attr(miri, ignore)] // Takes too long
934    fn test_codec_gzip() {
935        for level in GzipLevel::MINIMUM_LEVEL..=GzipLevel::MAXIMUM_LEVEL {
936            let level = GzipLevel::try_new(level).unwrap();
937            test_codec_with_size(CodecType::GZIP(level));
938            test_codec_without_size(CodecType::GZIP(level));
939        }
940    }
941
942    #[test]
943    #[cfg_attr(miri, ignore)] // Takes too long
944    fn test_codec_brotli() {
945        for level in BrotliLevel::MINIMUM_LEVEL..=BrotliLevel::MAXIMUM_LEVEL {
946            let level = BrotliLevel::try_new(level).unwrap();
947            test_codec_with_size(CodecType::BROTLI(level));
948            test_codec_without_size(CodecType::BROTLI(level));
949        }
950    }
951
952    #[test]
953    #[cfg_attr(miri, ignore)] // Takes too long
954    fn test_codec_lz4() {
955        test_codec_with_size(CodecType::LZ4);
956    }
957
958    #[test]
959    #[cfg_attr(miri, ignore)] // Zstd calls native C functions unsupported by Miri
960    fn test_codec_zstd() {
961        // since ZstdLevel::MINIMUM_LEVEL is a large negative number, we test a smaller range
962        for level in
963            std::iter::once(ZstdLevel::MINIMUM_LEVEL).chain(-100..=ZstdLevel::MAXIMUM_LEVEL)
964        {
965            let level = ZstdLevel::try_new(level).unwrap();
966            test_codec_with_size(CodecType::ZSTD(level));
967            test_codec_without_size(CodecType::ZSTD(level));
968        }
969    }
970
971    #[test]
972    #[cfg_attr(miri, ignore)] // Takes too long
973    fn test_codec_lz4_raw() {
974        test_codec_with_size(CodecType::LZ4_RAW);
975    }
976}