| // 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. |
| |
| |
| #ifndef IMPALA_UTIL_BIT_STREAM_UTILS_H |
| #define IMPALA_UTIL_BIT_STREAM_UTILS_H |
| |
| #include <boost/cstdint.hpp> |
| #include <string.h> |
| #include "common/compiler-util.h" |
| #include "common/logging.h" |
| #include "util/bit-packing.h" |
| #include "util/bit-util.h" |
| |
| namespace impala { |
| |
| /// Utility class to write bit/byte streams. This class can write data to either be |
| /// bit packed or byte aligned (and a single stream that has a mix of both). |
| /// This class does not allocate memory. |
| class BitWriter { |
| public: |
| /// buffer: buffer to write bits to. Buffer should be preallocated with |
| /// 'buffer_len' bytes. |
| BitWriter(uint8_t* buffer, int buffer_len) : |
| buffer_(buffer), |
| max_bytes_(buffer_len) { |
| Clear(); |
| } |
| |
| void Clear() { |
| buffered_values_ = 0; |
| byte_offset_ = 0; |
| bit_offset_ = 0; |
| } |
| |
| /// The number of current bytes written, including the current byte (i.e. may include a |
| /// fraction of a byte). Includes buffered values. |
| int bytes_written() const { return byte_offset_ + BitUtil::Ceil(bit_offset_, 8); } |
| uint8_t* buffer() const { return buffer_; } |
| int buffer_len() const { return max_bytes_; } |
| |
| /// Writes a value to buffered_values_, flushing to buffer_ if necessary. This is bit |
| /// packed. Returns false if there was not enough space. num_bits must be <= 64. |
| bool PutValue(uint64_t v, int num_bits); |
| |
| /// Writes v to the next aligned byte using num_bytes. If T is larger than num_bytes, the |
| /// extra high-order bytes will be ignored. Returns false if there was not enough space. |
| template<typename T> |
| bool PutAligned(T v, int num_bytes); |
| |
| /// Write an unsigned ULEB-128 encoded int to the buffer. Return false if there was not |
| /// enough room. The value is written byte aligned. For more details on ULEB-128: |
| /// https://en.wikipedia.org/wiki/LEB128 |
| /// UINT_T must be an unsigned integer type. |
| template<typename UINT_T> |
| bool PutUleb128(UINT_T v); |
| |
| /// Write a ZigZag encoded int to the buffer. Return false if there was not enough |
| /// room. The value is written byte aligned. For more details on ZigZag encoding: |
| /// https://developers.google.com/protocol-buffers/docs/encoding#signed-integers |
| /// INT_T must be a signed integer type. |
| template<typename INT_T> |
| bool PutZigZagInteger(INT_T v); |
| |
| /// Get a pointer to the next aligned byte and advance the underlying buffer |
| /// by num_bytes. |
| /// Returns NULL if there was not enough space. |
| uint8_t* GetNextBytePtr(int num_bytes = 1); |
| |
| /// Flushes all buffered values to the buffer. Call this when done writing to the buffer. |
| /// If 'align' is true, buffered_values_ is reset and any future writes will be written |
| /// to the next byte boundary. |
| void Flush(bool align=false); |
| |
| /// Maximum supported bitwidth for writer. |
| static const int MAX_BITWIDTH = 64; |
| |
| private: |
| uint8_t* buffer_; |
| int max_bytes_; |
| |
| /// Bit-packed values are initially written to this variable before being memcpy'd to |
| /// buffer_. This is faster than writing values byte by byte directly to buffer_. |
| uint64_t buffered_values_; |
| |
| int byte_offset_; // Offset in buffer_ |
| int bit_offset_; // Offset in buffered_values_ |
| }; |
| |
| /// Utility class to read bit/byte stream. This class can read bits or bytes that are |
| /// either byte aligned or not. It also has utilities to read multiple bytes in one |
| /// read (e.g. encoded int). Exposes a batch-oriented interface to allow efficient |
| /// processing of multiple values at a time. |
| class BatchedBitReader { |
| public: |
| /// 'buffer' is the buffer to read from. The buffer's length is 'buffer_len'. |
| /// Does not take ownership of the buffer. |
| BatchedBitReader(const uint8_t* buffer, int64_t buffer_len) { |
| Reset(buffer, buffer_len); |
| } |
| |
| BatchedBitReader() {} |
| |
| // The implicit copy constructor is left defined. If a BatchedBitReader is copied, the |
| // two copies do not share any state. Invoking functions on either copy continues |
| // reading from the current read position without modifying the state of the other |
| // copy. |
| |
| /// Resets the read to start reading from the start of 'buffer'. The buffer's |
| /// length is 'buffer_len'. Does not take ownership of the buffer. |
| void Reset(const uint8_t* buffer, int64_t buffer_len) { |
| DCHECK(buffer != nullptr); |
| DCHECK_GE(buffer_len, 0); |
| buffer_pos_ = buffer; |
| buffer_end_ = buffer + buffer_len; |
| } |
| |
| /// Gets up to 'num_values' bit-packed values, starting from the current byte in the |
| /// buffer and advance the read position. 'bit_width' must be <= 64. |
| /// If 'bit_width' * 'num_values' is not a multiple of 8, the trailing bytes are |
| /// skipped and the next UnpackBatch() call will start reading from the next byte. |
| /// |
| /// If the caller does not want to drop trailing bits, 'num_values' must be exactly the |
| /// total number of values the caller wants to read from a run of bit-packed values, or |
| /// 'bit_width' * 'num_values' must be a multiple of 8. This condition is always |
| /// satisfied if 'num_values' is a multiple of 32. |
| /// |
| /// The output type 'T' must be an unsigned integer. |
| /// |
| /// Returns the number of values read. |
| template<typename T> |
| int UnpackBatch(int bit_width, int num_values, T* v); |
| |
| /// Skip 'num_values_to_skip' bit-packed values. |
| /// 'num_values_to_skip * bit_width' is either divisible by 8, or |
| /// 'num_values_to_skip' equals to the count of the remaining bit-packed values. |
| bool SkipBatch(int bit_width, int num_values_to_skip); |
| |
| /// Unpack bit-packed values in the same way as UnpackBatch() and decode them using the |
| /// dictionary 'dict' with 'dict_len' entries. Return -1 if a decoding error is |
| /// encountered, i.e. if the bit-packed values are not valid indices in 'dict'. |
| /// Otherwise returns the number of values decoded. The values are written to 'v' with |
| /// a stride of 'stride' bytes. |
| template <typename T> |
| int UnpackAndDecodeBatch( |
| int bit_width, T* dict, int64_t dict_len, int num_values, T* v, int64_t stride); |
| |
| /// Reads an unpacked 'num_bytes'-sized value from the buffer and stores it in 'v'. T |
| /// needs to be a little-endian native type and big enough to store 'num_bytes'. |
| /// Returns false if there are not enough bytes left. |
| template<typename T> |
| bool GetBytes(int num_bytes, T* v); |
| |
| /// Read an unsigned ULEB-128 encoded int from the stream. The encoded int must start |
| /// at the beginning of a byte. Return false if there were not enough bytes in the |
| /// buffer or the int is invalid. For more details on ULEB-128: |
| /// https://en.wikipedia.org/wiki/LEB128 |
| /// UINT_T must be an unsigned integer type. |
| template<typename UINT_T> |
| bool GetUleb128(UINT_T* v); |
| |
| /// Read a ZigZag encoded int from the stream. The encoded int must start at the |
| /// beginning of a byte. Return false if there were not enough bytes in the buffer or |
| /// the int is invalid. For more details on ZigZag encoding: |
| /// https://developers.google.com/protocol-buffers/docs/encoding#signed-integers |
| /// INT_T must be a signed integer type. |
| template<typename INT_T> |
| bool GetZigZagInteger(INT_T* v); |
| |
| /// Returns the number of bytes left in the stream. |
| int bytes_left() { return buffer_end_ - buffer_pos_; } |
| |
| /// Maximum byte length of a vlq encoded integer of type T. |
| template <typename T> |
| static constexpr int max_vlq_byte_len() { |
| return BitUtil::Ceil(sizeof(T) * 8, 7); |
| } |
| |
| /// Maximum supported bitwidth for reader. |
| static const int MAX_BITWIDTH = BitPacking::MAX_BITWIDTH; |
| |
| private: |
| /// Current read position in the buffer. |
| const uint8_t* buffer_pos_ = nullptr; |
| |
| /// Pointer to the byte after the end of the buffer. |
| const uint8_t* buffer_end_ = nullptr; |
| }; |
| } |
| |
| #endif |