blob: 3a93a5fdf53aa8dd6474ab2f13e2b16f911acddc [file]
/**
* 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
*
* https://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 "DataFile.hh"
#include "Compiler.hh"
#include "Exception.hh"
#include <random>
#include <sstream>
#ifdef SNAPPY_CODEC_AVAILABLE
#include <snappy.h>
#endif
#ifdef ZSTD_CODEC_AVAILABLE
#include "ZstdCompressWrapper.hh"
#include "ZstdDecompressWrapper.hh"
#endif
#include <zlib.h>
namespace avro {
using std::copy;
using std::istringstream;
using std::ostringstream;
using std::string;
using std::unique_ptr;
using std::vector;
using std::array;
namespace {
const string AVRO_SCHEMA_KEY("avro.schema");
const string AVRO_CODEC_KEY("avro.codec");
const size_t minSyncInterval = 32;
const size_t maxSyncInterval = 1u << 30;
// Recommended by https://www.zlib.net/zlib_how.html
const size_t zlibBufGrowSize = 128 * 1024;
template<Codec codec>
struct codec_trait {
static std::string name() {
throw Exception("Unsupported codec: {}", static_cast<int>(codec));
}
static void validate(std::optional<int> level) {
throw Exception("Unsupported codec: {}", static_cast<int>(codec));
}
static bool available() {
throw Exception("Unsupported codec: {}", static_cast<int>(codec));
}
};
template<>
struct codec_trait<NULL_CODEC> {
static std::string name() {
return "null";
}
static void validate(std::optional<int> /*level*/) {}
static bool available() {
return true;
}
};
template<>
struct codec_trait<DEFLATE_CODEC> {
static std::string name() {
return "deflate";
}
static void validate(std::optional<int> level) {
if (!level.has_value()) {
return;
}
int levelValue = level.value();
if (levelValue < 0 || levelValue > 9) {
throw Exception("Invalid compression level {} for deflate codec. "
"Valid range is 0-9.",
levelValue);
}
}
static bool available() {
return true;
}
};
template<>
struct codec_trait<SNAPPY_CODEC> {
static std::string name() {
return "snappy";
}
static void validate(std::optional<int> /*level*/) {
}
static bool available() {
#ifdef SNAPPY_CODEC_AVAILABLE
return true;
#else
return false;
#endif
}
};
template<>
struct codec_trait<ZSTD_CODEC> {
static std::string name() {
return "zstandard";
}
static void validate(std::optional<int> level) {
if (!level.has_value()) {
return;
}
int levelValue = level.value();
if (levelValue < 1 || levelValue > 22) {
throw Exception("Invalid compression level {} for zstandard codec. "
"Valid range is 1-22.",
levelValue);
}
}
static bool available() {
#ifdef ZSTD_CODEC_AVAILABLE
return true;
#else
return false;
#endif
}
};
#define DISPATCH_CODEC_FUNC(codec, func, ...) \
switch (codec) { \
case NULL_CODEC: \
return codec_trait<NULL_CODEC>::func(__VA_ARGS__); \
case DEFLATE_CODEC: \
return codec_trait<DEFLATE_CODEC>::func(__VA_ARGS__); \
case SNAPPY_CODEC: \
return codec_trait<SNAPPY_CODEC>::func(__VA_ARGS__); \
case ZSTD_CODEC: \
return codec_trait<ZSTD_CODEC>::func(__VA_ARGS__); \
default: \
throw Exception("Unknown codec: {}", static_cast<int>(codec)); \
}
std::string getCodecName(Codec codec) {
DISPATCH_CODEC_FUNC(codec, name);
}
void validateCodec(Codec codec, std::optional<int> level) {
if (!isCodecAvailable(codec)) {
throw Exception("Codec {} is not available.", getCodecName(codec));
}
DISPATCH_CODEC_FUNC(codec, validate, level);
}
Codec getCodec(const std::string &name) {
if (name == codec_trait<NULL_CODEC>::name()) {
return NULL_CODEC;
} else if (name == codec_trait<DEFLATE_CODEC>::name()) {
return DEFLATE_CODEC;
} else if (name == codec_trait<SNAPPY_CODEC>::name()) {
return SNAPPY_CODEC;
} else if (name == codec_trait<ZSTD_CODEC>::name()) {
return ZSTD_CODEC;
} else {
throw Exception("Unknown codec name: {}", name);
}
}
} // namespace
bool isCodecAvailable(Codec codec) {
DISPATCH_CODEC_FUNC(codec, available);
}
#undef DISPATCH_CODEC_FUNC
DataFileWriterBase::DataFileWriterBase(const char *filename, const ValidSchema &schema, size_t syncInterval,
Codec codec, const Metadata &metadata,
std::optional<int> compressionLevel) : filename_(filename),
schema_(schema),
encoderPtr_(binaryEncoder()),
syncInterval_(syncInterval),
codec_(codec),
compressionLevel_(compressionLevel),
stream_(fileOutputStream(filename)),
buffer_(memoryOutputStream()),
sync_(makeSync()),
objectCount_(0),
metadata_(metadata),
lastSync_(0) {
init(schema, syncInterval, codec);
}
DataFileWriterBase::DataFileWriterBase(std::unique_ptr<OutputStream> outputStream, const ValidSchema &schema,
size_t syncInterval, Codec codec, const Metadata &metadata,
std::optional<int> compressionLevel) : filename_(),
schema_(schema),
encoderPtr_(binaryEncoder()),
syncInterval_(syncInterval),
codec_(codec),
compressionLevel_(compressionLevel),
stream_(std::move(outputStream)),
buffer_(memoryOutputStream()),
sync_(makeSync()),
objectCount_(0),
metadata_(metadata),
lastSync_(0) {
init(schema, syncInterval, codec);
}
void DataFileWriterBase::init(const ValidSchema &schema, size_t syncInterval, const Codec &codec) {
if (syncInterval < minSyncInterval || syncInterval > maxSyncInterval) {
throw Exception(
"Invalid sync interval: {}. Should be between {} and {}",
syncInterval, minSyncInterval, maxSyncInterval);
}
validateCodec(codec, compressionLevel_);
setMetadata(AVRO_CODEC_KEY, getCodecName(codec));
setMetadata(AVRO_SCHEMA_KEY, schema.toJson(false));
writeHeader();
encoderPtr_->init(*buffer_);
lastSync_ = stream_->byteCount();
}
DataFileWriterBase::~DataFileWriterBase() {
if (stream_) {
try {
close();
} catch (...) {}
}
}
void DataFileWriterBase::close() {
flush();
stream_.reset();
}
void DataFileWriterBase::sync() {
encoderPtr_->flush();
encoderPtr_->init(*stream_);
avro::encode(*encoderPtr_, objectCount_);
if (codec_ == NULL_CODEC) {
int64_t byteCount = buffer_->byteCount();
avro::encode(*encoderPtr_, byteCount);
encoderPtr_->flush();
std::unique_ptr<InputStream> in = memoryInputStream(*buffer_);
copy(*in, *stream_);
} else if (codec_ == DEFLATE_CODEC) {
std::vector<uint8_t> buf;
{
z_stream zs;
zs.zalloc = Z_NULL;
zs.zfree = Z_NULL;
zs.opaque = Z_NULL;
// Use Z_DEFAULT_COMPRESSION if no level specified
int effectiveLevel = compressionLevel_.value_or(Z_DEFAULT_COMPRESSION);
int ret = deflateInit2(&zs, effectiveLevel, Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY);
if (ret != Z_OK) {
throw Exception("Failed to initialize deflate, error: {}", ret);
}
std::unique_ptr<InputStream> input = memoryInputStream(*buffer_);
const uint8_t *data;
size_t len;
while (ret != Z_STREAM_END && input->next(&data, &len)) {
zs.avail_in = static_cast<uInt>(len);
zs.next_in = const_cast<Bytef *>(data);
bool flush = (zs.total_in + len) >= buffer_->byteCount();
do {
if (zs.total_out == buf.size()) {
buf.resize(buf.size() + zlibBufGrowSize);
}
zs.avail_out = static_cast<uInt>(buf.size() - zs.total_out);
zs.next_out = buf.data() + zs.total_out;
ret = deflate(&zs, flush ? Z_FINISH : Z_NO_FLUSH);
if (ret == Z_STREAM_END) {
break;
}
if (ret != Z_OK) {
throw Exception("Failed to deflate, error: {}", ret);
}
} while (zs.avail_out == 0);
}
buf.resize(zs.total_out);
(void) deflateEnd(&zs);
} // make sure all is flushed
std::unique_ptr<InputStream> in = memoryInputStream(buf.data(), buf.size());
int64_t byteCount = buf.size();
avro::encode(*encoderPtr_, byteCount);
encoderPtr_->flush();
copy(*in, *stream_);
#ifdef SNAPPY_CODEC_AVAILABLE
} else if (codec_ == SNAPPY_CODEC) {
std::vector<char> temp;
std::string compressed;
const uint8_t *data;
size_t len;
std::unique_ptr<InputStream> input = memoryInputStream(*buffer_);
while (input->next(&data, &len)) {
temp.insert(temp.end(), reinterpret_cast<const char *>(data),
reinterpret_cast<const char *>(data) + len);
}
// For Snappy, add the CRC32 checksum
auto checksum = crc32(0, reinterpret_cast<const Bytef *>(temp.data()),
static_cast<uInt>(temp.size()));
// Now compress
size_t compressed_size = snappy::Compress(
reinterpret_cast<const char *>(temp.data()), temp.size(),
&compressed);
temp.clear();
temp.insert(temp.end(), compressed.c_str(),
compressed.c_str() + compressed_size);
temp.push_back(static_cast<char>((checksum >> 24) & 0xFF));
temp.push_back(static_cast<char>((checksum >> 16) & 0xFF));
temp.push_back(static_cast<char>((checksum >> 8) & 0xFF));
temp.push_back(static_cast<char>(checksum & 0xFF));
std::unique_ptr<InputStream> in = memoryInputStream(
reinterpret_cast<const uint8_t *>(temp.data()), temp.size());
int64_t byteCount = temp.size();
avro::encode(*encoderPtr_, byteCount);
encoderPtr_->flush();
copy(*in, *stream_);
#endif
#ifdef ZSTD_CODEC_AVAILABLE
} else if (codec_ == ZSTD_CODEC) {
// Read all uncompressed data into a single buffer
std::vector<char> uncompressed;
const uint8_t *data;
size_t len;
std::unique_ptr<InputStream> input = memoryInputStream(*buffer_);
while (input->next(&data, &len)) {
uncompressed.insert(uncompressed.end(), reinterpret_cast<const char *>(data),
reinterpret_cast<const char *>(data) + len);
}
ZstdCompressWrapper zstdCompressWrapper;
std::vector<char> compressed = zstdCompressWrapper.compress(uncompressed, compressionLevel_);
std::unique_ptr<InputStream> in = memoryInputStream(
reinterpret_cast<const uint8_t *>(compressed.data()), compressed.size());
avro::encode(*encoderPtr_, static_cast<int64_t>(compressed.size()));
encoderPtr_->flush();
copy(*in, *stream_);
#endif
}
encoderPtr_->init(*stream_);
avro::encode(*encoderPtr_, sync_);
encoderPtr_->flush();
lastSync_ = stream_->byteCount();
buffer_ = memoryOutputStream();
encoderPtr_->init(*buffer_);
objectCount_ = 0;
}
void DataFileWriterBase::syncIfNeeded() {
encoderPtr_->flush();
if (buffer_->byteCount() >= syncInterval_) {
sync();
}
}
uint64_t DataFileWriterBase::getCurrentBlockStart() const {
return lastSync_;
}
void DataFileWriterBase::flush() {
sync();
}
DataFileSync DataFileWriterBase::makeSync() {
std::mt19937 random(static_cast<uint32_t>(time(nullptr)));
DataFileSync sync;
std::generate(sync.begin(), sync.end(), random);
return sync;
}
typedef array<uint8_t, 4> Magic;
static Magic magic = {{'O', 'b', 'j', '\x01'}};
void DataFileWriterBase::writeHeader() {
encoderPtr_->init(*stream_);
avro::encode(*encoderPtr_, magic);
avro::encode(*encoderPtr_, metadata_);
avro::encode(*encoderPtr_, sync_);
encoderPtr_->flush();
}
void DataFileWriterBase::setMetadata(const string &key, const string &value) {
vector<uint8_t> v(value.size());
copy(value.begin(), value.end(), v.begin());
metadata_[key] = v;
}
DataFileReaderBase::DataFileReaderBase(const char *filename) : filename_(filename), stream_(fileSeekableInputStream(filename)),
decoder_(binaryDecoder()), objectCount_(0), eof_(false),
codec_(NULL_CODEC), blockStart_(-1), blockEnd_(-1) {
readHeader();
}
DataFileReaderBase::DataFileReaderBase(std::unique_ptr<InputStream> inputStream) : stream_(std::move(inputStream)),
decoder_(binaryDecoder()), objectCount_(0), eof_(false), codec_(NULL_CODEC) {
readHeader();
}
void DataFileReaderBase::init() {
readerSchema_ = dataSchema_;
dataDecoder_ = binaryDecoder();
readDataBlock();
}
void DataFileReaderBase::init(const ValidSchema &readerSchema) {
readerSchema_ = readerSchema;
dataDecoder_ = (readerSchema_.toJson(true) != dataSchema_.toJson(true)) ? resolvingDecoder(dataSchema_, readerSchema_, binaryDecoder()) : binaryDecoder();
readDataBlock();
}
static void drain(InputStream &in) {
const uint8_t *p = nullptr;
size_t n = 0;
while (in.next(&p, &n));
}
char hex(unsigned int x) {
return static_cast<char>(x + (x < 10 ? '0' : ('a' - 10)));
}
std::ostream &operator<<(std::ostream &os, const DataFileSync &s) {
for (uint8_t i : s) {
os << hex(i / 16) << hex(i % 16) << ' ';
}
os << std::endl;
return os;
}
bool DataFileReaderBase::hasMore() {
for (;;) {
if (eof_) {
return false;
} else if (objectCount_ != 0) {
return true;
}
dataDecoder_->init(*dataStream_);
drain(*dataStream_);
DataFileSync s;
decoder_->init(*stream_);
avro::decode(*decoder_, s);
if (s != sync_) {
throw Exception("Sync mismatch");
}
readDataBlock();
}
}
class BoundedInputStream : public InputStream {
InputStream &in_;
size_t limit_;
bool next(const uint8_t **data, size_t *len) final {
if (limit_ != 0 && in_.next(data, len)) {
if (*len > limit_) {
in_.backup(*len - limit_);
*len = limit_;
}
limit_ -= *len;
return true;
}
return false;
}
void backup(size_t len) final {
in_.backup(len);
limit_ += len;
}
void skip(size_t len) final {
if (len > limit_) {
len = limit_;
}
in_.skip(len);
limit_ -= len;
}
size_t byteCount() const final {
return in_.byteCount();
}
public:
BoundedInputStream(InputStream &in, size_t limit) : in_(in), limit_(limit) {}
};
unique_ptr<InputStream> boundedInputStream(InputStream &in, size_t limit) {
return unique_ptr<InputStream>(new BoundedInputStream(in, limit));
}
void DataFileReaderBase::readDataBlock() {
decoder_->init(*stream_);
blockStart_ = stream_->byteCount();
const uint8_t *p = nullptr;
size_t n = 0;
if (!stream_->next(&p, &n)) {
eof_ = true;
return;
}
stream_->backup(n);
avro::decode(*decoder_, objectCount_);
int64_t byteCount;
avro::decode(*decoder_, byteCount);
decoder_->init(*stream_);
blockEnd_ = stream_->byteCount() + byteCount;
unique_ptr<InputStream> st = boundedInputStream(*stream_, static_cast<size_t>(byteCount));
if (codec_ == NULL_CODEC) {
dataDecoder_->init(*st);
dataStream_ = std::move(st);
#ifdef SNAPPY_CODEC_AVAILABLE
} else if (codec_ == SNAPPY_CODEC) {
uint32_t checksum = 0;
compressed_.clear();
uncompressed.clear();
const uint8_t *data;
size_t len;
while (st->next(&data, &len)) {
compressed_.insert(compressed_.end(), data, data + len);
}
len = compressed_.size();
if (len < 4)
throw Exception("Cannot read compressed data, expected at least 4 bytes, got " + std::to_string(len));
int b1 = compressed_[len - 4] & 0xFF;
int b2 = compressed_[len - 3] & 0xFF;
int b3 = compressed_[len - 2] & 0xFF;
int b4 = compressed_[len - 1] & 0xFF;
checksum = (b1 << 24) + (b2 << 16) + (b3 << 8) + (b4);
if (!snappy::Uncompress(reinterpret_cast<const char *>(compressed_.data()),
len - 4, &uncompressed)) {
throw Exception(
"Snappy Compression reported an error when decompressing");
}
auto c = crc32(0, reinterpret_cast<const Bytef *>(uncompressed.c_str()),
static_cast<uInt>(uncompressed.size()));
if (checksum != c) {
throw Exception(
"Checksum did not match for Snappy compression: Expected: {}, computed: {}",
checksum, c);
}
std::unique_ptr<InputStream> in = memoryInputStream(
reinterpret_cast<const uint8_t *>(uncompressed.c_str()),
uncompressed.size());
dataDecoder_->init(*in);
dataStream_ = std::move(in);
#endif
#ifdef ZSTD_CODEC_AVAILABLE
} else if (codec_ == ZSTD_CODEC) {
compressed_.clear();
uncompressed.clear();
const uint8_t *data;
size_t len;
while (st->next(&data, &len)) {
compressed_.insert(compressed_.end(), data, data + len);
}
ZstdDecompressWrapper zstdDecompressWrapper;
uncompressed = zstdDecompressWrapper.decompress(compressed_);
std::unique_ptr<InputStream> in = memoryInputStream(
reinterpret_cast<const uint8_t *>(uncompressed.data()),
uncompressed.size());
dataDecoder_->init(*in);
dataStream_ = std::move(in);
#endif
} else {
compressed_.clear();
uncompressed.clear();
{
z_stream zs;
zs.zalloc = Z_NULL;
zs.zfree = Z_NULL;
zs.opaque = Z_NULL;
zs.avail_in = 0;
zs.next_in = Z_NULL;
int ret = inflateInit2(&zs, /*windowBits=*/-15);
if (ret != Z_OK) {
throw Exception("Failed to initialize inflate, error: {}", ret);
}
const uint8_t *data;
size_t len;
while (ret != Z_STREAM_END && st->next(&data, &len)) {
zs.avail_in = static_cast<uInt>(len);
zs.next_in = const_cast<Bytef *>(data);
do {
if (zs.total_out == uncompressed.size()) {
uncompressed.resize(uncompressed.size() + zlibBufGrowSize);
}
zs.avail_out = static_cast<uInt>(uncompressed.size() - zs.total_out);
zs.next_out = reinterpret_cast<Bytef *>(uncompressed.data() + zs.total_out);
ret = inflate(&zs, Z_NO_FLUSH);
if (ret == Z_STREAM_END) {
break;
}
if (ret != Z_OK) {
throw Exception("Failed to inflate, error: {}", ret);
}
} while (zs.avail_out == 0);
}
uncompressed.resize(zs.total_out);
(void) inflateEnd(&zs);
}
std::unique_ptr<InputStream> in = memoryInputStream(
reinterpret_cast<const uint8_t *>(uncompressed.c_str()),
uncompressed.size());
dataDecoder_->init(*in);
dataStream_ = std::move(in);
}
}
void DataFileReaderBase::close() {
stream_.reset();
eof_ = true;
objectCount_ = 0;
blockStart_ = 0;
blockEnd_ = 0;
}
static string toString(const vector<uint8_t> &v) {
string result;
result.resize(v.size());
copy(v.begin(), v.end(), result.begin());
return result;
}
static ValidSchema makeSchema(const vector<uint8_t> &v) {
istringstream iss(toString(v));
ValidSchema vs;
compileJsonSchema(iss, vs);
return vs;
}
void DataFileReaderBase::readHeader() {
decoder_->init(*stream_);
Magic m;
avro::decode(*decoder_, m);
if (magic != m) {
throw Exception("Invalid data file. Magic does not match: "
+ filename_);
}
avro::decode(*decoder_, metadata_);
Metadata::const_iterator it = metadata_.find(AVRO_SCHEMA_KEY);
if (it == metadata_.end()) {
throw Exception("No schema in metadata");
}
dataSchema_ = makeSchema(it->second);
if (!readerSchema_.root()) {
readerSchema_ = dataSchema();
}
// Parse codec from metadata using codec_trait
it = metadata_.find(AVRO_CODEC_KEY);
if (it != metadata_.end()) {
const auto codecName = toString(it->second);
codec_ = getCodec(codecName);
if (!isCodecAvailable(codec_)) {
throw Exception("Codec {} is not available.", codecName);
}
} else {
codec_ = NULL_CODEC;
}
avro::decode(*decoder_, sync_);
decoder_->init(*stream_);
blockStart_ = stream_->byteCount();
}
void DataFileReaderBase::doSeek(int64_t position) {
if (auto *ss = dynamic_cast<SeekableInputStream *>(stream_.get())) {
if (!eof_) {
dataDecoder_->init(*dataStream_);
drain(*dataStream_);
}
decoder_->init(*stream_);
ss->seek(position);
eof_ = false;
} else {
throw Exception("seek not supported on non-SeekableInputStream");
}
}
void DataFileReaderBase::seek(int64_t position) {
doSeek(position);
readDataBlock();
}
void DataFileReaderBase::sync(int64_t position) {
doSeek(position);
DataFileSync sync_buffer;
const uint8_t *p = nullptr;
size_t n = 0;
size_t i = 0;
while (i < SyncSize) {
if (n == 0 && !stream_->next(&p, &n)) {
eof_ = true;
return;
}
size_t len = std::min(SyncSize - i, n);
memcpy(&sync_buffer[i], p, len);
p += len;
n -= len;
i += len;
}
for (;;) {
size_t j = 0;
for (; j < SyncSize; ++j) {
if (sync_[j] != sync_buffer[(i + j) % SyncSize]) {
break;
}
}
if (j == SyncSize) {
// Found the sync marker!
break;
}
if (n == 0 && !stream_->next(&p, &n)) {
eof_ = true;
return;
}
sync_buffer[i++ % SyncSize] = *p++;
--n;
}
stream_->backup(n);
readDataBlock();
}
bool DataFileReaderBase::pastSync(int64_t position) {
return !hasMore() || blockStart_ >= position + SyncSize;
}
int64_t DataFileReaderBase::previousSync() const {
return blockStart_;
}
} // namespace avro