arrow_avro/
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
18use crate::errors::AvroError;
19use arrow_schema::ArrowError;
20#[cfg(any(
21    feature = "deflate",
22    feature = "zstd",
23    feature = "bzip2",
24    feature = "xz"
25))]
26use std::io::{Read, Write};
27
28/// The metadata key used for storing the JSON encoded [`CompressionCodec`]
29pub const CODEC_METADATA_KEY: &str = "avro.codec";
30
31#[derive(Debug, Copy, Clone, Eq, PartialEq)]
32/// Supported compression codecs for Avro data
33///
34/// Avro supports multiple compression formats for data blocks.
35/// This enum represents the compression codecs available in this implementation.
36pub enum CompressionCodec {
37    /// Deflate compression (RFC 1951)
38    Deflate,
39    /// Snappy compression
40    Snappy,
41    /// ZStandard compression
42    ZStandard,
43    /// Bzip2 compression
44    Bzip2,
45    /// Xz compression
46    Xz,
47}
48
49impl CompressionCodec {
50    #[cfg_attr(
51        not(any(
52            feature = "deflate",
53            feature = "snappy",
54            feature = "zstd",
55            feature = "bzip2",
56            feature = "xz"
57        )),
58        expect(unused_variables)
59    )]
60    pub(crate) fn decompress(&self, block: &[u8]) -> Result<Vec<u8>, AvroError> {
61        match self {
62            #[cfg(feature = "deflate")]
63            CompressionCodec::Deflate => {
64                let mut decoder = flate2::read::DeflateDecoder::new(block);
65                let mut out = Vec::new();
66                decoder.read_to_end(&mut out)?;
67                Ok(out)
68            }
69            #[cfg(not(feature = "deflate"))]
70            CompressionCodec::Deflate => Err(AvroError::ParseError(
71                "Deflate codec requires deflate feature".to_string(),
72            )),
73            #[cfg(feature = "snappy")]
74            CompressionCodec::Snappy => {
75                // Each compressed block is followed by the 4-byte, big-endian CRC32
76                // checksum of the uncompressed data in the block.
77                let crc = &block[block.len() - 4..];
78                let block = &block[..block.len() - 4];
79
80                let mut decoder = snap::raw::Decoder::new();
81                let decoded = decoder
82                    .decompress_vec(block)
83                    .map_err(|e| AvroError::External(Box::new(e)))?;
84
85                let checksum = crc::Crc::<u32>::new(&crc::CRC_32_ISO_HDLC).checksum(&decoded);
86                if checksum != u32::from_be_bytes(crc.try_into().unwrap()) {
87                    return Err(AvroError::ParseError("Snappy CRC mismatch".to_string()));
88                }
89                Ok(decoded)
90            }
91            #[cfg(not(feature = "snappy"))]
92            CompressionCodec::Snappy => Err(AvroError::ParseError(
93                "Snappy codec requires snappy feature".to_string(),
94            )),
95
96            #[cfg(feature = "zstd")]
97            CompressionCodec::ZStandard => {
98                let mut decoder = zstd::Decoder::new(block)?;
99                let mut out = Vec::new();
100                decoder
101                    .read_to_end(&mut out)
102                    .map_err(|e| AvroError::External(Box::new(e)))?;
103                Ok(out)
104            }
105            #[cfg(not(feature = "zstd"))]
106            CompressionCodec::ZStandard => Err(AvroError::ParseError(
107                "ZStandard codec requires zstd feature".to_string(),
108            )),
109            #[cfg(feature = "bzip2")]
110            CompressionCodec::Bzip2 => {
111                let mut decoder = bzip2::read::BzDecoder::new(block);
112                let mut out = Vec::new();
113                decoder
114                    .read_to_end(&mut out)
115                    .map_err(|e| AvroError::External(Box::new(e)))?;
116                Ok(out)
117            }
118            #[cfg(not(feature = "bzip2"))]
119            CompressionCodec::Bzip2 => Err(AvroError::ParseError(
120                "Bzip2 codec requires bzip2 feature".to_string(),
121            )),
122            #[cfg(feature = "xz")]
123            CompressionCodec::Xz => {
124                let mut decoder = xz::read::XzDecoder::new(block);
125                let mut out = Vec::new();
126                decoder
127                    .read_to_end(&mut out)
128                    .map_err(|e| AvroError::External(Box::new(e)))?;
129                Ok(out)
130            }
131            #[cfg(not(feature = "xz"))]
132            CompressionCodec::Xz => Err(AvroError::ParseError(
133                "XZ codec requires xz feature".to_string(),
134            )),
135        }
136    }
137
138    #[cfg_attr(
139        not(any(
140            feature = "deflate",
141            feature = "snappy",
142            feature = "zstd",
143            feature = "bzip2",
144            feature = "xz"
145        )),
146        expect(unused_variables)
147    )]
148    pub(crate) fn compress(&self, data: &[u8]) -> Result<Vec<u8>, ArrowError> {
149        match self {
150            #[cfg(feature = "deflate")]
151            CompressionCodec::Deflate => {
152                let mut encoder =
153                    flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
154                encoder.write_all(data)?;
155                let compressed = encoder.finish()?;
156                Ok(compressed)
157            }
158            #[cfg(not(feature = "deflate"))]
159            CompressionCodec::Deflate => Err(ArrowError::ParseError(
160                "Deflate codec requires deflate feature".to_string(),
161            )),
162
163            #[cfg(feature = "snappy")]
164            CompressionCodec::Snappy => {
165                let mut encoder = snap::raw::Encoder::new();
166                // Allocate and compress in one step for efficiency
167                let mut compressed = encoder
168                    .compress_vec(data)
169                    .map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
170                // Compute CRC32 (ISO‑HDLC poly) of **uncompressed** data
171                let crc_val = crc::Crc::<u32>::new(&crc::CRC_32_ISO_HDLC).checksum(data);
172                compressed.extend_from_slice(&crc_val.to_be_bytes());
173                Ok(compressed)
174            }
175            #[cfg(not(feature = "snappy"))]
176            CompressionCodec::Snappy => Err(ArrowError::ParseError(
177                "Snappy codec requires snappy feature".to_string(),
178            )),
179
180            #[cfg(feature = "zstd")]
181            CompressionCodec::ZStandard => {
182                let mut encoder = zstd::Encoder::new(Vec::new(), 0)
183                    .map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
184                encoder.write_all(data)?;
185                let compressed = encoder
186                    .finish()
187                    .map_err(|e| ArrowError::ExternalError(Box::new(e)))?;
188                Ok(compressed)
189            }
190            #[cfg(not(feature = "zstd"))]
191            CompressionCodec::ZStandard => Err(ArrowError::ParseError(
192                "ZStandard codec requires zstd feature".to_string(),
193            )),
194
195            #[cfg(feature = "bzip2")]
196            CompressionCodec::Bzip2 => {
197                let mut encoder =
198                    bzip2::write::BzEncoder::new(Vec::new(), bzip2::Compression::default());
199                encoder.write_all(data)?;
200                let compressed = encoder.finish()?;
201                Ok(compressed)
202            }
203            #[cfg(not(feature = "bzip2"))]
204            CompressionCodec::Bzip2 => Err(ArrowError::ParseError(
205                "Bzip2 codec requires bzip2 feature".to_string(),
206            )),
207            #[cfg(feature = "xz")]
208            CompressionCodec::Xz => {
209                let mut encoder = xz::write::XzEncoder::new(Vec::new(), 6);
210                encoder.write_all(data)?;
211                let compressed = encoder.finish()?;
212                Ok(compressed)
213            }
214            #[cfg(not(feature = "xz"))]
215            CompressionCodec::Xz => Err(ArrowError::ParseError(
216                "XZ codec requires xz feature".to_string(),
217            )),
218        }
219    }
220}
arrow_avro/compression.rs

arrow_avro/
compression.rs
