| /** |
| * 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 |