src/parquet/compression.cc - parquet-cpp - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.

 #include "parquet/compression.h"

 #include <cstdint>
 #include <memory>
 #include <string>

 #include <brotli/decode.h>
 #include <brotli/encode.h>
 #include <snappy.h>
 #include <zlib.h>

 #include "parquet/exception.h"
 #include "parquet/types.h"

 namespace parquet {

 Codec::~Codec() {}

 std::unique_ptr<Codec> Codec::Create(Compression::type codec_type) {
   std::unique_ptr<Codec> result;
   switch (codec_type) {
     case Compression::UNCOMPRESSED:
       break;
     case Compression::SNAPPY:
       result.reset(new SnappyCodec());
       break;
     case Compression::GZIP:
       result.reset(new GZipCodec());
       break;
     case Compression::LZO:
       ParquetException::NYI("LZO codec not implemented");
       break;
     case Compression::BROTLI:
       result.reset(new BrotliCodec());
       break;
     default:
       ParquetException::NYI("Unrecognized codec");
       break;
   }
   return result;
 }

 // ----------------------------------------------------------------------
 // gzip implementation

 // These are magic numbers from zlib.h.  Not clear why they are not defined
 // there.

 // Maximum window size
 static constexpr int WINDOW_BITS = 15;

 // Output Gzip.
 static constexpr int GZIP_CODEC = 16;

 // Determine if this is libz or gzip from header.
 static constexpr int DETECT_CODEC = 32;

 class GZipCodec::GZipCodecImpl {
  public:
   explicit GZipCodecImpl(GZipCodec::Format format)
       : format_(format),
         compressor_initialized_(false),
         decompressor_initialized_(false) {}

   ~GZipCodecImpl() {
     EndCompressor();
     EndDecompressor();
   }

   void InitCompressor() {
     EndDecompressor();
     memset(&stream_, 0, sizeof(stream_));

     int ret;
     // Initialize to run specified format
     int window_bits = WINDOW_BITS;
     if (format_ == DEFLATE) {
       window_bits = -window_bits;
     } else if (format_ == GZIP) {
       window_bits += GZIP_CODEC;
     }
     if ((ret = deflateInit2(&stream_, Z_DEFAULT_COMPRESSION, Z_DEFLATED, window_bits, 9,
              Z_DEFAULT_STRATEGY)) != Z_OK) {
       throw ParquetException("zlib deflateInit failed: " + std::string(stream_.msg));
     }

     compressor_initialized_ = true;
   }

   void EndCompressor() {
     if (compressor_initialized_) { (void)deflateEnd(&stream_); }
     compressor_initialized_ = false;
   }

   void InitDecompressor() {
     EndCompressor();
     memset(&stream_, 0, sizeof(stream_));
     int ret;

     // Initialize to run either deflate or zlib/gzip format
     int window_bits = format_ == DEFLATE ? -WINDOW_BITS : WINDOW_BITS | DETECT_CODEC;
     if ((ret = inflateInit2(&stream_, window_bits)) != Z_OK) {
       throw ParquetException("zlib inflateInit failed: " + std::string(stream_.msg));
     }
     decompressor_initialized_ = true;
   }

   void EndDecompressor() {
     if (decompressor_initialized_) { (void)inflateEnd(&stream_); }
     decompressor_initialized_ = false;
   }

   void Decompress(int64_t input_length, const uint8_t* input, int64_t output_length,
       uint8_t* output) {
     if (!decompressor_initialized_) { InitDecompressor(); }
     if (output_length == 0) {
       // The zlib library does not allow *output to be NULL, even when output_length
       // is 0 (inflate() will return Z_STREAM_ERROR). We don't consider this an
       // error, so bail early if no output is expected. Note that we don't signal
       // an error if the input actually contains compressed data.
       return;
     }

     // Reset the stream for this block
     if (inflateReset(&stream_) != Z_OK) {
       throw ParquetException("zlib inflateReset failed: " + std::string(stream_.msg));
     }

     int ret = 0;
     // gzip can run in streaming mode or non-streaming mode.  We only
     // support the non-streaming use case where we present it the entire
     // compressed input and a buffer big enough to contain the entire
     // compressed output.  In the case where we don't know the output,
     // we just make a bigger buffer and try the non-streaming mode
     // from the beginning again.
     while (ret != Z_STREAM_END) {
       stream_.next_in = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(input));
       stream_.avail_in = input_length;
       stream_.next_out = reinterpret_cast<Bytef*>(output);
       stream_.avail_out = output_length;

       // We know the output size.  In this case, we can use Z_FINISH
       // which is more efficient.
       ret = inflate(&stream_, Z_FINISH);
       if (ret == Z_STREAM_END || ret != Z_OK) break;

       // Failure, buffer was too small
       std::stringstream ss;
       ss << "Too small a buffer passed to GZipCodec. InputLength=" << input_length
          << " OutputLength=" << output_length;
       throw ParquetException(ss.str());
     }

     // Failure for some other reason
     if (ret != Z_STREAM_END) {
       std::stringstream ss;
       ss << "GZipCodec failed: ";
       if (stream_.msg != NULL) ss << stream_.msg;
       throw ParquetException(ss.str());
     }
   }

   int64_t MaxCompressedLen(int64_t input_length, const uint8_t* input) {
     // Most be in compression mode
     if (!compressor_initialized_) { InitCompressor(); }
     // TODO(wesm): deal with zlib < 1.2.3 (see Impala codebase)
     return deflateBound(&stream_, static_cast<uLong>(input_length));
   }

   int64_t Compress(int64_t input_length, const uint8_t* input, int64_t output_length,
       uint8_t* output) {
     if (!compressor_initialized_) { InitCompressor(); }
     stream_.next_in = const_cast<Bytef*>(reinterpret_cast<const Bytef*>(input));
     stream_.avail_in = input_length;
     stream_.next_out = reinterpret_cast<Bytef*>(output);
     stream_.avail_out = output_length;

     int64_t ret = 0;
     if ((ret = deflate(&stream_, Z_FINISH)) != Z_STREAM_END) {
       if (ret == Z_OK) {
         // will return Z_OK (and stream.msg NOT set) if stream.avail_out is too
         // small
         throw ParquetException("zlib deflate failed, output buffer too small");
       }
       std::stringstream ss;
       ss << "zlib deflate failed: " << stream_.msg;
       throw ParquetException(ss.str());
     }

     if (deflateReset(&stream_) != Z_OK) {
       throw ParquetException("zlib deflateReset failed: " + std::string(stream_.msg));
     }

     // Actual output length
     return output_length - stream_.avail_out;
   }

  private:
   // zlib is stateful and the z_stream state variable must be initialized
   // before
   z_stream stream_;

   // Realistically, this will always be GZIP, but we leave the option open to
   // configure
   GZipCodec::Format format_;

   // These variables are mutually exclusive. When the codec is in "compressor"
   // state, compressor_initialized_ is true while decompressor_initialized_ is
   // false. When it's decompressing, the opposite is true.
   //
   // Indeed, this is slightly hacky, but the alternative is having separate
   // Compressor and Decompressor classes. If this ever becomes an issue, we can
   // perform the refactoring then
   bool compressor_initialized_;
   bool decompressor_initialized_;
 };

 GZipCodec::GZipCodec(Format format) {
   impl_.reset(new GZipCodecImpl(format));
 }

 GZipCodec::~GZipCodec() {}

 void GZipCodec::Decompress(
     int64_t input_length, const uint8_t* input, int64_t output_length, uint8_t* output) {
   return impl_->Decompress(input_length, input, output_length, output);
 }

 int64_t GZipCodec::MaxCompressedLen(int64_t input_length, const uint8_t* input) {
   return impl_->MaxCompressedLen(input_length, input);
 }

 int64_t GZipCodec::Compress(
     int64_t input_length, const uint8_t* input, int64_t output_length, uint8_t* output) {
   return impl_->Compress(input_length, input, output_length, output);
 }

 const char* GZipCodec::name() const {
   return "gzip";
 }

 // ----------------------------------------------------------------------
 // Snappy implementation

 void SnappyCodec::Decompress(
     int64_t input_len, const uint8_t* input, int64_t output_len, uint8_t* output_buffer) {
   if (!snappy::RawUncompress(reinterpret_cast<const char*>(input),
           static_cast<size_t>(input_len), reinterpret_cast<char*>(output_buffer))) {
     throw parquet::ParquetException("Corrupt snappy compressed data.");
   }
 }

 int64_t SnappyCodec::MaxCompressedLen(int64_t input_len, const uint8_t* input) {
   return snappy::MaxCompressedLength(input_len);
 }

 int64_t SnappyCodec::Compress(int64_t input_len, const uint8_t* input,
     int64_t output_buffer_len, uint8_t* output_buffer) {
   size_t output_len;
   snappy::RawCompress(reinterpret_cast<const char*>(input),
       static_cast<size_t>(input_len), reinterpret_cast<char*>(output_buffer),
       &output_len);
   return output_len;
 }

 // ----------------------------------------------------------------------
 // Brotli implementation

 void BrotliCodec::Decompress(
     int64_t input_len, const uint8_t* input, int64_t output_len, uint8_t* output_buffer) {
   size_t output_size = output_len;
   if (BrotliDecoderDecompress(input_len, input, &output_size, output_buffer) !=
       BROTLI_DECODER_RESULT_SUCCESS) {
     throw parquet::ParquetException("Corrupt brotli compressed data.");
   }
 }

 int64_t BrotliCodec::MaxCompressedLen(int64_t input_len, const uint8_t* input) {
   return BrotliEncoderMaxCompressedSize(input_len);
 }

 int64_t BrotliCodec::Compress(int64_t input_len, const uint8_t* input,
     int64_t output_buffer_len, uint8_t* output_buffer) {
   size_t output_len = output_buffer_len;
   // TODO: Make quality configurable. We use 8 as a default as it is the best
   //       trade-off for Parquet workload
   if (BrotliEncoderCompress(8, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, input_len,
           input, &output_len, output_buffer) == BROTLI_FALSE) {
     throw parquet::ParquetException("Brotli compression failure.");
   }
   return output_len;
 }

 }  // namespace parquet
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.

	#include "parquet/compression.h"

	#include <cstdint>
	#include <memory>
	#include <string>

	#include <brotli/decode.h>
	#include <brotli/encode.h>
	#include <snappy.h>
	#include <zlib.h>

	#include "parquet/exception.h"
	#include "parquet/types.h"

	namespace parquet {

	Codec::~Codec() {}

	std::unique_ptr<Codec> Codec::Create(Compression::type codec_type) {
	std::unique_ptr<Codec> result;
	switch (codec_type) {
	case Compression::UNCOMPRESSED:
	break;
	case Compression::SNAPPY:
	result.reset(new SnappyCodec());
	break;
	case Compression::GZIP:
	result.reset(new GZipCodec());
	break;
	case Compression::LZO:
	ParquetException::NYI("LZO codec not implemented");
	break;
	case Compression::BROTLI:
	result.reset(new BrotliCodec());
	break;
	default:
	ParquetException::NYI("Unrecognized codec");
	break;
	}
	return result;
	}

	// ----------------------------------------------------------------------
	// gzip implementation

	// These are magic numbers from zlib.h. Not clear why they are not defined
	// there.

	// Maximum window size
	static constexpr int WINDOW_BITS = 15;

	// Output Gzip.
	static constexpr int GZIP_CODEC = 16;

	// Determine if this is libz or gzip from header.
	static constexpr int DETECT_CODEC = 32;

	class GZipCodec::GZipCodecImpl {
	public:
	explicit GZipCodecImpl(GZipCodec::Format format)
	: format_(format),
	compressor_initialized_(false),
	decompressor_initialized_(false) {}

	~GZipCodecImpl() {
	EndCompressor();
	EndDecompressor();
	}

	void InitCompressor() {
	EndDecompressor();
	memset(&stream_, 0, sizeof(stream_));

	int ret;
	// Initialize to run specified format
	int window_bits = WINDOW_BITS;
	if (format_ == DEFLATE) {
	window_bits = -window_bits;
	} else if (format_ == GZIP) {
	window_bits += GZIP_CODEC;
	}
	if ((ret = deflateInit2(&stream_, Z_DEFAULT_COMPRESSION, Z_DEFLATED, window_bits, 9,
	Z_DEFAULT_STRATEGY)) != Z_OK) {
	throw ParquetException("zlib deflateInit failed: " + std::string(stream_.msg));
	}

	compressor_initialized_ = true;
	}

	void EndCompressor() {
	if (compressor_initialized_) { (void)deflateEnd(&stream_); }
	compressor_initialized_ = false;
	}

	void InitDecompressor() {
	EndCompressor();
	memset(&stream_, 0, sizeof(stream_));
	int ret;

	// Initialize to run either deflate or zlib/gzip format
	int window_bits = format_ == DEFLATE ? -WINDOW_BITS : WINDOW_BITS \| DETECT_CODEC;
	if ((ret = inflateInit2(&stream_, window_bits)) != Z_OK) {
	throw ParquetException("zlib inflateInit failed: " + std::string(stream_.msg));
	}
	decompressor_initialized_ = true;
	}

	void EndDecompressor() {
	if (decompressor_initialized_) { (void)inflateEnd(&stream_); }
	decompressor_initialized_ = false;
	}

	void Decompress(int64_t input_length, const uint8_t* input, int64_t output_length,
	uint8_t* output) {
	if (!decompressor_initialized_) { InitDecompressor(); }
	if (output_length == 0) {
	// The zlib library does not allow *output to be NULL, even when output_length
	// is 0 (inflate() will return Z_STREAM_ERROR). We don't consider this an
	// error, so bail early if no output is expected. Note that we don't signal
	// an error if the input actually contains compressed data.
	return;
	}

	// Reset the stream for this block
	if (inflateReset(&stream_) != Z_OK) {
	throw ParquetException("zlib inflateReset failed: " + std::string(stream_.msg));
	}

	int ret = 0;
	// gzip can run in streaming mode or non-streaming mode. We only
	// support the non-streaming use case where we present it the entire
	// compressed input and a buffer big enough to contain the entire
	// compressed output. In the case where we don't know the output,
	// we just make a bigger buffer and try the non-streaming mode
	// from the beginning again.
	while (ret != Z_STREAM_END) {
	stream_.next_in = const_cast<Bytef>(reinterpret_cast<const Bytef>(input));
	stream_.avail_in = input_length;
	stream_.next_out = reinterpret_cast<Bytef*>(output);
	stream_.avail_out = output_length;

	// We know the output size. In this case, we can use Z_FINISH
	// which is more efficient.
	ret = inflate(&stream_, Z_FINISH);
	if (ret == Z_STREAM_END \|\| ret != Z_OK) break;

	// Failure, buffer was too small
	std::stringstream ss;
	ss << "Too small a buffer passed to GZipCodec. InputLength=" << input_length
	<< " OutputLength=" << output_length;
	throw ParquetException(ss.str());
	}

	// Failure for some other reason
	if (ret != Z_STREAM_END) {
	std::stringstream ss;
	ss << "GZipCodec failed: ";
	if (stream_.msg != NULL) ss << stream_.msg;
	throw ParquetException(ss.str());
	}
	}

	int64_t MaxCompressedLen(int64_t input_length, const uint8_t* input) {
	// Most be in compression mode
	if (!compressor_initialized_) { InitCompressor(); }
	// TODO(wesm): deal with zlib < 1.2.3 (see Impala codebase)
	return deflateBound(&stream_, static_cast<uLong>(input_length));
	}

	int64_t Compress(int64_t input_length, const uint8_t* input, int64_t output_length,
	uint8_t* output) {
	if (!compressor_initialized_) { InitCompressor(); }
	stream_.next_in = const_cast<Bytef>(reinterpret_cast<const Bytef>(input));
	stream_.avail_in = input_length;
	stream_.next_out = reinterpret_cast<Bytef*>(output);
	stream_.avail_out = output_length;

	int64_t ret = 0;
	if ((ret = deflate(&stream_, Z_FINISH)) != Z_STREAM_END) {
	if (ret == Z_OK) {
	// will return Z_OK (and stream.msg NOT set) if stream.avail_out is too
	// small
	throw ParquetException("zlib deflate failed, output buffer too small");
	}
	std::stringstream ss;
	ss << "zlib deflate failed: " << stream_.msg;
	throw ParquetException(ss.str());
	}

	if (deflateReset(&stream_) != Z_OK) {
	throw ParquetException("zlib deflateReset failed: " + std::string(stream_.msg));
	}

	// Actual output length
	return output_length - stream_.avail_out;
	}

	private:
	// zlib is stateful and the z_stream state variable must be initialized
	// before
	z_stream stream_;

	// Realistically, this will always be GZIP, but we leave the option open to
	// configure
	GZipCodec::Format format_;

	// These variables are mutually exclusive. When the codec is in "compressor"
	// state, compressor_initialized_ is true while decompressor_initialized_ is
	// false. When it's decompressing, the opposite is true.
	//
	// Indeed, this is slightly hacky, but the alternative is having separate
	// Compressor and Decompressor classes. If this ever becomes an issue, we can
	// perform the refactoring then
	bool compressor_initialized_;
	bool decompressor_initialized_;
	};

	GZipCodec::GZipCodec(Format format) {
	impl_.reset(new GZipCodecImpl(format));
	}

	GZipCodec::~GZipCodec() {}

	void GZipCodec::Decompress(
	int64_t input_length, const uint8_t* input, int64_t output_length, uint8_t* output) {
	return impl_->Decompress(input_length, input, output_length, output);
	}

	int64_t GZipCodec::MaxCompressedLen(int64_t input_length, const uint8_t* input) {
	return impl_->MaxCompressedLen(input_length, input);
	}

	int64_t GZipCodec::Compress(
	int64_t input_length, const uint8_t* input, int64_t output_length, uint8_t* output) {
	return impl_->Compress(input_length, input, output_length, output);
	}

	const char* GZipCodec::name() const {
	return "gzip";
	}

	// ----------------------------------------------------------------------
	// Snappy implementation

	void SnappyCodec::Decompress(
	int64_t input_len, const uint8_t* input, int64_t output_len, uint8_t* output_buffer) {
	if (!snappy::RawUncompress(reinterpret_cast<const char*>(input),
	static_cast<size_t>(input_len), reinterpret_cast<char*>(output_buffer))) {
	throw parquet::ParquetException("Corrupt snappy compressed data.");
	}
	}

	int64_t SnappyCodec::MaxCompressedLen(int64_t input_len, const uint8_t* input) {
	return snappy::MaxCompressedLength(input_len);
	}

	int64_t SnappyCodec::Compress(int64_t input_len, const uint8_t* input,
	int64_t output_buffer_len, uint8_t* output_buffer) {
	size_t output_len;
	snappy::RawCompress(reinterpret_cast<const char*>(input),
	static_cast<size_t>(input_len), reinterpret_cast<char*>(output_buffer),
	&output_len);
	return output_len;
	}

	// ----------------------------------------------------------------------
	// Brotli implementation

	void BrotliCodec::Decompress(
	int64_t input_len, const uint8_t* input, int64_t output_len, uint8_t* output_buffer) {
	size_t output_size = output_len;
	if (BrotliDecoderDecompress(input_len, input, &output_size, output_buffer) !=
	BROTLI_DECODER_RESULT_SUCCESS) {
	throw parquet::ParquetException("Corrupt brotli compressed data.");
	}
	}

	int64_t BrotliCodec::MaxCompressedLen(int64_t input_len, const uint8_t* input) {
	return BrotliEncoderMaxCompressedSize(input_len);
	}

	int64_t BrotliCodec::Compress(int64_t input_len, const uint8_t* input,
	int64_t output_buffer_len, uint8_t* output_buffer) {
	size_t output_len = output_buffer_len;
	// TODO: Make quality configurable. We use 8 as a default as it is the best
	// trade-off for Parquet workload
	if (BrotliEncoderCompress(8, BROTLI_DEFAULT_WINDOW, BROTLI_DEFAULT_MODE, input_len,
	input, &output_len, output_buffer) == BROTLI_FALSE) {
	throw parquet::ParquetException("Brotli compression failure.");
	}
	return output_len;
	}

	} // namespace parquet