dict-encoding.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// Licensed 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_DICT_ENCODING_H
#define IMPALA_UTIL_DICT_ENCODING_H

#include <map>

#include <boost/foreach.hpp>
#include <boost/scoped_ptr.hpp>
#include <boost/unordered_map.hpp>

#include "exec/parquet-common.h"
#include "runtime/mem-pool.h"
#include "runtime/string-value.h"
#include "util/rle-encoding.h"
#include "util/runtime-profile.h"

namespace impala {


class DictEncoderBase {
 public:
  virtual ~DictEncoderBase() {
    DCHECK(buffered_indices_.empty());
  }

  virtual void WriteDict(uint8_t* buffer) = 0;

  virtual int num_entries() const = 0;

  void ClearIndices() { buffered_indices_.clear(); }

  int EstimatedDataEncodedSize() {
    return 1 + RleEncoder::MaxBufferSize(bit_width(), buffered_indices_.size());
  }

  int bit_width() const {
    if (UNLIKELY(num_entries() == 0)) return 0;
    if (UNLIKELY(num_entries() == 1)) return 1;
    return BitUtil::Log2(num_entries());
  }

  int WriteData(uint8_t* buffer, int buffer_len);

  int dict_encoded_size() { return dict_encoded_size_; }

 protected:
  DictEncoderBase(MemPool* pool)
    : dict_encoded_size_(0), pool_(pool) {
  }

  std::vector<int> buffered_indices_;

  int dict_encoded_size_;

  MemPool* pool_;
};

template<typename T>
class DictEncoder : public DictEncoderBase {
 public:
  DictEncoder(MemPool* pool, int encoded_value_size) :
      DictEncoderBase(pool), buckets_(HASH_TABLE_SIZE, Node::INVALID_INDEX),
      encoded_value_size_(encoded_value_size) { }

  int Put(const T& value);

  virtual void WriteDict(uint8_t* buffer);

  virtual int num_entries() const { return nodes_.size(); }

 private:
  enum { HASH_TABLE_SIZE = 1 << 16 };

  typedef uint16_t NodeIndex;

  std::vector<NodeIndex> buckets_;

  struct Node {
    Node(const T& v, const NodeIndex& n) : value(v), next(n) { }

    T value;

    NodeIndex next;

    enum { INVALID_INDEX = 40000 };
  };

  std::vector<Node> nodes_;

  int encoded_value_size_;

  inline uint32_t Hash(const T& value) const;

  int AddToTable(const T& value, NodeIndex* bucket);
};

class DictDecoderBase {
 public:
  void SetData(uint8_t* buffer, int buffer_len) {
    DCHECK_GT(buffer_len, 0);
    uint8_t bit_width = *buffer;
    DCHECK_GE(bit_width, 0);
    ++buffer;
    --buffer_len;
    data_decoder_.reset(new RleDecoder(buffer, buffer_len, bit_width));
  }

  virtual ~DictDecoderBase() {}

  virtual int num_entries() const = 0;

 protected:
  boost::scoped_ptr<RleDecoder> data_decoder_;
};

template<typename T>
class DictDecoder : public DictDecoderBase {
 public:
  DictDecoder(uint8_t* dict_buffer, int dict_len, int fixed_len_size);

  virtual int num_entries() const { return dict_.size(); }

  bool GetValue(T* value);

 private:
  std::vector<T> dict_;
};

template<typename T>
inline int DictEncoder<T>::Put(const T& value) {
  NodeIndex* bucket = &buckets_[Hash(value) & (HASH_TABLE_SIZE - 1)];
  NodeIndex i = *bucket;
  // Look for the value in the dictionary.
  while (i != Node::INVALID_INDEX) {
    const Node* n = &nodes_[i];
    if (LIKELY(n->value == value)) {
      // Value already in dictionary.
      buffered_indices_.push_back(i);
      return 0;
    }
    i = n->next;
  }
  // Value not found. Add it to the dictionary if there's space.
  i = nodes_.size();
  if (UNLIKELY(i >= Node::INVALID_INDEX)) return -1;
  buffered_indices_.push_back(i);
  return AddToTable(value, bucket);
}

template<typename T>
inline uint32_t DictEncoder<T>::Hash(const T& value) const {
  return HashUtil::Hash(&value, sizeof(value), 0);
}

template<>
inline uint32_t DictEncoder<StringValue>::Hash(const StringValue& value) const {
  return HashUtil::Hash(value.ptr, value.len, 0);
}

template<typename T>
inline int DictEncoder<T>::AddToTable(const T& value, NodeIndex* bucket) {
  DCHECK_GT(encoded_value_size_, 0);
  // Prepend the new node to this bucket's chain.
  nodes_.push_back(Node(value, *bucket));
  *bucket = nodes_.size() - 1;
  dict_encoded_size_ += encoded_value_size_;
  return encoded_value_size_;
}

template<>
inline int DictEncoder<StringValue>::AddToTable(const StringValue& value,
    NodeIndex* bucket) {
  char* ptr_copy = reinterpret_cast<char*>(pool_->Allocate(value.len));
  memcpy(ptr_copy, value.ptr, value.len);
  StringValue sv(ptr_copy, value.len);
  // Prepend the new node to this bucket's chain.
  nodes_.push_back(Node(sv, *bucket));
  *bucket = nodes_.size() - 1;
  int bytes_added = ParquetPlainEncoder::ByteSize(sv);
  dict_encoded_size_ += bytes_added;
  return bytes_added;
}

template<typename T>
inline bool DictDecoder<T>::GetValue(T* value) {
  DCHECK(data_decoder_.get() != NULL);
  int index;
  bool result = data_decoder_->Get(&index);
  if (!result) return false;
  if (index >= dict_.size()) return false;
  *value = dict_[index];
  return true;
}

template<>
inline bool DictDecoder<Decimal16Value>::GetValue(Decimal16Value* value) {
  DCHECK(data_decoder_.get() != NULL);
  int index;
  bool result = data_decoder_->Get(&index);
  if (!result) return false;
  if (index >= dict_.size()) return false;
  // Workaround for IMPALA-959. Use memcpy instead of '=' so addresses
  // do not need to be 16 byte aligned.
  uint8_t* addr = reinterpret_cast<uint8_t*>(&dict_[0]);
  addr = addr + index * sizeof(*value);
  memcpy(value, addr, sizeof(*value));
  return true;
}

template<typename T>
inline void DictEncoder<T>::WriteDict(uint8_t* buffer) {
  BOOST_FOREACH(const Node& node, nodes_) {
    buffer += ParquetPlainEncoder::Encode(buffer, encoded_value_size_, node.value);
  }
}

inline int DictEncoderBase::WriteData(uint8_t* buffer, int buffer_len) {
  // Write bit width in first byte
  *buffer = bit_width();
  ++buffer;
  --buffer_len;

  RleEncoder encoder(buffer, buffer_len, bit_width());
  BOOST_FOREACH(int index, buffered_indices_) {
    if (!encoder.Put(index)) return -1;
  }
  encoder.Flush();
  return 1 + encoder.len();
}

template<typename T>
inline DictDecoder<T>::DictDecoder(uint8_t* dict_buffer, int dict_len,
    int fixed_len_size) {
  uint8_t* end = dict_buffer + dict_len;
  while (dict_buffer < end) {
    T value;
    dict_buffer +=
        ParquetPlainEncoder::Decode(dict_buffer, fixed_len_size, &value);
    dict_.push_back(value);
  }
}

}
#endif