PARQUET-542: Support custom memory allocators

Added `MemoryAllocator` interface and a default implementation, ensured `MemPool` can use a custom allocator, added `Vector<T>` to replace `std::vector<T>`.

Author: Aliaksei Sandryhaila <aliaksei.sandryhaila@hp.com>

Closes #72 from asandryh/PARQUET-542 and squashes the following commits:

a740edb [Aliaksei Sandryhaila] Incorporated PR feedback.
6422e0d [Aliaksei Sandryhaila] Added MemoryAllocator interface and default implementation, ensured MemPool can use a custom allocator, added Vector<T> to replace std::vector<T>.

Author: Aliaksei Sandryhaila

src/parquet/api/io.h

@@ -21,6 +21,7 @@
 #include "parquet/exception.h"
 #include "parquet/util/buffer.h"
 #include "parquet/util/input.h"
+#include "parquet/util/mem-allocator.h"
 #include "parquet/util/output.h"
 
 #endif // PARQUET_API_IO_H

src/parquet/column/reader.cc

@@ -29,11 +29,12 @@
 namespace parquet_cpp {
 
 ColumnReader::ColumnReader(const ColumnDescriptor* descr,
-    std::unique_ptr<PageReader> pager)
+    std::unique_ptr<PageReader> pager, MemoryAllocator* allocator)
   : descr_(descr),
     pager_(std::move(pager)),
     num_buffered_values_(0),
-    num_decoded_values_(0) {}
+    num_decoded_values_(0),
+    allocator_(allocator) {}
 
 template <int TYPE>
 void TypedColumnReader<TYPE>::ConfigureDictionary(const DictionaryPage* page) {
@@ -59,7 +60,7 @@ void TypedColumnReader<TYPE>::ConfigureDictionary(const DictionaryPage* page) {
     // TODO(wesm): investigate whether this all-or-nothing decoding of the
     // dictionary makes sense and whether performance can be improved
 
-    auto decoder = std::make_shared<DictionaryDecoder<TYPE> >(descr_);
+    auto decoder = std::make_shared<DictionaryDecoder<TYPE> >(descr_, allocator_);
     decoder->SetDict(&dictionary);
     decoders_[encoding] = decoder;
   } else {
@@ -196,24 +197,26 @@ int64_t ColumnReader::ReadRepetitionLevels(int64_t batch_size, int16_t* levels)
 
 std::shared_ptr<ColumnReader> ColumnReader::Make(
     const ColumnDescriptor* descr,
-    std::unique_ptr<PageReader> pager) {
+    std::unique_ptr<PageReader> pager,
+    MemoryAllocator* allocator) {
   switch (descr->physical_type()) {
     case Type::BOOLEAN:
-      return std::make_shared<BoolReader>(descr, std::move(pager));
+      return std::make_shared<BoolReader>(descr, std::move(pager), allocator);
     case Type::INT32:
-      return std::make_shared<Int32Reader>(descr, std::move(pager));
+      return std::make_shared<Int32Reader>(descr, std::move(pager), allocator);
     case Type::INT64:
-      return std::make_shared<Int64Reader>(descr, std::move(pager));
+      return std::make_shared<Int64Reader>(descr, std::move(pager), allocator);
     case Type::INT96:
-      return std::make_shared<Int96Reader>(descr, std::move(pager));
+      return std::make_shared<Int96Reader>(descr, std::move(pager), allocator);
     case Type::FLOAT:
-      return std::make_shared<FloatReader>(descr, std::move(pager));
+      return std::make_shared<FloatReader>(descr, std::move(pager), allocator);
     case Type::DOUBLE:
-      return std::make_shared<DoubleReader>(descr, std::move(pager));
+      return std::make_shared<DoubleReader>(descr, std::move(pager), allocator);
     case Type::BYTE_ARRAY:
-      return std::make_shared<ByteArrayReader>(descr, std::move(pager));
+      return std::make_shared<ByteArrayReader>(descr, std::move(pager), allocator);
     case Type::FIXED_LEN_BYTE_ARRAY:
-      return std::make_shared<FixedLenByteArrayReader>(descr, std::move(pager));
+      return std::make_shared<FixedLenByteArrayReader>(descr,
+          std::move(pager), allocator);
     default:
       ParquetException::NYI("type reader not implemented");
   }

src/parquet/column/reader.h

@@ -30,15 +30,17 @@
 #include "parquet/exception.h"
 #include "parquet/schema/descriptor.h"
 #include "parquet/types.h"
+#include "parquet/util/mem-allocator.h"
 
 namespace parquet_cpp {
 
 class ColumnReader {
  public:
-  ColumnReader(const ColumnDescriptor*, std::unique_ptr<PageReader>);
+  ColumnReader(const ColumnDescriptor*, std::unique_ptr<PageReader>,
+      MemoryAllocator* allocator = default_allocator());
 
   static std::shared_ptr<ColumnReader> Make(const ColumnDescriptor*,
-      std::unique_ptr<PageReader>);
+      std::unique_ptr<PageReader>, MemoryAllocator* allocator = default_allocator());
 
   // Returns true if there are still values in this column.
   bool HasNext() {
@@ -95,6 +97,8 @@ class ColumnReader {
   // The number of values from the current data page that have been decoded
   // into memory
   int num_decoded_values_;
+
+  MemoryAllocator* allocator_;
 };
 
 // API to read values from a single column. This is the main client facing API.
@@ -104,8 +108,8 @@ class TypedColumnReader : public ColumnReader {
   typedef typename type_traits<TYPE>::value_type T;
 
   TypedColumnReader(const ColumnDescriptor* schema,
-      std::unique_ptr<PageReader> pager) :
-      ColumnReader(schema, std::move(pager)),
+      std::unique_ptr<PageReader> pager, MemoryAllocator* allocator) :
+      ColumnReader(schema, std::move(pager), allocator),
       current_decoder_(NULL) {
   }
 

src/parquet/column/scanner.cc

@@ -25,24 +25,25 @@
 namespace parquet_cpp {
 
 std::shared_ptr<Scanner> Scanner::Make(std::shared_ptr<ColumnReader> col_reader,
-    int64_t batch_size) {
+    int64_t batch_size, MemoryAllocator* allocator) {
   switch (col_reader->type()) {
     case Type::BOOLEAN:
-      return std::make_shared<BoolScanner>(col_reader, batch_size);
+      return std::make_shared<BoolScanner>(col_reader, batch_size, allocator);
     case Type::INT32:
-      return std::make_shared<Int32Scanner>(col_reader, batch_size);
+      return std::make_shared<Int32Scanner>(col_reader, batch_size, allocator);
     case Type::INT64:
-      return std::make_shared<Int64Scanner>(col_reader, batch_size);
+      return std::make_shared<Int64Scanner>(col_reader, batch_size, allocator);
     case Type::INT96:
-      return std::make_shared<Int96Scanner>(col_reader, batch_size);
+      return std::make_shared<Int96Scanner>(col_reader, batch_size, allocator);
     case Type::FLOAT:
-      return std::make_shared<FloatScanner>(col_reader, batch_size);
+      return std::make_shared<FloatScanner>(col_reader, batch_size, allocator);
     case Type::DOUBLE:
-      return std::make_shared<DoubleScanner>(col_reader, batch_size);
+      return std::make_shared<DoubleScanner>(col_reader, batch_size, allocator);
     case Type::BYTE_ARRAY:
-      return std::make_shared<ByteArrayScanner>(col_reader, batch_size);
+      return std::make_shared<ByteArrayScanner>(col_reader, batch_size, allocator);
     case Type::FIXED_LEN_BYTE_ARRAY:
-      return std::make_shared<FixedLenByteArrayScanner>(col_reader, batch_size);
+      return std::make_shared<FixedLenByteArrayScanner>(col_reader,
+          batch_size, allocator);
     default:
       ParquetException::NYI("type reader not implemented");
   }

src/parquet/column/scanner.h

@@ -29,6 +29,7 @@
 #include "parquet/exception.h"
 #include "parquet/schema/descriptor.h"
 #include "parquet/types.h"
+#include "parquet/util/mem-allocator.h"
 
 namespace parquet_cpp {
 
@@ -37,10 +38,12 @@ static constexpr int64_t DEFAULT_SCANNER_BATCH_SIZE = 128;
 class Scanner {
  public:
   explicit Scanner(std::shared_ptr<ColumnReader> reader,
-      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE) :
+      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE,
+      MemoryAllocator* allocator = default_allocator()) :
       batch_size_(batch_size),
       level_offset_(0),
       levels_buffered_(0),
+      value_buffer_(0, allocator),
       value_offset_(0),
       values_buffered_(0),
       reader_(reader) {
@@ -52,7 +55,8 @@ class Scanner {
   virtual ~Scanner() {}
 
   static std::shared_ptr<Scanner> Make(std::shared_ptr<ColumnReader> col_reader,
-      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE);
+      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE,
+      MemoryAllocator* allocator = default_allocator());
 
   virtual void PrintNext(std::ostream& out, int width) = 0;
 
@@ -78,7 +82,7 @@ class Scanner {
   int level_offset_;
   int levels_buffered_;
 
-  std::vector<uint8_t> value_buffer_;
+  OwnedMutableBuffer value_buffer_;
   int value_offset_;
   int64_t values_buffered_;
 
@@ -93,11 +97,12 @@ class TypedScanner : public Scanner {
   typedef typename type_traits<TYPE>::value_type T;
 
   explicit TypedScanner(std::shared_ptr<ColumnReader> reader,
-      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE) :
-      Scanner(reader, batch_size) {
+      int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE,
+      MemoryAllocator* allocator = default_allocator()) :
+      Scanner(reader, batch_size, allocator) {
     typed_reader_ = static_cast<TypedColumnReader<TYPE>*>(reader.get());
     int value_byte_size = type_traits<TYPE>::value_byte_size;
-    value_buffer_.resize(batch_size_ * value_byte_size);
+    value_buffer_.Resize(batch_size_ * value_byte_size);
     values_ = reinterpret_cast<T*>(&value_buffer_[0]);
   }
 

src/parquet/column/test-util.h

@@ -252,23 +252,23 @@ class DictionaryPageBuilder {
         if (TN == Type::FIXED_LEN_BYTE_ARRAY) {
           type_length = d->type_length();
         }
-        encoder_.reset(new DictEncoder<TC>(&pool_, type_length));
+        encoder_.reset(new DictEncoder<TC>(&pool_, default_allocator(), type_length));
   }
 
   ~DictionaryPageBuilder() {
     pool_.FreeAll();
   }
 
   shared_ptr<Buffer> AppendValues(const vector<TC>& values) {
-    shared_ptr<OwnedMutableBuffer> rle_indices = std::make_shared<OwnedMutableBuffer>();
     int num_values = values.size();
     // Dictionary encoding
     for (int i = 0; i < num_values; ++i) {
       encoder_->Put(values[i]);
     }
     num_dict_values_ = encoder_->num_entries();
     have_values_ = true;
-    rle_indices->Resize(sizeof(int) * encoder_->EstimatedDataEncodedSize());
+    shared_ptr<OwnedMutableBuffer> rle_indices = std::make_shared<OwnedMutableBuffer>(
+        sizeof(int) * encoder_->EstimatedDataEncodedSize());
     int actual_bytes = encoder_->WriteIndices(rle_indices->mutable_data(),
         rle_indices->size());
     rle_indices->Resize(actual_bytes);
@@ -277,8 +277,8 @@ class DictionaryPageBuilder {
   }
 
   shared_ptr<Buffer> WriteDict() {
-    shared_ptr<OwnedMutableBuffer> dict_buffer = std::make_shared<OwnedMutableBuffer>();
-    dict_buffer->Resize(encoder_->dict_encoded_size());
+    shared_ptr<OwnedMutableBuffer> dict_buffer = std::make_shared<OwnedMutableBuffer>(
+        encoder_->dict_encoded_size());
     encoder_->WriteDict(dict_buffer->mutable_data());
     return dict_buffer;
   }

src/parquet/encodings/decoder.h

@@ -22,6 +22,7 @@
 
 #include "parquet/exception.h"
 #include "parquet/types.h"
+#include "parquet/util/mem-allocator.h"
 
 namespace parquet_cpp {
 
@@ -54,8 +55,7 @@ class Decoder {
   const Encoding::type encoding() const { return encoding_; }
 
  protected:
-  explicit Decoder(const ColumnDescriptor* descr,
-      const Encoding::type& encoding)
+  explicit Decoder(const ColumnDescriptor* descr, const Encoding::type& encoding)
       : descr_(descr), encoding_(encoding), num_values_(0) {}
 
   // For accessing type-specific metadata, like FIXED_LEN_BYTE_ARRAY

src/parquet/encodings/delta-bit-pack-encoding.h

@@ -24,6 +24,7 @@
 
 #include "parquet/encodings/decoder.h"
 #include "parquet/util/bit-stream-utils.inline.h"
+#include "parquet/util/buffer.h"
 
 namespace parquet_cpp {
 
@@ -32,8 +33,10 @@ class DeltaBitPackDecoder : public Decoder<TYPE> {
  public:
   typedef typename type_traits<TYPE>::value_type T;
 
-  explicit DeltaBitPackDecoder(const ColumnDescriptor* descr)
-      : Decoder<TYPE>(descr, Encoding::DELTA_BINARY_PACKED) {
+  explicit DeltaBitPackDecoder(const ColumnDescriptor* descr,
+      MemoryAllocator* allocator = default_allocator())
+      : Decoder<TYPE>(descr, Encoding::DELTA_BINARY_PACKED),
+        delta_bit_widths_(0, allocator) {
     if (TYPE != Type::INT32 && TYPE != Type::INT64) {
       throw ParquetException("Delta bit pack encoding should only be for integer data.");
     }
@@ -61,7 +64,7 @@ class DeltaBitPackDecoder : public Decoder<TYPE> {
       ParquetException::EofException();
     }
     if (!decoder_.GetZigZagVlqInt(&last_value_)) ParquetException::EofException();
-    delta_bit_widths_.resize(num_mini_blocks_);
+    delta_bit_widths_.Resize(num_mini_blocks_);
 
     if (!decoder_.GetZigZagVlqInt(&min_delta_)) ParquetException::EofException();
     for (int i = 0; i < num_mini_blocks_; ++i) {
@@ -81,7 +84,7 @@ class DeltaBitPackDecoder : public Decoder<TYPE> {
     for (int i = 0; i < max_values; ++i) {
       if (UNLIKELY(values_current_mini_block_ == 0)) {
         ++mini_block_idx_;
-        if (mini_block_idx_ < delta_bit_widths_.size()) {
+        if (mini_block_idx_ < static_cast<size_t>(delta_bit_widths_.size())) {
           delta_bit_width_ = delta_bit_widths_[mini_block_idx_];
           values_current_mini_block_ = values_per_mini_block_;
         } else {
@@ -111,7 +114,7 @@ class DeltaBitPackDecoder : public Decoder<TYPE> {
 
   int32_t min_delta_;
   size_t mini_block_idx_;
-  std::vector<uint8_t> delta_bit_widths_;
+  OwnedMutableBuffer delta_bit_widths_;
   int delta_bit_width_;
 
   int32_t last_value_;

src/parquet/encodings/delta-byte-array-encoding.h

@@ -28,10 +28,11 @@ namespace parquet_cpp {
 
 class DeltaByteArrayDecoder : public Decoder<Type::BYTE_ARRAY> {
  public:
-  explicit DeltaByteArrayDecoder(const ColumnDescriptor* descr)
+  explicit DeltaByteArrayDecoder(const ColumnDescriptor* descr,
+      MemoryAllocator* allocator = default_allocator())
       : Decoder<Type::BYTE_ARRAY>(descr, Encoding::DELTA_BYTE_ARRAY),
-      prefix_len_decoder_(nullptr),
-      suffix_decoder_(nullptr) {
+      prefix_len_decoder_(nullptr, allocator),
+      suffix_decoder_(nullptr, allocator) {
   }
 
   virtual void SetData(int num_values, const uint8_t* data, int len) {

src/parquet/encodings/delta-length-byte-array-encoding.h

@@ -29,10 +29,10 @@ namespace parquet_cpp {
 
 class DeltaLengthByteArrayDecoder : public Decoder<Type::BYTE_ARRAY> {
  public:
-  explicit DeltaLengthByteArrayDecoder(const ColumnDescriptor* descr)
-      : Decoder<Type::BYTE_ARRAY>(descr,
-          Encoding::DELTA_LENGTH_BYTE_ARRAY),
-      len_decoder_(nullptr) {
+  explicit DeltaLengthByteArrayDecoder(const ColumnDescriptor* descr,
+      MemoryAllocator* allocator = default_allocator()) :
+      Decoder<Type::BYTE_ARRAY>(descr, Encoding::DELTA_LENGTH_BYTE_ARRAY),
+      len_decoder_(nullptr, allocator) {
   }
 
   virtual void SetData(int num_values, const uint8_t* data, int len) {